refactor(memory): update memory reference handling and agent orchestration

This commit is contained in:
jinli.yl 2026-01-28 01:47:49 +08:00
parent 2fffb847a0
commit afa4eb9114
9 changed files with 58 additions and 81 deletions

View file

@ -10,10 +10,6 @@ from ....core.utils import format_messages
class ReMeRetriever(BaseMemoryAgent):
"""Orchestrate multiple memory agents to retrieve information."""
def __init__(self, meta_memories: list[dict] | None = None, **kwargs):
super().__init__(**kwargs)
self.meta_memories: list[dict] = meta_memories or []
async def build_messages(self) -> list[Message]:
if self.context.get("query"):
context = self.context.query
@ -27,7 +23,7 @@ class ReMeRetriever(BaseMemoryAgent):
role=Role.SYSTEM,
content=self.prompt_format(
prompt_name="system_prompt",
meta_memory_info=await self.read_meta_memories(self.meta_memories),
meta_memory_info=self.meta_memory_info,
context=context.strip(),
),
),
@ -58,21 +54,14 @@ class ReMeRetriever(BaseMemoryAgent):
async def react(self, messages: list[Message], tools: list["BaseTool"], stage: str = ""):
"""Run single ReAct step - only one tool call iteration."""
success: bool = False
used_tools: list[BaseTool] = []
# Reasoning: LLM decides next action
assistant_message, should_act = await self._reasoning_step(messages, tools, step=0, stage=stage)
success = True
if should_act:
# Acting: execute tools and collect results (only once)
t_tools, tool_messages = await self._acting_step(assistant_message, tools, step=0, stage=stage)
used_tools.extend(t_tools)
messages.extend(tool_messages)
success = True
else:
# No tools requested
success = True
return used_tools, messages, success
@ -87,7 +76,6 @@ class ReMeRetriever(BaseMemoryAgent):
messages = []
tools = []
retrieved_nodes = []
for agent in agents:
answer.append(agent.response.answer)
success = success and agent.response.success

View file

@ -1,23 +1,23 @@
system_prompt: |
You are a Memory Orchestrator responsible for routing memory retrieval tasks to specialized agents based on the user query.
You are a Memory Orchestrator responsible for routing memory retrieval tasks to specialized agents based on the context.
# User Query
# Context
{context}
## Available Memory Agents
Each line indicates a specialized Memory Agent dedicated to storing and retrieving memories within a specific dimension <memory_type>(<memory_target>).
Format: "- <memory_type>(<memory_target>): <description>"
Each line indicates a specialized Memory Agent that is an expert for retrieving memories about a specific memory_target.
{meta_memory_info}
## Your Task
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
- Do NOT query agents that don't exist above
3. Multiple tasks can be specified to enable parallel retrieval from specialized agents
Analyze the context and delegate retrieval tasks to appropriate specialized agents:
1. Examine the context content and identify which memory_target(s) are relevant for retrieving information
2. For each relevant memory_target, delegate the retrieval task to its corresponding specialized agent
- The memory_target must **exactly match** existing entries in the "Available Memory Agents" listed above
- Do NOT delegate to agents that don't exist above
- Each memory_target should be assigned **only once** - do not duplicate assignments
3. Use the `delegate_task` tool **once** with all relevant memory_target(s) to enable parallel processing by specialized agents
Note: If the retrieved information is insufficient to answer the query, respond: "nothing found after thorough search."
Note: If the context contains no memorable information (e.g., simple greetings), return `<NO_MEMORY_NEEDED>`.
user_message: |
Please analyze the user query and retrieve relevant information from the appropriate existing agents.
Please analyze the context and delegate retrieval tasks to the appropriate specialized agents.

View file

@ -4,37 +4,30 @@ from ..base_memory_agent import BaseMemoryAgent
from ....core.enumeration import Role
from ....core.op import BaseTool
from ....core.schema import Message
from ....core.utils import format_messages
class ReMeSummarizer(BaseMemoryAgent):
"""Orchestrates multiple memory agents to summarize and store information across different memory types."""
def __init__(self, meta_memories: list[dict] | None = None, **kwargs):
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()
add_history_tool: BaseTool | None = self.pop_tool("add_history")
if add_history_tool is not None:
await add_history_tool.call(
messages=self.messages,
description=self.description,
service_context=self.service_context,
)
self.context.history_node = add_history_tool.context.history_node
context = self.context.description + "\n" + format_messages(self.context.messages)
messages = [
Message(
role=Role.SYSTEM,
content=self.prompt_format(
prompt_name="system_prompt",
meta_memory_info=await self.read_meta_memories(self.meta_memories),
context=self.context.history_node.content,
meta_memory_info=self.meta_memory_info,
context=context.strip(),
),
),
Message(
@ -66,21 +59,14 @@ class ReMeSummarizer(BaseMemoryAgent):
async def react(self, messages: list[Message], tools: list["BaseTool"], stage: str = ""):
"""Run single ReAct step - only one tool call iteration."""
success: bool = False
used_tools: list[BaseTool] = []
# Reasoning: LLM decides next action
assistant_message, should_act = await self._reasoning_step(messages, tools, step=0, stage=stage)
success = True
if should_act:
# Acting: execute tools and collect results (only once)
t_tools, tool_messages = await self._acting_step(assistant_message, tools, step=0, stage=stage)
used_tools.extend(t_tools)
messages.extend(tool_messages)
success = True
else:
# No tools requested
success = True
return used_tools, messages, success

View file

@ -1,23 +1,23 @@
system_prompt: |
You are a Memory Orchestrator responsible for routing memory tasks to specialized agents based on the context.
You are a Memory Orchestrator responsible for routing memory summarization 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: "- <memory_type>(<memory_target>): <description>"
Each line indicates a specialized Memory Agent that is an expert for summarizing memories about a specific memory_target.
{meta_memory_info}
## Your Task
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
- 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
Analyze the context and delegate summarization tasks to appropriate specialized agents:
1. Examine the context content and identify which memory_target(s) are relevant for storing information
2. For each relevant memory_target, delegate the summarization task to its corresponding specialized agent
- The memory_target must **exactly match** existing entries in the "Available Memory Agents" listed above
- Do NOT delegate to agents that don't exist above
- Each memory_target should be assigned **only once** - do not duplicate assignments
3. Use the `delegate_task` tool **once** with all relevant memory_target(s) to enable parallel processing by specialized agents
Note: If the context contains no memorable information (e.g., simple greetings), return `<NO_MEMORY_NEEDED>`.
user_message: |
Please analyze the context and route memory tasks to the appropriate existing agents.
Please analyze the context and delegate summarization tasks to the appropriate specialized agents.

View file

@ -38,6 +38,13 @@ class BaseReact(BaseOp):
"""Return available tools for the agent."""
return self.sub_ops
def pop_tool(self, name: str) -> "BaseTool | None":
"""Remove and return a tool from self.tools by name."""
for i, tool in enumerate(self.sub_ops):
if tool.tool_call.name == name:
return self.sub_ops.pop(i)
return None
async def build_messages(self) -> list[Message]:
"""Build initial message list from context query or messages."""
if self.context.get("query"):

View file

@ -116,7 +116,7 @@ class AddMemory(BaseMemoryTool):
"content": memory_content,
"when_to_use": when_to_use,
"message_time": message_time,
"ref_memory_id": self.history_node.memory_id,
"ref_memory_id": self.history_id,
"author": self.author,
"metadata": metadata,
})

View file

@ -71,9 +71,11 @@ class BaseMemoryTool(BaseTool, metaclass=ABCMeta):
raise ValueError("memory_target is not specified in context or memory_target_type_mapping!")
@property
def history_node(self) -> MemoryNode:
def history_id(self) -> str:
"""Get the history node from context."""
return self.context.history_node
if "history_node" in self.context:
return self.context.history_node.memory_id
return ""
@property
def retrieved_nodes(self) -> list[MemoryNode]:

View file

@ -33,16 +33,10 @@ class DelegateTask(BaseMemoryTool):
"properties": {
"tasks": {
"type": "array",
"description": "tasks to delegate to specific agents",
"description": "tasks to delegate to specific agents, each task is a memory_target",
"items": {
"type": "object",
"properties": {
"task_name": {
"type": "string",
"description": "task_name",
},
},
"required": ["task_name"],
"type": "string",
"description": "memory_target to delegate to specific agents",
},
},
},
@ -58,13 +52,13 @@ class DelegateTask(BaseMemoryTool):
# Submit tasks to agents
agent_list: list[BaseMemoryAgent] = []
for i, task in enumerate(tasks):
memory_type = self.memory_target_type_mapping[task]
for i, memory_target in enumerate(tasks):
memory_type = self.memory_target_type_mapping[memory_target]
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}
logger.info(f"Task {i}: {memory_type.value} agent for {memory_target}")
task_kwargs = {"memory_target": memory_target}
for k in ["query", "messages", "description", "history_node"]:
if k in self.context:
task_kwargs[k] = self.context[k]
@ -76,7 +70,7 @@ class DelegateTask(BaseMemoryTool):
for agent in agent_list:
results.append(f"Task: {agent.memory_target}\n{agent.response.answer}")
logger.info(f"Completed {len(results)} task(s)")
logger.info(f"Completed {len(results)} memory_target(s)")
return {
"answer": "\n\n".join(results),
"agents": agent_list,

View file

@ -82,7 +82,7 @@ class UpdateProfile(BaseMemoryTool):
# 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)
new_nodes = profile_handler.add_batch(profiles=profiles_to_add, ref_memory_id=self.history_id)
self.memory_nodes.extend(new_nodes)
added_count = len(new_nodes)