From 9f8ad20c5382391f7a7e16e975302d1c113ea4e2 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 23 Jan 2026 15:47:02 +0800 Subject: [PATCH] feat(memory): add ReAct agent and memory management tools --- reme/agent/memory/__init__.py | 0 reme/agent/memory/base_memory_agent.py | 10 + reme/agent/memory/default/__init__.py | 0 reme/agent/memory/default/reme_summarizer.py | 63 ++++++ .../agent/memory/default/reme_summarizer.yaml | 26 +++ reme/core/op/__init__.py | 2 + reme/core/op/base_op.py | 3 +- reme/core/op/base_react.py | 152 ++++++++++++++ reme/core/schema/message.py | 7 +- reme/core/schema/tool_call.py | 4 +- reme/tool/memory/vector/add_memory.py | 93 +++++++++ reme/tool/memory/vector/delete_memory.py | 45 +++++ reme/tool/memory/vector/retrieve_memory.py | 185 ++++++++++++++++++ reme/tool/memory/vector/update_memory.py | 100 ++++++++++ .../procedural_memory/summary/__init__.py | 0 reme_ai/mem_tool/hands_off_tool.py | 144 -------------- reme_ai/mem_tool/hands_off_tool.yaml | 19 -- reme_ai/mem_tool/read_local_memories.py | 54 ----- reme_ai/mem_tool/read_local_memories.yaml | 8 - reme_ai/mem_tool/write_local_memories.py | 57 ------ reme_ai/mem_tool/write_local_memories.yaml | 5 - 21 files changed, 682 insertions(+), 295 deletions(-) create mode 100644 reme/agent/memory/__init__.py create mode 100644 reme/agent/memory/base_memory_agent.py create mode 100644 reme/agent/memory/default/__init__.py create mode 100644 reme/agent/memory/default/reme_summarizer.py create mode 100644 reme/agent/memory/default/reme_summarizer.yaml create mode 100644 reme/core/op/base_react.py create mode 100644 reme/tool/memory/vector/add_memory.py create mode 100644 reme/tool/memory/vector/delete_memory.py create mode 100644 reme/tool/memory/vector/retrieve_memory.py create mode 100644 reme/tool/memory/vector/update_memory.py create mode 100644 reme/workflow/procedural_memory/summary/__init__.py delete mode 100644 reme_ai/mem_tool/hands_off_tool.py delete mode 100644 reme_ai/mem_tool/hands_off_tool.yaml delete mode 100644 reme_ai/mem_tool/read_local_memories.py delete mode 100644 reme_ai/mem_tool/read_local_memories.yaml delete mode 100644 reme_ai/mem_tool/write_local_memories.py delete mode 100644 reme_ai/mem_tool/write_local_memories.yaml diff --git a/reme/agent/memory/__init__.py b/reme/agent/memory/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme/agent/memory/base_memory_agent.py b/reme/agent/memory/base_memory_agent.py new file mode 100644 index 00000000..27d4d777 --- /dev/null +++ b/reme/agent/memory/base_memory_agent.py @@ -0,0 +1,10 @@ +"""Base memory agent for handling memory operations with tool-based reasoning.""" + +from abc import ABCMeta + +from ...core.enumeration import MemoryType +from ...core.op import BaseReact + + +class BaseMemoryAgent(BaseReact, metaclass=ABCMeta): + memory_type: MemoryType | None = None diff --git a/reme/agent/memory/default/__init__.py b/reme/agent/memory/default/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme/agent/memory/default/reme_summarizer.py b/reme/agent/memory/default/reme_summarizer.py new file mode 100644 index 00000000..d6b7e4eb --- /dev/null +++ b/reme/agent/memory/default/reme_summarizer.py @@ -0,0 +1,63 @@ +from loguru import logger + +from ..base_memory_agent import BaseMemoryAgent +from ....core.enumeration import Role, MemoryType +from ....core.schema import Message, MemoryNode +from ....core.utils import format_messages + + +class ReMeSummarizer(BaseMemoryAgent): + + def __init__(self, meta_memories: list[dict] | None = None, **kwargs): + super().__init__(**kwargs) + self.meta_memories: list[dict] = meta_memories or [] + + async def _read_meta_memories(self) -> str: + from ....tool.memory import ReadMetaMemory + meta_memory_info = ReadMetaMemory().format_memory_metadata(self.meta_memories) + logger.info(f"meta_memory_info={meta_memory_info}") + return meta_memory_info + + async def build_messages(self) -> list[Message]: + self.context.messages = [Message(**m) if isinstance(m, dict) else m for m in self.context.messages] + history_content = self.description + "\n" + format_messages(self.context.messages) + self.context.history_node = history_node = MemoryNode( + memory_type=MemoryType.HISTORY, + memory_target="", + when_to_use=history_content[:100], + content=history_content, + ref_memory_id="", + author=self.author, + metadata={}, + ) + + logger.info(f"Adding summary node: {history_node.model_dump_json(indent=2, exclude_none=True)}") + await self.vector_store.delete(history_node.memory_id) + await self.vector_store.insert([history_node.to_vector_node()]) + + messages = [ + Message( + role=Role.SYSTEM, + content=self.prompt_format( + prompt_name="system_prompt", + meta_memory_info=await self._read_meta_memories(), + context=history_node.content, + ), + ), + Message( + role=Role.USER, + content=self.get_prompt("user_message"), + ), + ] + + return messages + + async def _acting_step(self, assistant_message: Message, step: int, **kwargs) -> list[Message]: + return await super()._acting_step( + assistant_message, + step, + messages=self.context.messages, + history_node=self.context.history_node, + author=self.author, + **kwargs, + ) diff --git a/reme/agent/memory/default/reme_summarizer.yaml b/reme/agent/memory/default/reme_summarizer.yaml new file mode 100644 index 00000000..4cb6f54f --- /dev/null +++ b/reme/agent/memory/default/reme_summarizer.yaml @@ -0,0 +1,26 @@ +tool: | + Orchestrate memory updates across specialized memory agents. + +system_prompt: | + You are a Memory Orchestrator responsible for routing memory tasks to specialized agents based on the context. + + # Context + {context} + + ## Available Memory Agents + Each line indicates a specialized Memory Agent dedicated to deep summarization and updating of memories within a specific dimension (memory_type + memory_target). + Format: "- (): " + {meta_memory_info} + + ## Your Task + Use the `hands_off` 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 + - Do NOT create new agents or use memory_type/memory_target combinations that don't exist above + 3. Multiple tasks can be specified to enable parallel processing by specialized agents + + Note: If the context contains no memorable information (e.g., simple greetings), return ``. + +user_message: | + Please analyze the context and route memory tasks to the appropriate existing agents. diff --git a/reme/core/op/__init__.py b/reme/core/op/__init__.py index 53a6f379..3c86a522 100644 --- a/reme/core/op/__init__.py +++ b/reme/core/op/__init__.py @@ -2,6 +2,7 @@ from .base_op import BaseOp from .base_ray_op import BaseRayOp +from .base_react import BaseReact from .base_tool import BaseTool from .mcp_tool import MCPTool from .parallel_op import ParallelOp @@ -11,6 +12,7 @@ from ..context import R __all__ = [ "BaseOp", "BaseRayOp", + "BaseReact", "BaseTool", "MCPTool", "ParallelOp", diff --git a/reme/core/op/base_op.py b/reme/core/op/base_op.py index ef469d0f..9ed78187 100644 --- a/reme/core/op/base_op.py +++ b/reme/core/op/base_op.py @@ -21,6 +21,7 @@ from ..vector_store import BaseVectorStore class BaseOp(metaclass=ABCMeta): """Base operator class for LLM workflow execution and composition.""" + __alias_name__: str = "" def __new__(cls, *args, **kwargs): """Capture initialization arguments for object cloning.""" @@ -52,7 +53,7 @@ class BaseOp(metaclass=ABCMeta): **kwargs, ): """Initialize operator configurations and internal state.""" - self.name = name or camel_to_snake(self.__class__.__name__) + self.name = name or self.__alias_name__ or camel_to_snake(self.__class__.__name__) self.async_mode = async_mode self.language = language self.prompt = self._get_prompt_handler(prompt_name, prompt_path) diff --git a/reme/core/op/base_react.py b/reme/core/op/base_react.py new file mode 100644 index 00000000..e56d5f47 --- /dev/null +++ b/reme/core/op/base_react.py @@ -0,0 +1,152 @@ +"""Base memory agent for handling memory operations with tool-based reasoning.""" + +import asyncio + +from loguru import logger + +from . import BaseTool +from ..enumeration import Role +from ..op import BaseOp +from ..schema import Message + + +class BaseReact(BaseOp): + """ReAct agent that performs reasoning and acting cycles with tools.""" + + def __init__( + self, + tools: list[BaseTool], + tool_call_interval: float = 0, + max_steps: int = 10, + **kwargs, + ): + """Initialize ReAct agent with tools and execution parameters.""" + kwargs["sub_ops"] = tools or [] + super().__init__(**kwargs) + # Filter only BaseTool instances from sub_ops + self.sub_ops: list[BaseTool] = [t for t in self.sub_ops if isinstance(t, BaseTool)] + self.tool_call_interval: float = tool_call_interval + self.max_steps: int = max_steps + + @property + def tools(self) -> list[BaseTool]: + """Return available tools for the agent.""" + return self.sub_ops + + async def build_messages(self) -> list[Message]: + """Build initial message list from context query or messages.""" + if self.context.get("query"): + messages = [Message(role=Role.USER, content=self.context.query)] + elif self.context.get("messages"): + messages = [Message(**m) if isinstance(m, dict) else m for m in self.context.messages] + else: + raise ValueError("input must have either `query` or `messages`") + return messages + + async def _reasoning_step( + self, + messages: list[Message], + step: int, + stage: str = "", + **kwargs) -> tuple[Message, bool]: + """Execute one reasoning step where LLM decides whether to use tools.""" + # Get tool definitions for LLM + tool_calls = [t.tool_call for t in self.tools] + # Generate assistant response with potential tool calls + assistant_message: Message = await self.llm.chat(messages=messages, tools=tool_calls, **kwargs) + messages.append(assistant_message) + assistant_content: str = assistant_message.simple_dump(as_dict=False) + logger.info(f"[{self.__class__.__name__} {stage or ''} step{step + 1}] assistant={assistant_content}") + # Determine if tools should be called + should_act = bool(assistant_message.tool_calls) + return assistant_message, should_act + + async def _acting_step( + self, + assistant_message: Message, + step: int, + stage: str = "", + **kwargs + ) -> tuple[list[BaseTool], list[Message]]: + """Execute tool calls requested by the assistant and collect results.""" + tool_list: list[BaseTool] = [] + tool_messages: list[Message] = [] + + if not assistant_message.tool_calls: + return tool_list, tool_messages + + # Create tool name to tool instance mapping + tool_dict = {t.tool_call.name: t for t in self.tools} + for j, tool_call in enumerate(assistant_message.tool_calls): + prefix: str = f"[{self.__class__.__name__} {stage or ''} step{step + 1}.{j}]" + if tool_call.name not in tool_dict: + logger.warning(f"{prefix} unknown tool_call={tool_call.name}") + continue + + logger.info(f"{prefix} submit tool_calls={tool_call.simple_output_dump(as_dict=False)}") + + # Create independent tool copy with unique ID + tool_copy: BaseTool = tool_dict[tool_call.name].copy() + tool_copy.tool_call.id = tool_call.id + tool_list.append(tool_copy) + + # Create isolated kwargs for each tool call to avoid parameter conflicts + tool_kwargs = {**kwargs, **tool_call.argument_dict} + self.submit_async_task(tool_copy.call, **tool_kwargs) + if self.tool_call_interval > 0: + await asyncio.sleep(self.tool_call_interval) + + # Wait for all tool executions to complete + await self.join_async_tasks() + + # Collect tool results as messages + for j, tool in enumerate(tool_list): + tool_messages.append(Message( + role=Role.TOOL, + content=tool.response.answer, + tool_call_id=tool.tool_call.id, + )) + prefix: str = f"[{self.__class__.__name__} {stage or ''} step{step + 1}.{j}]" + logger.info(f"{prefix} join tool={tool.name} result={tool.response.answer}") + return tool_list, tool_messages + + async def react(self, messages: list[Message], stage: str = ""): + """Run ReAct loop alternating between reasoning and acting until completion.""" + success: bool = False + tools: list[BaseTool] = [] + for step in range(self.max_steps): + # Reasoning: LLM decides next action + assistant_message, should_act = await self._reasoning_step(messages, step=step, stage=stage) + + if not should_act: + # No tools requested, task complete + success = True + break + + # Acting: execute tools and collect results + t_tools, tool_messages = await self._acting_step(assistant_message, step=step, stage=stage) + tools.extend(t_tools) + messages.extend(tool_messages) + + return tools, messages, success + + async def execute(self): + """Execute the ReAct agent and return final results.""" + # Log available tools + for i, tool in enumerate(self.tools): + logger.info(f"[{self.__class__.__name__}] tool_call={tool.tool_call.simple_input_dump(as_dict=False)}") + + # Build and log initial messages + messages = await self.build_messages() + for i, message in enumerate(messages): + role = message.name or message.role + logger.info(f"[{self.__class__.__name__}] role={role} {message.simple_dump(as_dict=False)}") + + # Run ReAct loop + t_tools, messages, success = await self.react(messages) + return { + "answer": messages[-1].content if success else "", + "success": success, + "messages": messages, + "tools": t_tools, + } diff --git a/reme/core/schema/message.py b/reme/core/schema/message.py index 321dd8c8..54acde63 100644 --- a/reme/core/schema/message.py +++ b/reme/core/schema/message.py @@ -84,7 +84,7 @@ class Message(BaseModel): add_reasoning: bool = True, add_time_created: bool = False, add_metadata: bool = False, - enable_json_dump: bool = False, + as_dict: bool = True, ) -> dict | str: """Transforms the message into a simplified dictionary for standard APIs.""" result = {} @@ -109,10 +109,7 @@ class Message(BaseModel): if add_metadata: result["metadata"] = self.metadata - if enable_json_dump: - return json.dumps(result, ensure_ascii=False) - else: - return result + return result if as_dict else json.dumps(result, ensure_ascii=False) def format_message( self, diff --git a/reme/core/schema/tool_call.py b/reme/core/schema/tool_call.py index 22577817..f971937a 100644 --- a/reme/core/schema/tool_call.py +++ b/reme/core/schema/tool_call.py @@ -143,7 +143,7 @@ class ToolCall(BaseModel): "parameters": self.parameters.simple_input_dump(), }, } - return result if as_dict else json.dumps(result) + return result if as_dict else json.dumps(result, ensure_ascii=False) def simple_output_dump(self, as_dict: bool = True) -> dict | str: """Convert ToolCall to output format dictionary or JSON string for API responses. @@ -160,7 +160,7 @@ class ToolCall(BaseModel): }, "type": self.type, } - return result if as_dict else json.dumps(result) + return result if as_dict else json.dumps(result, ensure_ascii=False) @property def argument_dict(self) -> dict: diff --git a/reme/tool/memory/vector/add_memory.py b/reme/tool/memory/vector/add_memory.py new file mode 100644 index 00000000..78c6d07e --- /dev/null +++ b/reme/tool/memory/vector/add_memory.py @@ -0,0 +1,93 @@ +"""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 __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 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": { + "memory_content": { + "type": "string", + "description": "content of the memory.", + }, + "metadata": { + "type": "object", + "description": "metadata for the memory.", + } + }, + "required": ["memory_content"], + }, + }, + }, + "required": ["memories"], + }, + }, + ) + + def _extract_memory_data(self, mem_dict: dict) -> tuple[str, dict]: + """Extract memory content and metadata from dict""" + memory_content = mem_dict.get("memory_content", "") + raw_metadata = mem_dict.get("metadata", {}) + metadata = {key: str(value).strip() for key, value in raw_metadata.items() if value} + return memory_content, metadata + + def _build_memory_node(self, content: str, metadata: dict = None) -> MemoryNode: + """Build a memory node""" + return MemoryNode( + memory_type=self.memory_type, + memory_target=self.memory_target, + content=content, + author=self.author, + metadata=metadata or {}, + ) + + async def execute(self): + memory_nodes: list[MemoryNode] = [] + + memories: list[dict] = self.context.get("memories", []) + if not memories: + self.output = "No memories provided for addition." + return + + for mem in memories: + memory_content, metadata = self._extract_memory_data(mem) + if not memory_content: + logger.warning("Skipping memory with empty content") + continue + + memory_nodes.append(self._build_memory_node(memory_content, metadata=metadata)) + + if not memory_nodes: + self.output = "No valid memories provided for addition." + return + + 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=vector_ids) + await self.vector_store.insert(nodes=vector_nodes) + self.memory_nodes = memory_nodes + + self.output = f"Successfully added {len(memory_nodes)} memories to vector_store." + logger.info(self.output) diff --git a/reme/tool/memory/vector/delete_memory.py b/reme/tool/memory/vector/delete_memory.py new file mode 100644 index 00000000..a8f84596 --- /dev/null +++ b/reme/tool/memory/vector/delete_memory.py @@ -0,0 +1,45 @@ +"""Delete memory from vector store""" + +from loguru import logger + +from ..base_memory_tool import BaseMemoryTool +from ....core.schema import ToolCall + + +class DeleteMemory(BaseMemoryTool): + """Tool to delete memories from vector store""" + + 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": "delete multiple memories from vector store using their unique IDs.", + "parameters": { + "type": "object", + "properties": { + "memory_ids": { + "type": "array", + "description": "list of unique identifiers (memory_ids) of memories to delete.", + "items": {"type": "string"}, + }, + }, + "required": ["memory_ids"], + }, + }, + ) + + async def execute(self): + memory_ids = [m for m in self.context.get("memory_ids", []) if m] + + if not memory_ids: + self.output = "No valid memory IDs provided for deletion." + return + + await self.vector_store.delete(vector_ids=memory_ids) + self.memory_nodes = memory_ids + self.output = f"Successfully deleted {len(memory_ids)} memories from vector_store." + logger.info(self.output) diff --git a/reme/tool/memory/vector/retrieve_memory.py b/reme/tool/memory/vector/retrieve_memory.py new file mode 100644 index 00000000..18f602cc --- /dev/null +++ b/reme/tool/memory/vector/retrieve_memory.py @@ -0,0 +1,185 @@ +"""Retrieve memory from vector store""" + +from loguru import logger + +from ..base_memory_tool import BaseMemoryTool +from ....core.enumeration import MemoryType +from ....core.schema import ToolCall, MemoryNode, VectorNode +from ....core.utils import deduplicate_memories + + +class VectorRetrieveMemory(BaseMemoryTool): + """Tool to retrieve memories from vector store using similarity search""" + + def __init__( + self, + add_memory_type_target: bool = False, + top_k: int = 20, + enable_metadata: bool = True, + **kwargs, + ): + super().__init__(**kwargs) + self.add_memory_type_target: bool = add_memory_type_target + self.top_k: int = top_k + self.enable_metadata: bool = enable_metadata + + def _build_query_schema(self) -> tuple[dict, list[str]]: + """Build query schema for single/multiple retrieval""" + properties = {} + required = [] + + if self.add_memory_type_target: + properties["memory_type"] = { + "type": "string", + "description": "type of memory to search for.", + } + properties["memory_target"] = { + "type": "string", + "description": "target of memory to search within.", + } + required.extend(["memory_type", "memory_target"]) + + properties["query"] = { + "type": "string", + "description": "query text for vector similarity search.", + } + required.append("query") + + if self.enable_metadata: + properties["metadata"] = { + "type": "object", + "description": "optional metadata filters for narrowing search results.", + } + + return properties, required + + def _build_tool_call(self) -> ToolCall: + """Build and return the tool call schema""" + properties, required = self._build_query_schema() + return ToolCall( + **{ + "description": "retrieve memories using vector similarity search.", + "parameters": { + "type": "object", + "properties": properties, + "required": required, + }, + }, + ) + + def _build_multiple_tool_call(self) -> ToolCall: + """Build and return the multiple tool call schema""" + item_properties, item_required = self._build_query_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": { + "type": "object", + "properties": item_properties, + "required": item_required, + }, + }, + }, + "required": ["query_items"], + }, + }, + ) + + async def _retrieve_by_query( + self, + memory_type: str, + memory_target: str, + query: str, + metadata: dict | None = None, + ) -> list[MemoryNode]: + """Retrieve memories by query with filters""" + filter_dict = { + "memory_type": [memory_type], + "memory_target": [memory_target], + } + + if metadata: + for key, value in metadata.items(): + if value: + value = str(value).strip() + filter_dict[key] = [value] if not isinstance(value, list) else value + + nodes: list[VectorNode] = await self.vector_store.search(query=query, limit=self.top_k, filters=filter_dict) + + memory_nodes: list[MemoryNode] = [MemoryNode.from_vector_node(n) for n in nodes] + + filtered_memory_nodes = [ + m for m in memory_nodes if not (m.memory_type == MemoryType.TOOL and m.when_to_use != query) + ] + + return filtered_memory_nodes + + async def execute(self): + default_memory_type: str = self.context.get("memory_type", "") + default_memory_target: str = self.context.get("memory_target", "") + + if self.enable_multiple: + query_items: list[dict] = self.context.get("query_items", []) + if not query_items: + self.output = "No query items provided for retrieval." + return + else: + query = self.context.get("query", "") + if not query: + self.output = "No query provided for retrieval." + return + + query_items = [ + { + "memory_type": default_memory_type, + "memory_target": default_memory_target, + "query": query, + }, + ] + + query_items = [item for item in query_items if item.get("query")] + + if not query_items: + self.output = "No valid query texts provided for retrieval." + return + + memory_nodes: list[MemoryNode] = [] + for item in query_items: + memory_type = item.get("memory_type") or default_memory_type + memory_target = item.get("memory_target") or default_memory_target + metadata = item.get("metadata", {}) if self.enable_metadata else None + + if not memory_type or not memory_target: + logger.warning(f"Skipping query with missing memory_type or memory_target: {item}") + continue + + retrieved = await self._retrieve_by_query( + memory_type=memory_type, + memory_target=memory_target, + query=item["query"], + metadata=metadata, + ) + 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) + + self.memory_nodes = new_memory_nodes + + if not new_memory_nodes: + self.output = "No new memory_nodes found matching the query (duplicates removed)." + else: + self.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") diff --git a/reme/tool/memory/vector/update_memory.py b/reme/tool/memory/vector/update_memory.py new file mode 100644 index 00000000..07b1e41f --- /dev/null +++ b/reme/tool/memory/vector/update_memory.py @@ -0,0 +1,100 @@ +"""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 __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 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.", + }, + "memory_content": { + "type": "string", + "description": "new content of the memory.", + }, + "metadata": { + "type": "object", + "description": "metadata for the memory.", + } + }, + "required": ["memory_id", "memory_content", "metadata"], + }, + }, + }, + "required": ["memories"], + }, + }, + ) + + def _extract_memory_data(self, mem_dict: dict) -> tuple[str, str, dict]: + """Extract memory id, content and metadata from dict""" + memory_id = mem_dict.get("memory_id", "") + memory_content = mem_dict.get("memory_content", "") + raw_metadata = mem_dict.get("metadata", {}) + metadata = {key: str(value).strip() for key, value in raw_metadata.items() if value} + return memory_id, memory_content, metadata + + def _build_memory_node(self, content: str, metadata: dict = None) -> MemoryNode: + """Build a memory node""" + return MemoryNode( + memory_type=self.memory_type, + memory_target=self.memory_target, + content=content, + author=self.author, + metadata=metadata or {}, + ) + + async def execute(self): + old_memory_ids: list[str] = [] + new_memory_nodes: list[MemoryNode] = [] + + memories: list[dict] = self.context.get("memories", []) + if not memories: + self.output = "No memories provided for update." + return + + for mem in memories: + memory_id, memory_content, metadata = self._extract_memory_data(mem) + if not memory_id or not memory_content: + logger.warning(f"Skipping memory with missing id or content: {mem}") + continue + old_memory_ids.append(memory_id) + new_memory_nodes.append(self._build_memory_node(memory_content, metadata=metadata)) + + if not old_memory_ids or not new_memory_nodes: + self.output = "No valid memories provided for update." + return + + vector_nodes = [node.to_vector_node() for node in new_memory_nodes] + new_vector_ids = [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 = new_memory_nodes + + self.output = f"Successfully updated {len(new_memory_nodes)} memories in vector_store." + logger.info(self.output) diff --git a/reme/workflow/procedural_memory/summary/__init__.py b/reme/workflow/procedural_memory/summary/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme_ai/mem_tool/hands_off_tool.py b/reme_ai/mem_tool/hands_off_tool.py deleted file mode 100644 index 2cb28d60..00000000 --- a/reme_ai/mem_tool/hands_off_tool.py +++ /dev/null @@ -1,144 +0,0 @@ -"""Hands-off tool for distributing memory tasks to appropriate agents.""" - -import json -from typing import TYPE_CHECKING - -from loguru import logger - -from .base_memory_tool import BaseMemoryTool -from ..core.context import C -from ..core.enumeration import MemoryType - -if TYPE_CHECKING: - from ..mem_agent import BaseMemoryAgent - - -@C.register_op() -class HandsOffTool(BaseMemoryTool): - """Distribute memory tasks to appropriate agents based on memory_type.""" - - def __init__(self, memory_agents: list["BaseMemoryAgent"], **kwargs): - kwargs["sub_ops"] = memory_agents or [] - super().__init__(**kwargs) - from ..mem_agent import BaseMemoryAgent - - self.sub_ops: list[BaseMemoryAgent] = [a for a in self.sub_ops if isinstance(a, BaseMemoryAgent)] - - @property - def memory_agent_dict(self) -> dict[MemoryType, "BaseMemoryAgent"]: - """Returns a dictionary mapping memory types to their corresponding agents.""" - return {a.memory_type: a for a in self.sub_ops} - - def _build_item_schema(self) -> tuple[dict, list[str]]: - """Build shared schema properties and required fields for memory tasks.""" - properties = { - "memory_type": { - "type": "string", - "description": self.get_prompt("memory_type"), - "enum": [k.value for k in self.memory_agent_dict], - }, - "memory_target": { - "type": "string", - "description": self.get_prompt("memory_target"), - }, - } - required = ["memory_type", "memory_target"] - return properties, required - - def _build_parameters(self) -> dict: - """Build input schema for single memory task distribution.""" - properties, required = self._build_item_schema() - return { - "type": "object", - "properties": properties, - "required": required, - } - - def _build_multiple_parameters(self) -> dict: - """Build input schema for multiple memory task distribution.""" - item_properties, required_fields = self._build_item_schema() - return { - "type": "object", - "properties": { - "memory_tasks": { - "type": "array", - "description": self.get_prompt("memory_tasks"), - "items": { - "type": "object", - "properties": item_properties, - "required": required_fields, - }, - }, - }, - "required": ["memory_tasks"], - } - - @staticmethod - def _parse_memory_type_target(task: dict): - memory_type = task.get("memory_type", "") - memory_target = task.get("memory_target", "") - return {"memory_type": MemoryType(memory_type), "memory_target": memory_target} - - def _collect_tasks(self) -> list[dict]: - """Collect memory tasks from context based on enable_multiple flag.""" - tasks: list[dict] = [] - if self.enable_multiple: - memory_tasks: list[dict] = self.context.get("memory_tasks", []) - for task in memory_tasks: - tasks.append(self._parse_memory_type_target(task)) - else: - tasks.append(self._parse_memory_type_target(self.context)) - return tasks - - async def execute(self): - """Execute memory tasks by distributing to appropriate agents in parallel.""" - tasks = self._collect_tasks() - - if not tasks: - self.output = "No valid memory tasks to execute." - return - - # Submit tasks to corresponding agents - agent_list = [] - for i, task in enumerate(tasks): - memory_type: MemoryType = task["memory_type"] - memory_target: str = task["memory_target"] - - if memory_type not in self.memory_agent_dict: - logger.warning(f"No agent found for memory_type={memory_type}") - continue - - agent = self.memory_agent_dict[memory_type].copy() - agent_list.append([agent, memory_type, memory_target]) - - logger.info(f"Task {i}: Submitting {memory_type.value} agent for target={memory_target}") - self.submit_async_task( - agent.call, - query=self.context.get("query", ""), - messages=self.context.get("messages", []), - memory_type=memory_type, - memory_target=memory_target, - description=self.context.get("description"), - ref_memory_id=self.context.get("ref_memory_id", ""), - ) - - await self.join_async_tasks() - - # Collect results - results = [] - for i, (agent, memory_type, memory_target) in enumerate(agent_list): - result_str = str(agent.output) - if agent.memory_nodes: - self.memory_nodes.extend(agent.memory_nodes) - - results.append( - { - "memory_type": memory_type.value, - "memory_target": memory_target, - "result": result_str[:200] + ("..." if len(result_str) > 200 else ""), - }, - ) - logger.info(f"Task {i}: Completed {memory_type.value} agent for target={memory_target}") - - results_str = json.dumps(results, ensure_ascii=False, indent=2) - self.output = f"Successfully executed {len(results)} memory tasks:\n{results_str}" diff --git a/reme_ai/mem_tool/hands_off_tool.yaml b/reme_ai/mem_tool/hands_off_tool.yaml deleted file mode 100644 index 06a443d7..00000000 --- a/reme_ai/mem_tool/hands_off_tool.yaml +++ /dev/null @@ -1,19 +0,0 @@ -tool: | - Distribute a memory task to the appropriate agent based on memory_type. - Use this tool to hand off memory summarization to specialized agents. - Examples: summarizing user preferences, extracting procedural knowledge, or analyzing tool usage patterns. - -tool_multiple: | - Distribute multiple memory tasks to appropriate agents in parallel. - Use this tool to hand off multiple memory summarization tasks efficiently. - Each task will be processed by its corresponding specialized agent based on memory_type. - -memory_type: | - The type of memory to process. Determines which specialized agent handles the task. - -memory_target: | - The target entity for this memory. - This helps the agent focus on the specific subject of the memory task. - -memory_tasks: | - A list of memory tasks to distribute, each with memory_type and memory_target. diff --git a/reme_ai/mem_tool/read_local_memories.py b/reme_ai/mem_tool/read_local_memories.py deleted file mode 100644 index f239d5d1..00000000 --- a/reme_ai/mem_tool/read_local_memories.py +++ /dev/null @@ -1,54 +0,0 @@ -from loguru import logger - -from ..base_memory_tool import BaseMemoryTool -from ...core.context import C -from ...core.schema.memory_node import MemoryNode - - -@C.register_op() -class ReadLocalMemories(BaseMemoryTool): - - def __init__(self, **kwargs): - kwargs["enable_multiple"] = False - super().__init__(**kwargs) - - def _build_parameters(self) -> dict: - return { - "type": "object", - "properties": { - "memory_type": { - "type": "string", - "description": self.get_prompt("memory_type"), - }, - "memory_target": { - "type": "string", - "description": self.get_prompt("memory_target"), - }, - }, - "required": ["memory_type", "memory_target"], - } - - async def execute(self): - memory_type = self.context.get("memory_type", "") - memory_target = self.context.get("memory_target", "") - - if not memory_type or not memory_target: - self.output = "memory_type and memory_target are required." - return - - cache_key = f"{memory_type}_{memory_target}" - cached_data = self.meta_memory.load(cache_key, auto_clean=False) - - if not cached_data: - self.output = f"Local memory not found: {memory_type}_{memory_target}" - logger.info(self.output) - return - - memory_nodes = [MemoryNode(**node_data) for node_data in cached_data] - - if not memory_nodes: - self.output = f"No valid memory nodes found in {memory_type}_{memory_target}" - return - - self.output = memory_nodes - logger.info(f"Read {len(memory_nodes)} nodes from cache key: {cache_key}") diff --git a/reme_ai/mem_tool/read_local_memories.yaml b/reme_ai/mem_tool/read_local_memories.yaml deleted file mode 100644 index c8e155f6..00000000 --- a/reme_ai/mem_tool/read_local_memories.yaml +++ /dev/null @@ -1,8 +0,0 @@ -tool: | - Read memory nodes from local memory files. - -memory_type: | - The type of local memory to read. - -memory_target: | - The target identifier for the local memory. diff --git a/reme_ai/mem_tool/write_local_memories.py b/reme_ai/mem_tool/write_local_memories.py deleted file mode 100644 index f6336395..00000000 --- a/reme_ai/mem_tool/write_local_memories.py +++ /dev/null @@ -1,57 +0,0 @@ -from loguru import logger - -from ..base_memory_tool import BaseMemoryTool -from ...core.context import C -from ...core.schema.memory_node import MemoryNode - - -@C.register_op() -class WriteLocalMemories(BaseMemoryTool): - - def __init__(self, **kwargs): - kwargs["enable_multiple"] = True - super().__init__(**kwargs) - - def _build_multiple_parameters(self) -> dict: - return { - "type": "object", - "properties": { - "memory_nodes": { - "type": "array", - "description": self.get_prompt("memory_nodes"), - "items": { - "type": "object", - "description": "Memory node object", - }, - }, - }, - "required": ["memory_nodes"], - } - - async def execute(self): - memory_nodes = self.context.get("memory_nodes", []) - - if not memory_nodes: - self.output = "No memory nodes provided." - return - - memory_nodes = [MemoryNode(**node) if isinstance(node, dict) else node for node in memory_nodes] - - grouped = {} - for node in memory_nodes: - key = (node.memory_type.value, node.memory_target) - if key not in grouped: - grouped[key] = [] - grouped[key].append(node) - - written_keys = [] - - for (memory_type, memory_target), nodes in grouped.items(): - cache_key = f"{memory_type}_{memory_target}" - nodes_data = [node.model_dump() for node in nodes] - - self.meta_memory.save(cache_key, nodes_data) - written_keys.append(f"{memory_type}_{memory_target}") - logger.info(f"Saved {len(nodes)} nodes to cache key: {cache_key}") - - self.output = f"Successfully written local memories: {', '.join(written_keys)}" diff --git a/reme_ai/mem_tool/write_local_memories.yaml b/reme_ai/mem_tool/write_local_memories.yaml deleted file mode 100644 index 81615b1b..00000000 --- a/reme_ai/mem_tool/write_local_memories.yaml +++ /dev/null @@ -1,5 +0,0 @@ -tool_multiple: | - Write memory nodes to local memory files. - -memory_nodes: | - List of memory nodes to write to local files.