diff --git a/reme_ai/agent/react/agentic_retrieve_op.py b/reme_ai/agent/react/agentic_retrieve_op.py index ee29bd1a..5d8d0686 100644 --- a/reme_ai/agent/react/agentic_retrieve_op.py +++ b/reme_ai/agent/react/agentic_retrieve_op.py @@ -3,9 +3,11 @@ from typing import Dict, List from flowllm.core.context import C +from flowllm.core.enumeration import Role from flowllm.core.op import BaseAsyncToolOp from flowllm.core.schema import ToolCall, Message from flowllm.gallery.agent import ReactAgentOp +from loguru import logger @C.register_op() @@ -18,10 +20,11 @@ class AgenticRetrieveOp(ReactAgentOp): self, llm: str = "qwen3_30b_instruct", max_steps: int = 5, - add_think_tool: bool = False, **kwargs, ): - super().__init__(llm=llm, max_steps=max_steps, add_think_tool=add_think_tool, **kwargs) + """Initialize the agent runtime configuration.""" + super().__init__(llm=llm, **kwargs) + self.max_steps: int = max_steps def build_tool_call(self) -> ToolCall: """Expose metadata describing how to invoke the agent.""" @@ -112,3 +115,49 @@ class AgenticRetrieveOp(ReactAgentOp): messages = op.context.response.answer messages = [Message(**x) for x in messages] return messages + + async def execute_tool(self, op: BaseAsyncToolOp, tool_call: ToolCall): + """Execute a tool operation asynchronously using the provided tool call arguments.""" + self.submit_async_task(op.async_call, **tool_call.argument_dict) + + async def async_execute(self): + """Main execution loop that alternates LLM calls and tool invocations.""" + tool_op_dict = await self.build_tool_op_dict() + messages = await self.build_messages() + + for i in range(self.max_steps): + messages = await self.before_chat(messages) + + assistant_message: Message = await self.llm.achat( + messages=messages, + tools=[op.tool_call for op in tool_op_dict.values()], + ) + messages.append(assistant_message) + logger.info(f"round{i + 1}.assistant={assistant_message.model_dump_json()}") + + if not assistant_message.tool_calls: + break + + op_list: List[BaseAsyncToolOp] = [] + for j, tool_call in enumerate(assistant_message.tool_calls): + if tool_call.name not in tool_op_dict: + logger.exception(f"unknown tool_call.name={tool_call.name}") + continue + + logger.info(f"round{i + 1}.{j} submit tool_calls={tool_call.name} argument={tool_call.argument_dict}") + + op_copy: BaseAsyncToolOp = tool_op_dict[tool_call.name].copy() + op_copy.tool_call.id = tool_call.id + op_list.append(op_copy) + await self.execute_tool(op_copy, tool_call) + + await self.join_async_task() + + for j, op in enumerate(op_list): + tool_result = str(op.output) + tool_message = Message(role=Role.TOOL, content=tool_result, tool_call_id=op.tool_call.id) + messages.append(tool_message) + logger.info(f"round{i + 1}.{j} join tool_result={tool_result[:200]}...\n\n") + + self.set_output(messages[-1].content) + self.context.response.messages = messages