Criação langchain

requirements.txt

@@ -5,9 +5,11 @@ psycopg2-binary==2.9.10
 python-dotenv==1.1.0
 vanna==0.7.9
 chromadb
-google-generativeai==0.8.5
+google-genai
 google-cloud-aiplatform==1.96.0
 onnxruntime==1.22.0
+langchain-google-genai<1.0.0
+langchain
 
 pytest==7.4.0
 pytest-asyncio==0.21.1

src/api/controller/AskController.py

@@ -1,6 +1,8 @@
 from src.assets.pattern.singleton import SingletonMeta
 from src.api.models import Question, Response
 
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain.schema import SystemMessage, HumanMessage
 from google import genai
 from src.api.database.MyVanna import MyVanna
 
@@ -13,6 +15,15 @@ class AskController(metaclass=SingletonMeta):
     def __init__(self):
         self.client = genai.Client(api_key=GEMINI_API_KEY)
 
+        self.gen = ChatGoogleGenerativeAI(model=GEMINI_MODEL_NAME,
+            google_api_key=GEMINI_API_KEY,
+            temperature=0,
+            max_tokens=None,
+            timeout=None,
+            max_retries=2,
+            convert_system_message_to_human=True
+            )
+
         self.vn = MyVanna(config={
             'print_prompt': False, 
             'print_sql': False,
@@ -23,22 +34,31 @@ def __init__(self):
         self.vn.prepare()
 
     def ask(self, question: Question):
-        client = genai.Client(api_key=GEMINI_API_KEY)
-
-        vn = MyVanna(config={
-            'print_prompt': False, 
-            'print_sql': False,
-            'api_key': GEMINI_API_KEY,
-            'model_name': GEMINI_MODEL_NAME
-        })
 
         try:
-            sql_gerado = vn.generate_sql(question.question)
+
+            mensagem = [SystemMessage(content="""Você é um assistente em um sistema que de chat AI,
+                         seu trabalho é receber uma mensagem de um humano e tratar ela para ser processada
+                         por outra IA que possui falhas. Dentro dos tratamentos necessários estão:
+                         -Trocar datas que são dadas em formato que não sejam dias, por exemplo 3 mêses,
+                         e transformar em dias, 90 dias. Outro exemplo, 1 ano e 2 meses, trocar por 425 dias
+                         (serão considerados que os mêses separados terão 30 dias)
+                         -Quando usada a expressão "mudança" referente ao repositório, você trocara por "commit",
+                         por exemplo, "qual foi a ultima mudança feita no repositório?" será trocado por
+                         "qual foi o ultimo commit feito no repositório?)
+                         NÃO explique, NÃO confirme, NÃO dê exemplos. Apenas RETORNE a mensagem tratada.
+                         Mensagem a ser processada:"""),HumanMessage(content=question.question)]
+
+            ai_mensagem = self.gen.invoke(mensagem)
+
+            print(ai_mensagem.content)
+
+            sql_gerado = self.vn.generate_sql(ai_mensagem.content)
 
             if "SELECT" not in sql_gerado.upper():
                 return {"output": "Não consegui entender sua pergunta bem o suficiente para gerar uma resposta SQL válida."}
 
-            resultado = vn.run_sql(sql_gerado)
+            resultado = self.vn.run_sql(sql_gerado)
 
             if not resultado:
                 return {"output": "A consulta foi feita, mas não há dados correspondentes no banco."}
@@ -54,7 +74,8 @@ def ask(self, question: Question):
             Gere uma resposta clara e útil para o usuário, explicando o que o resultado significa.
             """
 
-            response = client.models.generate_content(
+
+            response = self.client.models.generate_content(
                 model=GEMINI_MODEL_NAME,
                 contents=prompt,
                 config={

src/api/database/MyVanna.py

@@ -15,12 +15,61 @@
 DB_PASSWORD = env["DB_PASSWORD"]
 DB_URL = env["DB_URL"]
 
-class MyVanna(ChromaDB_VectorStore, GoogleGeminiChat):
+class ChromaDB_VectorStoreReset(ChromaDB_VectorStore):
     def __init__(self, config=None):
         if config is None:
             config = {}
 
-        ChromaDB_VectorStore.__init__(self, config=config)
+        # Força o reset na inicialização
+        config["reset_on_init"] = config.get("reset_on_init", True)
+
+        super().__init__(config=config)
+
+        # Limpa as coleções após a inicialização padrão
+        if config["reset_on_init"]:
+            self._reset_collections()
+
+            # Recria as coleções vazias
+            collection_metadata = config.get("collection_metadata", None)
+            self.documentation_collection = self.chroma_client.get_or_create_collection(
+                name="documentation",
+                embedding_function=self.embedding_function,
+                metadata=collection_metadata,
+            )
+            self.ddl_collection = self.chroma_client.get_or_create_collection(
+                name="ddl",
+                embedding_function=self.embedding_function,
+                metadata=collection_metadata,
+            )
+            self.sql_collection = self.chroma_client.get_or_create_collection(
+                name="sql",
+                embedding_function=self.embedding_function,
+                metadata=collection_metadata,
+            )
+
+    def _reset_collections(self):
+        """Limpa todas as coleções existentes"""
+        try:
+            self.chroma_client.delete_collection("documentation")
+        except Exception:
+            pass
+
+        try:
+            self.chroma_client.delete_collection("ddl")
+        except Exception:
+            pass
+
+        try:
+            self.chroma_client.delete_collection("sql")
+        except Exception:
+            pass
+
+class MyVanna(ChromaDB_VectorStoreReset, GoogleGeminiChat):
+    def __init__(self, config=None):
+        if config is None:
+            config = {}
+
+        ChromaDB_VectorStoreReset.__init__(self, config=config)
 
         GEMINI_API_KEY = config.get('api_key')
         GEMINI_MODEL_NAME = config.get('model_name')
@@ -123,32 +172,134 @@ def prepare(self):
             password = DB_PASSWORD
         )
 
-        schema = self.get_schema()
-
-        self.train(ddl=self.schema)
+        self.train(ddl=self.get_schema())
 
         self.train(documentation="""
-        A tabela repository armazena os repositórios, identificados por um ID único e nome.
+Table: user_info
+
+    id: Bigint primary key with default value from sequence
+
+    login: Required username field (character varying)
+
+    html_url: Required profile URL field (text)
+
+Table: milestone
+
+    id: Bigint primary key with default value from sequence
+
+    repository_id: Associated repository ID (integer, required)
+
+    title: Milestone title (text, required)
+
+    description: Milestone description (text, optional)
+
+    number: Milestone number (integer, required)
+
+    state: Milestone state (character varying, required)
+
+    created_at: Creation timestamp with time zone
+
+    updated_at: Update timestamp with time zone
+
+    creator: Creator user ID (bigint, required)
+
+Table: repository
+
+    id: Integer primary key with default value from sequence
+
+    name: Repository name (character varying, required)
+
+Table: branch
+
+    id: Bigint primary key with default value from sequence
+
+    name: Branch name (character varying, required)
+
+    repository_id: Associated repository ID (integer, required)
+
+Table: issue
 
-        A tabela user contém informações dos usuários, como login e URL de perfil. É usada como referência em outras tabelas, como quem criou issues, pull requests e milestones.
+    id: Bigint primary key with default value from sequence
 
-        A tabela milestone representa marcos definidos nos repositórios, contendo título, descrição, número, estado (aberta ou fechada), datas de criação e atualização, o repositório ao qual pertence e o usuário criador.
+    title: Issue title (text, required)
 
-        A tabela issue armazena tarefas ou bugs reportados. Contém título, corpo, número, autor, repositório, milestone relacionada, datas e URL. As atribuições de usuários a uma issue são registradas em issue_assignees.
+    body: Issue body/description (text, optional)
 
-        A tabela pull_requests armazena os pull requests criados nos repositórios. Inclui título, corpo, número, estado, criador, repositório, milestone (opcional), datas e URL. Os responsáveis são registrados em pull_request_assignees.
+    number: Issue number (integer, required)
 
-        A tabela branch armazena os nomes de branches de cada repositório.
+    html_url: Issue URL (text, optional)
 
-        A tabela commits armazena cada commit feito. Cada registro contém o SHA, mensagem, autor (usuário), repositório, branch (opcional), data de criação e URL. Também há referência à tabela de usuários.
+    created_at: Creation timestamp with time zone
 
-        A tabela parents_commits representa a relação entre commits e seus pais (para commits com múltiplos ancestrais, como em merges). Usa o SHA do commit pai e o ID do commit filho.
+    updated_at: Update timestamp with time zone
 
-        A tabela issue_assignees relaciona múltiplos usuários a uma mesma issue, representando atribuições de tarefas. É uma tabela de junção entre issues e usuários.
+    created_by: Creator user ID (bigint, required)
 
-        A tabela pull_request_assignees relaciona múltiplos usuários a um pull request, permitindo registrar quem é responsável por revisar ou aprovar um PR.
+    repository_id: Associated repository ID (bigint, required)
 
-        O modelo garante integridade por meio de chaves estrangeiras, e unicidade de registros por restrições compostas (como número + repositório para issues, pull requests e milestones).
+    milestone_id: Associated milestone ID (bigint, optional)
+
+Table: pull_requests
+
+    id: Bigint primary key with default value from sequence
+
+    created_by: Creator user ID (bigint, required)
+
+    repository_id: Associated repository ID (integer, required)
+
+    number: Pull request number (integer, required)
+
+    state: Pull request state (character varying, required)
+
+    title: Pull request title (text, required)
+
+    body: Pull request body/description (text, optional)
+
+    html_url: Pull request URL (text, required)
+
+    created_at: Creation timestamp with time zone
+
+    updated_at: Update timestamp with time zone
+
+    milestone_id: Associated milestone ID (bigint, optional)
+
+Table: commits
+
+    id: Bigint primary key with default value from sequence
+
+    user_id: Author user ID (bigint, required)
+
+    branch_id: Associated branch ID (integer, optional)
+
+    pull_request_id: Associated pull request ID (bigint, optional)
+
+    created_at: Creation timestamp with time zone
+
+    message: Commit message (text, required)
+
+    sha: Commit SHA hash (character varying, required)
+
+    html_url: Commit URL (text, optional)
+
+Table: parents_commits
+
+    id: Integer primary key with default value from sequence
+
+    parent_sha: Parent commit SHA hash (character varying, required)
+
+    commit_id: Child commit ID (integer, required)
+
+Table: issue_assignees
+
+    issue_id: Issue ID (bigint, required, part of primary key)
+
+    user_id: Assigned user ID (bigint, required, part of primary key)
+
+Table: pull_request_assignees
+
+    pull_request_id: Pull request ID (bigint, required, part of primary key)
+
+    user_id: Assigned user ID (bigint, required, part of primary key)
         """)
 
         self.train(sql="""
@@ -232,3 +383,5 @@ def prepare(self):
         ORDER BY
             total_commits DESC;
         """)
+
+