add api endpoint for rating reference documents (#146)

1 files changed, 18 insertions, 1 deletions

diff --git a/gn3/api/llm.py b/gn3/api/llm.py
index 68e6acc..5fcc1f5 100644
--- a/gn3/api/llm.py
+++ b/gn3/api/llm.py
@@ -6,6 +6,8 @@ from flask import jsonify, request, Blueprint, current_app
 
 from gn3.llms.process import getGNQA
 
+from gn3.llms.process import rate_document
+
 GnQNA = Blueprint("GnQNA", __name__)
 
 
@@ -17,10 +19,11 @@ def gnqa():
 
     try:
         auth_token = current_app.config.get("FAHAMU_AUTH_TOKEN")
-        answer, refs = getGNQA(
+        task_id, answer, refs = getGNQA(
             query, auth_token)
 
         return jsonify({
+            "task_id": task_id,
             "query": query,
             "answer": answer,
             "references": refs
@@ -28,3 +31,17 @@ def gnqa():
 
     except Exception as error:
         return jsonify({"query": query, "error": "Internal server error"}), 500
+
+
+@GnQNA.route("/rating/<task_id>/<doc_id>/<int:rating>", methods=["POST"])
+def rating(task_id, doc_id, rating):
+    try:
+        results = rate_document(task_id, doc_id, rating,
+                                current_app.config.get("FAHAMU_AUTH_TOKEN"))
+        return jsonify({
+            **results,
+            "doc_id": doc_id,
+            "task_id": task_id,
+        }),
+    except Exception as error:
+        return jsonify({"error": str(error), doc_id: doc_id}), 500