flowllm 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

flowllm/storage/vector_store/base_vector_store.py

@@ -1,9 +1,12 @@
-from abc import ABC
+import asyncio
+from abc import ABC, abstractmethod
+from functools import partial
 from pathlib import Path
-from typing import List, Iterable
+from typing import List, Iterable, Dict, Any, Optional
 
 from pydantic import BaseModel, Field
 
+from flowllm.context.service_context import C
 from flowllm.embedding_model.base_embedding_model import BaseEmbeddingModel
 from flowllm.schema.vector_node import VectorNode
 
@@ -12,36 +15,117 @@ class BaseVectorStore(BaseModel, ABC):
     embedding_model: BaseEmbeddingModel | None = Field(default=None)
     batch_size: int = Field(default=1024)
 
+    @abstractmethod
     def exist_workspace(self, workspace_id: str, **kwargs) -> bool:
+        """Check if a workspace exists in the vector store."""
         raise NotImplementedError
 
-    def delete_workspace(self, workspace_id: str, **kwargs):
+    @abstractmethod
+    def delete_workspace(self, workspace_id: str, **kwargs) -> None:
+        """Delete a workspace from the vector store."""
         raise NotImplementedError
 
-    def create_workspace(self, workspace_id: str, **kwargs):
+    @abstractmethod
+    def create_workspace(self, workspace_id: str, **kwargs) -> None:
+        """Create a new workspace in the vector store."""
         raise NotImplementedError
 
-    def _iter_workspace_nodes(self, workspace_id: str, **kwargs) -> Iterable[VectorNode]:
+    @abstractmethod
+    def iter_workspace_nodes(self, workspace_id: str, callback_fn=None, **kwargs) -> Iterable[VectorNode]:
+        """Iterate over all nodes in a workspace."""
         raise NotImplementedError
 
-    def iter_workspace_nodes(self, workspace_id: str, **kwargs) -> Iterable[VectorNode]:
-        return self._iter_workspace_nodes(workspace_id, **kwargs)
-
-    def dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs):
+    @abstractmethod
+    def dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs) -> None:
+        """Dump workspace data to a file or path."""
         raise NotImplementedError
 
-    def load_workspace(self, workspace_id: str, path: str | Path = "", nodes: List[VectorNode] = None, callback_fn=None,
-                       **kwargs):
+    @abstractmethod
+    def load_workspace(self, workspace_id: str, path: str | Path = "", nodes: Optional[List[VectorNode]] = None,
+                       callback_fn=None, **kwargs) -> None:
+        """Load workspace data from a file or path, or from provided nodes."""
         raise NotImplementedError
 
-    def copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs):
+    @abstractmethod
+    def copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs) -> None:
+        """Copy one workspace to another."""
         raise NotImplementedError
 
-    def search(self, query: str, workspace_id: str, top_k: int = 1, **kwargs) -> List[VectorNode]:
+    @abstractmethod
+    def search(self, query: str, workspace_id: str, top_k: int = 1, filter_dict: Optional[Dict[str, Any]] = None,
+               **kwargs) -> List[VectorNode]:
+        """Search for similar vectors in the workspace."""
         raise NotImplementedError
 
-    def insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs):
+    @abstractmethod
+    def insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs) -> None:
+        """Insert nodes into the workspace."""
         raise NotImplementedError
 
-    def delete(self, node_ids: str | List[str], workspace_id: str, **kwargs):
+    @abstractmethod
+    def delete(self, node_ids: str | List[str], workspace_id: str, **kwargs) -> None:
+        """Delete nodes from the workspace by their IDs."""
         raise NotImplementedError
+
+    def close(self) -> None:
+        """Close the vector store and clean up resources. Default implementation does nothing."""
+        pass
+
+    """
+    Async versions of all methods
+    """
+
+    async def async_exist_workspace(self, workspace_id: str, **kwargs) -> bool:
+        """Async version of exist_workspace."""
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, partial(self.exist_workspace, workspace_id, **kwargs))
+
+    async def async_delete_workspace(self, workspace_id: str, **kwargs) -> None:
+        """Async version of delete_workspace."""
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, partial(self.delete_workspace, workspace_id, **kwargs))
+
+    async def async_create_workspace(self, workspace_id: str, **kwargs) -> None:
+        """Async version of create_workspace."""
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, partial(self.create_workspace, workspace_id, **kwargs))
+
+    async def async_iter_workspace_nodes(self, workspace_id: str, callback_fn=None, **kwargs) -> Iterable[VectorNode]:
+        """Async version of iter_workspace_nodes. Returns an iterable, not an async iterator."""
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, partial(self.iter_workspace_nodes, workspace_id,
+                                                                 callback_fn, **kwargs))
+
+    async def async_dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs):
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, partial(self.dump_workspace, workspace_id, path,
+                                                                 callback_fn, **kwargs))
+
+    async def async_load_workspace(self, workspace_id: str, path: str | Path = "", nodes: List[VectorNode] = None,
+                                   callback_fn=None, **kwargs):
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, partial(self.load_workspace, workspace_id, path, nodes,
+                                                                 callback_fn, **kwargs))
+
+    async def async_copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs):
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, partial(self.copy_workspace, src_workspace_id,
+                                                                 dest_workspace_id, **kwargs))
+
+    async def async_search(self, query: str, workspace_id: str, top_k: int = 1, filter_dict: dict = None,
+                           **kwargs) -> List[VectorNode]:
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, partial(self.search, query, workspace_id, top_k,
+                                                                 filter_dict, **kwargs))
+
+    async def async_insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs):
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, partial(self.insert, nodes, workspace_id, **kwargs))
+
+    async def async_delete(self, node_ids: str | List[str], workspace_id: str, **kwargs):
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, partial(self.delete, node_ids, workspace_id, **kwargs))
+
+    async def async_close(self):
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, self.close)

flowllm/storage/vector_store/chroma_vector_store.py

@@ -1,4 +1,7 @@
-from typing import List, Iterable
+import asyncio
+import os
+from functools import partial
+from typing import List, Iterable, Dict, Any, Optional
 
 import chromadb
 from chromadb import Collection
@@ -6,6 +9,9 @@ from chromadb.config import Settings
 from loguru import logger
 from pydantic import Field, PrivateAttr, model_validator
 
+# Disable ChromaDB telemetry to avoid PostHog warnings
+os.environ.setdefault("ANONYMIZED_TELEMETRY", "false")
+
 from flowllm.context.service_context import C
 from flowllm.schema.vector_node import VectorNode
 from flowllm.storage.vector_store.local_vector_store import LocalVectorStore
@@ -19,7 +25,12 @@ class ChromaVectorStore(LocalVectorStore):
 
     @model_validator(mode="after")
     def init_client(self):
-        self._client = chromadb.Client(Settings(persist_directory=self.store_dir))
+        # Disable telemetry to avoid PostHog warnings
+        settings = Settings(
+            persist_directory=self.store_dir,
+            anonymized_telemetry=False
+        )
+        self._client = chromadb.Client(settings)
         return self
 
     def _get_collection(self, workspace_id: str) -> Collection:
@@ -38,7 +49,8 @@ class ChromaVectorStore(LocalVectorStore):
     def create_workspace(self, workspace_id: str, **kwargs):
         self.collections[workspace_id] = self._client.get_or_create_collection(workspace_id)
 
-    def _iter_workspace_nodes(self, workspace_id: str, **kwargs) -> Iterable[VectorNode]:
+    def iter_workspace_nodes(self, workspace_id: str, callback_fn=None, **kwargs) -> Iterable[VectorNode]:
+        """Iterate over all nodes in a workspace."""
         collection: Collection = self._get_collection(workspace_id)
         results = collection.get()
         for i in range(len(results["ids"])):
@@ -46,23 +58,69 @@ class ChromaVectorStore(LocalVectorStore):
                               unique_id=results["ids"][i],
                               content=results["documents"][i],
                               metadata=results["metadatas"][i])
-            yield node
-
-    def search(self, query: str, workspace_id: str, top_k: int = 1, **kwargs) -> List[VectorNode]:
+            if callback_fn:
+                yield callback_fn(node)
+            else:
+                yield node
+
+    @staticmethod
+    def _build_chroma_filters(filter_dict: Optional[Dict[str, Any]] = None) -> Optional[Dict]:
+        """Build ChromaDB where clause from filter_dict"""
+        if not filter_dict:
+            return None
+
+        where_conditions = {}
+        for key, filter_value in filter_dict.items():
+            if isinstance(filter_value, dict):
+                # Range filter: {"gte": 1, "lte": 10}
+                range_conditions = {}
+                if "gte" in filter_value:
+                    range_conditions["$gte"] = filter_value["gte"]
+                if "lte" in filter_value:
+                    range_conditions["$lte"] = filter_value["lte"]
+                if "gt" in filter_value:
+                    range_conditions["$gt"] = filter_value["gt"]
+                if "lt" in filter_value:
+                    range_conditions["$lt"] = filter_value["lt"]
+                if range_conditions:
+                    where_conditions[key] = range_conditions
+            else:
+                # Term filter: direct value comparison
+                where_conditions[key] = filter_value
+
+        return where_conditions if where_conditions else None
+
+    def search(self, query: str, workspace_id: str, top_k: int = 1, filter_dict: Optional[Dict[str, Any]] = None,
+               **kwargs) -> List[VectorNode]:
         if not self.exist_workspace(workspace_id=workspace_id):
             logger.warning(f"workspace_id={workspace_id} is not exists!")
             return []
 
         collection: Collection = self._get_collection(workspace_id)
         query_vector = self.embedding_model.get_embeddings(query)
-        results = collection.query(query_embeddings=[query_vector], n_results=top_k)
+
+        # Build where clause from filter_dict
+        where_clause = self._build_chroma_filters(filter_dict)
+        
+        results = collection.query(
+            query_embeddings=[query_vector], 
+            n_results=top_k,
+            where=where_clause
+        )
+        
         nodes = []
         for i in range(len(results["ids"][0])):
             node = VectorNode(workspace_id=workspace_id,
                               unique_id=results["ids"][0][i],
                               content=results["documents"][0][i],
                               metadata=results["metadatas"][0][i])
+            # ChromaDB returns distances, convert to similarity score
+            if results.get("distances") and len(results["distances"][0]) > i:
+                distance = results["distances"][0][i]
+                # Convert distance to similarity (assuming cosine distance)
+                node.metadata["score"] = 1.0 - distance
             nodes.append(node)
+        
         return nodes
 
     def insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs):
@@ -95,6 +153,85 @@ class ChromaVectorStore(LocalVectorStore):
         collection.delete(ids=node_ids)
 
 
+    async def async_search(self, query: str, workspace_id: str, top_k: int = 1,
+                           filter_dict: Optional[Dict[str, Any]] = None, **kwargs) -> List[VectorNode]:
+        """Async version of search using async embedding and run_in_executor for ChromaDB operations"""
+        if not await self.async_exist_workspace(workspace_id=workspace_id):
+            logger.warning(f"workspace_id={workspace_id} is not exists!")
+            return []
+
+        # Use async embedding
+        query_vector = await self.embedding_model.get_embeddings_async(query)
+
+        # Build where clause from filter_dict
+        where_clause = self._build_chroma_filters(filter_dict)
+
+        # Execute ChromaDB query in thread pool
+        loop = asyncio.get_event_loop()
+        collection = await loop.run_in_executor(C.thread_pool, self._get_collection, workspace_id)
+        results = await loop.run_in_executor(
+            C.thread_pool,
+            partial(collection.query, query_embeddings=[query_vector], n_results=top_k, where=where_clause)
+        )
+
+        nodes = []
+        for i in range(len(results["ids"][0])):
+            node = VectorNode(workspace_id=workspace_id,
+                              unique_id=results["ids"][0][i],
+                              content=results["documents"][0][i],
+                              metadata=results["metadatas"][0][i])
+            # ChromaDB returns distances, convert to similarity score
+            if results.get("distances") and len(results["distances"][0]) > i:
+                distance = results["distances"][0][i]
+                # Convert distance to similarity (assuming cosine distance)
+                node.metadata["score"] = 1.0 - distance
+            nodes.append(node)
+        
+        return nodes
+
+    async def async_insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs):
+        """Async version of insert using async embedding and run_in_executor for ChromaDB operations"""
+        if not await self.async_exist_workspace(workspace_id=workspace_id):
+            await self.async_create_workspace(workspace_id=workspace_id)
+
+        if isinstance(nodes, VectorNode):
+            nodes = [nodes]
+
+        embedded_nodes = [node for node in nodes if node.vector]
+        not_embedded_nodes = [node for node in nodes if not node.vector]
+
+        # Use async embedding
+        now_embedded_nodes = await self.embedding_model.get_node_embeddings_async(not_embedded_nodes)
+
+        all_nodes = embedded_nodes + now_embedded_nodes
+
+        # Execute ChromaDB operations in thread pool
+        loop = asyncio.get_event_loop()
+        collection = await loop.run_in_executor(C.thread_pool, self._get_collection, workspace_id)
+        await loop.run_in_executor(
+            C.thread_pool,
+            partial(collection.add,
+                    ids=[n.unique_id for n in all_nodes],
+                    embeddings=[n.vector for n in all_nodes],
+                    documents=[n.content for n in all_nodes],
+                    metadatas=[n.metadata for n in all_nodes])
+        )
+
+    async def async_delete(self, node_ids: str | List[str], workspace_id: str, **kwargs):
+        """Async version of delete using run_in_executor for ChromaDB operations"""
+        if not await self.async_exist_workspace(workspace_id=workspace_id):
+            logger.warning(f"workspace_id={workspace_id} is not exists!")
+            return
+
+        if isinstance(node_ids, str):
+            node_ids = [node_ids]
+
+        # Execute ChromaDB operations in thread pool
+        loop = asyncio.get_event_loop()
+        collection = await loop.run_in_executor(C.thread_pool, self._get_collection, workspace_id)
+        await loop.run_in_executor(C.thread_pool, partial(collection.delete, ids=node_ids))
+
+
 def main():
     from flowllm.utils.common_utils import load_env
     from flowllm.embedding_model import OpenAICompatibleEmbeddingModel
@@ -160,6 +297,15 @@ def main():
         logger.info(r.model_dump(exclude={"vector"}))
     logger.info("=" * 20)
 
+    # Test filter_dict
+    logger.info("=" * 20 + " FILTER TEST " + "=" * 20)
+    filter_dict = {"node_type": "n1"}
+    results = chroma_store.search("What is AI?", top_k=5, workspace_id=workspace_id, filter_dict=filter_dict)
+    logger.info(f"Filtered results (node_type=n1): {len(results)} results")
+    for r in results:
+        logger.info(r.model_dump(exclude={"vector"}))
+    logger.info("=" * 20)
+
     node2_update = VectorNode(
         unique_id="node2",
         workspace_id=workspace_id,
@@ -184,6 +330,102 @@ def main():
     chroma_store.delete_workspace(workspace_id=workspace_id)
 
 
+async def async_main():
+    from flowllm.utils.common_utils import load_env
+    from flowllm.embedding_model import OpenAICompatibleEmbeddingModel
+
+    load_env()
+
+    embedding_model = OpenAICompatibleEmbeddingModel(dimensions=64, model_name="text-embedding-v4")
+    workspace_id = "chroma_async_test_index"
+
+    chroma_store = ChromaVectorStore(
+        embedding_model=embedding_model,
+        store_dir="./async_chroma_async_test_db"
+    )
+
+    # Clean up and create workspace
+    if await chroma_store.async_exist_workspace(workspace_id):
+        await chroma_store.async_delete_workspace(workspace_id)
+    await chroma_store.async_create_workspace(workspace_id)
+
+    sample_nodes = [
+        VectorNode(
+            unique_id="async_node1",
+            workspace_id=workspace_id,
+            content="Artificial intelligence is a technology that simulates human intelligence.",
+            metadata={
+                "node_type": "n1",
+                "category": "tech"
+            }
+        ),
+        VectorNode(
+            unique_id="async_node2",
+            workspace_id=workspace_id,
+            content="AI is the future of mankind.",
+            metadata={
+                "node_type": "n1",
+                "category": "tech"
+            }
+        ),
+        VectorNode(
+            unique_id="async_node3",
+            workspace_id=workspace_id,
+            content="I want to eat fish!",
+            metadata={
+                "node_type": "n2",
+                "category": "food"
+            }
+        ),
+        VectorNode(
+            unique_id="async_node4",
+            workspace_id=workspace_id,
+            content="The bigger the storm, the more expensive the fish.",
+            metadata={
+                "node_type": "n1",
+                "category": "food"
+            }
+        ),
+    ]
+
+    # Test async insert
+    await chroma_store.async_insert(sample_nodes, workspace_id=workspace_id)
+
+    logger.info("ASYNC TEST - " + "=" * 20)
+    # Test async search
+    results = await chroma_store.async_search("What is AI?", top_k=5, workspace_id=workspace_id)
+    for r in results:
+        logger.info(r.model_dump(exclude={"vector"}))
+    logger.info("=" * 20)
+
+    # Test async update (delete + insert)
+    node2_update = VectorNode(
+        unique_id="async_node2",
+        workspace_id=workspace_id,
+        content="AI is the future of humanity and technology.",
+        metadata={
+            "node_type": "n1",
+            "category": "tech",
+            "updated": True
+        }
+    )
+    await chroma_store.async_delete(node2_update.unique_id, workspace_id=workspace_id)
+    await chroma_store.async_insert(node2_update, workspace_id=workspace_id)
+
+    logger.info("ASYNC Updated Result:")
+    results = await chroma_store.async_search("fish?", top_k=10, workspace_id=workspace_id)
+    for r in results:
+        logger.info(r.model_dump(exclude={"vector"}))
+    logger.info("=" * 20)
+
+    # Clean up
+    await chroma_store.async_dump_workspace(workspace_id=workspace_id)
+    await chroma_store.async_delete_workspace(workspace_id=workspace_id)
+
+
 if __name__ == "__main__":
     main()
-    # launch with: python -m flowllm.storage.chroma_vector_store
+
+    # Run async test
+    logger.info("\n" + "=" * 50 + " ASYNC TESTS " + "=" * 50)
+    # asyncio.run(async_main())