From 4b5d7ca645a12040b03d979b3ddbb6dc281c7a87 Mon Sep 17 00:00:00 2001 From: Alex Zhang Date: Mon, 16 Feb 2026 17:19:25 -0500 Subject: [PATCH 1/2] ensure tool calls and reserved names aren't overwritten --- rlm/environments/base_env.py | 7 ++++--- rlm/environments/daytona_repl.py | 6 ++++++ rlm/environments/docker_repl.py | 6 ++++++ rlm/environments/e2b_repl.py | 6 ++++++ rlm/environments/local_repl.py | 26 +++++++++++++++++++++++- rlm/environments/modal_repl.py | 6 ++++++ rlm/environments/prime_repl.py | 6 ++++++ tests/test_local_repl.py | 35 ++++++++++++++++++++++++++++++++ 8 files changed, 94 insertions(+), 4 deletions(-) diff --git a/rlm/environments/base_env.py b/rlm/environments/base_env.py index 103199601..9b9e4ba53 100644 --- a/rlm/environments/base_env.py +++ b/rlm/environments/base_env.py @@ -8,8 +8,8 @@ # Custom Tools Support # ============================================================================= -# Reserved names that cannot be overridden by custom tools. -# These are the built-in REPL functions that all environments must provide. +# Reserved names: cannot be overridden by custom tools, and are restored after each +# code execution to prevent namespace corruption (e.g. context = "...", llm_query = ...). RESERVED_TOOL_NAMES: frozenset[str] = frozenset( { "llm_query", @@ -17,6 +17,7 @@ "FINAL_VAR", "SHOW_VARS", "context", + "history", } ) @@ -173,7 +174,7 @@ class SupportsCustomTools(Protocol): The following names cannot be used as custom tool names: - llm_query, llm_query_batched: Single LM completion functions (no tool access) - FINAL_VAR, SHOW_VARS: Built-in helper functions - - context: The input context variable + - context, history: The input context and conversation history variables EXAMPLE: custom_tools = { diff --git a/rlm/environments/daytona_repl.py b/rlm/environments/daytona_repl.py index 0e489c3c5..afc4700bc 100644 --- a/rlm/environments/daytona_repl.py +++ b/rlm/environments/daytona_repl.py @@ -349,6 +349,12 @@ def SHOW_VARS(): sys.stdout = old_stdout sys.stderr = old_stderr +# Restore scaffold aliases if overwritten by executed code +if "context_0" in _locals: + _locals["context"] = _locals["context_0"] +if "history_0" in _locals: + _locals["history"] = _locals["history_0"] + save_state(_locals) result = {{ diff --git a/rlm/environments/docker_repl.py b/rlm/environments/docker_repl.py index 2de8c2b26..926b52268 100644 --- a/rlm/environments/docker_repl.py +++ b/rlm/environments/docker_repl.py @@ -173,6 +173,12 @@ def SHOW_VARS(): finally: sys.stdout, sys.stderr = old_stdout, old_stderr +# Restore scaffold aliases if overwritten by executed code +if "context_0" in _locals: + _locals["context"] = _locals["context_0"] +if "history_0" in _locals: + _locals["history"] = _locals["history_0"] + save_state(_locals) print(json.dumps({{"stdout": stdout_buf.getvalue(), "stderr": stderr_buf.getvalue(), "locals": {{k: repr(v) for k, v in _locals.items() if not k.startswith("_")}}}}, ensure_ascii=False)) ''' diff --git a/rlm/environments/e2b_repl.py b/rlm/environments/e2b_repl.py index aa33975db..6cec83b4b 100644 --- a/rlm/environments/e2b_repl.py +++ b/rlm/environments/e2b_repl.py @@ -241,6 +241,12 @@ def FINAL_VAR(variable_name): sys.stdout = old_stdout sys.stderr = old_stderr +# Restore scaffold aliases if overwritten by executed code +if "context_0" in _locals: + _locals["context"] = _locals["context_0"] +if "history_0" in _locals: + _locals["history"] = _locals["history_0"] + save_state(_locals) result = {{ diff --git a/rlm/environments/local_repl.py b/rlm/environments/local_repl.py index cc712ec74..65a49697c 100644 --- a/rlm/environments/local_repl.py +++ b/rlm/environments/local_repl.py @@ -13,7 +13,12 @@ from rlm.core.comms_utils import LMRequest, send_lm_request, send_lm_request_batched from rlm.core.types import REPLResult, RLMChatCompletion -from rlm.environments.base_env import NonIsolatedEnv, extract_tool_value, validate_custom_tools +from rlm.environments.base_env import ( + RESERVED_TOOL_NAMES, + NonIsolatedEnv, + extract_tool_value, + validate_custom_tools, +) # ============================================================================= # Safe Builtins @@ -375,6 +380,22 @@ def _temp_cwd(self): finally: os.chdir(old_cwd) + def _restore_scaffold(self) -> None: + """Restore scaffold names after execution so overwrites (e.g. context = 'x') don't persist.""" + for name in RESERVED_TOOL_NAMES: + if name == "llm_query": + self.globals["llm_query"] = self._llm_query + elif name == "llm_query_batched": + self.globals["llm_query_batched"] = self._llm_query_batched + elif name == "FINAL_VAR": + self.globals["FINAL_VAR"] = self._final_var + elif name == "SHOW_VARS": + self.globals["SHOW_VARS"] = self._show_vars + elif name == "context" and "context_0" in self.locals: + self.locals["context"] = self.locals["context_0"] + elif name == "history" and "history_0" in self.locals: + self.locals["history"] = self.locals["history_0"] + def execute_code(self, code: str) -> REPLResult: """Execute code in the persistent namespace and return result.""" start_time = time.perf_counter() @@ -392,6 +413,9 @@ def execute_code(self, code: str) -> REPLResult: if key not in self.globals and not key.startswith("_"): self.locals[key] = value + # Restore scaffold so model overwrites (context = ..., llm_query = ...) don't persist + self._restore_scaffold() + stdout = stdout_buf.getvalue() stderr = stderr_buf.getvalue() except Exception as e: diff --git a/rlm/environments/modal_repl.py b/rlm/environments/modal_repl.py index a00da7b2d..4287a44c4 100644 --- a/rlm/environments/modal_repl.py +++ b/rlm/environments/modal_repl.py @@ -261,6 +261,12 @@ def SHOW_VARS(): sys.stdout = old_stdout sys.stderr = old_stderr +# Restore scaffold aliases if overwritten by executed code +if "context_0" in _locals: + _locals["context"] = _locals["context_0"] +if "history_0" in _locals: + _locals["history"] = _locals["history_0"] + save_state(_locals) result = {{ diff --git a/rlm/environments/prime_repl.py b/rlm/environments/prime_repl.py index 22aad1518..eeb9a28a2 100644 --- a/rlm/environments/prime_repl.py +++ b/rlm/environments/prime_repl.py @@ -260,6 +260,12 @@ def SHOW_VARS(): sys.stdout = old_stdout sys.stderr = old_stderr +# Restore scaffold aliases if overwritten by executed code +if "context_0" in _locals: + _locals["context"] = _locals["context_0"] +if "history_0" in _locals: + _locals["history"] = _locals["history_0"] + save_state(_locals) result = {{ diff --git a/tests/test_local_repl.py b/tests/test_local_repl.py index 787a5b54e..b16e43cbb 100644 --- a/tests/test_local_repl.py +++ b/tests/test_local_repl.py @@ -175,6 +175,41 @@ def test_list_context(self): repl.cleanup() +class TestLocalREPLScaffoldRestoration: + """Tests that overwriting scaffold names (context, llm_query, etc.) is reverted after each execution.""" + + def test_context_restored_after_overwrite(self): + """If the model does context = 'something', the next execution still sees the real context.""" + repl = LocalREPL(context_payload="original context content") + assert repl.locals["context"] == "original context content" + + repl.execute_code('context = "hijacked"') + assert repl.locals["context"] == "original context content" + + out = repl.execute_code("print(context)") + assert "original context content" in out.stdout + repl.cleanup() + + def test_llm_query_restored_after_overwrite(self): + """If the model does llm_query = lambda x: 'hijacked', the next execution still has real llm_query.""" + repl = LocalREPL() + repl.execute_code("llm_query = lambda x: 'hijacked'") + + repl.execute_code("r = llm_query('test')") + assert "Error" in repl.locals["r"] + repl.cleanup() + + def test_final_var_restored_after_overwrite(self): + """If the model overwrites FINAL_VAR, the next execution still has the real FINAL_VAR.""" + repl = LocalREPL() + repl.execute_code("answer = 42") + repl.execute_code("FINAL_VAR = lambda x: 'overwritten'") + + repl.execute_code("result = FINAL_VAR('answer')") + assert repl.locals["result"] == "42" + repl.cleanup() + + class TestLocalREPLCleanup: """Tests for cleanup behavior.""" From b679e8f2085b859d59a27e87ae65e0be76a2e2a3 Mon Sep 17 00:00:00 2001 From: Alex Zhang Date: Mon, 16 Feb 2026 17:45:19 -0500 Subject: [PATCH 2/2] add logger examples + extra logging func --- README.md | 13 +++--- examples/logger_example.py | 46 +++++++++++++++++++++ rlm/core/rlm.py | 5 +++ rlm/core/types.py | 9 ++++- rlm/logger/rlm_logger.py | 82 +++++++++++++++++++++++++------------- tests/test_types.py | 34 ++++++++++++++++ 6 files changed, 156 insertions(+), 33 deletions(-) create mode 100644 examples/logger_example.py diff --git a/README.md b/README.md index 75c0f8b80..d4fdb5fa9 100644 --- a/README.md +++ b/README.md @@ -135,17 +135,20 @@ If you use this code or repository in your research, please cite: } ``` +## Optional: Trajectory metadata and logging +`RLMChatCompletion` has an optional `metadata` field (default empty) that can hold the full trajectory (run config + all iterations and sub-calls) so you can reconstruct the run. Pass an `RLMLogger` to capture it: + +- **In-memory only** (trajectory on `completion.metadata`): `logger=RLMLogger()` (no `log_dir`). +- **Also save to disk** (JSONL for the visualizer): `logger=RLMLogger(log_dir="./logs")`. + ## Optional Debugging: Visualizing RLM Trajectories -We additionally provide a simple visualizer tool to examine and view the code, sub-LM, and root-LM calls of an RLM trajectory. To save log files (`.jsonl`) on every completion call that can be viewed in the visualizer, initialize the `RLMLogger` object and pass it into the `RLM` on initialization: +We provide a simple visualizer to inspect code, sub-LM, and root-LM calls. Use `RLMLogger(log_dir="./logs")` so each completion writes a `.jsonl` file: ```python from rlm.logger import RLMLogger from rlm import RLM logger = RLMLogger(log_dir="./logs") -rlm = RLM( - ... - logger=logger -) +rlm = RLM(..., logger=logger) ``` To run the visualizer locally, we use Node.js and shadcn/ui: diff --git a/examples/logger_example.py b/examples/logger_example.py new file mode 100644 index 000000000..aec7b85f0 --- /dev/null +++ b/examples/logger_example.py @@ -0,0 +1,46 @@ +""" +Logger example: same structure as quickstart, with a different prompt. +Shows trajectory capture in completion.metadata (and optional save to disk). +""" + +import os + +from dotenv import load_dotenv + +from rlm import RLM +from rlm.logger import RLMLogger + +load_dotenv() + +# In-memory only: trajectory on completion.metadata (no disk write) +logger = RLMLogger() +# To also save JSONL for the visualizer: logger = RLMLogger(log_dir="./logs") + +rlm = RLM( + backend="portkey", + backend_kwargs={ + "model_name": "@openai/gpt-5-nano", + "api_key": os.getenv("PORTKEY_API_KEY"), + }, + environment="local", + environment_kwargs={}, + max_depth=1, + logger=logger, + verbose=True, +) + +prompt = "Compute the sum 1 + 2 + 3 + 4 + 5 and print the result using a REPL block, then return it with FINAL_VAR." +result = rlm.completion(prompt) + +print("Response:", result.response) +if result.metadata: + traj = result.metadata + print("Trajectory: run_metadata +", len(traj.get("iterations", [])), "iterations") + print("Trajectory:", traj) + if traj.get("iterations"): + first = traj["iterations"][0] + print(" First iteration keys:", list(first.keys())) +else: + print("Trajectory: (none — no logger or metadata not captured)") + +print("Full response:", result) diff --git a/rlm/core/rlm.py b/rlm/core/rlm.py index 0cdf4010f..6c0646ec3 100644 --- a/rlm/core/rlm.py +++ b/rlm/core/rlm.py @@ -229,6 +229,9 @@ def completion( if self.depth >= self.max_depth: return self._fallback_answer(prompt) + if self.logger: + self.logger.clear_iterations() + with self._spawn_completion_context(prompt) as (lm_handler, environment): message_history = self._setup_prompt(prompt) @@ -283,6 +286,7 @@ def completion( response=final_answer, usage_summary=usage, execution_time=time_end - time_start, + metadata=self.logger.get_trajectory() if self.logger else None, ) # Format the iteration for the next prompt. @@ -310,6 +314,7 @@ def completion( response=final_answer, usage_summary=usage, execution_time=time_end - time_start, + metadata=self.logger.get_trajectory() if self.logger else None, ) def _completion_turn( diff --git a/rlm/core/types.py b/rlm/core/types.py index 195bd8be3..631051ea3 100644 --- a/rlm/core/types.py +++ b/rlm/core/types.py @@ -96,15 +96,21 @@ class RLMChatCompletion: response: str usage_summary: UsageSummary execution_time: float + metadata: dict | None = ( + None # Full trajectory (run_metadata + iterations) when logger captures it + ) def to_dict(self): - return { + out = { "root_model": self.root_model, "prompt": self.prompt, "response": self.response, "usage_summary": self.usage_summary.to_dict(), "execution_time": self.execution_time, } + if self.metadata is not None: + out["metadata"] = self.metadata + return out @classmethod def from_dict(cls, data: dict) -> "RLMChatCompletion": @@ -114,6 +120,7 @@ def from_dict(cls, data: dict) -> "RLMChatCompletion": response=data.get("response"), usage_summary=UsageSummary.from_dict(data.get("usage_summary")), execution_time=data.get("execution_time"), + metadata=data.get("metadata"), ) diff --git a/rlm/logger/rlm_logger.py b/rlm/logger/rlm_logger.py index 25d57d307..418e679ce 100644 --- a/rlm/logger/rlm_logger.py +++ b/rlm/logger/rlm_logger.py @@ -1,7 +1,8 @@ """ Logger for RLM iterations. -Writes RLMIteration data to JSON-lines files for analysis and debugging. +Captures run metadata and iterations in memory so they can be attached to +RLMChatCompletion.metadata. Optionally writes the same data to JSON-lines files. """ import json @@ -13,50 +14,77 @@ class RLMLogger: - """Logger that writes RLMIteration data to a JSON-lines file.""" + """ + Captures trajectory (run metadata + iterations) for each completion. + By default only captures in memory; set log_dir to also save to disk. - def __init__(self, log_dir: str, file_name: str = "rlm"): - self.log_dir = log_dir - os.makedirs(log_dir, exist_ok=True) - - timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") - run_id = str(uuid.uuid4())[:8] - self.log_file_path = os.path.join(log_dir, f"{file_name}_{timestamp}_{run_id}.jsonl") + - log_dir=None: trajectory is available via get_trajectory() and can be + attached to RLMChatCompletion.metadata (no disk write). + - log_dir="path": same capture plus appends to a JSONL file per run. + """ + def __init__(self, log_dir: str | None = None, file_name: str = "rlm"): + self._save_to_disk = log_dir is not None + self.log_dir = log_dir + self.log_file_path: str | None = None + if self._save_to_disk and log_dir: + os.makedirs(log_dir, exist_ok=True) + timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + run_id = str(uuid.uuid4())[:8] + self.log_file_path = os.path.join(log_dir, f"{file_name}_{timestamp}_{run_id}.jsonl") + + self._run_metadata: dict | None = None + self._iterations: list[dict] = [] self._iteration_count = 0 self._metadata_logged = False - def log_metadata(self, metadata: RLMMetadata): - """Log RLM metadata as the first entry in the file.""" + def log_metadata(self, metadata: RLMMetadata) -> None: + """Capture run metadata (and optionally write to file).""" if self._metadata_logged: return - entry = { - "type": "metadata", - "timestamp": datetime.now().isoformat(), - **metadata.to_dict(), - } - - with open(self.log_file_path, "a") as f: - json.dump(entry, f) - f.write("\n") - + self._run_metadata = metadata.to_dict() self._metadata_logged = True - def log(self, iteration: RLMIteration): - """Log an RLMIteration to the file.""" + if self._save_to_disk and self.log_file_path: + entry = { + "type": "metadata", + "timestamp": datetime.now().isoformat(), + **self._run_metadata, + } + with open(self.log_file_path, "a") as f: + json.dump(entry, f) + f.write("\n") + + def log(self, iteration: RLMIteration) -> None: + """Capture one iteration (and optionally append to file).""" self._iteration_count += 1 - entry = { "type": "iteration", "iteration": self._iteration_count, "timestamp": datetime.now().isoformat(), **iteration.to_dict(), } + self._iterations.append(entry) + + if self._save_to_disk and self.log_file_path: + with open(self.log_file_path, "a") as f: + json.dump(entry, f) + f.write("\n") - with open(self.log_file_path, "a") as f: - json.dump(entry, f) - f.write("\n") + def clear_iterations(self) -> None: + """Reset iterations for the next completion (trajectory is per completion).""" + self._iterations = [] + self._iteration_count = 0 + + def get_trajectory(self) -> dict | None: + """Return captured run_metadata + iterations for the current completion, or None if no metadata yet.""" + if self._run_metadata is None: + return None + return { + "run_metadata": self._run_metadata, + "iterations": list(self._iterations), + } @property def iteration_count(self) -> int: diff --git a/tests/test_types.py b/tests/test_types.py index 21780421e..847043f34 100644 --- a/tests/test_types.py +++ b/tests/test_types.py @@ -5,6 +5,7 @@ ModelUsageSummary, QueryMetadata, REPLResult, + RLMChatCompletion, RLMIteration, RLMMetadata, UsageSummary, @@ -156,6 +157,39 @@ def test_to_dict(self): assert d["iteration_time"] == 1.5 +class TestRLMChatCompletion: + """Tests for RLMChatCompletion.""" + + def test_metadata_default_none(self): + usage = UsageSummary(model_usage_summaries={}) + c = RLMChatCompletion( + root_model="gpt-4", + prompt="hi", + response="hello", + usage_summary=usage, + execution_time=1.0, + ) + assert c.metadata is None + d = c.to_dict() + assert "metadata" not in d + + def test_metadata_roundtrip(self): + usage = UsageSummary(model_usage_summaries={}) + trajectory = {"run_metadata": {"root_model": "gpt-4"}, "iterations": []} + c = RLMChatCompletion( + root_model="gpt-4", + prompt="hi", + response="hello", + usage_summary=usage, + execution_time=1.0, + metadata=trajectory, + ) + d = c.to_dict() + assert d["metadata"] == trajectory + c2 = RLMChatCompletion.from_dict(d) + assert c2.metadata == trajectory + + class TestQueryMetadata: """Tests for QueryMetadata."""