v1-fix

api.py

@@ -1,5 +1,6 @@
 from fastapi import FastAPI, Response, Body, Security
 from fastapi.security import APIKeyHeader
+from pydantic import BaseModel
 
 from conversation.conversation_store import ConversationStore
 from rag_langchain import LangChainRAG
@@ -15,31 +16,40 @@ rewrite_prompt_id = "first"
 default_llm = "gpt-4o 128k"
 
 
+class QModel(BaseModel):
+    q: str
+    retrieval_count: int = 10
+    temperature: str = "0.2"
+    llm: str = default_llm
+
+
+class EmoModel(BaseModel):
+    qid: str
+    helpfulness: str
+
+
 @api.get("/")
 def read_root():
     return "Empty"
 
 
 @api.post("/q")
-async def q(api_key: str = Security(api_key_header), json_body: dict = Body(...)):
+async def q(api_key: str = Security(api_key_header), json_body: QModel = Body(...)):
     # Verify the API key
     if not valid_api_key(api_key):
         return Response(status_code=401)
 
-    # Process the JSON body
-    data = json_body
-
     rag = LangChainRAG(
         config={
-            "retrieve_documents": data["retrieval_count"],
-            "temperature": data["temperature"],
+            "retrieve_documents": json_body.retrieval_count,
+            "temperature": json_body.temperature,
             "prompt_id": prompt_id,
             "check_prompt_id": check_prompt_id,
             "rewrite_prompt_id": rewrite_prompt_id
         }
     )
 
-    answer, check_result, sources = rag.rag_chain(data["q"], default_llm)
+    answer, check_result, sources = rag.rag_chain(json_body.q, json_body.llm)
 
     oid = conversation_store.save_content(
         q=q,
@@ -51,8 +61,8 @@ async def q(api_key: str = Security(api_key_header), json_body: dict = Body(...)
             "check_prompt_id": check_prompt_id,
             "rewrite_prompt_id": rewrite_prompt_id,
             "check_result": check_result,
-            "temperature": data["temperature"],
-            "retrieve_document_count": data["retrieval_count"],
+            "temperature": json_body.temperature,
+            "retrieve_document_count": json_body.retrieval_count,
         }
     )
 
@@ -67,14 +77,14 @@ async def q(api_key: str = Security(api_key_header), json_body: dict = Body(...)
 
 
 @api.post("/emo")
-async def emo(api_key: str = Security(api_key_header), json_body: dict = Body(...)):
+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"])
+    qa = conversation_store.get(json_body.qid)
     new_params = qa.params
-    new_params["user_grading"] = str(json_body["helpfulness"])
+    new_params["user_grading"] = str(json_body.helpfulness)
     conversation_store.update(
         oid=json_body["qid"],
         q=qa.conversation[0].q,

rag.py

@@ -16,7 +16,7 @@ from agent.Agent import Agent
 from agent.agents import chat_openai_llm, deepinfra_chat
 from conversation.conversation_store import ConversationStore
 from prompt.prompt_store import PromptStore
-from retrieval import retrieve, retrieve_with_rerank
+from retrieval import retrieve_with_rerank
 
 load_dotenv()
 

retrieval.py

@@ -1,13 +1,7 @@
-import datetime
-
 from langchain.retrievers import ContextualCompressionRetriever
 from langchain_cohere.rerank import CohereRerank
 from langchain_core.vectorstores import VectorStoreRetriever
 
-from emdedd.Embedding import Embedding
-from emdedd.embeddings import embed_zakonnik_prace
-from questions import questions
-
 
 def retrieve(embedding, q, retrieve_document_count):
     retriever: VectorStoreRetriever = embedding.get_vector_store().as_retriever(
@@ -68,43 +62,3 @@ def reranking_retriever(embedding, retrieve_document_count):
     #     print("    kontext: " + text.replace('\n', ' ').replace('\r', ' '))
     #
     # return context_doc
-
-
-def retrieve_test(name: str, embed_dict: dict[str, Embedding], emded: bool = False):
-    try:
-        result_file = open(name + "_retrieve_test.md", "a")
-        for embed_key, embedding in embed_dict.items():
-            if emded:
-                embed_zakonnik_prace(embedding)
-            print("--- Running on " + embed_key)
-            result_file.write("\n\n| " + embed_key + " | " + str(datetime.datetime.now()) + " |")
-            result_file.write("\n|-------|-----------|")
-            dobre: int = 0
-            for q in questions:
-                print(q)
-                context_doc = retrieve(embedding, q, 5)
-                for doc in context_doc:
-                    text = doc.page_content
-                    print("    kontext: " + text.replace('\n', ' ').replace('\r', ' '))
-                    result_file.write("\n| " + q + " | " + text.replace('\n', ' ').replace('\r', ' ') + " |")
-                    dobre = dobre + 1 if "§ 100" in text else dobre
-                    dobre = dobre + 1 if "§ 101" in text else dobre
-                    dobre = dobre + 1 if "§ 103" in text else dobre
-                    dobre = dobre + 1 if "§ 104" in text else dobre
-                    dobre = dobre + 1 if "§ 105" in text else dobre
-                    dobre = dobre + 1 if "§ 106" in text else dobre
-                    dobre = dobre + 1 if "§ 107" in text else dobre
-                    dobre = dobre + 1 if "§ 109" in text else dobre
-                    dobre = dobre + 1 if "§ 110" in text else dobre
-                    dobre = dobre + 1 if "§ 111" in text else dobre
-                    dobre = dobre + 1 if "§ 112" in text else dobre
-                    dobre = dobre + 1 if "§ 113" in text else dobre
-                    dobre = dobre + 1 if "§ 114" in text else dobre
-                    dobre = dobre + 1 if "§ 115" in text else dobre
-                    dobre = dobre + 1 if "§ 116" in text else dobre
-                    dobre = dobre + 1 if "§ 117" in text else dobre
-            result_file.write("\n| Dobre: | " + str(dobre) + " |")
-    finally:
-        result_file.write("\n\n")
-        result_file.flush()
-        result_file.close()