retriever funcionando e novo banco de questoes

.dockerignore

@@ -41,4 +41,5 @@ next-env.d.ts
 .yarn
 
 *venv
-como_nao_errar.txt
+como_nao_errar.txt
+server/venv

.gitignore

@@ -28,3 +28,5 @@ dist-ssr
 
 *.env
 *chroma_db
+como_nao_errar.txt
+server/venv

server/app.py

@@ -3,11 +3,14 @@ from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from fastapi.staticfiles import StaticFiles
 from services.generate_questions_service import GenerateQuestionsService
+
 # from data.load_data import retriever_pre
 
 generate_questions_service = GenerateQuestionsService()
 
+
 class Body(BaseModel):
+    school_subject: str
     subject: str
     difficultie: str
 
@@ -24,9 +27,10 @@ app.add_middleware(
 
 @app.post("/generate_questions")
 async def generate_questions(body: Body):
+    school_subject = body.school_subject
     subject = body.subject
     difficultie = body.difficultie
-    query = f"Quero que você gere questões de biologia, sendo do assunto: {subject} e sendo da dificuldade: {difficultie}."
+    query = f"Quero que você gere questões de {school_subject}, sendo do assunto: {subject} e sendo da dificuldade: {difficultie}."
     res = generate_questions_service.handle(f"""{query}""")
     return res
 

server/data/retriever.py

@@ -3,9 +3,11 @@ from langchain_community.document_loaders import TextLoader
 from langchain.vectorstores import Chroma
 from langchain.chains.query_constructor.base import AttributeInfo
 from langchain.retrievers.self_query.base import SelfQueryRetriever
+from langchain_text_splitters import CharacterTextSplitter
 from llm.gemini import Gemini
 from utils.questions_parser import parse_question
 
+
 class Retriever:
 
     _model = Gemini()
@@ -17,21 +19,14 @@ class Retriever:
 
         DATA_PATH = os.environ["DATA_PATH"]
 
-        self.data_loader = TextLoader(DATA_PATH, encoding="UTF-8").load()
-
-        self.questions = list(
-            map(lambda x: "##Questão" + x, self.data_loader[0].page_content.split("##Questão"))
-        )
-
-        self.docs = []
+        data_loader = TextLoader(DATA_PATH, encoding="UTF-8").load()
 
-        for question in self.questions:
-            try:
-                self.docs.append(parse_question(question))
-            except Exception as e:
-                print(e, question)
+        text_splitter = CharacterTextSplitter(chunk_size=1024, chunk_overlap=0)
+        docs = text_splitter.split_documents(data_loader)
 
-        self.vectorstore = Chroma.from_documents(self.docs, self._model.embeddings, persist_directory="./chroma_db")
+        self.vectorstore = Chroma.from_documents(
+            docs, self._model.embeddings, persist_directory="./chroma_db"
+        )
 
         self.metadata_field_info = [
             AttributeInfo(
@@ -58,6 +53,14 @@ class Retriever:
 
         document_content_description = "Questões de matérias do ensino médio."
 
+        db = Chroma.from_documents(docs, self._model.embeddings)
+
         self.retriever = SelfQueryRetriever.from_llm(
-            self._model.llm, self.vectorstore, document_content_description, self.metadata_field_info, verbose=True
+            self._model.llm,
+            self.vectorstore,
+            document_content_description,
+            self.metadata_field_info,
+            verbose=True,
         )
+
+        self.docs_retriever = db.as_retriever()

server/llm/gemini.py

@@ -84,4 +84,4 @@ class Gemini:
             )
         ]
 
-        self.parser = StructuredOutputParser.from_response_schemas(self.schemas)  
+        self.parser = StructuredOutputParser.from_response_schemas(self.schemas)

server/services/generate_questions_service.py

@@ -11,39 +11,43 @@ class GenerateQuestionsService:
     _model = Gemini()
 
     def handle(self, query: str):
-       
+
         rag_chain = {
             "context": self._retrieve.retriever | RunnableLambda(self._format_docs),
             "question": RunnablePassthrough(),
         } | RunnableLambda(self._get_questions)
+        response_rag = self._retrieve.docs_retriever
 
-        return rag_chain.invoke(query)
-    
+        rag_result = rag_chain.invoke(query)
+        retriever_result = response_rag.invoke(query)
 
-    def _get_questions(self, _dict):
+        print("RAG result:", rag_result)
+        print("Retriever result:", retriever_result)
 
-      question = _dict["question"]
-      context = _dict["context"]
-      messages = self._model.template.format_messages(
-          context=context,
-          question=question,
-          format_questions_instructions=self._model._format_questions_instructions,
-      )
+        return {"rag_result": rag_result, "retriever_result": retriever_result}
+
+    def _get_questions(self, _dict):
 
-      tries = 0
+        question = _dict["question"]
+        context = _dict["context"]
+        messages = self._model.template.format_messages(
+            context=context,
+            question=question,
+            format_questions_instructions=self._model._format_questions_instructions,
+        )
 
-      while tries < 3:
-          try:
-              chat = ChatGoogleGenerativeAI(model="gemini-pro")
-              response = chat.invoke(messages)
-              return self._model.parser.parse(response.content)
-          except Exception as e:
-              print(e)
-              tries += 1
+        tries = 0
 
-      return "Não foi possível gerar as questões."
+        while tries < 3:
+            try:
+                chat = ChatGoogleGenerativeAI(model="gemini-pro")
+                response = chat.invoke(messages)
+                return self._model.parser.parse(response.content)
+            except Exception as e:
+                print(e)
+                tries += 1
 
+        return "Não foi possível gerar as questões."
 
     def _format_docs(self, docs):
         return "\n\n".join(doc.page_content for doc in docs)
-