From dc4078d744f7b94bac6edefd48dcbcc9098a82c9 Mon Sep 17 00:00:00 2001
From: ManishMadan2882 <manishmadan321@gmail.com>
Date: Sun, 11 Aug 2024 21:26:30 +0530
Subject: [PATCH] migration(fixes): retriver/sharing endpoints

---
 application/api/answer/routes.py | 16 +++++++++-----
 application/api/user/routes.py   | 38 ++++++++++++++++++--------------
 2 files changed, 32 insertions(+), 22 deletions(-)

diff --git a/application/api/answer/routes.py b/application/api/answer/routes.py
index 85cc3afd..e2d7b6e8 100644
--- a/application/api/answer/routes.py
+++ b/application/api/answer/routes.py
@@ -82,9 +82,11 @@ def get_data_from_api_key(api_key):
 
     if "source" in data and isinstance(data["source"], DBRef):
         source_doc = db.dereference(data["source"])
-        data["source"] = str(source_doc._id)
+        data["source"] = str(source_doc["_id"])
         if "retriever" in source_doc:
             data["retriever"] = source_doc["retriever"]
+    else:
+        data["source"] = {}
     return data
 
 
@@ -357,10 +359,14 @@ def api_answer():
             data_key = get_data_from_api_key(data["api_key"])
             chunks = int(data_key["chunks"])
             prompt_id = data_key["prompt_id"]
-            source = data_key["source"]
+            source = {"active_docs": data_key["source"]}
+            retriever_name = data_key["retriever"]
             user_api_key = data["api_key"]
         elif "active_docs" in data:
-            source = data["active_docs"]
+            source = {"active_docs":data["active_docs"]}
+            user_api_key = None
+        else:
+            source = {}
             user_api_key = None
 
         prompt = get_prompt(prompt_id)
@@ -411,10 +417,10 @@ def api_search():
     if "api_key" in data:
         data_key = get_data_from_api_key(data["api_key"])
         chunks = int(data_key["chunks"])
-        source = data_key["source"]
+        source = {"active_docs":data_key["source"]}
         user_api_key = data_key["api_key"]
     elif "active_docs" in data:
-        source = data["active_docs"]
+        source = {"active_docs":data["active_docs"]}
         user_api_key = None
     else:
         source = {}
diff --git a/application/api/user/routes.py b/application/api/user/routes.py
index 7ce0b2e2..84831a65 100644
--- a/application/api/user/routes.py
+++ b/application/api/user/routes.py
@@ -489,26 +489,31 @@ def share_conversation():
         isPromptable = request.args.get("isPromptable").lower() == "true"
 
         conversation = conversations_collection.find_one({"_id": ObjectId(conversation_id)})
+        if(conversation is None):
+            raise Exception("Conversation does not exist")
         current_n_queries = len(conversation["queries"])
 
         ##generate binary representation of uuid
         explicit_binary = Binary.from_uuid(uuid.uuid4(), UuidRepresentation.STANDARD)
 
         if isPromptable:
-            source = "default" if "source" not in data else data["source"]
             prompt_id = "default" if "prompt_id" not in data else data["prompt_id"]
             chunks = "2" if "chunks" not in data else data["chunks"]
 
             name = conversation["name"] + "(shared)"
-            pre_existing_api_document = api_key_collection.find_one(
-                {
+            new_api_key_data =  {
                     "prompt_id": prompt_id,
                     "chunks": chunks,
-                    "source": DBRef("vectors", ObjectId(source)) if ObjectId.is_valid(source) else source,
                     "user": user,
                 }
+            if "source" in data and ObjectId.is_valid(data["source"]):
+                new_api_key_data["source"] = DBRef("vectors",ObjectId(data["source"]))
+            elif "retriever" in data:
+                new_api_key_data["retriever"] = data["retriever"]
+                 
+            pre_existing_api_document = api_key_collection.find_one(
+                new_api_key_data
             )
-            api_uuid = str(uuid.uuid4())
             if pre_existing_api_document:
                 api_uuid = pre_existing_api_document["key"]
                 pre_existing = shared_conversations_collections.find_one(
@@ -546,17 +551,16 @@ def share_conversation():
                     )
                     return jsonify({"success": True, "identifier": str(explicit_binary.as_uuid())})
             else:
-                api_key_collection.insert_one(
-                    {
-                        "name": name,
-                        "key": api_uuid,
-                        "source": DBRef("vectors", ObjectId(source)) if ObjectId.is_valid(source) else source,
-                        "user": user,
-                        "prompt_id": prompt_id,
-                        "chunks": chunks,
-                    }
-                )
-            shared_conversations_collections.insert_one(
+                
+                api_uuid = str(uuid.uuid4())
+                new_api_key_data["key"] = api_uuid
+                new_api_key_data["name"] = name
+                if "source" in data and ObjectId.is_valid(data["source"]):
+                    new_api_key_data["source"] = DBRef("vectors", ObjectId(data["source"]))
+                if "retriever" in data:
+                    new_api_key_data["retriever"] = data["retriever"]
+                api_key_collection.insert_one(new_api_key_data)
+                shared_conversations_collections.insert_one(
                 {
                     "uuid": explicit_binary,
                     "conversation_id": {
@@ -568,7 +572,7 @@ def share_conversation():
                     "user": user,
                     "api_key": api_uuid,
                 }
-            )
+              )
             ## Identifier as route parameter in frontend
             return (
                 jsonify({"success": True, "identifier": str(explicit_binary.as_uuid())}),