feat: refatorações e melhorias gerais

_utils/LLMs/LLM_class.py

@@ -1,6 +1,4 @@
-# from langchain_openai import OpenAI
 from typing import cast
-from openai import OpenAI
 from pydantic import SecretStr
 from setup.environment import default_model
 from setup.easy_imports import ChatOpenAI, ChatGoogleGenerativeAI
@@ -17,17 +15,17 @@ class LLM:
     # def create_GPT_model(self, model=default_model):
     #     return ChatOpen()
 
-    def deepseek(self):
+    def deepseek(self, model="deepseek-chat"):
         return ChatOpenAI(
             api_key=SecretStr(deepseek_api_key),
             base_url="https://api.deepseek.com/v1",
-            model="deepseek-chat",
+            model=model,
         )
 
-    def googleGemini(self):
+    def google_gemini(self, model="gemini-1.5-flash"):
         return ChatGoogleGenerativeAI(
             api_key=SecretStr(google_api_key),
-            model="gemini-1.5-flash",
+            model=model,
             temperature=0,
             max_tokens=None,
             timeout=None,

_utils/gerar_relatorio_modelo_usuario/{EnhancedDocumentSummarizer.py → GerarDocumento.py}

@@ -10,23 +10,23 @@ from setup.easy_imports import (
     PromptTemplate,
     BM25Okapi,
     Response,
+    HuggingFaceEmbeddings,
 )
 import logging
-import requests
 from _utils.gerar_relatorio_modelo_usuario.DocumentSummarizer_simples import (
     DocumentSummarizer,
 )
 from _utils.models.gerar_relatorio import (
     RetrievalConfig,
 )
-from modelos_usuarios.serializer import ModeloUsuarioSerializer
-from setup.environment import api_url
+from cohere import Client
+from _utils.splitters.Splitter_class import Splitter
 
-from asgiref.sync import sync_to_async
 
-
-class EnhancedDocumentSummarizer(DocumentSummarizer):
+class GerarDocumento:
     openai_api_key = os.environ.get("OPENAI_API_KEY", "")
+    cohere_api_key = os.environ.get("COHERE_API_KEY", "")
+    resumo_gerado = ""
 
     def __init__(
         self,
@@ -43,15 +43,6 @@ class EnhancedDocumentSummarizer(DocumentSummarizer):
         prompt_gerar_documento,
         reciprocal_rank_fusion,
     ):
-        super().__init__(
-            self.openai_api_key,
-            os.environ.get("COHERE_API_KEY", ""),
-            embedding_model,
-            chunk_size,
-            chunk_overlap,
-            num_k_rerank,
-            model_cohere_rerank,
-        )
         self.config = config
         self.logger = logging.getLogger(__name__)
         self.prompt_auxiliar = prompt_auxiliar
@@ -59,7 +50,13 @@ class EnhancedDocumentSummarizer(DocumentSummarizer):
         self.gpt_temperature = gpt_temperature
         self.prompt_gerar_documento = prompt_gerar_documento
         self.reciprocal_rank_fusion = reciprocal_rank_fusion
-        self.resumo_gerado = ""
+
+        self.openai_api_key = self.openai_api_key
+        self.cohere_client = Client(self.cohere_api_key)
+        self.embeddings = HuggingFaceEmbeddings(model_name=embedding_model)
+        self.num_k_rerank = num_k_rerank
+        self.model_cohere_rerank = model_cohere_rerank
+        self.splitter = Splitter(chunk_size, chunk_overlap)
 
         self.vector_store = VectorStore(embedding_model)
 
@@ -114,56 +111,73 @@ class EnhancedDocumentSummarizer(DocumentSummarizer):
             self.logger.error(f"Error in rank fusion retrieval: {str(e)}")
             raise
 
-    async def generate_enhanced_summary(
+    def rank_fusion_get_top_results(
         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]:
-        """Generate enhanced summary using both vector and BM25 retrieval"""
-        try:
-            # Get combined results using rank fusion
-            ranked_results = self.retrieve_with_rank_fusion(
-                vector_store, bm25, chunk_ids, query
+    ):
+        # Get combined results using rank fusion
+        ranked_results = self.retrieve_with_rank_fusion(
+            vector_store, bm25, chunk_ids, query
+        )
+
+        # Prepare context and track sources
+        contexts = []
+        sources = []
+
+        # Get full documents for top results
+        for chunk_id, score in ranked_results[: self.config.num_chunks]:
+            results = vector_store.get(
+                where={"chunk_id": chunk_id}, include=["documents", "metadatas"]
             )
 
-            # Prepare context and track sources
-            contexts = []
-            sources = []
+            if results["documents"]:
+                context = results["documents"][0]
+                metadata = results["metadatas"][0]
 
-            # Get full documents for top results
-            for chunk_id, score in ranked_results[: self.config.num_chunks]:
-                results = vector_store.get(
-                    where={"chunk_id": chunk_id}, include=["documents", "metadatas"]
+                contexts.append(context)
+                sources.append(
+                    {
+                        "content": context,
+                        "page": metadata["page"],
+                        "chunk_id": chunk_id,
+                        "relevance_score": score,
+                        "context": metadata.get("context", ""),
+                    }
                 )
 
-                if results["documents"]:
-                    context = results["documents"][0]
-                    metadata = results["metadatas"][0]
-
-                    contexts.append(context)
-                    sources.append(
-                        {
-                            "content": context,
-                            "page": metadata["page"],
-                            "chunk_id": chunk_id,
-                            "relevance_score": score,
-                            "context": metadata.get("context", ""),
-                        }
-                    )
+        return sources, contexts
 
-            if llm_ultimas_requests == "gpt-4o-mini":
-                llm = ChatOpenAI(
-                    temperature=self.gpt_temperature,
-                    model=self.gpt_model,
-                    api_key=SecretStr(self.openai_api_key),
-                )
-            elif llm_ultimas_requests == "deepseek-chat":
-                llm_instance = LLM()
-                llm = llm_instance.deepseek()
+    def select_model_for_last_requests(self, llm_ultimas_requests: str):
+        llm_instance = LLM()
+        if llm_ultimas_requests == "gpt-4o-mini":
+            llm = ChatOpenAI(
+                temperature=self.gpt_temperature,
+                model=self.gpt_model,
+                api_key=SecretStr(self.openai_api_key),
+            )
+        elif llm_ultimas_requests == "deepseek-chat":
+            llm = llm_instance.deepseek()
+        elif llm_ultimas_requests == "gemini-2.0-flash":
+            llm = llm_instance.google_gemini("gemini-2.0-flash")
+        return llm
+
+    async def gerar_documento_final(
+        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:
+            sources, contexts = self.rank_fusion_get_top_results(
+                vector_store, bm25, chunk_ids, query
+            )
 
+            llm = self.select_model_for_last_requests(llm_ultimas_requests)
             prompt_auxiliar = PromptTemplate(
                 template=self.prompt_auxiliar, input_variables=["context"]
             )
@@ -176,21 +190,23 @@ class EnhancedDocumentSummarizer(DocumentSummarizer):
 
             prompt_gerar_documento = PromptTemplate(
                 template=self.prompt_gerar_documento,
-                input_variables=["context"],
+                input_variables=["documento_gerado", "context"],
             )
 
-            documento_gerado = cast(
+            documento_gerado_final = cast(
                 str,
                 llm.invoke(
                     prompt_gerar_documento.format(
-                        context=self.resumo_gerado,
-                        # modelo_usuario=serializer.data["modelo"],
+                        documento_gerado=self.resumo_gerado,
+                        context="\n\n".join(contexts),
                     )
                 ).content,
             )
 
             # Split the response into paragraphs
-            summaries = [p.strip() for p in documento_gerado.split("\n\n") if p.strip()]
+            summaries = [
+                p.strip() for p in documento_gerado_final.split("\n\n") if p.strip()
+            ]
 
             # Create structured output
             structured_output = []

_utils/gerar_relatorio_modelo_usuario/utils.py

@@ -82,7 +82,7 @@ async def get_response_from_auxiliar_contextual_prompt(full_text_as_array: List[
     # )
 
     llms = LLM()
-    response_auxiliar_summary = await llms.googleGemini().ainvoke(
+    response_auxiliar_summary = await llms.google_gemini().ainvoke(
         [HumanMessage(content=prompt_auxiliar_summary)]
     )
 

_utils/ragas.py

@@ -1,6 +1,6 @@
 import os
 from langchain_community.document_loaders import PyPDFLoader
-from _utils.resumo_completo_cursor import EnhancedDocumentSummarizer, RetrievalConfig
+from _utils.resumo_completo_cursor import GerarDocumento, RetrievalConfig
 from rest_framework.response import Response
 from ragas import evaluate
 
@@ -13,7 +13,7 @@ from langchain.memory import SimpleMemory
 
 def test_ragas(serializer, listaPDFs):
 
-    # Step 2: Setup RetrievalConfig and EnhancedDocumentSummarizer
+    # Step 2: Setup RetrievalConfig and GerarDocumento
     config = RetrievalConfig(
         num_chunks=serializer["num_chunks_retrieval"],
         embedding_weight=serializer["embedding_weight"],
@@ -22,7 +22,7 @@ def test_ragas(serializer, listaPDFs):
         chunk_overlap=serializer["chunk_overlap"],
     )
 
-    summarizer = EnhancedDocumentSummarizer(
+    summarizer = GerarDocumento(
         openai_api_key=os.environ.get("OPENAI_API_KEY"),
         claude_api_key=os.environ.get("CLAUDE_API_KEY"),
         config=config,
@@ -74,7 +74,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.generate_enhanced_summary(
+        structured_summaries = summarizer.gerar_documento_final(
             vector_store, bm25, chunk_ids, query
         )
         return {"structured_summaries": structured_summaries}

_utils/resumo_completo_cursor.py

@@ -1,7 +1,7 @@
 import os
 from _utils.gerar_relatorio_modelo_usuario.prompts import prompt_auxiliar_SEM_CONTEXT
-from _utils.gerar_relatorio_modelo_usuario.EnhancedDocumentSummarizer import (
-    EnhancedDocumentSummarizer,
+from _utils.gerar_relatorio_modelo_usuario.GerarDocumento import (
+    GerarDocumento,
 )
 from _utils.gerar_relatorio_modelo_usuario.contextual_retriever import (
     ContextualRetriever,
@@ -59,7 +59,7 @@ async def get_llm_summary_answer_by_cursor_complete(
     )
 
     # Initialize enhanced summarizer
-    summarizer = EnhancedDocumentSummarizer(
+    summarizer = GerarDocumento(
         config=config,
         embedding_model=serializer["hf_embedding"],
         chunk_overlap=serializer["chunk_overlap"],
@@ -100,7 +100,7 @@ async def get_llm_summary_answer_by_cursor_complete(
 
     llm_ultimas_requests = serializer["llm_ultimas_requests"]
     # Generate enhanced summary
-    structured_summaries = await summarizer.generate_enhanced_summary(
+    structured_summaries = await summarizer.gerar_documento_final(
         vector_store, bm25, chunk_ids, llm_ultimas_requests, prompt_auxiliar_SEM_CONTEXT
     )
 

gerar_documento/serializer.py

@@ -46,7 +46,9 @@ class GerarDocumentoSerializer(ResumoCursorSerializer):
     id_modelo_do_usuario = serializers.IntegerField(required=False)
     should_have_contextual_chunks = serializers.BooleanField(default=False)  # type: ignore
     should_use_llama_parse = serializers.BooleanField(required=False, default=False)  # type: ignore
-    llm_ultimas_requests = serializers.CharField(required=False, default="gpt-4o-mini")
+    llm_ultimas_requests = serializers.CharField(
+        required=False, default="gemini-2.0-flash"
+    )
 
 
 class GerarDocumentoComPDFProprioSerializer(ResumoCursorSerializer):

tests/LLMs/test_LLM_class.py

@@ -0,0 +1,35 @@
+import pytest
+import os
+from _utils.LLMs.LLM_class import LLM
+from _utils.splitters.Splitter_class import Splitter
+from _utils.models.gerar_relatorio import (
+    DocumentChunk,
+)
+
+
+class TestSplitters:
+    llm = LLM()
+
+    @pytest.mark.asyncio
+    async def test_google_gemini_modelo_padrao(self, monkeypatch):
+        gemini = self.llm.google_gemini()
+        result = gemini.invoke("Qual a capital do Brasil?").content
+
+        assert isinstance(result, str)
+        assert len(result) > 0
+
+    @pytest.mark.asyncio
+    async def test_google_gemini_modelo_flash_2(self, monkeypatch):
+        gemini = self.llm.google_gemini("gemini-2.0-flash")
+        result = gemini.invoke("Qual a capital do Brasil?").content
+
+        assert isinstance(result, str)
+        assert len(result) > 0
+
+    @pytest.mark.asyncio
+    async def test_deepseek_modelo_padrao(self, monkeypatch):
+        deepseek = self.llm.deepseek()
+        result = deepseek.invoke("Qual a capital do Brasil?").content
+
+        assert isinstance(result, str)
+        assert len(result) > 0