data-science-document-ai 1.61.0__py3-none-any.whl → 1.61.2__py3-none-any.whl

{data_science_document_ai-1.61.0.dist-info → data_science_document_ai-1.61.2.dist-info}/METADATA

@@ -1,13 +1,11 @@
 Metadata-Version: 2.4
 Name: data-science-document-ai
-Version: 1.61.0
+Version: 1.61.2
 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.0.dist-info → data_science_document_ai-1.61.2.dist-info}/RECORD

@@ -1,32 +1,32 @@
-src/constants.py,sha256=E1P6CvQ8eoxzI5pDD_4HaNKLnx3SL94Uo4l6Xy5juxM,3545
+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=gcHZTuj_TRG4E0eTTpr7lUjTYKG9N4IiClvqiEoXGZ8,26591
+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/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=rhlxpDmFKhCvLQ9Lq4Ng7Dv6Bo-EbHzSwV6GVqAaewE,5862
+src/prompts/library/bookingConfirmation/evergreen/placeholders.json,sha256=49mmxxjExniMqLkT37zd5q8ILCLPGsfugKOMkR926kk,5854
 src/prompts/library/bookingConfirmation/evergreen/prompt.txt,sha256=dhcqzzTJ9ycMCc5fFznVxW9PwIYMwwxxJ344IIF0vNo,5020
-src/prompts/library/bookingConfirmation/hapag-lloyd/placeholders.json,sha256=nz3Lkmt97DIPIyH7Vd4ZbqF-DqCsnnZjX_ykywCWQc8,6145
+src/prompts/library/bookingConfirmation/hapag-lloyd/placeholders.json,sha256=WpooLsv-1ak4-_2HBmAKjyOCuYa3QYSzZg9AfyZKZrc,6137
 src/prompts/library/bookingConfirmation/hapag-lloyd/prompt.txt,sha256=UDUJC6zKujOc9mAoAJLwnxbsXVrTrMFD4alziwqf0Jc,5196
-src/prompts/library/bookingConfirmation/maersk/placeholders.json,sha256=7dv0WBYcomIW7xZD1VjHNPeDSZEXznbzouJX7n0e0K4,6364
-src/prompts/library/bookingConfirmation/maersk/prompt.txt,sha256=Vazuu0DxdWJpI9r-TrgUZjPmc8CP2jsgJgynf_jpYsU,5688
-src/prompts/library/bookingConfirmation/msc/placeholders.json,sha256=d9QXY_bLiTcw_y57ZutXPlZ-ufZ9T5ipddGN2qyEx1s,6094
-src/prompts/library/bookingConfirmation/msc/prompt.txt,sha256=OBoPd9I2fWifgx-IFstnxQ_g7OMnwZI9qTdxEl_uMW4,5882
-src/prompts/library/bookingConfirmation/oocl/placeholders.json,sha256=U9D1jX98kAA0x9yijGHmvYvRIBbt3B_3wIUvH024yKA,6364
-src/prompts/library/bookingConfirmation/oocl/prompt.txt,sha256=z0iJXIwgsnrZafav7qqDBffz_L_IieEyfoJ1vpKcnyk,4502
-src/prompts/library/bookingConfirmation/other/placeholders.json,sha256=zkpcKzOTClk-RxZJLivWPXNDSCf1eGmlkx3gDmTvG6I,7460
+src/prompts/library/bookingConfirmation/maersk/placeholders.json,sha256=JaOES8YCMudqRujHjNq4r81JHdz1geYRJ5clciySmvs,6356
+src/prompts/library/bookingConfirmation/maersk/prompt.txt,sha256=SyBsu0qjXhRzsPLgNfaftqWAtKo8eoBwsfU9XOZ55N4,5684
+src/prompts/library/bookingConfirmation/msc/placeholders.json,sha256=XAg1JcXiehzCgjWCCIS0F5zrVVlbOpQsjpPGG20ozrs,6086
+src/prompts/library/bookingConfirmation/msc/prompt.txt,sha256=8eiuALucSu9OVsHJNKQc07Cdggv95zyNDZAeSbzCwiI,5878
+src/prompts/library/bookingConfirmation/oocl/placeholders.json,sha256=9riPyLadreieLrz3Zb4Wg5CwwqF8C244UJy8UlNe01k,6356
+src/prompts/library/bookingConfirmation/oocl/prompt.txt,sha256=sAqtzDNwzL45-FCxtFlaa3WQt9gug2LwRaEHZm5WI6U,4498
+src/prompts/library/bookingConfirmation/other/placeholders.json,sha256=Epcgf8r9GHZ9Tn3iBtbLtmqrDfbxbwjyHs49FSlSA_0,7452
 src/prompts/library/bookingConfirmation/other/prompt.txt,sha256=EMmKLDX7ir96KcU7Ca12B7k5-crdioU-Sw3VebX33HY,7867
-src/prompts/library/bookingConfirmation/yangming/placeholders.json,sha256=2fu7q7J4QxrYfmZJlpAfDs9wIHKTheyysIw7i7UmOR0,6254
-src/prompts/library/bookingConfirmation/yangming/prompt.txt,sha256=c_qHnP7f2JTCumDRN0aQ9QpXmglLPRuaMEiDXq3wMJ0,4605
+src/prompts/library/bookingConfirmation/yangming/placeholders.json,sha256=VWx3nLRnMPtRjZZ9TSnjeLckbkGoCDY4pPPsMaIADXI,6246
+src/prompts/library/bookingConfirmation/yangming/prompt.txt,sha256=OG33AZ0DSLhWLFgpO5WCx1H1RumVmT7f_ln-QrBJsd8,4601
 src/prompts/library/bundeskasse/other/placeholders.json,sha256=6VvXUluc_shZNg8jb40JiaDjovNhuoEHpNSfsTzoRs4,4456
 src/prompts/library/bundeskasse/other/prompt.txt,sha256=prmk2BIGlSZ_1vhi5WcBUm91D5zGJJT3fbbLEQ3YeTk,3487
 src/prompts/library/commercialInvoice/other/placeholders.json,sha256=zUK2mg9MnHiEQRYF6VgTiUiq68WGy5f7_4qL63CWyR0,4700
@@ -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.0.dist-info/METADATA,sha256=dmTTQxd4p6VM9mW8O8AgZkuXoCXjClPdWJPXzFmVuRw,2152
-data_science_document_ai-1.61.0.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
-data_science_document_ai-1.61.0.dist-info/RECORD,,
+data_science_document_ai-1.61.2.dist-info/METADATA,sha256=eEkafeJHhdHNOAKuEAmx-3dmS33GHvACepPhKgL0XH4,2058
+data_science_document_ai-1.61.2.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
+data_science_document_ai-1.61.2.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,
@@ -93,8 +94,8 @@ formatting_rules = {
     "bookingConfirmation": {
         "pickUpDepotCode": "depot",
         "dropOffDepotCode": "depot",
-        "gateInTerminalCode": "terminal",
-        "pickUpTerminalCode": "terminal",
+        "gateInTerminal": "terminal",
+        "pickUpTerminal": "terminal",
     },
     "deliveryOrder": {"pickUpTerminal": "terminal", "EmptyContainerDepot": "depot"},
 }

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/common.py

@@ -723,16 +723,23 @@ async def format_all_entities(result, document_type_code, params, mime_type):
     if document_type_code in ["partnerInvoice", "bundeskasse"]:
         await process_partner_invoice(params, aggregated_data, document_type_code)
 
-    # TODO: This is a temporary change until the terminal codes are updated
-    if document_type_code == "bookingConfirmation":
-        if "gateInTerminalCode" in aggregated_data:
-            aggregated_data["gateInTerminal"] = aggregated_data.pop(
-                "gateInTerminalCode"
-            )
-        if "pickUpTerminalCode" in aggregated_data:
-            aggregated_data["pickUpTerminal"] = aggregated_data.pop(
-                "pickUpTerminalCode"
-            )
+    # Post process terminals with code and name
+    for key in list(aggregated_data):
+        if "terminal" not in key.lower():
+            continue
+
+        entity = aggregated_data.pop(key)
+        formatted = entity.get("formattedValue", {})
+
+        if formatted is None:
+            aggregated_data[key] = {**entity}
+            continue
+
+        aggregated_data[key] = {**entity, "formattedValue": formatted.get("id")}
+        aggregated_data[f"{key}Code"] = {
+            **entity,
+            "formattedValue": formatted.get("code"),
+        }
 
     logger.info("Data Extraction completed successfully")
     return aggregated_data

src/prompts/library/bookingConfirmation/evergreen/placeholders.json

@@ -11,12 +11,12 @@
       "nullable": true,
       "description": "It's a contract number between the carrier and Forto Logistics SE & Co KG. Shipment Id 'S' followed by 6, 7, or 8 digits e.g. S9486358 is not a contract number."
     },
-    "pickUpTerminalCode": {
+    "pickUpTerminal": {
       "type": "STRING",
       "nullable": true,
       "description": "The specific terminal for cargo pickup during the import shipment."
     },
-    "gateInTerminalCode": {
+    "gateInTerminal": {
       "type": "STRING",
       "nullable": true,
       "description": "The specific terminal where cargo is gated in especially Export terminal delivery address. E.g., FULL RETURN TO or Export terminal name."

src/prompts/library/bookingConfirmation/hapag-lloyd/placeholders.json

@@ -11,12 +11,12 @@
       "nullable": true,
       "description": "It's a contract number between the carrier and Forto Logistics SE & Co KG. It can be mentioned as Contract No., Contract Reference, or Quotation No. Shipment Id 'S' followed by 6, 7, or 8 digits e.g. S9486358 is not a contract number."
     },
-    "pickUpTerminalCode": {
+    "pickUpTerminal": {
       "type": "STRING",
       "nullable": true,
       "description": "The specific terminal for cargo pickup during the import shipment."
     },
-    "gateInTerminalCode": {
+    "gateInTerminal": {
       "type": "STRING",
       "nullable": true,
       "description": "The specific terminal where cargo is gated in especially Export terminal delivery address. E.g., Export terminal delivery address, Export terminal location, or Export terminal name."

src/prompts/library/bookingConfirmation/maersk/placeholders.json

@@ -11,12 +11,12 @@
       "nullable": true,
       "description": "It's a contract number between the carrier and Forto Logistics SE & Co KG. Shipment Id 'S' followed by 6, 7, or 8 digits e.g. S9486358 is not a contract number."
     },
-    "pickUpTerminalCode": {
+    "pickUpTerminal": {
       "type": "STRING",
       "nullable": true,
       "description": "The specific terminal for cargo pickup during the import shipment."
     },
-    "gateInTerminalCode": {
+    "gateInTerminal": {
       "type": "STRING",
       "nullable": true,
       "description": "The specific terminal where cargo is gated in especially Export terminal delivery address. E.g., Export terminal delivery address, Export terminal location, or Return Equip Delivery Terminal."

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

@@ -16,7 +16,7 @@ For Import Shipment: The loaded container / cargo arrives at a port of discharge
 <INSTRUCTIONS>
 - bookingNumber: A unique identifier for the booking.
 
-- gateInTerminalCode: The specific terminal where cargo is gated in. It can be called Return Equip Delivery Terminal and Location interception. This sometimes can be the same as portOfLoading of the First transportLeg.
+- gateInTerminal: The specific terminal where cargo is gated in. It can be called Return Equip Delivery Terminal and Location interception. This sometimes can be the same as portOfLoading of the First transportLeg.
 - gateInReference: A reference code for cargo entering the terminal. If not mentioned explicitly and gateInTerminal is extracted, then use bookingNumber as gateInReference.
 - pickUpTerminal: The specific terminal for cargo pickup. It can be found as Import pick up address(es),  PORT OF DISCHARGE (after the slash '/').
 - pickUpReference: A reference code for cargo pickup. If not mentioned explicitly and pickUpTerminal is extracted, then use bookingNumber as pickUpReference.

src/prompts/library/bookingConfirmation/msc/placeholders.json

@@ -11,12 +11,12 @@
       "nullable": true,
       "description": "It's a contract number between the carrier and Forto Logistics SE & Co KG. It can be found as Service Contract/Rate Reff. No. Shipment Id 'S' followed by 6, 7, or 8 digits e.g. S9486358 is not a contract number."
     },
-    "pickUpTerminalCode": {
+    "pickUpTerminal": {
       "type": "STRING",
       "nullable": true,
       "description": "The specific terminal for cargo pickup during the import shipment."
     },
-    "gateInTerminalCode": {
+    "gateInTerminal": {
       "type": "STRING",
       "nullable": true,
       "description": "The specific terminal where cargo is gated in especially Export terminal delivery address. It can be found as GATE IN AT TERMINAL/DEPOT"

src/prompts/library/bookingConfirmation/msc/prompt.txt

@@ -14,7 +14,7 @@ For Import Shipment: The loaded container / cargo arrives at a port of discharge
 <CONTEXT>
 
 <INSTRUCTIONS>
-- gateInTerminalCode: The specific terminal where cargo is gated in. It can be called Return Equip Delivery Terminal and Location interception. This sometimes can be the same as portOfLoading of the First transportLeg.
+- gateInTerminal: The specific terminal where cargo is gated in. It can be called Return Equip Delivery Terminal and Location interception. This sometimes can be the same as portOfLoading of the First transportLeg.
 - gateInReference: A reference code for cargo entering the terminal. If not mentioned explicitly and gateInTerminal is extracted, then use bookingNumber as gateInReference.
 - pickUpTerminal: The specific terminal for cargo pickup. It can be found as Import pick up address(es),  PORT OF DISCHARGE (after the slash '/').
 - pickUpReference: A reference code for cargo pickup. If not mentioned explicitly and pickUpTerminal is extracted, then use bookingNumber as pickUpReference.

src/prompts/library/bookingConfirmation/oocl/placeholders.json

@@ -11,12 +11,12 @@
       "nullable": true,
       "description": "It's a contract number between the carrier and Forto Logistics SE & Co KG. It can be referred as Rate Agreement Number or Contract No.. Shipment Id 'S' followed by 6, 7, or 8 digits e.g. S9486358 is not a contract number."
     },
-    "pickUpTerminalCode": {
+    "pickUpTerminal": {
       "type": "STRING",
       "nullable": true,
       "description": "The specific terminal for cargo pickup during the import shipment."
     },
-    "gateInTerminalCode": {
+    "gateInTerminal": {
       "type": "STRING",
       "nullable": true,
       "description": "The specific terminal where cargo is gated in especially Export terminal delivery address. E.g., FULL RETURN LOCATION"

src/prompts/library/bookingConfirmation/oocl/prompt.txt

@@ -19,7 +19,7 @@ For Import Shipment: The loaded container / cargo arrives at a port of discharge
 
 - bookingNumber: Extract the booking number. This information can be found at the top section of the document".
 
-- gateInTerminalCode: The specific terminal where cargo is gated in. It can be called as FULL RETURN LOCATION. This sometimes can be the same as portOfLoading of the First transportLeg.
+- gateInTerminal: The specific terminal where cargo is gated in. It can be called as FULL RETURN LOCATION. This sometimes can be the same as portOfLoading of the First transportLeg.
 - gateInReference: A reference code for cargo entering the terminal. If not mentioned explicitly and gateInTerminal is extracted, then use bookingNumber as gateInReference.
 - pickUpTerminal: The specific terminal for cargo pickup. It can be found as Import pick up address(es),  PORT OF DISCHARGE (after the slash '/').
 - pickUpReference: A reference code for cargo pickup. If not mentioned explicitly and pickUpTerminal is extracted, then use bookingNumber as pickUpReference.

src/prompts/library/bookingConfirmation/other/placeholders.json

@@ -11,12 +11,12 @@
       "nullable": true,
       "description": "It's a contract number between the carrier and Forto Logistics SE & Co KG. Shipment Id 'S' followed by 6, 7, or 8 digits e.g. S9486358 is not a contract number."
     },
-    "pickUpTerminalCode": {
+    "pickUpTerminal": {
       "type": "STRING",
       "nullable": true,
       "description": "The specific terminal for cargo pickup during the import shipment."
     },
-    "gateInTerminalCode": {
+    "gateInTerminal": {
       "type": "STRING",
       "nullable": true,
       "description": "The specific terminal where cargo is gated in especially Export terminal delivery address. It can be found as Gate In At Terminal, Gate in Terminal, Full Return To Terminal, Full Return Location, Cargo Delivery At, Delivery Terminal, Full Return CY, eyc.."

src/prompts/library/bookingConfirmation/yangming/placeholders.json

@@ -11,12 +11,12 @@
       "nullable": true,
       "description": "It's a contract number between the carrier and Forto Logistics SE & Co KG. Shipment Id 'S' followed by 6, 7, or 8 digits e.g. S9486358 is not a contract number."
     },
-    "pickUpTerminalCode": {
+    "pickUpTerminal": {
       "type": "STRING",
       "nullable": true,
       "description": "The specific terminal for cargo pickup during the import shipment."
     },
-    "gateInTerminalCode": {
+    "gateInTerminal": {
       "type": "STRING",
       "nullable": true,
       "description": "The specific terminal where cargo is gated in especially Export terminal delivery address. It is mentioned as Delivery Terminal."

src/prompts/library/bookingConfirmation/yangming/prompt.txt

@@ -17,7 +17,7 @@ For Import Shipment: The loaded container / cargo arrives at a port of discharge
 - Populate fields as defined in the response schema.
 - Use the data field description to understand the context of the data.
 
-- gateInTerminalCode: The specific terminal where cargo is gated in. It is mentioned as Delivery Terminal. This sometimes can be the same as portOfLoading of the First transportLeg.
+- gateInTerminal: The specific terminal where cargo is gated in. It is mentioned as Delivery Terminal. This sometimes can be the same as portOfLoading of the First transportLeg.
 - cyCutOff: The deadline for cargo to be delivered to the Container Yard. It can be found at Cargo Cut Off or FCL delivery cut-off.
 
 - transportLegs: Multiple Transport Legs entries may exist, capture all instances under "transportLegs". Make sure the order of the legs are important.

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"],
     )