From 2fffb847a0907cdbce20862c641c5baf72953ce9 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 28 Jan 2026 01:18:06 +0800 Subject: [PATCH] refactor(memory): restructure memory tools and handlers --- pyproject.toml | 2 +- reme/agent/memory/base_memory_agent.py | 55 ++-- reme/agent/memory/default/reme_retriever.py | 4 +- reme/agent/memory/default/reme_retriever.yaml | 2 +- reme/agent/memory/default/reme_summarizer.py | 16 +- .../agent/memory/default/reme_summarizer.yaml | 2 +- reme/core/op/base_op.py | 1 + reme/core/schema/memory_node.py | 32 +++ reme/reme.py | 241 +++++++++--------- reme/tool/memory/__init__.py | 42 ++- reme/tool/memory/{history => }/add_history.py | 13 +- reme/tool/memory/add_memory.py | 135 ++++++++++ reme/tool/memory/base_memory_tool.py | 36 ++- reme/tool/memory/delegate_task.py | 83 ++++++ .../tool/memory/{vector => }/delete_memory.py | 32 +-- reme/tool/memory/hands_off/__init__.py | 0 reme/tool/memory/hands_off/hands_off.py | 108 -------- reme/tool/memory/history/__init__.py | 0 reme/tool/memory/identity/__init__.py | 0 reme/tool/memory/identity/add_identity.py | 42 --- reme/tool/memory/identity/read_identity.py | 36 --- reme/tool/memory/memory_handler.py | 210 +++++++++++++++ reme/tool/memory/meta/__init__.py | 0 reme/tool/memory/meta/add_meta_memory.py | 89 ------- reme/tool/memory/meta/read_meta_memory.py | 75 ------ reme/tool/memory/profile_handler.py | 199 +++++++++++++++ .../tool/memory/{history => }/read_history.py | 6 +- reme/tool/memory/read_profile.py | 45 ++++ reme/tool/memory/retrieve_memory.py | 121 +++++++++ reme/tool/memory/retrieve_recent_memory.py | 52 ++++ reme/tool/memory/update_memory.py | 139 ++++++++++ reme/tool/memory/update_profile.py | 97 +++++++ reme/tool/memory/user_profile/__init__.py | 0 .../memory/user_profile/read_user_profile.py | 63 ----- .../user_profile/update_user_profile.py | 112 -------- reme/tool/memory/vector/__init__.py | 0 reme/tool/memory/vector/add_memory.py | 110 -------- reme/tool/memory/vector/retrieve_memory.py | 136 ---------- .../memory/vector/retrieve_recent_memory.py | 62 ----- reme/tool/memory/vector/update_memory.py | 126 --------- 40 files changed, 1323 insertions(+), 1201 deletions(-) rename reme/tool/memory/{history => }/add_history.py (77%) create mode 100644 reme/tool/memory/add_memory.py create mode 100644 reme/tool/memory/delegate_task.py rename reme/tool/memory/{vector => }/delete_memory.py (59%) delete mode 100644 reme/tool/memory/hands_off/__init__.py delete mode 100644 reme/tool/memory/hands_off/hands_off.py delete mode 100644 reme/tool/memory/history/__init__.py delete mode 100644 reme/tool/memory/identity/__init__.py delete mode 100644 reme/tool/memory/identity/add_identity.py delete mode 100644 reme/tool/memory/identity/read_identity.py create mode 100644 reme/tool/memory/memory_handler.py delete mode 100644 reme/tool/memory/meta/__init__.py delete mode 100644 reme/tool/memory/meta/add_meta_memory.py delete mode 100644 reme/tool/memory/meta/read_meta_memory.py create mode 100644 reme/tool/memory/profile_handler.py rename reme/tool/memory/{history => }/read_history.py (90%) create mode 100644 reme/tool/memory/read_profile.py create mode 100644 reme/tool/memory/retrieve_memory.py create mode 100644 reme/tool/memory/retrieve_recent_memory.py create mode 100644 reme/tool/memory/update_memory.py create mode 100644 reme/tool/memory/update_profile.py delete mode 100644 reme/tool/memory/user_profile/__init__.py delete mode 100644 reme/tool/memory/user_profile/read_user_profile.py delete mode 100644 reme/tool/memory/user_profile/update_user_profile.py delete mode 100644 reme/tool/memory/vector/__init__.py delete mode 100644 reme/tool/memory/vector/add_memory.py delete mode 100644 reme/tool/memory/vector/retrieve_memory.py delete mode 100644 reme/tool/memory/vector/retrieve_recent_memory.py delete mode 100644 reme/tool/memory/vector/update_memory.py diff --git a/pyproject.toml b/pyproject.toml index 47aa478d..63bedc49 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -76,6 +76,6 @@ Documentation = "https://reme.agentscope.io/" Repository = "https://github.com/agentscope-ai/ReMe" [project.scripts] -reme = "reme_ai.main:main" +reme = "reme.reme:main" # python -m build && twine upload dist/* diff --git a/reme/agent/memory/base_memory_agent.py b/reme/agent/memory/base_memory_agent.py index 6ad3499d..bdb824f4 100644 --- a/reme/agent/memory/base_memory_agent.py +++ b/reme/agent/memory/base_memory_agent.py @@ -1,9 +1,6 @@ """Base memory agent for handling memory operations with tool-based reasoning.""" from abc import ABCMeta -from typing import Literal - -from loguru import logger from ...core.enumeration import MemoryType from ...core.op import BaseReact @@ -15,40 +12,6 @@ class BaseMemoryAgent(BaseReact, metaclass=ABCMeta): memory_type: MemoryType | None = None - @staticmethod - async def read_meta_memories(meta_memories: list[dict]) -> str: - """Read and format meta memory information from the provided metadata list.""" - from ...tool.memory import ReadMetaMemory - - meta_memory_info = ReadMetaMemory().format_memory_metadata(meta_memories) - logger.info(f"meta_memory_info={meta_memory_info}") - return meta_memory_info - - async def read_user_profile(self, show_id: Literal["profile", "history"] = "profile") -> str: - """Read current user profile.""" - from ...tool.memory import ReadUserProfile - - read_tool = ReadUserProfile(show_id=show_id) - await read_tool.call(memory_target=self.memory_target, service_context=self.service_context) - return str(read_tool.response.answer) - - async def add_history_node(self) -> MemoryNode: - """Add history node""" - from ...tool.memory import AddHistory - - add_history_tool = AddHistory() - await add_history_tool.call( - messages=self.messages, - description=self.description, - service_context=self.service_context, - ) - return add_history_tool.context.history_node - - @property - def memory_target(self) -> str: - """memory_target""" - return self.context.get("memory_target", "") - @property def query(self) -> str: """query""" @@ -64,6 +27,11 @@ class BaseMemoryAgent(BaseReact, metaclass=ABCMeta): """description""" return self.context.get("description", "") + @property + def memory_target(self) -> str: + """memory_target""" + return self.context.memory_target + @property def history_node(self) -> MemoryNode: """Returns the history node.""" @@ -80,3 +48,16 @@ class BaseMemoryAgent(BaseReact, metaclass=ABCMeta): if "retrieved_nodes" not in self.context: self.context.retrieved_nodes = [] return self.context.retrieved_nodes + + @property + def memory_target_type_mapping(self) -> dict[str, MemoryType]: + """Get the memory target type mapping from context.""" + return self.context.memory_target_type_mapping + + @property + def meta_memory_info(self) -> str: + """Get the meta memory info from context.""" + lines = ["Format: - memory_target: memory_type memories about memory_target"] + for memory_target, memory_type in self.memory_target_type_mapping.items(): + lines.append(f"- {memory_target}: {memory_type} memories about {memory_target}") + return "\n".join(lines) diff --git a/reme/agent/memory/default/reme_retriever.py b/reme/agent/memory/default/reme_retriever.py index a0639f9b..b7a917f9 100644 --- a/reme/agent/memory/default/reme_retriever.py +++ b/reme/agent/memory/default/reme_retriever.py @@ -79,8 +79,8 @@ class ReMeRetriever(BaseMemoryAgent): async def execute(self): result = await super().execute() tools: list[BaseTool] = result["tools"] - hands_off_tool = tools[0] - agents: list[BaseMemoryAgent] = hands_off_tool.response.metadata["agents"] + delegate_task_tool = tools[0] + agents: list[BaseMemoryAgent] = delegate_task_tool.response.metadata["agents"] answer = [] success = True diff --git a/reme/agent/memory/default/reme_retriever.yaml b/reme/agent/memory/default/reme_retriever.yaml index 849f562e..11d7186e 100644 --- a/reme/agent/memory/default/reme_retriever.yaml +++ b/reme/agent/memory/default/reme_retriever.yaml @@ -10,7 +10,7 @@ system_prompt: | {meta_memory_info} ## Your Task - Use the `hands_off` tool to retrieve information from specialized agents: + Use the `delegate_task` tool to retrieve information from specialized agents: 1. Analyze the user query and identify which memory dimensions are relevant 2. Specify `memory_type` and `memory_target` for each retrieval task - The `memory_type` and `memory_target` must **exactly match** existing entries in the "Available Memory Agents" listed above diff --git a/reme/agent/memory/default/reme_summarizer.py b/reme/agent/memory/default/reme_summarizer.py index 38ab85d7..c9b4db81 100644 --- a/reme/agent/memory/default/reme_summarizer.py +++ b/reme/agent/memory/default/reme_summarizer.py @@ -13,6 +13,18 @@ class ReMeSummarizer(BaseMemoryAgent): super().__init__(**kwargs) self.meta_memories: list[dict] = meta_memories or [] + async def add_history_node(self) -> MemoryNode: + """Add history node""" + from ...tool.memory import AddHistory + + add_history_tool = AddHistory() + await add_history_tool.call( + messages=self.messages, + description=self.description, + service_context=self.service_context, + ) + return add_history_tool.context.history_node + async def build_messages(self) -> list[Message]: self.context.history_node = await self.add_history_node() @@ -75,8 +87,8 @@ class ReMeSummarizer(BaseMemoryAgent): async def execute(self): result = await super().execute() tools: list[BaseTool] = result["tools"] - hands_off_tool = tools[0] - agents: list[BaseMemoryAgent] = hands_off_tool.response.metadata["agents"] + delegate_task_tool = tools[0] + agents: list[BaseMemoryAgent] = delegate_task_tool.response.metadata["agents"] success = True messages = [] diff --git a/reme/agent/memory/default/reme_summarizer.yaml b/reme/agent/memory/default/reme_summarizer.yaml index d31e0213..d0a6748e 100644 --- a/reme/agent/memory/default/reme_summarizer.yaml +++ b/reme/agent/memory/default/reme_summarizer.yaml @@ -10,7 +10,7 @@ system_prompt: | {meta_memory_info} ## Your Task - Use the `hands_off` tool to distribute memory tasks to specialized agents: + Use the `delegate_task` tool to distribute memory tasks to specialized agents: 1. Analyze the context and identify which memory dimensions require updates 2. Specify `memory_type` and `memory_target` for each task - The `memory_type` and `memory_target` must **exactly match** existing entries in the "Available Memory Agents" listed above diff --git a/reme/core/op/base_op.py b/reme/core/op/base_op.py index 9c2dff71..6d32ddc0 100644 --- a/reme/core/op/base_op.py +++ b/reme/core/op/base_op.py @@ -115,6 +115,7 @@ class BaseOp(metaclass=ABCMeta): @property def service_context(self) -> ServiceContext: """Access the service context.""" + assert self.context, "Service context is not initialized!" return self.context.service_context @property diff --git a/reme/core/schema/memory_node.py b/reme/core/schema/memory_node.py index 67ed7c43..3cf096ac 100644 --- a/reme/core/schema/memory_node.py +++ b/reme/core/schema/memory_node.py @@ -36,6 +36,7 @@ class MemoryNode(BaseModel): memory_target: Target or topic this memory relates to. when_to_use: Condition description for vector retrieval. content: Actual memory content. + message_time: Time of the message that generated this memory. ref_memory_id: Reference to related raw history memory. time_created: Creation timestamp. time_modified: Last modification timestamp. @@ -49,6 +50,7 @@ class MemoryNode(BaseModel): memory_target: str = Field(default="", description="Target or topic of the memory") when_to_use: str = Field(default="", description="Condition description for vector retrieval") content: str = Field(default="", description="Actual memory content") + message_time: str = Field(default="", description="Time of the message that generated this memory") ref_memory_id: str = Field(default="", description="Reference to related raw history memory ID") time_created: str = Field(default_factory=get_now_time, description="Creation timestamp") @@ -123,6 +125,7 @@ class MemoryNode(BaseModel): metadata: dict[str, Any] = { "memory_type": self.memory_type.value, "memory_target": self.memory_target, + "message_time": self.message_time, "ref_memory_id": self.ref_memory_id, "time_created": self.time_created, "time_modified": self.time_modified, @@ -145,6 +148,34 @@ class MemoryNode(BaseModel): metadata=metadata, ) + def format( + self, + include_memory_id: bool = True, + include_when_to_use: bool = True, + include_content: bool = True, + include_message_time: bool = True, + ref_memory_id_key: str = "", + ) -> str: + """Format memory node as string with configurable fields.""" + line = "" + + if include_memory_id and self.memory_id: + line += f"memory_id={self.memory_id} " + + if include_message_time and self.message_time: + line += f"[{self.message_time}] " + + if include_when_to_use and self.when_to_use: + line += f"{self.when_to_use} " + + if include_content and self.content: + line += self.content.strip() + + if ref_memory_id_key and self.ref_memory_id: + line += f" {ref_memory_id_key}={self.ref_memory_id}" + + return line.strip() + @classmethod def from_vector_node(cls, node: VectorNode) -> "MemoryNode": """Reconstruct MemoryNode from VectorNode. @@ -189,6 +220,7 @@ class MemoryNode(BaseModel): memory_target=metadata.pop("memory_target", ""), when_to_use=when_to_use, content=content, + message_time=metadata.pop("message_time", ""), ref_memory_id=metadata.pop("ref_memory_id", ""), time_created=metadata.pop("time_created", ""), time_modified=metadata.pop("time_modified", ""), diff --git a/reme/reme.py b/reme/reme.py index 015b962b..0c4a0fc2 100644 --- a/reme/reme.py +++ b/reme/reme.py @@ -2,6 +2,9 @@ import asyncio import sys +from pathlib import Path + +from loguru import logger from .agent.memory.default import ReMeSummarizer, PersonalSummarizer, PersonalRetriever, ReMeRetriever from .config import ReMeConfigParser @@ -14,7 +17,8 @@ from .core.schema import Response, Message, MemoryNode, VectorNode from .core.token_counter import BaseTokenCounter from .core.utils import execute_stream_task, get_now_time from .core.vector_store import BaseVectorStore -from .tool.memory import UpdateUserProfile, RetrieveMemory, AddMemory, HandsOff, ReadHistory, ReadUserProfile +from .tool.memory import UpdateUserProfile, RetrieveMemory, AddMemory, DelegateTask, ReadHistory, ReadUserProfile, \ + ProfileHandler class ReMe: @@ -32,8 +36,38 @@ class ReMe: embedding_model: dict | None = None, vector_store: dict | None = None, token_counter: dict | None = None, + personal_memory_target: list[str] | None = None, + procedural_memory_target: list[str] | None = None, + tool_memory_target: list[str] | None = None, + profile_path: str = "reme_profile", + main_summary_version: str = "default", + personal_summary_version: str = "default", + procedural_summary_version: str = "default", + tool_summary_version: str = "default", + main_retrieve_version: str = "default", + personal_retrieve_version: str = "default", + procedural_retrieve_version: str = "default", + tool_retrieve_version: str = "default", **kwargs, ): + # MemoryTarget -> MemoryType + memory_target_type_mapping: dict[str, MemoryType] = {} + if personal_memory_target: + for name in personal_memory_target: + assert name not in memory_target_type_mapping, f"Memory target name {name} is already used." + memory_target_type_mapping[name] = MemoryType.PERSONAL + + if procedural_memory_target: + for name in procedural_memory_target: + assert name not in memory_target_type_mapping, f"Memory target name {name} is already used." + memory_target_type_mapping[name] = MemoryType.PROCEDURAL + + if tool_memory_target: + for name in tool_memory_target: + assert name not in memory_target_type_mapping, f"Memory target name {name} is already used." + memory_target_type_mapping[name] = MemoryType.TOOL + + # ServiceContext self.service_context = ServiceContext( *args, llm_api_key=llm_api_key, @@ -48,64 +82,73 @@ class ReMe: embedding_model=embedding_model, vector_store=vector_store, token_counter=token_counter, + memory_target_type_mapping=memory_target_type_mapping, **kwargs, ) + self.profile_path: str = profile_path + + # PromptHandler self.prompt_handler = PromptHandler(language=self.service_context.language) - async def __aenter__(self): - """Async context manager entry.""" - return self - - def __enter__(self): - """Context manager entry.""" - return self - - async def close(self): - """Close""" - return await self.service_context.close() - - def close_sync(self): - """Close synchronously""" - self.service_context.close_sync() - - async def __aexit__(self, exc_type=None, exc_val=None, exc_tb=None): - """Async context manager exit.""" - await self.close() - return False - - def __exit__(self, exc_type=None, exc_val=None, exc_tb=None): - """Context manager exit.""" - self.close_sync() - return False + # LLM & EmbeddingModel & VectorStore & TokenCounter + self.llm: BaseLLM | None = self.service_context.llms.get("default", None) + self.embedding_model: BaseEmbeddingModel | None = self.service_context.embedding_models.get("default", None) + self.vector_store: BaseVectorStore | None = self.service_context.vector_stores.get("default", None) + self.token_counter: BaseTokenCounter | None = self.service_context.token_counters.get("default", None) @property - def llm(self) -> BaseLLM: - """Return the default LLM instance from the service context.""" - return self.service_context.llms["default"] + def memory_target_type_mapping(self) -> dict[str, MemoryType]: + mapping = {} + if self.service_context.personal_memory_target: + for name in self.service_context.personal_memory_target: + assert name not in mapping, f"Memory target name {name} is already used." + mapping[name] = MemoryType.PERSONAL - @property - def embedding_model(self) -> BaseEmbeddingModel: - """Return the default embedding model instance from the service context.""" - return self.service_context.embedding_models["default"] + if self.service_context.procedural_memory_target: + for name in self.service_context.procedural_memory_target: + assert name not in mapping, f"Memory target name {name} is already used." + mapping[name] = MemoryType.PROCEDURAL + + if self.service_context.tool_memory_target: + for name in self.service_context.tool_memory_target: + assert name not in mapping, f"Memory target name {name} is already used." + mapping[name] = MemoryType.TOOL + return mapping + + def add_meta_memory(self, memory_type: str | MemoryType, memory_target: str): + memory_type = MemoryType(memory_type) + if memory_type is MemoryType.PERSONAL: + personal_memory_target = self.service_context.personal_memory_target + if memory_target not in personal_memory_target: + personal_memory_target.append(memory_target) + else: + logger.warning(f"Memory target {memory_target} is already added.") + + elif memory_type is MemoryType.PROCEDURAL: + procedural_memory_target = self.service_context.procedural_memory_target + if memory_target not in procedural_memory_target: + procedural_memory_target.append(memory_target) + else: + logger.warning(f"Memory target {memory_target} is already added.") + + elif memory_type is MemoryType.TOOL: + tool_memory_target = self.service_context.tool_memory_target + if memory_target not in tool_memory_target: + tool_memory_target.append(memory_target) + else: + logger.warning(f"Memory target {memory_target} is already added.") - @property - def vector_store(self) -> BaseVectorStore: - """Return the default vector store instance from the service context.""" - return self.service_context.vector_stores["default"] - @property - def token_counter(self) -> BaseTokenCounter: - """Return the default token counter instance from the service context.""" - return self.service_context.token_counters["default"] async def summary_memory( self, messages: list[Message | dict], description: str = "", - user_name: str | list[str] = "", + user_name: str = "", + task_name: str = "", + tool_name: str = "", enable_thinking_params: bool = False, - meta_memories: list[dict] = None, version: str = "default", return_dict: bool = False, **kwargs, @@ -133,7 +176,7 @@ class ReMe: reme_summarizer = ReMeSummarizer( meta_memories=meta_memories, tools=[ - HandsOff( + DelegateTask( memory_agents=[ PersonalSummarizer( tools=[ @@ -202,7 +245,7 @@ class ReMe: reme_retriever = ReMeRetriever( meta_memories=meta_memories, tools=[ - HandsOff( + DelegateTask( memory_agents=[ PersonalRetriever( tools=[ @@ -341,84 +384,10 @@ class ReMe: """Retrieve all memories from the vector store.""" return [node.to_memory_node() for node in await self.vector_store.list()] - async def get_profiles(self, user_name: str | list[str]) -> str | list[str]: - """Retrieve user profile(s) from the system for the specified user(s).""" - read_profile = ReadUserProfile(show_id="profile") - if isinstance(user_name, str): - return await read_profile.call(memory_target=user_name, service_context=self.service_context) - else: - return [ - await read_profile.call(memory_target=name, service_context=self.service_context) for name in user_name - ] - - async def add_profile( - self, - profile_key: str, - profile_value: str, - user_name: str, - update_time: str | None = None, - ) -> MemoryNode: - """Add user profile to ReMe system.""" - update_user_profile = UpdateUserProfile() - if update_time is None: - update_time = get_now_time() - - await update_user_profile.call( - profile_ids_to_delete=[], - profiles_to_add=[ - { - "update_time": update_time, - "profile_key": profile_key, - "profile_value": profile_value, - }, - ], - memory_target=user_name, - service_context=self.service_context, - ) - return update_user_profile.memory_nodes[0] - - async def update_profile( - self, - profile_id: str, - profile_key: str, - profile_value: str, - user_name: str, - update_time: str | None = None, - ) -> MemoryNode: - """Add user profile to ReMe system.""" - update_user_profile = UpdateUserProfile() - if update_time is None: - update_time = get_now_time() - - await update_user_profile.call( - profile_ids_to_delete=[profile_id], - profiles_to_add=[ - { - "update_time": update_time, - "profile_key": profile_key, - "profile_value": profile_value, - }, - ], - memory_target=user_name, - service_context=self.service_context, - ) - return update_user_profile.memory_nodes[0] - - async def delete_all_profiles(self, user_name: str | list[str]): - """Delete all user profiles from ReMe system.""" - if isinstance(user_name, str): - user_name = [user_name] - - read_profile = ReadUserProfile(show_id="profile") - update_profile = UpdateUserProfile() - for memory_target in user_name: - await read_profile.call(memory_target=memory_target, service_context=self.service_context) - profile_ids = [profile.memory_id for profile in read_profile.memory_nodes] - await update_profile.call( - profile_ids_to_delete=profile_ids, - memory_target=memory_target, - service_context=self.service_context, - ) + def get_profile_handler(self, user_name: str) -> ProfileHandler: + """Get the profile handler for the specified user.""" + profile_path = Path(self.profile_path) / self.vector_store.collection_name + return ProfileHandler(memory_target=user_name, profile_path=profile_path) async def context_offload(self): """working memory summary""" @@ -451,6 +420,32 @@ class ReMe: """Run the configured service (HTTP, MCP, or CMD).""" self.service_context.service.run() + async def __aenter__(self): + """Async context manager entry.""" + return self + + def __enter__(self): + """Context manager entry.""" + return self + + async def close(self): + """Close""" + return await self.service_context.close() + + def close_sync(self): + """Close synchronously""" + self.service_context.close_sync() + + async def __aexit__(self, exc_type=None, exc_val=None, exc_tb=None): + """Async context manager exit.""" + await self.close() + return False + + def __exit__(self, exc_type=None, exc_val=None, exc_tb=None): + """Context manager exit.""" + self.close_sync() + return False + def main(): """Main entry point for running ReMe from command line.""" diff --git a/reme/tool/memory/__init__.py b/reme/tool/memory/__init__.py index 9c8bbd11..52e86d80 100644 --- a/reme/tool/memory/__init__.py +++ b/reme/tool/memory/__init__.py @@ -1,40 +1,34 @@ """memory tools""" +from .add_history import AddHistory +from .add_memory import AddMemory from .base_memory_tool import BaseMemoryTool -from .hands_off.hands_off import HandsOff -from .history.add_history import AddHistory -from .history.read_history import ReadHistory -from .identity.add_identity import AddIdentity -from .identity.read_identity import ReadIdentity -from .meta.add_meta_memory import AddMetaMemory -from .meta.read_meta_memory import ReadMetaMemory -from .user_profile.read_user_profile import ReadUserProfile -from .user_profile.update_user_profile import UpdateUserProfile -from .vector.add_memory import AddMemory -from .vector.delete_memory import DeleteMemory -from .vector.retrieve_memory import RetrieveMemory -from .vector.retrieve_recent_memory import RetrieveRecentMemory -from .vector.update_memory import UpdateMemory +from .delegate_task import DelegateTask +from .delete_memory import DeleteMemory +from .profile_handler import ProfileHandler +from .read_history import ReadHistory +from .read_profile import ReadProfile +from .retrieve_memory import RetrieveMemory +from .retrieve_recent_memory import RetrieveRecentMemory +from .update_memory import UpdateMemory +from .update_profile import UpdateProfile from ...core import R __all__ = [ - "BaseMemoryTool", - "HandsOff", "AddHistory", - "ReadHistory", - "AddIdentity", - "ReadIdentity", - "AddMetaMemory", - "ReadMetaMemory", - "ReadUserProfile", - "UpdateUserProfile", "AddMemory", + "BaseMemoryTool", + "DelegateTask", "DeleteMemory", + "ProfileHandler", + "ReadHistory", + "ReadProfile", "RetrieveMemory", "RetrieveRecentMemory", "UpdateMemory", + "UpdateProfile", ] for name in __all__: tool_class = globals()[name] - R.op.register()(tool_class) + R.op.register()(tool_class) \ No newline at end of file diff --git a/reme/tool/memory/history/add_history.py b/reme/tool/memory/add_history.py similarity index 77% rename from reme/tool/memory/history/add_history.py rename to reme/tool/memory/add_history.py index d9527ff8..57a0097e 100644 --- a/reme/tool/memory/history/add_history.py +++ b/reme/tool/memory/add_history.py @@ -2,10 +2,10 @@ from loguru import logger -from ..base_memory_tool import BaseMemoryTool -from ....core.enumeration import MemoryType -from ....core.schema import ToolCall, MemoryNode, Message -from ....core.utils import format_messages +from .base_memory_tool import BaseMemoryTool +from ...core.enumeration import MemoryType +from ...core.schema import ToolCall, MemoryNode, Message +from ...core.utils import format_messages class AddHistory(BaseMemoryTool): @@ -31,10 +31,11 @@ class AddHistory(BaseMemoryTool): async def execute(self): """Execute the add history operation""" self.context.messages = [Message(**m) if isinstance(m, dict) else m for m in self.context.messages] - history_content: str = (self.context.description + "\n" + format_messages(self.context.messages)).strip() + history_content: str = self.context.description + "\n" + format_messages(self.context.messages) + history_content = history_content.strip() history_node = MemoryNode( memory_type=MemoryType.HISTORY, - when_to_use=history_content[:100], + when_to_use=history_content[:1024], content=history_content, author=self.author, ) diff --git a/reme/tool/memory/add_memory.py b/reme/tool/memory/add_memory.py new file mode 100644 index 00000000..8ecdf965 --- /dev/null +++ b/reme/tool/memory/add_memory.py @@ -0,0 +1,135 @@ +"""Add memory to vector store""" + +from loguru import logger + +from .base_memory_tool import BaseMemoryTool +from .memory_handler import MemoryHandler +from ...core.schema import ToolCall + + +class AddMemory(BaseMemoryTool): + """Tool to add memories to vector store""" + + def __init__( + self, + enable_memory_target: bool = False, + enable_when_to_use: bool = False, + **kwargs, + ): + super().__init__(**kwargs) + self.enable_memory_target: bool = enable_memory_target + self.enable_when_to_use: bool = enable_when_to_use + + def _build_memory_parameters(self) -> dict: + """Build the memory parameters schema based on enabled features.""" + properties = { + "message_time": { + "type": "string", + "description": "message time, e.g. '2020-01-01 00:00:00'", + }, + "memory_content": { + "type": "string", + "description": "content of the memory.", + }, + } + required = ["message_time", "memory_content"] + + if self.enable_when_to_use: + properties["when_to_use"] = { + "type": "string", + "description": "description of when to use this memory.", + } + required.append("when_to_use") + + if self.enable_memory_target: + properties["memory_target"] = { + "type": "string", + "description": "target memory type for this memory.", + } + required.append("memory_target") + + return { + "type": "object", + "properties": properties, + "required": required, + } + + def _build_tool_call(self) -> ToolCall: + return ToolCall( + **{ + "description": "add a memory to vector store for future retrieval.", + "parameters": self._build_memory_parameters(), + }, + ) + + def _build_multiple_tool_call(self) -> ToolCall: + return ToolCall( + **{ + "description": "add multiple memories to vector store for future retrieval.", + "parameters": { + "type": "object", + "properties": { + "memories": { + "type": "array", + "description": "list of memories to store.", + "items": self._build_memory_parameters(), + }, + }, + "required": ["memories"], + }, + }, + ) + + async def execute(self): + if self.enable_multiple: + memories = self.context.get("memories", []) + else: + memories = [self.context] + + # Group memories by memory_target if enabled + if self.enable_memory_target: + memories_by_target = {} + for mem in memories: + target = mem["memory_target"] + if target not in memories_by_target: + memories_by_target[target] = [] + memories_by_target[target].append(mem) + else: + memories_by_target = {self.memory_target: memories} + + # Process each memory_target group + all_memory_nodes = [] + for target, target_memories in memories_by_target.items(): + # Parse and prepare memory data + memory_dicts = [] + for mem in target_memories: + memory_content = mem.get("memory_content", "") + message_time = mem.get("message_time", "") + when_to_use = mem.get("when_to_use", "") if self.enable_when_to_use else "" + metadata = {} + try: + metadata["time_int"] = int(message_time.split(" ")[0].replace("-", "")) + except Exception: + logger.warning(f"Invalid message time format: {message_time}") + + memory_dicts.append({ + "content": memory_content, + "when_to_use": when_to_use, + "message_time": message_time, + "ref_memory_id": self.history_node.memory_id, + "author": self.author, + "metadata": metadata, + }) + + if memory_dicts: + handler = MemoryHandler(target, self.service_context) + memory_nodes = await handler.add_batch(memory_dicts) + all_memory_nodes.extend(memory_nodes) + + if not all_memory_nodes: + return "No valid memories provided." + + self.memory_nodes.extend(all_memory_nodes) + output = f"Successfully added {len(all_memory_nodes)} memories." + logger.info(output) + return output diff --git a/reme/tool/memory/base_memory_tool.py b/reme/tool/memory/base_memory_tool.py index 5a900177..0f253b98 100644 --- a/reme/tool/memory/base_memory_tool.py +++ b/reme/tool/memory/base_memory_tool.py @@ -1,12 +1,10 @@ """Base class for memory tool""" from abc import ABCMeta -from pathlib import Path from ...core.enumeration import MemoryType from ...core.op import BaseTool from ...core.schema import ToolCall, MemoryNode, ToolAttr -from ...core.utils import CacheHandler class BaseMemoryTool(BaseTool, metaclass=ABCMeta): @@ -16,13 +14,11 @@ class BaseMemoryTool(BaseTool, metaclass=ABCMeta): self, enable_multiple: bool = True, enable_thinking_params: bool = False, - local_memory_path: str = "./reme_local_memory", **kwargs, ): super().__init__(**kwargs) self.enable_multiple: bool = enable_multiple self.enable_thinking_params: bool = enable_thinking_params - self.local_memory_path: str = local_memory_path def _build_tool_call(self) -> ToolCall: """Build and return the tool call schema""" @@ -59,44 +55,44 @@ class BaseMemoryTool(BaseTool, metaclass=ABCMeta): parameters.required = ["thinking"] return self._tool_call - @property - def local_memory(self) -> CacheHandler: - """Create the meta memory cache handler.""" - return CacheHandler(Path(self.local_memory_path) / self.vector_store.collection_name) - @property def memory_type(self) -> MemoryType: """Get the memory type from context.""" - return MemoryType(self.context.get("memory_type")) + return self.memory_target_type_mapping[self.memory_target] @property def memory_target(self) -> str: """Get the memory target from context.""" - return self.context.get("memory_target", "") - - @property - def memory_cache_key(self) -> str: - """Get the memory cache key from context.""" - return f"{self.memory_type.value}_{self.memory_target}".replace(" ", "_").lower() + if "memory_target" in self.context: + return self.context.memory_target + elif len(self.memory_target_type_mapping) == 1: + return list(self.memory_target_type_mapping.keys())[0] + else: + raise ValueError("memory_target is not specified in context or memory_target_type_mapping!") @property def history_node(self) -> MemoryNode: """Get the history node from context.""" - return self.context.get("history_node") + return self.context.history_node @property def retrieved_nodes(self) -> list[MemoryNode]: """Get the retrieved nodes from context.""" - return self.context["retrieved_nodes"] + return self.context.retrieved_nodes @property def author(self) -> str: """Get the author from context.""" - return self.context.get("author", "") + return self.context.author @property def memory_nodes(self) -> list[MemoryNode | str]: """Get the memory nodes from context.""" if "memory_nodes" not in self.context: self.context.memory_nodes = [] - return self.context["memory_nodes"] + return self.context.memory_nodes + + @property + def memory_target_type_mapping(self) -> dict[str, MemoryType]: + """Get the memory target type mapping from context.""" + return self.context.memory_target_type_mapping diff --git a/reme/tool/memory/delegate_task.py b/reme/tool/memory/delegate_task.py new file mode 100644 index 00000000..13c1ff07 --- /dev/null +++ b/reme/tool/memory/delegate_task.py @@ -0,0 +1,83 @@ +"""Hands-off tool to delegate memory tasks to specific agents""" + +from loguru import logger + +from .base_memory_tool import BaseMemoryTool +from ...agent.memory import BaseMemoryAgent +from ...core.enumeration import MemoryType +from ...core.schema import ToolCall + + +class DelegateTask(BaseMemoryTool): + """Tool to delegate memory tasks to appropriate memory agents""" + + def __init__(self, memory_agents: list[BaseMemoryAgent] = None, **kwargs): + kwargs["enable_multiple"] = True + kwargs["sub_ops"] = memory_agents or [] + super().__init__(**kwargs) + self.sub_ops: list[BaseMemoryAgent] = [a for a in self.sub_ops if isinstance(a, BaseMemoryAgent)] + assert all(a.memory_type is not None for a in self.sub_ops) + + @property + def memory_agent_dict(self) -> dict[MemoryType, BaseMemoryAgent]: + """Map memory types to their corresponding agents""" + return {a.memory_type: a for a in self.sub_ops} + + def _build_multiple_tool_call(self) -> ToolCall: + """Build and return the multiple tool call schema""" + return ToolCall( + **{ + "description": "Delegate tasks to appropriate agents.", + "parameters": { + "type": "object", + "properties": { + "tasks": { + "type": "array", + "description": "tasks to delegate to specific agents", + "items": { + "type": "object", + "properties": { + "task_name": { + "type": "string", + "description": "task_name", + }, + }, + "required": ["task_name"], + }, + }, + }, + "required": ["tasks"], + }, + }, + ) + + async def execute(self): + # Deduplicate and validate tasks + tasks = self.context.get("tasks", []) + tasks = sorted(set(tasks)) + + # Submit tasks to agents + agent_list: list[BaseMemoryAgent] = [] + for i, task in enumerate(tasks): + memory_type = self.memory_target_type_mapping[task] + agent = self.memory_agent_dict[memory_type].copy() + agent_list.append(agent) + + logger.info(f"Task {i}: {memory_type.value} agent for {task}") + task_kwargs = {"memory_target": task} + for k in ["query", "messages", "description", "history_node"]: + if k in self.context: + task_kwargs[k] = self.context[k] + self.submit_async_task(agent.call, service_context=self.service_context, **task_kwargs) + await self.join_async_tasks() + + # Collect results + results = [] + for agent in agent_list: + results.append(f"Task: {agent.memory_target}\n{agent.response.answer}") + + logger.info(f"Completed {len(results)} task(s)") + return { + "answer": "\n\n".join(results), + "agents": agent_list, + } diff --git a/reme/tool/memory/vector/delete_memory.py b/reme/tool/memory/delete_memory.py similarity index 59% rename from reme/tool/memory/vector/delete_memory.py rename to reme/tool/memory/delete_memory.py index 94a22a75..e95cfa40 100644 --- a/reme/tool/memory/vector/delete_memory.py +++ b/reme/tool/memory/delete_memory.py @@ -2,15 +2,15 @@ from loguru import logger -from ..base_memory_tool import BaseMemoryTool -from ....core.schema import ToolCall +from .base_memory_tool import BaseMemoryTool +from .memory_handler import MemoryHandler +from ...core.schema import ToolCall class DeleteMemory(BaseMemoryTool): """Tool to delete memories from vector store""" def _build_tool_call(self) -> ToolCall: - """Build and return the single tool call schema""" return ToolCall( **{ "description": "delete a memory from vector store using its unique ID.", @@ -19,7 +19,7 @@ class DeleteMemory(BaseMemoryTool): "properties": { "memory_id": { "type": "string", - "description": "unique identifier (memory_id) of the memory to delete.", + "description": "memory_id of the memory to delete.", }, }, "required": ["memory_id"], @@ -28,7 +28,6 @@ class DeleteMemory(BaseMemoryTool): ) def _build_multiple_tool_call(self) -> ToolCall: - """Build and return the multiple tool call schema""" return ToolCall( **{ "description": "delete multiple memories from vector store using their unique IDs.", @@ -37,7 +36,7 @@ class DeleteMemory(BaseMemoryTool): "properties": { "memory_ids": { "type": "array", - "description": "list of unique identifiers (memory_ids) of memories to delete.", + "description": "memory_ids of memories to delete.", "items": {"type": "string"}, }, }, @@ -47,25 +46,14 @@ class DeleteMemory(BaseMemoryTool): ) async def execute(self): - memory_ids: list[str] = [] - - # Handle multiple memories (array format) - ids_from_array = self.context.get("memory_ids", []) - if ids_from_array: - memory_ids = [m for m in ids_from_array if m] - else: - memory_id = self.context.get("memory_id", "") - if memory_id: - memory_ids = [memory_id] - + memory_ids = self.context.get("memory_ids") or [] if not memory_ids: - output = "No valid memory IDs provided for deletion." - logger.info(output) - return output + memory_ids = [self.context.get("memory_id", "")] - await self.vector_store.delete(vector_ids=list(set(memory_ids))) + handler = MemoryHandler(self.memory_target, self.service_context) + await handler.delete(memory_ids) self.memory_nodes.extend(memory_ids) - output = f"Successfully deleted {len(memory_ids)} memories from vector_store." + output = f"Successfully deleted {len(memory_ids)} memories." logger.info(output) return output diff --git a/reme/tool/memory/hands_off/__init__.py b/reme/tool/memory/hands_off/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/reme/tool/memory/hands_off/hands_off.py b/reme/tool/memory/hands_off/hands_off.py deleted file mode 100644 index 45822381..00000000 --- a/reme/tool/memory/hands_off/hands_off.py +++ /dev/null @@ -1,108 +0,0 @@ -"""Hands-off tool to delegate memory tasks to specific agents""" - -from loguru import logger - -from ..base_memory_tool import BaseMemoryTool -from ....agent.memory import BaseMemoryAgent -from ....core.enumeration import MemoryType -from ....core.schema import ToolCall - - -class HandsOff(BaseMemoryTool): - """Tool to delegate memory tasks to appropriate memory agents""" - - def __init__(self, memory_agents: list[BaseMemoryAgent] = None, **kwargs): - kwargs["enable_multiple"] = True - kwargs["sub_ops"] = memory_agents or [] - super().__init__(**kwargs) - self.sub_ops: list[BaseMemoryAgent] = [ - a for a in self.sub_ops if isinstance(a, BaseMemoryAgent) and a.memory_type is not None - ] - - @property - def memory_agent_dict(self) -> dict[MemoryType, BaseMemoryAgent]: - """Map memory types to their corresponding agents""" - return {a.memory_type: a for a in self.sub_ops} - - def _build_multiple_tool_call(self) -> ToolCall: - """Build and return the multiple tool call schema""" - return ToolCall( - **{ - "description": "Delegate memory tasks to appropriate agents.", - "parameters": { - "type": "object", - "properties": { - "memory_tasks": { - "type": "array", - "description": "Memory tasks to delegate to specific agents", - "items": { - "type": "object", - "properties": { - "memory_type": { - "type": "string", - "description": "Memory type to handle", - "enum": [k.value for k in self.memory_agent_dict if k], - }, - "memory_target": { - "type": "string", - "description": "Target or context for the memory operation", - }, - }, - "required": ["memory_type", "memory_target"], - }, - }, - }, - "required": ["memory_tasks"], - }, - }, - ) - - async def execute(self): - # Deduplicate and validate tasks - tasks = [] - seen = set() - for task in self.context.get("memory_tasks", []): - memory_type = MemoryType(task.get("memory_type", "")) - memory_target = task.get("memory_target", "") - - task_key = (memory_type, memory_target) - if task_key in seen: - logger.info(f"Skip duplicate: {memory_type.value} - {memory_target}") - continue - seen.add(task_key) - - tasks.append({"memory_type": memory_type, "memory_target": memory_target}) - - if not tasks: - return "No valid memory tasks to execute." - - # Submit tasks to agents - agent_list: list[BaseMemoryAgent] = [] - for i, task in enumerate(tasks): - memory_type: MemoryType = task["memory_type"] - memory_target: str = task["memory_target"] - - agent = self.memory_agent_dict[memory_type].copy() - agent_list.append(agent) - - logger.info(f"Task {i}: {memory_type.value} agent for {memory_target}") - task_kwargs = {"memory_type": memory_type, "memory_target": memory_target} - for k in ["query", "messages", "description", "history_node"]: - if k in self.context: - task_kwargs[k] = self.context[k] - self.submit_async_task(agent.call, service_context=self.service_context, **task_kwargs) - - await self.join_async_tasks() - - # Collect results - results = [] - for agent in agent_list: - memory_type = agent.memory_type - memory_target = agent.memory_target - results.append(f"{memory_type.value}({memory_target}): {agent.response.answer}") - - logger.info(f"Completed {len(results)} task(s)") - return { - "answer": "\n".join(results), - "agents": agent_list, - } diff --git a/reme/tool/memory/history/__init__.py b/reme/tool/memory/history/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/reme/tool/memory/identity/__init__.py b/reme/tool/memory/identity/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/reme/tool/memory/identity/add_identity.py b/reme/tool/memory/identity/add_identity.py deleted file mode 100644 index 5f6ab984..00000000 --- a/reme/tool/memory/identity/add_identity.py +++ /dev/null @@ -1,42 +0,0 @@ -"""Add identity memory tool""" - -from loguru import logger - -from ..base_memory_tool import BaseMemoryTool -from ....core.schema import ToolCall - - -class AddIdentity(BaseMemoryTool): - """Tool to add or update agent identity memory""" - - def __init__(self, **kwargs): - kwargs["enable_multiple"] = False - super().__init__(**kwargs) - - def _build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "description": "add or update agent identity memory.", - "parameters": { - "type": "object", - "properties": { - "identity_memory": { - "type": "string", - "description": "Agent identity content, such as role, personality, or current state.", - }, - }, - "required": ["identity_memory"], - }, - }, - ) - - async def execute(self): - identity_memory = self.context.get("identity_memory", "") - - if not identity_memory: - logger.warning("No valid identity memory provided") - return "No valid identity memory provided for update." - - self.local_memory.save("identity_memory", identity_memory) - logger.info(f"Successfully updated identity memory: {identity_memory}") - return "Successfully updated identity memory." diff --git a/reme/tool/memory/identity/read_identity.py b/reme/tool/memory/identity/read_identity.py deleted file mode 100644 index 856aea55..00000000 --- a/reme/tool/memory/identity/read_identity.py +++ /dev/null @@ -1,36 +0,0 @@ -"""Read identity memory tool""" - -from loguru import logger - -from ..base_memory_tool import BaseMemoryTool -from ....core.schema import ToolCall - - -class ReadIdentity(BaseMemoryTool): - """Tool to read agent identity memory""" - - def __init__(self, **kwargs): - kwargs["enable_multiple"] = False - super().__init__(**kwargs) - - def _build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "description": "read agent identity memory.", - "parameters": { - "type": "object", - "properties": {}, - "required": [], - }, - }, - ) - - async def execute(self): - identity_memory = self.local_memory.load("identity_memory") - - if not identity_memory: - logger.info("No identity memory found") - return "No identity memory found." - - logger.info(f"Read identity memory: {identity_memory}") - return f"Identity\n{identity_memory}" diff --git a/reme/tool/memory/memory_handler.py b/reme/tool/memory/memory_handler.py new file mode 100644 index 00000000..78151325 --- /dev/null +++ b/reme/tool/memory/memory_handler.py @@ -0,0 +1,210 @@ +from ...core.context import ServiceContext +from ...core.enumeration import MemoryType +from ...core.schema import MemoryNode +from ...core.vector_store import BaseVectorStore + + +class MemoryHandler: + """Handler for managing memory nodes in the vector store.""" + + def __init__(self, memory_target: str, service_context: ServiceContext): + self.memory_target: str = memory_target + self.memory_type: MemoryType = service_context.memory_target_type_mapping[memory_target] + self.vector_store: BaseVectorStore = service_context.vector_stores["default"] + + async def add_batch(self, memories: list[dict]) -> list[MemoryNode]: + """Add multiple memory nodes and return their memory_ids.""" + # First, delete existing memory nodes if memory_ids are provided + memory_ids_to_delete = [mem.get("memory_id") for mem in memories if mem.get("memory_id")] + if memory_ids_to_delete: + await self.vector_store.delete(memory_ids_to_delete) + + # Create MemoryNode objects + memory_nodes = [ + MemoryNode( + memory_type=self.memory_type, + memory_target=self.memory_target, + content=mem.get("content", ""), + when_to_use=mem.get("when_to_use", ""), + message_time=mem.get("message_time", ""), + ref_memory_id=mem.get("ref_memory_id", ""), + author=mem.get("author", ""), + score=mem.get("score", 0.0), + metadata=mem.get("metadata", {}), + ) + for mem in memories + ] + + # Deduplicate memory_nodes by content (keep last occurrence) + memory_dict = {node.content: node for node in memory_nodes} + memory_nodes = list(memory_dict.values()) + + # Convert to VectorNodes and insert + vector_nodes = [node.to_vector_node() for node in memory_nodes] + await self.vector_store.insert(vector_nodes) + return memory_nodes + + async def add( + self, + content: str, + when_to_use: str = "", + message_time: str = "", + ref_memory_id: str = "", + author: str = "", + score: float = 0.0, + **kwargs, + ) -> MemoryNode: + """Add a single memory node and return its memory_id.""" + memory_dict = { + "content": content, + "when_to_use": when_to_use, + "message_time": message_time, + "ref_memory_id": ref_memory_id, + "author": author, + "score": score, + "metadata": kwargs, + } + memory_nodes = await self.add_batch([memory_dict]) + return memory_nodes[0] + + async def delete(self, memory_ids: str | list[str]): + """Delete multiple memory nodes by their memory_ids.""" + # Deduplicate if input is a list + if isinstance(memory_ids, list): + memory_ids = list(dict.fromkeys(memory_ids)) + await self.vector_store.delete(memory_ids) + + async def delete_all(self): + """Delete all memory nodes.""" + await self.vector_store.delete_all() + + async def update_batch(self, updates: list[dict]) -> list[MemoryNode]: + """Update multiple memory nodes with their memory_ids and new values using delete + add.""" + # Deduplicate updates by memory_id (keep last occurrence) + updates_dict = {upd["memory_id"]: upd for upd in updates} + updates = list(updates_dict.values()) + memory_ids = list(updates_dict.keys()) + + # Get existing nodes + vector_nodes = await self.vector_store.get(memory_ids) + if not isinstance(vector_nodes, list): + vector_nodes = [vector_nodes] + + # Update and convert back + updated_nodes: list[MemoryNode] = [] + for vector_node, update in zip(vector_nodes, updates): + memory_node = MemoryNode.from_vector_node(vector_node) + memory_node.memory_target = self.memory_target + memory_node.memory_type = self.memory_type + + if "content" in update: + memory_node.content = update["content"] + if "when_to_use" in update: + memory_node.when_to_use = update["when_to_use"] + if "message_time" in update: + memory_node.message_time = update["message_time"] + if "ref_memory_id" in update: + memory_node.ref_memory_id = update["ref_memory_id"] + if "author" in update: + memory_node.author = update["author"] + if "score" in update: + memory_node.score = update["score"] + if "metadata" in update: + memory_node.metadata.update(update["metadata"]) + updated_nodes.append(memory_node) + + # Delete old nodes first + await self.vector_store.delete(memory_ids) + + # Then add updated nodes + vector_nodes = [node.to_vector_node() for node in updated_nodes] + await self.vector_store.insert(vector_nodes) + + return updated_nodes + + async def update( + self, + memory_id: str, + content: str | None = None, + when_to_use: str | None = None, + message_time: str | None = None, + ref_memory_id: str | None = None, + author: str | None = None, + score: float | None = None, + **kwargs, + ) -> MemoryNode: + """Update a memory node's content, when_to_use, or other fields.""" + update_dict: dict = {"memory_id": memory_id} + if content is not None: + update_dict["content"] = content + if when_to_use is not None: + update_dict["when_to_use"] = when_to_use + if message_time is not None: + update_dict["message_time"] = message_time + if ref_memory_id is not None: + update_dict["ref_memory_id"] = ref_memory_id + if author is not None: + update_dict["author"] = author + if score is not None: + update_dict["score"] = score + if kwargs is not None: + update_dict["metadata"] = kwargs + + memory_nodes = await self.update_batch([update_dict]) + return memory_nodes[0] + + async def search( + self, + query: str | list[str], + limit: int = 5, + filters: dict | None = None, + **kwargs, + ) -> list[MemoryNode]: + """Search for similar memory nodes based on query text.""" + filters = filters or {} + filters["memory_type"] = self.memory_type.value + filters["memory_target"] = self.memory_target + + # Handle single query + if isinstance(query, str): + vector_nodes = await self.vector_store.search(query, limit=limit, filters=filters, **kwargs) + return [MemoryNode.from_vector_node(node) for node in vector_nodes] + + # Handle multiple queries: search each query with the same limit + seen_ids: dict[str, MemoryNode] = {} + + for q in query: + vector_nodes = await self.vector_store.search(q, limit=limit, filters=filters, **kwargs) + for vector_node in vector_nodes: + memory_node = MemoryNode.from_vector_node(vector_node) + if memory_node.memory_id not in seen_ids: + seen_ids[memory_node.memory_id] = memory_node + + return list(seen_ids.values()) + + async def batch_search(self, searches: list[dict]) -> list[MemoryNode]: + """Execute multiple search queries in batch and return deduplicated results.""" + seen_ids: dict[str, MemoryNode] = {} + + for search_params in searches: + search_result = await self.search(**search_params) + for memory_node in search_result: + if memory_node.memory_id not in seen_ids: + seen_ids[memory_node.memory_id] = memory_node + + return list(seen_ids.values()) + + async def list( + self, + filters: dict | None = None, + limit: int | None = None, + sort_key: str | None = None, + reverse: bool = True, + ) -> list[MemoryNode]: + """List memory nodes with optional filtering and sorting.""" + filters = filters or {} + filters["memory_type"] = self.memory_type.value + filters["memory_target"] = self.memory_target + + vector_nodes = await self.vector_store.list(filters=filters, limit=limit, sort_key=sort_key, reverse=reverse) + return [MemoryNode.from_vector_node(node) for node in vector_nodes] diff --git a/reme/tool/memory/meta/__init__.py b/reme/tool/memory/meta/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/reme/tool/memory/meta/add_meta_memory.py b/reme/tool/memory/meta/add_meta_memory.py deleted file mode 100644 index 43e9f7db..00000000 --- a/reme/tool/memory/meta/add_meta_memory.py +++ /dev/null @@ -1,89 +0,0 @@ -"""Add meta memory tool""" - -import json - -from loguru import logger - -from ..base_memory_tool import BaseMemoryTool -from ....core.enumeration import MemoryType -from ....core.schema import ToolCall - - -class AddMetaMemory(BaseMemoryTool): - """Tool to add memory metadata entries to meta storage""" - - def __init__(self, **kwargs): - kwargs["enable_multiple"] = True - super().__init__(**kwargs) - - def _build_multiple_tool_call(self) -> ToolCall: - """Build and return the multiple tool call schema""" - return ToolCall( - **{ - "description": "add memory metadata entries to register memory types and targets. " - "Before using, verify Main Agent's Meta Memory doesn't already contain the " - "same memory_type(memory_target) combinations.", - "parameters": { - "type": "object", - "properties": { - "meta_memories": { - "type": "array", - "description": "List of memory metadata entries to add", - "items": { - "type": "object", - "properties": { - "memory_type": { - "type": "string", - "description": "Type of memory: 'personal' for person-specific preferences, " - "'procedural' for how-to knowledge", - "enum": [MemoryType.PERSONAL.value, MemoryType.PROCEDURAL.value], - }, - "memory_target": { - "type": "string", - "description": "Target identifier, " - "e.g., person's name ('John') or domain ('deployment')", - }, - }, - "required": ["memory_type", "memory_target"], - }, - }, - }, - "required": ["meta_memories"], - }, - }, - ) - - async def execute(self): - existing_memories: list[dict] = self.local_memory.load("meta_memories") or [] - existing_set = {(m["memory_type"], m["memory_target"]) for m in existing_memories} - - # Filter and build new memories to add - new_memories: list[dict] = [] - meta_memories: list[dict] = self.context.get("meta_memories", []) - - for mem in meta_memories: - memory_type = mem.get("memory_type", "") - memory_target = mem.get("memory_target", "") - - # Check if valid and not duplicate - if ( - memory_type in [MemoryType.PERSONAL.value, MemoryType.PROCEDURAL.value] - and memory_target - and (memory_type, memory_target) not in existing_set - ): - new_memories.append({"memory_type": memory_type, "memory_target": memory_target}) - existing_set.add((memory_type, memory_target)) - - if not new_memories: - output = "No new meta memories to add (all entries already exist or invalid)." - logger.info(output) - return output - - # Merge, sort and save - all_memories = sorted(existing_memories + new_memories, key=lambda m: (m["memory_type"], m["memory_target"])) - self.local_memory.save("meta_memories", all_memories) - - # Format output - output = f"Successfully update meta memory entries: {json.dumps(new_memories, ensure_ascii=False)}" - logger.info(output) - return output diff --git a/reme/tool/memory/meta/read_meta_memory.py b/reme/tool/memory/meta/read_meta_memory.py deleted file mode 100644 index f1936394..00000000 --- a/reme/tool/memory/meta/read_meta_memory.py +++ /dev/null @@ -1,75 +0,0 @@ -"""Read meta memory tool""" - -from loguru import logger - -from ..base_memory_tool import BaseMemoryTool -from ....core.enumeration import MemoryType -from ....core.schema import ToolCall - - -class ReadMetaMemory(BaseMemoryTool): - """Tool to read memory metadata from meta storage""" - - TYPE_DESC_DICT = { - MemoryType.IDENTITY.value: "self-cognition memory storing agent's identity and state", - MemoryType.PERSONAL.value: "person-specific memory storing preferences and context", - MemoryType.PROCEDURAL.value: "procedural memory storing how-to knowledge and processes", - } - - def __init__(self, enable_identity_memory: bool = False, **kwargs): - kwargs["enable_multiple"] = False - super().__init__(**kwargs) - self.enable_identity_memory = enable_identity_memory - - def _build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "description": "read memory metadata registry to see what types of memories are being tracked.", - "parameters": { - "type": "object", - "properties": {}, - "required": [], - }, - }, - ) - - def format_memory_metadata(self, memories: list[dict[str, str]]) -> str: - """Format memory metadata into a readable string.""" - if not memories: - return "" - - lines = [] - for memory in memories: - memory_type = memory["memory_type"] - memory_target = memory["memory_target"] - description = self.TYPE_DESC_DICT[memory_type] - lines.append(f"- {memory_type}({memory_target}): {description}") - - return "\n".join(lines) - - async def execute(self): - # Load and filter meta memories - result = self.local_memory.load("meta_memories") - all_memories = result if result is not None else [] - - memories = [ - m for m in all_memories if m.get("memory_type") in [MemoryType.PERSONAL.value, MemoryType.PROCEDURAL.value] - ] - - if self.enable_identity_memory: - memories.append( - { - "memory_type": MemoryType.IDENTITY.value, - "memory_target": "self", - }, - ) - - # Format output - output = self.format_memory_metadata(memories) - if output: - logger.info(f"Retrieved {len(memories)} meta memory entries") - else: - output = "No memory metadata found." - logger.info(output) - - return output diff --git a/reme/tool/memory/profile_handler.py b/reme/tool/memory/profile_handler.py new file mode 100644 index 00000000..4a950171 --- /dev/null +++ b/reme/tool/memory/profile_handler.py @@ -0,0 +1,199 @@ +"""Profile Handler for managing user profiles in local memory""" +from pathlib import Path + +from loguru import logger + +from ...core.enumeration import MemoryType +from ...core.schema import MemoryNode +from ...core.utils import CacheHandler, deduplicate_memories + + +class ProfileHandler: + """User profile CRUD handler""" + + def __init__(self, profile_path: str | Path, memory_target: str, max_capacity: int = 100): + """init""" + self.memory_target: str = memory_target + self.cache_key: str = self.memory_target.replace(" ", "_").lower() + self.cache_handler: CacheHandler = CacheHandler(profile_path) + self.max_capacity: int = max_capacity + + def _load_nodes(self) -> list[MemoryNode]: + """Load profile nodes""" + cached_data = self.cache_handler.load(self.cache_key, auto_clean=False) + if not cached_data: + return [] + return [MemoryNode(**data) for data in cached_data] + + def _save_nodes(self, nodes: list[MemoryNode], apply_limits: bool = True): + """Save nodes with optional deduplication and capacity enforcement""" + if apply_limits: + nodes = deduplicate_memories(nodes) + + # Enforce capacity limit by removing the oldest profiles + if len(nodes) > self.max_capacity: + sorted_nodes = sorted(nodes, key=lambda n: n.message_time) + removed_count = len(sorted_nodes) - self.max_capacity + nodes = sorted_nodes[removed_count:] + logger.info( + f"Capacity limit reached: removed {removed_count} oldest profiles (kept {len(nodes)}/{self.max_capacity})") + + nodes_data = [node.model_dump(exclude_none=True) for node in nodes] + self.cache_handler.save(self.cache_key, nodes_data) + logger.info(f"Saved {len(nodes)} profiles to {self.cache_key}") + + def delete(self, profile_id: str | list[str]) -> bool | int: + """Delete profile by ID(s), returns True/False for single ID or count for batch delete""" + nodes = self._load_nodes() + original_count = len(nodes) + + # Batch delete mode + if isinstance(profile_id, list): + profile_ids_set = set(profile_id) + nodes = [n for n in nodes if n.memory_id not in profile_ids_set] + deleted_count = original_count - len(nodes) + + if deleted_count == 0: + logger.warning(f"No profiles found to delete from {len(profile_id)} IDs") + return 0 + + self._save_nodes(nodes, apply_limits=False) + logger.info(f"Batch deleted {deleted_count} profiles") + return deleted_count + + # Single delete mode + nodes = [n for n in nodes if n.memory_id != profile_id] + + if len(nodes) == original_count: + logger.warning(f"Profile {profile_id} not found") + return False + + self._save_nodes(nodes, apply_limits=False) + logger.info(f"Deleted profile {profile_id}") + return True + + def delete_all(self) -> int: + """Delete all profiles, returns count deleted""" + nodes = self._load_nodes() + count = len(nodes) + self._save_nodes([], apply_limits=False) + logger.info(f"Deleted all {count} profiles") + return count + + def add(self, message_time: str, profile_key: str, profile_value: str, ref_memory_id: str = "") -> MemoryNode: + """Add new profile, returns created MemoryNode""" + nodes = self._load_nodes() + + new_node = MemoryNode( + memory_type=MemoryType.PERSONAL, + memory_target=self.memory_target, + when_to_use=profile_key, + content=profile_value, + message_time=message_time, + ref_memory_id=ref_memory_id, + ) + + nodes.append(new_node) + self._save_nodes(nodes) + logger.info(f"Added profile: {profile_key}={profile_value}") + return new_node + + def add_batch(self, profiles: list[dict], ref_memory_id: str = "") -> list[MemoryNode]: + """Add multiple profiles in batch, returns list of created MemoryNodes""" + if not profiles: + return [] + + nodes = self._load_nodes() + + new_nodes = [ + MemoryNode( + memory_type=MemoryType.PERSONAL, + memory_target=self.memory_target, + when_to_use=p.get("profile_key", ""), + content=p.get("profile_value", ""), + message_time=p.get("message_time", ""), + ref_memory_id=ref_memory_id, + ) + for p in profiles + ] + + nodes.extend(new_nodes) + self._save_nodes(nodes) + logger.info(f"Batch added {len(new_nodes)} profiles") + return new_nodes + + def update(self, profile_id: str, message_time: str, profile_key: str, profile_value: str) -> MemoryNode | None: + """Update profile by ID, returns updated node or None if not found""" + nodes = self._load_nodes() + + target_node = None + for node in nodes: + if node.memory_id == profile_id: + node.when_to_use = profile_key + node.content = profile_value + node.message_time = message_time + target_node = node + break + + if target_node is None: + logger.warning(f"Profile {profile_id} not found") + return None + + self._save_nodes(nodes, apply_limits=False) + logger.info(f"Updated profile {profile_id}: {profile_key}={profile_value}") + return target_node + + def get_by(self, *, profile_id: str | None = None, profile_key: str | None = None) -> MemoryNode | None: + """Get profile by ID or key""" + if not profile_id and not profile_key: + raise ValueError("Must provide either profile_id or profile_key") + + nodes = self._load_nodes() + for node in nodes: + if profile_id and node.memory_id == profile_id: + return node + if profile_key and node.when_to_use == profile_key: + return node + return None + + def get_by_id(self, profile_id: str) -> MemoryNode | None: + """Get profile by ID (convenience method)""" + return self.get_by(profile_id=profile_id) + + def get_by_key(self, profile_key: str) -> MemoryNode | None: + """Get profile by key (convenience method)""" + return self.get_by(profile_key=profile_key) + + def get_all(self) -> list[MemoryNode]: + """Get all profiles, sorted by message_time""" + nodes = self._load_nodes() + nodes.sort(key=lambda n: n.message_time) + return nodes + + @staticmethod + def _format_node(node: MemoryNode, add_profile_id: bool = False, add_history_id: bool = False) -> str: + """Format a single node to string""" + parts = [] + + if add_profile_id: + parts.append(f"profile_id={node.memory_id}") + + if node.message_time: + parts.append(f"[{node.message_time}]") + + parts.append(f"{node.when_to_use}: {node.content}") + + if add_history_id: + parts.append(f"history_id={node.ref_memory_id}") + + return " ".join(parts) + + def read_all(self, add_profile_id: bool = False, add_history_id: bool = False) -> str: + """Read all profiles and return formatted string""" + nodes = self.get_all() + formatted_profiles = [ + self._format_node(node, add_profile_id, add_history_id) + for node in nodes + ] + logger.info(f"Read {len(formatted_profiles)} profiles from {self.cache_key}") + return "\n".join(formatted_profiles).strip() diff --git a/reme/tool/memory/history/read_history.py b/reme/tool/memory/read_history.py similarity index 90% rename from reme/tool/memory/history/read_history.py rename to reme/tool/memory/read_history.py index 4da6918e..500d4bd9 100644 --- a/reme/tool/memory/history/read_history.py +++ b/reme/tool/memory/read_history.py @@ -2,8 +2,8 @@ from loguru import logger -from ..base_memory_tool import BaseMemoryTool -from ....core.schema import MemoryNode, ToolCall +from .base_memory_tool import BaseMemoryTool +from ...core.schema import MemoryNode, ToolCall class ReadHistory(BaseMemoryTool): @@ -36,7 +36,7 @@ class ReadHistory(BaseMemoryTool): nodes = await self.vector_store.get(vector_ids=[history_id]) if not nodes: - output = f"No history: {history_id}" + output = f"No history_id={history_id} data." logger.warning(output) return output diff --git a/reme/tool/memory/read_profile.py b/reme/tool/memory/read_profile.py new file mode 100644 index 00000000..b3286194 --- /dev/null +++ b/reme/tool/memory/read_profile.py @@ -0,0 +1,45 @@ +"""Read user profile tool""" +from pathlib import Path + +from loguru import logger + +from .base_memory_tool import BaseMemoryTool +from .profile_handler import ProfileHandler +from ...core.schema import ToolCall + + +class ReadProfile(BaseMemoryTool): + """Tool to read all user profiles""" + + def __init__(self, profile_path: str, **kwargs): + kwargs["enable_multiple"] = False + super().__init__(**kwargs) + self.profile_path: str = profile_path + + def _build_tool_call(self) -> ToolCall: + """Build and return the tool call schema""" + return ToolCall( + **{ + "description": "Read all user profiles.", + "parameters": { + "type": "object", + "properties": {}, + "required": [], + }, + }, + ) + + async def execute(self): + profile_handler = ProfileHandler( + profile_path=Path(self.profile_path) / self.vector_store.collection_name, + memory_target=self.memory_target, + ) + + profiles_str = profile_handler.read_all() + if not profiles_str: + output = "No profiles found." + logger.info(output) + return output + + logger.info(f"Successfully read profiles") + return profiles_str diff --git a/reme/tool/memory/retrieve_memory.py b/reme/tool/memory/retrieve_memory.py new file mode 100644 index 00000000..f8845405 --- /dev/null +++ b/reme/tool/memory/retrieve_memory.py @@ -0,0 +1,121 @@ +"""Retrieve memory from vector store""" + +from loguru import logger + +from .base_memory_tool import BaseMemoryTool +from .memory_handler import MemoryHandler +from ...core.schema import ToolCall, MemoryNode +from ...core.utils import deduplicate_memories + + +class RetrieveMemory(BaseMemoryTool): + """Tool to retrieve memories using similarity search""" + + def __init__(self, top_k: int = 20, enable_memory_target: bool = False, **kwargs): + super().__init__(**kwargs) + self.top_k: int = top_k + self.enable_memory_target: bool = enable_memory_target + + def _build_query_parameters(self) -> dict: + """Build the query parameters schema based on enabled features.""" + properties = { + "query": { + "type": "string", + "description": "query text for vector similarity search.", + }, + "time_range": { + "type": "string", + "description": "optional time range filter. Format: '20200101' or '20200101,20200102'", + }, + } + required = ["query"] + + if self.enable_memory_target: + properties["memory_target"] = { + "type": "string", + "description": "target memory type to search in.", + } + required.append("memory_target") + + return { + "type": "object", + "properties": properties, + "required": required, + } + + def _build_tool_call(self) -> ToolCall: + return ToolCall( + **{ + "description": "retrieve memories using vector similarity search.", + "parameters": self._build_query_parameters(), + }, + ) + + def _build_multiple_tool_call(self) -> ToolCall: + return ToolCall( + **{ + "description": "retrieve memories using multiple queries with vector similarity search.", + "parameters": { + "type": "object", + "properties": { + "query_items": { + "type": "array", + "description": "list of query items for vector similarity search.", + "items": self._build_query_parameters(), + }, + }, + "required": ["query_items"], + }, + }, + ) + + async def execute(self): + if self.enable_multiple: + query_items = self.context.get("query_items", []) + else: + query_items = [self.context] + + queries_by_target: dict[str, list[dict]] = {} + for item in query_items: + if self.enable_memory_target: + target = item["memory_target"] + else: + target = self.memory_target + if target not in queries_by_target: + queries_by_target[target] = [] + + filters = {} + time_range = item.get("time_range") + if time_range: + time_range = time_range.strip() + if "," in time_range: + start, end = time_range.split(",") + filters = {"time_int": [int(start.strip()), int(end.strip())]} + else: + filters = {"time_int": [int(time_range), int(time_range)]} + + queries_by_target[target].append({ + "query": item["query"], + "limit": self.top_k, + "filters": filters, + }) + + # Execute batch searches for each target + memory_nodes: list[MemoryNode] = [] + for target, searches in queries_by_target.items(): + handler = MemoryHandler(target, self.service_context) + nodes = await handler.batch_search(searches) + memory_nodes.extend(nodes) + + memory_nodes = deduplicate_memories(memory_nodes) + retrieved_ids = {n.memory_id for n in self.retrieved_nodes if n.memory_id} + new_nodes = [n for n in memory_nodes if n.memory_id not in retrieved_ids] + self.retrieved_nodes.extend(new_nodes) + + if not new_nodes: + output = "No new memories found." + else: + output = "\n".join([n.format(ref_memory_id_key="history_id") for n in new_nodes]) + + logger.info(f"Retrieved {len(memory_nodes)} memories, {len(new_nodes)} new after deduplication") + return output diff --git a/reme/tool/memory/retrieve_recent_memory.py b/reme/tool/memory/retrieve_recent_memory.py new file mode 100644 index 00000000..05b1a777 --- /dev/null +++ b/reme/tool/memory/retrieve_recent_memory.py @@ -0,0 +1,52 @@ +"""Retrieve most recent memories from vector store""" + +from loguru import logger + +from .base_memory_tool import BaseMemoryTool +from .memory_handler import MemoryHandler +from ...core.schema import ToolCall, MemoryNode +from ...core.utils import deduplicate_memories + + +class RetrieveRecentMemory(BaseMemoryTool): + """Tool to retrieve most recent memories sorted by time""" + + def __init__(self, top_k: int = 20, **kwargs): + kwargs["enable_multiple"] = False + super().__init__(**kwargs) + self.top_k: int = top_k + + def _build_tool_call(self) -> ToolCall: + return ToolCall( + **{ + "description": "retrieve the most recent memories sorted by message time (newest first).", + "parameters": { + "type": "object", + "properties": {}, + "required": [], + }, + }, + ) + + async def execute(self): + handler = MemoryHandler(self.memory_target, self.service_context) + + memory_nodes: list[MemoryNode] = await handler.list( + limit=self.top_k, + sort_key="message_time", + reverse=True, + ) + memory_nodes = deduplicate_memories(memory_nodes) + + retrieved_ids = {n.memory_id for n in self.retrieved_nodes if n.memory_id} + new_nodes = [n for n in memory_nodes if n.memory_id not in retrieved_ids] + self.retrieved_nodes.extend(new_nodes) + self.memory_nodes.extend(new_nodes) + + if not new_nodes: + output = "No new memories found." + else: + output = "\n".join([n.format(ref_memory_id_key="history_id") for n in new_nodes]) + + logger.info(f"Retrieved {len(memory_nodes)} memories, {len(new_nodes)} new after deduplication") + return output diff --git a/reme/tool/memory/update_memory.py b/reme/tool/memory/update_memory.py new file mode 100644 index 00000000..1370e180 --- /dev/null +++ b/reme/tool/memory/update_memory.py @@ -0,0 +1,139 @@ +"""Update memory in vector store""" + +from loguru import logger + +from .base_memory_tool import BaseMemoryTool +from .memory_handler import MemoryHandler +from ...core.schema import ToolCall + + +class UpdateMemory(BaseMemoryTool): + """Tool to update memories in vector store""" + + def __init__( + self, + enable_memory_target: bool = False, + enable_when_to_use: bool = False, + **kwargs, + ): + super().__init__(**kwargs) + self.enable_memory_target: bool = enable_memory_target + self.enable_when_to_use: bool = enable_when_to_use + + def _build_update_parameters(self) -> dict: + """Build the update parameters schema based on enabled features.""" + properties = { + "memory_id": { + "type": "string", + "description": "unique identifier of memory to update.", + }, + "message_time": { + "type": "string", + "description": "message time, e.g. '2020-01-01 00:00:00'", + }, + "memory_content": { + "type": "string", + "description": "new content of the memory.", + }, + } + required = ["memory_id", "message_time", "memory_content"] + + if self.enable_when_to_use: + properties["when_to_use"] = { + "type": "string", + "description": "description of when to use this memory.", + } + required.append("when_to_use") + + if self.enable_memory_target: + properties["memory_target"] = { + "type": "string", + "description": "target memory type for this memory.", + } + required.append("memory_target") + + return { + "type": "object", + "properties": properties, + "required": required, + } + + def _build_tool_call(self) -> ToolCall: + return ToolCall( + **{ + "description": "update a memory in vector store by replacing old memory with new content.", + "parameters": self._build_update_parameters(), + }, + ) + + def _build_multiple_tool_call(self) -> ToolCall: + return ToolCall( + **{ + "description": "update multiple memories in vector store by replacing old memories with new content.", + "parameters": { + "type": "object", + "properties": { + "memories": { + "type": "array", + "description": "list of memory update objects.", + "items": self._build_update_parameters(), + }, + }, + "required": ["memories"], + }, + }, + ) + + async def execute(self): + if self.enable_multiple: + memories = self.context.get("memories", []) + else: + memories = [self.context] + + # Group memories by memory_target if enabled + if self.enable_memory_target: + memories_by_target = {} + for mem in memories: + target = mem["memory_target"] + if target not in memories_by_target: + memories_by_target[target] = [] + memories_by_target[target].append(mem) + else: + memories_by_target = {self.memory_target: memories} + + # Process each memory_target group + all_memory_nodes = [] + for target, target_memories in memories_by_target.items(): + # Parse and prepare update data + update_dicts = [] + for mem in target_memories: + memory_content = mem.get("memory_content", "") + message_time = mem.get("message_time", "") + when_to_use = mem.get("when_to_use", "") if self.enable_when_to_use else "" + metadata = {} + try: + metadata["time_int"] = int(message_time.split(" ")[0].replace("-", "")) + except Exception: + logger.warning(f"Invalid message time format: {message_time}") + + update_dicts.append({ + "memory_id": mem.get("memory_id", ""), + "content": memory_content, + "when_to_use": when_to_use, + "message_time": message_time, + "author": self.author, + "metadata": metadata, + }) + + if update_dicts: + handler = MemoryHandler(target, self.service_context) + memory_nodes = await handler.update_batch(update_dicts) + all_memory_nodes.extend(memory_nodes) + + if not all_memory_nodes: + return "No valid memories provided." + + self.memory_nodes.extend(all_memory_nodes) + output = f"Successfully updated {len(all_memory_nodes)} memories." + logger.info(output) + return output diff --git a/reme/tool/memory/update_profile.py b/reme/tool/memory/update_profile.py new file mode 100644 index 00000000..2032964d --- /dev/null +++ b/reme/tool/memory/update_profile.py @@ -0,0 +1,97 @@ +"""Update user profile tool""" +from pathlib import Path + +from loguru import logger + +from .base_memory_tool import BaseMemoryTool +from .profile_handler import ProfileHandler +from ...core.schema import ToolCall + + +class UpdateProfile(BaseMemoryTool): + """Tool to update user profile by adding or removing profile entries""" + + def __init__(self, profile_path: str, **kwargs): + kwargs["enable_multiple"] = True + super().__init__(**kwargs) + + self.profile_path: str = profile_path + + def _build_multiple_tool_call(self) -> ToolCall: + """Build and return the multiple tool call schema""" + return ToolCall( + **{ + "description": "update user profile by removing and adding profile entries.", + "parameters": { + "type": "object", + "properties": { + "profile_ids_to_delete": { + "type": "array", + "description": "List of profile IDs to delete", + "items": { + "type": "string" + }, + }, + "profiles_to_add": { + "type": "array", + "description": "List of profiles to add", + "items": { + "type": "object", + "properties": { + "message_time": { + "type": "string", + "description": "Message time, e.g. '2020-01-01 00:00:00'", + }, + "profile_key": { + "type": "string", + "description": "Profile key or category, e.g. 'name'", + }, + "profile_value": { + "type": "string", + "description": "Profile value or content, e.g. 'John Smith'", + }, + }, + "required": ["message_time", "profile_key", "profile_value"], + }, + }, + }, + "required": ["profile_ids_to_delete", "profiles_to_add"], + }, + }, + ) + + async def execute(self): + profile_handler = ProfileHandler( + profile_path=Path(self.profile_path) / self.vector_store.collection_name, + memory_target=self.memory_target, + ) + + # Get parameters + profile_ids_to_delete = self.context.get("profile_ids_to_delete", []) + profile_ids_to_delete = sorted(set([pid for pid in profile_ids_to_delete if pid])) + profiles_to_add = self.context.get("profiles_to_add", []) + + if not profile_ids_to_delete and not profiles_to_add: + return "No profiles to remove or add, operation completed." + + # Delete profiles using ProfileHandler (batch mode) + removed_count = 0 + if profile_ids_to_delete: + removed_count = profile_handler.delete(profile_ids_to_delete) + + # Add new profiles using ProfileHandler (batch mode) + added_count = 0 + if profiles_to_add: + new_nodes = profile_handler.add_batch(profiles=profiles_to_add, ref_memory_id=self.history_node.memory_id) + self.memory_nodes.extend(new_nodes) + added_count = len(new_nodes) + + # Build output message + operations = [] + if removed_count > 0: + operations.append(f"removed {removed_count} old profiles.") + if added_count > 0: + operations.append(f"added {added_count} new profiles.") + operations.append("Operation completed.") + logger.info("\n".join(operations)) + return "\n".join(operations) diff --git a/reme/tool/memory/user_profile/__init__.py b/reme/tool/memory/user_profile/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/reme/tool/memory/user_profile/read_user_profile.py b/reme/tool/memory/user_profile/read_user_profile.py deleted file mode 100644 index 637b2070..00000000 --- a/reme/tool/memory/user_profile/read_user_profile.py +++ /dev/null @@ -1,63 +0,0 @@ -"""Read user profile tool""" - -from typing import Literal - -from loguru import logger - -from ..base_memory_tool import BaseMemoryTool -from ....core.enumeration import MemoryType -from ....core.schema import ToolCall, MemoryNode - - -class ReadUserProfile(BaseMemoryTool): - """Tool to read user profile from local memory""" - - def __init__(self, show_id: Literal["profile", "history"] = "profile", **kwargs): - kwargs["enable_multiple"] = False - super().__init__(**kwargs) - self.show_id = show_id - - def _build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "description": "read user profile.", - "parameters": { - "type": "object", - "properties": {}, - "required": [], - }, - }, - ) - - async def execute(self): - self.context.memory_type = MemoryType.PERSONAL - cached_data = self.local_memory.load(self.memory_cache_key, auto_clean=False) - - if not cached_data: - logger.info(f"No cached data found for {self.memory_cache_key}") - return "" - - nodes = [MemoryNode(**data) for data in cached_data] - self.memory_nodes.clear() - self.memory_nodes.extend(nodes) - nodes.sort(key=lambda n: n.metadata.get("update_time", "")) - - formatted_profiles = [] - for node in nodes: - parts = [] - if self.show_id == "profile": - parts.append(f"profile_id={node.memory_id}") - - if update_time := node.metadata.get("update_time"): - parts.append(f"update_time={update_time}") - - parts.append(f"{node.when_to_use}: {node.content}") - - if self.show_id == "history": - parts.append(f"history_id={node.ref_memory_id}") - - formatted_profiles.append(" ".join(parts)) - - logger.info(f"Read {len(formatted_profiles)} profiles from cache key: {self.memory_cache_key}") - - return "\n".join(formatted_profiles).strip() diff --git a/reme/tool/memory/user_profile/update_user_profile.py b/reme/tool/memory/user_profile/update_user_profile.py deleted file mode 100644 index 1c568c5d..00000000 --- a/reme/tool/memory/user_profile/update_user_profile.py +++ /dev/null @@ -1,112 +0,0 @@ -"""Update user profile tool""" - -from loguru import logger - -from ..base_memory_tool import BaseMemoryTool -from ....core.enumeration import MemoryType -from ....core.schema import ToolCall, MemoryNode -from ....core.utils import deduplicate_memories - - -class UpdateUserProfile(BaseMemoryTool): - """Tool to update user profile by adding or removing profile entries""" - - def __init__(self, **kwargs): - kwargs["enable_multiple"] = True - super().__init__(**kwargs) - - def _build_multiple_tool_call(self) -> ToolCall: - """Build and return the multiple tool call schema""" - return ToolCall( - **{ - "description": "update user profile by adding or removing profile entries.", - "parameters": { - "type": "object", - "properties": { - "profile_ids_to_delete": { - "type": "array", - "description": "List of profile IDs to delete", - "items": {"type": "string"}, - }, - "profiles_to_add": { - "type": "array", - "description": "List of profiles to add", - "items": { - "type": "object", - "properties": { - "update_time": { - "type": "string", - "description": "Update time, e.g. '2020-01-01 00:00:00'", - }, - "profile_key": { - "type": "string", - "description": "Profile key or category, e.g. 'name'", - }, - "profile_value": { - "type": "string", - "description": "Profile value or content, e.g. 'John Smith'", - }, - }, - "required": ["update_time", "profile_key", "profile_value"], - }, - }, - }, - "required": ["profile_ids_to_delete", "profiles_to_add"], - }, - }, - ) - - async def execute(self): - # Get and deduplicate profile IDs to delete - self.context.memory_type = MemoryType.PERSONAL - - profile_ids_to_delete = self.context.get("profile_ids_to_delete", []) - profile_ids_to_delete = list(dict.fromkeys([pid for pid in profile_ids_to_delete if pid])) - profiles_to_add = self.context.get("profiles_to_add", []) - - if not profile_ids_to_delete and not profiles_to_add: - return "No profiles to remove or add. Operation completed." - - # Load existing profiles from local memory - cached_data = self.local_memory.load(self.memory_cache_key, auto_clean=False) - existing_nodes = [MemoryNode(**data) for data in cached_data] if cached_data else [] - - # Remove profiles - removed_count = 0 - if profile_ids_to_delete: - original_count = len(existing_nodes) - existing_nodes = [n for n in existing_nodes if n.memory_id not in profile_ids_to_delete] - removed_count = original_count - len(existing_nodes) - logger.info(f"Removed {removed_count} profiles.") - - # Add new profiles - new_nodes = [] - if profiles_to_add: - for profile in profiles_to_add: - node = MemoryNode( - memory_type=self.memory_type, - memory_target=self.memory_target, - when_to_use=profile.get("profile_key", ""), - content=profile.get("profile_value", ""), - ref_memory_id=self.history_node.memory_id, - author=self.author, - metadata={"update_time": profile.get("update_time", "")}, - ) - new_nodes.append(node) - logger.info(f"Added {len(new_nodes)} new profiles.") - - # Deduplicate and save updated profiles - self.memory_nodes.extend(new_nodes) - updated_nodes = deduplicate_memories(existing_nodes + new_nodes) - nodes_data = [node.model_dump(exclude_none=True) for node in updated_nodes] - self.local_memory.save(self.memory_cache_key, nodes_data) - - # Build output message - operations = [] - if removed_count > 0: - operations.append(f"removed {removed_count} old profiles.") - if len(new_nodes) > 0: - operations.append(f"added {len(new_nodes)} new profiles.") - operations.append("Operation completed.") - logger.info("\n".join(operations)) - return "\n".join(operations) diff --git a/reme/tool/memory/vector/__init__.py b/reme/tool/memory/vector/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/reme/tool/memory/vector/add_memory.py b/reme/tool/memory/vector/add_memory.py deleted file mode 100644 index 5d6b8528..00000000 --- a/reme/tool/memory/vector/add_memory.py +++ /dev/null @@ -1,110 +0,0 @@ -"""Add memory to vector store""" - -from loguru import logger - -from ..base_memory_tool import BaseMemoryTool -from ....core.schema import ToolCall, MemoryNode - - -class AddMemory(BaseMemoryTool): - """Tool to add memories to vector store""" - - def _build_tool_call(self) -> ToolCall: - """Build and return the single tool call schema""" - return ToolCall( - **{ - "description": "add a memory to vector store for future retrieval.", - "parameters": { - "type": "object", - "properties": { - "conversation_time": { - "type": "string", - "description": "conversation time, e.g. '2020-01-01 00:00:00'", - }, - "memory_content": { - "type": "string", - "description": "content of the memory.", - }, - }, - "required": ["conversation_time", "memory_content"], - }, - }, - ) - - def _build_multiple_tool_call(self) -> ToolCall: - """Build and return the multiple tool call schema""" - return ToolCall( - **{ - "description": "add multiple memories to vector store for future retrieval.", - "parameters": { - "type": "object", - "properties": { - "memories": { - "type": "array", - "description": "list of memories to store.", - "items": { - "type": "object", - "properties": { - "conversation_time": { - "type": "string", - "description": "conversation time, e.g. '2020-01-01 00:00:00'", - }, - "memory_content": { - "type": "string", - "description": "content of the memory.", - }, - }, - "required": ["conversation_time", "memory_content"], - }, - }, - }, - "required": ["memories"], - }, - }, - ) - - def _create_memory_node(self, data: dict) -> MemoryNode: - """Create a MemoryNode from a dictionary.""" - memory_content = data.get("memory_content", "") - conversation_time = data.get("conversation_time", "") - metadata: dict = {"conversation_time": conversation_time} - - try: - metadata["time_int"] = int(conversation_time.split(" ")[0].replace("-", "")) - except Exception: - logger.warning(f"Invalid conversation time format. {conversation_time}") - - return MemoryNode( - memory_type=self.memory_type, - memory_target=self.memory_target, - content=memory_content, - author=self.author, - ref_memory_id=self.history_node.memory_id, - metadata=metadata, - ) - - async def execute(self): - memory_nodes: list[MemoryNode] = [] - memories: list[dict] = self.context.get("memories", []) - - if not memories: - memory_nodes.append(self._create_memory_node(self.context)) - else: - for mem in memories: - memory_nodes.append(self._create_memory_node(mem)) - - if not memory_nodes: - output = "No valid memories provided for addition." - logger.info(output) - return output - - vector_nodes = [node.to_vector_node() for node in memory_nodes] - vector_ids: list[str] = [node.vector_id for node in vector_nodes] - - await self.vector_store.delete(vector_ids=list(set(vector_ids))) - await self.vector_store.insert(nodes=vector_nodes) - self.memory_nodes.extend(memory_nodes) - - output = f"Successfully added {len(memory_nodes)} memories to vector_store." - logger.info(output) - return output diff --git a/reme/tool/memory/vector/retrieve_memory.py b/reme/tool/memory/vector/retrieve_memory.py deleted file mode 100644 index 7c95b538..00000000 --- a/reme/tool/memory/vector/retrieve_memory.py +++ /dev/null @@ -1,136 +0,0 @@ -"""Retrieve memory from vector store""" - -from loguru import logger - -from ..base_memory_tool import BaseMemoryTool -from ....core.schema import ToolCall, MemoryNode, VectorNode -from ....core.utils import deduplicate_memories - - -class RetrieveMemory(BaseMemoryTool): - """Tool to retrieve memories from vector store using similarity search""" - - def __init__(self, top_k: int = 20, **kwargs): - super().__init__(**kwargs) - self.top_k: int = top_k - - @staticmethod - def _build_query_parameters() -> dict: - """Build query parameters schema for retrieval""" - return { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "query text for vector similarity search.", - }, - "time_range": { - "type": "string", - "description": "optional time range filter. " - "Format: single date '20200101' or range '20200101,20200102'", - }, - }, - "required": ["query"], - } - - def _build_tool_call(self) -> ToolCall: - """Build and return the tool call schema""" - return ToolCall( - **{ - "description": "retrieve memories using vector similarity search.", - "parameters": self._build_query_parameters(), - }, - ) - - def _build_multiple_tool_call(self) -> ToolCall: - """Build and return the multiple tool call schema""" - return ToolCall( - **{ - "description": "retrieve memories using multiple queries with vector similarity search.", - "parameters": { - "type": "object", - "properties": { - "query_items": { - "type": "array", - "description": "list of query items for vector similarity search.", - "items": self._build_query_parameters(), - }, - }, - "required": ["query_items"], - }, - }, - ) - - async def _retrieve_by_query( - self, - memory_type: str, - memory_target: str, - query: str, - time_range: str | None = None, - ) -> list[MemoryNode]: - """Retrieve memories by query with filters""" - filter_dict: dict = { - "memory_type": memory_type, - "memory_target": memory_target, - } - - if time_range: - time_range = time_range.strip() - if "," in time_range: - parts = time_range.split(",") - start_time = int(parts[0].strip()) - end_time = int(parts[1].strip()) - filter_dict["time_int"] = [start_time, end_time] - else: - single_time = int(time_range) - filter_dict["time_int"] = [single_time, single_time] - - nodes: list[VectorNode] = await self.vector_store.search(query=query, limit=self.top_k, filters=filter_dict) - return [MemoryNode.from_vector_node(n) for n in nodes] - - async def execute(self): - memory_type: str = self.memory_type.value - memory_target: str = self.memory_target - - if self.enable_multiple: - query_items: list[dict] = self.context.get("query_items", []) - else: - query_items: list[dict] = [ - { - "query": self.context.get("query", ""), - "time_range": self.context.get("time_range", ""), - }, - ] - - query_items = [item for item in query_items if item.get("query")] - memory_nodes: list[MemoryNode] = [] - for item in query_items: - retrieved = await self._retrieve_by_query( - memory_type=memory_type, - memory_target=memory_target, - query=item["query"], - time_range=item.get("time_range", ""), - ) - memory_nodes.extend(retrieved) - - memory_nodes = deduplicate_memories(memory_nodes) - retrieved_memory_ids = {node.memory_id for node in self.retrieved_nodes if node.memory_id} - new_memory_nodes = [node for node in memory_nodes if node.memory_id not in retrieved_memory_ids] - self.retrieved_nodes.extend(new_memory_nodes) - - if not new_memory_nodes: - output = "No new memory_nodes found matching the query (duplicates removed)." - else: - outputs = [] - for node in new_memory_nodes: - line = "" - if "conversation_time" in node.metadata and node.metadata["conversation_time"]: - line += f"conversation_time={node.metadata['conversation_time']} " - line += node.content.strip() + " " - if node.ref_memory_id: - line += f"history_id={node.ref_memory_id}" - outputs.append(line.strip()) - output = "\n".join(outputs) - - logger.info(f"Retrieved {len(memory_nodes)} memory_nodes, {len(new_memory_nodes)} new after deduplication") - return output diff --git a/reme/tool/memory/vector/retrieve_recent_memory.py b/reme/tool/memory/vector/retrieve_recent_memory.py deleted file mode 100644 index c7c6bc7f..00000000 --- a/reme/tool/memory/vector/retrieve_recent_memory.py +++ /dev/null @@ -1,62 +0,0 @@ -"""Retrieve most recent memories from vector store""" - -from loguru import logger - -from ..base_memory_tool import BaseMemoryTool -from ....core.schema import ToolCall, MemoryNode, VectorNode -from ....core.utils import deduplicate_memories - - -class RetrieveRecentMemory(BaseMemoryTool): - """Tool to retrieve most recent memories sorted by conversation time""" - - def __init__(self, top_k: int = 20, **kwargs): - kwargs["enable_multiple"] = False - super().__init__(**kwargs) - self.top_k: int = top_k - - def _build_tool_call(self) -> ToolCall: - """Build and return the tool call schema""" - return ToolCall( - **{ - "description": "retrieve the most recent memories sorted by conversation time (newest first).", - "parameters": { - "type": "object", - "properties": {}, - "required": [], - }, - }, - ) - - async def _retrieve_recent(self) -> list[MemoryNode]: - """Retrieve recent memories sorted by conversation_time descending""" - filter_dict = { - "memory_type": self.memory_type.value, - "memory_target": self.memory_target, - } - - nodes: list[VectorNode] = await self.vector_store.list( - filters=filter_dict, - limit=self.top_k, - sort_key="conversation_time", - reverse=True, - ) - - return [MemoryNode.from_vector_node(n) for n in nodes] - - async def execute(self): - memory_nodes: list[MemoryNode] = await self._retrieve_recent() - memory_nodes = deduplicate_memories(memory_nodes) - - retrieved_memory_ids = {node.memory_id for node in self.retrieved_nodes if node.memory_id} - new_memory_nodes = [node for node in memory_nodes if node.memory_id not in retrieved_memory_ids] - self.retrieved_nodes.extend(new_memory_nodes) - self.memory_nodes.extend(new_memory_nodes) - - if not new_memory_nodes: - output = "No new memory_nodes found (duplicates removed)." - else: - output = "\n".join([m.format_memory() for m in new_memory_nodes]) - - logger.info(f"Retrieved {len(memory_nodes)} memory_nodes, {len(new_memory_nodes)} new after deduplication") - return output diff --git a/reme/tool/memory/vector/update_memory.py b/reme/tool/memory/vector/update_memory.py deleted file mode 100644 index ca52582b..00000000 --- a/reme/tool/memory/vector/update_memory.py +++ /dev/null @@ -1,126 +0,0 @@ -"""Update memory in vector store""" - -from loguru import logger - -from ..base_memory_tool import BaseMemoryTool -from ....core.schema import ToolCall, MemoryNode - - -class UpdateMemory(BaseMemoryTool): - """Tool to update memories in vector store""" - - def _build_tool_call(self) -> ToolCall: - """Build and return the single tool call schema""" - return ToolCall( - **{ - "description": "update a memory in vector store by replacing old memory with new content.", - "parameters": { - "type": "object", - "properties": { - "memory_id": { - "type": "string", - "description": "unique identifier of memory to update.", - }, - "conversation_time": { - "type": "string", - "description": "conversation time, e.g. '2020-01-01 00:00:00'", - }, - "memory_content": { - "type": "string", - "description": "new content of the memory.", - }, - }, - "required": ["memory_id", "conversation_time", "memory_content"], - }, - }, - ) - - def _build_multiple_tool_call(self) -> ToolCall: - """Build and return the multiple tool call schema""" - return ToolCall( - **{ - "description": "update multiple memories in vector store by replacing old memories with new content.", - "parameters": { - "type": "object", - "properties": { - "memories": { - "type": "array", - "description": "list of memory update objects.", - "items": { - "type": "object", - "properties": { - "memory_id": { - "type": "string", - "description": "unique identifier of memory to update.", - }, - "conversation_time": { - "type": "string", - "description": "conversation time, e.g. '2020-01-01 00:00:00'", - }, - "memory_content": { - "type": "string", - "description": "new content of the memory.", - }, - }, - "required": ["memory_id", "conversation_time", "memory_content"], - }, - }, - }, - "required": ["memories"], - }, - }, - ) - - def _create_memory_node(self, data: dict) -> tuple[str, MemoryNode]: - """Create a MemoryNode from a dictionary.""" - memory_id = data.get("memory_id", "") - memory_content = data.get("memory_content", "") - conversation_time = data.get("conversation_time", "") - metadata: dict = {"conversation_time": conversation_time} - - try: - metadata["time_int"] = int(conversation_time.split(" ")[0].replace("-", "")) - except Exception: - logger.warning(f"Invalid conversation time format. {conversation_time}") - - memory_node = MemoryNode( - memory_type=self.memory_type, - memory_target=self.memory_target, - content=memory_content, - author=self.author, - metadata=metadata, - ) - - return memory_id, memory_node - - async def execute(self): - old_memory_ids: list[str] = [] - memory_nodes: list[MemoryNode] = [] - memories: list[dict] = self.context.get("memories", []) - - if not memories: - old_id, node = self._create_memory_node(self.context) - old_memory_ids.append(old_id) - memory_nodes.append(node) - else: - for mem in memories: - old_id, node = self._create_memory_node(mem) - old_memory_ids.append(old_id) - memory_nodes.append(node) - - if not memory_nodes: - output = "No valid memories provided for update." - logger.info(output) - return output - - vector_nodes = [node.to_vector_node() for node in memory_nodes] - new_vector_ids: list[str] = [node.vector_id for node in vector_nodes] - - all_ids_to_delete = list(set(old_memory_ids + new_vector_ids)) - await self.vector_store.delete(vector_ids=all_ids_to_delete) - await self.vector_store.insert(nodes=vector_nodes) - self.memory_nodes.extend(memory_nodes) - - output = f"Successfully updated {len(memory_nodes)} memories in vector_store." - logger.info(output) - return output