data-science-document-ai 1.41.0__tar.gz → 1.43.6__tar.gz

{data_science_document_ai-1.41.0 → data_science_document_ai-1.43.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: data-science-document-ai
-Version: 1.41.0
+Version: 1.43.6
 Summary: "Document AI repo for data science"
 Author: Naomi Nguyen
 Author-email: naomi.nguyen@forto.com
@@ -38,7 +38,7 @@ Requires-Dist: pdf2image (>=1.17.0,<2.0.0)
 Requires-Dist: pgzip (>=0.3.5,<0.4.0)
 Requires-Dist: pyarrow (==16.1.0)
 Requires-Dist: pymupdf (>=1.23.26,<2.0.0)
-Requires-Dist: pypdf2 (>=3.0.1,<4.0.0)
+Requires-Dist: pypdf (>=6.1.2,<7.0.0)
 Requires-Dist: python-dotenv (>=1.0.0,<2.0.0)
 Requires-Dist: python-multipart (>=0.0.7,<0.0.8)
 Requires-Dist: rapidfuzz (>=3.12.2,<4.0.0)

{data_science_document_ai-1.41.0 → data_science_document_ai-1.43.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "data-science-document-ai"
-version = "1.41.0"
+version = "1.43.6"
 description = "\"Document AI repo for data science\""
 authors = ["Naomi Nguyen <naomi.nguyen@forto.com>", "Kumar Rajendrababu <kumar.rajendrababu@forto.com>", "Igor Tonko <igor.tonko@forto.com>", "Osman Demirel <osman.demirel@forto.com>"]
 packages = [
@@ -48,7 +48,7 @@ rapidfuzz = "^3.12.2"
 fuzzywuzzy = "^0.18.0"
 nltk = "^3.9.1"
 pgzip = "^0.3.5"
-pypdf2 = "^3.0.1"
+pypdf = "^6.1.2"
 
 [tool.poetry.dev-dependencies]
 jupyter = "^1.0.0"

{data_science_document_ai-1.41.0 → data_science_document_ai-1.43.6}/src/constants.py

@@ -50,8 +50,6 @@ project_parameters = {
     "model_selector": {
         "stable": {
             "bookingConfirmation": 1,
-            "packingList": 0,
-            "commercialInvoice": 0,
             "finalMbL": 0,
             "draftMbl": 0,
             "arrivalNotice": 0,

{data_science_document_ai-1.41.0 → data_science_document_ai-1.43.6}/src/docai.py

@@ -3,11 +3,16 @@ import re
 
 from google.cloud import documentai
 
-from src.io import delete_folder_from_bucket, logger, upload_pdf_to_bucket
+from src.io import (
+    delete_folder_from_bucket,
+    get_gcp_labels,
+    logger,
+    upload_pdf_to_bucket,
+)
 from src.utils import cache_on_disk
 
 
-async def _process_pdf_w_docai(image_content, client, processor_name):
+async def _process_pdf_w_docai(image_content, client, processor_name, doc_type=None):
     """Process the PDF using Document AI.
 
     Args:
@@ -15,6 +20,7 @@ async def _process_pdf_w_docai(image_content, client, processor_name):
         client: The Document AI client.
         processor_name (str): The name of the processor to be used.
                             e.g.: projects/{project_id}/locations/{location}/processor/{processor_id}
+        doc_type (str, optional): Document type for cost tracking labels.
 
     Returns:
         The processed document.
@@ -24,10 +30,11 @@ async def _process_pdf_w_docai(image_content, client, processor_name):
         content=image_content, mime_type="application/pdf"
     )
 
-    # Configure the process request
+    # Configure the process request with labels for cost tracking
     request = documentai.ProcessRequest(
         name=processor_name,
         raw_document=raw_document,  # field_mask=field_mask
+        labels=get_gcp_labels(doc_type=doc_type),
     )
     result = await cache_on_disk(client.process_document, request=request)
 
@@ -35,7 +42,7 @@ async def _process_pdf_w_docai(image_content, client, processor_name):
 
 
 async def _batch_process_pdf_w_docai(
-    params, image_content, client, processor_name, timeout=1200
+    params, image_content, client, processor_name, timeout=1200, doc_type=None
 ):
     """Process the PDF using Document AI Batch Process API.
 
@@ -45,6 +52,7 @@ async def _batch_process_pdf_w_docai(
         processor_name (str): The name of the processor to be used.
                             e.g.: projects/{project_id}/locations/{location}/processor/{processor_id}
         timeout (int, optional): The timeout in seconds. Defaults to 1200.
+        doc_type (str, optional): Document type for cost tracking labels.
 
     Returns:
         The processed document.
@@ -72,11 +80,12 @@ async def _batch_process_pdf_w_docai(
     # Where to write results
     output_config = documentai.DocumentOutputConfig(gcs_output_config=gcs_output_config)
 
-    # The full resource name of the processor
+    # The full resource name of the processor with labels for cost tracking
     request = documentai.BatchProcessRequest(
         name=processor_name,
         input_documents=input_config,
         document_output_config=output_config,
+        labels=get_gcp_labels(doc_type=doc_type),
     )
 
     # BatchProcess returns a Long Running Operation (LRO)

{data_science_document_ai-1.41.0 → data_science_document_ai-1.43.6}/src/docai_processor_config.yaml

@@ -13,20 +13,6 @@ model_config:
           author: "igor.tonko@forto.com"
           created_date: ""
 
-    packingList:
-      - id: "d967005bd9d45aeb"
-        details:
-          display_name: "doc_cap_packingList"
-          author: "kumar.rajendrababu@forto.com"
-          created_date: ""
-
-    commercialInvoice:
-      - id: "7d37236207f75758"
-        details:
-          display_name: "doc_cap_commercialInvoice"
-          author: "kumar.rajendrababu@forto.com"
-          created_date: ""
-
     finalMbL:
       - id: "1eda2f22d64b1b89"
         details:

{data_science_document_ai-1.41.0 → data_science_document_ai-1.43.6}/src/excel_processing.py

@@ -2,23 +2,28 @@
 # flake8: noqa: E402
 import logging
 
+from ddtrace import tracer
+
 from src.postprocessing.common import llm_prediction_to_tuples
 
 logger = logging.getLogger(__name__)
 
 import asyncio
-import json
 
 import numpy as np
 import pandas as pd
 
 from src.llm import prompt_excel_extraction
-from src.utils import generate_schema_structure, get_excel_sheets
+from src.prompts.prompt_library import prompt_library
+from src.utils import estimate_page_count, generate_schema_structure, get_excel_sheets
 
 
-async def extract_data_from_sheet(params, sheet_name, sheet, response_schema):
+async def extract_data_from_sheet(
+    llm_client, sheet_name, sheet, response_schema, doc_type=None
+):
     logger.info(f"Processing sheet: {sheet_name}")
-    excel_content = pd.DataFrame(sheet.values)
+    excel_content = pd.DataFrame(sheet.values).dropna(how="all", axis=1)
+
     # Convert to Markdown format for the LLM model
     worksheet = (
         "This is from a excel. Pay attention to the cell position:\n"
@@ -29,9 +34,10 @@ async def extract_data_from_sheet(params, sheet_name, sheet, response_schema):
     prompt_docai = prompt_excel_extraction(worksheet)
 
     try:
-        result = await params["LlmClient"].get_unified_json_genai(
+        result = await llm_client.get_unified_json_genai(
             prompt_docai,
             response_schema=response_schema,
+            doc_type=doc_type,
         )
     except Exception as e:
         result = {}
@@ -45,6 +51,7 @@ async def extract_data_from_excel(
     input_doc_type,
     file_content,
     mime_type,
+    llm_client,
 ):
     """Extract data from the Excel file.
 
@@ -53,6 +60,7 @@ async def extract_data_from_excel(
         input_doc_type (str): The type of the document.
         file_content (bytes): The content of the Excel file to process.
         mime_type (str): The MIME type of the file.
+        llm_client: The LLM client to use for data extraction.
 
     Returns:
         formatted_data (list): A list of dictionaries containing the extracted data.
@@ -61,22 +69,42 @@ async def extract_data_from_excel(
 
     """
     # Generate the response structure
-    response_schema = generate_schema_structure(params, input_doc_type)
+    response_schema = (
+        prompt_library.library[input_doc_type]["other"]["placeholders"]
+        if input_doc_type
+        in [
+            "partnerInvoice",
+            "customsInvoice",
+            "bundeskasse",
+            "commercialInvoice",
+            "packingList",
+        ]
+        else generate_schema_structure(params, input_doc_type)
+    )
 
     # Load the Excel file and get ONLY the "visible" sheet names
     sheets, workbook = get_excel_sheets(file_content, mime_type)
 
+    # Track the number of sheets in dd-trace
+    span = tracer.current_span()
+    if span:
+        estimated_page_counts = [
+            estimate_page_count(workbook[sheet]) for sheet in sheets
+        ]
+        est_page_count = sum(estimated_page_counts)
+        span.set_metric("est_page_count", est_page_count)
+
     # Excel files may contain multiple sheets. Extract data from each sheet
     sheet_extract_tasks = [
         extract_data_from_sheet(
-            params, sheet_name, workbook[sheet_name], response_schema
+            llm_client,
+            sheet_name,
+            workbook[sheet_name],
+            response_schema,
+            doc_type=input_doc_type,
         )
         for sheet_name in sheets
     ]
     extracted_data = {k: v for k, v in await asyncio.gather(*sheet_extract_tasks)}
 
-    # Convert LLM prediction dictionary to tuples of (value, page_number).
-    extracted_data = llm_prediction_to_tuples(extracted_data)
-    stored_data = json.dumps(extracted_data)
-
-    return extracted_data, stored_data, params["gemini_params"]["model_id"]
+    return extracted_data, extracted_data, llm_client.model_id

{data_science_document_ai-1.41.0 → data_science_document_ai-1.43.6}/src/io.py

@@ -12,13 +12,55 @@ from pathlib import Path
 from google.cloud import bigquery, storage
 
 
+def get_gcp_labels(**extra_labels):
+    """Generate standardized GCP labels for cost tracking.
+
+    Args:
+        **extra_labels: Additional custom labels
+
+    Returns:
+        dict: Labels dictionary with keys normalized (lowercase, hyphens, max 63 chars)
+    """
+    project_name = os.getenv("PROJECT_NAME")
+
+    # If not set, detect once and cache it
+    if not project_name:
+        # Try pyproject.toml first
+        try:
+            import toml
+
+            pyproject_path = Path(__file__).parent.parent / "pyproject.toml"
+            if pyproject_path.exists():
+                config = toml.load(pyproject_path)
+                project_name = config.get("tool", {}).get("poetry", {}).get("name")
+        except Exception:
+            pass
+
+        # Fallback to unknown
+        if not project_name:
+            project_name = "unknown"
+
+        # Cache it
+        os.environ["PROJECT_NAME"] = project_name
+
+    labels = {
+        "ds-project-name": project_name.lower(),
+        "ds-env": os.getenv("CLUSTER", "local").lower(),
+    }
+
+    # Add any extra labels
+    labels.update({k.lower(): str(v).lower() for k, v in extra_labels.items()})
+
+    return labels
+
+
 def get_bq_client(params):
     """Get Google BigQuery client."""
     bq_client = bigquery.Client(project=params["g_ai_project_name"])
     job_config = bigquery.QueryJobConfig(
         allow_large_results=True,
         # flatten_results=True,
-        labels={"project-name": params["project_name"]},
+        labels=get_gcp_labels(),
     )
     return bq_client, job_config
 
@@ -112,3 +154,6 @@ def download_dir_from_bucket(bucket, directory_cloud, directory_local) -> bool:
         Path(directory).mkdir(parents=True, exist_ok=True)
         blob.download_to_filename(directory_local / Path(blob.name))
     return result
+
+
+# type: ignore

{data_science_document_ai-1.41.0 → data_science_document_ai-1.43.6}/src/llm.py

@@ -15,6 +15,7 @@ from vertexai.generative_models import (
     Part,
 )
 
+from src.io import get_gcp_labels
 from src.utils import cache_on_disk
 
 
@@ -69,6 +70,7 @@ class LlmClient:
         document: str = None,
         response_schema: dict = None,
         response_mime_type: str = "application/json",
+        doc_type: str = None,
     ):
         """Ask the Gemini model a question.
 
@@ -76,6 +78,7 @@ class LlmClient:
             prompt (str): The prompt to send to the model.
             document (str, optional): An optional document to provide context.
             response_schema (dict, optional): Defines a specific response schema for the model.
+            doc_type (str, optional): Document type for cost tracking labels.
 
         Returns:
             str: The response from the model.
@@ -96,12 +99,13 @@ class LlmClient:
             # Prepare inputs for the model
             inputs = [document, prompt] if document else prompt
 
-            # Generate the response
+            # Generate the response with labels for cost tracking
             model_response = await cache_on_disk(
                 self.geminy_client.generate_content_async,
                 contents=inputs,
                 generation_config=config,
                 safety_settings=self.safety_config,
+                labels=get_gcp_labels(doc_type=doc_type),
             )
 
             response_text = model_response.text
@@ -113,7 +117,7 @@ class LlmClient:
             return "{}"
 
     async def get_unified_json_genai(
-        self, prompt, document=None, response_schema=None, model="gemini"
+        self, prompt, document=None, response_schema=None, model="gemini", doc_type=None
     ):
         """Send a prompt to a Google Cloud AI Platform model and returns the generated json.
 
@@ -122,6 +126,7 @@ class LlmClient:
             document: Content of the PDF document
             response_schema: The schema to use for the response
             model (str): The model to use for the response ["gemini" or "chatGPT"]. Default is "gemini".
+            doc_type (str, optional): Document type for cost tracking labels.
 
         Returns:
             dict: The generated json from the model.
@@ -131,7 +136,9 @@ class LlmClient:
             response = await self.ask_chatgpt(prompt, document, response_schema)
         else:
             # Default to Gemini
-            response = await self.ask_gemini(prompt, document, response_schema)
+            response = await self.ask_gemini(
+                prompt, document, response_schema, doc_type=doc_type
+            )
 
         try:
             return json.loads(response)

{data_science_document_ai-1.41.0 → data_science_document_ai-1.43.6}/src/pdf_processing.py

@@ -9,12 +9,17 @@ logger = logging.getLogger(__name__)
 import asyncio
 from collections import defaultdict
 
+from ddtrace import tracer
 from fastapi import HTTPException
 from google.cloud.documentai_v1 import Document as docaiv1_document
 
 from src.docai import _batch_process_pdf_w_docai, _process_pdf_w_docai
 from src.excel_processing import extract_data_from_excel
-from src.postprocessing.common import format_all_entities, remove_none_values, llm_prediction_to_tuples
+from src.postprocessing.common import (
+    format_all_entities,
+    llm_prediction_to_tuples,
+    remove_none_values,
+)
 from src.postprocessing.postprocess_booking_confirmation import (
     postprocess_booking_confirmation,
 )
@@ -28,14 +33,17 @@ from src.prompts.prompt_library import prompt_library
 from src.utils import (
     extract_top_pages,
     generate_schema_structure,
+    get_pdf_page_count,
     get_processor_name,
     run_background_tasks,
+    transform_schema_strings,
     validate_based_on_schema,
-    transform_schema_strings
 )
 
 
-async def process_file_w_docai(params, image_content, client, processor_name):
+async def process_file_w_docai(
+    params, image_content, client, processor_name, doc_type=None
+):
     """
     Process a file using Document AI.
 
@@ -44,6 +52,7 @@ async def process_file_w_docai(params, image_content, client, processor_name):
         image_content (bytes): The file to be processed. It can be bytes object.
         client: The Document AI client.
         processor_name (str): The name of the processor to be used.
+        doc_type (str, optional): Document type for cost tracking labels.
 
     Returns:
         The processed document.
@@ -55,7 +64,9 @@ async def process_file_w_docai(params, image_content, client, processor_name):
 
     try:
         logger.info("Processing document...")
-        result = await _process_pdf_w_docai(image_content, client, processor_name)
+        result = await _process_pdf_w_docai(
+            image_content, client, processor_name, doc_type=doc_type
+        )
     except Exception as e:
         if e.reason == "PAGE_LIMIT_EXCEEDED":
             logger.warning(
@@ -64,7 +75,7 @@ async def process_file_w_docai(params, image_content, client, processor_name):
             # Process the document in batch method (offline processing)
             try:
                 result = await _batch_process_pdf_w_docai(
-                    params, image_content, client, processor_name
+                    params, image_content, client, processor_name, doc_type=doc_type
                 )
             except Exception as batch_e:
                 logger.error(f"Error processing document {batch_e}.")
@@ -94,7 +105,7 @@ async def extract_data_from_pdf_w_docai(
         )
 
     result = await process_file_w_docai(
-        params, file_content, processor_client, processor_name
+        params, file_content, processor_client, processor_name, doc_type=input_doc_type
     )
 
     # Create an entity object to store the result in gcs
@@ -105,16 +116,22 @@ async def extract_data_from_pdf_w_docai(
     # Extract entities from the result
     for entity in result.entities:
         value = (
-            {child.type_: (child.mention_text,
-                           child.page_anchor.page_refs[0].page
-                           if hasattr(child.page_anchor.page_refs[0], "page")
-                           else 0)
-             for child in entity.properties}
+            {
+                child.type_: (
+                    child.mention_text,
+                    child.page_anchor.page_refs[0].page
+                    if hasattr(child.page_anchor.page_refs[0], "page")
+                    else 0,
+                )
+                for child in entity.properties
+            }
             if entity.properties
-            else (entity.mention_text,
-                  entity.page_anchor.page_refs[0].page
-                  if hasattr(entity.page_anchor.page_refs[0], "page")
-                  else 0)
+            else (
+                entity.mention_text,
+                entity.page_anchor.page_refs[0].page
+                if hasattr(entity.page_anchor.page_refs[0], "page")
+                else 0,
+            )
         )
         aggregated_data[entity.type_].append(value)
 
@@ -145,7 +162,9 @@ async def extract_data_from_pdf_w_docai(
     return aggregated_data, result_for_store, processor_version
 
 
-async def identify_carrier(document, llm_client, prompt, response_schema):
+async def identify_carrier(
+    document, llm_client, prompt, response_schema, doc_type=None
+):
     """Identify the carrier from the Booking Confirmation document."""
 
     result = await llm_client.ask_gemini(
@@ -153,6 +172,7 @@ async def identify_carrier(document, llm_client, prompt, response_schema):
         document=document,
         response_schema=response_schema,
         response_mime_type="text/x.enum",
+        doc_type=doc_type,
     )
 
     if result:
@@ -180,6 +200,7 @@ async def process_file_w_llm(params, file_content, input_doc_type, llm_client):
         if input_doc_type == "bundeskasse"
         else file_content
     )
+    number_of_pages = get_pdf_page_count(file_content)
 
     # convert file_content to required document
     document = llm_client.prepare_document_for_gemini(file_content)
@@ -187,7 +208,14 @@ async def process_file_w_llm(params, file_content, input_doc_type, llm_client):
     # get the schema placeholder from the Doc AI and generate the response structure
     response_schema = (
         prompt_library.library[input_doc_type]["other"]["placeholders"]
-        if input_doc_type in ["partnerInvoice", "customsInvoice", "bundeskasse"]
+        if input_doc_type
+        in [
+            "partnerInvoice",
+            "customsInvoice",
+            "bundeskasse",
+            "commercialInvoice",
+            "packingList",
+        ]
         else generate_schema_structure(params, input_doc_type)
     )
 
@@ -209,7 +237,11 @@ async def process_file_w_llm(params, file_content, input_doc_type, llm_client):
 
         # identify carrier for customized prompting
         carrier = await identify_carrier(
-            document, llm_client, carrier_prompt, carrier_schema
+            document,
+            llm_client,
+            carrier_prompt,
+            carrier_schema,
+            doc_type=input_doc_type,
         )
 
     if input_doc_type == "bookingConfirmation":
@@ -224,12 +256,22 @@ async def process_file_w_llm(params, file_content, input_doc_type, llm_client):
         # get the related prompt from predefined prompt library
         prompt = prompt_library.library[input_doc_type][carrier]["prompt"]
 
+        # Update schema to extract value-page_number pairs
+        if number_of_pages > 1:
+            response_schema = transform_schema_strings(response_schema)
+
+            # Update the prompt to instruct LLM to include page numbers
+            prompt += "\nFor each field, provide the page number where the information was found. The page numbering starts from 0."
+
         # generate the result with LLM (gemini)
         result = await llm_client.get_unified_json_genai(
-            prompt=prompt, document=document, response_schema=response_schema
+            prompt=prompt,
+            document=document,
+            response_schema=response_schema,
+            doc_type=input_doc_type,
         )
 
-        result = llm_prediction_to_tuples(result)
+        result = llm_prediction_to_tuples(result, number_of_pages)
 
         return result
     return {}
@@ -309,15 +351,9 @@ async def extract_data_by_doctype(
     processor_client,
     if_use_docai,
     if_use_llm,
+    llm_client,
     isBetaTest=False,
 ):
-    # Select LLM client (Using 2.5 Flash model for Bundeskasse)
-    llm_client = (
-        params["LlmClient_Flash"]
-        if input_doc_type == "bundeskasse"
-        else params["LlmClient"]
-    )
-
     async def extract_w_docai():
         return await extract_data_from_pdf_w_docai(
             params=params,
@@ -366,6 +402,7 @@ async def data_extraction_manual_flow(
     meta,
     processor_client,
     schema_client,
+    use_default_logging=False,
 ):
     """
     Process a PDF file and extract data from it.
@@ -386,6 +423,15 @@ async def data_extraction_manual_flow(
     """
     # Get the start time for processing
     start_time = asyncio.get_event_loop().time()
+
+    # Select LLM client (Using 2.5 Pro model only for PI and customsInvoice)
+    llm_client = (
+        params["LlmClient_Flash"]
+        if meta.documentTypeCode not in ["customsInvoice", "partnerInvoice"]
+        else params["LlmClient"]
+    )
+
+    page_count = None
     # Validate the file type
     if mime_type == "application/pdf":
         # Enable Doc Ai only for certain document types.
@@ -407,8 +453,10 @@ async def data_extraction_manual_flow(
             processor_client,
             if_use_docai=if_use_docai,
             if_use_llm=if_use_llm,
+            llm_client=llm_client,
             isBetaTest=False,
         )
+        page_count = get_pdf_page_count(file_content)
 
     elif "excel" in mime_type or "spreadsheet" in mime_type:
         # Extract data from the Excel file
@@ -417,8 +465,19 @@ async def data_extraction_manual_flow(
             input_doc_type=meta.documentTypeCode,
             file_content=file_content,
             mime_type=mime_type,
+            llm_client=llm_client,
         )
 
+        # Get sheet count from dd-trace span (set in extract_data_from_excel)
+        # Note: we use the span metric instead of len(extracted_data) because
+        # some sheets may fail extraction and not appear in extracted_data
+        span = tracer.current_span()
+        page_count = span.get_metric("est_page_count") if span else len(extracted_data)
+        if page_count > 100:
+            logger.warning(
+                f"Check logic. Count of sheets in excel file is weirdly large: {page_count}"
+            )
+
     else:
         raise HTTPException(
             status_code=400,
@@ -426,7 +485,7 @@ async def data_extraction_manual_flow(
         )
     # Create the result dictionary with the extracted data
     extracted_data = await format_all_entities(
-        extracted_data, meta.documentTypeCode, params
+        extracted_data, meta.documentTypeCode, params, mime_type
     )
     result = {
         "id": meta.id,
@@ -441,7 +500,9 @@ async def data_extraction_manual_flow(
     logger.info(f"Time taken to process the document: {round(elapsed_time, 4)} seconds")
 
     # Schedule background tasks without using FastAPI's BackgroundTasks
-    if os.getenv("CLUSTER") != "ode":  # skip data export to bigquery in ODE environment
+    if (
+        os.getenv("CLUSTER") != "ode"
+    ) & use_default_logging:  # skip data export to bigquery in ODE environment
         asyncio.create_task(
             run_background_tasks(
                 params,
@@ -452,6 +513,7 @@ async def data_extraction_manual_flow(
                 processor_version,
                 mime_type,
                 elapsed_time,
+                page_count,
             )
         )
     return result