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

src/pdf_processing.py

@@ -9,6 +9,7 @@ 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
 
@@ -31,9 +32,10 @@ from src.postprocessing.postprocess_partner_invoice import (
 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,
+    split_pdf_into_chunks,
     transform_schema_strings,
     validate_based_on_schema,
 )
@@ -193,38 +195,29 @@ async def process_file_w_llm(params, file_content, input_doc_type, llm_client):
         result (dict): The structured data extracted from the document, formatted as JSON.
     """
     # Bundeskasse invoices contains all the required information in the first 3 pages.
-    file_content = (
-        extract_top_pages(file_content, num_pages=5)
-        if input_doc_type == "bundeskasse"
-        else file_content
-    )
-
-    # convert file_content to required document
-    document = llm_client.prepare_document_for_gemini(file_content)
+    if input_doc_type == "bundeskasse":
+        file_content = extract_top_pages(file_content, num_pages=5)
 
-    # 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"]
-        else generate_schema_structure(params, input_doc_type)
-    )
+    number_of_pages = get_pdf_page_count(file_content)
+    logger.info(f"processing {input_doc_type} with {number_of_pages} pages...")
 
     carrier = "other"
-    if (
-        "preprocessing" in prompt_library.library.keys()
-        and "carrier" in prompt_library.library["preprocessing"].keys()
-        and input_doc_type
-        in prompt_library.library["preprocessing"]["carrier"]["placeholders"].keys()
-    ):
-        carrier_schema = prompt_library.library["preprocessing"]["carrier"][
-            "placeholders"
-        ][input_doc_type]
+    carrier_schema = (
+        prompt_library.library.get("preprocessing", {})
+        .get("carrier", {})
+        .get("placeholders", {})
+        .get(input_doc_type)
+    )
 
+    if carrier_schema:
         carrier_prompt = prompt_library.library["preprocessing"]["carrier"]["prompt"]
         carrier_prompt = carrier_prompt.replace(
             "DOCUMENT_TYPE_PLACEHOLDER", input_doc_type
         )
 
+        # convert file_content to required document
+        document = llm_client.prepare_document_for_gemini(file_content)
+
         # identify carrier for customized prompting
         carrier = await identify_carrier(
             document,
@@ -234,30 +227,122 @@ async def process_file_w_llm(params, file_content, input_doc_type, llm_client):
             doc_type=input_doc_type,
         )
 
-    if input_doc_type == "bookingConfirmation":
-        response_schema = prompt_library.library[input_doc_type][carrier][
-            "placeholders"
-        ]
-
+    # Select prompt
     if (
-        input_doc_type in prompt_library.library.keys()
-        and carrier in prompt_library.library[input_doc_type].keys()
+        input_doc_type not in prompt_library.library
+        or carrier not in prompt_library.library[input_doc_type]
     ):
-        # get the related prompt from predefined prompt library
-        prompt = prompt_library.library[input_doc_type][carrier]["prompt"]
-
-        # generate the result with LLM (gemini)
-        result = await llm_client.get_unified_json_genai(
-            prompt=prompt,
-            document=document,
-            response_schema=response_schema,
-            doc_type=input_doc_type,
+        return {}
+
+    # get the related prompt from predefined prompt library
+    prompt = prompt_library.library[input_doc_type][carrier]["prompt"]
+
+    # get the schema placeholder
+    response_schema = prompt_library.library[input_doc_type][carrier]["placeholders"]
+
+    # Add page-number extraction for moderately large docs
+    use_chunking = number_of_pages >= params["chunk_after"]
+
+    # Update schema and prompt to extract value-page_number pairs
+    if not use_chunking and number_of_pages > 1:
+        response_schema = transform_schema_strings(response_schema)
+        prompt += "\nFor each field, provide the page number where the information was found. The page numbering starts from 0."
+
+    tasks = []
+    # Process in chunks if number of pages exceeds threshold and Process all chunks concurrently
+    for chunk in (
+        split_pdf_into_chunks(file_content, chunk_size=params["chunk_size"])
+        if use_chunking
+        else [file_content]
+    ):
+        tasks.append(
+            process_chunk_with_retry(
+                chunk,
+                prompt,
+                response_schema,
+                llm_client,
+                input_doc_type,
+            )
         )
 
-        result = llm_prediction_to_tuples(result)
+    results = await asyncio.gather(*tasks, return_exceptions=True)
 
-        return result
-    return {}
+    if use_chunking:
+        return merge_llm_results(results, response_schema)
+    else:
+        return llm_prediction_to_tuples(results[0], number_of_pages=number_of_pages)
+
+
+async def process_chunk_with_retry(
+    chunk_content, prompt, response_schema, llm_client, input_doc_type, retries=2
+):
+    """Process a chunk with retries in case of failure."""
+    for attempt in range(1, retries + 1):
+        try:
+            return await process_chunk(
+                chunk_content=chunk_content,
+                prompt=prompt,
+                response_schema=response_schema,
+                llm_client=llm_client,
+                input_doc_type=input_doc_type,
+            )
+        except Exception as e:
+            logger.error(f"Chunk failed on attempt {attempt}: {e}")
+            if attempt == retries:
+                raise
+            await asyncio.sleep(1)  # small backoff
+
+
+async def process_chunk(
+    chunk_content, prompt, response_schema, llm_client, input_doc_type
+):
+    """Process a chunk with Gemini."""
+    document = llm_client.prepare_document_for_gemini(chunk_content)
+    return await llm_client.get_unified_json_genai(
+        prompt=prompt,
+        document=document,
+        response_schema=response_schema,
+        doc_type=input_doc_type,
+    )
+
+
+def merge_llm_results(results, response_schema):
+    """Merge LLM results from multiple chunks."""
+    merged = {}
+    for i, result in enumerate(results):
+        if not isinstance(result, dict):
+            continue
+        # Add page number to all values coming from this chunk
+        result = llm_prediction_to_tuples(result, number_of_pages=1, page_number=i)
+
+        # Merge the result into the final merged dictionary
+        for key, value in result.items():
+            field_type = (
+                response_schema["properties"].get(key, {}).get("type", "").upper()
+            )
+
+            if key not in merged:
+                if field_type == "ARRAY":
+                    # append the values as a list
+                    merged[key] = (
+                        value if isinstance(value, list) else ([value] if value else [])
+                    )
+                else:
+                    merged[key] = value
+                continue
+
+            if field_type == "ARRAY":
+                # append list contents across chunks
+                if isinstance(value, list):
+                    merged[key].extend(value)
+                else:
+                    merged[key].append(value)
+
+            # take first non-null value only
+            if merged[key] in (None, "", [], {}):
+                merged[key] = value
+
+    return merged
 
 
 async def extract_data_from_pdf_w_llm(params, input_doc_type, file_content, llm_client):
@@ -268,8 +353,7 @@ async def extract_data_from_pdf_w_llm(params, input_doc_type, file_content, llm_
     # Add currency from the amount field
     if input_doc_type in ["commercialInvoice"]:
         result = postprocessing_commercial_invoice(result, params, input_doc_type)
-    elif input_doc_type == "bookingConfirmation":
-        result = postprocess_booking_confirmation(result)
+
     return result, llm_client.model_id
 
 
@@ -288,13 +372,14 @@ def combine_llm_results_w_doc_ai(
     Returns:
         combined result
     """
-    result = doc_ai.copy()
-    llm = remove_none_values(llm)
-    if not llm:
+    result = remove_none_values(llm)
+
+    docAi = doc_ai.copy()
+    if not docAi:
         return result
 
     # Merge top-level keys
-    result.update({k: v for k, v in llm.items() if k not in result})
+    result.update({k: v for k, v in docAi.items() if k not in result})
 
     if (
         input_doc_type
@@ -302,28 +387,28 @@ def combine_llm_results_w_doc_ai(
         and keys_to_combine
     ):
         result.update(
-            {key: llm.get(key) for key in keys_to_combine if key in llm.keys()}
+            {key: docAi.get(key) for key in keys_to_combine if key in docAi.keys()}
         )
         return result
 
     # Handle specific key-based merging logic for multiple keys
     if keys_to_combine:
         for key in keys_to_combine:
-            if key in llm.keys():
+            if key in docAi.keys():
                 # Merge the list of dictionaries
-                # If the length of the LLM list is less than the Doc AI result, replace with the LLM list
-                if len(llm[key]) < len(result[key]):
-                    result[key] = llm[key]
+                # If the length of the docAi list is less than the LLM result, replace with the docAi list
+                if len(docAi[key]) < len(result[key]):
+                    result[key] = docAi[key]
                 else:
-                    # If the length of the LLM list is greater than or equal to the Doc AI result,
+                    # If the length of the docAi list is greater than or equal to the LLM result,
                     # add & merge the dictionaries
-                    if isinstance(llm[key], list):
-                        for i in range(len(llm[key])):
+                    if isinstance(docAi[key], list):
+                        for i in range(len(docAi[key])):
                             if i == len(result[key]):
-                                result[key].append(llm[key][i])
+                                result[key].append(docAi[key][i])
                             else:
-                                for sub_key in llm[key][i].keys():
-                                    result[key][i][sub_key] = llm[key][i][sub_key]
+                                for sub_key in docAi[key][i].keys():
+                                    result[key][i][sub_key] = docAi[key][i][sub_key]
     return result
 
 
@@ -334,15 +419,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,
@@ -391,6 +470,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.
@@ -411,15 +491,26 @@ 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":
+        if_use_docai = params["if_use_docai"]
+
         # Enable Doc Ai only for certain document types.
-        if_use_docai = (
-            True if meta.documentTypeCode in params["model_config"]["stable"] else False
-        )
-        if_use_llm = (
-            True if meta.documentTypeCode in params["key_to_combine"].keys() else False
-        )
+        if params["if_use_docai"]:
+            if_use_docai = (
+                True
+                if meta.documentTypeCode in params["model_config"]["stable"]
+                else False
+            )
 
         (
             extracted_data,
@@ -431,9 +522,11 @@ async def data_extraction_manual_flow(
             meta.documentTypeCode,
             processor_client,
             if_use_docai=if_use_docai,
-            if_use_llm=if_use_llm,
+            if_use_llm=params["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
@@ -442,8 +535,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,
@@ -451,7 +555,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,
@@ -466,7 +570,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,
@@ -477,6 +583,7 @@ async def data_extraction_manual_flow(
                 processor_version,
                 mime_type,
                 elapsed_time,
+                page_count,
             )
         )
     return result