ws-bom-robot-app 0.0.37__py3-none-any.whl → 0.0.103__py3-none-any.whl

ws_bom_robot_app/llm/vector_store/db/base.py

@@ -0,0 +1,193 @@
+from abc import ABC, abstractmethod
+from typing import Any, Optional, List, Dict
+import asyncio
+from langchain_core.documents import Document
+from langchain_core.embeddings import Embeddings
+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:
+      """
+      A strategy interface for managing vector databases. It caches and retrieves vector
+      stores, providing mechanisms for creating them, retrieving them, and invoking
+      document searches.
+      Attributes:
+        _CACHE (dict[str, VectorStore]):
+          A dictionary that caches loaded VectoreStore(e.g. Faiss,Chroma,qDrant) indexes keyed by their storage IDs.
+      Methods:
+        create(embeddings, documents, storage_id, **kwargs):
+          Asynchronously create a vector store using the provided embeddings,
+          documents, and a unique storage identifier. Returns the created
+          store's ID or None if creation fails.
+        get_loader(embeddings, storage_id, **kwargs):
+          Retrieve a vector store loader based on the provided embeddings
+          and storage identifier. This loader can be used to perform
+          further operations like retrieving documents.
+        get_retriever(embeddings, storage_id, search_type, search_kwargs, **kwargs):
+          Retrieve a VectorStoreRetriever for searching documents. Supports
+          different search methods (e.g., similarity, mmr) and employs the
+          appropriate strategy based on the search_type argument.
+        supports_self_query():
+          Indicates whether this strategy supports self-querying functionality.
+          By default, returns True.
+        _get_self_query_retriever(llm, store, description, metadata):
+          Creates a SelfQueryRetriever using the specified language model,
+          vector store, document description, and metadata. Used internally
+          for self-querying when supported.
+        invoke(embeddings, storage_id, query, search_type, search_kwargs, **kwargs):
+          Asynchronously searches for documents based on a query. Depending
+          on arguments and available metadata, either uses a self-query
+          retriever or falls back to other retrieval methods (e.g., mixed
+          similarity and mmr).
+        _remove_duplicates(docs):
+          Removes duplicate documents by checking their page content,
+          returning a list with unique results.
+        _combine_search(retrievers, query):
+          Asynchronously invokes multiple retrievers in parallel, then merges
+          their results while removing duplicates.
+      """
+    MAX_TOKENS_PER_BATCH = 300_000 * 0.8
+    def __init__(self):
+        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 > VectorDBStrategy.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:
+            del self._CACHE[key]
+
+    @abstractmethod
+    async def create(
+        self,
+        embeddings: Embeddings,
+        documents: List[Document],
+        storage_id: str,
+        **kwargs
+    ) -> Optional[str]:
+        pass
+
+    @abstractmethod
+    def get_loader(
+        self,
+        embeddings: Embeddings,
+        storage_id: str,
+        **kwargs
+    ) -> VectorStore:
+        pass
+
+    def get_retriever(
+        self,
+        embeddings: Embeddings,
+        storage_id: str,
+        search_type: str,
+        search_kwargs: Dict[str, Any],
+        **kwargs
+    ) -> VectorStoreRetriever:
+        return self.get_loader(embeddings, storage_id).as_retriever(
+            search_type=search_type,
+            search_kwargs=search_kwargs
+        )
+
+    def supports_self_query(self) -> bool:
+        return True
+
+    @staticmethod
+    def _get_self_query_retriever(llm:BaseChatModel,store:VectorStore,description:str, metadata: list[AttributeInfo]) -> SelfQueryRetriever:
+        return SelfQueryRetriever.from_llm(
+            llm=llm,
+            vectorstore=store,
+            document_contents=description,
+            metadata_field_info=metadata,
+            enable_limit=True,
+            verbose=True
+        )
+
+    async def invoke(
+        self,
+        embeddings: Embeddings,
+        storage_id: str,
+        query: str,
+        search_type: str,
+        search_kwargs: Dict[str, Any],
+        **kwargs
+    ) -> List[Document]:
+        if self.supports_self_query():
+            if "app_tool" in kwargs and "llm" in kwargs:
+                from ws_bom_robot_app.llm.tools.tool_manager import LlmAppTool
+                app_tool: LlmAppTool = kwargs["app_tool"]
+                _description,_metadata=app_tool.get_vector_filtering()
+                if _description and _metadata:
+                    llm: BaseChatModel = kwargs["llm"]
+                    retriever = VectorDBStrategy._get_self_query_retriever(llm,self.get_loader(embeddings, storage_id),_description,_metadata)
+                    return await retriever.ainvoke(query, config={"source": kwargs.get("source", "retriever")})
+        if search_type == "mixed":
+            similarity_retriever = self.get_retriever(embeddings, storage_id, "similarity", search_kwargs)
+            mmr_kwargs = {
+                "k": search_kwargs.get("k", 4),
+                "fetch_k": search_kwargs.get("fetch_k", 20),
+                "lambda_mult": search_kwargs.get("lambda_mult", 0.2),
+            }
+            mmr_retriever = self.get_retriever(embeddings, storage_id, "mmr", mmr_kwargs)
+            return await VectorDBStrategy._combine_search([similarity_retriever, mmr_retriever], query)
+        retriever = self.get_retriever(embeddings, storage_id, search_type, search_kwargs)
+        return await retriever.ainvoke(query, config={"source": kwargs.get("source", "retriever")})
+
+    @staticmethod
+    def _remove_empty_documents(docs: List[Document]) -> List[Document]:
+        return [doc for doc in docs if doc.page_content and doc.page_content.strip()]
+    @staticmethod
+    def _remove_duplicates(docs: List[Document]) -> List[Document]:
+        seen = set()
+        return [doc for doc in docs if not (doc.page_content in seen or seen.add(doc.page_content))]
+    @staticmethod
+    async def _combine_search(
+        retrievers: List[VectorStoreRetriever],
+        query: str
+    ) -> List[Document]:
+        tasks = [retriever.ainvoke(query, config={"source": "custom source"}) for retriever in retrievers]
+        return VectorDBStrategy._remove_duplicates([doc for res in await asyncio.gather(*tasks) for doc in res])

ws_bom_robot_app/llm/vector_store/db/chroma.py

@@ -0,0 +1,97 @@
+from langchain_chroma import Chroma as CHROMA
+from langchain_core.documents import Document
+from typing import Any, Optional
+import asyncio, gc, logging
+from langchain_core.embeddings import Embeddings
+from ws_bom_robot_app.llm.utils.chunker import DocumentChunker
+from ws_bom_robot_app.llm.vector_store.db.base import VectorDBStrategy
+
+
+class Chroma(VectorDBStrategy):
+    """
+    A strategy class for interacting with a Chroma-based vector store implementation.
+    This class provides methods to create a Chroma vector store from a list of documents
+    and retrieve an existing Chroma instance. The vector store can be used to perform
+    operations such as embedding documents, persisting them to a storage directory, and
+    later loading them for retrieval tasks.
+    Attributes:
+      _CACHE (dict[str, CHROMA]): A cache to store and reuse Chroma instances.
+    Methods:
+      create(embeddings, documents, storage_id, **kwargs):
+        Creates a new Chroma instance after chunking the provided documents
+        and embedding them. Persists the vector store in the given storage directory.
+        If any error occurs during creation, logs the error and returns None.
+        Args:
+          embeddings (Embeddings): The embeddings strategy used to embed documents.
+          documents (list[Document]): The list of documents to be chunked and embedded.
+          storage_id (str): The directory where the Chroma vector store should be persisted.
+          **kwargs: Additional keyword arguments.
+        Returns:
+          Optional[str]: The storage ID if creation is successful; otherwise, None.
+      get_loader(embeddings, storage_id, **kwargs):
+        Retrieves a Chroma instance from the cache if it exists;
+        otherwise, creates and caches a new instance using the given embeddings and storage ID.
+        Args:
+          embeddings (Embeddings): The embeddings strategy used to create or load the Chroma instance.
+          storage_id (str): The directory where the Chroma vector store is persisted.
+          **kwargs: Additional keyword arguments.
+        Returns:
+          CHROMA: The retrieved or newly created Chroma instance.
+    """
+    def __init__(self):
+        super().__init__()
+
+    async def create(
+        self,
+        embeddings: Embeddings,
+        documents: list[Document],
+        storage_id: str,
+        **kwargs
+    ) -> Optional[str]:
+        try:
+            documents = self._remove_empty_documents(documents)
+            chunked_docs = DocumentChunker.chunk(documents)
+            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, chunked_docs, _instance
+            gc.collect()
+
+    def get_loader(
+        self,
+        embeddings: Embeddings,
+        storage_id: str,
+        **kwargs
+    ) -> CHROMA:
+        if storage_id not in self._CACHE:
+            self._CACHE[storage_id] = CHROMA(
+                collection_name="default",
+                embedding_function=embeddings,
+                persist_directory=storage_id
+            )
+        return self._CACHE[storage_id]

ws_bom_robot_app/llm/vector_store/db/faiss.py

@@ -0,0 +1,91 @@
+from langchain_community.vectorstores.faiss import FAISS
+from langchain_core.documents import Document
+from typing import Any, Optional
+import asyncio, gc, logging
+from langchain_core.embeddings import Embeddings
+from ws_bom_robot_app.llm.utils.chunker import DocumentChunker
+from ws_bom_robot_app.llm.vector_store.db.base import VectorDBStrategy
+
+class Faiss(VectorDBStrategy):
+    """
+    Faiss is a vector database strategy that leverages a FAISS index to store and retrieve
+    vectorized documents. It provides methods for creating a new FAISS index and for
+    loading an existing index from a local directory, with an internal caching mechanism
+    to optimize repeated retrievals.
+    Methods:
+      create(
+        Asynchronously creates a FAISS index from the given documents, using the
+        provided embeddings, then saves it locally under the specified storage ID.
+        Returns the storage ID if successful, or None otherwise.
+      get_loader(
+        Retrieves a FAISS index associated with a given storage ID. If this index
+        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,
+        documents: list[Document],
+        storage_id: str,
+        **kwargs
+    ) -> Optional[str]:
+        try:
+            documents = self._remove_empty_documents(documents)
+            chunked_docs = DocumentChunker.chunk(documents)
+            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, chunked_docs, _instance
+            gc.collect()
+
+    def get_loader(
+        self,
+        embeddings: Embeddings,
+        storage_id: str,
+        **kwargs
+    ) -> FAISS:
+        if storage_id not in self._CACHE:
+            self._CACHE[storage_id] = FAISS.load_local(
+                folder_path=storage_id,
+                embeddings=embeddings,
+                allow_dangerous_deserialization=True
+            )
+        return self._CACHE[storage_id]
+
+    def supports_self_query(self) -> bool:
+        return False

ws_bom_robot_app/llm/vector_store/db/manager.py

@@ -0,0 +1,15 @@
+from ws_bom_robot_app.llm.vector_store.db.base import VectorDBStrategy
+from ws_bom_robot_app.llm.vector_store.db.chroma import Chroma
+from ws_bom_robot_app.llm.vector_store.db.faiss import Faiss
+from ws_bom_robot_app.llm.vector_store.db.qdrant import Qdrant
+
+class VectorDbManager:
+  _list: dict[str, VectorDBStrategy] = {
+    "chroma": Chroma(),
+    "faiss": Faiss(),
+    "qdrant": Qdrant()
+  }
+
+  @classmethod
+  def get_strategy(cls, name: str) -> VectorDBStrategy:
+      return cls._list.get(name.lower(), Faiss())

ws_bom_robot_app/llm/vector_store/db/qdrant.py

@@ -0,0 +1,73 @@
+from langchain_qdrant import QdrantVectorStore as QDRANT, FastEmbedSparse, RetrievalMode
+from qdrant_client import QdrantClient
+from langchain_core.documents import Document
+from typing import Any, Optional
+import asyncio, gc, logging, os
+from langchain_core.embeddings import Embeddings
+from ws_bom_robot_app.llm.utils.chunker import DocumentChunker
+from ws_bom_robot_app.llm.vector_store.db.base import VectorDBStrategy
+
+
+class Qdrant(VectorDBStrategy):
+    async def create(
+        self,
+        embeddings: Embeddings,
+        documents: list[Document],
+        storage_id: str,
+        **kwargs
+    ) -> Optional[str]:
+        try:
+            documents = self._remove_empty_documents(documents)
+            chunked_docs = DocumentChunker.chunk(documents)
+            batches = self._batch_documents_by_tokens(chunked_docs)
+            logging.info(f"documents: {len(documents)}, after chunking: {len(chunked_docs)}, processing batches: {len(batches)}")
+            _instance: QDRANT = None            
+            if not os.path.exists(storage_id):
+                os.makedirs(storage_id)
+
+            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(
+                    QDRANT.from_documents,
+                    documents=batch,
+                    embedding=embeddings,
+                    sparse_embedding=kwargs['sparse_embedding'] if 'sparse_embedding' in kwargs else FastEmbedSparse(),
+                    collection_name="default",
+                    path=storage_id,
+                    retrieval_mode=RetrievalMode.HYBRID
+                )
+                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 {Qdrant.__name__} index with {len(chunked_docs)} total documents")
+            return storage_id                
+        except Exception as e:
+            logging.error(f"{Qdrant.__name__} create error: {e}")
+            raise e
+        finally:
+            del documents, chunked_docs, _instance
+            gc.collect()
+
+    def get_loader(
+        self,
+        embeddings: Embeddings,
+        storage_id: str,
+        **kwargs
+    ) -> QDRANT:
+        if storage_id not in self._CACHE:
+            self._CACHE[storage_id] = QDRANT(
+                client=QdrantClient(path=storage_id),
+                collection_name="default",
+                embedding=embeddings,
+                sparse_embedding=FastEmbedSparse(),
+                retrieval_mode=RetrievalMode.HYBRID,
+            )
+        return self._CACHE[storage_id]