From a0cc23de5797f7d7baced531396e11c2d508ceaf Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 11 Jun 2025 23:13:46 +0800 Subject: [PATCH] add your own agent --- README.md | 8 + cookbook/__init__.py | 0 cookbook/simple_agent/quick_start.md | 162 ++++++++++++++++++ cookbook/simple_agent/your_own_agent.py | 106 ++++++++++++ .../simple_agent/your_own_agent_enhanced.py | 53 ++++++ .../simple_agent/your_own_agent_prompt.yaml | 23 +++ cookbook/vector_store/elasticsearch.md | 43 +++++ experiencemaker/em_service.py | 86 +++++++--- .../agent_wrapper/agent_wrapper_mixin.py | 3 +- .../module/agent_wrapper/simple_agent.py | 13 +- .../agent_wrapper/simple_agent_prompt.yaml | 7 +- .../agent_wrapper/simple_agent_wrapper.py | 8 +- .../base_context_generator.py | 13 +- .../simple_context_generator.py | 12 +- .../module/summarizer/base_summarizer.py | 9 +- experiencemaker/schema/experience.py | 4 +- experiencemaker/schema/request.py | 1 + experiencemaker/storage/base_vector_store.py | 23 ++- experiencemaker/storage/es_vector_store.py | 100 +++++++---- experiencemaker/storage/file_vector_store.py | 94 ++++++---- experiencemaker/utils/util_function.py | 8 +- 21 files changed, 640 insertions(+), 136 deletions(-) create mode 100644 cookbook/__init__.py create mode 100644 cookbook/simple_agent/quick_start.md create mode 100644 cookbook/simple_agent/your_own_agent.py create mode 100644 cookbook/simple_agent/your_own_agent_enhanced.py create mode 100644 cookbook/simple_agent/your_own_agent_prompt.yaml create mode 100644 cookbook/vector_store/elasticsearch.md diff --git a/README.md b/README.md index 8dde8646..ef9fd418 100644 --- a/README.md +++ b/README.md @@ -41,3 +41,11 @@ docker run -p 9200:9200 \ ``` # run module service + +[] 注册函数 + +[] 更新workspace_id + index_name +[] trajectory 加 reward +[] summary去重，context加LLM eval +[] + diff --git a/cookbook/__init__.py b/cookbook/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/cookbook/simple_agent/quick_start.md b/cookbook/simple_agent/quick_start.md new file mode 100644 index 00000000..8db8233f --- /dev/null +++ b/cookbook/simple_agent/quick_start.md @@ -0,0 +1,162 @@ +# Quick Start + +## Hello Experience Maker +Here is a simple user guide for ExperienceMaker. + +### Step0: Preparation Work + +#### Prepare LLM & EMBEDDING_MODEL +We need to prepare the API services for the LLM and the Embedding model. +Since we are using an OpenAI-compatible service, we only need to write the `OPENAI_API_KEY` and `OPENAI_BASE_URL` into the environment. +```shell +export OPENAI_API_KEY="sk-xxx" +export OPENAI_BASE_URL="xxx" +``` + +#### Prepare Vector Store +If you want to use vector store, you need to set up a vector database. Don't forget to set up the `ES_HOSTS`. +- Elasticsearch [quick start](../vector_store/elasticsearch.md) + +#### Prepare Your Own Agent +Assume you have a runnable agent. +Here, we use a basic LLM combined with a simple react framework including three tools(code, web_search, terminate) as an example. +```python +class YourOwnAgent(...): + ... + def think(self, **kwargs) -> bool: + ... + + def act(self, **kwargs): + ... + + def run(self, query: str, previous_experience: str): + ... +``` + +Here is an [example code](./your_own_agent.py) with [prompt](./your_own_agent_prompt.yaml). To use this simple agent, you will need to set up `DASHSCOPE_API_KEY`. +```shell +export DASHSCOPE_API_KEY="sk-xxx" +``` + +### Step1: Start Experience Maker Http Service +- Install dependencies. +```shell +pip install . +``` + +- We start our context and summary services in `simple` mode. Next, you can use standard HTTP interfaces to call the services. +```shell +pip install . +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", "index_name": "your_own_agent"}' \ + --context_generator='{"backend": "simple", "retrieve_top_k": 1}' \ + --summarizer='{"backend": "simple"}' +``` + + + + +### Step2: Implement AgentWrapper + +In order to utilize the **context generator** and **summarizer** capabilities of experiencemaker, please inherit from **MxcAgent** and **BaseAgentWrapperMixin** to implement the AgentWrapper. + +Here, you need to customize two parts: +- how to integrate the content message(insight) generated by the `self.context_generator` into the context. +- implement the execute function to output the trajectory. + +Below is a simple example of integrating **trajectory-level insight** into the context. + +```python +from experiencemaker.core.module.agent_wrapper.base_agent_wrapper import BaseAgentWrapperMixin + +class MxcAgentWrapper(MxcAgent, BaseAgentWrapperMixin): + def execute(self, query: str, **kwargs) -> Trajectory: + trajectory = Trajectory(steps=messages, query=query) + context_msg = self.context_generator.execute(trajectory=trajectory) + new_query = f""" +previous insight: +{context_msg.content} +Please consider the helpful parts from these in answering the question, to make the response more comprehensive and substantial. + +user query: +{query} + """.strip() + + messages = self.run(new_query, **kwargs) + return Trajectory(query=query, steps=messages, answer=messages[-1].content, done=True) + +``` + +### Step3: Run AgentRunner with insight + +Once you have completed the implementation of the AgentWrapper class, you will be able to utilize the capabilities of +experiencemaker. + +Here is an example using **SimpleAgentRunner**. +We first executed two historical tasks, then summarized the experience and made it persistent. +Finally, we utilized the historical experience in a new task. + +[insights demo](./insight.json) + + +```python +from experiencemaker.core.module.runner.simple_agent_runner import SimpleAgentRunner + + + +mxc_agent_wrapper = MxcAgentWrapper(llm=OpenAICompatibleBaseLLM(model_name="qwen3-32b", temperature=0.0001), + max_steps=10, + tools=[CodeTool(), DashscopeSearchTool(), TerminateTool()]) +agent_runner = SimpleAgentRunner(agent_wrapper=mxc_agent_wrapper, summarizer="default", context_generator="default") + +# historical tasks +agent_runner.rollout_trajectory(query="Analyze the company Tesla.") +agent_runner.rollout_trajectory(query="Analyze the company Apple.") + +# summary insights and store them +agent_runner.summary_and_store() + +# run agent with historical insights +trajectory = agent_runner.rollout_trajectory(query="Analyze the company Xiaomi Corporation.") +``` + + +### Step4: Evaluation(Optional) + +If we have a reward function that allows us to compare the performance before and after adding context, we can try this +part. + +Use `run_agent` to obtain the answer from the original agent (answer1), and use `run_agent_wrapper` to get the answer +with added insights and experience (answer2). + +Here, the reward function is used to compare and score the two answers. The `reward.reward_value` indicates the win rate +of answer2. + +```python +# task +query = "Analyze Xiaomi Corporation." + +# run agent +agent = MxcAgent(llm=OpenAICompatibleBaseLLM(model_name="qwen3-32b", temperature=0.0001), + max_steps=10, + tools=[CodeTool(), DashscopeSearchTool(), TerminateTool()]) +messages = agent.run(query=query) +answer1 = messages[-1].content + +# agent runner: Assume we already have some historical experience. +mxc_agent_wrapper = MxcAgentWrapper(llm=OpenAICompatibleBaseLLM(model_name="qwen3-32b", temperature=0.0001), + max_steps=10, + tools=[CodeTool(), DashscopeSearchTool(), TerminateTool()]) +agent_runner = SimpleAgentRunner(agent_wrapper=mxc_agent_wrapper, summarizer="default", context_generator="default") +trajectory = agent_runner.rollout_trajectory(query=query) +answer2 = trajectory.answer + +# pair-wise LLM evaluation +from experiencemaker.core.module.reward_fn.simple_reward_fn import SimpleRewardFn +reward_fn = SimpleRewardFn(llm=OpenAICompatibleBaseLLM(model_name="qwen3-32b", temperature=0.0001)) +reward = reward_fn.execute(query=query, answer1=answer1, answer2=answer2, eval_times=5) +print(f"final reward={reward.reward_value}") +``` \ No newline at end of file diff --git a/cookbook/simple_agent/your_own_agent.py b/cookbook/simple_agent/your_own_agent.py new file mode 100644 index 00000000..b812b75c --- /dev/null +++ b/cookbook/simple_agent/your_own_agent.py @@ -0,0 +1,106 @@ +import datetime +from pathlib import Path +from typing import List + +from loguru import logger +from pydantic import Field, BaseModel + +from experiencemaker.utils.util_function import load_env_keys +load_env_keys("../../.env") + +from experiencemaker.model import OpenAICompatibleBaseLLM +from experiencemaker.model.base_llm import BaseLLM +from experiencemaker.module.prompt.prompt_mixin import PromptMixin +from experiencemaker.schema.trajectory import Message, ActionMessage, ToolCall, StateMessage +from experiencemaker.tool import CodeTool, DashscopeSearchTool, TerminateTool +from experiencemaker.tool.base_tool import BaseTool + + +class AgentContext(BaseModel): + current_step: int = Field(default=-1) + query: str = Field(default="") + messages: List[Message] = Field(default_factory=list) + metadata: dict = Field(default_factory=dict) + has_terminate_tool: bool = Field(default=False) + + +class YourOwnAgent(PromptMixin): + llm: BaseLLM | None = Field(default=None) + max_steps: int = Field(default=10) + tools: List[BaseTool] = [CodeTool(), DashscopeSearchTool(), TerminateTool()] + prompt_file_path: Path = Path(__file__).parent / "Your_own_agent_prompt.yaml" + + def think(self, context: AgentContext): + now_time = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S') + tool_names = [x.name for x in self.tools] + + if context.current_step == 0: + user_prompt = self.prompt_format(prompt_name="role_prompt", + time=now_time, + tools=", ".join(tool_names), + query=context.query) + + elif context.has_terminate_tool: + user_prompt = self.prompt_format(prompt_name="final_prompt", query=context.query) + + else: + user_prompt = self.prompt_format(prompt_name="next_prompt", query=context.query) + + context.messages.append(Message(content=user_prompt)) + logger.info(f"step.{context.current_step} user_prompt={user_prompt}") + + if context.has_terminate_tool: + action_msg: ActionMessage = self.llm.chat(context.messages) + + else: + action_msg: ActionMessage = self.llm.chat(context.messages, tools=self.tools) + for tool in action_msg.tool_calls: + if tool.name == "terminate": + context.has_terminate_tool = True + break + + context.messages.append(action_msg) + action_msg_context: str = action_msg.content + "\n\n" + action_msg.reasoning_content + logger.info(f"step.{context.current_step} action_msg_context={action_msg_context} " + f"tool_calls={action_msg.tool_calls}") + return True if action_msg.tool_calls else False + + def act(self, context: AgentContext): + action_msg = context.messages[-1] + assert isinstance(action_msg, ActionMessage) + + tool_dict = {tool.name: tool for tool in self.tools} + + new_tool_calls: List[ToolCall] = [] + for tool_call in action_msg.tool_calls: + if tool_call.name not in tool_dict: + continue + + new_tool_call = tool_call.model_copy(deep=True) + tool = tool_dict[tool_call.name] + new_tool_call.result = tool.execute(**tool_call.argument_dict) + new_tool_calls.append(new_tool_call) + + state_msg = StateMessage(tool_calls=new_tool_calls) + context.messages.append(state_msg) + logger.info(f"step.{context.current_step} state_msg_context={state_msg.content}") + + def run(self, query: str) -> List[Message]: + context: AgentContext = AgentContext(query=query) + + for i in range(self.max_steps): + context.current_step = i + + should_act: bool = self.think(context) + if should_act: + self.act(context) + else: + break + return context.messages + + +if __name__ == "__main__": + agent = YourOwnAgent(llm=OpenAICompatibleBaseLLM(model_name="qwen3-32b", temperature=0.6)) + messages = agent.run(query="Analyze Xiaomi Corporation.") + answer = messages[-1].content + logger.info(answer) diff --git a/cookbook/simple_agent/your_own_agent_enhanced.py b/cookbook/simple_agent/your_own_agent_enhanced.py new file mode 100644 index 00000000..f252c30f --- /dev/null +++ b/cookbook/simple_agent/your_own_agent_enhanced.py @@ -0,0 +1,53 @@ +from loguru import logger +from pydantic import Field, model_validator + +from cookbook.simple_agent.your_own_agent import YourOwnAgent +from experiencemaker.em_client import EMClient +from experiencemaker.model import OpenAICompatibleBaseLLM +from experiencemaker.schema.request import ContextGeneratorRequest, SummarizerRequest +from experiencemaker.schema.response import ContextGeneratorResponse, SummarizerResponse +from experiencemaker.schema.trajectory import Trajectory + + +class YourOwnAgentEnhanced(YourOwnAgent): + em_client: EMClient | None = Field(default=None) + + @model_validator(mode="after") + def init_client(self): + self.em_client = EMClient(base_url="http://0.0.0.0:8001") + return self + + def summary_experience(self, query: str): + messages = self.run(query) + trajectory: Trajectory = Trajectory(query=query, steps=messages, answer=messages[-1].content, done=True) + + request: SummarizerRequest = SummarizerRequest(trajectories=[trajectory]) + response: SummarizerResponse = self.em_client.call_summarizer(request) + for experience in response.experiences: + logger.info(experience.model_dump_json()) + + def run_with_experience(self, query: str): + trajectory: Trajectory = Trajectory(query=query) + request: ContextGeneratorRequest = ContextGeneratorRequest(trajectory=trajectory) + response: ContextGeneratorResponse = self.em_client.call_context_generator(request) + new_query = f"{response.context_msg.content}\n\nUser Question\n{query}" + logger.info(f"new query={new_query}") + messages = self.run(new_query) + + trajectory.steps = messages + trajectory.answer = messages[-1].content + trajectory.done = True + trajectory.metadata["experience"] = response.context_msg.content + return trajectory + + def execute(self): + self.summary_experience(query="Analyze the company Tesla.") + self.summary_experience(query="Analyze the company Apple.") + + return self.run_with_experience(query="Analyze Xiaomi Corporation.") + + +if __name__ == "__main__": + agent = YourOwnAgentEnhanced(llm=OpenAICompatibleBaseLLM(model_name="qwen3-32b", temperature=0.6)) + traj = agent.execute() + logger.info(traj.model_dump_json(indent=2)) diff --git a/cookbook/simple_agent/your_own_agent_prompt.yaml b/cookbook/simple_agent/your_own_agent_prompt.yaml new file mode 100644 index 00000000..eaa1cc81 --- /dev/null +++ b/cookbook/simple_agent/your_own_agent_prompt.yaml @@ -0,0 +1,23 @@ +role_prompt: | + You are a helpful assistant named BeyondAgent. + The current time is {time}. + + Please proactively choose the most suitable tool or combination of tools based on the user's question, including {tools} etc. + For complex tasks, you can break down the problem step by step and use different tools to solve it incrementally. + Please determine the response language based on the language of the user's question. + + {query} + +next_prompt: | + User's question + {query} + + Plan the most suitable tool or combination of tools based on the context and the user's question. + For complex tasks, break down the problem and use different tools step by step to solve it, don't give up easily. + Try calling the same tool multiple times with different parameters to obtain information from various perspectives. + If the task is completed and the user's question can now be answered, use the **terminate** tool. + +final_prompt: | + Please integrate the context and provide a complete answer to the user's question: + {query} + diff --git a/cookbook/vector_store/elasticsearch.md b/cookbook/vector_store/elasticsearch.md new file mode 100644 index 00000000..6b2b71ba --- /dev/null +++ b/cookbook/vector_store/elasticsearch.md @@ -0,0 +1,43 @@ +## Elasticsearch Vector Store +If a vector database is involved, you will need an Elasticsearch environment. You can refer to the following steps. + +### Install Docker Desktop +If you don’t have Docker installed, download and install [Docker Desktop](https://www.docker.com/products/docker-desktop) for your operating system. + +### Set up Elasticsearch +You can choose one of the following three options. + +#### All in One Script +To set up [Elasticsearch](https://www.elastic.co/docs/solutions/search/run-elasticsearch-locally) and Kibana locally, run the start-local script in the command line: +```shell +curl -fsSL https://elastic.co/start-local | sh +``` + +#### Docker Run Image with 4GB Memory +manually download and load the image. Here, we take `elasticsearch-wolfi:9.0.0` as an example: +```shell +docker pull docker.elastic.co/elasticsearch/elasticsearch-wolfi:9.0.0 +docker run -p 9200:9200 \ + --memory='4GB' \ + -e "discovery.type=single-node" \ + -e "xpack.security.enabled=false" \ + -e "xpack.license.self_generated.type=trial" \ + docker.elastic.co/elasticsearch/elasticsearch-wolfi:9.0.0 +``` + +#### Docker Run Image with Http Host +```shell +docker pull docker.elastic.co/elasticsearch/elasticsearch-wolfi:9.0.0 +docker run -p 8200:9200 \ + -e "discovery.type=single-node" \ + -e "xpack.security.enabled=false" \ + -e "xpack.license.self_generated.type=trial" \ + -e "http.host=0.0.0.0" \ + docker.elastic.co/elasticsearch/elasticsearch-wolfi:9.0.0 +``` + +### Inject Environment Variables +Inject variables of the `ES_HOSTS` into the environment where you use Elasticsearch. +```shell +export ES_HOSTS=http://localhost:9200 +``` \ No newline at end of file diff --git a/experiencemaker/em_service.py b/experiencemaker/em_service.py index d03ea18a..66f7b5b3 100644 --- a/experiencemaker/em_service.py +++ b/experiencemaker/em_service.py @@ -1,4 +1,5 @@ import argparse +import copy import json import types from typing import List @@ -30,7 +31,6 @@ from experiencemaker.storage.base_vector_store import BaseVectorStore class EMService(BaseModel): - workspace_id: str = Field(default="") host: str = Field(default="0.0.0.0") port: int = Field(default=8001) timeout_keep_alive: int = Field(default=600000) @@ -43,6 +43,8 @@ class EMService(BaseModel): context_generator: BaseContextGenerator | None = Field(default=None) summarizer: BaseSummarizer | None = Field(default=None) + origin_config: dict = Field(default_factory=dict) + @staticmethod def init_llm(llm_config: dict) -> BaseLLM: backend = llm_config.pop("backend", None) @@ -115,7 +117,7 @@ class EMService(BaseModel): embedding_model=data.get("embedding_model")) context_generator: BaseContextGenerator = CONTEXT_GENERATOR_REGISTRY[backend]( - **context_generator_config, llm=llm, vector_store=vector_store, workspace_id=data.get("workspace_id", "")) + **context_generator_config, llm=llm, vector_store=vector_store) logger.info(f"context_generator is inited with backend={backend} params={context_generator_config}") return context_generator @@ -130,7 +132,7 @@ class EMService(BaseModel): vector_store = cls.get_vector_store(summarizer_config, vector_store=data.get("vector_store"), embedding_model=data.get("embedding_model")) summarizer: BaseSummarizer = SUMMARIZER_REGISTRY[backend]( - **summarizer_config, llm=llm, vector_store=vector_store, workspace_id=data.get("workspace_id", "")) + **summarizer_config, llm=llm, vector_store=vector_store) logger.info(f"summarizer is inited with backend={backend} params={summarizer_config}") return summarizer @@ -143,14 +145,12 @@ class EMService(BaseModel): llm = cls.get_llm(agent_wrapper_config, llm=data.get("llm")) agent_wrapper: AgentWrapperMixin = AGENT_WRAPPER_REGISTRY[backend]( - **agent_wrapper_config, llm=llm, context_generator=data.get("context_generator"), - workspace_id=data.get("workspace_id", "")) + **agent_wrapper_config, llm=llm, context_generator=data.get("context_generator")) logger.info(f"agent_wrapper is inited with backend={backend} params={agent_wrapper_config}") return agent_wrapper - @model_validator(mode="before") # noqa @classmethod - def init_modules(cls, data: dict): + def init_class_by_config(cls, data: dict): try: if "llm" in data: data["llm"] = cls.init_llm(data["llm"]) @@ -170,24 +170,62 @@ class EMService(BaseModel): if "agent_wrapper" in data: data["agent_wrapper"] = cls.init_agent_wrapper(data["agent_wrapper"], data) + except Exception as e: logger.exception(e.args) return data + @model_validator(mode="before") # noqa + @classmethod + def init_modules(cls, data: dict): + origin_config = copy.deepcopy(data) + data = cls.init_class_by_config(data) + data["origin_config"] = origin_config + return data + def call_agent_wrapper(self, request: AgentWrapperRequest) -> AgentWrapperResponse: - assert self.agent_wrapper_ is not None, "agent_wrapper must be provided." - trajectory: Trajectory = self.agent_wrapper_.execute(request.query, **request.metadata) + if "em_config" in request.metadata: + new_config = copy.deepcopy(self.origin_config) + new_config.update(request.metadata["em_config"]) + data = EMService.init_class_by_config(new_config) + agent_wrapper = data["agent_wrapper"] + else: + assert self.agent_wrapper is not None, "agent_wrapper must be provided." + agent_wrapper = self.agent_wrapper + + trajectory: Trajectory = agent_wrapper.execute(query=request.query, + workspace_id=request.workspace_id, + **request.metadata) return AgentWrapperResponse(trajectory=trajectory) def call_context_generator(self, request: ContextGeneratorRequest) -> ContextGeneratorResponse: - assert self.context_generator_ is not None, "context_generator must be provided." - context_msg: ContextMessage = self.context_generator_.execute(request.trajectory, **request.metadata) + if "em_config" in request.metadata: + new_config = copy.deepcopy(self.origin_config) + new_config.update(request.metadata["em_config"]) + data = EMService.init_class_by_config(new_config) + context_generator = data["context_generator"] + else: + assert self.context_generator is not None, "context_generator must be provided." + context_generator = self.context_generator + + context_msg: ContextMessage = context_generator.execute(trajectory=request.trajectory, + workspace_id=request.workspace_id, + **request.metadata) return ContextGeneratorResponse(context_msg=context_msg) def call_summarizer(self, request: SummarizerRequest) -> SummarizerResponse: - assert self.summarizer_ is not None, "summarizer must be provided." - experiences: List[Experience] = self.summarizer_.execute(request.trajectories, request.return_experience, - **request.metadata) + if "em_config" in request.metadata: + new_config = copy.deepcopy(self.origin_config) + new_config.update(request.metadata["em_config"]) + data = EMService.init_class_by_config(new_config) + summarizer = data["summarizer"] + else: + assert self.summarizer is not None, "summarizer must be provided." + summarizer = self.summarizer + + experiences: List[Experience] = summarizer.execute(trajectories=request.trajectories, + workspace_id=request.workspace_id, + **request.metadata) return SummarizerResponse(experiences=experiences) @@ -238,12 +276,14 @@ if __name__ == "__main__": timeout_keep_alive=service.timeout_keep_alive, limit_concurrency=service.limit_concurrency) -# launch with: -# 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", "index_name": "naive_agent"}' \ -# --agent_wrapper='{"backend": "simple", "max_steps": 10}' \ -# --context_generator='{"backend": "simple", "retrieve_top_k": 1}' \ -# --summarizer='{"backend": "simple"}' +""" +launch with: +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", "index_name": "simple_agent"}' \ + --agent_wrapper='{"backend": "simple", "max_steps": 10}' \ + --context_generator='{"backend": "simple", "retrieve_top_k": 1}' \ + --summarizer='{"backend": "simple"}' +""" diff --git a/experiencemaker/module/agent_wrapper/agent_wrapper_mixin.py b/experiencemaker/module/agent_wrapper/agent_wrapper_mixin.py index a432834f..c0ac7316 100644 --- a/experiencemaker/module/agent_wrapper/agent_wrapper_mixin.py +++ b/experiencemaker/module/agent_wrapper/agent_wrapper_mixin.py @@ -10,7 +10,6 @@ from experiencemaker.schema.trajectory import Trajectory class AgentWrapperMixin(BaseModel, ABC): context_generator: BaseContextGenerator | None = Field(default=None) llm: BaseLLM | None = Field(default=None) - workspace_id: str = Field(default="") - def execute(self, query: str, **kwargs) -> Trajectory: + def execute(self, query: str, workspace_id: str = None, **kwargs) -> Trajectory: raise NotImplementedError diff --git a/experiencemaker/module/agent_wrapper/simple_agent.py b/experiencemaker/module/agent_wrapper/simple_agent.py index b2db9b24..98d6f6bd 100644 --- a/experiencemaker/module/agent_wrapper/simple_agent.py +++ b/experiencemaker/module/agent_wrapper/simple_agent.py @@ -12,10 +12,9 @@ from experiencemaker.tool import CodeTool, DashscopeSearchTool, TerminateTool from experiencemaker.tool.base_tool import BaseTool -class SimpleAgentContext(BaseModel): +class AgentContext(BaseModel): current_step: int = Field(default=-1) query: str = Field(default="") - previous_experience: str = Field(default="") messages: List[Message] = Field(default_factory=list) metadata: dict = Field(default_factory=dict) has_terminate_tool: bool = Field(default=False) @@ -27,16 +26,14 @@ class SimpleAgent(PromptMixin): tools: List[BaseTool] = [CodeTool(), DashscopeSearchTool(), TerminateTool()] prompt_file_path: Path = Path(__file__).parent / "simple_agent_prompt.yaml" - def think(self, context: SimpleAgentContext): + def think(self, context: AgentContext): now_time = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S') tool_names = [x.name for x in self.tools] if context.current_step == 0: user_prompt = self.prompt_format(prompt_name="role_prompt", - experience_tag=False if context.previous_experience else True, time=now_time, tools=", ".join(tool_names), - previous_insight=context.previous_experience, query=context.query) elif context.has_terminate_tool: @@ -64,7 +61,7 @@ class SimpleAgent(PromptMixin): f"tool_calls={action_msg.tool_calls}") return True if action_msg.tool_calls else False - def act(self, context: SimpleAgentContext): + def act(self, context: AgentContext): action_msg = context.messages[-1] assert isinstance(action_msg, ActionMessage) @@ -84,8 +81,8 @@ class SimpleAgent(PromptMixin): context.messages.append(state_msg) logger.info(f"step.{context.current_step} state_msg_context={state_msg.content}") - def run(self, query: str, previous_experience: str) -> List[Message]: - context: SimpleAgentContext = SimpleAgentContext(query=query, previous_experience=previous_experience) + def run(self, query: str) -> List[Message]: + context: AgentContext = AgentContext(query=query) for i in range(self.max_steps): context.current_step = i diff --git a/experiencemaker/module/agent_wrapper/simple_agent_prompt.yaml b/experiencemaker/module/agent_wrapper/simple_agent_prompt.yaml index 514c2ebe..eaa1cc81 100644 --- a/experiencemaker/module/agent_wrapper/simple_agent_prompt.yaml +++ b/experiencemaker/module/agent_wrapper/simple_agent_prompt.yaml @@ -5,12 +5,7 @@ role_prompt: | Please proactively choose the most suitable tool or combination of tools based on the user's question, including {tools} etc. For complex tasks, you can break down the problem step by step and use different tools to solve it incrementally. Please determine the response language based on the language of the user's question. - [experience_tag] - [experience_tag]Previous Experience - [experience_tag]{previous_experience} - [experience_tag]Please consider the helpful parts from these in answering the question, to make the response more comprehensive and substantial. - - User's question + {query} next_prompt: | diff --git a/experiencemaker/module/agent_wrapper/simple_agent_wrapper.py b/experiencemaker/module/agent_wrapper/simple_agent_wrapper.py index 0a984897..eee172b3 100644 --- a/experiencemaker/module/agent_wrapper/simple_agent_wrapper.py +++ b/experiencemaker/module/agent_wrapper/simple_agent_wrapper.py @@ -5,12 +5,12 @@ from experiencemaker.schema.trajectory import Trajectory class SimpleAgentWrapper(SimpleAgent, AgentWrapperMixin): - def execute(self, query: str, **kwargs) -> Trajectory: + def execute(self, query: str, workspace_id: str = None, **kwargs) -> Trajectory: trajectory = Trajectory(query=query) - context_msg = self.context_generator.execute(trajectory=trajectory) - previous_experience = context_msg.content + context_msg = self.context_generator.execute(trajectory=trajectory, workspace_id=workspace_id) + new_query = f"{context_msg.content}\n\nUser Question\n{query}" - messages = self.run(query, previous_experience) + messages = self.run(new_query) trajectory.steps = messages trajectory.answer = messages[-1].content diff --git a/experiencemaker/module/context_generator/base_context_generator.py b/experiencemaker/module/context_generator/base_context_generator.py index c22b6509..fe1cb58a 100644 --- a/experiencemaker/module/context_generator/base_context_generator.py +++ b/experiencemaker/module/context_generator/base_context_generator.py @@ -12,12 +12,12 @@ from experiencemaker.storage.base_vector_store import BaseVectorStore class BaseContextGenerator(BaseModel, ABC): vector_store: BaseVectorStore | None = Field(default=None) llm: BaseLLM | None = Field(default=None) - workspace_id: str = Field(default="") def _build_retrieve_query(self, trajectory: Trajectory, **kwargs) -> str: raise NotImplementedError - def _retrieve_by_query(self, trajectory: Trajectory, query: str, **kwargs) -> List[VectorStoreNode]: + def _retrieve_by_query(self, trajectory: Trajectory, query: str, workspace_id: str, retrieve_top_k: int, + **kwargs) -> List[VectorStoreNode]: raise NotImplementedError def _generate_context_message(self, @@ -26,8 +26,13 @@ class BaseContextGenerator(BaseModel, ABC): **kwargs) -> ContextMessage: raise NotImplementedError - def execute(self, trajectory: Trajectory, **kwargs) -> ContextMessage: + def execute(self, trajectory: Trajectory, workspace_id: str = None, retrieve_top_k: int = 1, + **kwargs) -> ContextMessage: query: str = self._build_retrieve_query(trajectory, **kwargs) - nodes: List[VectorStoreNode] = self._retrieve_by_query(trajectory, query, **kwargs) + nodes: List[VectorStoreNode] = self._retrieve_by_query(trajectory=trajectory, + query=query, + workspace_id=workspace_id, + retrieve_top_k=retrieve_top_k, + **kwargs) context_msg: ContextMessage = self._generate_context_message(trajectory, nodes, **kwargs) return context_msg diff --git a/experiencemaker/module/context_generator/simple_context_generator.py b/experiencemaker/module/context_generator/simple_context_generator.py index cc104e9c..cd2b4ec9 100644 --- a/experiencemaker/module/context_generator/simple_context_generator.py +++ b/experiencemaker/module/context_generator/simple_context_generator.py @@ -1,7 +1,5 @@ from typing import List -from pydantic import Field - from experiencemaker.module.context_generator.base_context_generator import BaseContextGenerator from experiencemaker.schema.experience import Experience from experiencemaker.schema.trajectory import Trajectory, ContextMessage @@ -9,7 +7,6 @@ from experiencemaker.schema.vector_store_node import VectorStoreNode class SimpleContextGenerator(BaseContextGenerator): - retrieve_top_k: int = Field(default=5) def _build_retrieve_query(self, trajectory: Trajectory, **kwargs) -> str: query = "" @@ -17,11 +14,12 @@ class SimpleContextGenerator(BaseContextGenerator): query = trajectory.query return query - def _retrieve_by_query(self, trajectory: Trajectory, query: str, **kwargs) -> List[VectorStoreNode]: + def _retrieve_by_query(self, trajectory: Trajectory, query: str, workspace_id: str, retrieve_top_k: int, + **kwargs) -> List[VectorStoreNode]: if not query: return [] - return self.vector_store.retrieve_by_query(query=query, top_k=self.retrieve_top_k) + return self.vector_store.retrieve_by_query(query=query, top_k=retrieve_top_k, index_name=workspace_id, **kwargs) def _generate_context_message(self, trajectory: Trajectory, @@ -30,11 +28,13 @@ class SimpleContextGenerator(BaseContextGenerator): if not nodes: return ContextMessage(content="") - content = "" + content = "Previous Experience\n" for node in nodes: experience: Experience = Experience.from_vector_store_node(node) if not experience.experience_content: continue content += f"- {experience.experience_desc} {experience.experience_content}\n" + content += "Please consider the helpful parts from these in answering the question, to make the response more comprehensive and substantial." + return ContextMessage(content=content.strip()) diff --git a/experiencemaker/module/summarizer/base_summarizer.py b/experiencemaker/module/summarizer/base_summarizer.py index d668f1a9..40e00e26 100644 --- a/experiencemaker/module/summarizer/base_summarizer.py +++ b/experiencemaker/module/summarizer/base_summarizer.py @@ -13,17 +13,14 @@ from experiencemaker.storage.base_vector_store import BaseVectorStore class BaseSummarizer(BaseModel, ABC): vector_store: BaseVectorStore | None = Field(default=None) llm: BaseLLM | None = Field(default=None) - workspace_id: str = Field(default="") def _extract_experiences(self, trajectories: List[Trajectory], **kwargs) -> List[Experience]: raise NotImplementedError - def execute(self, trajectories: List[Trajectory], return_experience: bool = True, **kwargs) -> List[Experience]: + def execute(self, trajectories: List[Trajectory], workspace_id: str = None, **kwargs) -> List[Experience]: experiences: List[Experience] = self._extract_experiences(trajectories, **kwargs) nodes: List[VectorStoreNode] = [x.to_vector_store_node() for x in experiences] - self.vector_store.insert(nodes, **kwargs) + self.vector_store.insert(nodes, index_name=workspace_id, **kwargs) - if return_experience: - return experiences - return [] + return experiences diff --git a/experiencemaker/schema/experience.py b/experiencemaker/schema/experience.py index 017a6fb1..61cf663a 100644 --- a/experiencemaker/schema/experience.py +++ b/experiencemaker/schema/experience.py @@ -35,12 +35,14 @@ class Experience(BaseModel): metadata: dict = { "experience_role": self.experience_role, "experience_content": self.experience_content, - "experience_function": self.experience_function.model_dump(), "experience_score": self.experience_score, "experience_created_time": self.experience_created_time, "experience_modified_time": self.experience_modified_time, "metadata": self.metadata, } + if self.experience_function: + metadata["experience_function"] = self.experience_function.model_dump(), + return VectorStoreNode( unique_id=self.experience_id, workspace_id=self.experience_workspace_id, diff --git a/experiencemaker/schema/request.py b/experiencemaker/schema/request.py index 57e23fcc..5ffe659b 100644 --- a/experiencemaker/schema/request.py +++ b/experiencemaker/schema/request.py @@ -8,6 +8,7 @@ from experiencemaker.schema.trajectory import Trajectory class BaseRequest(BaseModel, ABC): metadata: dict = Field(default_factory=dict) + workspace_id: str = Field(default="") class AgentWrapperRequest(BaseRequest): diff --git a/experiencemaker/storage/base_vector_store.py b/experiencemaker/storage/base_vector_store.py index 9674c6cb..9f6afb54 100644 --- a/experiencemaker/storage/base_vector_store.py +++ b/experiencemaker/storage/base_vector_store.py @@ -9,18 +9,31 @@ from experiencemaker.schema.vector_store_node import VectorStoreNode class BaseVectorStore(BaseModel, ABC): embedding_model: BaseEmbeddingModel = Field(default=...) + index_name: str = Field(default=...) - def insert(self, nodes: VectorStoreNode | List[VectorStoreNode], **kwargs): + def exist_index(self, index_name: str = None) -> bool: raise NotImplementedError - def update(self, nodes: VectorStoreNode | List[VectorStoreNode], **kwargs): + def delete_index(self, index_name: str = None): raise NotImplementedError - def delete_by_id(self, unique_id: str, **kwargs): + def create_index(self, index_name: str = None): raise NotImplementedError - def retrieve_by_id(self, unique_id: str, **kwargs) -> VectorStoreNode | None: + def exist_id(self, unique_id: str, index_name: str = None): raise NotImplementedError - def retrieve_by_query(self, query: str, top_k: int = 3, **kwargs) -> List[VectorStoreNode]: + def insert(self, nodes: VectorStoreNode | List[VectorStoreNode], index_name: str = None, **kwargs): + raise NotImplementedError + + def update(self, nodes: VectorStoreNode | List[VectorStoreNode], index_name: str = None, **kwargs): + raise NotImplementedError + + def delete_by_id(self, unique_id: str, index_name: str = None, **kwargs): + raise NotImplementedError + + def retrieve_by_id(self, unique_id: str, index_name: str = None, **kwargs) -> VectorStoreNode | None: + raise NotImplementedError + + def retrieve_by_query(self, query: str, top_k: int = 1, index_name: str = None, **kwargs) -> List[VectorStoreNode]: raise NotImplementedError diff --git a/experiencemaker/storage/es_vector_store.py b/experiencemaker/storage/es_vector_store.py index 3dcba383..2efa5a10 100644 --- a/experiencemaker/storage/es_vector_store.py +++ b/experiencemaker/storage/es_vector_store.py @@ -13,7 +13,6 @@ from experiencemaker.storage.base_vector_store import BaseVectorStore class EsVectorStore(BaseVectorStore): hosts: str | List[str] = Field(default_factory=lambda: os.getenv("ES_HOSTS", "http://localhost:9200")) - index_name: str = Field(default=...) basic_auth: str | Tuple[str, str] | None = Field(default=None) bulk_chunk_size: int = Field(default=512) retrieve_filters: List[dict] = [] @@ -28,13 +27,23 @@ class EsVectorStore(BaseVectorStore): self._client = Elasticsearch(hosts=hosts, basic_auth=self.basic_auth) return self - def delete_index(self): - if self._client.indices.exists(index=self.index_name): - self._client.indices.delete(index=self.index_name) + 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 create_index(self): - if self._client.indices.exists(index=self.index_name): - logger.warning(f"index_name={self.index_name} is already exists!") + 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 = { @@ -51,10 +60,12 @@ class EsVectorStore(BaseVectorStore): } } - return self._client.indices.create(index=self.index_name, body=index) + return self._client.indices.create(index=index_name, body=index) - def refresh_index(self): - self._client.indices.refresh(index=self.index_name) + 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: @@ -64,12 +75,17 @@ class EsVectorStore(BaseVectorStore): node.metadata["_score"] = doc["_score"] - 1 return node - def exist_id(self, doc_id: str): - return self._client.exists(index=self.index_name, id=doc_id) + 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 node2doc(self, node: VectorStoreNode, add_op_type: bool = False) -> dict: doc: dict = { - "_index": self.index_name, + "_index": index_name, "_id": node.unique_id, "_source": { "workspace_id": node.workspace_id, @@ -80,13 +96,12 @@ class EsVectorStore(BaseVectorStore): } if add_op_type: - doc["_op_type"] = "update" if self.exist_id(node.unique_id) else "index", + 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): @@ -97,14 +112,20 @@ class EsVectorStore(BaseVectorStore): 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, **kwargs): + 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] @@ -112,40 +133,59 @@ class EsVectorStore(BaseVectorStore): 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] + 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() + 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) - def update(self, nodes: VectorStoreNode | List[VectorStoreNode], refresh_index: bool = True, **kwargs): if isinstance(nodes, VectorStoreNode): nodes = [nodes] nodes = self.embedding_model.get_node_embeddings(nodes) - docs = [self.node2doc(node, True) for node in 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() + self.refresh_index(index_name) - def delete_by_id(self, unique_id: str, **kwargs): - return self._client.delete(index=self.index_name, id=unique_id, **kwargs) + 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 - def retrieve_by_id(self, unique_id: str, **kwargs) -> VectorStoreNode | None: try: - doc = self._client.get(index=self.index_name, id=unique_id, **kwargs) + doc = self._client.get(index=index_name, id=unique_id, **kwargs) return self.doc2node(doc) except Exception as e: - logger.warning(f"retrieve_by_id unique_id={unique_id} is not found with error={e.args}") + 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 = 3, **kwargs) -> List[VectorStoreNode]: + 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 + query_vector = self.embedding_model.get_embeddings(query) body = { @@ -160,7 +200,7 @@ class EsVectorStore(BaseVectorStore): }, "size": top_k } - response = self._client.search(index=self.index_name, body=body, **kwargs) + response = self._client.search(index=index_name, body=body, **kwargs) nodes: List[VectorStoreNode] = [] for doc in response['hits']['hits']: diff --git a/experiencemaker/storage/file_vector_store.py b/experiencemaker/storage/file_vector_store.py index d06c91a6..6c14e5d8 100644 --- a/experiencemaker/storage/file_vector_store.py +++ b/experiencemaker/storage/file_vector_store.py @@ -13,15 +13,13 @@ from experiencemaker.storage.base_vector_store import BaseVectorStore class FileVectorStore(BaseVectorStore): - store_dir: str = Field(default="./") - index_name: str = Field(default=...) + 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" @@ -29,57 +27,75 @@ class FileVectorStore(BaseVectorStore): self.index_path.touch(exist_ok=True) return self - def delete_index(self): - with self._thread_lock: - if self.index_path.exists() and self.index_path.is_file(): - self.index_path.unlink() + 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 create_index(self): + def exist_index(self, index_name: str = None) -> bool: + index_path = self.get_index_path(index_name) with self._thread_lock: - if not self.index_path.exists(): - self.index_path.touch(exist_ok=True) + 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) - def load(self) -> List[VectorStoreNode]: nodes = [] with self._thread_lock: - with open(self.index_path) as f: + with open(index_path) as f: for line in f: if line.strip(): nodes.append(VectorStoreNode(**json.loads(line))) return nodes - def _load(self) -> List[VectorStoreNode]: - nodes = [] - with self._thread_lock: - with open(self.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) - def _dump(self, nodes: List[VectorStoreNode]): with self._thread_lock: - with open(self.index_path, "w") as f: + 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): - nodes = self._load() + 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], **kwargs): - return self.update(nodes, **kwargs) + 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 - def update(self, nodes: VectorStoreNode | List[VectorStoreNode], **kwargs): 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() + exist_nodes: List[VectorStoreNode] = self._load(index_name=index_name) for node in exist_nodes: all_node_dict[node.unique_id] = node @@ -90,22 +106,26 @@ class FileVectorStore(BaseVectorStore): all_node_dict[node.unique_id] = node - self._dump(list(all_node_dict.values())) - logger.info(f"update nodes.size={len(nodes)} all.size={len(all_node_dict)} update_cnt={update_cnt}") + 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, **kwargs): - nodes = self._load() + 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) + 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, **kwargs) -> VectorStoreNode | None: - nodes = self._load() + 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 @@ -123,9 +143,9 @@ 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 = 3, **kwargs) -> List[VectorStoreNode]: + 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() + nodes: List[VectorStoreNode] = self._load(index_name=index_name) for node in nodes: node.metadata["score"] = self.calculate_similarity(query_vector, node.vector) diff --git a/experiencemaker/utils/util_function.py b/experiencemaker/utils/util_function.py index 08fde209..45b827e6 100644 --- a/experiencemaker/utils/util_function.py +++ b/experiencemaker/utils/util_function.py @@ -13,11 +13,11 @@ def get_html_match_content(content: str, key: str): return None -def load_env_keys(): - if os.path.exists(".env"): - with open(".env") as f: +def load_env_keys(file_path: str = ".env"): + if os.path.exists(file_path): + with open(file_path) as f: config = json.load(f) for k, v in config.items(): os.environ[k] = v else: - logger.warning(".env file not found~") + logger.warning(f"{file_path} file not found~")