data-science-document-ai 1.42.5__py3-none-any.whl → 1.56.1__py3-none-any.whl

src/utils.py

@@ -6,38 +6,29 @@ import json
 import os
 import pickle
 from datetime import datetime
-from typing import Literal
+from typing import Any, Dict, List, Literal, Optional
 
+import httpx
+import numpy as np
 import openpyxl
 import pandas as pd
-import requests
 from google.cloud import documentai_v1beta3 as docu_ai_beta
-from PyPDF2 import PdfReader, PdfWriter
+from pypdf import PdfReader, PdfWriter
 
-from src.io import get_storage_client, logger
+from src.io import bq_logs, get_storage_client, logger
 
 
-def bq_logs(data_to_insert, params):
-    """Insert logs into Google BigQuery.
+def get_pdf_page_count(pdf_bytes):
+    """Get the number of pages in a PDF document efficiently.
 
     Args:
-        data_to_insert (list): The data to insert into BigQuery.
-        params (dict): The parameters dictionary.
+        pdf_bytes (bytes): The PDF content as bytes.
+
+    Returns:
+        int: The number of pages in the PDF.
     """
-    # Use the pre-initialized BigQuery client
-    bq_client = params["bq_client"]
-    # Get the table string
-    table_string = f"{params['g_ai_project_name']}.{params['g_ai_gbq_db_schema']}.{params['g_ai_gbq_db_table_out']}"
-
-    logger.info(f"Log table: {table_string}")
-    # Insert the rows into the table
-    insert_logs = bq_client.insert_rows_json(table_string, data_to_insert)
-
-    # Check if there were any errors inserting the rows
-    if not insert_logs:
-        logger.info("New rows have been added.")
-    else:
-        logger.info("Errors occurred while inserting rows: ", insert_logs)
+    reader = PdfReader(io.BytesIO(pdf_bytes))
+    return len(reader.pages)
 
 
 async def get_data_set_schema_from_docai(
@@ -164,6 +155,7 @@ async def run_background_tasks(
     processor_version,
     mime_type,
     elapsed_time=None,
+    page_count=None,
 ):
     """
     Run background tasks asynchronously.
@@ -177,6 +169,7 @@ async def run_background_tasks(
         processor_version: The processor version used to extract the data.
         mime_type: The MIME type of the document.
         elapsed_time: The time taken to process the document.
+        page_count (int, optional): The number of pages in the document.
 
     Returns:
         None
@@ -185,13 +178,8 @@ async def run_background_tasks(
 
     await loop.run_in_executor(None, store_json_in_gcs, params, doc_id, store_data)
 
-    # Keep the page count as 1 for Excel files.
-    page_count = 1
-    # calculate the number of pages processed for PDFs
-    try:
-        if mime_type == "application/pdf":
-            page_count = len(json.loads(store_data.encode("utf-8"))["pages"])
-    except AttributeError:
+    # Use the passed page_count or default to 0 if not provided
+    if page_count is None:
         page_count = 0
 
     # Log the request in BigQuery
@@ -303,9 +291,6 @@ def generate_schema_structure(params, input_doc_type):
             "type": "string",
         }
 
-    # update schema to extract value-page_number pairs
-    response_schema = transform_schema_strings(response_schema)
-
     return response_schema
 
 
@@ -375,9 +360,9 @@ def extract_top_pages(pdf_bytes, num_pages=4):
     return output.getvalue()
 
 
-def get_tms_mappings(
-    input_list: list[str], embedding_type: str, llm_ports: list[str] = None
-):
+async def get_tms_mappings(
+    input_list: List[str], embedding_type: str, llm_ports: Optional[List[str]] = None
+) -> Dict[str, Any]:
     """Get TMS mappings for the given values.
 
     Args:
@@ -387,39 +372,66 @@ def get_tms_mappings(
         llm_ports (list[str], optional): List of LLM ports to use. Defaults to None.
 
     Returns:
-        dict: A dictionary with the mapping results.
+        dict or string: A dictionary or a string with the mapping results.
     """
-    # To test the API locally, port-forward the embedding service in the sandbox to 8080:80
-    # If you want to launch uvicorn from the tms-embedding repo, then use --port 8080 in the config file
     base_url = (
         "http://0.0.0.0:8080/"
         if os.getenv("CLUSTER") is None
         else "http://tms-mappings.api.svc.cluster.local./"
     )
 
+    # Ensure clean inputs
+    if not input_list:
+        return {}
+
     # Ensure input_list is a list
     if not isinstance(input_list, list):
         input_list = [input_list]
 
     # Always send a dict with named keys
     payload = {embedding_type: input_list}
+
     if llm_ports:
         payload["llm_ports"] = llm_ports if isinstance(llm_ports, list) else [llm_ports]
 
     # Make the POST request to the TMS mappings API
-    url = f"{base_url}/{embedding_type}"
-    response = requests.post(url=url, json=payload)
+    url = f"{base_url}{embedding_type}"
 
-    if response.status_code != 200:
-        logger.error(
-            f"Error from TMS mappings API: {response.status_code} - {response.text}"
-        )
+    # Use a timeout so the code doesn't hang forever
+    timeout = httpx.Timeout(60.0, connect=10.0)
+
+    async with httpx.AsyncClient(timeout=timeout) as client:
+        try:
+            response = await client.post(url, json=payload)
+            response.raise_for_status()
 
-    formatted_values = (
-        response.json().get("response", {}).get("data", {}).get(input_list[0], None)
+            # Structure expected: {"response": {"data": {"desc1": "code1", "desc2": "code2"}}}
+            return response.json().get("response", {}).get("data", {})
+
+        except httpx.HTTPStatusError as exc:
+            logger.error(
+                f"Error from TMS mappings API: {exc.response.status_code} - {exc.response.text}"
+            )
+            return {}
+
+
+async def batch_fetch_all_mappings(container_types, terminals, depots):
+    """Batch fetch all mappings for container types, terminals, and depots."""
+    # run batch calls concurrently
+    results = await asyncio.gather(
+        get_tms_mappings(list(container_types), "container_types"),
+        get_tms_mappings(list(terminals), "terminals"),
+        get_tms_mappings(list(depots), "depots"),
     )
 
-    return formatted_values
+    batch_container_map, batch_terminal_map, batch_depot_map = results
+
+    # Convert lists of tuples to dicts if necessary
+    return (
+        dict(batch_container_map or {}),
+        dict(batch_terminal_map or {}),
+        dict(batch_depot_map or {}),
+    )
 
 
 def transform_schema_strings(schema):
@@ -435,12 +447,23 @@ def transform_schema_strings(schema):
     Returns:
         dict: The transformed schema dictionary.
     """
-    # Base case: if the current schema definition is for a string
-    if isinstance(schema, dict) and schema.get("type").upper() == "STRING":
-        new_schema = {
+    if not isinstance(schema, dict):
+        return schema
+
+    schema_type = schema.get("type")
+    if not schema_type:
+        return schema
+
+    # Base case: STRING → OBJECT (only if not already transformed)
+    if schema_type.upper() == "STRING":
+        return {
             "type": "OBJECT",
             "properties": {
-                "value": {"type": "STRING"},
+                "value": {
+                    "type": "STRING",
+                    "nullable": schema.get("nullable", False),
+                    "description": schema.get("description", ""),
+                },
                 "page_number": {
                     "type": "STRING",
                     "description": "Number of a page where the value was found in the document starting from 0.",
@@ -449,26 +472,55 @@ def transform_schema_strings(schema):
             "required": [],
         }
 
-        # Preserve original properties like nullable and description on the new 'value' key
-        if "nullable" in schema:
-            new_schema["properties"]["value"]["nullable"] = schema["nullable"]
-        if "description" in schema:
-            new_schema["properties"]["value"]["description"] = schema["description"]
+    # Skip already transformed OBJECT (has both 'value' & 'page_number')
+    if (
+        schema_type.upper() == "OBJECT"
+        and "properties" in schema
+        and {"value", "page_number"}.issubset(schema["properties"].keys())
+    ):
+        return schema
+
+    # Recursive case for OBJECT
+    if schema_type.upper() == "OBJECT" and "properties" in schema:
+        new_schema = schema.copy()
+        new_schema["properties"] = {
+            k: transform_schema_strings(v) for k, v in schema["properties"].items()
+        }
+        return new_schema
 
+    # Recursive case for ARRAY
+    if schema_type.upper() == "ARRAY" and "items" in schema:
+        new_schema = schema.copy()
+        new_schema["items"] = transform_schema_strings(schema["items"])
         return new_schema
 
-    # Recursive case: if the schema is a dictionary
-    elif isinstance(schema, dict) and schema.get("type").upper() == "OBJECT":
-        transformed_schema = schema.copy()
-        for key, value in schema.get("properties").items():
-            transformed_schema["properties"][key] = transform_schema_strings(value)
-        return transformed_schema
+    return schema
 
-    # Recursive case: if the schema is a list
-    elif isinstance(schema, dict) and schema.get("type").upper() == "ARRAY":
-        schema["items"] = transform_schema_strings(schema["items"])
-        return schema
 
-    # Base case: for non-dict/list values (e.g., None, bool, str)
+def estimate_page_count(sheet):
+    """Assuming a page is 10 columns x 50 rows."""
+    if hasattr(sheet, "shape"):
+        pg_cnt = sheet.shape[0] * sheet.shape[1]
+    elif hasattr(sheet, "max_row"):
+        pg_cnt = sheet.max_column * sheet.max_row
     else:
-        return schema
+        return None
+    return np.ceil(pg_cnt / 500)
+
+
+def split_pdf_into_chunks(file_content: bytes, chunk_size: int = 1):
+    """Split PDF into smaller page chunks."""
+    pdf = PdfReader(io.BytesIO(file_content))
+    total_pages = len(pdf.pages)
+
+    # TODO: update the chunk_size based on doc length. However, it breaks the page number extraction logic.
+    for i in range(0, total_pages, chunk_size):
+        writer = PdfWriter()
+        for j in range(i, min(i + chunk_size, total_pages)):
+            writer.add_page(pdf.pages[j])
+
+        buffer = io.BytesIO()
+        writer.write(buffer)
+        buffer.seek(0)
+
+        yield buffer.getvalue()

src/prompts/library/draftMbl/hapag-lloyd/prompt.txt

@@ -1,45 +0,0 @@
-You are a document entity extraction specialist. Given a document, the explained datapoint need to extract.
-
-blNumber: Bill of Lading number.
-voyage: The journey or route code taken by the vessel.
-portOfLoading: The port where cargo is loaded.
-portOfDischarge:  The port where cargo is unloaded.
-bookingNumber: A unique identifier for the booking.
-containers:
-    containerType: Type of the shipping container, usually related to it's size.
-    grossWeight: Total weight of the cargo, including the tare weight of the container.
-    measurements: Dimensions of the cargo (length, width, height) for freight calculations.
-    packageQuantity: package quantity.
-    packageType: Type of packaging used (e.g., cartons, pallets, barrels).
-    containerNumber: Unique ID for tracking the shipping container.
-    sealNumber: Number of the container's seal.
-vessel: The name of the vessel.
-
-
-Your task is to extract the text value of the following entities and page numbers starting from 0 where the value was found in the document:
-
-
-Keywords for datapoints:
-- blNumber: Bill of Lading number, bill of landing no., swb-no., b/l no.
-- voyage: voyage, voy. no, voyage-no.
-- portOfLoading: port of loading, pol, from.]
-- portOfDischarge: port of discharge, pod, delivery, to
-- bookingNumber: Our reference, booking no., carrier reference
-- containers:
-    - containerType: x 40' container
-    - grossWeight: gross weight
-    - measurements: Dimensions of the cargo (length, width, height) for freight calculations
-    - packageQuantity: package quantity, number and kind of packages
-    - packageType: Type of packaging used (e.g., cartons, pallets, barrels), number and kind of packages, description of goods
-    - containerNumber: container number, cntr. nos., it is a combination of 4 letters and 7 digits separated by space right above 'SEAL'
-    - sealNumber: seal number, seal nos., shipper seal, seal.
-- vessel: vessel
-
-
-You must apply the following rules:
-- The JSON schema must be followed during the extraction.
-- The values must only include text found in the document
-- Do not normalize any entity value.
-- If 'sealNumber' is not found don't add it to the result.
-- Validate the JSON make sure it is a valid JSON ! No extra text, no missing comma!
-- Add an escape character (backwards slash) in from of all quotes in values

src/prompts/library/draftMbl/maersk/prompt.txt

@@ -1,19 +0,0 @@
-Extract the following information from the sea waybill document.
-Your task is to extract the text value of the following entities and page numbers starting from 0 where the value was found in the document:
-
-
- **blNumber:** Find the value labeled as "B/L No.".
- **voyage:** Get the "Voyage No." value.
- **portOfLoading:** Find the value in the "Port of Loading" field.
- **portOfDischarge:** Extract the text from the "Port of Discharge" field.
- **bookingNumber:** Look for the value associated with "Booking No.".
- **containers:**
-    The document may contain multiple containers listed within the section "PARTICULARS FURNISHED BY SHIPPER" under the line starting with "Kind of Packages; Description of goods; Marks and Numbers; Container No./Seal No.".  Look for container information that starts with a line that includes "Container Said to Contain" and continues until the next instance of "Container Said to Contain" or the end of the section. For each container, extract the following:
-    * **containerType:** Extract the container type information. It is usually a combination of numbers, the word "DRY", and may include additional characters. It is found on the same line as the container number.
-    * **grossWeight:** Find the value corresponding to the "gross weight" of the container. It is usually represented in KGS and is found on the same line as the container number.
-    * **measurements:** Find the value corresponding to the "measurement" of the container. It is usually represented in CBM and is found on the same line as the container number.
-    * **packageQuantity:** Extract the "package quantity" information. It is usually a whole number and precedes the text "PACKAGE".  All container information will be on the same line as the "package quantity".
-    * **packageType:** Extract the value from the "Kind of Packages" field.
-    * **containerNumber:**  Find the container number. It starts with "MRKU" and is followed by a sequence of digits. It is found on the same line as the text "Container Said to Contain".
-    * **sealNumber:** Get the "Shipper Seal" value. It follows after the text "Shipper Seal :".
- **vessel:** Extract the text from the field "Vessel".

src/prompts/library/finalMbL/hapag-lloyd/prompt.txt

@@ -1,44 +0,0 @@
-You are a document entity extraction specialist. Given a document, the explained datapoint need to extract.
-
-blNumber: Bill of Lading number.
-voyage: The journey or route code taken by the vessel.
-portOfLoading: The port where cargo is loaded.
-portOfDischarge:  The port where cargo is unloaded.
-bookingNumber: A unique identifier for the booking.
-containers:
-    containerType: Type of the shipping container, usually related to it's size.
-    grossWeight: Total weight of the cargo, including the tare weight of the container.
-    measurements: Dimensions of the cargo (length, width, height) for freight calculations.
-    packageQuantity: package quantity.
-    packageType: Type of packaging used (e.g., cartons, pallets, barrels).
-    containerNumber: Unique ID for tracking the shipping container.
-    sealNumber: Number of the container's seal.
-vessel: The name of the vessel.
-
-
-Your task is to extract the text value of the following entities and page numbers starting from 0 where the value was found in the document:
-
-Keywords for datapoints:
-- blNumber: Bill of Lading number, bill of landing no., swb-no., b/l no.
-- voyage: voyage, voy. no, voyage-no.
-- portOfLoading: port of loading, pol, from.]
-- portOfDischarge: port of discharge, pod, delivery, to
-- bookingNumber: Our reference, booking no., carrier reference
-- containers:
-    - containerType: x 40' container
-    - grossWeight: gross weight
-    - measurements: Dimensions of the cargo (length, width, height) for freight calculations
-    - packageQuantity: package quantity, number and kind of packages
-    - packageType: Type of packaging used (e.g., cartons, pallets, barrels), number and kind of packages, description of goods
-    - containerNumber: container number, cntr. nos., it is a combination of 4 letters and 7 digits separated by space right above 'SEAL'
-    - sealNumber: seal number, seal nos., shipper seal, seal.
-- vessel: vessel
-
-
-You must apply the following rules:
-- The JSON schema must be followed during the extraction.
-- The values must only include text found in the document
-- Do not normalize any entity value.
-- If 'sealNumber' is not found don't add it to the result.
-- Validate the JSON make sure it is a valid JSON ! No extra text, no missing comma!
-- Add an escape character (backwards slash) in from of all quotes in values

src/prompts/library/finalMbL/maersk/prompt.txt

@@ -1,19 +0,0 @@
-Extract the following information from the sea waybill document.
-Your task is to extract the text value of the following entities and page numbers starting from 0 where the value was found in the document:
-
-
- **blNumber:** Find the value labeled as "B/L No.".
- **voyage:** Get the "Voyage No." value.
- **portOfLoading:** Find the value in the "Port of Loading" field.
- **portOfDischarge:** Extract the text from the "Port of Discharge" field.
- **bookingNumber:** Look for the value associated with "Booking No.".
- **containers:**
-    The document may contain multiple containers listed within the section "PARTICULARS FURNISHED BY SHIPPER" under the line starting with "Kind of Packages; Description of goods; Marks and Numbers; Container No./Seal No.".  Look for container information that starts with a line that includes "Container Said to Contain" and continues until the next instance of "Container Said to Contain" or the end of the section. For each container, extract the following:
-    * **containerType:** Extract the container type information. It is usually a combination of numbers, the word "DRY", and may include additional characters. It is found on the same line as the container number.
-    * **grossWeight:** Find the value corresponding to the "gross weight" of the container. It is usually represented in KGS and is found on the same line as the container number.
-    * **measurements:** Find the value corresponding to the "measurement" of the container. It is usually represented in CBM and is found on the same line as the container number.
-    * **packageQuantity:** Extract the "package quantity" information. It is usually a whole number and precedes the text "PACKAGE".  All container information will be on the same line as the "package quantity".
-    * **packageType:** Extract the value from the "Kind of Packages" field.
-    * **containerNumber:**  Find the container number. It starts with "MRKU" and is followed by a sequence of digits. It is found on the same line as the text "Container Said to Contain".
-    * **sealNumber:** Get the "Shipper Seal" value. It follows after the text "Shipper Seal :".
- **vessel:** Extract the text from the field "Vessel".