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

flowllm/storage/vector_store/es_vector_store.py

@@ -1,8 +1,9 @@
+import asyncio
 import os
-from typing import List, Tuple, Iterable
+from typing import List, Tuple, Iterable, Dict, Any, Optional
 
-from elasticsearch import Elasticsearch
-from elasticsearch.helpers import bulk
+from elasticsearch import Elasticsearch, AsyncElasticsearch
+from elasticsearch.helpers import bulk, async_bulk
 from loguru import logger
 from pydantic import Field, PrivateAttr, model_validator
 
@@ -15,14 +16,15 @@ from flowllm.storage.vector_store.local_vector_store import LocalVectorStore
 class EsVectorStore(LocalVectorStore):
     hosts: str | List[str] = Field(default_factory=lambda: os.getenv("FLOW_ES_HOSTS", "http://localhost:9200"))
     basic_auth: str | Tuple[str, str] | None = Field(default=None)
-    retrieve_filters: List[dict] = []
     _client: Elasticsearch = PrivateAttr()
+    _async_client: AsyncElasticsearch = PrivateAttr()
 
     @model_validator(mode="after")
     def init_client(self):
         if isinstance(self.hosts, str):
             self.hosts = [self.hosts]
         self._client = Elasticsearch(hosts=self.hosts, basic_auth=self.basic_auth)
+        self._async_client = AsyncElasticsearch(hosts=self.hosts, basic_auth=self.basic_auth)
         logger.info(f"Elasticsearch client initialized with hosts: {self.hosts}")
         return self
 
@@ -48,10 +50,16 @@ class EsVectorStore(LocalVectorStore):
         }
         return self._client.indices.create(index=workspace_id, body=body)
 
-    def _iter_workspace_nodes(self, workspace_id: str, max_size: int = 10000, **kwargs) -> Iterable[VectorNode]:
+    def iter_workspace_nodes(self, workspace_id: str, callback_fn=None, max_size: int = 10000, **kwargs) -> Iterable[
+        VectorNode]:
+        """Iterate over all nodes in a workspace."""
         response = self._client.search(index=workspace_id, body={"query": {"match_all": {}}, "size": max_size})
         for doc in response['hits']['hits']:
-            yield self.doc2node(doc, workspace_id)
+            node = self.doc2node(doc, workspace_id)
+            if callback_fn:
+                yield callback_fn(node)
+            else:
+                yield node
 
     def refresh(self, workspace_id: str):
         self._client.indices.refresh(index=workspace_id)
@@ -65,35 +73,51 @@ class EsVectorStore(LocalVectorStore):
             node.metadata["score"] = doc["_score"] - 1
         return node
 
-    def add_term_filter(self, key: str, value):
-        if key:
-            self.retrieve_filters.append({"term": {key: value}})
-        return self
-
-    def add_range_filter(self, key: str, gte=None, lte=None):
-        if key:
-            if gte is not None and lte is not None:
-                self.retrieve_filters.append({"range": {key: {"gte": gte, "lte": lte}}})
-            elif gte is not None:
-                self.retrieve_filters.append({"range": {key: {"gte": gte}}})
-            elif lte is not None:
-                self.retrieve_filters.append({"range": {key: {"lte": lte}}})
-        return self
-
-    def clear_filter(self):
-        self.retrieve_filters.clear()
-        return self
+    @staticmethod
+    def _build_es_filters(filter_dict: Optional[Dict[str, Any]] = None) -> List[Dict]:
+        """Build Elasticsearch filter clauses from filter_dict"""
+        if not filter_dict:
+            return []
 
-    def search(self, query: str, workspace_id: str, top_k: int = 1, **kwargs) -> List[VectorNode]:
+        filters = []
+        for key, filter_value in filter_dict.items():
+            # Handle nested keys by prefixing with metadata.
+            es_key = f"metadata.{key}" if not key.startswith("metadata.") else key
+
+            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:
+                    filters.append({"range": {es_key: range_conditions}})
+            else:
+                # Term filter: direct value comparison
+                filters.append({"term": {es_key: filter_value}})
+
+        return filters
+
+    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 []
 
         query_vector = self.embedding_model.get_embeddings(query)
+
+        # Build filters from filter_dict
+        es_filters = self._build_es_filters(filter_dict)
+
         body = {
             "query": {
                 "script_score": {
-                    "query": {"bool": {"must": self.retrieve_filters}},
+                    "query": {"bool": {"must": es_filters}} if es_filters else {"match_all": {}},
                     "script": {
                         "source": "cosineSimilarity(params.query_vector, 'vector') + 1.0",
                         "params": {"query_vector": query_vector},
@@ -106,9 +130,10 @@ class EsVectorStore(LocalVectorStore):
 
         nodes: List[VectorNode] = []
         for doc in response['hits']['hits']:
-            nodes.append(self.doc2node(doc, workspace_id))
+            node = self.doc2node(doc, workspace_id)
+            node.metadata["score"] = doc["_score"] - 1  # Adjust score since we added 1.0
+            nodes.append(node)
 
-        self.retrieve_filters.clear()
         return nodes
 
     def insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, refresh: bool = True, **kwargs):
@@ -160,6 +185,134 @@ class EsVectorStore(LocalVectorStore):
         if refresh:
             self.refresh(workspace_id=workspace_id)
 
+
+    # Async methods using native Elasticsearch async APIs
+    async def async_exist_workspace(self, workspace_id: str, **kwargs) -> bool:
+        """Async version of exist_workspace using native ES async client"""
+        return await self._async_client.indices.exists(index=workspace_id)
+
+    async def async_delete_workspace(self, workspace_id: str, **kwargs):
+        """Async version of delete_workspace using native ES async client"""
+        return await self._async_client.indices.delete(index=workspace_id, **kwargs)
+
+    async def async_create_workspace(self, workspace_id: str, **kwargs):
+        """Async version of create_workspace using native ES async client"""
+        body = {
+            "mappings": {
+                "properties": {
+                    "workspace_id": {"type": "keyword"},
+                    "content": {"type": "text"},
+                    "metadata": {"type": "object"},
+                    "vector": {
+                        "type": "dense_vector",
+                        "dims": self.embedding_model.dimensions
+                    }
+                }
+            }
+        }
+        return await self._async_client.indices.create(index=workspace_id, body=body)
+
+    async def async_refresh(self, workspace_id: str):
+        """Async version of refresh using native ES async client"""
+        await self._async_client.indices.refresh(index=workspace_id)
+
+    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 native ES async client and async embedding"""
+        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 filters from filter_dict
+        es_filters = self._build_es_filters(filter_dict)
+
+        body = {
+            "query": {
+                "script_score": {
+                    "query": {"bool": {"must": es_filters}} if es_filters else {"match_all": {}},
+                    "script": {
+                        "source": "cosineSimilarity(params.query_vector, 'vector') + 1.0",
+                        "params": {"query_vector": query_vector},
+                    }
+                }
+            },
+            "size": top_k
+        }
+        response = await self._async_client.search(index=workspace_id, body=body, **kwargs)
+
+        nodes: List[VectorNode] = []
+        for doc in response['hits']['hits']:
+            node = self.doc2node(doc, workspace_id)
+            node.metadata["score"] = doc["_score"] - 1  # Adjust score since we added 1.0
+            nodes.append(node)
+
+        return nodes
+
+    async def async_insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, refresh: bool = True,
+                           **kwargs):
+        """Async version of insert using native ES async client and async embedding"""
+        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)
+
+        docs = [
+            {
+                "_op_type": "index",
+                "_index": workspace_id,
+                "_id": node.unique_id,
+                "_source": {
+                    "workspace_id": workspace_id,
+                    "content": node.content,
+                    "metadata": node.metadata,
+                    "vector": node.vector
+                }
+            } for node in embedded_nodes + now_embedded_nodes]
+
+        status, error = await async_bulk(self._async_client, docs, chunk_size=self.batch_size, **kwargs)
+        logger.info(f"async insert docs.size={len(docs)} status={status} error={error}")
+
+        if refresh:
+            await self.async_refresh(workspace_id=workspace_id)
+
+    async def async_delete(self, node_ids: str | List[str], workspace_id: str, refresh: bool = True, **kwargs):
+        """Async version of delete using native ES async client"""
+        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]
+
+        actions = [
+            {
+                "_op_type": "delete",
+                "_index": workspace_id,
+                "_id": node_id
+            } for node_id in node_ids]
+
+        status, error = await async_bulk(self._async_client, actions, chunk_size=self.batch_size, **kwargs)
+        logger.info(f"async delete actions.size={len(actions)} status={status} error={error}")
+
+        if refresh:
+            await self.async_refresh(workspace_id=workspace_id)
+
+    def close(self):
+        self._client.close()
+
+    async def async_close(self):
+        await self._async_client.close()
+
 def main():
     from flowllm.utils.common_utils import load_env
     from flowllm.embedding_model import OpenAICompatibleEmbeddingModel
@@ -207,21 +360,124 @@ def main():
 
     es.insert(sample_nodes, workspace_id=workspace_id, refresh=True)
 
-    logger.info("=" * 20)
-    results = es.add_term_filter(key="metadata.node_type", value="n1") \
-        .search("What is AI?", top_k=5, workspace_id=workspace_id)
+    logger.info("=" * 20 + " FILTER TEST " + "=" * 20)
+    filter_dict = {"node_type": "n1"}
+    results = es.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)
 
-    logger.info("=" * 20)
+    logger.info("=" * 20 + " UNFILTERED TEST " + "=" * 20)
     results = es.search("What is AI?", top_k=5, workspace_id=workspace_id)
+    logger.info(f"Unfiltered results: {len(results)} results")
     for r in results:
         logger.info(r.model_dump(exclude={"vector"}))
     logger.info("=" * 20)
     es.dump_workspace(workspace_id=workspace_id)
     es.delete_workspace(workspace_id=workspace_id)
 
+    es.close()
+
+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 = "async_rag_nodes_index"
+    hosts = "http://11.160.132.46:8200"
+
+    # Use async context manager to ensure proper cleanup
+    es = EsVectorStore(hosts=hosts, embedding_model=embedding_model)
+    # Clean up and create workspace
+    if await es.async_exist_workspace(workspace_id=workspace_id):
+        await es.async_delete_workspace(workspace_id=workspace_id)
+    await es.async_create_workspace(workspace_id=workspace_id)
+
+    sample_nodes = [
+        VectorNode(
+            unique_id="async_es_node1",
+            workspace_id=workspace_id,
+            content="Artificial intelligence is a technology that simulates human intelligence.",
+            metadata={
+                "node_type": "n1",
+            }
+        ),
+        VectorNode(
+            unique_id="async_es_node2",
+            workspace_id=workspace_id,
+            content="AI is the future of mankind.",
+            metadata={
+                "node_type": "n1",
+            }
+        ),
+        VectorNode(
+            unique_id="async_es_node3",
+            workspace_id=workspace_id,
+            content="I want to eat fish!",
+            metadata={
+                "node_type": "n2",
+            }
+        ),
+        VectorNode(
+            unique_id="async_es_node4",
+            workspace_id=workspace_id,
+            content="The bigger the storm, the more expensive the fish.",
+            metadata={
+                "node_type": "n1",
+            }
+        ),
+    ]
+
+    # Test async insert
+    await es.async_insert(sample_nodes, workspace_id=workspace_id, refresh=True)
+
+    logger.info("ASYNC TEST - " + "=" * 20)
+    # Test async search with filter
+    filter_dict = {"node_type": "n1"}
+    results = await es.async_search("What is AI?", top_k=5, workspace_id=workspace_id, filter_dict=filter_dict)
+    for r in results:
+        logger.info(r.model_dump(exclude={"vector"}))
+    logger.info("=" * 20)
+
+    # Test async search without filter
+    logger.info("ASYNC TEST WITHOUT FILTER - " + "=" * 20)
+    results = await es.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_es_node2",
+        workspace_id=workspace_id,
+        content="AI is the future of humanity and technology.",
+        metadata={
+            "node_type": "n1",
+            "updated": True
+        }
+    )
+    await es.async_delete(node2_update.unique_id, workspace_id=workspace_id, refresh=True)
+    await es.async_insert(node2_update, workspace_id=workspace_id, refresh=True)
+
+    logger.info("ASYNC Updated Result:")
+    results = await es.async_search("fish?", workspace_id=workspace_id, top_k=10)
+    for r in results:
+        logger.info(r.model_dump(exclude={"vector"}))
+    logger.info("=" * 20)
+
+    # Clean up
+    await es.async_dump_workspace(workspace_id=workspace_id)
+    await es.async_delete_workspace(workspace_id=workspace_id)
+
+    await es.async_close()
+
 
 if __name__ == "__main__":
     main()
+
+    # Run async test
+    logger.info("\n" + "=" * 50 + " ASYNC TESTS " + "=" * 50)
+    asyncio.run(async_main())

flowllm/storage/vector_store/local_vector_store.py

@@ -1,8 +1,10 @@
+import asyncio
 import fcntl
 import json
 import math
+from functools import partial
 from pathlib import Path
-from typing import List, Iterable
+from typing import List, Iterable, Optional, Dict, Any
 
 from loguru import logger
 from pydantic import Field, model_validator
@@ -15,7 +17,7 @@ from flowllm.storage.vector_store.base_vector_store import BaseVectorStore
 
 @C.register_vector_store("local")
 class LocalVectorStore(BaseVectorStore):
-    store_dir: str = Field(default="./file_vector_store")
+    store_dir: str = Field(default="./local_vector_store")
 
     @model_validator(mode="after")
     def init_client(self):
@@ -88,8 +90,9 @@ class LocalVectorStore(BaseVectorStore):
     def create_workspace(self, workspace_id: str, **kwargs):
         self._dump_to_path(nodes=[], workspace_id=workspace_id, path=self.store_path, **kwargs)
 
-    def _iter_workspace_nodes(self, workspace_id: str, **kwargs) -> Iterable[VectorNode]:
-        for i, node in enumerate(self._load_from_path(path=self.store_path, workspace_id=workspace_id, **kwargs)):
+    def iter_workspace_nodes(self, workspace_id: str, callback_fn=None, **kwargs):
+        for node in self._load_from_path(path=self.store_path, workspace_id=workspace_id, callback_fn=callback_fn,
+                                         **kwargs):
             yield node
 
     def dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs):
@@ -97,7 +100,8 @@ class LocalVectorStore(BaseVectorStore):
             logger.warning(f"workspace_id={workspace_id} is not exist!")
             return {}
 
-        return self._dump_to_path(nodes=self._iter_workspace_nodes(workspace_id=workspace_id, **kwargs),
+        return self._dump_to_path(
+            nodes=self.iter_workspace_nodes(workspace_id=workspace_id, callback_fn=callback_fn, **kwargs),
                                   workspace_id=workspace_id,
                                   path=path,
                                   callback_fn=callback_fn,
@@ -129,7 +133,7 @@ class LocalVectorStore(BaseVectorStore):
 
         nodes = []
         node_size = 0
-        for node in self._iter_workspace_nodes(workspace_id=src_workspace_id, **kwargs):
+        for node in self.iter_workspace_nodes(workspace_id=src_workspace_id, **kwargs):
             nodes.append(node)
             node_size += 1
             if len(nodes) >= self.batch_size:
@@ -140,6 +144,39 @@ class LocalVectorStore(BaseVectorStore):
             self.insert(nodes=nodes, workspace_id=dest_workspace_id, **kwargs)
         return {"size": node_size}
 
+    @staticmethod
+    def _matches_filters(node: VectorNode, filter_dict: dict = None) -> bool:
+        """Check if a node matches all filters in filter_dict"""
+        if not filter_dict:
+            return True
+
+        for key, filter_value in filter_dict.items():
+            # Navigate nested keys (e.g., "metadata.node_type")
+            value = node.metadata
+            for key_part in key.split('.'):
+                if isinstance(value, dict) and key_part in value:
+                    value = value[key_part]
+                else:
+                    return False  # Key not found
+
+            # Handle different filter types
+            if isinstance(filter_value, dict):
+                # Range filter: {"gte": 1, "lte": 10}
+                if "gte" in filter_value and value < filter_value["gte"]:
+                    return False
+                if "lte" in filter_value and value > filter_value["lte"]:
+                    return False
+                if "gt" in filter_value and value <= filter_value["gt"]:
+                    return False
+                if "lt" in filter_value and value >= filter_value["lt"]:
+                    return False
+            else:
+                # Term filter: direct value comparison
+                if value != filter_value:
+                    return False
+        
+        return True
+
     @staticmethod
     def calculate_similarity(query_vector: List[float], node_vector: List[float]):
         assert query_vector, f"query_vector is empty!"
@@ -152,12 +189,15 @@ class LocalVectorStore(BaseVectorStore):
         norm_v2 = math.sqrt(sum(y ** 2 for y in node_vector))
         return dot_product / (norm_v1 * norm_v2)
 
-    def search(self, query: str, workspace_id: str, top_k: int = 1, **kwargs) -> List[VectorNode]:
+    def search(self, query: str, workspace_id: str, top_k: int = 1, filter_dict: Optional[Dict[str, Any]] = None,
+               **kwargs) -> List[VectorNode]:
         query_vector = self.embedding_model.get_embeddings(query)
         nodes: List[VectorNode] = []
         for node in self._load_from_path(path=self.store_path, workspace_id=workspace_id, **kwargs):
-            node.metadata["score"] = self.calculate_similarity(query_vector, node.vector)
-            nodes.append(node)
+            # Apply filters
+            if self._matches_filters(node, filter_dict):
+                node.metadata["score"] = self.calculate_similarity(query_vector, node.vector)
+                nodes.append(node)
 
         nodes = sorted(nodes, key=lambda x: x.metadata["score"], reverse=True)
         return nodes[:top_k]
@@ -203,6 +243,66 @@ class LocalVectorStore(BaseVectorStore):
         self._dump_to_path(nodes=all_nodes, workspace_id=workspace_id, path=self.store_path, **kwargs)
         logger.info(f"delete workspace_id={workspace_id} before_size={before_size} after_size={after_size}")
 
+    # Override async methods for better performance with file I/O
+    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 embedding model async capabilities"""
+        query_vector = await self.embedding_model.get_embeddings_async(query)
+
+        # Load nodes asynchronously
+        loop = asyncio.get_event_loop()
+        nodes_iter = await loop.run_in_executor(
+            C.thread_pool,
+            partial(self._load_from_path, path=self.store_path, workspace_id=workspace_id, **kwargs)
+        )
+
+        nodes: List[VectorNode] = []
+        for node in nodes_iter:
+            # Apply filters
+            if self._matches_filters(node, filter_dict):
+                node.metadata["score"] = self.calculate_similarity(query_vector, node.vector)
+                nodes.append(node)
+
+        nodes = sorted(nodes, key=lambda x: x.metadata["score"], reverse=True)
+        return nodes[:top_k]
+
+    async def async_insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs):
+        """Async version of insert using embedding model async capabilities"""
+        if isinstance(nodes, VectorNode):
+            nodes = [nodes]
+
+        # Use async embedding
+        nodes = await self.embedding_model.get_node_embeddings_async(nodes)
+
+        # Load existing nodes asynchronously
+        loop = asyncio.get_event_loop()
+        exist_nodes_iter = await loop.run_in_executor(
+            C.thread_pool,
+            partial(self._load_from_path, path=self.store_path, workspace_id=workspace_id)
+        )
+
+        all_node_dict = {}
+        exist_nodes: List[VectorNode] = list(exist_nodes_iter)
+        for node in exist_nodes:
+            all_node_dict[node.unique_id] = node
+
+        update_cnt = 0
+        for node in nodes:
+            if node.unique_id in all_node_dict:
+                update_cnt += 1
+            all_node_dict[node.unique_id] = node
+
+        # Dump to path asynchronously
+        await loop.run_in_executor(
+            C.thread_pool,
+            partial(self._dump_to_path, nodes=list(all_node_dict.values()),
+                    workspace_id=workspace_id, path=self.store_path, **kwargs)
+        )
+
+        logger.info(f"update workspace_id={workspace_id} nodes.size={len(nodes)} all.size={len(all_node_dict)} "
+                    f"update_cnt={update_cnt}")
+
+
 def main():
     from flowllm.utils.common_utils import load_env
     from flowllm.embedding_model import OpenAICompatibleEmbeddingModel
@@ -253,10 +353,107 @@ def main():
     for r in results:
         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 = client.search("What is AI?", workspace_id=workspace_id, top_k=5, 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)
     client.dump_workspace(workspace_id)
 
     client.delete_workspace(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 = "async_rag_nodes_index"
+    client = LocalVectorStore(embedding_model=embedding_model, store_dir="./async_file_vector_store")
+
+    # Clean up and create workspace
+    if await client.async_exist_workspace(workspace_id):
+        await client.async_delete_workspace(workspace_id)
+    await client.async_create_workspace(workspace_id)
+
+    sample_nodes = [
+        VectorNode(
+            unique_id="async_local_node1",
+            workspace_id=workspace_id,
+            content="Artificial intelligence is a technology that simulates human intelligence.",
+            metadata={
+                "node_type": "n1",
+            }
+        ),
+        VectorNode(
+            unique_id="async_local_node2",
+            workspace_id=workspace_id,
+            content="AI is the future of mankind.",
+            metadata={
+                "node_type": "n1",
+            }
+        ),
+        VectorNode(
+            unique_id="async_local_node3",
+            workspace_id=workspace_id,
+            content="I want to eat fish!",
+            metadata={
+                "node_type": "n2",
+            }
+        ),
+        VectorNode(
+            unique_id="async_local_node4",
+            workspace_id=workspace_id,
+            content="The bigger the storm, the more expensive the fish.",
+            metadata={
+                "node_type": "n1",
+            }
+        ),
+    ]
+
+    # Test async insert
+    await client.async_insert(sample_nodes, workspace_id)
+
+    logger.info("ASYNC TEST - " + "=" * 20)
+    # Test async search
+    results = await client.async_search("What is AI?", workspace_id=workspace_id, top_k=5)
+    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_local_node2",
+        workspace_id=workspace_id,
+        content="AI is the future of humanity and technology.",
+        metadata={
+            "node_type": "n1",
+            "updated": True
+        }
+    )
+    await client.async_delete(node2_update.unique_id, workspace_id=workspace_id)
+    await client.async_insert(node2_update, workspace_id=workspace_id)
+
+    logger.info("ASYNC Updated Result:")
+    results = await client.async_search("fish?", workspace_id=workspace_id, top_k=10)
+    for r in results:
+        logger.info(r.model_dump(exclude={"vector"}))
+    logger.info("=" * 20)
+
+    # Clean up
+    await client.async_dump_workspace(workspace_id)
+    await client.async_delete_workspace(workspace_id)
+
+
 if __name__ == "__main__":
     main()
+
+    # Run async test
+    logger.info("\n" + "=" * 50 + " ASYNC TESTS " + "=" * 50)
+    asyncio.run(async_main())