ws-bom-robot-app 0.0.18__tar.gz → 0.0.19__tar.gz

{ws_bom_robot_app-0.0.18 → ws_bom_robot_app-0.0.19}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ws_bom_robot_app
-Version: 0.0.18
+Version: 0.0.19
 Summary: A FastAPI application serving ws bom/robot/llm platform ai.
 Home-page: https://github.com/websolutespa/bom
 Author: Websolute Spa

{ws_bom_robot_app-0.0.18 → ws_bom_robot_app-0.0.19}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="ws_bom_robot_app",
-    version="0.0.18",
+    version="0.0.19",
     description="A FastAPI application serving ws bom/robot/llm platform ai.",
     long_description=open("README.md", encoding='utf-8').read(),
     long_description_content_type="text/markdown",

{ws_bom_robot_app-0.0.18 → ws_bom_robot_app-0.0.19}/ws_bom_robot_app/llm/utils/faiss_helper.py

@@ -3,19 +3,32 @@ from langchain_core.documents import Document
 from langchain_core.vectorstores.base import VectorStoreRetriever
 from langchain_openai import OpenAIEmbeddings
 from typing import Any
-import asyncio
+import asyncio, gc, logging
 from pydantic import SecretStr
-from ws_bom_robot_app.util import timer,atimer
 
 class FaissHelper():
     _embedding_model = "text-embedding-3-small"
     _CACHE: dict[str, FAISS] = {}
     @staticmethod
     #@timer
-    async def create(documents: list[Document], folder_path: str, api_key: SecretStr) -> None:
-      embeddings = OpenAIEmbeddings(api_key=api_key, model=FaissHelper._embedding_model)
-      faiss_instance = await asyncio.to_thread(FAISS.from_documents, documents, embeddings)
-      await asyncio.to_thread(faiss_instance.save_local, folder_path)
+    async def create(documents: list[Document], folder_path: str, api_key: SecretStr, return_folder_path:bool = False) -> str | None:
+      try:
+        embeddings = OpenAIEmbeddings(api_key=api_key, model=FaissHelper._embedding_model)
+        faiss_instance = await asyncio.to_thread(FAISS.from_documents, documents, embeddings)
+        await asyncio.to_thread(faiss_instance.save_local, folder_path)
+        del faiss_instance, embeddings
+        gc.collect()
+        if return_folder_path:
+          return folder_path
+        return None
+      except Exception as e:
+        logging.error(f"Failed to create Faiss instance: {e}")
+        return None
+      finally:
+        if 'documents' in locals():
+          del documents
+          gc.collect()
+
     @staticmethod
     #@timer
     def get_loader(folder_path:str,api_key:SecretStr) -> FAISS:

{ws_bom_robot_app-0.0.18 → ws_bom_robot_app-0.0.19}/ws_bom_robot_app/llm/vector_store/generator.py

@@ -1,5 +1,5 @@
-import os, shutil, traceback
-import asyncio
+import os, gc, shutil, traceback
+import asyncio, aiofiles, aiofiles.os
 from fastapi import HTTPException
 from fastapi.responses import StreamingResponse
 from langchain_core.documents import Document
@@ -7,7 +7,13 @@ from ws_bom_robot_app.llm.vector_store.loader.base import Loader
 from ws_bom_robot_app.llm.models.api import RulesRequest, KbRequest, VectorDbResponse
 from ws_bom_robot_app.llm.vector_store.integration.manager import IntegrationManager
 from ws_bom_robot_app.llm.utils.faiss_helper import FaissHelper
-from ws_bom_robot_app.util import timer
+from ws_bom_robot_app.config import config
+from ws_bom_robot_app.llm.models.kb import load_endpoints
+from ws_bom_robot_app.llm.utils.download import download_files
+
+async def _cleanup_directory(directory_path: str):
+  if os.path.exists(directory_path):
+    await asyncio.to_thread(shutil.rmtree, directory_path)
 
 #@timer
 async def rules(rq: RulesRequest) -> VectorDbResponse:
@@ -17,21 +23,17 @@ async def rules(rq: RulesRequest) -> VectorDbResponse:
   store_path = os.path.join(_config.robot_data_folder, _config.robot_data_db_folder, _config.robot_data_db_folder_store, db_name)
   try:
     await FaissHelper.create([Document(page_content=rule, metadata={"source": "rules"}) for rule in rq.rules], store_path, api_key) #type: ignore
-    #FAISS.from_texts([json.dumps({"rule": rule}) for rule in rq.rules], OpenAIEmbeddings(api_key=api_key)).save_local(store_path) #type: ignore
     db_file_path = shutil.make_archive(os.path.join(_config.robot_data_folder, _config.robot_data_db_folder, _config.robot_data_db_folder_out, db_name), "zip", store_path)
     return VectorDbResponse(file = os.path.basename(db_file_path))
   except Exception as e:
-    if os.path.exists(store_path):
-      shutil.rmtree(store_path)
+    await _cleanup_directory(store_path)
     return VectorDbResponse(success = False, error = str(e))
+  finally:
+    gc.collect()
 
 #@atimer
 async def kb(rq: KbRequest) -> VectorDbResponse:
-  import aiofiles
-  import aiofiles.os
   os.environ['MPLCONFIGDIR'] = './tmp/.matplotlib'
-  from ws_bom_robot_app.llm.models.kb import load_endpoints
-  from ws_bom_robot_app.llm.utils.download import download_files
   api_key = rq.api_key()
   _config = rq.config()
   db_name = rq.out_name()
@@ -44,12 +46,12 @@ async def kb(rq: KbRequest) -> VectorDbResponse:
     await aiofiles.os.makedirs(src_path, exist_ok=True)
     await aiofiles.os.makedirs(working_path, exist_ok=True)
 
-  loaders = Loader(working_path)
-  filter_file_extensions = loaders.managed_file_extensions()
-
+  documents: list[Document] = []
   # Download/copy all files
   if rq.files:
     try:
+      loaders = Loader(working_path)
+      filter_file_extensions = loaders.managed_file_extensions()
       files_to_download = [file for file in rq.files if not os.path.exists(os.path.join(src_path, os.path.basename(file)))]
       if files_to_download:
         await download_files(
@@ -60,15 +62,21 @@ async def kb(rq: KbRequest) -> VectorDbResponse:
         async with aiofiles.open(os.path.join(src_path, os.path.basename(file)), 'rb') as src_file:
           async with aiofiles.open(os.path.join(working_path, os.path.basename(file)), 'wb') as dest_file:
             await dest_file.write(await src_file.read())
+      #load files
+      try:
+        documents.extend(await loaders.load())
+      except Exception as e:
+        tb = traceback.format_exc()
+        return VectorDbResponse(success = False, error = f"File loader failure: {e} | {tb}")
     except Exception as e:
+      await _cleanup_directory(working_path)
       return VectorDbResponse(success = False, error = f"Failed to download file {e}")
 
-  documents: list[Document] = []
-
   if rq.endpoints:
     try:
       documents.extend(await load_endpoints(rq.endpoints, working_path))
     except Exception as e:
+      await _cleanup_directory(working_path)
       tb = traceback.format_exc()
       return VectorDbResponse(success = False, error = f"Endpoint failure: {e} | {tb}")
 
@@ -85,33 +93,32 @@ async def kb(rq: KbRequest) -> VectorDbResponse:
       for docs in integration_documents:
         documents.extend(docs)
     except Exception as e:
+      await _cleanup_directory(working_path)
       tb = traceback.format_exc()
       return VectorDbResponse(success=False, error=f"Integration failure: {e} | {tb}")
 
-  try:
-    documents.extend(await loaders.load())
-  except Exception as e:
-    tb = traceback.format_exc()
-    return VectorDbResponse(success = False, error = f"Failed to load files {e} | {tb}")
-  finally:
-    await aiofiles.os.wrap(shutil.rmtree)(working_path)
+  #cleanup
+  await _cleanup_directory(working_path)
 
   if documents and len(documents) > 0:
     try:
       store_path = os.path.join(_config.robot_data_folder, _config.robot_data_db_folder, _config.robot_data_db_folder_store, db_name)
-      # create db
-      await FaissHelper.create(documents,store_path,api_key) #type: ignore
-      # zip db
-      db_file_path = await aiofiles.os.wrap(shutil.make_archive)(os.path.join(_config.robot_data_folder, _config.robot_data_db_folder, _config.robot_data_db_folder_out, db_name), "zip", store_path)
+      db_file_path = await aiofiles.os.wrap(shutil.make_archive)(
+          os.path.join(_config.robot_data_folder, _config.robot_data_db_folder, _config.robot_data_db_folder_out, db_name),
+          "zip",
+          await FaissHelper.create(documents, store_path, api_key, return_folder_path=True)
+      )
       return VectorDbResponse(file = os.path.basename(db_file_path))
     except Exception as e:
+      await _cleanup_directory(store_path)
       return VectorDbResponse(success = False, error = str(e))
+    finally:
+      del documents
+      gc.collect()
   else:
     return VectorDbResponse(success = False, error = "No documents found in the knowledgebase folder")
 
 async def kb_stream_file(filename: str):
-    from ws_bom_robot_app.config import config
-
     file_path = os.path.join(config.robot_data_folder, config.robot_data_db_folder, config.robot_data_db_folder_out, filename)
     if not os.path.isfile(file_path):
         raise HTTPException(status_code=404, detail="File not found")

{ws_bom_robot_app-0.0.18 → ws_bom_robot_app-0.0.19}/ws_bom_robot_app/llm/vector_store/integration/confluence.py

@@ -1,5 +1,4 @@
 import asyncio
-import os
 from ws_bom_robot_app.llm.vector_store.integration.base import IntegrationStrategy, UnstructuredIngest
 from unstructured_ingest.connector.confluence import SimpleConfluenceConfig, ConfluenceAccessConfig
 from unstructured_ingest.runner import ConfluenceRunner
@@ -42,6 +41,6 @@ class Confluence(IntegrationStrategy):
       )
     runner.run()
   async def load(self) -> list[Document]:
-      self.run()
-      await asyncio.sleep(1)
+      await asyncio.to_thread(self.run)
       return await self.__loader.load()
+

{ws_bom_robot_app-0.0.18 → ws_bom_robot_app-0.0.19}/ws_bom_robot_app/llm/vector_store/integration/github.py

@@ -1,5 +1,4 @@
 import asyncio
-import os
 from typing import Optional, Union
 from ws_bom_robot_app.llm.vector_store.integration.base import IntegrationStrategy, UnstructuredIngest
 from unstructured_ingest.connector.git import GitAccessConfig
@@ -43,6 +42,5 @@ class Github(IntegrationStrategy):
       )
     runner.run()
   async def load(self) -> list[Document]:
-      self.run()
-      await asyncio.sleep(1)
+      await asyncio.to_thread(self.run)
       return await self.__loader.load()

{ws_bom_robot_app-0.0.18 → ws_bom_robot_app-0.0.19}/ws_bom_robot_app/llm/vector_store/integration/jira.py

@@ -1,5 +1,4 @@
 import asyncio
-import os
 from ws_bom_robot_app.llm.vector_store.integration.base import IntegrationStrategy, UnstructuredIngest
 from unstructured_ingest.connector.jira import SimpleJiraConfig, JiraAccessConfig
 from unstructured_ingest.runner import JiraRunner
@@ -44,6 +43,5 @@ class Jira(IntegrationStrategy):
       )
     runner.run()
   async def load(self) -> list[Document]:
-      self.run()
-      await asyncio.sleep(1)
+      await asyncio.to_thread(self.run)
       return await self.__loader.load()

{ws_bom_robot_app-0.0.18 → ws_bom_robot_app-0.0.19}/ws_bom_robot_app/llm/vector_store/loader/base.py

@@ -9,7 +9,8 @@ from langchain_core.documents import Document
 from langchain_unstructured import UnstructuredLoader
 from pydantic import BaseModel
 from ws_bom_robot_app.llm.vector_store.loader.json_loader import JsonLoader
-from ws_bom_robot_app.util import timer
+import gc, logging
+
 class LoaderConfig(BaseModel):
   loader: type[BaseLoader]
   kwargs: Optional[dict[str, Any]] = {}
@@ -108,6 +109,18 @@ class Loader():
 
   #@timer
   async def load(self) -> list[Document]:
+    MAX_RETRIES = 3
     loaders = MergedDataLoader(self.__directory_loader())
-    return await loaders.aload()
-    #return await asyncio.to_thread(loaders.load)
+    try:
+      for attempt in range(MAX_RETRIES):
+        try:
+          return await loaders.aload()
+          #return await [doc async for doc in loaders.alazy_load()]
+        except Exception as e:
+          logging.warning(f"Attempt {attempt+1} load document  failed: {e}")
+          await asyncio.sleep(1)
+          if attempt == MAX_RETRIES - 1:
+            logging.error(f"Failed to load documents: {e}")
+            return []
+    finally:
+      gc.collect()

{ws_bom_robot_app-0.0.18 → ws_bom_robot_app-0.0.19}/ws_bom_robot_app/task_manager.py

@@ -4,7 +4,7 @@ from math import floor
 import asyncio, os
 from datetime import datetime, timedelta
 from enum import Enum
-from typing import Annotated, Coroutine, Literal, TypeVar, Optional, Dict, Union, Any
+from typing import Annotated, Coroutine, Literal, TypeVar, Optional, Dict, Union, Any, Callable
 from pydantic import BaseModel, ConfigDict, Field, computed_field
 from uuid import uuid4
 from fastapi import APIRouter, HTTPException
@@ -171,7 +171,7 @@ class TaskManagerStrategy(ABC):
     def task_cleanup_rule(self, task: TaskEntry) -> bool:
         return task.status.status in {"completed", "failure"} and datetime.fromisoformat(task.status.metadata.end_at) < datetime.now() - timedelta(days=config.robot_task_retention_days)
 
-    def task_done_callback(self, task_entry: TaskEntry, headers: TaskHeader | None = None) -> callable:
+    def task_done_callback(self, task_entry: TaskEntry, headers: TaskHeader | None = None) -> Callable:
         def callback(task: asyncio.Task):
             try:
                 result = task_entry.task.result()
@@ -188,9 +188,9 @@ class TaskManagerStrategy(ABC):
                 del self.running_tasks[task_entry.id]
                 #notify webhooks
                 if headers and headers.x_ws_bom_webhooks:
-                      asyncio.create_task(
-                          WebhookNotifier().notify_webhook(task_entry.status,headers.x_ws_bom_webhooks)
-                          )
+                  asyncio.create_task(
+                      WebhookNotifier().notify_webhook(task_entry.status,headers.x_ws_bom_webhooks)
+                      )
         return callback
 
     def create_task_entry(self, coroutine: asyncio.coroutines, headers: TaskHeader | None = None) -> TaskEntry:

{ws_bom_robot_app-0.0.18 → ws_bom_robot_app-0.0.19}/ws_bom_robot_app/util.py

@@ -76,6 +76,32 @@ def timer(func):
 def atimer(func):
     return _get_timer_wrapper(is_async=True)(func)
 
+#profiler
+def memory_leak_detector(func):
+    import tracemalloc, gc, sys
+    async def wrapper(*args, **kwargs):
+        # start tracking
+        tracemalloc.start()
+        initial_snapshot = tracemalloc.take_snapshot()
+        # run
+        result = await func(*args, **kwargs)
+        # take final snapshot
+        final_snapshot = tracemalloc.take_snapshot()
+        # compare snapshots
+        top_stats = final_snapshot.compare_to(initial_snapshot, 'lineno')
+        print(f"\nMemory Leak Analysis for {func.__name__}:")
+        for stat in top_stats[:10]:
+            print(stat)
+        # uncollectable objects
+        print("\n[ Uncollectable Objects ]")
+        print(gc.garbage)
+        print("\nGarbage Collector Stats:")
+        print(gc.get_stats())
+        # stop tracking
+        tracemalloc.stop()
+        return result
+    return wrapper
+
 
 
 

{ws_bom_robot_app-0.0.18 → ws_bom_robot_app-0.0.19}/ws_bom_robot_app.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ws_bom_robot_app
-Version: 0.0.18
+Version: 0.0.19
 Summary: A FastAPI application serving ws bom/robot/llm platform ai.
 Home-page: https://github.com/websolutespa/bom
 Author: Websolute Spa