refactor: melhorando a função gerar_documento e a classe GerarDocumento

_utils/Utils_Class.py

@@ -8,6 +8,3 @@ from _utils.Handle_Files_Class import HandleFilesClass
 class UtilsClass:
     lista_pdfs: List[str] | None = None
     handle_files = HandleFilesClass()
-
-    def configure_lista_pdfs(self, lista_pdfs):
-        self.lista_pdfs = lista_pdfs

_utils/axiom_logs.py

@@ -0,0 +1,27 @@
+from dataclasses import dataclass
+
+from setup.logging import Axiom
+
+
+@dataclass
+class AxiomLogs:
+    axiom: Axiom
+
+    def texto_completo_pdf(self, full_text_as_array):
+        self.axiom.send_axiom(
+            f"INÍCIO DO TEXTO COMPLETO DOS PDFS: {full_text_as_array[0:5]}"
+        )
+
+    def resumo_inicial_processo(self, response_auxiliar_summary):
+        self.axiom.send_axiom(
+            f"RESUMO INICIAL DO PROCESSO: {response_auxiliar_summary}"
+        )
+
+    def inicio_requisicao_contextual(self):
+        self.axiom.send_axiom("COMEÇANDO A FAZER AS REQUISIÇÕES DO CONTEXTUAL")
+
+    def fim_requisicao_contextual(self):
+        self.axiom.send_axiom("TERMINOU DE FAZER TODAS AS REQUISIÇÕES DO CONTEXTUAL")
+
+    def chunks_inicialmente(self, chunks_processados):
+        self.axiom.send_axiom(f"CHUNKS PROCESSADOS INICIALMENTE: {chunks_processados}")

_utils/custom_exception_handler.py

@@ -44,7 +44,7 @@ def custom_exception_handler(exc, context):
     return response
 
 
-def custom_exception_handler_wihout_api_handler(
+def custom_exception_handler_without_api_handler(
     error, serializer: Union[GerarDocumentoSerializerData, Any], axiom_instace: Axiom
 ):
     bahia_tz = pytz.timezone("America/Bahia")

_utils/gerar_documento.py

@@ -2,12 +2,11 @@ import os
 from langchain_core.messages import HumanMessage
 from typing import Any, Union, cast
 from _utils.Utils_Class import UtilsClass
+from _utils.axiom_logs import AxiomLogs
 from _utils.langchain_utils.LLM_class import LLM
 from _utils.bubble_integrations.enviar_resposta_final import enviar_resposta_final
-from _utils.custom_exception_handler import custom_exception_handler_wihout_api_handler
-from _utils.gerar_documento_utils.prompts import (
-    prompt_gerar_query_dinamicamente,
-)
+from _utils.custom_exception_handler import custom_exception_handler_without_api_handler
+from rest_framework.response import Response
 from _utils.gerar_documento_utils.GerarDocumento import (
     GerarDocumento,
 )
@@ -50,120 +49,44 @@ async def gerar_documento(
 ):
     try:
         axiom = axiom_instance.send_axiom
+        ax = AxiomLogs(axiom_instance)
         utils = UtilsClass()
-        handle_files = utils.handle_files
-        contextual_retriever = ContextualRetriever(serializer)
-        summarizer = GerarDocumento(serializer, axiom_instance)
-
-        params = {
-            "listaPDFs": listaPDFs,
-            "splitterObject": summarizer.splitter,
-            "should_use_llama_parse": serializer.should_use_llama_parse,
-            "isBubble": isBubble,
-        }
+        summarizer = GerarDocumento(serializer, isBubble, axiom_instance)
 
-        all_PDFs_chunks, full_text_as_array = (
-            await handle_files.get_full_text_and_all_PDFs_chunks(**params)
-        )
-        axiom(f"INÍCIO DO TEXTO COMPLETO DOS PDFS: {full_text_as_array[0:5]}")
+        all_PDFs_chunks, full_text_as_array = await summarizer.get_text_and_pdf_chunks()
 
         is_contextualized_chunk = serializer.should_have_contextual_chunks
 
-        if is_contextualized_chunk:
-            response_auxiliar_summary = (
-                await get_response_from_auxiliar_contextual_prompt(full_text_as_array)
-            )
-            axiom(f"RESUMO INICIAL DO PROCESSO: {response_auxiliar_summary}")
-
-            axiom("COMEÇANDO A FAZER AS REQUISIÇÕES DO CONTEXTUAL")
-            contextualized_chunks = await contextual_retriever.contextualize_all_chunks(
-                all_PDFs_chunks, response_auxiliar_summary, axiom_instance
-            )
-            axiom("TERMINOU DE FAZER TODAS AS REQUISIÇÕES DO CONTEXTUAL")
-            chunks_processados = contextualized_chunks
-            axiom(f"CHUNKS PROCESSADOS INICIALMENTE: {chunks_processados}")
-        else:
-            chunks_processados = all_PDFs_chunks
-
-        llm = LLM()
-        prompt_para_gerar_query_dinamico = prompt_gerar_query_dinamicamente(
-            cast(str, response_auxiliar_summary)
+        response_auxiliar_summary = await get_response_from_auxiliar_contextual_prompt(
+            full_text_as_array
         )
+        summarizer.resumo_auxiliar = response_auxiliar_summary
+        ax.resumo_inicial_processo(response_auxiliar_summary)
 
-        axiom("COMEÇANDO REQUISIÇÃO PARA GERAR O QUERY DINAMICAMENTE DO VECTOR STORE")
-        query_gerado_dinamicamente_para_o_vector_store = (
-            await llm.google_gemini_ainvoke(
-                prompt_para_gerar_query_dinamico, "gemini-2.0-flash"
-            )
-        )
+        await summarizer.generate_chunks_processados()
 
-        axiom(
-            f"query_gerado_dinamicamente_para_o_vector_store: {query_gerado_dinamicamente_para_o_vector_store.content}",
-        )
+        await summarizer.generate_query_for_vector_store()
 
-        # Create enhanced vector store and BM25 index
-        vector_store, bm25, chunk_ids = (
-            summarizer.vector_store.create_enhanced_vector_store(
-                chunks_processados, is_contextualized_chunk, axiom_instance
-            )
-        )
+        await summarizer.create_enhanced_vector_store()
 
-        llm_ultimas_requests = serializer.llm_ultimas_requests
-        axiom("COMEÇANDO A FAZER ÚLTIMA REQUISIÇÃO")
-        structured_summaries = await summarizer.gerar_documento_final(
-            vector_store,
-            bm25,
-            chunk_ids,
-            llm_ultimas_requests,
-            cast(
-                str, query_gerado_dinamicamente_para_o_vector_store.content
-            ),  # prompt_auxiliar_SEM_CONTEXT,
-        )
-        axiom("TERMINOU DE FAZER A ÚLTIMA REQUISIÇÃO")
+        structured_summaries = await summarizer.do_last_requests()
 
         if not isinstance(structured_summaries, list):
-            from rest_framework.response import Response
-
             return Response({"erro": structured_summaries})
 
-        texto_completo = summarizer.resumo_gerado + "\n\n"
-
-        for x in structured_summaries:
-            texto_completo = texto_completo + x["content"] + "\n"
-            x["source"]["text"] = x["source"]["text"][0:200]
-            x["source"]["context"] = x["source"]["context"][0:200]
+        await summarizer.generate_complete_text()
 
-        texto_completo_como_html = convert_markdown_to_HTML(texto_completo).replace(
-            "resposta_segunda_etapa:", "<br><br>"
-        )
-        axiom(f"texto_completo_como_html: {texto_completo_como_html}")
-
-        if is_contextualized_chunk:
-            prompt_titulo_do_documento = response_auxiliar_summary
-        else:
-            prompt_titulo_do_documento = texto_completo_como_html
-        titulo_do_documento = await generate_document_title(
-            cast(str, prompt_titulo_do_documento)
-        )
+        await summarizer.get_document_title()
 
         if isBubble:
-            axiom("COMEÇANDO A REQUISIÇÃO FINAL PARA O BUBBLE")
-            enviar_resposta_final(
-                serializer.doc_id,  # type: ignore
-                serializer.form_response_id,  # type: ignore
-                serializer.version,  # type: ignore
-                texto_completo_como_html,
-                False,
-                cast(str, titulo_do_documento),
-            )
-            axiom("TERMINOU A REQUISIÇÃO FINAL PARA O BUBBLE")
+            await summarizer.send_to_bubble()
 
         return {
-            "texto_completo": texto_completo_como_html,
-            "titulo_do_documento": titulo_do_documento,
+            "texto_completo": summarizer.texto_completo_como_html,
+            "titulo_do_documento": summarizer.titulo_do_documento,
             "resultado": structured_summaries,
             "parametros-utilizados": gerar_resposta_compilada(serializer),
         }
     except Exception as e:
-        custom_exception_handler_wihout_api_handler(e, serializer, axiom_instance)
+        custom_exception_handler_without_api_handler(e, serializer, axiom_instance)
         raise

_utils/gerar_documento_utils/GerarDocumento.py

@@ -3,10 +3,20 @@ import os
 from typing import Any, List, Dict, Literal, Tuple, Optional, Union, cast
 
 from pydantic import SecretStr
+from _utils.Utils_Class import UtilsClass
+from _utils.axiom_logs import AxiomLogs
+from _utils.bubble_integrations.enviar_resposta_final import enviar_resposta_final
+from _utils.gerar_documento_utils.contextual_retriever import ContextualRetriever
+from _utils.gerar_documento_utils.llm_calls import agemini_answer
+from _utils.gerar_documento_utils.prompts import (
+    prompt_gerar_query_dinamicamente,
+    prompt_para_gerar_titulo,
+)
 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 _utils.utils import convert_markdown_to_HTML
 from gerar_documento.serializer import (
     GerarDocumentoComPDFProprioSerializerData,
     GerarDocumentoSerializerData,
@@ -21,6 +31,8 @@ from setup.easy_imports import (
 )
 import logging
 from _utils.models.gerar_documento import (
+    ContextualizedChunk,
+    DocumentChunk,
     RetrievalConfig,
 )
 from cohere import Client
@@ -97,14 +109,28 @@ class GerarDocumentoUtils:
 
 
 class GerarDocumento:
-    resumo_gerado = ""
+    lista_pdfs: List[str]
+    should_use_llama_parse: bool
+    all_PDFs_chunks: List[DocumentChunk]
+    full_text_as_array: List[str]
+    isBubble: bool
+    chunks_processados: List[ContextualizedChunk] | List[DocumentChunk]
+    resumo_auxiliar: str
     gerar_documento_utils = GerarDocumentoUtils()
+    utils = UtilsClass()
+    llm = LLM()
+    enhanced_vector_store: tuple[Chroma, BM25Okapi, List[str]]
+    query_gerado_dinamicamente_para_o_vector_store: str
+    structured_output: List[Any]
+    texto_completo_como_html: str
+    titulo_do_documento: str
 
     def __init__(
         self,
         serializer: Union[
             GerarDocumentoSerializerData, GerarDocumentoComPDFProprioSerializerData, Any
         ],
+        isBubble: bool,
         axiom_instance: Axiom,
     ):
         self.config = self.gerar_documento_utils.create_retrieval_config(serializer)
@@ -113,6 +139,11 @@ class GerarDocumento:
         self.gpt_model = serializer.model
         self.gpt_temperature = serializer.gpt_temperature
         self.prompt_gerar_documento = serializer.prompt_gerar_documento
+        self.should_use_llama_parse = serializer.should_use_llama_parse
+        self.isBubble = isBubble
+        self.is_contextualized_chunk = serializer.should_have_contextual_chunks
+        self.contextual_retriever = ContextualRetriever(serializer)
+        self.llm_ultimas_requests = serializer.llm_ultimas_requests
 
         self.cohere_client = Client(cohere_api_key)
         self.embeddings = HuggingFaceEmbeddings(model_name=serializer.hf_embedding)
@@ -124,6 +155,71 @@ class GerarDocumento:
 
         self.vector_store = VectorStore(serializer.hf_embedding)
         self.axiom_instance: Axiom = axiom_instance
+        self.ax = AxiomLogs(axiom_instance)
+
+    async def get_text_and_pdf_chunks(self):
+        all_PDFs_chunks, full_text_as_array = (
+            await self.utils.handle_files.get_full_text_and_all_PDFs_chunks(
+                self.lista_pdfs,
+                self.splitter,
+                self.should_use_llama_parse,
+                self.isBubble,
+            )
+        )
+        self.ax.texto_completo_pdf(full_text_as_array)
+
+        self.all_PDFs_chunks = all_PDFs_chunks
+        self.full_text_as_array = full_text_as_array
+        return all_PDFs_chunks, full_text_as_array
+
+    async def generate_chunks_processados(self):
+        if self.is_contextualized_chunk:
+            self.ax.inicio_requisicao_contextual()
+            contextualized_chunks = (
+                await self.contextual_retriever.contextualize_all_chunks(
+                    self.all_PDFs_chunks, self.resumo_auxiliar, self.axiom_instance
+                )
+            )
+            self.ax.fim_requisicao_contextual()
+
+        chunks_processados = (
+            contextualized_chunks
+            if self.is_contextualized_chunk
+            else self.all_PDFs_chunks
+        )
+        self.chunks_processados = chunks_processados
+        self.ax.chunks_inicialmente(chunks_processados)
+        return chunks_processados
+
+    async def generate_query_for_vector_store(self):
+        prompt_para_gerar_query_dinamico = prompt_gerar_query_dinamicamente(
+            cast(str, self.resumo_auxiliar)
+        )
+
+        self.axiom_instance.send_axiom(
+            "COMEÇANDO REQUISIÇÃO PARA GERAR O QUERY DINAMICAMENTE DO VECTOR STORE"
+        )
+        response = await self.llm.google_gemini_ainvoke(
+            prompt_para_gerar_query_dinamico, "gemini-2.0-flash"
+        )
+
+        self.query_gerado_dinamicamente_para_o_vector_store = cast(
+            str, response.content
+        )
+
+        self.axiom_instance.send_axiom(
+            f"query_gerado_dinamicamente_para_o_vector_store: {self.query_gerado_dinamicamente_para_o_vector_store}",
+        )
+
+        return self.query_gerado_dinamicamente_para_o_vector_store
+
+    async def create_enhanced_vector_store(self):
+        vector_store, bm25, chunk_ids = self.vector_store.create_enhanced_vector_store(
+            self.chunks_processados, self.is_contextualized_chunk, self.axiom_instance  # type: ignore
+        )
+
+        self.enhanced_vector_store = vector_store, bm25, chunk_ids
+        return vector_store, bm25, chunk_ids
 
     def retrieve_with_rank_fusion(
         self, vector_store: Chroma, bm25: BM25Okapi, chunk_ids: List[str], query: str
@@ -236,17 +332,18 @@ class GerarDocumento:
             llm = llm_instance.google_gemini("gemini-2.5-flash-preview-04-17")
         return llm
 
-    async def gerar_documento_final(
+    async def do_last_requests(
         self,
-        vector_store: Chroma,
-        bm25: BM25Okapi,
-        chunk_ids: List[str],
-        llm_ultimas_requests: str,
-        query: str = "Summarize the main points of this document",
     ) -> List[Dict]:
         try:
+            self.axiom_instance.send_axiom("COMEÇANDO A FAZER ÚLTIMA REQUISIÇÃO")
+            vector_store, bm25, chunk_ids = self.enhanced_vector_store
+
             sources, contexts = self.rank_fusion_get_top_results(
-                vector_store, bm25, chunk_ids, query
+                vector_store,
+                bm25,
+                chunk_ids,
+                self.query_gerado_dinamicamente_para_o_vector_store,
             )
 
             prompt_gerar_documento = PromptTemplate(
@@ -254,6 +351,7 @@ class GerarDocumento:
                 input_variables=["context"],
             )
 
+            llm_ultimas_requests = self.llm_ultimas_requests
             llm = self.select_model_for_last_requests(llm_ultimas_requests)  # type: ignore
             prompt_instance = Prompt()
             context_do_prompt_primeira_etapa = "\n\n".join(contexts)
@@ -311,12 +409,56 @@ class GerarDocumento:
             structured_output = self.gerar_documento_utils.criar_output_estruturado(
                 summaries, sources
             )
+
+            self.axiom_instance.send_axiom("TERMINOU DE FAZER A ÚLTIMA REQUISIÇÃO")
+            self.structured_output = structured_output
             return structured_output
 
         except Exception as e:
             self.logger.error(f"Error generating enhanced summary: {str(e)}")
             raise
 
+    async def generate_complete_text(self):
+        texto_completo = "\n\n"
+
+        for x in self.structured_output:
+            texto_completo = texto_completo + x["content"] + "\n"
+            x["source"]["text"] = x["source"]["text"][0:200]
+            x["source"]["context"] = x["source"]["context"][0:200]
+
+        self.texto_completo_como_html = convert_markdown_to_HTML(
+            texto_completo
+        ).replace("resposta_segunda_etapa:", "<br><br>")
+
+        self.axiom_instance.send_axiom(
+            f"texto_completo_como_html: {self.texto_completo_como_html}"
+        )
+
+    async def get_document_title(self):
+        if self.is_contextualized_chunk:
+            resumo_para_gerar_titulo = self.resumo_auxiliar
+        else:
+            resumo_para_gerar_titulo = self.texto_completo_como_html
+
+        prompt = prompt_para_gerar_titulo(resumo_para_gerar_titulo)
+        response = await agemini_answer(prompt, "gemini-2.0-flash-lite")
+        self.titulo_do_documento = response
+        return self.titulo_do_documento
+
+    async def send_to_bubble(self):
+        self.axiom_instance.send_axiom("COMEÇANDO A REQUISIÇÃO FINAL PARA O BUBBLE")
+
+        enviar_resposta_final(
+            serializer.doc_id,  # type: ignore
+            serializer.form_response_id,  # type: ignore
+            serializer.version,  # type: ignore
+            self.texto_completo_como_html,
+            False,
+            cast(str, self.titulo_do_documento),
+        )
+
+        self.axiom_instance.send_axiom("TERMINOU A REQUISIÇÃO FINAL PARA O BUBBLE")
+
     async def checar_se_resposta_vazia_do_documento_final(
         self, llm_ultimas_requests: str, prompt: str
     ):

_utils/gerar_documento_utils/prompts.py

@@ -239,3 +239,7 @@ def prompt_gerar_query_dinamicamente(resumo_do_processo: str):
 
 - *Importante:** Sua resposta final deve ser *somente* a string da query.
 """
+
+
+def prompt_para_gerar_titulo(resumo_para_gerar_titulo: str):
+    return f"Você é um assistente jurídico e irá receber abaixo o resumo de um documento jurídico. Quero que você gere um título para este documento. Mande como resposta apenas o título gerado, nada mais. Aqui está um título de exemplo pra você se basear ao criar um novo: <titulo_de_exemplo>Ação Penal por Furto Qualificado nº 0002269-86.2009.805.0032<titulo_de_exemplo>\n\nSegue abaixo o resumo do documento jurídico:\n{resumo_para_gerar_titulo}"

_utils/gerar_documento_utils/utils.py

@@ -7,6 +7,7 @@ from _utils.langchain_utils.Splitter_class import Splitter
 from _utils.langchain_utils.LLM_class import LLM
 from _utils.gerar_documento_utils.prompts import (
     create_prompt_auxiliar_do_contextual_prompt,
+    prompt_para_gerar_titulo,
 )
 
 from _utils.models.gerar_documento import DocumentChunk
@@ -102,6 +103,6 @@ def split_text_by_tokens(full_text: str):
 
 
 async def generate_document_title(resumo_para_gerar_titulo: str):
-    prompt = f"Você é um assistente jurídico e irá receber abaixo o resumo de um documento jurídico. Quero que você gere um título para este documento. Mande como resposta apenas o título gerado, nada mais. Aqui está um título de exemplo pra você se basear ao criar um novo: <titulo_de_exemplo>Ação Penal por Furto Qualificado nº 0002269-86.2009.805.0032<titulo_de_exemplo>\n\nSegue abaixo o resumo do documento jurídico:\n{resumo_para_gerar_titulo}"
+    prompt = prompt_para_gerar_titulo(resumo_para_gerar_titulo)
     response = await agemini_answer(prompt, "gemini-2.0-flash-lite")
     return response

_utils/ragas.py

@@ -76,7 +76,7 @@ def test_ragas(serializer, listaPDFs):
 
     def generate_summary(vector_store, bm25, chunk_ids, query, summarizer):
         """Generates an enhanced summary using the vector store and BM25 index."""
-        structured_summaries = summarizer.gerar_documento_final(
+        structured_summaries = summarizer.do_last_requests(
             vector_store, bm25, chunk_ids, query
         )
         return {"structured_summaries": structured_summaries}

gerar_documento/views.py

@@ -2,7 +2,7 @@ from typing import Any, Dict, cast
 from langchain.prompts import PromptTemplate
 from _utils.Utils_Class import UtilsClass
 from _utils.bubble_integrations.enviar_resposta_final import enviar_resposta_final
-from _utils.custom_exception_handler import custom_exception_handler_wihout_api_handler
+from _utils.custom_exception_handler import custom_exception_handler_without_api_handler
 from _utils.gerar_documento_utils.GerarDocumento import GerarDocumento
 from _utils.langchain_utils.LLM_class import LLM
 from _utils.gerar_documento_utils.utils import (
@@ -143,7 +143,9 @@ class GerarEmentaView(AsyncAPIView):
                     data = cast(Dict[str, Any], serializer.validated_data)
                     self.serializer = data
 
-                    gerar_documento_instance = GerarDocumento(obj, self.axiom_instance)
+                    gerar_documento_instance = GerarDocumento(
+                        obj, True, self.axiom_instance
+                    )
 
                     listaPDFs = [l["link_arquivo"] for l in data["files"]]
 
@@ -210,7 +212,7 @@ class GerarEmentaView(AsyncAPIView):
                     )
                 except Exception as e:
                     print(f"ERRO GERAR EMENTA: {e}")
-                    custom_exception_handler_wihout_api_handler(
+                    custom_exception_handler_without_api_handler(
                         e, serializer, self.axiom_instance
                     )
                     raise
@@ -250,7 +252,7 @@ class GerarEmentaComPDFProprioView(AsyncAPIView):
             serializer_obj = serializer.get_obj()
 
             gerar_documento_instance = GerarDocumento(
-                serializer_obj, self.axiom_instance
+                serializer_obj, False, self.axiom_instance
             )
 
             listaPDFs = handle_pdf_files_from_serializer(