ws-bom-robot-app 0.0.99__py3-none-any.whl → 0.0.101__py3-none-any.whl

ws_bom_robot_app/llm/vector_store/integration/thron.py

@@ -1,236 +1,236 @@
-import asyncio, logging, aiohttp
-from ws_bom_robot_app.llm.vector_store.integration.base import IntegrationStrategy
-from langchain_core.documents import Document
-from ws_bom_robot_app.llm.vector_store.loader.base import Loader
-from typing import List, Union, Optional
-from pydantic import BaseModel, Field, AliasChoices
-import json
-import os
-import platform
-import pandas as pd
-from io import BytesIO
-
-# Fix for Windows event loop issue with aiodns
-if platform.system() == 'Windows':
-    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
-
-class ThronParams(BaseModel):
-  """
-  ThronParams is a model that defines the parameters required for Thron integration.
-
-  Attributes:
-    app_id (str): The application ID for Thron.
-    client_id (str): The client ID for Thron.
-    client_secret (str): The client secret for Thron.
-  """
-  organization_name: str = Field(validation_alias=AliasChoices("organizationName","organization_name"))
-  attribute_fields: Optional[List[str]] = Field(default=None, validation_alias=AliasChoices("attributeFields","attribute_fields"))
-  client_id: str = Field(validation_alias=AliasChoices("clientId","client_id"))
-  client_secret: str = Field(validation_alias=AliasChoices("clientSecret","client_secret"))
-
-class Thron(IntegrationStrategy):
-  def __init__(self, knowledgebase_path: str, data: dict[str, Union[str,int,list]]):
-    super().__init__(knowledgebase_path, data)
-    self.__token = None
-    self.__data = ThronParams.model_validate(self.data)
-
-  def working_subdirectory(self) -> str:
-    return 'thron'
-
-  async def __ensure_token(self) -> bool:
-    """Ensure we have a valid token, getting one if needed."""
-    if not self.__token:
-      self.__token = await self.__get_auth_token()
-    return self.__token is not None
-
-  def __convert_xlsx_to_csv(self, file_content: bytes) -> bool:
-    """Convert XLSX file content to CSV and save to working directory."""
-    try:
-      df = pd.read_excel(BytesIO(file_content))
-      csv_path = os.path.join(self.working_directory, 'thron_export.csv')
-      df.to_csv(csv_path, index=False, encoding='utf-8')
-      return True
-    except Exception as e:
-      logging.error(f"Error converting XLSX to CSV: {e}")
-      return False
-
-  async def run(self) -> None:
-    _run_id = await self.__get_data()
-    if _run_id:
-      await self.__fetch_exported_file(_run_id)
-
-  async def load(self) -> list[Document]:
-    await self.run()
-    await asyncio.sleep(1)
-    return await Loader(self.working_directory).load()
-
-  async def __get_auth_token(self) -> str:
-    """
-    Get authentication token from Thron API.
-
-    Returns:
-      str: The access token if successful, None otherwise.
-    """
-    try:
-      async with aiohttp.ClientSession() as session:
-        auth_data = {
-          "grant_type": "client_credentials",
-          "client_id": self.__data.client_id,
-          "client_secret": self.__data.client_secret
-        }
-        headers = {
-          "accept": "application/json",
-          "Content-Type": "application/x-www-form-urlencoded"
-        }
-        async with session.post(f"https://{self.__data.organization_name}.thron.com/api/v1/authentication/oauth2/token", data=auth_data, headers=headers) as response:
-          result = await response.json()
-          return result.get("access_token", "")
-    except Exception as e:
-      logging.error(f"Error fetching Thron auth token: {e}")
-      return None
-
-  async def __refresh_token(self) -> bool:
-    """Refresh the authentication token and update the instance variable."""
-    try:
-      new_token = await self.__get_auth_token()
-      if new_token:
-        self.__token = new_token
-        logging.info("Thron authentication token refreshed successfully.")
-        return True
-      else:
-        logging.error("Failed to refresh Thron authentication token.")
-        return False
-    except Exception as e:
-      logging.error(f"Error refreshing Thron auth token: {e}")
-      return False
-
-  async def __get_data(self) -> str:
-    """
-    Initiates a data export request to Thron API.
-
-    Returns:
-      str: The export ID if successful, None otherwise.
-    """
-    max_retries = 2
-    retry_count = 0
-
-    while retry_count < max_retries:
-      try:
-        if not await self.__ensure_token():
-          logging.error("Failed to obtain Thron authentication token.")
-          return {}
-
-        async with aiohttp.ClientSession() as session:
-          headers = {
-            "accept": "application/json",
-            "Authorization": f"Bearer {self.__token}"
-          }
-          payload = {"attributes": self.__data.attribute_fields or [],"assetsBy":"CODE","type":"CODES","format":"XLSX","locales":[],"systemAttributes":["family","master","variation","variationGroup","hierarchyLevel"]}
-          async with session.post(f"https://{self.__data.organization_name}.thron.com/api/v1/product-sync/exports", headers=headers, json=payload) as response:
-            # Check for authentication errors
-            if response.status == 401:
-              logging.warning("Authentication failed in __get_data, attempting to refresh token...")
-              if await self.__refresh_token():
-                retry_count += 1
-                continue
-              else:
-                logging.error("Token refresh failed in __get_data.")
-                return None
-
-            if response.status not in range(200, 300):
-              logging.error(f"API request failed with status {response.status}")
-              return None
-
-            result = await response.json()
-            return result.get("id", None)
-
-      except Exception as e:
-        logging.error(f"Error fetching Thron product data (attempt {retry_count + 1}): {e}")
-        if retry_count < max_retries - 1:
-          if await self.__refresh_token():
-            retry_count += 1
-            continue
-        retry_count += 1
-
-    logging.error(f"Failed to fetch Thron product data after {max_retries} attempts.")
-    return {}
-
-
-  async def __fetch_exported_file(self, export_id: str) -> bool:
-    """
-    Fetches the exported file from Thron API using the provided export ID.
-    Polls the export status until it's processed, then downloads the XLSX file
-    and converts it to CSV format in the working directory.
-
-    Args:
-      export_id (str): The ID of the export to fetch.
-
-    Returns:
-      bool: True if file was successfully downloaded and converted, False otherwise.
-    """
-    max_retries = 2
-    retry_count = 0
-
-    while retry_count < max_retries:
-      try:
-        # Ensure we have a token
-        if not await self.__ensure_token():
-          logging.error("Failed to obtain Thron authentication token.")
-          return {}
-
-        async with aiohttp.ClientSession() as session:
-          headers = {
-            "accept": "application/json",
-            "Authorization": f"Bearer {self.__token}"
-          }
-
-          # Polling until status is PROCESSED
-          while True:
-            async with session.get(f"https://{self.__data.organization_name}.thron.com/api/v1/product-sync/exports/{export_id}", headers=headers) as response:
-              # Check for authentication errors
-              if response.status == 401:
-                logging.warning("Authentication failed, attempting to refresh token...")
-                if await self.__refresh_token():
-                  headers["Authorization"] = f"Bearer {self.__token}"
-                  continue
-                else:
-                  logging.error("Token refresh failed, aborting request.")
-                  return {}
-
-              if response.status != 200:
-                logging.error(f"API request failed with status {response.status}")
-                break
-
-              result = await response.json()
-              if result.get("status") == "PROCESSED":
-                download_uri = result.get("downloadUri")
-                if download_uri:
-                  async with session.get(download_uri) as file_response:
-                    if file_response.status == 200:
-                      # Download XLSX file
-                      file_content = await file_response.read()
-                      return self.__convert_xlsx_to_csv(file_content)
-
-                    elif file_response.status == 401:
-                      logging.warning("Authentication failed during file download, attempting to refresh token...")
-                      if await self.__refresh_token():
-                        retry_count += 1
-                        break
-                      else:
-                        logging.error("Token refresh failed during file download.")
-                        return False
-                break
-
-              await asyncio.sleep(5)
-        return False
-
-      except Exception as e:
-        logging.error(f"Error fetching exported data (attempt {retry_count + 1}): {e}")
-        if retry_count < max_retries - 1:
-          if await self.__refresh_token():
-            retry_count += 1
-            continue
-        retry_count += 1
-
-    logging.error(f"Failed to fetch exported data after {max_retries} attempts.")
-    return False
+import asyncio, logging, aiohttp
+from ws_bom_robot_app.llm.vector_store.integration.base import IntegrationStrategy
+from langchain_core.documents import Document
+from ws_bom_robot_app.llm.vector_store.loader.base import Loader
+from typing import List, Union, Optional
+from pydantic import BaseModel, Field, AliasChoices
+import json
+import os
+import platform
+import pandas as pd
+from io import BytesIO
+
+# Fix for Windows event loop issue with aiodns
+if platform.system() == 'Windows':
+    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+
+class ThronParams(BaseModel):
+  """
+  ThronParams is a model that defines the parameters required for Thron integration.
+
+  Attributes:
+    app_id (str): The application ID for Thron.
+    client_id (str): The client ID for Thron.
+    client_secret (str): The client secret for Thron.
+  """
+  organization_name: str = Field(validation_alias=AliasChoices("organizationName","organization_name"))
+  attribute_fields: Optional[List[str]] = Field(default=None, validation_alias=AliasChoices("attributeFields","attribute_fields"))
+  client_id: str = Field(validation_alias=AliasChoices("clientId","client_id"))
+  client_secret: str = Field(validation_alias=AliasChoices("clientSecret","client_secret"))
+
+class Thron(IntegrationStrategy):
+  def __init__(self, knowledgebase_path: str, data: dict[str, Union[str,int,list]]):
+    super().__init__(knowledgebase_path, data)
+    self.__token = None
+    self.__data = ThronParams.model_validate(self.data)
+
+  def working_subdirectory(self) -> str:
+    return 'thron'
+
+  async def __ensure_token(self) -> bool:
+    """Ensure we have a valid token, getting one if needed."""
+    if not self.__token:
+      self.__token = await self.__get_auth_token()
+    return self.__token is not None
+
+  def __convert_xlsx_to_csv(self, file_content: bytes) -> bool:
+    """Convert XLSX file content to CSV and save to working directory."""
+    try:
+      df = pd.read_excel(BytesIO(file_content))
+      csv_path = os.path.join(self.working_directory, 'thron_export.csv')
+      df.to_csv(csv_path, index=False, encoding='utf-8')
+      return True
+    except Exception as e:
+      logging.error(f"Error converting XLSX to CSV: {e}")
+      return False
+
+  async def run(self) -> None:
+    _run_id = await self.__get_data()
+    if _run_id:
+      await self.__fetch_exported_file(_run_id)
+
+  async def load(self) -> list[Document]:
+    await self.run()
+    await asyncio.sleep(1)
+    return await Loader(self.working_directory).load()
+
+  async def __get_auth_token(self) -> str:
+    """
+    Get authentication token from Thron API.
+
+    Returns:
+      str: The access token if successful, None otherwise.
+    """
+    try:
+      async with aiohttp.ClientSession() as session:
+        auth_data = {
+          "grant_type": "client_credentials",
+          "client_id": self.__data.client_id,
+          "client_secret": self.__data.client_secret
+        }
+        headers = {
+          "accept": "application/json",
+          "Content-Type": "application/x-www-form-urlencoded"
+        }
+        async with session.post(f"https://{self.__data.organization_name}.thron.com/api/v1/authentication/oauth2/token", data=auth_data, headers=headers) as response:
+          result = await response.json()
+          return result.get("access_token", "")
+    except Exception as e:
+      logging.error(f"Error fetching Thron auth token: {e}")
+      return None
+
+  async def __refresh_token(self) -> bool:
+    """Refresh the authentication token and update the instance variable."""
+    try:
+      new_token = await self.__get_auth_token()
+      if new_token:
+        self.__token = new_token
+        logging.info("Thron authentication token refreshed successfully.")
+        return True
+      else:
+        logging.error("Failed to refresh Thron authentication token.")
+        return False
+    except Exception as e:
+      logging.error(f"Error refreshing Thron auth token: {e}")
+      return False
+
+  async def __get_data(self) -> str:
+    """
+    Initiates a data export request to Thron API.
+
+    Returns:
+      str: The export ID if successful, None otherwise.
+    """
+    max_retries = 2
+    retry_count = 0
+
+    while retry_count < max_retries:
+      try:
+        if not await self.__ensure_token():
+          logging.error("Failed to obtain Thron authentication token.")
+          return {}
+
+        async with aiohttp.ClientSession() as session:
+          headers = {
+            "accept": "application/json",
+            "Authorization": f"Bearer {self.__token}"
+          }
+          payload = {"attributes": self.__data.attribute_fields or [],"assetsBy":"CODE","type":"CODES","format":"XLSX","locales":[],"systemAttributes":["family","master","variation","variationGroup","hierarchyLevel"]}
+          async with session.post(f"https://{self.__data.organization_name}.thron.com/api/v1/product-sync/exports", headers=headers, json=payload) as response:
+            # Check for authentication errors
+            if response.status == 401:
+              logging.warning("Authentication failed in __get_data, attempting to refresh token...")
+              if await self.__refresh_token():
+                retry_count += 1
+                continue
+              else:
+                logging.error("Token refresh failed in __get_data.")
+                return None
+
+            if response.status not in range(200, 300):
+              logging.error(f"API request failed with status {response.status}")
+              return None
+
+            result = await response.json()
+            return result.get("id", None)
+
+      except Exception as e:
+        logging.error(f"Error fetching Thron product data (attempt {retry_count + 1}): {e}")
+        if retry_count < max_retries - 1:
+          if await self.__refresh_token():
+            retry_count += 1
+            continue
+        retry_count += 1
+
+    logging.error(f"Failed to fetch Thron product data after {max_retries} attempts.")
+    return {}
+
+
+  async def __fetch_exported_file(self, export_id: str) -> bool:
+    """
+    Fetches the exported file from Thron API using the provided export ID.
+    Polls the export status until it's processed, then downloads the XLSX file
+    and converts it to CSV format in the working directory.
+
+    Args:
+      export_id (str): The ID of the export to fetch.
+
+    Returns:
+      bool: True if file was successfully downloaded and converted, False otherwise.
+    """
+    max_retries = 2
+    retry_count = 0
+
+    while retry_count < max_retries:
+      try:
+        # Ensure we have a token
+        if not await self.__ensure_token():
+          logging.error("Failed to obtain Thron authentication token.")
+          return {}
+
+        async with aiohttp.ClientSession() as session:
+          headers = {
+            "accept": "application/json",
+            "Authorization": f"Bearer {self.__token}"
+          }
+
+          # Polling until status is PROCESSED
+          while True:
+            async with session.get(f"https://{self.__data.organization_name}.thron.com/api/v1/product-sync/exports/{export_id}", headers=headers) as response:
+              # Check for authentication errors
+              if response.status == 401:
+                logging.warning("Authentication failed, attempting to refresh token...")
+                if await self.__refresh_token():
+                  headers["Authorization"] = f"Bearer {self.__token}"
+                  continue
+                else:
+                  logging.error("Token refresh failed, aborting request.")
+                  return {}
+
+              if response.status != 200:
+                logging.error(f"API request failed with status {response.status}")
+                break
+
+              result = await response.json()
+              if result.get("status") == "PROCESSED":
+                download_uri = result.get("downloadUri")
+                if download_uri:
+                  async with session.get(download_uri) as file_response:
+                    if file_response.status == 200:
+                      # Download XLSX file
+                      file_content = await file_response.read()
+                      return self.__convert_xlsx_to_csv(file_content)
+
+                    elif file_response.status == 401:
+                      logging.warning("Authentication failed during file download, attempting to refresh token...")
+                      if await self.__refresh_token():
+                        retry_count += 1
+                        break
+                      else:
+                        logging.error("Token refresh failed during file download.")
+                        return False
+                break
+
+              await asyncio.sleep(5)
+        return False
+
+      except Exception as e:
+        logging.error(f"Error fetching exported data (attempt {retry_count + 1}): {e}")
+        if retry_count < max_retries - 1:
+          if await self.__refresh_token():
+            retry_count += 1
+            continue
+        retry_count += 1
+
+    logging.error(f"Failed to fetch exported data after {max_retries} attempts.")
+    return False

ws_bom_robot_app/llm/vector_store/loader/base.py

@@ -117,11 +117,14 @@ class Loader():
     #region log
     import warnings
     warnings.filterwarnings("ignore", message=".*pin_memory.*no accelerator is found.*")
+    warnings.filterwarnings("ignore", category=UserWarning)
     log_msg_to_ignore = [
         "Going to convert document batch...",
         "Initializing pipeline for",
         "Accelerator device:",
         "detected formats:",
+        "The text detection result is empty",
+        "RapidOCR returned empty result!",       
     ]
     class MessageFilter(logging.Filter):
         def __init__(self, patterns):
@@ -139,8 +142,11 @@ class Loader():
         'docling.document_converter',
         'docling.datamodel',
         'docling.datamodel.document',
+        'docling.models',
+        'docling.models.rapidocr_model',
         'docling.utils.accelerator_utils',
-        'unstructured'
+        'unstructured',
+        'RapidOCR'
     ]
     for logger_name in loggers_to_filter:
         logging.getLogger(logger_name).addFilter(message_filter)

ws_bom_robot_app/llm/vector_store/loader/docling.py

@@ -8,31 +8,39 @@ from docling.datamodel.accelerator_options import AcceleratorDevice, Accelerator
 from docling.datamodel.pipeline_options import PdfPipelineOptions, TableStructureOptions, TableFormerMode
 from langchain_community.document_loaders import UnstructuredFileLoader
 from ws_bom_robot_app.llm.vector_store.db.base import VectorDBStrategy
-from docling.datamodel.pipeline_options import TableStructureOptions, TableFormerMode, EasyOcrOptions
+from docling.datamodel.pipeline_options import TableStructureOptions, TableFormerMode, RapidOcrOptions
 from docling.document_converter import DocumentConverter, PdfFormatOption, ImageFormatOption
 
 def _doclingConverter() -> DocumentConverter:
-  pipeline_options = PdfPipelineOptions(
-      accelerator_options=AcceleratorOptions(
+  _pipeline_config = {
+      "accelerator_options": AcceleratorOptions(
           device=AcceleratorDevice.AUTO,
           cuda_use_flash_attention2=False,
       ),
-      table_structure_options=TableStructureOptions(mode=TableFormerMode.FAST),
-      #do_ocr=False,
-      ocr_options=EasyOcrOptions()
-  )
+      "table_structure_options": TableStructureOptions(mode=TableFormerMode.ACCURATE),
+  }
+  _base_pipeline_options = PdfPipelineOptions(
+      **_pipeline_config,
+      do_ocr=False)  
+  _ocr_pipeline_options = PdfPipelineOptions(
+      **_pipeline_config,
+      ocr_options=RapidOcrOptions(
+         print_verbose=False,         
+         text_score=0.5,
+         #rapidocr_params={"det_use_cuda": True}
+         ))  
   doc_converter = DocumentConverter(
-      format_options={
+      format_options={          
           InputFormat.PDF: PdfFormatOption(
-              pipeline_options=pipeline_options,
+              pipeline_options=_base_pipeline_options,
           ),
           InputFormat.IMAGE: ImageFormatOption(
-              pipeline_options=pipeline_options,
+              pipeline_options=_ocr_pipeline_options,
           ),
       }
   )
-  doc_converter.initialize_pipeline(InputFormat.PDF)
-  doc_converter.initialize_pipeline(InputFormat.IMAGE)
+  for frm in [InputFormat.PDF, InputFormat.IMAGE]:
+     doc_converter.initialize_pipeline(frm)
   return doc_converter
 
 class DoclingLoader(BaseLoader):

ws_bom_robot_app/llm/vector_store/loader/json_loader.py

@@ -1,25 +1,25 @@
-import json
-from typing import Optional
-from langchain_core.documents import Document
-from langchain_community.document_loaders.base import BaseLoader
-
-class JsonLoader(BaseLoader):
-  def __init__(self, file_path: str, meta_fields:Optional[list[str]] = [],encoding: Optional[str] = "utf-8"):
-    self.file_path = file_path
-    self.meta_fields = meta_fields
-    self.encoding = encoding
-
-  def load(self) -> list[Document]:
-    with open(self.file_path, "r", encoding=self.encoding) as file:
-      data = json.load(file)
-    _list = data if isinstance(data, list) else [data]
-    return [
-      Document(
-        page_content=json.dumps(item),
-        metadata={
-          "source": self.file_path,
-          **{field: item.get(field) for field in self.meta_fields if item.get(field)}
-        }
-      )
-      for item in _list
-    ]
+import json
+from typing import Optional
+from langchain_core.documents import Document
+from langchain_community.document_loaders.base import BaseLoader
+
+class JsonLoader(BaseLoader):
+  def __init__(self, file_path: str, meta_fields:Optional[list[str]] = [],encoding: Optional[str] = "utf-8"):
+    self.file_path = file_path
+    self.meta_fields = meta_fields
+    self.encoding = encoding
+
+  def load(self) -> list[Document]:
+    with open(self.file_path, "r", encoding=self.encoding) as file:
+      data = json.load(file)
+    _list = data if isinstance(data, list) else [data]
+    return [
+      Document(
+        page_content=json.dumps(item),
+        metadata={
+          "source": self.file_path,
+          **{field: item.get(field) for field in self.meta_fields if item.get(field)}
+        }
+      )
+      for item in _list
+    ]