Parsing code to OOP

server/app.py

@@ -1,10 +1,11 @@
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
-from inference import rag_chain
 from pydantic import BaseModel
 from fastapi.staticfiles import StaticFiles
-from data.load_data import retriever_pre
+from services.generate_questions_service import GenerateQuestionsService
+# from data.load_data import retriever_pre
 
+generate_questions_service = GenerateQuestionsService()
 
 class Body(BaseModel):
     subject: str
@@ -26,10 +27,8 @@ async def generate_questions(body: Body):
     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}."
-    res = rag_chain.invoke(f"""{query}""")
-    return {
-        "res": res,
-    }
+    res = generate_questions_service.handle(f"""{query}""")
+    return res
 
 
 app.mount("/", StaticFiles(directory="static", html=True), name="static")

server/data/load_data.py

@@ -1,75 +0,0 @@
-import os
-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 llm.gemini import gemini_embeddings, llm
-from utils.questions_parser import parse_question
-
-
-try:
-    vectorstore = Chroma(
-        persist_directory="./chroma_db", embedding_function=gemini_embeddings
-    )
-
-except Exception as e:
-
-    print(e)
-
-    if "DATA_PATH" not in os.environ:
-        raise ValueError("DATA_PATH environment variable is not set")
-
-    DATA_PATH = os.environ["DATA_PATH"]
-
-    data_loader = TextLoader(DATA_PATH, encoding="UTF-8").load()
-
-    questions = list(
-        map(lambda x: "##Questão" + x, data_loader[0].page_content.split("##Questão"))
-    )
-
-    docs = []
-
-    for question in questions:
-        try:
-            docs.append(parse_question(question))
-        except Exception as e:
-            print(e, question)
-
-    db = Chroma.from_documents(docs, gemini_embeddings)
-    vectorstore = Chroma.from_documents(
-        documents=docs, embedding=gemini_embeddings, persist_directory="./chroma_db"
-    )
-
-    vectorstore_disk = Chroma(
-        persist_directory="./chroma_db", embedding_function=gemini_embeddings
-    )
-
-
-metadata_field_info = [
-    AttributeInfo(
-        name="topico",
-        description="A materia escolar da qual a questão pertence.",
-        type="string",
-    ),
-    AttributeInfo(
-        name="assunto",
-        description="O assunto da materia fornecida anteriormente.",
-        type="string",
-    ),
-    AttributeInfo(
-        name="dificuldade",
-        description="O nivel de dificuldade para resolver a questao.",
-        type="string",
-    ),
-    AttributeInfo(
-        name="tipo",
-        description="O tipo da questao. Pode ser ou Multipla Escolha ou Justificativa",
-        type="string",
-    ),
-]
-
-document_content_description = "Questões de matérias do ensino médio."
-
-retriever = SelfQueryRetriever.from_llm(
-    llm, vectorstore, document_content_description, metadata_field_info, verbose=True
-)

server/data/retriever.py

@@ -0,0 +1,63 @@
+import os
+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 llm.gemini import Gemini
+from utils.questions_parser import parse_question
+
+class Retriever:
+
+    _model = Gemini()
+
+    def __init__(self):
+
+        if "DATA_PATH" not in os.environ:
+            raise ValueError("DATA_PATH environment variable is not set")
+
+        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 = []
+
+        for question in self.questions:
+            try:
+                self.docs.append(parse_question(question))
+            except Exception as e:
+                print(e, question)
+
+        self.vectorstore = Chroma.from_documents(self.docs, self._model.embeddings, persist_directory="./chroma_db")
+
+        self.metadata_field_info = [
+            AttributeInfo(
+                name="topico",
+                description="A materia escolar da qual a questão pertence.",
+                type="string",
+            ),
+            AttributeInfo(
+                name="assunto",
+                description="O assunto da materia fornecida anteriormente.",
+                type="string",
+            ),
+            AttributeInfo(
+                name="dificuldade",
+                description="O nivel de dificuldade para resolver a questao.",
+                type="string",
+            ),
+            AttributeInfo(
+                name="tipo",
+                description="O tipo da questao. Pode ser ou Multipla Escolha ou Justificativa",
+                type="string",
+            ),
+        ]
+
+        document_content_description = "Questões de matérias do ensino médio."
+
+        self.retriever = SelfQueryRetriever.from_llm(
+            self._model.llm, self.vectorstore, document_content_description, self.metadata_field_info, verbose=True
+        )

server/inference.py

@@ -1,43 +0,0 @@
-from langchain.schema.runnable import RunnablePassthrough
-from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain_core.runnables import RunnableLambda
-
-from llm.gemini import (
-    questions_template,
-    format_questions_instructions,
-    questions_parser,
-)
-from data.load_data import retriever
-
-
-def get_questions(_dict):
-    question = _dict["question"]
-    context = _dict["context"]
-    messages = questions_template.format_messages(
-        context=context,
-        question=question,
-        format_questions_instructions=format_questions_instructions,
-    )
-
-    tries = 0
-
-    while tries < 3:
-        try:
-            chat = ChatGoogleGenerativeAI(model="gemini-pro")
-            response = chat.invoke(messages)
-            return questions_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(docs):
-    return "\n\n".join(doc.page_content for doc in docs)
-
-
-rag_chain = {
-    "context": retriever | RunnableLambda(format_docs),
-    "question": RunnablePassthrough(),
-} | RunnableLambda(get_questions)

server/llm/gemini.py

@@ -5,78 +5,83 @@ from langchain import PromptTemplate, LLMChain
 from langchain.output_parsers import ResponseSchema, StructuredOutputParser
 from langchain.prompts import ChatPromptTemplate
 
-if "GOOGLE_API_KEY" not in os.environ:
-    raise ValueError("GOOGLE_API_KEY environment variable is not set")
 
-llm_prompt_template = """Olá, sou uma IA treinada para gerar conteúdo educacional. Por favor, gere cinco questões de múltipla escolha sobre o seguinte tema:
-Instruções para cada questão:
-- Crie uma questão clara e relevante para o tema.
-- Forneça cinco opções de resposta, rotuladas de A) a E).
-- Apenas uma das opções de resposta deve ser correta.
-- Indique a resposta correta ao final de cada questão.
+class Gemini:
 
-Exemplo de uma questão:
-Tema: Fotossíntese
+    _llm_prompt_template = """
+    Olá, sou uma IA treinada para gerar conteúdo educacional. Por favor, gere cinco questões de múltipla escolha sobre o seguinte tema:
+    Instruções para cada questão:
+    - Crie uma questão clara e relevante para o tema.
+    - Forneça cinco opções de resposta, rotuladas de A) a E).
+    - Apenas uma das opções de resposta deve ser correta.
+    - Indique a resposta correta ao final de cada questão.
 
-Questão:
-Qual é o pigmento primário responsável pela fotossíntese nas plantas?
+    Exemplo de uma questão:
+    Tema: Fotossíntese
 
-Opções de Resposta:
-A) Clorofila
-B) Hemoglobina
-C) Mioglobina
-D) Citocromo
-E) Queratina
+    Questão:
+    Qual é o pigmento primário responsável pela fotossíntese nas plantas?
 
-Resposta Correta:
-A) Clorofila
+    Opções de Resposta:
+    A) Clorofila
+    B) Hemoglobina
+    C) Mioglobina
+    D) Citocromo
+    E) Queratina
 
-Context: {context}
-Question: {question}
-Answer:
+    Resposta Correta:
+    A) Clorofila
 
+    Context: {context}
+    Question: {question}
+    Answer:
 
-{format_questions_instructions}
-"""
 
-llm_prompt = PromptTemplate.from_template(llm_prompt_template)
+    {format_questions_instructions}
+    """
 
-gemini_embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
-
-llm = ChatGoogleGenerativeAI(model="gemini-pro", temperature=0.7, top_p=1)
-
-questions_template = ChatPromptTemplate.from_template(template=llm_prompt_template)
-
-questions_chain = LLMChain(llm=llm, prompt=questions_template)
-
-questions_schema = ResponseSchema(
-    name="questions",
-    description="""Give the questions in json as an array""",
-)
-
-questions_schemas = [questions_schema]
-
-questions_parser = StructuredOutputParser.from_response_schemas(questions_schemas)
-format_questions_instructions = questions_parser.get_format_instructions()
-format_questions_instructions = """
-The output should be a markdown code snippet formatted in the following schema, including the leading and trailing "```json" and "```":
-```json
-{
-	"questions": [
+    _format_questions_instructions = """
+    The output should be a markdown code snippet formatted in the following schema, including the leading and trailing "```json" and "```":
+    ```json
     {
-      question: "Qual é o pigmento primário responsável pela fotossíntese nas plantas?",
-      options: ["A) Clorofila",
-                "B) Hemoglobina",
-                "C) Mioglobina",
-                "D) Citocromo",
-                "E) Queratina"],
-      answer: "A"
-    }
-  ]
-```
-}"""
-
-
-
-
-
+      "questions": [
+        {
+          question: "Qual é o pigmento primário responsável pela fotossíntese nas plantas?",
+          options: ["A) Clorofila",
+                    "B) Hemoglobina",
+                    "C) Mioglobina",
+                    "D) Citocromo",
+                    "E) Queratina"],
+          answer: "A"
+        }
+      ]
+    ```
+    }"""
+
+    def __init__(self):
+
+        if "GOOGLE_API_KEY" not in os.environ:
+            raise ValueError("GOOGLE_API_KEY environment variable is not set")
+
+        self.llm_prompt = PromptTemplate.from_template(self._llm_prompt_template)
+
+        self.embeddings_model = "models/embedding-001"
+        self.model = "gemini-pro"
+
+        self.embeddings = GoogleGenerativeAIEmbeddings(model=self.embeddings_model)
+        self.llm = ChatGoogleGenerativeAI(model=self.model, temperature=0.7, top_p=1)
+
+        self.template = ChatPromptTemplate.from_template(
+            template=self._llm_prompt_template
+        )
+
+        self.chain = LLMChain(llm=self.llm, prompt=self.template)
+
+        self.schemas = [
+            ResponseSchema(
+                name="questions",
+                description="""Give the questions in json as an array""",
+            )
+        ]
+
+        self.parser = StructuredOutputParser.from_response_schemas(self.schemas)  

server/services/generate_questions_service.py

@@ -0,0 +1,49 @@
+from langchain_core.runnables import RunnableLambda
+from langchain.schema.runnable import RunnablePassthrough
+from data.retriever import Retriever
+from langchain_google_genai import ChatGoogleGenerativeAI
+from llm.gemini import Gemini
+
+
+class GenerateQuestionsService:
+
+    _retrieve = Retriever()
+    _model = Gemini()
+
+    def handle(self, query: str):
+       
+        rag_chain = {
+            "context": self._retrieve.retriever | RunnableLambda(self._format_docs),
+            "question": RunnablePassthrough(),
+        } | RunnableLambda(self._get_questions)
+
+        return rag_chain.invoke(query)
+    
+
+    def _get_questions(self, _dict):
+
+      question = _dict["question"]
+      context = _dict["context"]
+      messages = self._model.template.format_messages(
+          context=context,
+          question=question,
+          format_questions_instructions=self._model._format_questions_instructions,
+      )
+
+      tries = 0
+
+      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)
+    

web/index.html

@@ -5,7 +5,7 @@
   <meta charset="UTF-8" />
   <link rel="icon" type="image/svg+xml" href="/vite.svg" />
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>Vite + React</title>
+  <title>Perguntaí</title>
   <link rel="preconnect" href="https://fonts.googleapis.com">
   <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
   <link href="https://fonts.googleapis.com/css2?family=Inter:[email protected]&display=swap" rel="stylesheet">