mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-08 22:21:15 +00:00
294 lines
14 KiB
Python
294 lines
14 KiB
Python
"""
|
||
Async React agent operator tailored for retrieval workflows.
|
||
|
||
This module implements a ReAct (Reasoning + Acting) agent pattern that combines
|
||
language model reasoning with tool execution. The agent iteratively:
|
||
1. Processes and manages conversation context (compaction/compression)
|
||
2. Generates reasoning and tool calls via LLM
|
||
3. Executes tools (e.g., file search, reading)
|
||
4. Incorporates tool results back into the conversation
|
||
5. Repeats until a final answer is reached or max_steps is exceeded
|
||
|
||
The agent is specifically designed for RAG (Retrieval-Augmented Generation) workflows,
|
||
providing context management capabilities to handle long conversations efficiently.
|
||
|
||
Context management is controlled via ``working_summary_mode`` and
|
||
``compact_ratio_threshold`` parameters, which are forwarded to
|
||
``MessageOffloadOp``. ``working_summary_mode`` selects between:
|
||
- ``compact`` – only compact verbose tool messages by storing full content externally
|
||
and keeping short previews in the context.
|
||
- ``compress`` – only apply LLM-based compression to generate a compact state snapshot.
|
||
- ``auto`` – first run compaction, then optionally run compression if the
|
||
compaction ratio is not sufficient (default).
|
||
|
||
``compact_ratio_threshold`` is only used in ``auto`` mode and defines the compaction
|
||
ratio (tokens after compaction divided by original tokens) above which an additional
|
||
LLM-based compression pass is applied. It defaults to ``0.75``.
|
||
"""
|
||
|
||
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 loguru import logger
|
||
|
||
|
||
@C.register_op()
|
||
class AgenticRetrieveOp(BaseAsyncToolOp):
|
||
"""
|
||
ReAct agent that exposes RAG-friendly tools and context policies.
|
||
|
||
This agent implements the ReAct pattern where the LLM alternates between:
|
||
- Reasoning: Analyzing the problem and deciding what to do
|
||
- Acting: Executing tools to gather information
|
||
- Observing: Processing tool results and updating understanding
|
||
|
||
The agent supports sophisticated context management to handle long conversations
|
||
through compaction (storing large tool messages externally) and compression
|
||
(LLM-based summarization of message history).
|
||
|
||
Context management behavior is configured via ``working_summary_mode`` and
|
||
``compact_ratio_threshold`` (see module docstring for details). These options are
|
||
passed to ``MessageOffloadOp`` to control whether the agent only compacts tool
|
||
messages, only compresses history, or applies an automatic compaction-then-
|
||
compression pipeline.
|
||
|
||
Available tools:
|
||
- GrepOp: Search for patterns in files
|
||
- ReadFileOp: Read file contents
|
||
|
||
The agent automatically manages context size by offloading large messages and
|
||
compressing conversation history when token limits are approached.
|
||
"""
|
||
|
||
file_path: str = __file__
|
||
|
||
def __init__(
|
||
self,
|
||
llm: str = "qwen3_30b_instruct",
|
||
max_steps: int = 20,
|
||
**kwargs,
|
||
):
|
||
"""
|
||
Initialize the agent runtime configuration.
|
||
|
||
Args:
|
||
llm: Identifier for the language model to use. Defaults to "qwen3_30b_instruct".
|
||
max_steps: Maximum number of reasoning-action cycles before stopping.
|
||
Each cycle includes: context management -> LLM reasoning -> tool execution.
|
||
Defaults to 5 steps.
|
||
**kwargs: Additional arguments passed to the base BaseAsyncToolOp class.
|
||
"""
|
||
super().__init__(llm=llm, **kwargs)
|
||
# Maximum number of ReAct iterations (reasoning + tool execution cycles)
|
||
self.max_steps: int = max_steps
|
||
|
||
def build_tool_call(self) -> ToolCall:
|
||
"""
|
||
Expose metadata describing how to invoke the agent.
|
||
|
||
This method defines the tool schema that other components use to invoke
|
||
this agent. It specifies all input parameters including conversation messages
|
||
and context management configuration.
|
||
|
||
Returns:
|
||
ToolCall: A schema object describing the agent's interface, including
|
||
all parameters for message handling and context management.
|
||
"""
|
||
return ToolCall(
|
||
**{
|
||
"description": "A React agent that answers user queries.",
|
||
"input_schema": {
|
||
"messages": {
|
||
"type": "array",
|
||
"description": "messages",
|
||
"required": True,
|
||
},
|
||
"working_summary_mode": {
|
||
"type": "string",
|
||
"description": "summary strategy: 'compact' only compacts large tool messages, 'compress' "
|
||
"only applies LLM-based compression, 'auto' first compacts then optionally compresses when "
|
||
"reduction is insufficient. Defaults to 'auto'.",
|
||
"required": False,
|
||
"enum": ["compact", "compress", "auto"],
|
||
},
|
||
"compact_ratio_threshold": {
|
||
"type": "number",
|
||
"description": "Only used in 'auto' mode. Threshold for compaction (tokens after compaction "
|
||
"divided by original tokens). When the ratio is greater than this value, an additional "
|
||
"LLM-based compression pass is triggered. Defaults to 0.75.",
|
||
"required": False,
|
||
},
|
||
"max_total_tokens": {
|
||
"type": "integer",
|
||
"description": "Maximum token threshold for triggering compression/compaction. For compaction "
|
||
"this is total tokens; for compression this excludes keep_recent_count and "
|
||
"system messages. Defaults to 20000.",
|
||
"required": False,
|
||
},
|
||
"max_tool_message_tokens": {
|
||
"type": "integer",
|
||
"description": "Maximum token count per tool message before compaction applies. Exceeding "
|
||
"messages store full content externally with a preview in context. Defaults "
|
||
"to 2000.",
|
||
"required": False,
|
||
},
|
||
"group_token_threshold": {
|
||
"type": "integer",
|
||
"description": "Maximum tokens per compression group for LLM-based compression. None/0 "
|
||
"compresses all messages together. Oversized messages form their own group. "
|
||
"Used in 'compress' or 'auto' mode.",
|
||
"required": False,
|
||
},
|
||
"keep_recent_count": {
|
||
"type": "integer",
|
||
"description": "Number of recent messages preserved without compression/compaction. Defaults "
|
||
"to 1 for compaction and 2 for compression.",
|
||
"required": False,
|
||
},
|
||
"store_dir": {
|
||
"type": "string",
|
||
"description": "Directory for storing offloaded contents. Required for compaction/compression "
|
||
"to save full tool messages and compressed groups.",
|
||
"required": False,
|
||
},
|
||
"chat_id": {
|
||
"type": "string",
|
||
"description": "Chat session identifier for naming stored files. Defaults to auto-generated "
|
||
"UUID if omitted.",
|
||
"required": False,
|
||
},
|
||
},
|
||
},
|
||
)
|
||
|
||
async def async_execute(self):
|
||
"""
|
||
Main execution loop implementing the ReAct (Reasoning + Acting) pattern.
|
||
|
||
This method orchestrates the iterative agent workflow:
|
||
1. Context Management: Compact/compress conversation history if needed
|
||
2. Reasoning: LLM analyzes context and decides on actions (may include tool calls)
|
||
3. Acting: Execute requested tools in parallel
|
||
4. Observing: Incorporate tool results back into conversation
|
||
5. Repeat until final answer (no tool calls) or max_steps reached
|
||
|
||
The loop continues until:
|
||
- The LLM produces a response without tool calls (final answer)
|
||
- Maximum number of steps (max_steps) is reached
|
||
|
||
Each iteration is called a "round" and represents one reasoning-action cycle.
|
||
"""
|
||
# Import tool operators that the agent can use
|
||
from reme_ai.retrieve.working import GrepOp, ReadFileOp, BatchWriteFileOp
|
||
from reme_ai.summary.working import MessageOffloadOp
|
||
|
||
# Initialize available tools for the agent
|
||
# GrepOp: Search for patterns/text in files (useful for code search)
|
||
grep_op = GrepOp(language=self.language)
|
||
# ReadFileOp: Read and return file contents
|
||
read_file_op = ReadFileOp(language=self.language)
|
||
|
||
# Create a dictionary mapping tool names to their operator instances
|
||
# This allows quick lookup when the LLM requests a specific tool
|
||
tool_op_dict: Dict[str, BaseAsyncToolOp] = {
|
||
grep_op.tool_call.name: grep_op,
|
||
read_file_op.tool_call.name: read_file_op,
|
||
}
|
||
|
||
# Convert input messages to Message objects for processing
|
||
messages = [Message(**x) for x in self.context.messages]
|
||
|
||
# Extract context management parameters from input, excluding messages
|
||
# These will be passed to the context management pipeline
|
||
context_kwargs = self.input_dict.copy()
|
||
context_kwargs.pop("messages", None)
|
||
|
||
# Main ReAct loop: iterate up to max_steps times
|
||
for i in range(self.max_steps):
|
||
# Step 1: Context Management Phase
|
||
# Create a pipeline: MessageOffloadOp (compacts/compresses) -> BatchWriteFileOp (saves offloaded content)
|
||
# The >> operator chains these operations together
|
||
op = MessageOffloadOp() >> BatchWriteFileOp()
|
||
|
||
# Apply context management to current message history
|
||
# This may compact large tool messages or compress old messages based on context_manage_mode
|
||
await op.async_call(messages=[x.simple_dump() for x in messages], **context_kwargs)
|
||
|
||
# Update messages with the processed/optimized version from context management
|
||
# Large messages may now reference external files instead of containing full content
|
||
logger.info(f"round{i + 1}.offload={op.context.response.answer}")
|
||
messages = [Message(**x) for x in op.context.response.answer]
|
||
|
||
# Step 2: Reasoning Phase
|
||
# LLM analyzes the conversation context and decides what to do
|
||
# It may generate a final answer or request tool calls to gather more information
|
||
assistant_message: Message = await self.llm.achat(
|
||
messages=messages, # Current conversation history (possibly optimized)
|
||
tools=[op.tool_call for op in tool_op_dict.values()], # Available tools the LLM can use
|
||
)
|
||
|
||
# Add the LLM's response to the conversation history
|
||
messages.append(assistant_message)
|
||
logger.info(f"round{i + 1}.assistant={assistant_message.model_dump_json()}")
|
||
|
||
# Step 3: Check if we have a final answer
|
||
# If the LLM didn't request any tools, it has provided a final answer
|
||
# Exit the loop as the agent's work is complete
|
||
if not assistant_message.tool_calls:
|
||
break
|
||
|
||
# Step 4: Acting Phase - Execute requested tools
|
||
# The LLM requested one or more tools to gather information
|
||
# We'll execute them in parallel for efficiency
|
||
op_list: List[BaseAsyncToolOp] = []
|
||
|
||
# Process each tool call requested by the LLM
|
||
for j, tool_call in enumerate(assistant_message.tool_calls):
|
||
# Validate that the requested tool exists in our available tools
|
||
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}")
|
||
|
||
# Create a copy of the tool operator for this specific tool call
|
||
# Each tool call needs its own operator instance with the correct call ID
|
||
op_copy: BaseAsyncToolOp = tool_op_dict[tool_call.name].copy()
|
||
op_copy.tool_call.id = tool_call.id # Match the ID from LLM's tool call request
|
||
op_list.append(op_copy)
|
||
|
||
# Submit the tool execution as an async task (runs in parallel with other tools)
|
||
# The op_copy instance is used to execute the tool with the specific call ID and arguments
|
||
self.submit_async_task(op_copy.async_call, **tool_call.argument_dict)
|
||
|
||
# Wait for all submitted tool executions to complete
|
||
# This ensures we have all results before proceeding
|
||
await self.join_async_task()
|
||
|
||
# Step 5: Observing Phase - Incorporate tool results into conversation
|
||
# Process each completed tool execution and add results to message history
|
||
for j, op in enumerate(op_list):
|
||
# Extract the tool execution result as a string
|
||
tool_result = str(op.output)
|
||
|
||
# Create a tool message with the result, linked to the original tool call via ID
|
||
# This allows the LLM to associate results with the specific tool calls it made
|
||
tool_message = Message(role=Role.TOOL, content=tool_result, tool_call_id=op.tool_call.id)
|
||
messages.append(tool_message)
|
||
|
||
# Log the tool result (truncated to first 200 chars for readability)
|
||
logger.info(f"round{i + 1}.{j} join tool_result={tool_result[:200]}...\n\n")
|
||
|
||
# Loop continues: next iteration will process these tool results,
|
||
# manage context again if needed, and let LLM reason about the new information
|
||
|
||
# After loop completes, set the final output
|
||
# The last message should be the LLM's final answer (without tool calls)
|
||
self.set_output(messages[-1].content)
|
||
|
||
# Store the complete conversation history in the context response
|
||
# This includes all reasoning steps, tool calls, and tool results
|
||
self.context.response.metadata["messages"] = [x.simple_dump(add_reasoning=True) for x in messages]
|