ws-bom-robot-app 0.0.73__py3-none-any.whl → 0.0.75__py3-none-any.whl

ws_bom_robot_app/llm/vector_store/db/base.py

@@ -7,6 +7,7 @@ from langchain_core.language_models import BaseChatModel
 from langchain_core.vectorstores.base import VectorStoreRetriever, VectorStore
 from langchain.retrievers import SelfQueryRetriever
 from langchain.chains.query_constructor.schema import AttributeInfo
+import tiktoken
 
 class VectorDBStrategy(ABC):
     class VectorDBStrategy:
@@ -49,6 +50,52 @@ class VectorDBStrategy(ABC):
           Asynchronously invokes multiple retrievers in parallel, then merges
           their results while removing duplicates.
       """
+    def __init__(self):
+        self.max_tokens_per_batch = 300_000 * 0.8  # conservative limit below 300k openai limit: https://platform.openai.com/docs/api-reference/embeddings/create
+        try:
+            self.encoding = tiktoken.get_encoding("cl100k_base")  # text-embedding-3-small, text-embedding-3-large: https://cookbook.openai.com/examples/how_to_count_tokens_with_tiktoken
+        except Exception:
+            self.encoding = None
+
+    def _count_tokens(self, text: str) -> int:
+        """Count tokens in text using tiktoken or fallback estimation"""
+        if self.encoding:
+            try:
+                return len(self.encoding.encode(text))
+            except Exception:
+                pass
+        # fallback: rough estimation (1 token ≈ 4 characters)
+        return len(text) // 4
+
+    def _batch_documents_by_tokens(self, documents: list[Document]) -> list[list[Document]]:
+      """Split documents into batches based on token count"""
+      if not documents:
+        return []
+      batches = []
+      current_batch = []
+      current_token_count = 0
+
+      for doc in documents:
+          doc_tokens = self._count_tokens(doc.page_content)
+          # check if adding this document exceeds the limit
+          if current_token_count + doc_tokens > self.max_tokens_per_batch:
+              # start new batch if current batch is not empty
+              if current_batch:
+                  batches.append(current_batch)
+              # reset current batch
+              current_batch = [doc]
+              current_token_count = doc_tokens  # reset to current doc's tokens
+          else:
+              # add to current batch
+              current_batch.append(doc)
+              current_token_count += doc_tokens
+
+      # add final batch if not empty
+      if current_batch:
+          batches.append(current_batch)
+
+      return batches
+
     _CACHE: dict[str, VectorStore] = {}
     def _clear_cache(self, key: str):
         if key in self._CACHE:

ws_bom_robot_app/llm/vector_store/db/chroma.py

@@ -38,6 +38,9 @@ class Chroma(VectorDBStrategy):
         Returns:
           CHROMA: The retrieved or newly created Chroma instance.
     """
+    def __init__(self):
+        super().__init__()
+
     async def create(
         self,
         embeddings: Embeddings,
@@ -47,19 +50,35 @@ class Chroma(VectorDBStrategy):
     ) -> Optional[str]:
         try:
             chunked_docs = DocumentChunker.chunk(documents)
-            await asyncio.to_thread(
-                CHROMA.from_documents,
-                documents=chunked_docs,
-                embedding=embeddings,
-                persist_directory=storage_id
-            )
-            self._clear_cache(storage_id)
+            batches = self._batch_documents_by_tokens(chunked_docs)
+            logging.info(f"documents: {len(documents)}, after chunking: {len(chunked_docs)}, processing batches: {len(batches)}")
+            _instance: CHROMA = None
+            for i, batch in enumerate(batches):
+                batch_tokens = sum(self._count_tokens(doc.page_content) for doc in batch)
+                logging.info(f"processing batch {i+1}/{len(batches)} with {len(batch)} docs ({batch_tokens:,} tokens)")
+                # create instance from first batch
+                if _instance is None:
+                    _instance = await asyncio.to_thread(
+                    CHROMA.from_documents,
+                    documents=batch,
+                    embedding=embeddings,
+                    persist_directory=storage_id
+                )
+                else:
+                    # merge to existing instance
+                    await _instance.aadd_documents(batch)
+                # add a small delay to avoid rate limiting
+                if i < len(batches) - 1:  # except last batch
+                    await asyncio.sleep(1)
+            if _instance:
+                self._clear_cache(storage_id)
+                logging.info(f"Successfully created {Chroma.__name__} index with {len(chunked_docs)} total documents")
             return storage_id
         except Exception as e:
             logging.error(f"{Chroma.__name__} create error: {e}")
             raise e
         finally:
-            del documents
+            del documents, chunked_docs, _instance
             gc.collect()
 
     def get_loader(

ws_bom_robot_app/llm/vector_store/db/faiss.py

@@ -22,6 +22,9 @@ class Faiss(VectorDBStrategy):
         was previously loaded and cached, it returns the cached instance; otherwise,
         it loads the index from local storage and caches it for subsequent use.
     """
+    def __init__(self):
+        super().__init__()
+
     async def create(
         self,
         embeddings: Embeddings,
@@ -31,19 +34,42 @@ class Faiss(VectorDBStrategy):
     ) -> Optional[str]:
         try:
             chunked_docs = DocumentChunker.chunk(documents)
-            _instance = await asyncio.to_thread(
-                FAISS.from_documents,
-                chunked_docs,
-                embeddings
-            )
-            await asyncio.to_thread(_instance.save_local, storage_id)
-            self._clear_cache(storage_id)
+            batches = self._batch_documents_by_tokens(chunked_docs)
+            logging.info(f"documents: {len(documents)}, after chunking: {len(chunked_docs)}, processing batches: {len(batches)}")
+            _instance: FAISS = None
+            for i, batch in enumerate(batches):
+                batch_tokens = sum(self._count_tokens(doc.page_content) for doc in batch)
+                logging.info(f"processing batch {i+1}/{len(batches)} with {len(batch)} docs ({batch_tokens:,} tokens)")
+                # init
+                _batch_instance = await asyncio.to_thread(
+                    FAISS.from_documents,
+                    batch,
+                    embeddings
+                )
+                # create instance from first batch
+                if _instance is None:
+                    _instance = _batch_instance
+                else:
+                    # merge to existing instance
+                    await asyncio.to_thread(
+                        _instance.merge_from,
+                        _batch_instance
+                    )
+                del _batch_instance
+                gc.collect()
+                # add a small delay to avoid rate limiting
+                if i < len(batches) - 1:  # except last batch
+                    await asyncio.sleep(1)
+            if _instance:
+                await asyncio.to_thread(_instance.save_local, storage_id)
+                self._clear_cache(storage_id)
+                logging.info(f"Successfully created {Faiss.__name__} index with {len(chunked_docs)} total documents")
             return storage_id
         except Exception as e:
             logging.error(f"{Faiss.__name__} create error: {e}")
             raise e
         finally:
-            del documents, _instance
+            del documents, chunked_docs, _instance
             gc.collect()
 
     def get_loader(

ws_bom_robot_app/llm/vector_store/generator.py

@@ -1,137 +1,137 @@
-import os, gc, shutil, logging, traceback
-import asyncio, aiofiles, aiofiles.os
-from fastapi import HTTPException
-from fastapi.responses import StreamingResponse
-from langchain_core.documents import Document
-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.vector_store.db.manager import VectorDbManager
-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:
-  _config = rq.config()
-  db_name = rq.out_name()
-  store_path = os.path.join(_config.robot_data_folder, _config.robot_data_db_folder, _config.robot_data_db_folder_store, db_name)
-  try:
-    await VectorDbManager.get_strategy(rq.vector_type).create(rq.embeddings(),[Document(page_content=rule, metadata={"source": "rules"}) for rule in rq.rules], 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), vector_type=rq.vector_type)
-  except Exception as e:
-    try:
-      await _cleanup_directory(store_path)
-    finally:
-      return VectorDbResponse(success = False, error = str(e))
-  finally:
-    gc.collect()
-
-#@atimer
-async def kb(rq: KbRequest) -> VectorDbResponse:
-  os.environ['MPLCONFIGDIR'] = './tmp/.matplotlib'
-  _config = rq.config()
-  db_name = rq.out_name()
-  src_path = os.path.join(_config.robot_data_folder, _config.robot_data_db_folder, _config.robot_data_db_folder_src)
-  working_path = os.path.join(src_path, db_name)
-
-  if all([not rq.files,not rq.endpoints,not rq.integrations]):
-    return VectorDbResponse(success = False, error = "No files, endpoints or integrations provided")
-  else:
-    await aiofiles.os.makedirs(src_path, exist_ok=True)
-    await aiofiles.os.makedirs(working_path, exist_ok=True)
-
-  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(
-          [f"{_config.robot_cms_host}/{_config.robot_cms_kb_folder}/{os.path.basename(file)}" for file in files_to_download if any([file.endswith(ext) for ext in filter_file_extensions])],
-          src_path, authorization=_config.robot_cms_auth)
-      # copy files to working tmp folder
-      for file in rq.files:
-        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()
-        _error = f"File loader failure: {e} | {tb}"
-        logging.warning(_error)
-        return VectorDbResponse(success = False, error = _error)
-    except Exception as e:
-      await _cleanup_directory(working_path)
-      return VectorDbResponse(success = False, error = f"Failed to download file {e}")
-
-  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()
-      _error = f"Endpoint failure: {e} | {tb}"
-      logging.warning(_error)
-      return VectorDbResponse(success = False, error = _error)
-
-  if rq.integrations:
-    tasks = []
-    for integration in rq.integrations:
-      tasks.append(
-        IntegrationManager
-        .get_strategy(integration.type.lower(), working_path, integration.__pydantic_extra__) #type: ignore
-        .load()
-      )
-    try:
-      integration_documents = await asyncio.gather(*tasks)
-      for docs in integration_documents:
-        documents.extend(docs)
-    except Exception as e:
-      await _cleanup_directory(working_path)
-      tb = traceback.format_exc()
-      _error = f"Integration failure: {e} | {tb}"
-      logging.warning(_error)
-      return VectorDbResponse(success=False, error=_error)
-
-  #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)
-      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 VectorDbManager.get_strategy(rq.vector_type).create(rq.embeddings(), documents, store_path, return_folder_path=True)
-      )
-      return VectorDbResponse(file = os.path.basename(db_file_path), vector_type=rq.vector_type)
-    except Exception as e:
-      await _cleanup_directory(store_path)
-      return VectorDbResponse(success = False, error = str(e))
-    finally:
-      del documents
-      gc.collect()
-  else:
-    _error = "No documents found in the knowledgebase folder"
-    logging.warning(_error)
-    return VectorDbResponse(success = False, error = _error)
-
-async def kb_stream_file(filename: str):
-    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")
-    def iter_file():
-        with open(file_path, mode="rb") as file:
-            while chunk := file.read(1024*8):
-                yield chunk
-    return StreamingResponse(iter_file(), media_type="application/octet-stream", headers={"Content-Disposition": f"attachment; filename={filename}"})
+import os, gc, shutil, logging, traceback
+import asyncio, aiofiles, aiofiles.os
+from fastapi import HTTPException
+from fastapi.responses import StreamingResponse
+from langchain_core.documents import Document
+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.vector_store.db.manager import VectorDbManager
+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:
+  _config = rq.config()
+  db_name = rq.out_name()
+  store_path = os.path.join(_config.robot_data_folder, _config.robot_data_db_folder, _config.robot_data_db_folder_store, db_name)
+  try:
+    await VectorDbManager.get_strategy(rq.vector_type).create(rq.embeddings(),[Document(page_content=rule, metadata={"source": "rules"}) for rule in rq.rules], 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), vector_type=rq.vector_type)
+  except Exception as e:
+    try:
+      await _cleanup_directory(store_path)
+    finally:
+      return VectorDbResponse(success = False, error = str(e))
+  finally:
+    gc.collect()
+
+#@atimer
+async def kb(rq: KbRequest) -> VectorDbResponse:
+  os.environ['MPLCONFIGDIR'] = './tmp/.matplotlib'
+  _config = rq.config()
+  db_name = rq.out_name()
+  src_path = os.path.join(_config.robot_data_folder, _config.robot_data_db_folder, _config.robot_data_db_folder_src)
+  working_path = os.path.join(src_path, db_name)
+
+  if all([not rq.files,not rq.endpoints,not rq.integrations]):
+    return VectorDbResponse(success = False, error = "No files, endpoints or integrations provided")
+  else:
+    await aiofiles.os.makedirs(src_path, exist_ok=True)
+    await aiofiles.os.makedirs(working_path, exist_ok=True)
+
+  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(
+          [f"{_config.robot_cms_host}/{_config.robot_cms_kb_folder}/{os.path.basename(file)}" for file in files_to_download if any([file.endswith(ext) for ext in filter_file_extensions])],
+          src_path, authorization=_config.robot_cms_auth)
+      # copy files to working tmp folder
+      for file in rq.files:
+        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()
+        _error = f"File loader failure: {e} | {tb}"
+        logging.warning(_error)
+        return VectorDbResponse(success = False, error = _error)
+    except Exception as e:
+      await _cleanup_directory(working_path)
+      return VectorDbResponse(success = False, error = f"Failed to download file {e}")
+
+  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()
+      _error = f"Endpoint failure: {e} | {tb}"
+      logging.warning(_error)
+      return VectorDbResponse(success = False, error = _error)
+
+  if rq.integrations:
+    tasks = []
+    for integration in rq.integrations:
+      tasks.append(
+        IntegrationManager
+        .get_strategy(integration.type.lower(), working_path, integration.__pydantic_extra__) #type: ignore
+        .load()
+      )
+    try:
+      integration_documents = await asyncio.gather(*tasks)
+      for docs in integration_documents:
+        documents.extend(docs)
+    except Exception as e:
+      await _cleanup_directory(working_path)
+      tb = traceback.format_exc()
+      _error = f"Integration failure: {e} | {tb}"
+      logging.warning(_error)
+      return VectorDbResponse(success=False, error=_error)
+
+  #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)
+      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 VectorDbManager.get_strategy(rq.vector_type).create(rq.embeddings(), documents, store_path, return_folder_path=True)
+      )
+      return VectorDbResponse(file = os.path.basename(db_file_path), vector_type=rq.vector_type)
+    except Exception as e:
+      await _cleanup_directory(store_path)
+      return VectorDbResponse(success = False, error = str(e))
+    finally:
+      del documents
+      gc.collect()
+  else:
+    _error = "No documents found in the knowledgebase folder"
+    logging.warning(_error)
+    return VectorDbResponse(success = False, error = _error)
+
+async def kb_stream_file(filename: str):
+    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")
+    def iter_file():
+        with open(file_path, mode="rb") as file:
+            while chunk := file.read(1024*8):
+                yield chunk
+    return StreamingResponse(iter_file(), media_type="application/octet-stream", headers={"Content-Disposition": f"attachment; filename={filename}"})