From 2ff75fd580c07718c780addfeb67ceb9b7eebae5 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Thu, 20 Nov 2025 17:01:29 +0800 Subject: [PATCH] feat(agent): add agentic retrieve operator with RAG-friendly tools --- cookbook/context_demo.py | 7 ++ reme_ai/agent/react/__init__.py | 2 + reme_ai/agent/react/agentic_retrieve_op.py | 113 ++++++++++++++++++ .../agent/react/agentic_retrieve_prompt.yaml | 5 + reme_ai/agent/react/simple_react_op.py | 2 +- reme_ai/config/default.yaml | 5 +- reme_ai/context/file_tool/grep_prompt.yaml | 10 ++ .../context/file_tool/read_file_prompt.yaml | 8 ++ reme_ai/context/offload/context_compact_op.py | 6 +- .../context/offload/context_compress_op.py | 8 +- test/http_client_test.py | 23 +++- 11 files changed, 182 insertions(+), 7 deletions(-) create mode 100644 cookbook/context_demo.py create mode 100644 reme_ai/agent/react/agentic_retrieve_op.py create mode 100644 reme_ai/agent/react/agentic_retrieve_prompt.yaml diff --git a/cookbook/context_demo.py b/cookbook/context_demo.py new file mode 100644 index 00000000..c57eb933 --- /dev/null +++ b/cookbook/context_demo.py @@ -0,0 +1,7 @@ +def main(): + ... + + +if __name__ == "__main__": + def main(): + ... diff --git a/reme_ai/agent/react/__init__.py b/reme_ai/agent/react/__init__.py index afb89a82..11adaed4 100644 --- a/reme_ai/agent/react/__init__.py +++ b/reme_ai/agent/react/__init__.py @@ -4,8 +4,10 @@ This module provides ReAct (Reasoning and Acting) agent implementations for answering user queries through iterative reasoning and search actions. """ +from .agentic_retrieve_op import AgenticRetrieveOp from .simple_react_op import SimpleReactOp __all__ = [ + "AgenticRetrieveOp", "SimpleReactOp", ] diff --git a/reme_ai/agent/react/agentic_retrieve_op.py b/reme_ai/agent/react/agentic_retrieve_op.py new file mode 100644 index 00000000..0e792617 --- /dev/null +++ b/reme_ai/agent/react/agentic_retrieve_op.py @@ -0,0 +1,113 @@ +"""Async React agent operator tailored for retrieval workflows.""" + +from typing import Dict, List + +from flowllm.core.context import C +from flowllm.core.op import BaseAsyncToolOp +from flowllm.core.schema import ToolCall, Message +from flowllm.gallery.agent import ReactAgentOp + + +@C.register_op() +class AgenticRetrieveOp(ReactAgentOp): + """React agent that exposes RAG-friendly tools and context policies.""" + + file_path: str = __file__ + + def __init__( + 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) + + def build_tool_call(self) -> ToolCall: + """Expose metadata describing how to invoke the agent.""" + return ToolCall( + **{ + "description": "A React agent that answers user queries.", + "input_schema": { + "messages": { + "type": "array", + "description": "messages", + "required": True, + }, + "context_manage_mode": { + "type": "string", + "description": "Context management mode: 'compact' (only compacts tool messages), 'compress' " + "(only LLM-based compression), 'auto' (compaction first then compression if " + "needed). Defaults to 'auto'.", + "required": True, + "enum": ["compact", "compress", "auto"], + }, + "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, + }, + }, + }, + ) + + def build_tool_op_dict(self) -> dict: + """Collect available tool operators from the execution context.""" + from reme_ai.context.file_tool import GrepOp, ReadFileOp + + grep_op = GrepOp(language=self.language) + read_file_op = ReadFileOp(language=self.language) + tool_op_dict: Dict[str, BaseAsyncToolOp] = { + grep_op.tool_call.name: grep_op, + read_file_op.tool_call.name: read_file_op, + } + + return tool_op_dict + + def build_messages(self) -> List[Message]: + """Build the initial message history for the LLM.""" + return self.context.messages + + async def before_chat(self, messages: List[Message]): + """Run context offload to trim prior messages before invoking the agent.""" + from reme_ai.context.offload import ContextOffloadOp + + op = ContextOffloadOp() + await op.async_call(**self.input_dict) + messages = op.context.response.answer + messages = [Message(**x) for x in messages] + return messages diff --git a/reme_ai/agent/react/agentic_retrieve_prompt.yaml b/reme_ai/agent/react/agentic_retrieve_prompt.yaml new file mode 100644 index 00000000..84e65a6c --- /dev/null +++ b/reme_ai/agent/react/agentic_retrieve_prompt.yaml @@ -0,0 +1,5 @@ +system_prompt: | + You are a helpful assistant. The current time is {time}. + +system_prompt_zh: | + 你是一个有用的助手。当前时间是 {time}。 diff --git a/reme_ai/agent/react/simple_react_op.py b/reme_ai/agent/react/simple_react_op.py index ae1157ce..8c04e87d 100644 --- a/reme_ai/agent/react/simple_react_op.py +++ b/reme_ai/agent/react/simple_react_op.py @@ -8,7 +8,7 @@ reasoning and search actions. import asyncio from flowllm.core.context import C, FlowContext -from flowllm.gallery import ReactSearchOp +from flowllm.gallery.agent import ReactSearchOp @C.register_op() diff --git a/reme_ai/config/default.yaml b/reme_ai/config/default.yaml index 4c8fdc81..52330320 100644 --- a/reme_ai/config/default.yaml +++ b/reme_ai/config/default.yaml @@ -156,6 +156,9 @@ flow: description: "user query" required: true + agentic_retrieve: + flow_content: AgenticRetrieveOp() + context_offload: flow_content: ContextOffloadOp() >> BatchWriteFileOp() description: "Manages context window limits by compacting tool messages and compressing conversation history. First compacts large tool messages by storing full content in external files, then applies LLM-based compression if compaction ratio exceeds threshold. This helps reduce token usage while preserving important information." @@ -205,7 +208,7 @@ flow: context_manage_mode: type: string description: "Context management mode: 'compact' only applies compaction to tool messages, 'compress' only applies LLM-based compression, 'auto' applies compaction first then compression if compaction ratio exceeds threshold. Defaults to 'auto'." - required: false + required: true enum: ["compact", "compress", "auto"] max_total_tokens: type: integer diff --git a/reme_ai/context/file_tool/grep_prompt.yaml b/reme_ai/context/file_tool/grep_prompt.yaml index 0c9a256b..0f49feac 100644 --- a/reme_ai/context/file_tool/grep_prompt.yaml +++ b/reme_ai/context/file_tool/grep_prompt.yaml @@ -1,15 +1,25 @@ tool_desc: | + 一个支持完整正则语法(如 "log.*Error"、"function\\s+\\w+")的高效搜索工具,可用 glob 过滤文件并限制返回结果,适合跨文件快速定位代码或文本内容。 +tool_desc_zh: | A powerful search tool for finding patterns in files using regular expressions. Supports full regex syntax (e.g., "log.*Error", "function\\s+\\w+"), glob pattern filtering, and result limiting. Ideal for searching code or text content across multiple files. pattern: | + 需要在文件内容中匹配的正则表达式模式。 +pattern_zh: | The regular expression pattern to search for in file contents. path: | + 可选:要执行搜索的目录,默认为当前工作目录。 +path_zh: | Optional: The directory to search in. Defaults to current working directory. glob: | + 可选:用于过滤目标文件的 glob 模式(如 "*.js"、"*.{ts,tsx}")。 +glob_zh: | Optional: Glob pattern to filter files (e.g., "*.js", "*.{ts,tsx}"). limit: | + 可选:设置最多返回的匹配行数;未指定时会返回所有匹配。 +limit_zh: | Optional: Maximum number of matching lines to return. Shows all matches if not specified. diff --git a/reme_ai/context/file_tool/read_file_prompt.yaml b/reme_ai/context/file_tool/read_file_prompt.yaml index 52f6a466..eb083182 100644 --- a/reme_ai/context/file_tool/read_file_prompt.yaml +++ b/reme_ai/context/file_tool/read_file_prompt.yaml @@ -1,12 +1,20 @@ tool_desc: | + 读取指定文件内容的工具;对文本文件可通过 offset 与 limit 获取特定行区间,便于分页浏览大文件。 +tool_desc_zh: | Reads and returns the content of a specified file. For text files, it can read specific line ranges using the 'offset' and 'limit' parameters. Use offset and limit to paginate through large files. absolute_path: | + 必填:待读取文件的绝对路径(如 "/home/user/project/file.txt"),不支持相对路径。 +absolute_path_zh: | The absolute path to the file to read (e.g., '/home/user/project/file.txt'). Relative paths are not supported. You must provide an absolute path. offset: | + 可选:文本文件起始读取的 0 基行号;需与 limit 同时使用,适合分页查看大文件。 +offset_zh: | Optional: For text files, the 0-based line number to start reading from. Requires 'limit' to be set. Use for paginating through large files. limit: | + 可选:文本文件最多读取的行数;与 offset 配合实现分页,若仅设 offset 则会读到文件末尾。 +limit_zh: | Optional: For text files, maximum number of lines to read. Use with 'offset' to paginate through large files. If omitted when offset is provided, reads from offset to the end of the file. diff --git a/reme_ai/context/offload/context_compact_op.py b/reme_ai/context/offload/context_compact_op.py index 7bd878a1..6e4031b8 100644 --- a/reme_ai/context/offload/context_compact_op.py +++ b/reme_ai/context/offload/context_compact_op.py @@ -48,11 +48,13 @@ class ContextCompactOp(BaseAsyncOp): assert max_total_tokens > 0, "max_total_tokens must be greater than 0" assert max_tool_message_tokens > 0, "max_tool_message_tokens must be greater than 0" assert preview_char_length >= 0, "preview_char_length must be greater than 0" - assert keep_recent_count > 0, "keep_recent_count must be greater than 0" + assert keep_recent_count >= 0, "keep_recent_count must be greater than 0" # Convert context messages to Message objects messages = [Message(**x) for x in self.context.messages] - messages_to_compress = [x for x in messages if x.role is not Role.SYSTEM][:-keep_recent_count] + messages_to_compress = [x for x in messages if x.role is not Role.SYSTEM] + if keep_recent_count > 0: + messages_to_compress = messages_to_compress[:-keep_recent_count] # If nothing to compress after filtering, return original messages if not messages_to_compress: diff --git a/reme_ai/context/offload/context_compress_op.py b/reme_ai/context/offload/context_compress_op.py index 218f6e21..7713ee18 100644 --- a/reme_ai/context/offload/context_compress_op.py +++ b/reme_ai/context/offload/context_compress_op.py @@ -265,8 +265,12 @@ class ContextCompressOp(BaseAsyncOp): system_message = system_message[0] messages_without_system = [x for x in messages if x.role is not Role.SYSTEM] - messages_to_compress = messages_without_system[:-keep_recent_count] - recent_messages = messages_without_system[-keep_recent_count:] + if keep_recent_count > 0: + messages_to_compress = messages_without_system[:-keep_recent_count] + recent_messages = messages_without_system[-keep_recent_count:] + else: + messages_to_compress = messages_without_system + recent_messages = [] # If nothing to compress after filtering, return original messages if not messages_to_compress: diff --git a/test/http_client_test.py b/test/http_client_test.py index 326f2cf9..acd1254c 100644 --- a/test/http_client_test.py +++ b/test/http_client_test.py @@ -129,6 +129,26 @@ async def run3(session): print(json.dumps(result, ensure_ascii=False)) +async def run4(session): + workspace_id = "default4" + + async with session.post( + f"{base_url}/agentic_retrieve", + json={ + "messages": [ + {"role": "user", "content": "hello" * 10000}, + ], + "workspace_id": workspace_id, + "context_manage_mode": "auto", + "keep_recent_count": 0, + "max_total_tokens": 10000, + }, + headers={"Content-Type": "application/json"}, + ) as response: + result = await response.json() + print(json.dumps(result, ensure_ascii=False)) + + async def main(): async with aiohttp.ClientSession() as session: @@ -137,7 +157,8 @@ async def main(): # await run1(session) # await run2(session) - await run3(session) + # await run3(session) + await run4(session) if __name__ == "__main__":