feat: pequenas melhorias

_utils/gerar_relatorio_modelo_usuario/EnhancedDocumentSummarizer.py

@@ -20,15 +20,12 @@ from _utils.models.gerar_relatorio import (
 )
 from modelos_usuarios.serializer import ModeloUsuarioSerializer
 from setup.environment import api_url
-from _utils.gerar_relatorio_modelo_usuario.contextual_retriever import (
-    ContextualRetriever,
-)
+
 from asgiref.sync import sync_to_async
 
 
 class EnhancedDocumentSummarizer(DocumentSummarizer):
     openai_api_key = os.environ.get("OPENAI_API_KEY", "")
-    claude_api_key = os.environ.get("CLAUDE_API_KEY", "")
 
     def __init__(
         self,
@@ -38,7 +35,6 @@ class EnhancedDocumentSummarizer(DocumentSummarizer):
         chunk_overlap,
         num_k_rerank,
         model_cohere_rerank,
-        claude_context_model,
         prompt_auxiliar,
         gpt_model,
         gpt_temperature,
@@ -56,14 +52,10 @@ class EnhancedDocumentSummarizer(DocumentSummarizer):
             model_cohere_rerank,
         )
         self.config = config
-        self.contextual_retriever = ContextualRetriever(
-            config, self.claude_api_key, claude_context_model
-        )
         self.logger = logging.getLogger(__name__)
         self.prompt_auxiliar = prompt_auxiliar
         self.gpt_model = gpt_model
         self.gpt_temperature = gpt_temperature
-        # self.id_modelo_do_usuario = id_modelo_do_usuario
         self.prompt_gerar_documento = prompt_gerar_documento
         self.reciprocal_rank_fusion = reciprocal_rank_fusion
         self.resumo_gerado = ""

_utils/gerar_relatorio_modelo_usuario/contextual_retriever.py

@@ -1,33 +1,16 @@
 import os
-
-from _utils.LLMs.LLM_class import LLM
-from _utils.gerar_relatorio_modelo_usuario.prompts import (
-    prompt_auxiliar_do_contextual_prompt,
-    create_prompt_auxiliar_do_contextual_prompt,
-)
-from _utils.bubble_integrations.obter_arquivo import get_pdf_from_bubble
-from _utils.chains.Chain_class import Chain
 from _utils.gerar_relatorio_modelo_usuario.utils import (
+    get_response_from_auxiliar_contextual_prompt,
     validate_many_chunks_in_one_request,
 )
-from _utils.handle_files import return_document_list_with_llama_parser
-from _utils.prompts.Prompt_class import Prompt
-from _utils.splitters.Splitter_class import Splitter
-from setup.easy_imports import PyPDFLoader
-from langchain_openai import ChatOpenAI
 from typing import List, Dict, Tuple, Optional, cast
 from anthropic import Anthropic, AsyncAnthropic
 import logging
 from langchain.schema import Document
 from llama_index import Document as Llama_Index_Document
 import asyncio
-from langchain.prompts import PromptTemplate
 from typing import List
-from multiprocessing import Process, Barrier, Queue
 from dataclasses import dataclass
-from langchain_core.messages import HumanMessage
-from asgiref.sync import sync_to_async
-from setup.easy_imports import ChatPromptTemplate, ChatOpenAI
 
 from _utils.gerar_relatorio_modelo_usuario.llm_calls import aclaude_answer, agpt_answer
 from _utils.gerar_relatorio_modelo_usuario.prompts import contextual_prompt
@@ -36,161 +19,30 @@ from _utils.models.gerar_relatorio import (
     DocumentChunk,
     RetrievalConfig,
 )
-from _utils.prompts.Prompt_class import prompt as prompt_obj
 
 lista_contador = []
 
 
 class ContextualRetriever:
-    def __init__(
-        self, config: RetrievalConfig, claude_api_key: str, claude_context_model: str
-    ):
+
+    def __init__(self, config: RetrievalConfig, claude_context_model: str):
         self.config = config
-        # self.claude_client = Anthropic(api_key=claude_api_key)
-        self.claude_client = AsyncAnthropic(api_key=claude_api_key)
         self.logger = logging.getLogger(__name__)
         self.bm25 = None
         self.claude_context_model = claude_context_model
 
-    async def contextualize_all_chunks(
-        self, full_text_as_array: List[str], chunks: List[DocumentChunk]
-    ) -> List[ContextualizedChunk]:
-        """Add context to all chunks"""
-        contextualized_chunks = []
-        full_text = ""
-        for x in full_text_as_array:
-            full_text += x
-
-        prompt_auxiliar_summary = create_prompt_auxiliar_do_contextual_prompt(full_text)
-
-        print("\n\n\nprompt_auxiliar_summary[0:500]: ", prompt_auxiliar_summary[0:500])
-
-        # Claude comentado pois o limite de tokens estava sendo passado pela requisição e dava erro
-        # response_auxiliar_summary = await aclaude_answer(
-        #     self.claude_client, self.claude_context_model, prompt_auxiliar_summary
-        # )
-
-        llms = LLM()
-        response_auxiliar_summary = await llms.googleGemini().ainvoke(
-            [HumanMessage(content=prompt_auxiliar_summary)]
-        )
-
-        print("\n\n\n\nresponse_auxiliar_summary: ", response_auxiliar_summary.content)
-
-        lista_de_listas_cada_com_20_chunks = [
-            chunks[i : i + 20] for i in range(0, len(chunks), 20)
-        ]
-        print(
-            "lista_de_listas_cada_com_20_chunks: ", lista_de_listas_cada_com_20_chunks
-        )
-
-        async with asyncio.TaskGroup() as tg:
-            tasks = [
-                tg.create_task(
-                    self.create_contextualized_chunk(
-                        chunk, full_text_as_array, response_auxiliar_summary.content
-                    )
-                )
-                # for chunk in chunks # ORIGINAL
-                for chunk in lista_de_listas_cada_com_20_chunks
-            ]
-
-        # contextualized_chunks = [task.result() for task in tasks]
-        contextualized_chunks = []
-        for task in tasks:
-            # print("\n\ntask", task)
-            # print("\n\ntask.result()", task.result())
-
-            contextualized_chunks = contextualized_chunks + task.result()
-
-        return contextualized_chunks
-
-    # ORIGINAL
-    # async def create_contextualized_chunk(
-    #     self, chunk, single_page_text, response_auxiliar_summary
-    # ):
-    #     lista_contador.append(0)
-    #     print("contador: ", len(lista_contador))
-    #     page_number = chunk.page_number - 1
-    #     page_content = single_page_text[page_number].page_content
-
-    #     context = await self.llm_generate_context(
-    #         page_content, chunk, response_auxiliar_summary
-    #     )
-    #     print("context: ", context)
-    #     return ContextualizedChunk(
-    #         content=chunk.content,
-    #         page_number=chunk.page_number,
-    #         chunk_id=chunk.chunk_id,
-    #         start_char=chunk.start_char,
-    #         end_char=chunk.end_char,
-    #         context=context,
-    #     )
-
-    async def create_contextualized_chunk(
-        self, chunks: List[DocumentChunk], single_page_text, response_auxiliar_summary
-    ):
-
-        lista_contador.append(0)
-        print("contador: ", len(lista_contador))
-        # all_pages_contents = ""
-        # contador = 1
-        # for chunk in chunks:
-        #     page_number = chunk.page_number - 1
-        #     page_content = single_page_text[page_number].page_content
-
-        #     all_pages_contents += page_content
-        #     contador += 1
-
-        result = await self.llm_generate_context(chunks, response_auxiliar_summary)
-
-        lista_chunks = []
-        for index, chunk in enumerate(chunks):
-            lista_chunks.append(
-                ContextualizedChunk(
-                    contextual_summary=result[index][2],
-                    content=chunk.content,
-                    page_number=chunk.page_number,
-                    id_do_processo=int(result[index][0]),
-                    chunk_id=chunk.chunk_id,
-                    start_char=chunk.start_char,
-                    end_char=chunk.end_char,
-                    context=result[index][1],
-                )
-            )
-
-        return lista_chunks
-
-    # ORIGINAL
-    # async def llm_generate_context(
-    #     self, page_text: str, chunk: DocumentChunk, resumo_auxiliar
-    # ) -> str:
-    #     """Generate contextual description using ChatOpenAI"""
-    #     try:
-    #         print("COMEÇOU A REQUISIÇÃO")
-    #         prompt = contextual_prompt(page_text, resumo_auxiliar, chunk.content)
-    #         # response = await aclaude_answer(
-    #         #     self.claude_client, self.claude_context_model, prompt
-    #         # )
-
-    #         # response = await agpt_answer(prompt)
-    #         llms = LLM()
-    #         response = await llms.deepseek().ainvoke([HumanMessage(content=prompt)])
-    #         return cast(str, response.content)
-    #     except Exception as e:
-    #         self.logger.error(
-    #             f"Context generation failed for chunk {chunk.chunk_id}: {str(e)}"
-    #         )
-    #         return ""
+        self.claude_api_key = os.environ.get("CLAUDE_API_KEY", "")
+        self.claude_client = AsyncAnthropic(api_key=self.claude_api_key)
+        # self.claude_client = Anthropic(api_key=claude_api_key)
 
-    async def llm_generate_context(
-        self, chunks: List[DocumentChunk], resumo_auxiliar  # , page_text: str
+    async def llm_call_uma_lista_de_chunks(
+        self, lista_com_20_chunks: List[DocumentChunk], resumo_auxiliar
     ) -> str:
         """Generate contextual description using ChatOpenAI"""
         contador = 1
         all_chunks_contents = ""
 
-        for chunk in chunks:
+        for chunk in lista_com_20_chunks:
             all_chunks_contents += chunk.content
             all_chunks_contents += f"\n\n CHUNK {contador}:\n"
             contador += 1
@@ -203,7 +55,9 @@ class ContextualRetriever:
             # )
 
             for attempt in range(4):
-                print(f"\n\nTENTATIVA FORMATAÇÃO CHUNKS NÚMERO {attempt}")
+                print(
+                    f"\n\nTENTATIVA FORMATAÇÃO CHUNKS NÚMERO {attempt}: {all_chunks_contents[0:500]}"
+                )
                 raw_response = await agpt_answer(prompt)
                 response = cast(str, raw_response)
                 # llms = LLM()
@@ -211,7 +65,6 @@ class ContextualRetriever:
                 # return cast(str, response.content)
 
                 matches = validate_many_chunks_in_one_request(response)
-                # Convert matches to the desired format
 
                 if matches:
                     result = [
@@ -224,62 +77,61 @@ class ContextualRetriever:
             self.logger.error(f"Context generation failed for chunks .... : {str(e)}")
             return ""
 
-    # def gerar_resumo_auxiliar_do_contextual_embedding(self):
-    #     prompt = Prompt().create_prompt_template(
-    #         "", prompt_auxiliar_do_contextual_prompt
-    #     )
-    #     Chain(prompt, ChatOpenAI())
-    #     return
-
+    async def contextualize_uma_lista_de_chunks(
+        self, lista_com_20_chunks: List[DocumentChunk], response_auxiliar_summary
+    ):
+        lista_contador.append(0)
+        print("contador: ", len(lista_contador))
 
-# Primeira função chamada do arquivo
-async def contextualize_chunk_based_on_serializer(
-    serializer, contextual_retriever: ContextualRetriever, pages, all_PDFs_chunks
-):
-    if serializer["should_have_contextual_chunks"]:
-        contextualized_chunks = await contextual_retriever.contextualize_all_chunks(
-            pages, all_PDFs_chunks
+        result = await self.llm_call_uma_lista_de_chunks(
+            lista_com_20_chunks, response_auxiliar_summary
         )
-        chunks_passados = contextualized_chunks
-        is_contextualized_chunk = True
-    else:
-        chunks_passados = all_PDFs_chunks
-        is_contextualized_chunk = False
-
-    return chunks_passados, is_contextualized_chunk
 
+        lista_chunks = []
+        for index, chunk in enumerate(lista_com_20_chunks):
+            lista_chunks.append(
+                ContextualizedChunk(
+                    contextual_summary=result[index][2],
+                    content=chunk.content,
+                    page_number=chunk.page_number,
+                    id_do_processo=int(result[index][0]),
+                    chunk_id=chunk.chunk_id,
+                    start_char=chunk.start_char,
+                    end_char=chunk.end_char,
+                    context=result[index][1],
+                )
+            )
 
-async def get_full_text_and_all_PDFs_chunks(
-    listaPDFs: List[str],
-    splitterObject: Splitter,
-    should_use_llama_parse: bool,
-    isBubble: bool,
-):
-    all_PDFs_chunks = []
+        return lista_chunks
 
-    pages: List[str] = []
+    async def contextualize_all_chunks(
+        self,
+        all_PDFs_chunks: List[DocumentChunk],
+        response_auxiliar_summary,
+    ) -> List[ContextualizedChunk]:
+        """Add context to all chunks"""
 
-    # Load and process document
-    for pdf_path in listaPDFs:
-        # if isBubble:
-        #     pages = pages + await get_pdf_from_bubble(pdf_path, should_use_llama_parse)
-        # else:
-        #     if should_use_llama_parse:
-        #         pages = pages + await return_document_list_with_llama_parser(pdf_path)
-        #     else:
-        #         pages = pages + PyPDFLoader(pdf_path).load()
+        lista_de_listas_cada_com_20_chunks = [
+            all_PDFs_chunks[i : i + 20] for i in range(0, len(all_PDFs_chunks), 20)
+        ]
 
-        chunks, pages = await splitterObject.load_and_split_document(
-            pdf_path, should_use_llama_parse, isBubble
-        )
-        all_PDFs_chunks = all_PDFs_chunks + chunks
-    # Get full text for contextualization
-    # loader = PyPDFLoader(pdf_path)
+        async with asyncio.TaskGroup() as tg:
+            tasks = [
+                tg.create_task(
+                    self.contextualize_uma_lista_de_chunks(
+                        lista_com_20_chunks,
+                        response_auxiliar_summary,
+                    )
+                )
+                for lista_com_20_chunks in lista_de_listas_cada_com_20_chunks
+            ]
 
-    # full_text = ""
-    # full_text = " ".join([page.page_content for page in pages])
+        # contextualized_chunks = [task.result() for task in tasks]
+        contextualized_chunks = []
+        for task in tasks:
+            contextualized_chunks = contextualized_chunks + task.result()
 
-    return all_PDFs_chunks, pages  # , full_text
+        return contextualized_chunks
 
 
 # Código comentado abaixo é para ler as páginas ao redor da página atual do chunk

_utils/gerar_relatorio_modelo_usuario/utils.py

@@ -1,5 +1,12 @@
 from typing import List, Tuple
 from langchain_core.documents import Document
+from langchain_core.messages import HumanMessage
+
+from _utils.splitters.Splitter_class import Splitter
+from _utils.LLMs.LLM_class import LLM
+from _utils.gerar_relatorio_modelo_usuario.prompts import (
+    create_prompt_auxiliar_do_contextual_prompt,
+)
 
 
 def gerar_resposta_compilada(serializer):
@@ -69,3 +76,51 @@ def validate_many_chunks_in_one_request(response: str):
     if len(matches) == 0:
         return False
     return matches_as_list
+
+
+# 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]):
+    full_text = ""
+    for x in full_text_as_array:
+        full_text += x
+
+    prompt_auxiliar_summary = create_prompt_auxiliar_do_contextual_prompt(full_text)
+
+    print("\n\n\nprompt_auxiliar_summary[0:500]: ", prompt_auxiliar_summary[0:500])
+
+    # Claude comentado pois o limite de tokens estava sendo passado pela requisição e dava erro
+    # response_auxiliar_summary = await aclaude_answer(
+    #     self.claude_client, self.claude_context_model, prompt_auxiliar_summary
+    # )
+
+    llms = LLM()
+    response_auxiliar_summary = await llms.googleGemini().ainvoke(
+        [HumanMessage(content=prompt_auxiliar_summary)]
+    )
+
+    print(
+        "\n\n\n\nresponse_auxiliar_summary.content[0:500]: ",
+        response_auxiliar_summary.content[0:500],
+    )
+
+    return response_auxiliar_summary.content
+
+
+async def get_full_text_and_all_PDFs_chunks(
+    listaPDFs: List[str],
+    splitterObject: Splitter,
+    should_use_llama_parse: bool,
+    isBubble: bool,
+):
+    all_PDFs_chunks = []
+
+    pages: List[str] = []
+
+    # Load and process document
+    for pdf_path in listaPDFs:
+        chunks, pages = await splitterObject.load_and_split_document(
+            pdf_path, should_use_llama_parse, isBubble
+        )
+        all_PDFs_chunks = all_PDFs_chunks + chunks
+
+    return all_PDFs_chunks, pages

_utils/resumo_completo_cursor.py

@@ -4,10 +4,13 @@ from _utils.gerar_relatorio_modelo_usuario.EnhancedDocumentSummarizer import (
     EnhancedDocumentSummarizer,
 )
 from _utils.gerar_relatorio_modelo_usuario.contextual_retriever import (
-    contextualize_chunk_based_on_serializer,
+    ContextualRetriever,
+)
+from _utils.gerar_relatorio_modelo_usuario.utils import (
+    gerar_resposta_compilada,
     get_full_text_and_all_PDFs_chunks,
+    get_response_from_auxiliar_contextual_prompt,
 )
-from _utils.gerar_relatorio_modelo_usuario.utils import gerar_resposta_compilada
 from _utils.models.gerar_relatorio import (
     RetrievalConfig,
 )
@@ -51,6 +54,10 @@ async def get_llm_summary_answer_by_cursor_complete(
         chunk_overlap=serializer["chunk_overlap"],
     )
 
+    contextual_retriever = ContextualRetriever(
+        config, serializer["claude_context_model"]
+    )
+
     # Initialize enhanced summarizer
     summarizer = EnhancedDocumentSummarizer(
         config=config,
@@ -59,29 +66,35 @@ async def get_llm_summary_answer_by_cursor_complete(
         chunk_size=serializer["chunk_size"],
         num_k_rerank=serializer["num_k_rerank"],
         model_cohere_rerank=serializer["model_cohere_rerank"],
-        claude_context_model=serializer["claude_context_model"],
         prompt_auxiliar=serializer["prompt_auxiliar"],
         gpt_model=serializer["model"],
         gpt_temperature=serializer["gpt_temperature"],
-        # id_modelo_do_usuario=serializer["id_modelo_do_usuario"],
         prompt_gerar_documento=serializer["prompt_gerar_documento"],
         reciprocal_rank_fusion=reciprocal_rank_fusion,
     )
 
-    allPdfsChunks, pages = await get_full_text_and_all_PDFs_chunks(
+    all_PDFs_chunks, full_text_as_array = await get_full_text_and_all_PDFs_chunks(
         listaPDFs, summarizer.splitter, serializer["should_use_llama_parse"], isBubble
     )
 
-    chunks_passados, is_contextualized_chunk = (
-        await contextualize_chunk_based_on_serializer(
-            serializer, summarizer.contextual_retriever, pages, allPdfsChunks
+    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
         )
-    )
+
+        contextualized_chunks = await contextual_retriever.contextualize_all_chunks(
+            all_PDFs_chunks, response_auxiliar_summary
+        )
+        chunks_processados = contextualized_chunks
+    else:
+        chunks_processados = all_PDFs_chunks
 
     # Create enhanced vector store and BM25 index
     vector_store, bm25, chunk_ids = (
         summarizer.vector_store.create_enhanced_vector_store(
-            chunks_passados, is_contextualized_chunk
+            chunks_processados, is_contextualized_chunk
         )
     )
 

tests/gerar_relatorio_modelo_usuario/test_contextual_retriever.py

@@ -0,0 +1,2 @@
+class TestContextualRetriever:
+    pass