diff --git a/reme_ai/retrieve/working/grep_op.py b/reme_ai/retrieve/working/grep_op.py index 3d20df18..e995723a 100644 --- a/reme_ai/retrieve/working/grep_op.py +++ b/reme_ai/retrieve/working/grep_op.py @@ -5,7 +5,6 @@ It enables efficient content-based search using regular expressions, with suppor for glob pattern filtering and result limiting. """ -import asyncio from pathlib import Path from typing import List @@ -13,6 +12,7 @@ from flowllm.core.context import C from flowllm.core.op import BaseAsyncToolOp from flowllm.core.schema import ToolCall from loguru import logger +from reme_ai.utils.op_utils import run_shell_command @C.register_op() @@ -81,20 +81,12 @@ class GrepOp(BaseAsyncToolOp): logger.info(f"Running grep command: {' '.join(cmd)}") - # Execute grep using an async subprocess - process = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - stdout, stderr = await process.communicate() + # Execute grep using run_shell_command + stdout, stderr, returncode = await run_shell_command(cmd, timeout=None) - if process.returncode not in (0, 1): - # grep returns 1 when no matches are found; treat other codes as errors - err_msg = stderr.decode("utf-8", errors="ignore").strip() - raise RuntimeError(f"grep failed with code {process.returncode}: {err_msg}") + assert returncode in (0, 1), f"grep failed with code {returncode}: {stderr.strip()}" - output_text = stdout.decode("utf-8", errors="ignore").strip() + output_text = stdout.strip() # Return raw grep output if not output_text: diff --git a/reme_ai/retrieve/working/read_file_op.py b/reme_ai/retrieve/working/read_file_op.py index c6d4aaaa..9516798d 100644 --- a/reme_ai/retrieve/working/read_file_op.py +++ b/reme_ai/retrieve/working/read_file_op.py @@ -4,13 +4,13 @@ This module provides a tool operation for reading file contents. It supports reading entire files or specific line ranges for large files. """ -import asyncio from pathlib import Path from typing import Optional from flowllm.core.context import C from flowllm.core.op import BaseAsyncToolOp from flowllm.core.schema import ToolCall +from reme_ai.utils.op_utils import run_shell_command @C.register_op() @@ -76,15 +76,10 @@ class ReadFileOp(BaseAsyncToolOp): end_line = offset + limit cmd = ["sed", "-n", f"{start_line},{end_line}p", str(file_path_obj)] - process = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=30) + stdout, stderr, returncode = await run_shell_command(cmd, timeout=30) - assert process.returncode == 0, f"sed command failed: {stderr.decode()}" - content = stdout.decode().rstrip("\n") + assert returncode == 0, f"sed command failed: {stderr}" + content = stdout.rstrip("\n") self.set_output(content) async def async_default_execute(self, e: Exception = None, **_kwargs): diff --git a/reme_ai/utils/op_utils.py b/reme_ai/utils/op_utils.py index a4d86645..c6e439b7 100644 --- a/reme_ai/utils/op_utils.py +++ b/reme_ai/utils/op_utils.py @@ -4,15 +4,50 @@ This module provides helper functions for merging messages, parsing JSON respons extracting trajectory context, and parsing insight updates. """ +import asyncio import json import re -from typing import List +from typing import List, Optional, Tuple from flowllm.core.schema import Message, Trajectory from flowllm.core.utils import merge_messages_content as merge_messages_content_flowllm from loguru import logger +async def run_shell_command( + cmd: List[str], + timeout: Optional[float] = 30, +) -> Tuple[str, str, int]: + """Run a shell command asynchronously. + + Args: + cmd: Command and arguments as a list. + timeout: Timeout in seconds. None for no timeout. + + Returns: + Tuple of (stdout, stderr, returncode). + + Raises: + asyncio.TimeoutError: If command execution exceeds timeout. + """ + process = await asyncio.create_subprocess_exec( + *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + + if timeout: + stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=timeout) + else: + stdout, stderr = await process.communicate() + + return ( + stdout.decode("utf-8", errors="ignore"), + stderr.decode("utf-8", errors="ignore"), + process.returncode, + ) + + def merge_messages_content(messages: List[Message | dict]) -> str: """Merge content from a list of messages into a single string. diff --git a/test_op/test_agentic_retrieve_op.py b/test_op/test_agentic_retrieve_op.py index 7b0f0c1c..7233d2dd 100644 --- a/test_op/test_agentic_retrieve_op.py +++ b/test_op/test_agentic_retrieve_op.py @@ -70,8 +70,9 @@ async def test_agentic_retrieve_basic(): ] # llm = "qwen3_coder_plus" - llm = "qwen3_30b_instruct" - # llm = "qwen3_coder_30b_instruct" + # llm = "qwen3_30b_instruct" + # llm = "qwen3_30b_thinking" + llm = "qwen3_coder_30b_instruct" # llm = "qwen3_max_instruct" op = AgenticRetrieveOp(llm=llm)