diff --git a/v1/tool/base_tool.py b/v1/tool/base_tool.py index 1f5edb85..0dc09edf 100644 --- a/v1/tool/base_tool.py +++ b/v1/tool/base_tool.py @@ -21,6 +21,7 @@ class BaseTool(BaseModel, ABC): def reset(self): self.arguments.clear() + self.cached_result.clear() self.success = True def _execute(self, **kwargs): diff --git a/v1/vector_store/base_vector_store.py b/v1/vector_store/base_vector_store.py index 1cf1e2fc..d13586be 100644 --- a/v1/vector_store/base_vector_store.py +++ b/v1/vector_store/base_vector_store.py @@ -1,42 +1,101 @@ +import fcntl +import json from abc import ABC -from typing import List +from pathlib import Path +from typing import List, Iterable +from loguru import logger from pydantic import BaseModel, Field +from tqdm import tqdm -from v1.model.base_embedding_model import BaseEmbeddingModel +from v1.embedding_model.base_embedding_model import BaseEmbeddingModel from v1.schema.vector_node import VectorNode class BaseVectorStore(BaseModel, ABC): embedding_model: BaseEmbeddingModel = Field(default=...) + @staticmethod + def _load_from_path(path: str | Path, workspace_id: str, **kwargs) -> Iterable[VectorNode]: + workspace_path = Path(path) / f"{workspace_id}.jsonl" + if workspace_path.exists(): + with workspace_path.open() as f: + fcntl.flock(f, fcntl.LOCK_SH) + try: + for line in tqdm(f, desc="load from path"): + if line.strip(): + yield VectorNode(**json.loads(line.strip(), **kwargs)) + + finally: + fcntl.flock(f, fcntl.LOCK_UN) + + @staticmethod + def _dump_to_path(nodes: Iterable[VectorNode], workspace_id: str, path: str | Path = "", + ensure_ascii: bool = False, **kwargs): + dump_path: Path = Path(path) + dump_path.mkdir(parents=True, exist_ok=True) + dump_file = dump_path / f"{workspace_id}.jsonl" + + with dump_file.open("w") as f: + fcntl.flock(f, fcntl.LOCK_EX) + try: + for node in tqdm(nodes, desc="dump to path"): + f.write(json.dumps(node.model_dump(), ensure_ascii=ensure_ascii, **kwargs)) + f.write("\n") + finally: + fcntl.flock(f, fcntl.LOCK_UN) + def exist_workspace(self, workspace_id: str, **kwargs) -> bool: raise NotImplementedError + def _delete_workspace(self, workspace_id: str, **kwargs): + raise NotImplementedError + def delete_workspace(self, workspace_id: str, **kwargs): + if self.exist_workspace(workspace_id, **kwargs): + self._delete_workspace(workspace_id, **kwargs) + + def _create_workspace(self, workspace_id: str, **kwargs): raise NotImplementedError def create_workspace(self, workspace_id: str, **kwargs): + if self.exist_workspace(workspace_id, **kwargs): + logger.warning(f"workspace={workspace_id} exists~") + return + self._create_workspace(workspace_id, **kwargs) + + def _iter_workspace_nodes(self, workspace_id: str, max_size: int = 10000, **kwargs) -> Iterable[VectorNode]: raise NotImplementedError - def dump_workspace(self, workspace_id: str, **kwargs): - raise NotImplementedError + def dump_workspace(self, workspace_id: str, path: str | Path = "", **kwargs): + self._dump_to_path(nodes=self._iter_workspace_nodes(workspace_id, **kwargs), + workspace_id=workspace_id, + path=path, **kwargs) - def load_workspace(self, workspace_id: str, **kwargs): - raise NotImplementedError + def load_workspace(self, workspace_id: str, path: str | Path = "", nodes: List[VectorNode] = None, **kwargs): + self.create_workspace(workspace_id=workspace_id, **kwargs) + all_nodes: List[VectorNode] = [] + all_nodes.extend(nodes) + for node in self._load_from_path(path=path, workspace_id=workspace_id, **kwargs): + all_nodes.append(node) + self.insert(nodes=all_nodes, workspace_id=workspace_id, **kwargs) def retrieve_by_query(self, query: str, workspace_id: str, top_k: int = 1, **kwargs) -> List[VectorNode]: raise NotImplementedError - def retrieve_by_id(self, unique_id: str, workspace_id: str = None, **kwargs) -> VectorNode | None: - raise NotImplementedError - def insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs): raise NotImplementedError def update(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs): raise NotImplementedError + """ + unimportant + """ + + def retrieve_by_id(self, unique_id: str, workspace_id: str = None, **kwargs) -> VectorNode | None: + raise NotImplementedError + def exist_id(self, unique_id: str, workspace_id: str = None, **kwargs) -> bool: raise NotImplementedError diff --git a/v1/vector_store/es_vector_store.py b/v1/vector_store/es_vector_store.py index 7c259532..f587dc76 100644 --- a/v1/vector_store/es_vector_store.py +++ b/v1/vector_store/es_vector_store.py @@ -1,14 +1,15 @@ import os -from typing import List, Tuple +from typing import List, Tuple, Iterable from elasticsearch import Elasticsearch from elasticsearch.helpers import bulk from loguru import logger from pydantic import Field, PrivateAttr, model_validator -from v1.model.openai_compatible_embedding_model import OpenAICompatibleEmbeddingModel -from v1.schema.vector_store_node import VectorStoreNode -from v1.storage.base_vector_store import BaseVectorStore, VECTOR_STORE_REGISTRY +from v1.embedding_model.openai_compatible_embedding_model import OpenAICompatibleEmbeddingModel +from v1.schema.vector_node import VectorNode +from v1.vector_store import VECTOR_STORE_REGISTRY +from v1.vector_store.base_vector_store import BaseVectorStore @VECTOR_STORE_REGISTRY.register("elasticsearch") @@ -28,26 +29,14 @@ class EsVectorStore(BaseVectorStore): self._client = Elasticsearch(hosts=hosts, basic_auth=self.basic_auth) return self - def exist_index(self, index_name: str = None) -> bool: - if index_name is None: - index_name = self.index_name - return self._client.indices.exists(index=index_name) + def exist_workspace(self, workspace_id: str, **kwargs) -> bool: + return self._client.indices.exists(index=workspace_id) - def delete_index(self, index_name: str = None): - if index_name is None: - index_name = self.index_name - if self._client.indices.exists(index=index_name): - self._client.indices.delete(index=index_name) + def _delete_workspace(self, workspace_id: str, **kwargs): + self._client.indices.delete(index=workspace_id, **kwargs) - def create_index(self, index_name: str = None): - if index_name is None: - index_name = self.index_name - - if self._client.indices.exists(index=index_name): - logger.warning(f"index_name={index_name} is already exists!") - return None - - index = { + def _create_workspace(self, workspace_id: str, **kwargs): + body = { "mappings": { "properties": { "workspace_id": {"type": "keyword"}, @@ -61,32 +50,37 @@ class EsVectorStore(BaseVectorStore): } } - return self._client.indices.create(index=index_name, body=index) + return self._client.indices.create(index=workspace_id, body=body) - def refresh_index(self, index_name: str = None): - if index_name is None: - index_name = self.index_name - self._client.indices.refresh(index=index_name) + def _iter_workspace_nodes(self, workspace_id: str, max_size: int = 10000, **kwargs) -> Iterable[VectorNode]: + response = self._client.search( + index=workspace_id, + body={"query": {"match_all": {}}}, + scroll='5m', + size=max_size + ) + + for doc in response['hits']['hits']: + yield self.doc2node(doc) + + def refresh(self, workspace_id: str): + self._client.indices.refresh(index=workspace_id) @staticmethod - def doc2node(doc) -> VectorStoreNode: - node = VectorStoreNode(**doc["_source"]) + def doc2node(doc) -> VectorNode: + node = VectorNode(**doc["_source"]) node.unique_id = doc["_id"] if "_score" in doc: node.metadata["_score"] = doc["_score"] - 1 return node - def exist_id(self, unique_id: str, index_name: str = None): - if index_name is None: - index_name = self.index_name - return self._client.exists(index=index_name, id=unique_id) - - def node2doc(self, node: VectorStoreNode, add_op_type: bool = False, index_name: str = None) -> dict: - if index_name is None: - index_name = self.index_name + def exist_id(self, unique_id: str, workspace_id: str = None, **kwargs) -> bool: + response = self._client.exists(index=workspace_id, id=unique_id) + return response.body + def node2doc(self, node: VectorNode, add_op_type: bool = False) -> dict: doc: dict = { - "_index": index_name, + "_index": node.workspace_id, "_id": node.unique_id, "_source": { "workspace_id": node.workspace_id, @@ -97,7 +91,7 @@ class EsVectorStore(BaseVectorStore): } if add_op_type: - doc["_op_type"] = "update" if self.exist_id(node.unique_id, index_name) else "index", + doc["_op_type"] = "update" if self.exist_id(node.unique_id, node.workspace_id) else "index", return doc def add_term_filter(self, key: str, value): @@ -119,84 +113,12 @@ class EsVectorStore(BaseVectorStore): self.retrieve_filters.clear() return self - def insert(self, nodes: VectorStoreNode | List[VectorStoreNode], refresh_index: bool = True, index_name: str = None, - **kwargs): - if index_name is None: - index_name = self.index_name - - if not self.exist_index(index_name): - self.create_index(index_name) - - if isinstance(nodes, VectorStoreNode): - nodes = [nodes] - - embedded_nodes = [node for node in nodes if node.vector] - not_embedded_nodes = [node for node in nodes if not node.vector] - now_embedded_nodes = self.embedding_model.get_node_embeddings(not_embedded_nodes) - - docs = [self.node2doc(node, False, index_name) for node in embedded_nodes + now_embedded_nodes] - status, error = bulk(self._client, docs, chunk_size=self.bulk_chunk_size, **kwargs) - logger.info(f"insert sample.size={len(nodes)} status={status} error={error}") - - if refresh_index: - self.refresh_index(index_name) - - def update(self, nodes: VectorStoreNode | List[VectorStoreNode], refresh_index: bool = True, index_name: str = None, - **kwargs): - if index_name is None: - index_name = self.index_name - - if not self.exist_index(index_name): - self.create_index(index_name) - - if isinstance(nodes, VectorStoreNode): - nodes = [nodes] - - nodes = self.embedding_model.get_node_embeddings(nodes) - docs = [self.node2doc(node, True, index_name) for node in nodes] - status, error = bulk(self._client, docs, chunk_size=self.bulk_chunk_size, **kwargs) - update_size = sum([1 if doc["_op_type"] == "update" else 0 for doc in docs]) - insert_size = len(docs) - update_size - logger.info(f"update update_size={update_size} insert_size={insert_size} status={status} error={error}") - - if refresh_index: - self.refresh_index(index_name) - - def delete_by_id(self, unique_id: str, index_name: str = None, **kwargs): - if index_name is None: - index_name = self.index_name - - if not self.exist_index(index_name): - self.create_index(index_name) - - return self._client.delete(index=index_name, id=unique_id, **kwargs) - - def retrieve_by_id(self, unique_id: str, index_name: str = None, **kwargs) -> VectorStoreNode | None: - if index_name is None: - index_name = self.index_name - - if not self.exist_index(index_name): - logger.warning(f"index_name={index_name} is not exists!") - return None - - try: - doc = self._client.get(index=index_name, id=unique_id, **kwargs) - return self.doc2node(doc) - - except Exception as e: - logger.warning(f"{index_name} retrieve_by_id unique_id={unique_id} is not found with error={e.args}") - return None - - def retrieve_by_query(self, query: str, top_k: int = 1, index_name: str = None, **kwargs) -> List[VectorStoreNode]: - if index_name is None: - index_name = self.index_name - - if not self.exist_index(index_name): - logger.warning(f"index_name={index_name} is not exists!") + def retrieve_by_query(self, query: str, workspace_id: str, top_k: int = 1, **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) - body = { "query": { "script_score": { @@ -209,51 +131,82 @@ class EsVectorStore(BaseVectorStore): }, "size": top_k } - response = self._client.search(index=index_name, body=body, **kwargs) + response = self._client.search(index=workspace_id, body=body, **kwargs) - nodes: List[VectorStoreNode] = [] + nodes: List[VectorNode] = [] for doc in response['hits']['hits']: nodes.append(self.doc2node(doc)) self.retrieve_filters.clear() return nodes + def insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, refresh: bool = True, **kwargs): + self.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] + now_embedded_nodes = self.embedding_model.get_node_embeddings(not_embedded_nodes) + + docs = [self.node2doc(node, False) for node in embedded_nodes + now_embedded_nodes] + status, error = bulk(self._client, docs, chunk_size=self.bulk_chunk_size, **kwargs) + logger.info(f"insert sample.size={len(nodes)} status={status} error={error}") + + if refresh: + self.refresh(workspace_id=workspace_id) + + def update(self, nodes: VectorNode | List[VectorNode], workspace_id: str, refresh: bool = True, **kwargs): + self.create_workspace(workspace_id=workspace_id) + if isinstance(nodes, VectorNode): + nodes = [nodes] + + nodes = self.embedding_model.get_node_embeddings(nodes) + docs = [self.node2doc(node, True) for node in nodes] + status, error = bulk(self._client, docs, chunk_size=self.bulk_chunk_size, **kwargs) + update_size = sum([1 if doc["_op_type"] == "update" else 0 for doc in docs]) + insert_size = len(docs) - update_size + logger.info(f"update update_size={update_size} insert_size={insert_size} status={status} error={error}") + + if refresh: + self.refresh(workspace_id=workspace_id) def main(): from experiencemaker.utils.util_function import load_env_keys load_env_keys() + load_env_keys("../../.env") embedding_model = OpenAICompatibleEmbeddingModel(dimensions=1024) - index_name = "rag_nodes_index" + workspace_id = "rag_nodes_index" hosts = "http://11.160.132.46:8200" - es = EsVectorStore(hosts=hosts, embedding_model=embedding_model, index_name=index_name) - es.delete_index() - es.create_index() + es = EsVectorStore(hosts=hosts, embedding_model=embedding_model) + es.delete_workspace(workspace_id=workspace_id) + es.create_workspace(workspace_id=workspace_id) sample_nodes = [ - VectorStoreNode( - workspace_id="w1", + VectorNode( + workspace_id=workspace_id, content="Artificial intelligence is a technology that simulates human intelligence.", metadata={ "node_type": "n1", } ), - VectorStoreNode( - workspace_id="w1", + VectorNode( + workspace_id=workspace_id, content="AI is the future of mankind.", metadata={ "node_type": "n1", } ), - VectorStoreNode( - workspace_id="w1", + VectorNode( + workspace_id=workspace_id, content="I want to eat fish!", metadata={ "node_type": "n2", } ), - VectorStoreNode( - workspace_id="w2", + VectorNode( + workspace_id=workspace_id, content="The bigger the storm, the more expensive the fish.", metadata={ "node_type": "n1", @@ -261,28 +214,22 @@ def main(): ), ] - es.insert(sample_nodes, refresh_index=True) + es.insert(sample_nodes, workspace_id=workspace_id, refresh=True) logger.info("=" * 20) - results = es.add_term_filter(key="workspace_id", value="w1") \ - .add_term_filter(key="metadata.node_type", value="n1") \ - .retrieve_by_query("What is AI?", top_k=5) + results = es.add_term_filter(key="metadata.node_type", value="n1") \ + .retrieve_by_query("What is AI?", top_k=5, workspace_id=workspace_id) for r in results: logger.info(r.model_dump(exclude={"vector"})) logger.info("=" * 20) logger.info("=" * 20) - results = es.add_term_filter(key="workspace_id", value="w1") \ - .retrieve_by_query("What is AI?", top_k=5) + results = es.retrieve_by_query("What is AI?", top_k=5, workspace_id=workspace_id) for r in results: logger.info(r.model_dump(exclude={"vector"})) logger.info("=" * 20) - logger.info("=" * 20) - results = es.retrieve_by_query("What is AI?", top_k=5) - for r in results: - logger.info(r.model_dump(exclude={"vector"})) - logger.info("=" * 20) + es.delete_workspace(workspace_id=workspace_id) if __name__ == "__main__": diff --git a/v1/vector_store/file_vector_store.py b/v1/vector_store/file_vector_store.py index 900ca3eb..d24cff35 100644 --- a/v1/vector_store/file_vector_store.py +++ b/v1/vector_store/file_vector_store.py @@ -1,136 +1,45 @@ -import json import math -import threading from pathlib import Path -from typing import List, Any +from typing import List, Iterable from loguru import logger -from pydantic import Field, model_validator, PrivateAttr +from pydantic import Field, model_validator -from experiencemaker.model.openai_compatible_embedding_model import OpenAICompatibleEmbeddingModel -from experiencemaker.schema.vector_store_node import VectorStoreNode -from experiencemaker.storage.base_vector_store import BaseVectorStore, VECTOR_STORE_REGISTRY +from v1.embedding_model.openai_compatible_embedding_model import OpenAICompatibleEmbeddingModel +from v1.schema.vector_node import VectorNode +from v1.vector_store import VECTOR_STORE_REGISTRY +from v1.vector_store.base_vector_store import BaseVectorStore @VECTOR_STORE_REGISTRY.register("local_file") class FileVectorStore(BaseVectorStore): store_dir: str = Field(default="./file_vector_store") - index_path: Path | None = Field(default=None) - _thread_lock: Any = PrivateAttr() @model_validator(mode="after") def init_client(self): - self._thread_lock = threading.Lock() store_path = Path(self.store_dir) store_path.mkdir(parents=True, exist_ok=True) - self.index_path = store_path / f"{self.index_name}.jsonl" - if not self.index_path.exists(): - self.index_path.touch(exist_ok=True) return self - def get_index_path(self, index_name: str = None) -> Path: - if index_name is None: - index_path = self.index_path - else: - store_path = Path(self.store_dir) - index_path = store_path / f"{self.index_name}.jsonl" - if not index_path.exists(): - index_path.touch(exist_ok=True) - return index_path + @property + def store_path(self) -> Path: + return Path(self.store_dir) - def exist_index(self, index_name: str = None) -> bool: - index_path = self.get_index_path(index_name) - with self._thread_lock: - return index_path.exists() + def exist_workspace(self, workspace_id: str, **kwargs) -> bool: + return (self.store_path / f"{workspace_id}.jsonl").exists() - def delete_index(self, index_name: str = None): - index_path = self.get_index_path(index_name) - with self._thread_lock: - if index_path.exists() and index_path.is_file(): - index_path.unlink() + def _delete_workspace(self, workspace_id: str, **kwargs): + workspace_path = self.store_path / f"{workspace_id}.jsonl" + if workspace_path.is_file(): + workspace_path.unlink() - def create_index(self, index_name: str = None): - index_path = self.get_index_path(index_name) - with self._thread_lock: - if not index_path.exists(): - index_path.touch(exist_ok=True) + def _create_workspace(self, workspace_id: str, **kwargs): + self._dump_to_path(nodes=[], workspace_id=workspace_id, path=self.store_path, **kwargs) - def _load(self, index_name: str = None) -> List[VectorStoreNode]: - index_path = self.get_index_path(index_name) - - nodes = [] - with self._thread_lock: - with open(index_path) as f: - for line in f: - if line.strip(): - nodes.append(VectorStoreNode(**json.loads(line))) - return nodes - - def _dump(self, nodes: List[VectorStoreNode], index_name: str = None): - index_path = self.get_index_path(index_name) - - with self._thread_lock: - with open(index_path, "w") as f: - for doc in nodes: - f.write(doc.model_dump_json() + "\n") - - def exist_id(self, unique_id: str, index_name: str = None): - nodes = self._load(index_name=index_name) - for node in nodes: - if node.unique_id == unique_id: - return True - return False - - def insert(self, nodes: VectorStoreNode | List[VectorStoreNode], index_name: str = None, **kwargs): - if index_name is None: - index_name = self.index_name - - return self.update(nodes, index_name=index_name, **kwargs) - - def update(self, nodes: VectorStoreNode | List[VectorStoreNode], index_name: str = None, **kwargs): - if index_name is None: - index_name = self.index_name - - if isinstance(nodes, VectorStoreNode): - nodes = [nodes] - - all_node_dict = {} - nodes: List[VectorStoreNode] = self.embedding_model.get_node_embeddings(nodes) - exist_nodes: List[VectorStoreNode] = self._load(index_name=index_name) - 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 - - self._dump(list(all_node_dict.values()), index_name=index_name) - logger.info( - f"update {index_name} nodes.size={len(nodes)} all.size={len(all_node_dict)} update_cnt={update_cnt}") - - def delete_by_id(self, unique_id: str, index_name: str = None, **kwargs): - if index_name is None: - index_name = self.index_name - - nodes = self._load(index_name=index_name) - dump_nodes: List[VectorStoreNode] = [] - for node in nodes: - if node.unique_id != unique_id: - dump_nodes.append(node) - - if len(dump_nodes) < len(nodes): - self._dump(dump_nodes, index_name=index_name) - logger.info(f"delete_by_id unique_id={unique_id}") - - def retrieve_by_id(self, unique_id: str, index_name: str = None, **kwargs) -> VectorStoreNode | None: - nodes = self._load(index_name=index_name) - for node in nodes: - if node.unique_id == unique_id: - return node - return None + def _iter_workspace_nodes(self, workspace_id: str, max_size: int = 10000, **kwargs) -> Iterable[VectorNode]: + for i, node in enumerate(self._load_from_path(path=self.store_path, workspace_id=workspace_id, **kwargs)): + if i < max_size: + yield node @staticmethod def calculate_similarity(query_vector: List[float], node_vector: List[float]): @@ -144,50 +53,79 @@ class FileVectorStore(BaseVectorStore): norm_v2 = math.sqrt(sum(y ** 2 for y in node_vector)) return dot_product / (norm_v1 * norm_v2) - def retrieve_by_query(self, query: str, top_k: int = 1, index_name: str = None, **kwargs) -> List[VectorStoreNode]: + def retrieve_by_query(self, query: str, workspace_id: str, top_k: int = 1, **kwargs) -> List[VectorNode]: query_vector = self.embedding_model.get_embeddings(query) - nodes: List[VectorStoreNode] = self._load(index_name=index_name) - for node in nodes: + 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) nodes = sorted(nodes, key=lambda x: x.metadata["score"], reverse=True) return nodes[:top_k] + def insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs): + return self.update(nodes=nodes, workspace_id=workspace_id, **kwargs) + + def update(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs): + if isinstance(nodes, VectorNode): + nodes = [nodes] + + all_node_dict = {} + nodes: List[VectorNode] = self.embedding_model.get_node_embeddings(nodes) + exist_nodes: List[VectorNode] = list(self._load_from_path(path=self.store_path, workspace_id=workspace_id)) + 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 + + 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 experiencemaker.utils.util_function import load_env_keys load_env_keys() + load_env_keys("../../.env") embedding_model = OpenAICompatibleEmbeddingModel(dimensions=1024) - index_name = "rag_nodes_index" - client = FileVectorStore(embedding_model=embedding_model, index_name=index_name) - client.delete_index() - client.create_index() + workspace_id = "rag_nodes_index" + client = FileVectorStore(embedding_model=embedding_model) + client.delete_workspace(workspace_id) + client.create_workspace(workspace_id) sample_nodes = [ - VectorStoreNode( - workspace_id="w1", + VectorNode( + workspace_id=workspace_id, content="Artificial intelligence is a technology that simulates human intelligence.", metadata={ "node_type": "n1", } ), - VectorStoreNode( - workspace_id="w1", + VectorNode( + workspace_id=workspace_id, content="AI is the future of mankind.", metadata={ "node_type": "n1", } ), - VectorStoreNode( - workspace_id="w1", + VectorNode( + workspace_id=workspace_id, content="I want to eat fish!", metadata={ "node_type": "n2", } ), - VectorStoreNode( - workspace_id="w2", + VectorNode( + workspace_id=workspace_id, content="The bigger the storm, the more expensive the fish.", metadata={ "node_type": "n1", @@ -195,14 +133,17 @@ def main(): ), ] - client.insert(sample_nodes) + client.insert(sample_nodes, workspace_id) logger.info("=" * 20) - results = client.retrieve_by_query("What is AI?", top_k=5) + results = client.retrieve_by_query("What is AI?", workspace_id=workspace_id, top_k=5) for r in results: logger.info(r.model_dump(exclude={"vector"})) logger.info("=" * 20) + client.delete_workspace(workspace_id) + client.dump_workspace(workspace_id) + if __name__ == "__main__": main()