mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-15 23:31:05 +00:00
94 lines
3.2 KiB
Python
94 lines
3.2 KiB
Python
"""ReMe retriever agent that orchestrates multiple memory agents to retrieve 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 ReMeRetriever(BaseMemoryAgent):
|
|
"""Orchestrate multiple memory agents to retrieve information."""
|
|
|
|
async def build_messages(self) -> list[Message]:
|
|
if self.context.get("query"):
|
|
context = self.context.query
|
|
elif self.context.get("messages"):
|
|
context = self.description + "\n" + format_messages(self.context.messages)
|
|
else:
|
|
raise ValueError("input must have either `query` or `messages`")
|
|
|
|
return [
|
|
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"),
|
|
),
|
|
]
|
|
|
|
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,
|
|
query=self.query,
|
|
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"]
|
|
|
|
answer = []
|
|
success = True
|
|
messages = []
|
|
tools_result = []
|
|
retrieved_nodes = []
|
|
|
|
if tools:
|
|
delegate_task_tool = tools[0]
|
|
agents: list[BaseMemoryAgent] = delegate_task_tool.response.metadata["agents"]
|
|
for agent in agents:
|
|
answer.append(agent.response.answer)
|
|
success = success and agent.response.success
|
|
messages.extend(agent.response.metadata["messages"])
|
|
tools_result.extend(agent.response.metadata["tools"])
|
|
retrieved_nodes.extend(agent.response.metadata["retrieved_nodes"])
|
|
|
|
return {
|
|
"answer": "\n".join(answer),
|
|
"success": True,
|
|
"messages": messages,
|
|
"tools": tools_result,
|
|
"retrieved_nodes": retrieved_nodes,
|
|
}
|