feat: adicionando OCR em casos de PDFs com problema

_utils/gerar_documento.py

@@ -54,101 +54,91 @@ async def gerar_documento(
         # Initialize enhanced summarizer
         summarizer = GerarDocumento(serializer, axiom_instance)
 
-        all_PDFs_chunks, full_text_as_array, vertex_response = (
-            await get_full_text_and_all_PDFs_chunks(
-                listaPDFs,
-                summarizer.splitter,
-                serializer.should_use_llama_parse,
-                isBubble,
-            )
+        all_PDFs_chunks, full_text_as_array = await get_full_text_and_all_PDFs_chunks(
+            listaPDFs,
+            summarizer.splitter,
+            serializer.should_use_llama_parse,
+            isBubble,
         )
         axiom_instance.send_axiom(
             f"INÍCIO DO TEXTO COMPLETO DOS PDFS: {full_text_as_array[0:5]}"
         )
 
-        if not vertex_response:
-            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_instance.send_axiom(
-                    f"RESUMO INICIAL DO PROCESSO: {response_auxiliar_summary}"
-                )
-
-                axiom_instance.send_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_instance.send_axiom(
-                    "TERMINOU DE FAZER TODAS AS REQUISIÇÕES DO CONTEXTUAL"
-                )
-                chunks_processados = contextualized_chunks
-                axiom_instance.send_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)
-            )
+        is_contextualized_chunk = serializer.should_have_contextual_chunks
 
-            axiom_instance.send_axiom(
-                "COMEÇANDO REQUISIÇÃO PARA GERAR O QUERY DINAMICAMENTE DO VECTOR STORE"
+        if is_contextualized_chunk:
+            response_auxiliar_summary = (
+                await get_response_from_auxiliar_contextual_prompt(full_text_as_array)
             )
-            query_gerado_dinamicamente_para_o_vector_store = (
-                await llm.google_gemini_ainvoke(
-                    prompt_para_gerar_query_dinamico, "gemini-2.0-flash"
-                )
+            axiom_instance.send_axiom(
+                f"RESUMO INICIAL DO PROCESSO: {response_auxiliar_summary}"
             )
 
+            axiom_instance.send_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_instance.send_axiom(
+                "TERMINOU DE FAZER TODAS AS REQUISIÇÕES DO CONTEXTUAL"
+            )
+            chunks_processados = contextualized_chunks
             axiom_instance.send_axiom(
-                f"query_gerado_dinamicamente_para_o_vector_store: {query_gerado_dinamicamente_para_o_vector_store.content}",
+                f"CHUNKS PROCESSADOS INICIALMENTE: {chunks_processados}"
             )
+        else:
+            chunks_processados = all_PDFs_chunks
+
+        if len(chunks_processados) == 0:
+            chunks_processados = all_PDFs_chunks
+        llm = LLM()
+        prompt_para_gerar_query_dinamico = prompt_gerar_query_dinamicamente(
+            cast(str, response_auxiliar_summary)
+        )
 
-            # 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
-                )
+        axiom_instance.send_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"
             )
+        )
+
+        axiom_instance.send_axiom(
+            f"query_gerado_dinamicamente_para_o_vector_store: {query_gerado_dinamicamente_para_o_vector_store.content}",
+        )
 
-            llm_ultimas_requests = serializer.llm_ultimas_requests
-            axiom_instance.send_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,
+        # 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
             )
-            axiom_instance.send_axiom("TERMINOU DE FAZER A ÚLTIMA REQUISIÇÃO")
+        )
 
-            if not isinstance(structured_summaries, list):
-                from rest_framework.response import Response
+        llm_ultimas_requests = serializer.llm_ultimas_requests
+        axiom_instance.send_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_instance.send_axiom("TERMINOU DE FAZER A ÚLTIMA REQUISIÇÃO")
 
-                return Response({"erro": structured_summaries})
+        if not isinstance(structured_summaries, list):
+            from rest_framework.response import Response
 
-            texto_completo = summarizer.resumo_gerado + "\n\n"
+            return Response({"erro": structured_summaries})
 
-            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]
+        texto_completo = summarizer.resumo_gerado + "\n\n"
 
-        else:
-            axiom_instance.send_axiom("FOI UTILIZADO O VERTEX AI DO GOOGLE")
+        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]
 
         texto_completo_como_html = convert_markdown_to_HTML(texto_completo).replace(
             "resposta_segunda_etapa:", "<br><br>"

_utils/gerar_documento_utils/GerarDocumento.py

@@ -4,7 +4,7 @@ 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.LLM_class import LLM, Google_llms
 from _utils.langchain_utils.Prompt_class import Prompt
 from _utils.langchain_utils.Vector_store_class import VectorStore
 from gerar_documento.serializer import (
@@ -26,7 +26,6 @@ from _utils.models.gerar_documento import (
 from cohere import Client
 from _utils.langchain_utils.Splitter_class import Splitter
 import time
-
 from setup.logging import Axiom
 
 

_utils/gerar_documento_utils/utils.py

@@ -106,13 +106,11 @@ async def get_full_text_and_all_PDFs_chunks(
     splitterObject: Splitter,
     should_use_llama_parse: bool,
     isBubble: bool,
-) -> Tuple[List[DocumentChunk], List[str], Union[None, str]]:
+) -> Tuple[List[DocumentChunk], List[str]]:
     all_PDFs_chunks: List[DocumentChunk] = []
 
     pages: List[str] = []
 
-    vertex_response = None  # Só terá valor se for necessário usar Vertex da Google para enviar o pdf e gerar resposta
-
     # Load and process document
     for pdf_path in listaPDFs:
         chunks, pages = await splitterObject.load_and_split_document(
@@ -120,14 +118,7 @@ async def get_full_text_and_all_PDFs_chunks(
         )
         all_PDFs_chunks = all_PDFs_chunks + chunks
 
-    if len(pages) == 0 or len(all_PDFs_chunks) == 0:
-        llm = LLM()
-        prompt = create_prompt_auxiliar_do_contextual_prompt(None)
-        vertex_response = await llm.google_gemini_vertex_ainvoke(
-            prompt, listaPDFs, "gemini-2.0-flash"
-        )
-
-    return all_PDFs_chunks, pages, vertex_response
+    return all_PDFs_chunks, pages
 
 
 async def generate_document_title(resumo_para_gerar_titulo: str):

_utils/google_integration/google_cloud.py

@@ -2,10 +2,12 @@ import os
 from google.cloud import storage
 
 GCP_PROJECT = "gen-lang-client-0350149082"
+GCP_REGION = "us-central1"
+DOCUMENT_API_ID = "b34a20d22dee16bb"
+GCS_BUCKET_NAME = "vella-pdfs"
 
 
 def upload_to_gcs(LOCAL_PDF_PATH: str) -> str:
-    GCS_BUCKET_NAME = "vella-pdfs"
 
     # Path in GCS
     GCS_DESTINATION_BLOB_NAME = "gemini_uploads/" + os.path.basename(LOCAL_PDF_PATH)

_utils/langchain_utils/Splitter_class.py

@@ -1,3 +1,5 @@
+import os
+import time
 from _utils.bubble_integrations.obter_arquivo import get_pdf_from_bubble
 from _utils.handle_files import return_document_list_with_llama_parser
 from _utils.langchain_utils.splitter_util import (
@@ -18,6 +20,16 @@ from _utils.models.gerar_documento import (
     DocumentChunk,
 )
 import uuid
+import json
+from _utils.google_integration.google_cloud import (
+    DOCUMENT_API_ID,
+    GCP_PROJECT,
+    GCP_REGION,
+    GCS_BUCKET_NAME,
+    upload_to_gcs,
+)
+from google.cloud import documentai
+from google.cloud import storage
 
 
 class Splitter:
@@ -34,7 +46,10 @@ class Splitter:
         self.chunk_metadata = {}  # Store chunk metadata for tracing
 
     async def load_and_split_document(
-        self, pdf_path: str, should_use_llama_parse: bool, isBubble: bool
+        self,
+        pdf_path: str,
+        should_use_llama_parse: bool,
+        isBubble: bool,
     ):
         """Load PDF and split into chunks with metadata"""
         # loader = PyPDFLoader(pdf_path)
@@ -144,6 +159,11 @@ class Splitter:
             # char_count += len(text)
         print("TERMINOU DE ORGANIZAR PDFS EM CHUNKS")
 
+        if len(pages) == 0 or len(chunks) == 0:
+            text = await self.getOCRFromGoogleDocumentAPI(pdf_path)
+            chunks = self.load_and_split_text(text)  # type: ignore
+            chunks_of_string_only = [chunk.content for chunk in chunks]
+
         return chunks, chunks_of_string_only
 
     def load_and_split_text(self, text: str) -> List[DocumentChunk]:
@@ -185,3 +205,132 @@ class Splitter:
         char_count += len(text)
 
         return chunks
+
+    async def getOCRFromGoogleDocumentAPI(self, pdf_path: str):
+
+        pdf_gcs_uri = upload_to_gcs(pdf_path)
+
+        GCS_OUTPUT_PREFIX = "documentai_output/"
+        # GCS_INPUT_URI = f"gs://{GCS_BUCKET_NAME}/{f"gemini_uploads/{pdf_gcs_uri}"}"
+        GCS_INPUT_URI = pdf_gcs_uri
+        GCS_OUTPUT_URI = f"gs://{GCS_BUCKET_NAME}/{GCS_OUTPUT_PREFIX}"
+
+        docai_client = documentai.DocumentProcessorServiceClient()
+
+        processor_name = docai_client.processor_path(
+            project=GCP_PROJECT, location="us", processor=DOCUMENT_API_ID
+        )
+
+        gcs_document = documentai.GcsDocument(
+            gcs_uri=GCS_INPUT_URI,
+            mime_type="application/pdf",  # Mime type is specified here for GcsDocument
+        )
+
+        gcs_documents = documentai.GcsDocuments(documents=[gcs_document])
+
+        # 3. Create the BatchDocumentsInputConfig
+        input_config = documentai.BatchDocumentsInputConfig(gcs_documents=gcs_documents)
+        # Note: If GCS_INPUT_URI was a prefix for multiple files, you'd use GcsPrefix:
+        # gcs_prefix = documentai.GcsPrefix(gcs_uri_prefix=GCS_INPUT_URI_PREFIX)
+        # input_config = documentai.BatchDocumentsInputConfig(gcs_prefix=gcs_prefix, mime_type="application/pdf")
+
+        # 4. Create the DocumentOutputConfig
+        # GCS_OUTPUT_URI should be a gs:// URI prefix where the output JSONs will be stored
+        output_config = documentai.DocumentOutputConfig(
+            gcs_output_config=documentai.DocumentOutputConfig.GcsOutputConfig(
+                gcs_uri=GCS_OUTPUT_URI
+            )
+        )
+
+        # 5. Construct the BatchProcessRequest
+        request = documentai.BatchProcessRequest(
+            name=processor_name,
+            input_documents=input_config,  # Use 'input_documents'
+            document_output_config=output_config,  # Use 'document_output_config'
+        )
+
+        # Submit the batch process request (this is a long-running operation)
+        operation = docai_client.batch_process_documents(request)
+
+        print("Batch processing operation started. Waiting for completion...")
+        while not operation.done():
+            time.sleep(15)  # Wait for 30 seconds before checking again
+            print("Waiting...")
+
+        print("Batch processing operation finished.")
+
+        # --- Download the results from GCS ---
+        storage_client = storage.Client(
+            project=GCP_PROJECT
+        )  # Uses GOOGLE_APPLICATION_CREDENTIALS/ADC
+        bucket = storage_client.bucket(GCS_BUCKET_NAME)
+
+        output_blobs = storage_client.list_blobs(
+            GCS_BUCKET_NAME, prefix=GCS_OUTPUT_PREFIX
+        )
+
+        downloaded_files_texts = []
+        try:
+            for blob in output_blobs:
+                # Document AI adds suffixes and subdirectories. Look for the actual JSON output files.
+                # The exact naming depends on the processor and options. Common pattern is ending with .json
+                if blob.name.endswith(".json"):
+                    local_download_path = os.path.basename(
+                        blob.name
+                    )  # Download to current directory with blob name
+                    print(f"Downloading {blob.name} to {local_download_path}...")
+                    blob.download_to_filename(local_download_path)
+
+                    with open(local_download_path, "r", encoding="utf-8") as f:
+                        document_data = json.load(f)
+
+                    # The top-level 'text' field contains the concatenated plain text.
+                    if "text" in document_data and document_data["text"] is not None:
+                        raw_text = document_data["text"]
+                        print(f"\n--- Raw Text Extracted from {blob.name} ---")
+                        # Print only a snippet or process as needed
+                        print(
+                            raw_text[:1000] + "..."
+                            if len(raw_text) > 1000
+                            else raw_text
+                        )
+                        print("--------------------------------------------")
+
+                        return raw_text
+
+                        # Optional: Store the text. If you processed a batch of files,
+                        # you might want to associate the text with the original file name.
+                        # Document AI metadata might link output JSONs back to input files.
+                        # For simplicity here, let's just show the extraction.
+                        # If you know it was a single input PDF, this is all the text.
+                        # If it was multiple, you'd need a mapping or process each JSON.
+
+                    else:
+                        print(
+                            f"Warning: 'text' field not found in {blob.name} or is empty."
+                        )
+
+                    # Optional: Read and print a snippet of the JSON content
+                    # with open(local_download_path, 'r', encoding='utf-8') as f:
+                    #     data = json.load(f)
+                    #     # Print some extracted text, for example (structure varies by processor)
+                    #     if 'text' in data:
+                    #         print(f"Extracted text snippet: {data['text'][:500]}...") # Print first 500 chars
+                    #     elif 'entities' in data:
+                    #         print(f"Number of entities found: {len(data['entities'])}")
+                    #     else:
+                    #         print("Output JSON structure not immediately recognizable.")
+                    # break # Uncomment if you only expect/need to process the first output file
+
+            if len(downloaded_files_texts) == 0 or not downloaded_files_texts:
+                print("No JSON output files found in the specified output location.")
+
+        except Exception as e:
+            print(f"Error listing or downloading output files: {e}")
+
+        print("\nProcess complete.")
+        if downloaded_files_texts:
+            print(f"Downloaded output file(s): {', '.join(downloaded_files_texts)}")
+            print("These files contain the OCR results in JSON format.")
+        else:
+            print("No output files were successfully downloaded.")

_utils/langchain_utils/Vector_store_class.py

@@ -22,6 +22,8 @@ class VectorStore:
         axiom_instance: Axiom,
     ) -> Tuple[Chroma, BM25Okapi, List[str]]:
         """Create vector store and BM25 index with contextualized chunks"""
+        contador_erro = 0
+
         try:
             # Prepare texts with context
             if is_contextualized_chunk:
@@ -69,5 +71,9 @@ class VectorStore:
             return vector_store, bm25, chunk_ids
 
         except Exception as e:
+            contador_erro += 1
+            if contador_erro >= 2:
+                raise Exception(f"Error creating enhanced vector store: {str(e)}")
+
             self.logger.error(f"Error creating enhanced vector store: {str(e)}")
-            raise Exception(f"Error creating enhanced vector store: {str(e)}")
+            return self.create_enhanced_vector_store(chunks, False, axiom_instance)