migration(fixes): retriver/sharing endpoints

This commit is contained in:
ManishMadan2882
2024-08-11 21:26:30 +05:30
parent 1eb168be55
commit dc4078d744
2 changed files with 32 additions and 22 deletions

View File

@@ -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 = {}

View File

@@ -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())}),