diff --git a/experiencemaker/schema/experience.py b/experiencemaker/schema/experience.py index 397fb960..576344c7 100644 --- a/experiencemaker/schema/experience.py +++ b/experiencemaker/schema/experience.py @@ -5,7 +5,7 @@ from uuid import uuid4 from pydantic import BaseModel, Field from experiencemaker.schema.vector_store_node import VectorStoreNode - +# https://docs.trychroma.com/docs/overview/getting-started class ExperienceFunctionArg(BaseModel): arg_name: str = Field(default=..., description="argument name") diff --git a/experiencemaker/tool/base_tool.py b/experiencemaker/tool/base_tool.py index 4b3ffae5..6243fd51 100644 --- a/experiencemaker/tool/base_tool.py +++ b/experiencemaker/tool/base_tool.py @@ -13,7 +13,6 @@ class BaseTool(BaseModel, ABC): arguments: dict = Field(default_factory=dict, description="execute arguments") enable_cache: bool = Field(default=False, description="whether to cache the tool result") - # TODO add cache expire cached_result: dict = Field(default_factory=dict, description="tool execution result") max_retries: int = Field(default=3, description="max retries") diff --git a/v1/main.py b/v1/app.py similarity index 81% rename from v1/main.py rename to v1/app.py index 2011ece8..d19166c9 100644 --- a/v1/main.py +++ b/v1/app.py @@ -4,10 +4,10 @@ from concurrent.futures.thread import ThreadPoolExecutor import uvicorn from fastapi import FastAPI +from v1.config.config_parser import ConfigParser from v1.em_service import EMService from v1.schema.request import RetrieverRequest, SummarizerRequest, VectorStoreRequest, AgentRequest from v1.schema.response import RetrieverResponse, SummarizerResponse, VectorStoreResponse, AgentResponse -from v1.utils.config_parser import ConfigParser app = FastAPI() config_parser = ConfigParser(sys.argv[1:]) @@ -17,28 +17,28 @@ thread_pool: ThreadPoolExecutor = ThreadPoolExecutor(max_workers=global_app_conf @app.post('/retriever', response_model=RetrieverResponse) def call_retriever(request: RetrieverRequest): - app_config = config_parser.get_app_config(**request.params) + app_config = config_parser.get_app_config(**request.config) ems = EMService(app_config=app_config, thread_pool=thread_pool) return ems.call_retriever(request) @app.post('/summarizer', response_model=SummarizerResponse) def call_summarizer(request: SummarizerRequest): - app_config = config_parser.get_app_config(**request.params) + app_config = config_parser.get_app_config(**request.config) ems = EMService(app_config=app_config, thread_pool=thread_pool) return ems.call_summarizer(request) @app.post('/vector_store', response_model=VectorStoreResponse) def call_vector_store(request: VectorStoreRequest): - app_config = config_parser.get_app_config(**request.params) + app_config = config_parser.get_app_config(**request.config) ems = EMService(app_config=app_config, thread_pool=thread_pool) return ems.call_vector_store(request) @app.post('/agent', response_model=AgentResponse) def call_agent(request: AgentRequest): - app_config = config_parser.get_app_config(**request.params) + app_config = config_parser.get_app_config(**request.config) ems = EMService(app_config=app_config, thread_pool=thread_pool) return ems.call_agent(request) @@ -48,4 +48,5 @@ if __name__ == "__main__": host=global_app_config.http_service.host, port=global_app_config.http_service.port, timeout_keep_alive=global_app_config.http_service.timeout_keep_alive, - limit_concurrency=global_app_config.http_service.limit_concurrency) + limit_concurrency=global_app_config.http_service.limit_concurrency, + workers=global_app_config.http_service.workers) diff --git a/v1/utils/config_parser.py b/v1/config/config_parser.py similarity index 53% rename from v1/utils/config_parser.py rename to v1/config/config_parser.py index 6f03b610..bde3359f 100644 --- a/v1/utils/config_parser.py +++ b/v1/config/config_parser.py @@ -1,4 +1,5 @@ import json +from pathlib import Path from loguru import logger from omegaconf import OmegaConf, DictConfig @@ -6,24 +7,29 @@ from omegaconf import OmegaConf, DictConfig from v1.schema.app_config import AppConfig -class ConfigParser(object): +class ConfigParser: + default_config_name: str = "demo_config.yaml" def __init__(self, args: list): - assert len(args) >= 1, "The `args` require at least one argument." - assert "config_path=" in args[0], f"The 0th args must be config_path, for example: `config_path=XXX`." - config_path: str = args[0].replace("config_path=", "") - yaml_config = OmegaConf.load(config_path) - + # step1: default config self.app_config: DictConfig = OmegaConf.structured(AppConfig) + + # step2: load from config yaml file + cli_config: DictConfig = OmegaConf.from_dotlist(args) + config_path = cli_config.get("config_path") + if config_path: + config_path = Path(config_path) + else: + config_path = Path(__file__).parent / self.default_config_name + logger.info(f"load config from path={config_path}") + yaml_config = OmegaConf.load(config_path) self.app_config = OmegaConf.merge(self.app_config, yaml_config) - if args[1:]: - cli_config = OmegaConf.from_dotlist(args[1:]) - self.app_config = OmegaConf.merge(self.app_config, cli_config) + # merge cli config + self.app_config = OmegaConf.merge(self.app_config, cli_config) app_config_dict = OmegaConf.to_container(self.app_config, resolve=True) - app_config_str = json.dumps(app_config_dict, indent=2, ensure_ascii=False) - logger.info(f"app_config_str={app_config_str}") + logger.info(f"app_config_str={json.dumps(app_config_dict, indent=2, ensure_ascii=False)}") def get_app_config(self, **kwargs) -> AppConfig: app_config = self.app_config.copy() diff --git a/v1/config.yaml b/v1/config/demo_config.yaml similarity index 67% rename from v1/config.yaml rename to v1/config/demo_config.yaml index 24d8aa5f..af3aa8c8 100644 --- a/v1/config.yaml +++ b/v1/config/demo_config.yaml @@ -1,33 +1,20 @@ # demo config.yaml -#python -m experiencemaker.em_service \ -# --port=8001 \ -# --llm='{"backend": "openai_compatible", "model_name": "qwen3-32b", "temperature": 0.6}' \ -# --embedding_model='{"backend": "openai_compatible", "model_name": "text-embedding-v4", "dimensions": 1024}' \ -# --vector_store='{"backend": "elasticsearch"}' \ - http_service: host: "0.0.0.0" port: 8001 timeout_keep_alive: 600 limit_concurrency: 64 - -# -http_service.port=8001 + workers: 2 thread_pool: max_workers: 20 - api: step_retriever: mock1_op->mock2_op->mock3_op step_summarizer: mock1_op->[mock4_op->mock2_op|mock5_op]->mock3_op vector_store: mock6_op -# -api.step_retriever=mock1_op->[mock4_op->mock2_op|mock5_op]->mock3_op - - -# -op.mock1_op.a=1 - op: mock1_op: backend: mock1_op diff --git a/v1/config_handler.py b/v1/config_handler.py deleted file mode 100644 index 852e868b..00000000 --- a/v1/config_handler.py +++ /dev/null @@ -1,4 +0,0 @@ -class ConfigHandler(object): - - def __init__(self): - pass diff --git a/v1/em_service.py b/v1/em_service.py index 91d5c419..d57bd81e 100644 --- a/v1/em_service.py +++ b/v1/em_service.py @@ -1,5 +1,7 @@ from concurrent.futures import ThreadPoolExecutor +from loguru import logger + from v1.pipeline.pipeline import Pipeline from v1.pipeline.pipeline_context import PipelineContext from v1.schema.app_config import AppConfig @@ -33,30 +35,46 @@ class EMService(object): self.context.set_context("request", request) response = RetrieverResponse() self.context.set_context("response", response) - pipeline = Pipeline(pipeline=self.context.app_config.api.retriever, context=self.context) - pipeline.execute_pipeline() + try: + Pipeline(pipeline=self.context.app_config.api.retriever, context=self.context)() + except Exception as e: + logger.exception(f"call_retriever encounter error={e.args}") + response.success = False + response.metadata["error"] = str(e) return response def call_summarizer(self, request: SummarizerRequest) -> SummarizerResponse: self.context.set_context("request", request) response = SummarizerResponse() self.context.set_context("request", response) - pipeline = Pipeline(pipeline=self.context.app_config.api.summarizer, context=self.context) - pipeline.execute_pipeline() + try: + Pipeline(pipeline=self.context.app_config.api.summarizer, context=self.context)() + except Exception as e: + logger.exception(f"call_summarizer encounter error={e.args}") + response.success = False + response.metadata["error"] = str(e) return response def call_vector_store(self, request: VectorStoreRequest) -> VectorStoreResponse: self.context.set_context("request", request) response = VectorStoreResponse() self.context.set_context("request", response) - pipeline = Pipeline(pipeline=self.context.app_config.api.vector_store, context=self.context) - pipeline.execute_pipeline() + try: + Pipeline(pipeline=self.context.app_config.api.vector_store, context=self.context)() + except Exception as e: + logger.exception(f"call_vector_store encounter error={e.args}") + response.success = False + response.metadata["error"] = str(e) return response def call_agent(self, request: AgentRequest) -> AgentResponse: self.context.set_context("request", request) response = AgentResponse() self.context.set_context("request", response) - pipeline = Pipeline(pipeline=self.context.app_config.api.agent, context=self.context) - pipeline.execute_pipeline() + try: + Pipeline(pipeline=self.context.app_config.api.agent, context=self.context)() + except Exception as e: + logger.exception(f"call_agent encounter error={e.args}") + response.success = False + response.metadata["error"] = str(e) return response diff --git a/v1/embedding_model/__init__.py b/v1/embedding_model/__init__.py new file mode 100644 index 00000000..8a2426d3 --- /dev/null +++ b/v1/embedding_model/__init__.py @@ -0,0 +1,3 @@ +from experiencemaker.utils.registry import Registry + +EMBEDDING_MODEL_REGISTRY = Registry() \ No newline at end of file diff --git a/v1/embedding_model/base_embedding_model.py b/v1/embedding_model/base_embedding_model.py new file mode 100644 index 00000000..0c80405f --- /dev/null +++ b/v1/embedding_model/base_embedding_model.py @@ -0,0 +1,72 @@ +from abc import ABC +from typing import List + +from loguru import logger +from pydantic import BaseModel, Field + +from v1.schema.vector_node import VectorNode + + +class BaseEmbeddingModel(BaseModel, ABC): + model_name: str = Field(default=..., description="model name") + dimensions: int = Field(default=..., description="dimensions") + max_retries: int = Field(default=3, description="max retries") + raise_exception: bool = Field(default=True, description="raise exception") + + def _get_embeddings(self, input_text: str | List[str]): + """ + Get the embedding vector based on the input text. + This is an abstract method, and its concrete implementation must be provided in a subclass to generate the embedding vector for the given text. + Args: + input_text (str | List[str]): The input text, which can be a single string or a list of strings. + Raises: + NotImplementedError: If the method is not implemented in the subclass. + """ + raise NotImplementedError + + def get_embeddings(self, input_text: str | List[str]): + """ + Retrieves embeddings for the input text. + + This function attempts to obtain embeddings for the given input text. It will retry a maximum number of times in case of failure. + + Parameters: + - input_text (str | List[str]): The input text, which can be a single string or a list of strings. + + Returns: + - embeddings: The embeddings for the input text. Returns None if the maximum number of retries is reached and no successful result is obtained. + """ + # Attempt to get embeddings, with a maximum number of retries set + for i in range(self.max_retries): + try: + # Attempt to get embeddings, return immediately if successful + return self._get_embeddings(input_text) + + except Exception as e: + # Log exception information when an error occurs + logger.exception(f"embedding model name={self.model_name} encounter error with e={e.args}") + + # If the maximum number of retries is reached and raise_exception is set to True, re-throw the exception + if i == self.max_retries - 1 and self.raise_exception: + raise e + + return None + + def get_node_embeddings(self, nodes: VectorNode | List[VectorNode]): + if isinstance(nodes, VectorNode): + nodes.vector = self.get_embeddings(nodes.content) + return nodes + + elif isinstance(nodes, list): + max_batch_size = 10 # text-embedding-v4 batch size should not be larger than 10 + embeddings = [emb for i in range(0, len(nodes), max_batch_size) for emb in + self.get_embeddings(input_text=[node.content for node in nodes[i:i + max_batch_size]])] + if len(embeddings) != len(nodes): + logger.warning(f"embeddings.size={len(embeddings)} <> nodes.size={len(nodes)}") + else: + for node, embedding in zip(nodes, embeddings): + node.vector = embedding + return nodes + + else: + raise RuntimeError(f"unsupported type={type(nodes)}") diff --git a/v1/embedding_model/openai_compatible_embedding_model.py b/v1/embedding_model/openai_compatible_embedding_model.py new file mode 100644 index 00000000..f7897a12 --- /dev/null +++ b/v1/embedding_model/openai_compatible_embedding_model.py @@ -0,0 +1,92 @@ +import os +from typing import Literal, List + +from openai import OpenAI +from pydantic import Field, PrivateAttr, model_validator + +from v1.embedding_model import EMBEDDING_MODEL_REGISTRY +from v1.embedding_model.base_embedding_model import BaseEmbeddingModel + + +@EMBEDDING_MODEL_REGISTRY.register("openai_compatible") +class OpenAICompatibleEmbeddingModel(BaseEmbeddingModel): + api_key: str = Field(default_factory=lambda: os.getenv("OPENAI_API_KEY"), description="api key") + base_url: str = Field(default_factory=lambda: os.getenv("OPENAI_BASE_URL"), description="base url") + model_name: str = Field(default="text-embedding-v4", description="model name") + dimensions: int = Field(default=1024, description="dimensions") + encoding_format: Literal["float", "base64"] = Field(default="float", description="encoding_format") + _client: OpenAI = PrivateAttr() + + @model_validator(mode="after") + def init_client(self): + """ + Initialize the OpenAI client after model validation. + + This method is called after the model's data has been validated, + ensuring that all necessary attributes are correctly set before + initializing the OpenAI client. It creates an instance of the OpenAI + client using the provided API key and base URL, storing it in the + + Returns: + self: Returns the instance of the current class for method chaining. + """ + self._client = OpenAI(api_key=self.api_key, base_url=self.base_url) + return self + + def _get_embeddings(self, input_text: str | List[str]): + """ + Generate embeddings for the input text. + + This method accepts either a single string or a list of strings as input, + and returns the corresponding embeddings based on the specified model. + + Parameters: + - input_text (str | List[str]): The input text, which can be a single string or a list of strings. + + Returns: + - List[float]: If the input is a single string, returns a list of floating-point numbers representing the embedding. + - List[List[float]]: If the input is a list of strings, returns a list where each item is a list of floating-point numbers representing the embedding of each string. + - Raises RuntimeError: If the input type is unsupported. + """ + + # Create embeddings using the specified model, input text, dimensions, and encoding format + completion = self._client.embeddings.create( + model=self.model_name, + input=input_text, + dimensions=self.dimensions, + encoding_format=self.encoding_format + ) + + # Determine the type of input and process accordingly + if isinstance(input_text, str): + # If the input is a single string, return the embedding of that string + return completion.data[0].embedding + + elif isinstance(input_text, list): + # If the input is a list of strings, initialize a list to hold the embeddings of each string + result_emb = [[] for _ in range(len(input_text))] + # Iterate through the generated embeddings and assign them to the corresponding positions in the result list + for emb in completion.data: + result_emb[emb.index] = emb.embedding + return result_emb + + else: + # If the input type is neither a string nor a list of strings, throw an exception + raise RuntimeError(f"unsupported type={type(input_text)}") + + +def main(): + from experiencemaker.utils.util_function import load_env_keys + load_env_keys() + + model = OpenAICompatibleEmbeddingModel(dimensions=64, model_name="text-embedding-v4") + res1 = model.get_embeddings( + "The clothes are of good quality and look good, definitely worth the wait. I love them.") + res2 = model.get_embeddings(["aa", "bb"]) + print(res1) + print(res2) + + +if __name__ == "__main__": + main() + # launch with: python -m experiencemaker.model.openai_compatible_embedding_model diff --git a/v1/faiss_vector_store.py b/v1/faiss_vector_store.py deleted file mode 100644 index 8f5fe893..00000000 --- a/v1/faiss_vector_store.py +++ /dev/null @@ -1,70 +0,0 @@ -import faiss -import numpy as np -from typing import List -from pydantic import Field, model_validator, PrivateAttr -from experiencemaker.model.base_embedding_model import BaseEmbeddingModel -from experiencemaker.schema.vector_store_node import VectorStoreNode -from experiencemaker.storage.base_vector_store import BaseVectorStore, VECTOR_STORE_REGISTRY - -@VECTOR_STORE_REGISTRY.register("faiss") -class FaissVectorStore(BaseVectorStore): - dim: int = Field(default=768) - _index: faiss.IndexFlatIP = PrivateAttr() - _id_map: dict = PrivateAttr(default_factory=dict) - _vectors: list = PrivateAttr(default_factory=list) - - @model_validator(mode="after") - def init_faiss(self): - self._index = faiss.IndexFlatIP(self.dim) - self._id_map = {} - self._vectors = [] - return self - - def exist_index(self, index_name: str | None = None) -> bool: - return hasattr(self, '_index') - - def create_index(self, index_name: str | None = None): - self._index = faiss.IndexFlatIP(self.dim) - self._id_map = {} - self._vectors = [] - - def delete_index(self, index_name: str | None = None): - self._index = faiss.IndexFlatIP(self.dim) - self._id_map = {} - self._vectors = [] - - def exist_id(self, unique_id: str, index_name: str | None = None): - return unique_id in self._id_map - - def insert(self, nodes: VectorStoreNode | List[VectorStoreNode], index_name: str | None = None, **kwargs): - if isinstance(nodes, VectorStoreNode): - nodes = [nodes] - nodes = self.embedding_model.get_node_embeddings(nodes) - for node in nodes: - vec = np.array(node.vector, dtype=np.float32) - self._index.add(vec.reshape(1, -1)) - self._id_map[len(self._vectors)] = node.unique_id - self._vectors.append(node) - - def update(self, nodes: VectorStoreNode | List[VectorStoreNode], index_name: str | None = None, **kwargs): - self.delete_index() - all_nodes = self._vectors + (nodes if isinstance(nodes, list) else [nodes]) - self.insert(all_nodes) - - def delete_by_id(self, unique_id: str, index_name: str | None = None, **kwargs): - idx_to_remove = [i for i, node in enumerate(self._vectors) if node.unique_id == unique_id] - if idx_to_remove: - self._vectors = [node for node in self._vectors if node.unique_id != unique_id] - self.create_index() - self.insert(self._vectors) - - def retrieve_by_id(self, unique_id: str, index_name: str | None = None, **kwargs) -> VectorStoreNode | None: - for node in self._vectors: - if node.unique_id == unique_id: - return node - return None - - def retrieve_by_query(self, query: str, top_k: int = 1, index_name: str | None = None, **kwargs) -> List[VectorStoreNode]: - query_vec = np.array(self.embedding_model.get_embeddings(query), dtype=np.float32).reshape(1, -1) - D, I = self._index.search(query_vec, top_k) - return [self._vectors[i] for i in I[0] if i >= 0] \ No newline at end of file diff --git a/v1/llm/__init__.py b/v1/llm/__init__.py new file mode 100644 index 00000000..c6895799 --- /dev/null +++ b/v1/llm/__init__.py @@ -0,0 +1,3 @@ +from v1.utils.registry import Registry + +LLM_REGISTRY = Registry() \ No newline at end of file diff --git a/v1/llm/base_llm.py b/v1/llm/base_llm.py new file mode 100644 index 00000000..805223d1 --- /dev/null +++ b/v1/llm/base_llm.py @@ -0,0 +1,52 @@ +import time +from abc import ABC +from typing import List, Literal + +from loguru import logger, Message +from pydantic import Field, BaseModel + +from v1.tool.base_tool import BaseTool + + +class BaseLLM(BaseModel, ABC): + model_name: str = Field(...) + + seed: int = Field(default=42) + top_p: float | None = Field(default=None) + # stream: bool = Field(default=True) + stream_options: dict = Field(default={"include_usage": True}) + temperature: float = Field(default=0.0000001) + presence_penalty: float | None = Field(default=None) + enable_thinking: bool = Field(default=True, description="whether the current mode is the reasoning model, " + "or whether Qwen3's reasoning mode is currently enabled.") + tool_choice: Literal["none", "auto", "required"] = Field(default="auto", description="tool choice") + parallel_tool_calls: bool = Field(default=True) + + max_retries: int = Field(default=5, description="max retries") + raise_exception: bool = Field(default=True, description="raise exception") + + def stream_chat(self, messages: List[Message], tools: List[BaseTool] = None, **kwargs): + raise NotImplementedError + + def stream_print(self, messages: List[Message], tools: List[BaseTool] = None, **kwargs): + raise NotImplementedError + + def _chat(self, messages: List[Message], tools: List[BaseTool] = None, **kwargs) -> ActionMessage: + raise NotImplementedError + + def chat(self, messages: List[Message], tools: List[BaseTool] = None, **kwargs) -> ActionMessage | None: + for i in range(self.max_retries): + try: + # Attempt to execute the chat logic + return self._chat(messages, tools, **kwargs) + + except Exception as e: + # Log exceptions during the chat process + logger.exception(f"chat with model={self.model_name} encounter error with e={e.args}") + time.sleep(1 + i) + + # If the maximum number of retries is reached and raise_exception is set to True, then re-throw the exception + if i == self.max_retries - 1 and self.raise_exception: + raise e + + return None diff --git a/v1/pipeline/pipeline.py b/v1/pipeline/pipeline.py index f8147885..a2853229 100644 --- a/v1/pipeline/pipeline.py +++ b/v1/pipeline/pipeline.py @@ -10,7 +10,7 @@ from v1.pipeline.pipeline_context import PipelineContext from v1.utils.timer import Timer, timer -class Pipeline(object): +class Pipeline: seq_symbol: str = "->" parallel_symbol: str = "|" @@ -66,7 +66,8 @@ class Pipeline(object): logger.info(f"stage_{i}: {op}") elif isinstance(pipeline, list): - for op_list in zip_longest(*[self._parse_sub_pipeline(x) for x in pipeline], fillvalue="-"): + parallel_pipeline = [self._parse_sub_pipeline(x) for x in pipeline] + for op_list in zip_longest(*parallel_pipeline, fillvalue="-"): i += 1 logger.info(f"stage{i}: {' | '.join(op_list)}") @@ -74,7 +75,7 @@ class Pipeline(object): raise ValueError(f"unknown pipeline.type={type(pipeline)}") @timer() - def execute_pipeline(self, enable_print: bool = True): + def __call__(self, enable_print: bool = True): if enable_print: self.print_pipeline() diff --git a/v1/pipeline/pipeline_context.py b/v1/pipeline/pipeline_context.py index 9a705178..dffed290 100644 --- a/v1/pipeline/pipeline_context.py +++ b/v1/pipeline/pipeline_context.py @@ -16,3 +16,11 @@ class PipelineContext(object): def get_context(self, key: str): return self.context.get(key) + + @property + def request(self): + return self.get_context("request") + + @property + def response(self): + return self.get_context("response") diff --git a/v1/schema/app_config.py b/v1/schema/app_config.py index bcc05937..dfcfe876 100644 --- a/v1/schema/app_config.py +++ b/v1/schema/app_config.py @@ -8,6 +8,7 @@ class HttpServiceConfig: port: int = field(default=8001) timeout_keep_alive: int = field(default=600) limit_concurrency: int = field(default=64) + workers: int | None = field(default=None) @dataclass @@ -54,6 +55,7 @@ class VectorStoreConfig: class AppConfig: + config_path: str = field(default="") http_service: HttpServiceConfig = field(default_factory=HttpServiceConfig) thread_pool: ThreadPoolConfig = field(default_factory=ThreadPoolConfig) api: APIConfig = field(default_factory=APIConfig) diff --git a/v1/schema/experience.py b/v1/schema/experience.py index 4e8382cf..6796a697 100644 --- a/v1/schema/experience.py +++ b/v1/schema/experience.py @@ -5,7 +5,7 @@ from uuid import uuid4 from pydantic import BaseModel, Field -from v1.schema.vector_store_node import VectorStoreNode +from v1.schema.vector_node import VectorNode class ExperienceMeta(BaseModel): @@ -14,26 +14,31 @@ class ExperienceMeta(BaseModel): modified_time: str = Field(default_factory=lambda: datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")) extra_info: dict | None = Field(default=None) + def update_modified_time(self): + self.modified_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") + class BaseExperienceNode(BaseModel, ABC): + workspace_id: str = Field(default="") + experience_id: str = Field(default_factory=lambda: uuid4().hex) experience_type: str = Field(default="text") - workspace_id: str = Field(default="") - experience_meta: ExperienceMeta | None = Field(default=None) - def to_vector_store_node(self) -> VectorStoreNode: - ... - - @classmethod - def from_vector_store_node(cls, node: VectorStoreNode) -> "TextExperienceNode": - ... - - -class TextExperienceNode(BaseExperienceNode): - experience_type: str = Field(default="text") when_to_use: str = Field(default="") content: str | bytes = Field(default="") score: float | None = Field(default=None) + metadata: ExperienceMeta = Field(default_factory=ExperienceMeta) + + def to_vector_node(self) -> VectorNode: + raise NotImplementedError + + @classmethod + def from_vector_node(cls, node: VectorNode) -> "BaseExperienceNode": + raise NotImplementedError + + +class TextExperienceNode(BaseExperienceNode): + ... class FunctionArg(BaseModel): @@ -49,28 +54,16 @@ class Function(BaseModel): class FuncExperienceNode(BaseExperienceNode): - """ - TODO - """ - experience_type: str = Field(default="func") + experience_type: str = Field(default="function") functions: List[Function] = Field(default_factory=list) class PersonalExperienceNode(BaseExperienceNode): - """ - TODO: memory node from MemoryScope - """ experience_type: str = Field(default="personal") person: str = Field(default="") topic: str = Field(default="") - content: str | bytes = Field(default="") class KnowledgeExperienceNode(BaseExperienceNode): - """ - TODO - """ experience_type: str = Field(default="knowledge") topic: str = Field(default="") - content: str | bytes = Field(default="") - score: float | None = Field(default=None) diff --git a/v1/schema/request.py b/v1/schema/request.py index 939fc0e5..cba484e8 100644 --- a/v1/schema/request.py +++ b/v1/schema/request.py @@ -7,8 +7,8 @@ from v1.schema.message import Message, Trajectory class BaseRequest(BaseModel, ABC): - workspace_id: str = Field(default="") - params: dict = Field(default_factory=dict) + workspace_id: str = Field(default=...) + config: dict = Field(default_factory=dict) metadata: dict | None = Field(default=None) diff --git a/v1/schema/response.py b/v1/schema/response.py index c7c369e2..9ec455be 100644 --- a/v1/schema/response.py +++ b/v1/schema/response.py @@ -12,12 +12,12 @@ class BaseResponse(BaseModel): class RetrieverResponse(BaseResponse): - experience_nodes: list[BaseExperienceNode] = Field(default_factory=list) + experience_nodes: List[BaseExperienceNode] = Field(default_factory=list) experience_merged: str = Field(default="") class SummarizerResponse(BaseResponse): - experience_nodes: list[BaseExperienceNode] = Field(default_factory=list) + experience_nodes: List[BaseExperienceNode] = Field(default_factory=list) class VectorStoreResponse(BaseResponse): diff --git a/v1/schema/vector_store_node.py b/v1/schema/vector_node.py similarity index 90% rename from v1/schema/vector_store_node.py rename to v1/schema/vector_node.py index 49ce532d..2244b814 100644 --- a/v1/schema/vector_store_node.py +++ b/v1/schema/vector_node.py @@ -3,7 +3,7 @@ from uuid import uuid4 from pydantic import BaseModel, Field -class VectorStoreNode(BaseModel): +class VectorNode(BaseModel): unique_id: str = Field(default_factory=lambda: uuid4().hex) workspace_id: str = Field(default="") content: str = Field(default="") diff --git a/v1/sqlite_vector_store.py b/v1/sqlite_vector_store.py deleted file mode 100644 index f447b1f7..00000000 --- a/v1/sqlite_vector_store.py +++ /dev/null @@ -1,113 +0,0 @@ -import sqlite3 -import json -from typing import List -from pydantic import Field, model_validator, PrivateAttr -from experiencemaker.model.base_embedding_model import BaseEmbeddingModel -from experiencemaker.schema.vector_store_node import VectorStoreNode -from experiencemaker.storage.base_vector_store import BaseVectorStore, VECTOR_STORE_REGISTRY - -@VECTOR_STORE_REGISTRY.register("sqlite") -class SQLiteVectorStore(BaseVectorStore): - db_path: str = Field(default="./vector_store.db") - _conn: sqlite3.Connection = PrivateAttr() - - @model_validator(mode="after") - def init_db(self): - self._conn = sqlite3.connect(self.db_path, check_same_thread=False) - self._conn.execute(f""" - CREATE TABLE IF NOT EXISTS {self.index_name} ( - unique_id TEXT PRIMARY KEY, - workspace_id TEXT, - content TEXT, - metadata TEXT, - vector TEXT - ) - """) - self._conn.commit() - return self - - def exist_index(self, index_name: str | None = None) -> bool: - index = index_name or self.index_name - cursor = self._conn.execute("SELECT name FROM sqlite_master WHERE type='table' AND name=?", (index,)) - return cursor.fetchone() is not None - - def create_index(self, index_name: str | None = None): - index = index_name or self.index_name - self._conn.execute(f""" - CREATE TABLE IF NOT EXISTS {index} ( - unique_id TEXT PRIMARY KEY, - workspace_id TEXT, - content TEXT, - metadata TEXT, - vector TEXT - ) - """) - self._conn.commit() - - def delete_index(self, index_name: str | None = None): - index = index_name or self.index_name - self._conn.execute(f"DROP TABLE IF EXISTS {index}") - self._conn.commit() - - def exist_id(self, unique_id: str, index_name: str | None = None): - index = index_name or self.index_name - cursor = self._conn.execute(f"SELECT 1 FROM {index} WHERE unique_id=?", (unique_id,)) - return cursor.fetchone() is not None - - def insert(self, nodes: VectorStoreNode | List[VectorStoreNode], index_name: str | None = None, **kwargs): - index = index_name or self.index_name - if isinstance(nodes, VectorStoreNode): - nodes = [nodes] - nodes = self.embedding_model.get_node_embeddings(nodes) - for node in nodes: - self._conn.execute(f"REPLACE INTO {index} (unique_id, workspace_id, content, metadata, vector) VALUES (?, ?, ?, ?, ?)", - (node.unique_id, node.workspace_id, node.content, json.dumps(node.metadata), json.dumps(node.vector))) - self._conn.commit() - - def update(self, nodes: VectorStoreNode | List[VectorStoreNode], index_name: str | None = None, **kwargs): - self.insert(nodes, index_name=index_name) - - def delete_by_id(self, unique_id: str, index_name: str | None = None, **kwargs): - index = index_name or self.index_name - self._conn.execute(f"DELETE FROM {index} WHERE unique_id=?", (unique_id,)) - self._conn.commit() - - def retrieve_by_id(self, unique_id: str, index_name: str | None = None, **kwargs) -> VectorStoreNode | None: - index = index_name or self.index_name - cursor = self._conn.execute(f"SELECT unique_id, workspace_id, content, metadata, vector FROM {index} WHERE unique_id=?", (unique_id,)) - row = cursor.fetchone() - if row: - return VectorStoreNode( - unique_id=row[0], - workspace_id=row[1], - content=row[2], - metadata=json.loads(row[3]), - vector=json.loads(row[4]) - ) - return None - - def retrieve_by_query(self, query: str, top_k: int = 1, index_name: str | None = None, **kwargs) -> List[VectorStoreNode]: - index = index_name or self.index_name - query_vec = self.embedding_model.get_embeddings(query) - cursor = self._conn.execute(f"SELECT unique_id, workspace_id, content, metadata, vector FROM {index}") - results = [] - for row in cursor: - node = VectorStoreNode( - unique_id=row[0], - workspace_id=row[1], - content=row[2], - metadata=json.loads(row[3]), - vector=json.loads(row[4]) - ) - node.metadata["score"] = self._cosine_similarity(query_vec, node.vector) - results.append(node) - results.sort(key=lambda x: x.metadata["score"], reverse=True) - return results[:top_k] - - @staticmethod - def _cosine_similarity(vec1, vec2): - import math - dot = sum(x * y for x, y in zip(vec1, vec2)) - norm1 = math.sqrt(sum(x * x for x in vec1)) - norm2 = math.sqrt(sum(y * y for y in vec2)) - return dot / (norm1 * norm2) if norm1 and norm2 else 0.0 \ No newline at end of file diff --git a/v1/tool/__init__.py b/v1/tool/__init__.py new file mode 100644 index 00000000..7ecc150c --- /dev/null +++ b/v1/tool/__init__.py @@ -0,0 +1,3 @@ +from v1.utils.registry import Registry + +TOOL_REGISTRY = Registry() \ No newline at end of file diff --git a/v1/tool/base_tool.py b/v1/tool/base_tool.py new file mode 100644 index 00000000..1f5edb85 --- /dev/null +++ b/v1/tool/base_tool.py @@ -0,0 +1,79 @@ +from abc import ABC + +from loguru import logger +from pydantic import BaseModel, Field + + +class BaseTool(BaseModel, ABC): + tool_id: str = Field(default="") + name: str = Field(..., description="tool name") + description: str = Field(..., description="tool description") + tool_type: str = Field(default="function") + parameters: dict = Field(default_factory=dict, description="tool parameters") + arguments: dict = Field(default_factory=dict, description="execute arguments") + + enable_cache: bool = Field(default=False, description="whether to cache the tool result") + cached_result: dict = Field(default_factory=dict, description="tool execution result") + + max_retries: int = Field(default=3, description="max retries") + raise_exception: bool = Field(default=True, description="raise exception") + success: bool = Field(default=True, description="whether the tool executed successfully") + + def reset(self): + self.arguments.clear() + self.success = True + + def _execute(self, **kwargs): + raise NotImplementedError + + def execute(self, **kwargs): + cache_id = "" + if self.enable_cache: + cache_id = self.get_cache_id(**kwargs) + if cache_id in self.cached_result: + return self.cached_result[cache_id] + + for i in range(self.max_retries): + try: + if self.enable_cache: + self.cached_result[cache_id] = self._execute(**kwargs) + return self.cached_result[cache_id] + + else: + return self._execute(**kwargs) + + except Exception as e: + logger.exception(f"using tool.name={self.name} encounter error with e={e.args}") + if i == self.max_retries - 1 and self.raise_exception: + raise e + + return None + + + def simple_dump(self) -> dict: + """ + It may be in other different tool params formats; different versions are completed here. + """ + return { + "type": self.tool_type, + self.tool_type: { + "name": self.name, + "description": self.description, + "parameters": self.parameters, + }, + } + + @property + def input_schema(self) -> dict: + return self.parameters.get("properties", {}) + + @property + def output_schema(self) -> dict: + raise NotImplementedError + + def refresh(self): + # for mcp + raise NotImplementedError + + def get_cache_id(self, **kwargs) -> str: + raise NotImplementedError diff --git a/v1/tt.py b/v1/tt.py deleted file mode 100644 index 8485b126..00000000 --- a/v1/tt.py +++ /dev/null @@ -1,15 +0,0 @@ -class Pipe(): - - def __init__(self): - pass - - - -if __name__ == "__main__": - line = "a->[c->d->d|e->f]->[f->g]->h" - ops = [] - for sub_line in line.split("["): - for sub_line2 in sub_line.split("]"): - ops.append(sub_line2) - - print(ops) \ No newline at end of file diff --git a/v1/vector_store/__init__.py b/v1/vector_store/__init__.py new file mode 100644 index 00000000..4243e5c6 --- /dev/null +++ b/v1/vector_store/__init__.py @@ -0,0 +1,3 @@ +from v1.utils.registry import Registry + +VECTOR_STORE_REGISTRY = Registry() \ No newline at end of file diff --git a/v1/vector_store/base_vector_store.py b/v1/vector_store/base_vector_store.py new file mode 100644 index 00000000..1cf1e2fc --- /dev/null +++ b/v1/vector_store/base_vector_store.py @@ -0,0 +1,44 @@ +from abc import ABC +from typing import List + +from pydantic import BaseModel, Field + +from v1.model.base_embedding_model import BaseEmbeddingModel +from v1.schema.vector_node import VectorNode + + +class BaseVectorStore(BaseModel, ABC): + embedding_model: BaseEmbeddingModel = Field(default=...) + + def exist_workspace(self, workspace_id: str, **kwargs) -> bool: + raise NotImplementedError + + def delete_workspace(self, workspace_id: str, **kwargs): + raise NotImplementedError + + def create_workspace(self, workspace_id: str, **kwargs): + raise NotImplementedError + + def dump_workspace(self, workspace_id: str, **kwargs): + raise NotImplementedError + + def load_workspace(self, workspace_id: str, **kwargs): + raise NotImplementedError + + 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 + + def exist_id(self, unique_id: str, workspace_id: str = None, **kwargs) -> bool: + raise NotImplementedError + + def delete_id(self, unique_id: str, workspace_id: str = None, **kwargs): + raise NotImplementedError diff --git a/v1/vector_store/es_vector_store.py b/v1/vector_store/es_vector_store.py new file mode 100644 index 00000000..7c259532 --- /dev/null +++ b/v1/vector_store/es_vector_store.py @@ -0,0 +1,290 @@ +import os +from typing import List, Tuple + +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 + + +@VECTOR_STORE_REGISTRY.register("elasticsearch") +class EsVectorStore(BaseVectorStore): + hosts: str | List[str] = Field(default_factory=lambda: os.getenv("ES_HOSTS", "http://localhost:9200")) + basic_auth: str | Tuple[str, str] | None = Field(default=None) + bulk_chunk_size: int = Field(default=512) + retrieve_filters: List[dict] = [] + _client: Elasticsearch = PrivateAttr() + + @model_validator(mode="after") + def init_client(self): + if isinstance(self.hosts, str): + hosts = [self.hosts] + else: + hosts = self.hosts + 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 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 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 = { + "mappings": { + "properties": { + "workspace_id": {"type": "keyword"}, + "content": {"type": "text"}, + "metadata": {"type": "object"}, + "vector": { + "type": "dense_vector", + "dims": self.embedding_model.dimensions + } + } + } + } + + return self._client.indices.create(index=index_name, body=index) + + 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) + + @staticmethod + def doc2node(doc) -> VectorStoreNode: + node = VectorStoreNode(**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 + + doc: dict = { + "_index": index_name, + "_id": node.unique_id, + "_source": { + "workspace_id": node.workspace_id, + "content": node.content, + "metadata": node.metadata, + "vector": node.vector + } + } + + if add_op_type: + doc["_op_type"] = "update" if self.exist_id(node.unique_id, index_name) else "index", + return doc + + 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 + + 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!") + return [] + + query_vector = self.embedding_model.get_embeddings(query) + + body = { + "query": { + "script_score": { + "query": {"bool": {"must": self.retrieve_filters}}, + "script": { + "source": "cosineSimilarity(params.query_vector, 'vector') + 1.0", + "params": {"query_vector": query_vector}, + } + } + }, + "size": top_k + } + response = self._client.search(index=index_name, body=body, **kwargs) + + nodes: List[VectorStoreNode] = [] + for doc in response['hits']['hits']: + nodes.append(self.doc2node(doc)) + + self.retrieve_filters.clear() + return nodes + + +def main(): + from experiencemaker.utils.util_function import load_env_keys + load_env_keys() + + embedding_model = OpenAICompatibleEmbeddingModel(dimensions=1024) + index_name = "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() + + sample_nodes = [ + VectorStoreNode( + workspace_id="w1", + content="Artificial intelligence is a technology that simulates human intelligence.", + metadata={ + "node_type": "n1", + } + ), + VectorStoreNode( + workspace_id="w1", + content="AI is the future of mankind.", + metadata={ + "node_type": "n1", + } + ), + VectorStoreNode( + workspace_id="w1", + content="I want to eat fish!", + metadata={ + "node_type": "n2", + } + ), + VectorStoreNode( + workspace_id="w2", + content="The bigger the storm, the more expensive the fish.", + metadata={ + "node_type": "n1", + } + ), + ] + + es.insert(sample_nodes, refresh_index=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) + 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) + 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) + + +if __name__ == "__main__": + main() + # launch with: python -m experiencemaker.storage.es_vector_store diff --git a/v1/vector_store/file_vector_store.py b/v1/vector_store/file_vector_store.py new file mode 100644 index 00000000..900ca3eb --- /dev/null +++ b/v1/vector_store/file_vector_store.py @@ -0,0 +1,209 @@ +import json +import math +import threading +from pathlib import Path +from typing import List, Any + +from loguru import logger +from pydantic import Field, model_validator, PrivateAttr + +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 + + +@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 + + 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 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 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 _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 + + @staticmethod + def calculate_similarity(query_vector: List[float], node_vector: List[float]): + assert query_vector, f"query_vector is empty!" + assert node_vector, f"node_vector is empty!" + assert len(query_vector) == len(node_vector), \ + f"query_vector.size={len(query_vector)} node_vector.size={len(node_vector)}" + + dot_product = sum(x * y for x, y in zip(query_vector, node_vector)) + norm_v1 = math.sqrt(sum(x ** 2 for x in query_vector)) + 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]: + query_vector = self.embedding_model.get_embeddings(query) + nodes: List[VectorStoreNode] = self._load(index_name=index_name) + for node in nodes: + node.metadata["score"] = self.calculate_similarity(query_vector, node.vector) + + nodes = sorted(nodes, key=lambda x: x.metadata["score"], reverse=True) + return nodes[:top_k] + + +def main(): + from experiencemaker.utils.util_function import load_env_keys + load_env_keys() + + 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() + + sample_nodes = [ + VectorStoreNode( + workspace_id="w1", + content="Artificial intelligence is a technology that simulates human intelligence.", + metadata={ + "node_type": "n1", + } + ), + VectorStoreNode( + workspace_id="w1", + content="AI is the future of mankind.", + metadata={ + "node_type": "n1", + } + ), + VectorStoreNode( + workspace_id="w1", + content="I want to eat fish!", + metadata={ + "node_type": "n2", + } + ), + VectorStoreNode( + workspace_id="w2", + content="The bigger the storm, the more expensive the fish.", + metadata={ + "node_type": "n1", + } + ), + ] + + client.insert(sample_nodes) + + logger.info("=" * 20) + results = client.retrieve_by_query("What is AI?", top_k=5) + for r in results: + logger.info(r.model_dump(exclude={"vector"})) + logger.info("=" * 20) + + +if __name__ == "__main__": + main() + # launch with: python -m experiencemaker.storage.file_vector_store