data-science-document-ai 1.41.0__py3-none-any.whl → 1.42.1__py3-none-any.whl

{data_science_document_ai-1.41.0.dist-info → data_science_document_ai-1.42.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: data-science-document-ai
-Version: 1.41.0
+Version: 1.42.1
 Summary: "Document AI repo for data science"
 Author: Naomi Nguyen
 Author-email: naomi.nguyen@forto.com

{data_science_document_ai-1.41.0.dist-info → data_science_document_ai-1.42.1.dist-info}/RECORD

@@ -1,13 +1,13 @@
 src/constants.py,sha256=TF_UblovdXZnKIb1lnyJwUqQncJCbzBVihoelI6foSU,3579
 src/constants_sandbox.py,sha256=Iu6HdjCoNSmOX0AwoL9qUQkhq_ZnIN5U9e-Q2UfNuGc,547
-src/docai.py,sha256=AepGdF3ZuSGkujLpewX393FgOBMy-e4sEudiGKho5EA,5280
+src/docai.py,sha256=dHuR0ehVjUi1CnoNvdp_yxJtpU_HFXqAZ61ywdz7BEo,5655
 src/docai_processor_config.yaml,sha256=qOMmCIORpLQ_D-ytvejXxFvER0e0uGYuzPVdZBGv4Pc,2105
-src/excel_processing.py,sha256=wArdSxwxdgyj5WVgVTaWsVSmF7z5zK6rq-bUKGENmo4,2660
-src/io.py,sha256=IXz4wWqiHa9mnHNgtrC6X9M2lItYp9eu6rHCThUIh5c,3585
-src/llm.py,sha256=aEK3rL8XvY7CakvkOJQmcHpEKwZRd8PPrLrzHiO-GFk,7827
+src/excel_processing.py,sha256=jBL6h5T3fJ4uM_rFiV8c0yWAy8Tt3V3RFtBBqb8ztfo,2744
+src/io.py,sha256=tOJpMyI-mP1AaXKG4UFudH47MHWzjWBgVahFJUcjGfs,4749
+src/llm.py,sha256=OE4IEIqcM-hYK9U7e0x1rAfcqdpeo4iXPHBp64L5Qz0,8199
 src/log_setup.py,sha256=RhHnpXqcl-ii4EJzRt47CF2R-Q3YPF68tepg_Kg7tkw,2895
-src/pdf_processing.py,sha256=GNHQl_ryyVOHu3FK39XzPJzOCrn01NNW3E2HO43Ot_c,15836
-src/postprocessing/common.py,sha256=Vj_NohcgWZRCzipnPGeM-rg11wdDJ-wwCR12QeE6qOY,21451
+src/pdf_processing.py,sha256=dxsYvNnONAjzS-T7K5aSo89rz7QcdW3ZDfeuFyeCeII,16294
+src/postprocessing/common.py,sha256=lc95nGvy-KrFFQyX2X3ABMjrx1xVYDjuTBgeAXQTcuU,21570
 src/postprocessing/postprocess_booking_confirmation.py,sha256=nK32eDiBNbauyQz0oCa9eraysku8aqzrcoRFoWVumDU,4827
 src/postprocessing/postprocess_commercial_invoice.py,sha256=3I8ijluTZcOs_sMnFZxfkAPle0UFQ239EMuvZfDZVPg,1028
 src/postprocessing/postprocess_partner_invoice.py,sha256=cM4te4qjOI_bXyrF8Zhb6X7eNf5aMKoRaPCFfqFv-98,11538
@@ -49,9 +49,9 @@ src/prompts/library/preprocessing/carrier/placeholders.json,sha256=1UmrQNqBEsjLI
 src/prompts/library/preprocessing/carrier/prompt.txt,sha256=NLvRZQCZ6aWC1yTr7Q93jK5z7Vi_b4HBaiFYYnIsO-w,134
 src/prompts/library/shippingInstruction/other/prompt.txt,sha256=dT2e-dPuvuz0rVYpwmok_1dWQ2Oa8Qy9NGZ6CCLOUI4,1468
 src/prompts/prompt_library.py,sha256=jPxybNPPGH7mzonqtAOqmw5WcT-RtbGP0pvMqqP22hg,2760
-src/setup.py,sha256=kPSZosrICfaGZeDaajr40Ha7Ok4XK4fo_uq35Omiwr0,7128
+src/setup.py,sha256=M-p5c8M9ejKcSZ9N86VtmtPc4TYLxe1_4_dxf6jpfVc,7262
 src/tms.py,sha256=UXbIo1QE--hIX6NZi5Qyp2R_CP338syrY9pCTPrfgnE,1741
-src/utils.py,sha256=SIEThJlaXWGoWV7236iNoAlabCPNge5oTBpDywTxJw0,15968
-data_science_document_ai-1.41.0.dist-info/METADATA,sha256=MmKbqDbe9voabVucTrE-GoM192GMqFgD09_KNvp6Wsg,2153
-data_science_document_ai-1.41.0.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
-data_science_document_ai-1.41.0.dist-info/RECORD,,
+src/utils.py,sha256=nU69zR3TB7IZmCc19DD8H27Riek8GJAldmhJjCSwNEE,16090
+data_science_document_ai-1.42.1.dist-info/METADATA,sha256=nsGhuml2YNlNF7s7aRUJPpY8psKss8wiLcIavpVInjs,2153
+data_science_document_ai-1.42.1.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
+data_science_document_ai-1.42.1.dist-info/RECORD,,

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)

src/excel_processing.py

@@ -16,7 +16,9 @@ from src.llm import prompt_excel_extraction
 from src.utils import generate_schema_structure, get_excel_sheets
 
 
-async def extract_data_from_sheet(params, sheet_name, sheet, response_schema):
+async def extract_data_from_sheet(
+    params, sheet_name, sheet, response_schema, doc_type=None
+):
     logger.info(f"Processing sheet: {sheet_name}")
     excel_content = pd.DataFrame(sheet.values)
     # Convert to Markdown format for the LLM model
@@ -32,6 +34,7 @@ async def extract_data_from_sheet(params, sheet_name, sheet, response_schema):
         result = await params["LlmClient"].get_unified_json_genai(
             prompt_docai,
             response_schema=response_schema,
+            doc_type=doc_type,
         )
     except Exception as e:
         result = {}
@@ -69,7 +72,11 @@ async def extract_data_from_excel(
     # 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
+            params,
+            sheet_name,
+            workbook[sheet_name],
+            response_schema,
+            doc_type=input_doc_type,
         )
         for sheet_name in sheets
     ]
@@ -77,6 +84,5 @@ async def extract_data_from_excel(
 
     # 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, params["gemini_params"]["model_id"]

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

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)

src/pdf_processing.py

@@ -14,7 +14,11 @@ 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,
 )
@@ -30,12 +34,14 @@ from src.utils import (
     generate_schema_structure,
     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 +50,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 +62,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 +73,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 +103,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 +114,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 +160,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 +170,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:
@@ -209,7 +227,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":
@@ -226,7 +248,10 @@ async def process_file_w_llm(params, file_content, input_doc_type, llm_client):
 
         # 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)

src/postprocessing/common.py

@@ -389,7 +389,9 @@ async def format_label(entity_k, entity_value, document_type_code, params):
     formatted_value = None
 
     if entity_key.startswith("port"):
-        formatted_value = await get_port_code_ai(entity_value, llm_client)
+        formatted_value = await get_port_code_ai(
+            entity_value, llm_client, doc_type=document_type_code
+        )
 
     elif (entity_key == "containertype") or (entity_key == "containersize"):
         formatted_value = get_tms_mappings(entity_value, "container_types")
@@ -484,14 +486,14 @@ async def format_label(entity_k, entity_value, document_type_code, params):
     return entity_k, result
 
 
-async def get_port_code_ai(port: str, llm_client):
+async def get_port_code_ai(port: str, llm_client, doc_type=None):
     """Get port code using AI model."""
-    port_llm = await get_port_code_llm(port, llm_client)
+    port_llm = await get_port_code_llm(port, llm_client, doc_type=doc_type)
 
     return get_tms_mappings(port, "ports", port_llm)
 
 
-async def get_port_code_llm(port: str, llm_client):
+async def get_port_code_llm(port: str, llm_client, doc_type=None):
     if (
         "postprocessing" in prompt_library.library.keys()
         and "port_code" in prompt_library.library["postprocessing"].keys()
@@ -518,7 +520,7 @@ async def get_port_code_llm(port: str, llm_client):
         }
 
         response = await llm_client.get_unified_json_genai(
-            prompt, response_schema=response_schema, model="chatgpt"
+            prompt, response_schema=response_schema, model="chatgpt", doc_type=doc_type
         )
         try:
             mapped_port = response["port"]

src/setup.py

@@ -18,7 +18,7 @@ from src.constants import project_parameters
 from src.constants_sandbox import project_parameters_sandbox
 
 # Parent repos are imported without .
-from src.io import download_dir_from_bucket, get_storage_client, logger
+from src.io import download_dir_from_bucket, get_bq_client, get_storage_client, logger
 from src.llm import LlmClient
 
 
@@ -118,6 +118,10 @@ def setup_params(args=None):
 
     params = setup_docai_client_and_path(params)
 
+    # Set up BigQuery client for logging
+    bq_client, _ = get_bq_client(params)
+    params["bq_client"] = bq_client
+
     # Set up Vertex AI for text embeddings
     setup_vertexai(params)
 

src/utils.py

@@ -14,7 +14,7 @@ import requests
 from google.cloud import documentai_v1beta3 as docu_ai_beta
 from PyPDF2 import PdfReader, PdfWriter
 
-from src.io import get_bq_client, get_storage_client, logger
+from src.io import get_storage_client, logger
 
 
 def bq_logs(data_to_insert, params):
@@ -24,8 +24,8 @@ def bq_logs(data_to_insert, params):
         data_to_insert (list): The data to insert into BigQuery.
         params (dict): The parameters dictionary.
     """
-    # Get the BigQuery client
-    bq_client, config = get_bq_client(params)
+    # 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']}"
 
@@ -139,7 +139,12 @@ def store_json_in_gcs(
         bucket = storage_client.bucket(params.get("doc_ai_bucket_name"))
         full_object_name = folder_path + document_id
         blob = bucket.blob(full_object_name)
-        blob.upload_from_string(json_data, content_type="application/json")
+
+        # Convert dict to JSON string if needed
+        json_string = (
+            json.dumps(json_data) if isinstance(json_data, dict) else json_data
+        )
+        blob.upload_from_string(json_string, content_type="application/json")
 
         logger.info(
             f"JSON object stored successfully in gs://{params.get('doc_ai_bucket_name')}/{full_object_name}"  # noqa
@@ -435,15 +440,13 @@ def transform_schema_strings(schema):
         new_schema = {
             "type": "OBJECT",
             "properties": {
-                "value": {
-                    "type": "STRING"
-                },
+                "value": {"type": "STRING"},
                 "page_number": {
                     "type": "STRING",
-                    "description": "Number of a page where the value was found in the document starting from 0."
-                }
+                    "description": "Number of a page where the value was found in the document starting from 0.",
+                },
             },
-            "required": []
+            "required": [],
         }
 
         # Preserve original properties like nullable and description on the new 'value' key
@@ -468,4 +471,4 @@ def transform_schema_strings(schema):
 
     # Base case: for non-dict/list values (e.g., None, bool, str)
     else:
-        return schema
+        return schema