feat: implement JWT authentication and token management in frontend and backend

Author: siiddhantt

application/api/answer/routes.py

@@ -122,6 +122,7 @@ def save_conversation(
     source_log_docs,
     tool_calls,
     llm,
+    decoded_token,
     index=None,
     api_key=None,
 ):
@@ -180,7 +181,7 @@ def save_conversation(
 
         completion = llm.gen(model=gpt_model, messages=messages_summary, max_tokens=30)
         conversation_data = {
-            "user": "local",
+            "user": decoded_token.get("sub"),
             "date": datetime.datetime.utcnow(),
             "name": completion,
             "queries": [
@@ -221,6 +222,7 @@ def complete_stream(
     retriever,
     conversation_id,
     user_api_key,
+    decoded_token,
     isNoneDoc=False,
     index=None,
     should_save_conversation=True,
@@ -271,6 +273,7 @@ def complete_stream(
                 source_log_docs,
                 tool_calls,
                 llm,
+                decoded_token,
                 index,
                 api_key=user_api_key,
             )
@@ -286,7 +289,7 @@ def complete_stream(
             {
                 "action": "stream_answer",
                 "level": "info",
-                "user": "local",
+                "user": decoded_token.get("sub"),
                 "api_key": user_api_key,
                 "question": question,
                 "response": response_full,
@@ -381,15 +384,21 @@ def post(self):
                 source = {"active_docs": data_key.get("source")}
                 retriever_name = data_key.get("retriever", retriever_name)
                 user_api_key = data["api_key"]
+                decoded_token = {"sub": data_key.get("user")}
 
             elif "active_docs" in data:
                 source = {"active_docs": data["active_docs"]}
                 retriever_name = get_retriever(data["active_docs"]) or retriever_name
                 user_api_key = None
+                decoded_token = request.decoded_token
 
             else:
                 source = {}
                 user_api_key = None
+                decoded_token = request.decoded_token
+
+            if not decoded_token:
+                return bad_request(401, "Unauthorized")
 
             logger.info(
                 f"/stream - request_data: {data}, source: {source}",
@@ -429,6 +438,7 @@ def post(self):
                     retriever=retriever,
                     conversation_id=conversation_id,
                     user_api_key=user_api_key,
+                    decoded_token=decoded_token,
                     isNoneDoc=data.get("isNoneDoc"),
                     index=index,
                     should_save_conversation=save_conv,
@@ -521,13 +531,21 @@ def post(self):
                 source = {"active_docs": data_key.get("source")}
                 retriever_name = data_key.get("retriever", retriever_name)
                 user_api_key = data["api_key"]
+                decoded_token = {"sub": data_key.get("user")}
+
             elif "active_docs" in data:
                 source = {"active_docs": data["active_docs"]}
                 retriever_name = get_retriever(data["active_docs"]) or retriever_name
                 user_api_key = None
+                decoded_token = request.decoded_token
+
             else:
                 source = {}
                 user_api_key = None
+                decoded_token = request.decoded_token
+
+            if not decoded_token:
+                return bad_request(401, "Unauthorized")
 
             prompt = get_prompt(prompt_id)
 
@@ -614,6 +632,7 @@ def post(self):
                     source_log_docs,
                     tool_calls,
                     llm,
+                    decoded_token,
                     api_key=user_api_key,
                 )
             )
@@ -623,7 +642,7 @@ def post(self):
                 {
                     "action": "api_answer",
                     "level": "info",
-                    "user": "local",
+                    "user": decoded_token.get("sub"),
                     "api_key": user_api_key,
                     "question": question,
                     "response": response_full,
@@ -692,12 +711,17 @@ def post(self):
                 chunks = int(data_key.get("chunks", 2))
                 source = {"active_docs": data_key.get("source")}
                 user_api_key = data["api_key"]
+                decoded_token = {"sub": data_key.get("user")}
+
             elif "active_docs" in data:
                 source = {"active_docs": data["active_docs"]}
                 user_api_key = None
+                decoded_token = request.decoded_token
+
             else:
                 source = {}
                 user_api_key = None
+                decoded_token = request.decoded_token
 
             logger.info(
                 f"/api/answer - request_data: {data}, source: {source}",
@@ -723,7 +747,7 @@ def post(self):
                 {
                     "action": "api_search",
                     "level": "info",
-                    "user": "local",
+                    "user": decoded_token.get("sub"),
                     "api_key": user_api_key,
                     "question": question,
                     "sources": docs,

application/api/user/routes.py

@@ -15,7 +15,6 @@
 from application.agents.tools.tool_manager import ToolManager
 
 from application.api.user.tasks import ingest, ingest_remote
-
 from application.core.mongo_db import MongoDB
 from application.core.settings import settings
 from application.extensions import api
@@ -110,11 +109,18 @@ class GetConversations(Resource):
         description="Retrieve a list of the latest 30 conversations (excluding API key conversations)",
     )
     def get(self):
+        decoded_token = request.decoded_token
+        if not decoded_token:
+            return make_response(jsonify({"success": False}), 401)
         try:
-            conversations = conversations_collection.find(
-                {"api_key": {"$exists": False}}
-            ).sort("date", -1).limit(30)
-            
+            conversations = (
+                conversations_collection.find(
+                    {"api_key": {"$exists": False}, "user": decoded_token.get("sub")}
+                )
+                .sort("date", -1)
+                .limit(30)
+            )
+
             list_conversations = [
                 {"id": str(conversation["_id"]), "name": conversation["name"]}
                 for conversation in conversations
@@ -132,6 +138,9 @@ class GetSingleConversation(Resource):
         params={"id": "The conversation ID"},
     )
     def get(self):
+        decoded_token = request.decoded_token
+        if not decoded_token:
+            return make_response(jsonify({"success": False}), 401)
         conversation_id = request.args.get("id")
         if not conversation_id:
             return make_response(
@@ -140,7 +149,7 @@ def get(self):
 
         try:
             conversation = conversations_collection.find_one(
-                {"_id": ObjectId(conversation_id)}
+                {"_id": ObjectId(conversation_id), "user": decoded_token.get("sub")}
             )
             if not conversation:
                 return make_response(jsonify({"status": "not found"}), 404)
@@ -227,7 +236,7 @@ def post(self):
                     {
                         "$unset": {
                             f"queries.{data['question_index']}.feedback": "",
-                            f"queries.{data['question_index']}.feedback_timestamp": ""
+                            f"queries.{data['question_index']}.feedback_timestamp": "",
                         }
                     },
                 )
@@ -240,8 +249,12 @@ def post(self):
                     },
                     {
                         "$set": {
-                            f"queries.{data['question_index']}.feedback": data["feedback"],
-                            f"queries.{data['question_index']}.feedback_timestamp": datetime.datetime.now(datetime.timezone.utc)
+                            f"queries.{data['question_index']}.feedback": data[
+                                "feedback"
+                            ],
+                            f"queries.{data['question_index']}.feedback_timestamp": datetime.datetime.now(
+                                datetime.timezone.utc
+                            ),
                         }
                     },
                 )
@@ -1211,7 +1224,13 @@ class GetMessageAnalytics(Resource):
                 required=False,
                 description="Filter option for analytics",
                 default="last_30_days",
-                enum=["last_hour", "last_24_hour", "last_7_days", "last_15_days", "last_30_days"],
+                enum=[
+                    "last_hour",
+                    "last_24_hour",
+                    "last_7_days",
+                    "last_15_days",
+                    "last_30_days",
+                ],
             ),
         },
     )
@@ -1244,35 +1263,30 @@ def post(self):
         else:
             if filter_option in ["last_7_days", "last_15_days", "last_30_days"]:
                 filter_days = (
-                    6 if filter_option == "last_7_days"
-                    else 14 if filter_option == "last_15_days"
-                    else 29
+                    6
+                    if filter_option == "last_7_days"
+                    else 14 if filter_option == "last_15_days" else 29
                 )
             else:
                 return make_response(
                     jsonify({"success": False, "message": "Invalid option"}), 400
                 )
             start_date = end_date - datetime.timedelta(days=filter_days)
             start_date = start_date.replace(hour=0, minute=0, second=0, microsecond=0)
-            end_date = end_date.replace(hour=23, minute=59, second=59, microsecond=999999)
+            end_date = end_date.replace(
+                hour=23, minute=59, second=59, microsecond=999999
+            )
             group_format = "%Y-%m-%d"
 
         try:
             pipeline = [
                 # Initial match for API key if provided
-                {
-                    "$match": {
-                        "api_key": api_key if api_key else {"$exists": False}
-                    }
-                },
+                {"$match": {"api_key": api_key if api_key else {"$exists": False}}},
                 {"$unwind": "$queries"},
                 # Match queries within the time range
                 {
                     "$match": {
-                        "queries.timestamp": {
-                            "$gte": start_date,
-                            "$lte": end_date
-                        }
+                        "queries.timestamp": {"$gte": start_date, "$lte": end_date}
                     }
                 },
                 # Group by formatted timestamp
@@ -1281,14 +1295,14 @@ def post(self):
                         "_id": {
                             "$dateToString": {
                                 "format": group_format,
-                                "date": "$queries.timestamp"
+                                "date": "$queries.timestamp",
                             }
                         },
-                        "count": {"$sum": 1}
+                        "count": {"$sum": 1},
                     }
                 },
                 # Sort by timestamp
-                {"$sort": {"_id": 1}}
+                {"$sort": {"_id": 1}},
             ]
 
             message_data = conversations_collection.aggregate(pipeline)
@@ -1511,11 +1525,21 @@ def post(self):
         if filter_option == "last_hour":
             start_date = end_date - datetime.timedelta(hours=1)
             group_format = "%Y-%m-%d %H:%M:00"
-            date_field = {"$dateToString": {"format": group_format, "date": "$queries.feedback_timestamp"}}
+            date_field = {
+                "$dateToString": {
+                    "format": group_format,
+                    "date": "$queries.feedback_timestamp",
+                }
+            }
         elif filter_option == "last_24_hour":
             start_date = end_date - datetime.timedelta(hours=24)
             group_format = "%Y-%m-%d %H:00"
-            date_field = {"$dateToString": {"format": group_format, "date": "$queries.feedback_timestamp"}}
+            date_field = {
+                "$dateToString": {
+                    "format": group_format,
+                    "date": "$queries.feedback_timestamp",
+                }
+            }
         else:
             if filter_option in ["last_7_days", "last_15_days", "last_30_days"]:
                 filter_days = (
@@ -1533,13 +1557,21 @@ def post(self):
                 hour=23, minute=59, second=59, microsecond=999999
             )
             group_format = "%Y-%m-%d"
-            date_field = {"$dateToString": {"format": group_format, "date": "$queries.feedback_timestamp"}}
+            date_field = {
+                "$dateToString": {
+                    "format": group_format,
+                    "date": "$queries.feedback_timestamp",
+                }
+            }
 
         try:
             match_stage = {
                 "$match": {
-                    "queries.feedback_timestamp": {"$gte": start_date, "$lte": end_date},
-                    "queries.feedback": {"$exists": True}
+                    "queries.feedback_timestamp": {
+                        "$gte": start_date,
+                        "$lte": end_date,
+                    },
+                    "queries.feedback": {"$exists": True},
                 }
             }
             if api_key: