data-science-document-ai 1.61.1__py3-none-any.whl → 1.61.3__py3-none-any.whl

{data_science_document_ai-1.61.1.dist-info → data_science_document_ai-1.61.3.dist-info}/METADATA

@@ -1,13 +1,11 @@
 Metadata-Version: 2.4
 Name: data-science-document-ai
-Version: 1.61.1
+Version: 1.61.3
 Summary: "Document AI repo for data science"
 Author: Naomi Nguyen
 Author-email: naomi.nguyen@forto.com
-Requires-Python: >=3.9,<3.12
+Requires-Python: >=3.11,<3.12
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Requires-Dist: apscheduler (>=3.10.4,<4.0.0)
 Requires-Dist: db-dtypes (>=1.2.0,<2.0.0)
@@ -24,8 +22,8 @@ Requires-Dist: google-cloud-bigquery-storage (>=2.20.0,<3.0.0)
 Requires-Dist: google-cloud-documentai (>=2.23.0,<3.0.0)
 Requires-Dist: google-cloud-storage (>=2.9.0,<3.0.0)
 Requires-Dist: google-cloud-vision (>=3.7.1,<4.0.0)
+Requires-Dist: google-genai (>=1.61.0,<2.0.0)
 Requires-Dist: gspread (>=6.1.0,<7.0.0)
-Requires-Dist: httpx (>=0.26.0,<0.27.0)
 Requires-Dist: jupyter (>=1.0.0,<2.0.0)
 Requires-Dist: kubernetes (>=30.1.0,<31.0.0)
 Requires-Dist: nltk (>=3.9.1,<4.0.0)

{data_science_document_ai-1.61.1.dist-info → data_science_document_ai-1.61.3.dist-info}/RECORD

@@ -1,16 +1,16 @@
-src/constants.py,sha256=JjSdPW8wbo_sI6NgKHZxVsMBaiWAJs4o018fmOfygg4,3537
+src/constants.py,sha256=qXsKAObEJdfEBL5gEqKNk9mVMe0HsSPxpQCVfscn898,3575
 src/constants_sandbox.py,sha256=Iu6HdjCoNSmOX0AwoL9qUQkhq_ZnIN5U9e-Q2UfNuGc,547
 src/docai.py,sha256=dHuR0ehVjUi1CnoNvdp_yxJtpU_HFXqAZ61ywdz7BEo,5655
 src/docai_processor_config.yaml,sha256=ZewXqbyiftzmVAaS08BoNp5trY6WXx3HMWDfPwmKfaI,256
 src/excel_processing.py,sha256=TRgAzSHvL1WKbUgjHtpXL701bPhiWGH7kk3S6e1UPaA,3074
 src/io.py,sha256=rYjXVLlriEacw1uNuPIYhg12bXNu48Qs9GYMY2YcVTE,5563
-src/llm.py,sha256=a7UYA4ITUNjzct_2fHgM-bma_XWc28VC0FV71g9tnUI,7137
+src/llm.py,sha256=I2UOCY1I4LjiMUHs11N-hzrG1f5U2Oma4K5O3q578E0,6762
 src/log_setup.py,sha256=RhHnpXqcl-ii4EJzRt47CF2R-Q3YPF68tepg_Kg7tkw,2895
 src/pdf_processing.py,sha256=Fx-Glb9niEUU3WUCrBZ02ZYV-E2vWoUM0ifN7-0A1Q4,19961
 src/postprocessing/common.py,sha256=uFaJYpctS4vr-0Z3InRyfRZcEar0UWpcTxdB_TDCJ5E,26671
 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=WuaTQK5D09dV_QNrh29ZoKX9IvQn2Ub-WnAMyRjCsvI,14240
+src/postprocessing/postprocess_partner_invoice.py,sha256=ijBh8LRN23HWdotqMw6uf0ro1IaiL_obBiiikMwXYXk,14768
 src/prompts/library/arrivalNotice/other/placeholders.json,sha256=1vzly1amgyKt3jr2JJQbb24kNZsnI289iduvoUo5dJU,3061
 src/prompts/library/arrivalNotice/other/prompt.txt,sha256=QNuU-BvMA8VbdupVNapad4O3WmCotH5cKNxImRMbKDk,2906
 src/prompts/library/bookingConfirmation/evergreen/placeholders.json,sha256=49mmxxjExniMqLkT37zd5q8ILCLPGsfugKOMkR926kk,5854
@@ -52,9 +52,9 @@ src/prompts/library/preprocessing/carrier/prompt.txt,sha256=C2ExzNXXxX1ZU8yGQNah
 src/prompts/library/shippingInstruction/other/placeholders.json,sha256=eK4AeMfORkGMWVYcqH7NjB56Zb4swHTvcQD5UQbTryg,6374
 src/prompts/library/shippingInstruction/other/prompt.txt,sha256=CbrqlKMtB-sVY-8E460KP1KNmz169YVPMrH3-uEldPg,2135
 src/prompts/prompt_library.py,sha256=VJWHeXN-s501C2GiidIIvQQuZdU6T1R27hE2dKBiI40,2555
-src/setup.py,sha256=8-vZWjC8Iwa3xxdk3iR4412VCjtNtgzVqkXcFon7UBE,7309
+src/setup.py,sha256=prJNY3N2qu14ttfCYnvRtpCqdHQeSW34zXwmuAGWX_M,7472
 src/tms.py,sha256=UXbIo1QE--hIX6NZi5Qyp2R_CP338syrY9pCTPrfgnE,1741
 src/utils.py,sha256=Ow5_Jals88o8mbZ1BoHfZpHZoCfig_UQb5aalH-mpWE,17278
-data_science_document_ai-1.61.1.dist-info/METADATA,sha256=knwcr94ybymLp_FaGl4Pu9uRJj-gVWh2lrDFhiPEddg,2152
-data_science_document_ai-1.61.1.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
-data_science_document_ai-1.61.1.dist-info/RECORD,,
+data_science_document_ai-1.61.3.dist-info/METADATA,sha256=Lq6ZHTHHCJJg5zOaI7awo2Uf4Mx1ZC0FgG7XmXKaETk,2058
+data_science_document_ai-1.61.3.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
+data_science_document_ai-1.61.3.dist-info/RECORD,,

src/constants.py

@@ -9,7 +9,7 @@ project_parameters = {
     "g_ai_project_id": "738250249861",
     "g_api_endpoint": "eu-documentai.googleapis.com",
     "g_location": "eu",
-    "g_region": "europe-west1",
+    "g_region": "global",
     # Google Cloud Storage
     "doc_ai_bucket_project_name": "forto-data-science-production",
     "doc_ai_bucket_name": "ds-document-capture",
@@ -60,7 +60,8 @@ project_parameters = {
         "top_p": 0.8,
         "top_k": 40,
         "seed": 42,
-        "model_id": "gemini-2.5-pro",
+        "model_id": "gemini-3-flash-preview",
+        "thinking_level": "medium",
     },
     "gemini_flash_params": {
         "temperature": 0,

src/llm.py

@@ -3,17 +3,11 @@ import logging
 
 logger = logging.getLogger(__name__)
 
-import base64
 import json
 
+from google import genai
+from google.genai import types
 from openai import AsyncOpenAI as OpenAI
-from vertexai.generative_models import (
-    GenerationConfig,
-    GenerativeModel,
-    HarmBlockThreshold,
-    HarmCategory,
-    Part,
-)
 
 from src.io import get_gcp_labels
 from src.utils import cache_on_disk
@@ -24,8 +18,10 @@ from src.utils import cache_on_disk
 class LlmClient:
     """A client for interacting with large language models (LLMs)."""
 
-    def __init__(self, openai_key=None, parameters=None):
+    def __init__(self, openai_key=None, parameters=None, genai_client=None):
         """Initialize the LLM client."""
+        self.genai_client = genai_client
+
         # Initialize the model parameters
         self.model_params = {
             "temperature": parameters.get("temperature", 0),
@@ -35,31 +31,24 @@ class LlmClient:
             "seed": parameters.get("seed", 42),
         }
         self.model_id = parameters.get("model_id", "gemini-2.5-flash")
-        # Initialize the safety configuration
-        self.safety_config = {
-            HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
-            HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
-            HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
-            HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
-        }
-        # Initialize the Gemini client
-        self.geminy_client = self._initialize_gemini()
+        self.thinking_level = parameters.get("thinking_level")
+
+        # Initialize the safety configuration (new format: list of SafetySetting objects)
+        self.safety_settings = [
+            types.SafetySetting(
+                category="HARM_CATEGORY_DANGEROUS_CONTENT", threshold="OFF"
+            ),
+            types.SafetySetting(category="HARM_CATEGORY_HARASSMENT", threshold="OFF"),
+            types.SafetySetting(category="HARM_CATEGORY_HATE_SPEECH", threshold="OFF"),
+            types.SafetySetting(
+                category="HARM_CATEGORY_SEXUALLY_EXPLICIT", threshold="OFF"
+            ),
+        ]
+
         if openai_key is not None:
             # Initialize the ChatGPT client
             self.chatgpt_client = self._create_client_chatgpt(openai_key)
 
-    def _initialize_gemini(self):
-        """Ask the Gemini model a question.
-
-        Returns:
-            str: The response from the model.
-        """
-        # Initialize the model if it is not already initialized
-        model_gen = GenerativeModel(model_name=self.model_id)
-        self.model_config = GenerationConfig(**self.model_params)
-
-        return model_gen
-
     def _create_client_chatgpt(self, openai_key):
         client = OpenAI(api_key=openai_key)
         return client
@@ -67,7 +56,7 @@ class LlmClient:
     async def ask_gemini(
         self,
         prompt: str,
-        document: str = None,
+        document: types.Part = None,
         response_schema: dict = None,
         response_mime_type: str = "application/json",
         doc_type: str = None,
@@ -76,7 +65,7 @@ class LlmClient:
 
         Args:
             prompt (str): The prompt to send to the model.
-            document (str, optional): An optional document to provide context.
+            document (types.Part, 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.
 
@@ -84,33 +73,35 @@ class LlmClient:
             str: The response from the model.
         """
         try:
+            # Build config with all parameters
+            config_params = {
+                **self.model_params,
+                "safety_settings": self.safety_settings,
+                "labels": get_gcp_labels(doc_type=doc_type),
+            }
+
+            # Add thinking config for Gemini 3 models
+            if self.thinking_level:
+                config_params["thinking_config"] = types.ThinkingConfig(
+                    thinking_level=self.thinking_level
+                )
 
-            # Start with the default model configuration
-            config = self.model_config
-
-            # Add response_schema if provided. This is only supported for Gemini 1.5 Flash & Pro models
             if response_schema is not None:
-                config = GenerationConfig(
-                    response_schema=response_schema,
-                    response_mime_type=response_mime_type,
-                    **self.model_params,
-                )
+                config_params["response_schema"] = response_schema
+                config_params["response_mime_type"] = response_mime_type
 
-            # Prepare inputs for the model
-            inputs = [document, prompt] if document else prompt
+            config = types.GenerateContentConfig(**config_params)
+            contents = [document, prompt] if document else prompt
 
-            # Generate the response with labels for cost tracking
+            # Use async client
             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),
+                self.genai_client.aio.models.generate_content,
+                model=self.model_id,
+                contents=contents,
+                config=config,
             )
 
-            response_text = model_response.text
-
-            return response_text
+            return model_response.text
 
         except Exception as e:
             logger.error(f"Failed to generate response: {str(e)}")
@@ -146,25 +137,20 @@ class LlmClient:
             logger.error(e)
             return {}
 
-    def prepare_document_for_gemini(self, file_content):
-        """Prepare a document from file content by encoding it to base64.
+    def prepare_document_for_gemini(self, file_content: bytes) -> types.Part:
+        """Prepare a document from file content for the Gemini model.
 
         Args:
             file_content (bytes): The binary content of the file to be processed.
 
         Returns:
-            Part: A document object ready for processing by the language model.
+            types.Part: A document object ready for processing by the language model.
         """
-        # Convert binary file to base64
-        pdf_base64 = base64.b64encode(file_content).decode("utf-8")
-
-        # Create the document for the model
-        document = Part.from_data(
-            mime_type="application/pdf", data=base64.b64decode(pdf_base64)
+        return types.Part.from_bytes(
+            data=file_content,
+            mime_type="application/pdf",
         )
 
-        return document
-
     async def ask_chatgpt(self, prompt: str, document=None, response_schema=None):
         """Ask the chatgpt model a question.
 

src/postprocessing/postprocess_partner_invoice.py

@@ -119,16 +119,22 @@ def post_process_bundeskasse(aggregated_data):
         if is_forto_recipient(line_item):
             is_recipient_forto = True
 
-    update_recipient_and_vendor(aggregated_data, is_recipient_forto)
+    update_aggregated_data_fields(aggregated_data, is_recipient_forto)
 
 
-def update_recipient_and_vendor(aggregated_data, is_recipient_forto):
-    """Update the recipient and vendor information in the aggregated data."""
+def update_aggregated_data_fields(aggregated_data, is_recipient_forto):
+    """Update the recipient, vendor and deferredDutyPayer information in the aggregated data."""
     # Check if the "recipientName" and "recipientAddress" keys exist
     keys_to_init = ["recipientName", "recipientAddress", "vendorName", "vendorAddress"]
     for key in keys_to_init:
         aggregated_data.setdefault(key, {"formattedValue": "", "documentValue": ""})
 
+    # Update the vendor details always to Bundeskasse Trier
+    aggregated_data["vendorName"]["formattedValue"] = "Bundeskasse Trier"
+    aggregated_data["vendorAddress"][
+        "formattedValue"
+    ] = "Dasbachstraße 15, 54292 Trier, Germany"
+
     if is_recipient_forto:
         # Update the aggregated data with the recipient information
         aggregated_data["recipientName"][
@@ -138,11 +144,17 @@ def update_recipient_and_vendor(aggregated_data, is_recipient_forto):
             "formattedValue"
         ] = "Schönhauser Allee 9, 10119 Berlin, Germany"
 
-    # Update the vendor details always to Bundeskasse Trier
-    aggregated_data["vendorName"]["formattedValue"] = "Bundeskasse Trier"
-    aggregated_data["vendorAddress"][
-        "formattedValue"
-    ] = "Dasbachstraße 15, 54292 Trier, Germany"
+        # Update the defferDutyPayer to Forto Logistics SE & Co KG 'DE789147263644738 - Fort'
+        lineitems = aggregated_data.get("lineItem", [])
+
+        for lineitem in lineitems:
+            if "deferredDutyPayer" in lineitem:
+                lineitem["deferredDutyPayer"] = {
+                    "documentValue": lineitem.get("deferredDutyPayer", {}).get(
+                        "documentValue"
+                    ),
+                    "formattedValue": "DE789147263644738",
+                }
 
 
 def select_unique_bank_account(bank_account):

src/setup.py

@@ -5,8 +5,8 @@ import random
 import time
 
 import toml
-import vertexai
 import yaml
+from google import genai
 from google.api_core.client_options import ClientOptions
 from google.cloud import documentai
 from google.cloud import documentai_v1beta3 as docai_beta
@@ -134,10 +134,14 @@ def setup_params(args=None):
 
     # Set up LLM clients
     params["LlmClient"] = LlmClient(
-        openai_key=os.getenv("OPENAI_KEY"), parameters=params["gemini_params"]
+        openai_key=os.getenv("OPENAI_KEY"),
+        parameters=params["gemini_params"],
+        genai_client=params["genai_client"],
     )
     params["LlmClient_Flash"] = LlmClient(
-        openai_key=os.getenv("OPENAI_KEY"), parameters=params["gemini_flash_params"]
+        openai_key=os.getenv("OPENAI_KEY"),
+        parameters=params["gemini_flash_params"],
+        genai_client=params["genai_client"],
     )
 
     # Load lookup data from GCS bucket
@@ -167,8 +171,9 @@ def setup_docai_client_and_path(params):
 
 
 def setup_vertexai(params):
-    """Initialize the Vertex AI with the specified project and location."""
-    vertexai.init(
+    """Initialize the Google GenAI client with Vertex AI configuration."""
+    params["genai_client"] = genai.Client(
+        vertexai=True,
         project=params["g_ai_project_name"],
         location=params["g_region"],
     )