ReMe/reme/memory/vector_based/reme_summarizer.py

98 lines
3.3 KiB
Python

"""ReMe summarizer agent that orchestrates multiple memory agents to summarize information."""
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."""
async def build_messages(self) -> list[Message]:
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,
author=self.author,
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=self.meta_memory_info,
context=context.strip(),
),
),
Message(
role=Role.USER,
content=self.get_prompt("user_message"),
),
]
return messages
async def _acting_step(
self,
assistant_message: Message,
tools: list[BaseTool],
step: int,
stage: str = "",
**kwargs,
) -> tuple[list[BaseTool], list[Message]]:
return await super()._acting_step(
assistant_message,
tools,
step,
description=self.description,
messages=self.messages,
history_node=self.history_node,
author=self.author,
**kwargs,
)
async def react(self, messages: list[Message], tools: list["BaseTool"], stage: str = ""):
"""Run single ReAct step - only one tool call iteration."""
used_tools: list[BaseTool] = []
assistant_message, should_act = await self._reasoning_step(messages, tools, step=0, stage=stage)
success = True
if should_act:
t_tools, tool_messages = await self._acting_step(assistant_message, tools, step=0, stage=stage)
used_tools.extend(t_tools)
messages.extend(tool_messages)
return used_tools, messages, success
async def execute(self):
result = await super().execute()
tools: list[BaseTool] = result["tools"]
success = True
messages = []
tools_result = []
memory_nodes = []
if tools:
delegate_task_tool = tools[0]
agents: list[BaseMemoryAgent] = delegate_task_tool.response.metadata["agents"]
for agent in agents:
success = success and agent.response.success
messages.extend(agent.response.metadata["messages"])
tools_result.extend(agent.response.metadata["tools"])
memory_nodes.extend(agent.response.metadata["memory_nodes"])
return {
"answer": memory_nodes,
"success": True,
"messages": messages,
"tools": tools_result,
}