async

api.py

@@ -1,4 +1,6 @@
-from bson import ObjectId
+import os
+
+from dotenv import load_dotenv
 from fastapi import FastAPI, Response, Body, Security
 from fastapi.security import APIKeyHeader
 from pydantic import BaseModel, model_validator
@@ -8,6 +10,9 @@ import json
 from conversation.conversation_store import ConversationStore
 from rag_langchain import LangChainRAG
 
+load_dotenv()
+api_keys = [os.environ["API_API_KEY"]]
+
 api = FastAPI()
 
 conversation_store = ConversationStore()
@@ -59,7 +64,6 @@ async def read_root():
 
 @api.post("/qa", response_model=AModel)
 async def qa(api_key: str = Security(api_key_header), data: QModel = Body(...)):
-    # Verify the API key
     if not valid_api_key(api_key):
         return Response(status_code=401)
 
@@ -73,7 +77,7 @@ async def qa(api_key: str = Security(api_key_header), data: QModel = Body(...)):
         }
     )
 
-    answer, check_result, sources = rag.rag_chain(data.q, data.llm)
+    answer, check_result, sources = await rag.rag_chain(data.q, data.llm)
 
     oid = conversation_store.save_content(
         q=data.q,
@@ -100,21 +104,20 @@ async def qa(api_key: str = Security(api_key_header), data: QModel = Body(...)):
 
 @api.post("/emo")
 async def emo(api_key: str = Security(api_key_header), json_body: EmoModel = Body(...)):
-    # Verify the API key
     if not valid_api_key(api_key):
         return Response(status_code=401)
 
-    qa = conversation_store.get(json_body.qid)
-    new_params = qa.params
+    conversation = conversation_store.get(json_body.qid)
+    new_params = conversation.params
     new_params["user_grading"] = str(json_body.helpfulness)
     conversation_store.update(
         oid=json_body["qid"],
-        q=qa.conversation[0].q,
-        a=qa.conversation[0].a,
-        sources=qa.conversation[0].sources,
+        q=conversation.conversation[0].q,
+        a=conversation.conversation[0].a,
+        sources=conversation.conversation[0].sources,
         params=new_params
     )
 
 
-def valid_api_key(api_key: str):
-    return api_key == "your_secret_api_key"
+def valid_api_key(api_key: str) -> bool:
+    return api_key in api_keys

rag.py

@@ -1,7 +1,7 @@
 import datetime
 import os
 import traceback
-from typing import Any
+from typing import Any, Coroutine
 
 from dotenv import load_dotenv
 from langchain.chains import LLMChain
@@ -91,9 +91,9 @@ def rag_with_rerank_check_multi_query_retriever(agent: Agent, q: str, retrieve_d
     return answer, check_result, context_doc
 
 
-def rag_chain(agent: Agent, q: str, retrieve_document_count: int, prompt: str,
-              check_prompt: str):
-    result = create_retrieval_chain(
+async def rag_chain(agent: Agent, q: str, retrieve_document_count: int, prompt: str,
+                    check_prompt: str):
+    result = await create_retrieval_chain(
         retriever=hyde_2_retrieval(agent, retrieve_document_count),
         combine_docs_chain=create_stuff_documents_chain(
             llm=agent.llm,
@@ -103,7 +103,7 @@ def rag_chain(agent: Agent, q: str, retrieve_document_count: int, prompt: str,
             ),
             document_prompt=PromptTemplate(input_variables=[], template="page_content")
         )
-    ).invoke(
+    ).ainvoke(
         input={
             "question": q,
             "input": q,
@@ -208,7 +208,7 @@ def hyde_2_retrieval(agent, retrieve_document_count):
         llm=agent.llm,
         retriever=agent.embedding.get_vector_store().as_retriever(
             search_type="similarity",
-            search_kwargs={"k": min(retrieve_document_count * 10, 500)}
+            search_kwargs={"k": min(retrieve_document_count * 10, 300)}
         ),
         prompt=PromptTemplate(
             input_variables=["question"],
@@ -224,7 +224,7 @@ def hyde_2_retrieval(agent, retrieve_document_count):
         llm=agent.llm,
         retriever=agent.embedding.get_vector_store().as_retriever(
             search_type="similarity",
-            search_kwargs={"k": min(retrieve_document_count * 10, 500)}
+            search_kwargs={"k": min(retrieve_document_count * 10, 300)}
         ),
         prompt=PromptTemplate(
             input_variables=["question"],
@@ -240,7 +240,7 @@ def hyde_2_retrieval(agent, retrieve_document_count):
         llm=agent.llm,
         retriever=agent.embedding.get_vector_store().as_retriever(
             search_type="similarity",
-            search_kwargs={"k": min(retrieve_document_count * 10, 500)}
+            search_kwargs={"k": min(retrieve_document_count * 10, 300)}
         ),
         prompt=PromptTemplate(
             input_variables=["question"],

rag_langchain.py

@@ -102,13 +102,13 @@ class LangChainRAG:
     def get_llms(self):
         return self.llms.keys()
 
-    def rag_chain(self, query, llm_choice):
+    async def rag_chain(self, query, llm_choice):
         print("Using " + llm_choice)
 
         # answer, check_result, context_doc = rag_with_rerank_check_rewrite_hyde(
         # answer, check_result, context_doc = rag_with_rerank_check_multi_query_retriever(
         # answer, check_result, context_doc = vanilla_rag_chain(
-        answer, check_result, context_doc = rag_chain(
+        answer, check_result, context_doc = await rag_chain(
             Agent(embedding=self.embedding, llm=self.llms[llm_choice]),
             query,
             self.config["retrieve_documents"],