Merge branch 'tests' of https://github.com/luanpoppe/vella-backend into feat-refatoracoes-gerais

.env.example

@@ -11,4 +11,6 @@ LLAMA_CLOUD_API_KEY_PEIXE=""
 DEEPSEEKK_API_KEY=""
 GOOGLE_API_KEY_PEIXE=""
 SENTRY_DSN=""
-AMBIENTE="testes"
+AMBIENTE="testes"
+GOOGLE_APPLICATION_CREDENTIALS="" # Só é necessário em ambiente de desenvolvimento que não esteja usando docker
+GCP_CREDENTIALS_JSON_CONTENT="Conteúdo inteiro do arquivo vella_gcp_luan_credentials.json" # Em produção, tem que conter todo o conteúdo do arquivo de credentials. Localmente, não precisa existir

Dockerfile

@@ -3,6 +3,10 @@ FROM python:3.12
 # Instalação necessária para converter arquivos .doc
 RUN apt-get update && apt-get install -y antiword
 
+# Copy the entrypoint script and make it executable
+COPY entrypoint.sh /entrypoint.sh
+RUN chmod +x /entrypoint.sh
+
 RUN useradd -m -u 1000 user
 USER user
 ENV PATH="/home/user/.local/bin:$PATH"
@@ -23,6 +27,10 @@ RUN pip install --no-cache-dir -r requirements.txt
 RUN python manage.py collectstatic --noinput
 
 RUN pip install uvicorn
+
+# Set the entrypoint to our script
+ENTRYPOINT ["/entrypoint.sh"]
+
 CMD ["uvicorn", "setup.asgi:application", "--host", "0.0.0.0", "--port", "7860"]
 
 # ENTRYPOINT ["python", "manage.py", "runserver"]

_utils/gerar_documento_utils/GerarDocumento.py

@@ -13,7 +13,7 @@ from _utils.gerar_documento_utils.prompts import (
     prompt_para_gerar_titulo,
 )
 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 _utils.utils import convert_markdown_to_HTML

_utils/gerar_documento_utils/prompts.py

@@ -1,4 +1,14 @@
-def create_prompt_auxiliar_do_contextual_prompt(PROCESSO_JURIDICO: str):
+def create_prompt_auxiliar_do_contextual_prompt(PROCESSO_JURIDICO: str | None = None):
+    if PROCESSO_JURIDICO:
+        adicionar_ao_prompt = f"""
+1.  **Análise Completa:** Leia e analise todo o conteúdo do processo fornecido.
+<processo_juridico>
+{PROCESSO_JURIDICO}
+</processo_juridico>"""
+    else:
+        adicionar_ao_prompt = """
+1.  **Análise Completa:** Leia e analise todo o conteúdo do processo fornecido como PDF."""
+
     return f"""
 <prompt>
 <persona>
@@ -46,10 +56,7 @@ Seu objetivo é analisar o processo jurídico fornecido e gerar um relatório co
 <instrucoes>
 Siga estritamente os passos abaixo:
 
-1.  **Análise Completa:** Leia e analise todo o conteúdo do processo fornecido.
-    <processo_juridico>
-    {PROCESSO_JURIDICO}
-    </processo_juridico>
+{adicionar_ao_prompt}
 
 2.  **Identificação e Listagem de Peças:** Identifique quais das peças listadas na `<tarefa>` estão presentes no texto. Liste **apenas** as encontradas na tag `<pecas_identificadas>`.
 

_utils/google_integration/google_cloud.py

@@ -0,0 +1,28 @@
+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:
+
+    # Path in GCS
+    GCS_DESTINATION_BLOB_NAME = "gemini_uploads/" + os.path.basename(LOCAL_PDF_PATH)
+
+    """Uploads a file to a GCS bucket and returns its URI."""
+    storage_client = storage.Client(
+        project=GCP_PROJECT,
+    )
+    bucket = storage_client.bucket(GCS_BUCKET_NAME)
+    blob = bucket.blob(GCS_DESTINATION_BLOB_NAME)
+
+    print(
+        f"Uploading {LOCAL_PDF_PATH} to gs://{GCS_BUCKET_NAME}/{GCS_DESTINATION_BLOB_NAME}..."
+    )
+    blob.upload_from_filename(LOCAL_PDF_PATH)
+    gcs_uri = f"gs://{GCS_BUCKET_NAME}/{GCS_DESTINATION_BLOB_NAME}"
+    print(f"File uploaded to {gcs_uri}")
+    return gcs_uri

_utils/langchain_utils/LLM_class.py

@@ -1,9 +1,10 @@
-from typing import Literal, cast
+from typing import List, Literal, cast
 from pydantic import SecretStr
-from setup.environment import default_model
+from _utils.google_integration.google_cloud import GCP_PROJECT, upload_to_gcs
 from setup.easy_imports import ChatOpenAI, ChatGoogleGenerativeAI
 import os
 from langchain_core.messages import HumanMessage
+from langchain_google_vertexai import ChatVertexAI
 
 deepseek_api_key = cast(str, os.environ.get("DEEPSEEKK_API_KEY"))
 google_api_key = cast(str, os.environ.get("GOOGLE_API_KEY_PEIXE"))
@@ -75,3 +76,44 @@ class LLM:
             raise Exception(
                 "Failed to generate the final document after 5 retries and the fallback attempt with chat-gpt-4o-mini."
             ) from e
+
+    async def google_gemini_vertex_ainvoke(
+        self,
+        prompt: str,
+        list_of_pdfs: List[str],
+        model: Google_llms = "gemini-2.5-flash-preview-04-17",
+        max_retries: int = 3,
+    ) -> str | None:
+        message_parts = [
+            {"type": "text", "text": prompt},
+        ]
+        for pdf in list_of_pdfs:
+            pdf_gcs_uri = upload_to_gcs(pdf)
+            message_parts.append(
+                {
+                    # This structure is used for file references via URI
+                    "type": "media",
+                    "mime_type": "application/pdf",  # <-- mime_type moved up
+                    "file_uri": pdf_gcs_uri,  # <-- file_uri moved up
+                }
+            )
+
+        for attempt in range(max_retries):
+            try:
+                llm = ChatVertexAI(
+                    model_name=model,
+                    project=GCP_PROJECT,
+                    location="us-central1",
+                    temperature=0,
+                )
+                response = await llm.ainvoke(
+                    [HumanMessage(content=message_parts)]  # type: ignore
+                )
+
+                if isinstance(response.content, list):
+                    response.content = "\n".join(response.content)  # type: ignore
+
+                return response.content  # type: ignore
+            except Exception as e:
+                model = "gemini-2.0-flash"
+                print(f"Attempt {attempt + 1} failed with error: {e}")

_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)

entrypoint.sh

@@ -0,0 +1,26 @@
+#!/bin/sh
+set -e # Exit immediately if a command exits with a non-zero status.
+
+CREDENTIALS_FILE_PATH="/app/vella_gcp_luan_credentials.json"
+
+# Check if the GCP_CREDENTIALS_JSON_CONTENT secret is provided
+if [ -n "$GCP_CREDENTIALS_JSON_CONTENT" ]; then
+  echo "GCP_CREDENTIALS_JSON_CONTENT secret found. Writing to $CREDENTIALS_FILE_PATH"
+  # Use printf to preserve newlines and special characters correctly
+  printf "%s" "$GCP_CREDENTIALS_JSON_CONTENT" > "$CREDENTIALS_FILE_PATH"
+  export GOOGLE_APPLICATION_CREDENTIALS="$CREDENTIALS_FILE_PATH"
+  echo "GOOGLE_APPLICATION_CREDENTIALS set to $CREDENTIALS_FILE_PATH"
+
+  # Optional: Add a check to see if the file looks like JSON (basic check)
+  if command -v jq >/dev/null && jq -e . "$CREDENTIALS_FILE_PATH" >/dev/null 2>&1; then
+    echo "Credentials file appears to be valid JSON."
+  else
+    echo "Warning: Credentials file may not be valid JSON. Content:"
+    # cat "$CREDENTIALS_FILE_PATH" # Print the content for debugging
+  fi
+
+else
+  echo "Warning: GCP_CREDENTIALS_JSON_CONTENT secret not found. GCP services might not authenticate."
+fi
+
+exec "$@"