v1-fix

agent/Agent.py

@@ -1,5 +1,3 @@
-from langchain.embeddings import CacheBackedEmbeddings
-from langchain.storage import LocalFileStore
 from langchain_core.language_models import BaseChatModel
 
 from emdedd.Embedding import Embedding

api.py

@@ -1,6 +1,8 @@
 from fastapi import FastAPI, Response, Body, Security
 from fastapi.security import APIKeyHeader
-from pydantic import BaseModel
+from pydantic import BaseModel, model_validator
+from typing import List
+import json
 
 from conversation.conversation_store import ConversationStore
 from rag_langchain import LangChainRAG
@@ -22,37 +24,58 @@ class QModel(BaseModel):
     temperature: str = "0.2"
     llm: str = default_llm
 
+    @classmethod
+    @model_validator(mode='before')
+    def validate_to_json(cls, value):
+        if isinstance(value, str):
+            return cls(**json.loads(value))
+        return value
+
+
+class AModel(BaseModel):
+    q: str
+    a: str
+    sources: List[str]
+    oid: str
+
 
 class EmoModel(BaseModel):
     qid: str
     helpfulness: str
 
+    @classmethod
+    @model_validator(mode='before')
+    def validate_to_json(cls, value):
+        if isinstance(value, str):
+            return cls(**json.loads(value))
+        return value
+
 
 @api.get("/")
-def read_root():
+async def read_root():
     return "Empty"
 
 
-@api.post("/q")
-async def q(api_key: str = Security(api_key_header), json_body: QModel = Body(...)):
+@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)
 
     rag = LangChainRAG(
         config={
-            "retrieve_documents": json_body.retrieval_count,
-            "temperature": json_body.temperature,
+            "retrieve_documents": data.retrieval_count,
+            "temperature": data.temperature,
             "prompt_id": prompt_id,
             "check_prompt_id": check_prompt_id,
             "rewrite_prompt_id": rewrite_prompt_id
         }
     )
 
-    answer, check_result, sources = rag.rag_chain(json_body.q, json_body.llm)
+    answer, check_result, sources = rag.rag_chain(data.q, data.llm)
 
     oid = conversation_store.save_content(
-        q=q,
+        q=data.q,
         a=answer,
         sources=list(map(lambda doc: doc.page_content, sources)),
         params=
@@ -61,19 +84,17 @@ async def q(api_key: str = Security(api_key_header), json_body: QModel = Body(..
             "check_prompt_id": check_prompt_id,
             "rewrite_prompt_id": rewrite_prompt_id,
             "check_result": check_result,
-            "temperature": json_body.temperature,
-            "retrieve_document_count": json_body.retrieval_count,
+            "temperature": data.temperature,
+            "retrieve_document_count": data.retrieval_count,
         }
     )
 
-    return Response(
-        status_code=200,
-        content={
-            "response": answer,
-            "sources": list(map(lambda doc: doc.page_content, sources)),
-            "qid": oid
-        }
-    )
+    return AModel(
+                a=answer,
+                q=data.q,
+                sources=list(map(lambda doc: doc.page_content, sources)),
+                oid=oid
+            )
 
 
 @api.post("/emo")

emdedd/MongoEmbedding.py

@@ -1,12 +1,11 @@
 from dataclasses import dataclass
 
+from bson.objectid import ObjectId
 from langchain.embeddings import CacheBackedEmbeddings
-from langchain.storage import LocalFileStore
-from langchain_community.vectorstores.mongodb_atlas import MongoDBAtlasVectorSearch
 from langchain_core.embeddings import Embeddings
 from langchain_core.stores import InMemoryStore
+from langchain_mongodb import MongoDBAtlasVectorSearch
 from pymongo import MongoClient
-from bson.objectid import ObjectId
 
 from emdedd.Embedding import Embedding
 
@@ -83,9 +82,7 @@ class MongoEmbedding(Embedding):
         )
 
     def search(self, query, search_type, doc_count):
-        vector_store = self.get_vector_store()
-        retriever = vector_store.as_retriever(
+        return self.get_vector_store().as_retriever(
             search_type="similarity",
             search_kwargs={"k": doc_count}
-        )
-        return retriever.get_relevant_documents(query=query)
+        ).get_relevant_documents(query=query)

rag.py

@@ -10,7 +10,7 @@ from langchain.chains.retrieval import create_retrieval_chain
 from langchain.retrievers import MultiQueryRetriever, MergerRetriever, ContextualCompressionRetriever, EnsembleRetriever
 from langchain_cohere import CohereRerank
 from langchain_core.documents import Document
-from langchain_core.prompts import PromptTemplate
+from langchain_core.prompts import PromptTemplate, BasePromptTemplate
 
 from agent.Agent import Agent
 from agent.agents import chat_openai_llm, deepinfra_chat
@@ -23,12 +23,12 @@ load_dotenv()
 conversation_store = ConversationStore()
 prompt_store = PromptStore()
 
-grammar_check_1 = prompt_store.get_by_name("gramar_check_1")
-rewrite_hyde_1 = prompt_store.get_by_name("rewrite_hyde_1")
-rewrite_hyde_2 = prompt_store.get_by_name("rewrite_hyde_2")
-rewrite_1 = prompt_store.get_by_name("rewrite_1")
-rewrite_2 = prompt_store.get_by_name("rewrite_2")
-rewrite_hyde = prompt_store.get_by_name("rewrite_hyde")
+grammar_check_1 = prompt_store.get_by_name("gramar_check_1").text
+rewrite_hyde_1 = prompt_store.get_by_name("rewrite_hyde_1").text
+rewrite_hyde_2 = prompt_store.get_by_name("rewrite_hyde_2").text
+rewrite_1 = prompt_store.get_by_name("rewrite_1").text
+rewrite_2 = prompt_store.get_by_name("rewrite_2").text
+rewrite_hyde = prompt_store.get_by_name("rewrite_hyde").text
 
 
 def replace_nl(input: str) -> str:
@@ -52,26 +52,6 @@ def rewrite(agent: Agent, q: str, prompt: str) -> list[str]:
     return [x for x in questions if ("##" not in x and len(str(x).strip()) > 0)]
 
 
-def rag_with_rerank_check_rewrite(agent: Agent, q: str, retrieve_document_count: int, prompt: str, check_prompt: str,
-                                  rewrite_prompt: str):
-    rewritten_list: list[str] = rewrite(agent, q, rewrite_prompt)
-
-    if len(rewritten_list) == 0:
-        return "Neviem, nemám podklady!", "", ""
-
-    context_doc = retrieve_subqueries(agent, retrieve_document_count, rewritten_list)
-
-    if len(context_doc) == 0:
-        return "Neviem, nemám kontext!", "", ""
-
-    result = answer_pipeline(agent, context_doc, prompt, q)
-    answer = result["text"]
-
-    check_result = check_pipeline(answer, check_prompt, context_doc, q)
-
-    return answer, check_result, context_doc
-
-
 def rag_with_rerank_check_rewrite_hyde(agent: Agent, q: str, retrieve_document_count: int, prompt: str,
                                        check_prompt: str,
                                        rewrite_prompt: str):
@@ -120,7 +100,48 @@ def rag_chain(agent: Agent, q: str, retrieve_document_count: int, prompt: str,
             prompt=PromptTemplate(
                 input_variables=["context", "question", "actual_date"],
                 template=prompt
-            )
+            ),
+            document_prompt=PromptTemplate(input_variables=[], template="page_content")
+        )
+    ).invoke(
+        input={
+            "question": q,
+            "input": q,
+            "actual_date": datetime.date.today().isoformat()
+        }
+    )
+
+    print(result)
+
+    check_result = check_pipeline(result["answer"], check_prompt, result["context"], q)
+
+    print(check_result)
+
+    return result["answer"], check_result, result["context"]
+
+
+def vanilla_rag_chain(agent: Agent, q: str, retrieve_document_count: int, prompt: str,
+                      check_prompt: str):
+    retriever = ContextualCompressionRetriever(
+        base_compressor=(CohereRerank(
+            model="rerank-multilingual-v3.0",
+            top_n=retrieve_document_count
+        )),
+        base_retriever=(agent.embedding.get_vector_store().as_retriever(
+            search_type="similarity",
+            search_kwargs={"k": min(retrieve_document_count * 10, 500)},
+        ))
+    )
+
+    result = create_retrieval_chain(
+        retriever=retriever,
+        combine_docs_chain=create_stuff_documents_chain(
+            llm=agent.llm,
+            prompt=PromptTemplate(
+                input_variables=["context", "question", "actual_date"],
+                template=prompt
+            ),
+            document_prompt=PromptTemplate(input_variables=[], template="page_content")
         )
     ).invoke(
         input={
@@ -130,8 +151,12 @@ def rag_chain(agent: Agent, q: str, retrieve_document_count: int, prompt: str,
         }
     )
 
+    print(result)
+
     check_result = check_pipeline(result["answer"], check_prompt, result["context"], q)
 
+    print(check_result)
+
     return result["answer"], check_result, result["context"]
 
 

rag_langchain.py

@@ -4,6 +4,7 @@ from dotenv import load_dotenv
 from gptcache import Cache
 from gptcache.manager.factory import manager_factory
 from gptcache.processor.pre import get_prompt
+from langchain.globals import set_debug
 from langchain.retrievers import ContextualCompressionRetriever
 from langchain_cohere import CohereRerank, CohereEmbeddings
 from langchain_community.cache import GPTCache
@@ -19,10 +20,13 @@ from agent.agents import deepinfra_chat, \
 from emdedd.Embedding import Embedding
 from emdedd.MongoEmbedding import EmbeddingDbConnection, MongoEmbedding
 from prompt.prompt_store import PromptStore
-from rag import rag_chain
+from rag import vanilla_rag_chain, rag_chain
 
 load_dotenv()
 
+# set_verbose(True)
+set_debug(True)
+
 
 class LangChainRAG:
     embedding: Embedding
@@ -89,24 +93,23 @@ class LangChainRAG:
 
         self.retriever = ContextualCompressionRetriever(
             base_compressor=CohereRerank(model="rerank-multilingual-v3.0", top_n=os.getenv("retrieve_documents")),
-            base_retriever=self.get_vector_store_mongodb().as_retriever(
+            base_retriever=self.embedding.get_vector_store().as_retriever(
                 search_type="similarity",
                 search_kwargs={"k": config["retrieve_documents"] * 10}
             )
         )
 
-    def get_vector_store_mongodb(self):
-        return self.embedding[0].get_vector_store()
-
     def get_llms(self):
         return self.llms.keys()
 
-    def rag_chain(self, query, choice):
-        # answer, check_result, context_doc = rag_with_rerank_check_rewrite(
+    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(
-            Agent(embedding=self.embedding[0], llm=self.llms[choice]),
+            Agent(embedding=self.embedding, llm=self.llms[llm_choice]),
             query,
             self.config["retrieve_documents"],
             self.prompt_store.get_by_name(self.config["prompt_id"]).text,

requirements.txt

@@ -6,6 +6,7 @@ langchain-mistralai
 langchain-cohere
 langchain-google-genai
 langchain-together
+langchain-mongodb
 
 fitz
 pypdf