Merge branch 'feat-adicionar-etapas-ao-gerar-documento' of https://github.com/luanpoppe/vella-backend into tests

_utils/gerar_documento.py

@@ -53,7 +53,7 @@ async def gerar_documento(
         contextual_retriever = ContextualRetriever(serializer)
 
         # Initialize enhanced summarizer
-        summarizer = GerarDocumento(serializer)
+        summarizer = GerarDocumento(serializer, axiom_instance)
 
         all_PDFs_chunks, full_text_as_array = await get_full_text_and_all_PDFs_chunks(
             listaPDFs,

_utils/gerar_relatorio_modelo_usuario/GerarDocumento.py

@@ -1,8 +1,11 @@
+from dataclasses import dataclass
 import os
 from typing import Any, List, Dict, Literal, Tuple, Optional, Union, cast
 
 from pydantic import SecretStr
+from _utils.langchain_utils.Chain_class import Chain
 from _utils.langchain_utils.LLM_class import LLM
+from _utils.langchain_utils.Prompt_class import Prompt
 from _utils.langchain_utils.Vector_store_class import VectorStore
 from gerar_documento.serializer import (
     GerarDocumentoComPDFProprioSerializerData,
@@ -27,6 +30,8 @@ from cohere import Client
 from _utils.langchain_utils.Splitter_class import Splitter
 import time
 
+from setup.logging import Axiom
+
 
 def reciprocal_rank_fusion(result_lists, weights=None):
     """Combine multiple ranked lists using reciprocal rank fusion"""
@@ -47,16 +52,51 @@ def reciprocal_rank_fusion(result_lists, weights=None):
     return sorted_results
 
 
+@dataclass
+class GerarDocumentoUtils:
+    def criar_output_estruturado(self, summaries: List[str | Any], sources: Any):
+        structured_output = []
+        for idx, summary in enumerate(summaries):
+            source_idx = min(idx, len(sources) - 1)
+            structured_output.append(
+                {
+                    "content": summary,
+                    "source": {
+                        "page": sources[source_idx]["page"],
+                        "text": sources[source_idx]["content"][:200] + "...",
+                        "context": sources[source_idx]["context"],
+                        "relevance_score": sources[source_idx]["relevance_score"],
+                        "chunk_id": sources[source_idx]["chunk_id"],
+                    },
+                }
+            )
+
+        return structured_output
+
+    def ultima_tentativa_requisicao(self, prompt_gerar_documento_formatado):
+        llm = LLM()
+        resposta = llm.open_ai().invoke(prompt_gerar_documento_formatado)
+        documento_gerado = resposta.content.strip()  # type: ignore
+        if not documento_gerado:
+            raise Exception(
+                "Falha ao tentar gerar o documento final por 5 tentativas e também ao tentar na última tentativa com o chat-gpt 4o mini."
+            )
+        else:
+            return documento_gerado
+
+
 class GerarDocumento:
     openai_api_key = os.environ.get("OPENAI_API_KEY", "")
     cohere_api_key = os.environ.get("COHERE_API_KEY", "")
     resumo_gerado = ""
+    gerar_documento_utils = GerarDocumentoUtils()
 
     def __init__(
         self,
         serializer: Union[
             GerarDocumentoSerializerData, GerarDocumentoComPDFProprioSerializerData, Any
         ],
+        axiom_instance: Axiom,
     ):
         self.config = RetrievalConfig(
             num_chunks=serializer.num_chunks_retrieval,
@@ -77,8 +117,11 @@ class GerarDocumento:
         self.num_k_rerank = serializer.num_k_rerank
         self.model_cohere_rerank = serializer.model_cohere_rerank
         self.splitter = Splitter(serializer.chunk_size, serializer.chunk_overlap)
+        self.prompt_gerar_documento_etapa_2 = serializer.prompt_gerar_documento_etapa_2
+        self.prompt_gerar_documento_etapa_3 = serializer.prompt_gerar_documento_etapa_3
 
         self.vector_store = VectorStore(serializer.hf_embedding)
+        self.axiom_instance: Axiom = axiom_instance
 
     def retrieve_with_rank_fusion(
         self, vector_store: Chroma, bm25: BM25Okapi, chunk_ids: List[str], query: str
@@ -205,17 +248,18 @@ class GerarDocumento:
                 input_variables=["context"],
             )
 
+            llm = self.select_model_for_last_requests(llm_ultimas_requests)  # type: ignore
+            prompt_instance = Prompt()
             documento_gerado = ""
             tentativas = 0
+            context_do_prompt_primeira_etapa = "\n\n".join(contexts)
+            prompt_primeira_etapa = prompt_gerar_documento.format(
+                context=context_do_prompt_primeira_etapa,
+            )
 
             while tentativas < 5 and not documento_gerado:
                 tentativas += 1
-                llm = self.select_model_for_last_requests(llm_ultimas_requests)  # type: ignore
-                resposta = llm.invoke(
-                    prompt_gerar_documento.format(
-                        context="\n\n".join(contexts),
-                    )
-                )
+                resposta = llm.invoke(prompt_primeira_etapa)
                 if hasattr(resposta, "content") and resposta.content.strip():  # type: ignore
                     documento_gerado = resposta.content.strip()  # type: ignore
                 else:
@@ -223,44 +267,48 @@ class GerarDocumento:
                     time.sleep(5)
 
             if not documento_gerado:
-                llm = self.select_model_for_last_requests("gpt-4o-mini")
-                resposta = llm.invoke(
-                    prompt_gerar_documento.format(
-                        context="\n\n".join(contexts),
-                    )
+                self.axiom_instance.send_axiom(
+                    "TENTANDO GERAR DOCUMENTO FINAL COM GPT 4o-mini COMO ÚLTIMA TENTATIVA"
                 )
-                documento_gerado = resposta.content.strip()  # type: ignore
-                if not documento_gerado:
-                    raise Exception(
-                        "Falha ao tentar gerar o documento final por 5 tentativas e também ao tentar na última tentativa com o chat-gpt 4o mini."
+                documento_gerado = (
+                    self.gerar_documento_utils.ultima_tentativa_requisicao(
+                        prompt_primeira_etapa
                     )
+                )
 
-            # Split the response into paragraphs
-            summaries = [p.strip() for p in documento_gerado.split("\n\n") if p.strip()]
+            resposta_primeira_etapa = documento_gerado
 
-            # Create structured output
-            structured_output = []
-            for idx, summary in enumerate(summaries):
-                source_idx = min(idx, len(sources) - 1)
-                structured_output.append(
-                    {
-                        "content": summary,
-                        "source": {
-                            "page": sources[source_idx]["page"],
-                            "text": sources[source_idx]["content"][:200] + "...",
-                            "context": sources[source_idx]["context"],
-                            "relevance_score": sources[source_idx]["relevance_score"],
-                            "chunk_id": sources[source_idx]["chunk_id"],
-                        },
-                    }
+            if self.prompt_gerar_documento_etapa_2:
+                self.axiom_instance.send_axiom("GERANDO DOCUMENTO - COMEÇANDO ETAPA 2")
+                prompt_etapa_2 = prompt_instance.create_and_invoke_prompt(
+                    self.prompt_gerar_documento_etapa_2,
+                    dynamic_dict={"context": context_do_prompt_primeira_etapa},
+                )
+                documento_gerado = llm.invoke(prompt_etapa_2).content
+                resposta_segunda_etapa = documento_gerado
+                self.axiom_instance.send_axiom(f"RESULTADO ETAPA 2: {documento_gerado}")
+
+            if self.prompt_gerar_documento_etapa_3:
+                self.axiom_instance.send_axiom("GERANDO DOCUMENTO - COMEÇANDO ETAPA 3")
+                prompt_etapa_3 = prompt_instance.create_and_invoke_prompt(
+                    self.prompt_gerar_documento_etapa_3,
+                    dynamic_dict={
+                        "context": f"{resposta_primeira_etapa}\n\n{resposta_segunda_etapa}"
+                    },
                 )
+                documento_gerado = llm.invoke(prompt_etapa_3).content
+                self.axiom_instance.send_axiom(f"RESULTADO ETAPA 3: {documento_gerado}")
 
+            # Split the response into paragraphs
+            summaries = [
+                p.strip() for p in documento_gerado.split("\n\n") if p.strip()  # type: ignore
+            ]
+
+            structured_output = self.gerar_documento_utils.criar_output_estruturado(
+                summaries, sources
+            )
             return structured_output
 
         except Exception as e:
             self.logger.error(f"Error generating enhanced summary: {str(e)}")
             raise
-
-    async def validar_conteudo_documento_final(self):
-        documento_gerado = ""
-        tentativas = 0

gerar_documento/serializer.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass, field
-from typing import List, Optional
+from email.policy import default
+from typing import List, Optional, Union
 from rest_framework import serializers
 from _utils.gerar_relatorio_modelo_usuario.prompts import (
     prompt_gerar_documento,
@@ -64,12 +65,11 @@ class GerarDocumentoSerializer(GerarDocumentoInitialSerializer):
         required=False, default="version-test"
     )  # Será o valor utilizado dentro da URL da requisição pro Bubble
 
-    # prompt_auxiliar = serializers.CharField(
-    #     required=False, default=prompt_auxiliar_padrao
-    # )
     prompt_gerar_documento = serializers.CharField(
         required=False, default=prompt_gerar_documento
     )
+    prompt_gerar_documento_etapa_2 = serializers.CharField(required=False)
+    prompt_gerar_documento_etapa_3 = serializers.CharField(required=False)
     user_message = serializers.CharField(required=False, default=user_message)
     num_chunks_retrieval = serializers.IntegerField(default=20)
     embedding_weight = serializers.FloatField(default=0.5)
@@ -101,6 +101,8 @@ class GerarDocumentoSerializer(GerarDocumentoInitialSerializer):
 
 @dataclass
 class GerarDocumentoSerializerData(GerarDocumentoInitialSerializerData):
+    prompt_gerar_documento_etapa_2: Union[str, None] = None
+    prompt_gerar_documento_etapa_3: Union[str, None] = None
     files: List[FileInfoSerializerData]
     bubble_editor_version: str = "version-test"
     prompt_gerar_documento: str = ""
@@ -132,6 +134,8 @@ class GerarDocumentoComPDFProprioSerializer(GerarDocumentoInitialSerializer):
     prompt_gerar_documento = serializers.CharField(
         required=False, default=prompt_gerar_documento
     )
+    prompt_gerar_documento_etapa_2 = serializers.CharField(required=False)
+    prompt_gerar_documento_etapa_3 = serializers.CharField(required=False)
     user_message = serializers.CharField(required=False, default=user_message)
     num_chunks_retrieval = serializers.IntegerField(default=20)
     embedding_weight = serializers.FloatField(default=0.5)
@@ -158,6 +162,8 @@ class GerarDocumentoComPDFProprioSerializer(GerarDocumentoInitialSerializer):
 
 @dataclass
 class GerarDocumentoComPDFProprioSerializerData(GerarDocumentoInitialSerializerData):
+    prompt_gerar_documento_etapa_2: Union[str, None] = None
+    prompt_gerar_documento_etapa_3: Union[str, None] = None
     prompt_gerar_documento: Optional[str] = field(default=None)
     user_message: Optional[str] = field(default=None)
     num_chunks_retrieval: int = field(default=20)