From 53fb42104e0ca4baaa3804bb64cc114e5f40a705 Mon Sep 17 00:00:00 2001 From: Jon Saad-Falcon <41205309+jonsaadfalcon@users.noreply.github.com> Date: Tue, 2 Jun 2026 19:31:08 -0700 Subject: [PATCH] feat(rlm): expose real tool calls inside the REPL (#481) * rlm: expose real tool calls inside the repl * rlm: wrap long TypeError message in repl * style(rlm): collapse over-wrapped TypeError to satisfy ruff format The cherry-picked RLM tool-call work left a `raise TypeError(\n message\n)` that `ruff format --check` rejects (the PR's own lint-fix commit broke format). Collapse to `raise TypeError(message)`. No behavior change. Co-Authored-By: Claude Opus 4.8 (1M context) --------- Co-authored-by: Eddie Richter Co-authored-by: Claude Opus 4.8 (1M context) --- src/openjarvis/agents/rlm.py | 92 +++++++++++++++++++++- src/openjarvis/agents/rlm_repl.py | 64 ++++++++++++++++ tests/agents/test_rlm.py | 123 ++++++++++++++++++++++++++++++ tests/agents/test_rlm_repl.py | 45 +++++++++++ 4 files changed, 323 insertions(+), 1 deletion(-) diff --git a/src/openjarvis/agents/rlm.py b/src/openjarvis/agents/rlm.py index 2b4fc2ad..c7bc8fa6 100644 --- a/src/openjarvis/agents/rlm.py +++ b/src/openjarvis/agents/rlm.py @@ -8,8 +8,9 @@ context and makes recursive sub-LM calls via ``llm_query()``/``llm_batch()``. from __future__ import annotations +import json import re -from typing import Any, List, Optional +from typing import Any, Dict, List, Optional from openjarvis.agents._stubs import AgentContext, AgentResult, ToolUsingAgent from openjarvis.agents.prompt_loader import ( @@ -35,6 +36,12 @@ RLM_SYSTEM_PROMPT = ( "prompt and get a response.\n" "- `llm_batch(prompts: list[str]) -> list[str]` — Call a " "sub-LM with multiple prompts.\n" + "- `tool_call(tool_name: str, args: dict) -> str` — Execute an " + "OpenJarvis tool and return its textual output.\n" + "- `read_file(path: str, max_lines: int = 120) -> str` — Read the " + "first N lines of a file through the real file_read tool.\n" + "- `read_file_chunk(path: str, start_line: int, end_line: int) -> str` " + "— Read only a bounded line range from a file.\n" "- `FINAL(value)` — Terminate and return `value` as the " "final answer.\n" "- `FINAL_VAR(var_name: str)` — Terminate and return the " @@ -59,6 +66,14 @@ RLM_SYSTEM_PROMPT = ( '`FINAL(answer_value)` or `FINAL_VAR("var_name")`.\n' "5. If you can answer directly without code, just respond " "with text (no code block).\n\n" + "6. For file access, web access, or any external side effect, use " + "the injected tool helpers (for example " + '`file_read(path="...")` or `tool_call("file_read", {{"path": "..."}})`). ' + "Do NOT use open(), subprocess, urllib, or direct OS access.\n\n" + "7. For file-analysis tasks, do NOT dump whole files. Prefer " + "`read_file(path, max_lines=...)` or " + "`read_file_chunk(path, start_line, end_line)` and summarize " + "incrementally.\n\n" "## Strategy Tips\n\n" "- Split long text into paragraphs or sections, summarize " "each with `llm_query()`.\n" @@ -167,9 +182,12 @@ class RLMAgent(ToolUsingAgent): system_prompt = prompt_template # Create REPL with sub-LM callbacks + self._repl_tool_results: List[ToolResult] = [] repl = RLMRepl( llm_query_fn=self._make_sub_query, llm_batch_fn=self._make_batch_query, + tool_call_fn=self._execute_tool_from_repl if self._executor else None, + tool_arg_names=self._tool_arg_names(), max_output_chars=self._max_output_chars, ) @@ -177,6 +195,9 @@ class RLMAgent(ToolUsingAgent): ctx_text = self._resolve_context(context) if ctx_text: repl.set_variable("context", ctx_text) + if self._executor is not None: + repl.set_variable("read_file", self._repl_read_file) + repl.set_variable("read_file_chunk", self._repl_read_file_chunk) # Build conversation messages = self._build_messages( @@ -227,6 +248,10 @@ class RLMAgent(ToolUsingAgent): # Execute code in REPL output = repl.execute(code) + if self._repl_tool_results: + all_tool_results.extend(self._repl_tool_results) + self._repl_tool_results = [] + # Record as tool result tool_result = ToolResult( tool_name="rlm_repl", @@ -333,6 +358,71 @@ class RLMAgent(ToolUsingAgent): """ return [self._make_sub_query(p) for p in prompts] + def _tool_arg_names(self) -> Dict[str, Optional[str]]: + """Build a best-effort primary-arg map for injected tool helpers.""" + arg_names: Dict[str, Optional[str]] = {} + for tool in self._tools: + props = tool.spec.parameters.get("properties", {}) + required = tool.spec.parameters.get("required", []) + primary = None + if required: + primary = required[0] + elif props: + primary = next(iter(props.keys())) + arg_names[tool.spec.name] = primary + return arg_names + + def _execute_tool_from_repl(self, tool_name: str, params: Dict[str, Any]) -> str: + """Execute a real OpenJarvis tool from within the REPL.""" + if self._executor is None: + raise RuntimeError(f"Tool '{tool_name}' is not available") + + tc = ToolCall( + id=f"rlm_tool_{len(getattr(self, '_repl_tool_results', []))}", + name=tool_name, + arguments=json.dumps(params), + ) + tr = self._executor.execute(tc) + getattr(self, "_repl_tool_results", []).append(tr) + return tr.content + + def _repl_read_file(self, path: str, max_lines: int = 120) -> str: + """Read a bounded number of lines from a file via the real tool.""" + params: Dict[str, Any] = {"path": path} + try: + max_lines_int = int(max_lines) + except (TypeError, ValueError): + max_lines_int = 120 + if max_lines_int > 0: + params["max_lines"] = max_lines_int + return self._execute_tool_from_repl("file_read", params) + + def _repl_read_file_chunk( + self, + path: str, + start_line: int, + end_line: int, + ) -> str: + """Read a bounded line range from a file via the file_read tool. + + The underlying tool only supports a head-style max_lines cap, so this + helper reads up to ``end_line`` and then slices the requested range in + Python. This is still much cheaper than asking the model to dump the + entire file into the next turn. + """ + try: + start = max(1, int(start_line)) + end = max(start, int(end_line)) + except (TypeError, ValueError): + start, end = 1, 120 + + content = self._execute_tool_from_repl( + "file_read", + {"path": path, "max_lines": end}, + ) + lines = content.splitlines(keepends=True) + return "".join(lines[start - 1 : end]) + # ------------------------------------------------------------------ # Helpers # ------------------------------------------------------------------ diff --git a/src/openjarvis/agents/rlm_repl.py b/src/openjarvis/agents/rlm_repl.py index dd9f58e3..161283c1 100644 --- a/src/openjarvis/agents/rlm_repl.py +++ b/src/openjarvis/agents/rlm_repl.py @@ -58,12 +58,15 @@ class RLMRepl: self, llm_query_fn: Optional[Callable[[str], str]] = None, llm_batch_fn: Optional[Callable[[List[str]], List[str]]] = None, + tool_call_fn: Optional[Callable[[str, Dict[str, Any]], str]] = None, + tool_arg_names: Optional[Dict[str, Optional[str]]] = None, *, max_output_chars: int = 10000, ) -> None: self._max_output_chars = max_output_chars self._terminated = False self._final_value: Any = None + self._tool_call_fn = tool_call_fn # Build namespace self._namespace: Dict[str, Any] = {} @@ -89,6 +92,17 @@ class RLMRepl: self._namespace["llm_query"] = llm_query_fn if llm_batch_fn is not None: self._namespace["llm_batch"] = llm_batch_fn + if tool_call_fn is not None: + self._namespace["tool_call"] = self._tool_call + + # Expose each tool as a direct Python helper (e.g. + # file_read(path="...")) so the model does not need to invent + # pseudo-tool syntax or fall back to blocked file I/O. + for tool_name, primary_arg in (tool_arg_names or {}).items(): + self._namespace[tool_name] = self._make_tool_wrapper( + tool_name, + primary_arg, + ) # ------------------------------------------------------------------ # Termination helpers @@ -129,6 +143,56 @@ class RLMRepl: # Execution # ------------------------------------------------------------------ + def _tool_call(self, tool_name: str, *args: Any, **kwargs: Any) -> str: + """Execute an injected OpenJarvis tool from within the REPL. + + Supported forms: + - ``tool_call("file_read", {"path": "foo.txt"})`` + - ``tool_call("file_read", path="foo.txt")`` + """ + if self._tool_call_fn is None: + raise RuntimeError("tool_call is not available in this REPL") + + if args: + if len(args) != 1 or kwargs: + raise TypeError( + "tool_call expects either a single dict argument or keyword args" + ) + if not isinstance(args[0], dict): + raise TypeError("tool_call positional argument must be a dict") + params = dict(args[0]) + else: + params = dict(kwargs) + + return self._tool_call_fn(tool_name, params) + + def _make_tool_wrapper( + self, + tool_name: str, + primary_arg: Optional[str], + ) -> Callable[..., str]: + """Return a Python helper that dispatches to ``tool_call``.""" + + def _wrapper(*args: Any, **kwargs: Any) -> str: + if kwargs: + params = dict(kwargs) + elif len(args) == 1 and primary_arg is not None: + params = {primary_arg: args[0]} + elif len(args) == 1 and isinstance(args[0], dict): + params = dict(args[0]) + elif not args: + params = {} + else: + message = ( + f"{tool_name} expects keyword args or a single " + f"{primary_arg!r} argument" + ) + raise TypeError(message) + return self._tool_call(tool_name, params) + + _wrapper.__name__ = tool_name + return _wrapper + def security_check(self, code: str) -> Optional[str]: """Check code for dangerous patterns. Returns error message or None.""" for pattern in _BLOCKED_PATTERNS: diff --git a/tests/agents/test_rlm.py b/tests/agents/test_rlm.py index 132e6246..059aa062 100644 --- a/tests/agents/test_rlm.py +++ b/tests/agents/test_rlm.py @@ -40,6 +40,47 @@ class _CalcStub(BaseTool): return ToolResult(tool_name="calculator", content=str(val), success=True) +class _FileReadStub(BaseTool): + tool_id = "file_read" + + @property + def spec(self) -> ToolSpec: + return ToolSpec( + name="file_read", + description="Read a file.", + parameters={ + "type": "object", + "properties": {"path": {"type": "string"}}, + "required": ["path"], + }, + ) + + def execute(self, **params) -> ToolResult: + path = params.get("path", "") + max_lines = params.get("max_lines") + if path == "rust/Cargo.toml": + content = '[workspace]\nmembers = ["a", "b", "c"]\n' + if max_lines is not None: + lines = content.splitlines(keepends=True) + content = "".join(lines[: int(max_lines)]) + return ToolResult( + tool_name="file_read", + content=content, + success=True, + ) + if path == "long.txt": + content = "line1\nline2\nline3\nline4\nline5\n" + if max_lines is not None: + lines = content.splitlines(keepends=True) + content = "".join(lines[: int(max_lines)]) + return ToolResult( + tool_name="file_read", + content=content, + success=True, + ) + return ToolResult(tool_name="file_read", content="", success=False) + + def _make_engine(content: str = "Final answer.") -> MagicMock: """Engine that returns plain content (no code block).""" engine = MagicMock() @@ -373,6 +414,88 @@ class TestRLMSubLMWithTools: assert result.content == "The answer is 4." +class TestRLMDirectToolBridge: + def test_root_repl_can_use_file_read_tool_directly(self): + engine = MagicMock() + engine.engine_id = "mock" + engine.generate.side_effect = [ + { + "content": ( + "```python\n" + 'content = file_read("rust/Cargo.toml")\n' + "print(content)\n" + "FINAL('read ok')\n" + "```" + ), + "usage": { + "prompt_tokens": 5, + "completion_tokens": 20, + "total_tokens": 25, + }, + "model": "test-model", + "finish_reason": "stop", + } + ] + agent = RLMAgent(engine, "test-model", tools=[_FileReadStub()]) + result = agent.run("Read Cargo") + assert result.content == "read ok" + assert any(tr.tool_name == "file_read" for tr in result.tool_results) + assert any( + tr.tool_name == "rlm_repl" and "members" in tr.content + for tr in result.tool_results + ) + + def test_root_repl_can_use_bounded_read_helper(self): + engine = MagicMock() + engine.engine_id = "mock" + engine.generate.side_effect = [ + { + "content": ( + "```python\n" + 'snippet = read_file("long.txt", max_lines=2)\n' + "FINAL(snippet)\n" + "```" + ), + "usage": { + "prompt_tokens": 5, + "completion_tokens": 20, + "total_tokens": 25, + }, + "model": "test-model", + "finish_reason": "stop", + } + ] + agent = RLMAgent(engine, "test-model", tools=[_FileReadStub()]) + result = agent.run("Read file head") + assert result.content == "line1\nline2\n" + assert any(tr.tool_name == "file_read" for tr in result.tool_results) + + def test_root_repl_can_use_file_chunk_helper(self): + engine = MagicMock() + engine.engine_id = "mock" + engine.generate.side_effect = [ + { + "content": ( + "```python\n" + 'snippet = read_file_chunk("long.txt", 2, 4)\n' + "FINAL(snippet)\n" + "```" + ), + "usage": { + "prompt_tokens": 5, + "completion_tokens": 20, + "total_tokens": 25, + }, + "model": "test-model", + "finish_reason": "stop", + } + ] + agent = RLMAgent(engine, "test-model", tools=[_FileReadStub()]) + result = agent.run("Read file chunk") + assert result.content == "line2\nline3\nline4\n" + assert any(tr.tool_name == "file_read" for tr in result.tool_results) + + class TestRLMBlockedCode: def test_blocked_code_returns_error(self): engine = MagicMock() diff --git a/tests/agents/test_rlm_repl.py b/tests/agents/test_rlm_repl.py index dd96f899..c438b992 100644 --- a/tests/agents/test_rlm_repl.py +++ b/tests/agents/test_rlm_repl.py @@ -169,6 +169,51 @@ class TestRLMReplCallbacks: output = repl.execute("llm_query('test')") assert "NameError" in output + def test_tool_call_callback(self): + calls = [] + + def mock_tool(tool_name, params): + calls.append((tool_name, params)) + return "tool-output" + + repl = RLMRepl(tool_call_fn=mock_tool) + out = repl.execute("result = tool_call('calculator', {'expression': '2+2'})") + assert out == "" + assert calls == [("calculator", {"expression": "2+2"})] + assert repl.get_variable("result") == "tool-output" + + def test_direct_tool_wrapper_with_kwargs(self): + calls = [] + + def mock_tool(tool_name, params): + calls.append((tool_name, params)) + return "file contents" + + repl = RLMRepl( + tool_call_fn=mock_tool, + tool_arg_names={"file_read": "path"}, + ) + out = repl.execute('result = file_read(path="README.md")') + assert out == "" + assert calls == [("file_read", {"path": "README.md"})] + assert repl.get_variable("result") == "file contents" + + def test_direct_tool_wrapper_with_positional_arg(self): + calls = [] + + def mock_tool(tool_name, params): + calls.append((tool_name, params)) + return "cargo toml" + + repl = RLMRepl( + tool_call_fn=mock_tool, + tool_arg_names={"file_read": "path"}, + ) + out = repl.execute('result = file_read("rust/Cargo.toml")') + assert out == "" + assert calls == [("file_read", {"path": "rust/Cargo.toml"})] + assert repl.get_variable("result") == "cargo toml" + class TestRLMReplOutput: """Output truncation and error handling."""