feat: melhorando refatoração do gerar_documento

_utils/gerar_documento.py

@@ -1,30 +1,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 typing import Any, Union
 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,
 )
-from _utils.gerar_documento_utils.contextual_retriever import (
-    ContextualRetriever,
-)
-from _utils.gerar_documento_utils.utils import (
-    generate_document_title,
-    gerar_resposta_compilada,
-    get_response_from_auxiliar_contextual_prompt,
-)
-from _utils.models.gerar_documento import (
-    RetrievalConfig,
-)
-import markdown
 
-from _utils.langchain_utils.Prompt_class import Prompt
-from _utils.utils import convert_markdown_to_HTML
 from gerar_documento.serializer import (
     GerarDocumentoComPDFProprioSerializer,
     GerarDocumentoComPDFProprioSerializerData,
@@ -48,20 +29,12 @@ async def gerar_documento(
     isBubble=False,
 ):
     try:
-        axiom = axiom_instance.send_axiom
-        ax = AxiomLogs(axiom_instance)
-        utils = UtilsClass()
         summarizer = GerarDocumento(serializer, isBubble, axiom_instance)
+        summarizer.lista_pdfs = listaPDFs
 
-        all_PDFs_chunks, full_text_as_array = await summarizer.get_text_and_pdf_chunks()
-
-        is_contextualized_chunk = serializer.should_have_contextual_chunks
+        await summarizer.get_text_and_pdf_chunks()
 
-        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)
+        await summarizer.get_response_from_auxiliar_contextual_prompt()
 
         await summarizer.generate_chunks_processados()
 
@@ -85,7 +58,7 @@ async def gerar_documento(
             "texto_completo": summarizer.texto_completo_como_html,
             "titulo_do_documento": summarizer.titulo_do_documento,
             "resultado": structured_summaries,
-            "parametros-utilizados": gerar_resposta_compilada(serializer),
+            "parametros-utilizados": summarizer.gerar_resposta_compilada(),
         }
     except Exception as e:
         custom_exception_handler_without_api_handler(e, serializer, axiom_instance)

_utils/gerar_documento_utils/GerarDocumento.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-import os
+from langchain_core.messages import HumanMessage
 from typing import Any, List, Dict, Literal, Tuple, Optional, Union, cast
 
 from pydantic import SecretStr
@@ -9,6 +9,7 @@ from _utils.bubble_integrations.enviar_resposta_final import enviar_resposta_fin
 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 (
+    create_prompt_auxiliar_do_contextual_prompt,
     prompt_gerar_query_dinamicamente,
     prompt_para_gerar_titulo,
 )
@@ -40,6 +41,7 @@ from _utils.langchain_utils.Splitter_class import Splitter
 import time
 from setup.tokens import openai_api_key, cohere_api_key
 from setup.logging import Axiom
+import tiktoken
 
 
 def reciprocal_rank_fusion(result_lists, weights=None):
@@ -124,6 +126,10 @@ class GerarDocumento:
     structured_output: List[Any]
     texto_completo_como_html: str
     titulo_do_documento: str
+    encoding_tiktoken = tiktoken.get_encoding("cl100k_base")
+    serializer: Union[
+        GerarDocumentoSerializerData, GerarDocumentoComPDFProprioSerializerData, Any
+    ]
 
     def __init__(
         self,
@@ -133,6 +139,7 @@ class GerarDocumento:
         isBubble: bool,
         axiom_instance: Axiom,
     ):
+        self.serializer = serializer
         self.config = self.gerar_documento_utils.create_retrieval_config(serializer)
         self.logger = logging.getLogger(__name__)
         # self.prompt_auxiliar = prompt_auxiliar
@@ -188,8 +195,10 @@ class GerarDocumento:
             else self.all_PDFs_chunks
         )
         self.chunks_processados = chunks_processados
-        self.ax.chunks_inicialmente(chunks_processados)
-        return chunks_processados
+        if len(self.chunks_processados) == 0:
+            self.chunks_processados = self.all_PDFs_chunks
+        self.ax.chunks_inicialmente(self.chunks_processados)
+        return self.chunks_processados
 
     async def generate_query_for_vector_store(self):
         prompt_para_gerar_query_dinamico = prompt_gerar_query_dinamicamente(
@@ -542,3 +551,77 @@ class GerarDocumento:
             self.axiom_instance.send_axiom(f"RESULTADO ETAPA 3: {documento_gerado}")
 
         return texto_final_juntando_as_etapas
+
+    # Esta função gera a resposta que será usada em cada um das requisições de cada chunk
+    async def get_response_from_auxiliar_contextual_prompt(self):
+        llms = LLM()
+        responses = []
+
+        current_chunk = []
+        current_token_count = 0
+        chunk_counter = 1
+
+        for part in self.full_text_as_array:
+            part_tokens = len(self.encoding_tiktoken.encode(part))
+
+            # Check if adding this part would EXCEED the limit
+            if current_token_count + part_tokens > 600000:
+                # Process the accumulated chunk before it exceeds the limit
+                chunk_text = "".join(current_chunk)
+                print(
+                    f"\nProcessing chunk {chunk_counter} with {current_token_count} tokens"
+                )
+
+                prompt = create_prompt_auxiliar_do_contextual_prompt(chunk_text)
+                response = await llms.google_gemini().ainvoke(
+                    [HumanMessage(content=prompt)]
+                )
+                responses.append(response.content)
+
+                # Start new chunk with current part
+                current_chunk = [part]
+                current_token_count = part_tokens
+                chunk_counter += 1
+            else:
+                # Safe to add to current chunk
+                current_chunk.append(part)
+                current_token_count += part_tokens
+
+        # Process the final remaining chunk
+        if current_chunk:
+            chunk_text = "".join(current_chunk)
+            print(
+                f"\nProcessing final chunk {chunk_counter} with {current_token_count} tokens"
+            )
+            prompt = create_prompt_auxiliar_do_contextual_prompt(chunk_text)
+            response = await llms.google_gemini().ainvoke(
+                [HumanMessage(content=prompt)]
+            )
+            responses.append(response.content)
+
+        self.resumo_auxiliar = "".join(responses)
+        self.ax.resumo_inicial_processo(self.resumo_auxiliar)
+
+        return self.resumo_auxiliar
+
+    def gerar_resposta_compilada(self):
+        serializer = self.serializer
+        return {
+            "num_chunks_retrieval": serializer.num_chunks_retrieval,
+            "embedding_weight": serializer.embedding_weight,
+            "bm25_weight": serializer.bm25_weight,
+            "context_window": serializer.context_window,
+            "chunk_overlap": serializer.chunk_overlap,
+            "num_k_rerank": serializer.num_k_rerank,
+            "model_cohere_rerank": serializer.model_cohere_rerank,
+            "more_initial_chunks_for_reranking": serializer.more_initial_chunks_for_reranking,
+            "claude_context_model": serializer.claude_context_model,
+            "gpt_temperature": serializer.gpt_temperature,
+            "user_message": serializer.user_message,
+            "model": serializer.model,
+            "hf_embedding": serializer.hf_embedding,
+            "chunk_size": serializer.chunk_size,
+            "chunk_overlap": serializer.chunk_overlap,
+            # "prompt_auxiliar": serializer.prompt_auxiliar,
+            "prompt_gerar_documento": serializer.prompt_gerar_documento[0:200],
+        }

_utils/gerar_documento_utils/utils.py

@@ -1,92 +1,10 @@
-from typing import Any, List, Tuple, Union
-from langchain_core.documents import Document
-from langchain_core.messages import HumanMessage
-
 from _utils.gerar_documento_utils.llm_calls import agemini_answer
-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
-from gerar_documento.serializer import GerarDocumentoSerializerData
+from _utils.gerar_documento_utils.prompts import prompt_para_gerar_titulo
 import tiktoken
 
 encoding = tiktoken.get_encoding("cl100k_base")
 
 
-def gerar_resposta_compilada(serializer: Union[GerarDocumentoSerializerData, Any]):
-    return {
-        "num_chunks_retrieval": serializer.num_chunks_retrieval,
-        "embedding_weight": serializer.embedding_weight,
-        "bm25_weight": serializer.bm25_weight,
-        "context_window": serializer.context_window,
-        "chunk_overlap": serializer.chunk_overlap,
-        "num_k_rerank": serializer.num_k_rerank,
-        "model_cohere_rerank": serializer.model_cohere_rerank,
-        "more_initial_chunks_for_reranking": serializer.more_initial_chunks_for_reranking,
-        "claude_context_model": serializer.claude_context_model,
-        "gpt_temperature": serializer.gpt_temperature,
-        "user_message": serializer.user_message,
-        "model": serializer.model,
-        "hf_embedding": serializer.hf_embedding,
-        "chunk_size": serializer.chunk_size,
-        "chunk_overlap": serializer.chunk_overlap,
-        # "prompt_auxiliar": serializer.prompt_auxiliar,
-        "prompt_gerar_documento": serializer.prompt_gerar_documento[0:200],
-    }
-
-
-# Esta função gera a resposta que será usada em cada um das requisições de cada chunk
-async def get_response_from_auxiliar_contextual_prompt(full_text_as_array: List[str]):
-    llms = LLM()
-    responses = []
-
-    current_chunk = []
-    current_token_count = 0
-    chunk_counter = 1
-
-    for part in full_text_as_array:
-        part_tokens = len(encoding.encode(part))
-
-        # Check if adding this part would EXCEED the limit
-        if current_token_count + part_tokens > 600000:
-            # Process the accumulated chunk before it exceeds the limit
-            chunk_text = "".join(current_chunk)
-            print(
-                f"\nProcessing chunk {chunk_counter} with {current_token_count} tokens"
-            )
-
-            prompt = create_prompt_auxiliar_do_contextual_prompt(chunk_text)
-            response = await llms.google_gemini().ainvoke(
-                [HumanMessage(content=prompt)]
-            )
-            responses.append(response.content)
-
-            # Start new chunk with current part
-            current_chunk = [part]
-            current_token_count = part_tokens
-            chunk_counter += 1
-        else:
-            # Safe to add to current chunk
-            current_chunk.append(part)
-            current_token_count += part_tokens
-
-    # Process the final remaining chunk
-    if current_chunk:
-        chunk_text = "".join(current_chunk)
-        print(
-            f"\nProcessing final chunk {chunk_counter} with {current_token_count} tokens"
-        )
-        prompt = create_prompt_auxiliar_do_contextual_prompt(chunk_text)
-        response = await llms.google_gemini().ainvoke([HumanMessage(content=prompt)])
-        responses.append(response.content)
-
-    return "".join(responses)
-
-
 def split_text_by_tokens(full_text: str):
     tokens = encoding.encode(full_text)
     max_tokens = 600000

gerar_documento/views.py

@@ -7,7 +7,6 @@ from _utils.gerar_documento_utils.GerarDocumento import GerarDocumento
 from _utils.langchain_utils.LLM_class import LLM
 from _utils.gerar_documento_utils.utils import (
     generate_document_title,
-    gerar_resposta_compilada,
     split_text_by_tokens,
 )
 from _utils.langchain_utils.Prompt_class import Prompt