Skip to content

fix(rlm): preserve input mutations across iterations #67

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
isaacbmiller wants to merge 1 commit into
base: review-base/stanford-main-10014
Choose a base branch
from
Draft

fix(rlm): preserve input mutations across iterations #67

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 12 additions & 19 deletions dspy/predict/rlm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -376,15 +376,10 @@ def _validate_inputs(self, input_args: dict[str, Any]) -> None:
if missing:
raise ValueError(f"Missing required inputs: {sorted(missing)}")

def _prepare_serializable_vars(
def _initialize_inputs(
self, input_args: dict[str, Any], repl: CodeInterpreter,
) -> dict[str, Any]:
"""Inject SandboxSerializable values into the interpreter.

For each SandboxSerializable value in input_args, serializes it and
executes setup + assignment code in the interpreter. Returns the
remaining non-serializable args (for per-iteration use).
"""
) -> None:
"""Initialize the interpreter namespace once for this RLM run."""
repl.start()
regular_args = {}
for name, value in input_args.items():
Expand Down Expand Up @@ -416,7 +411,8 @@ def _prepare_serializable_vars(
code_lines.append(assignment)
repl.execute("\n".join(code_lines), variables=payload_vars)

return regular_args
if regular_args:
repl.execute("pass", variables=regular_args)

# =========================================================================
# CodeInterpreter Lifecycle
Expand Down Expand Up @@ -578,11 +574,10 @@ def _execute_code(
self,
repl: CodeInterpreter,
code: str,
input_args: dict[str, Any],
) -> Any:
"""Execute code in the interpreter, returning the result or an error string."""
try:
return repl.execute(code, variables=dict(input_args))
return repl.execute(code)
except (CodeInterpreterError, SyntaxError) as e:
return f"[Error] {e}"

Expand All @@ -592,7 +587,6 @@ def _execute_iteration(
variables: list[REPLVariable],
history: REPLHistory,
iteration: int,
input_args: dict[str, Any],
output_field_names: list[str],
) -> Prediction | REPLHistory:
"""Execute one iteration. Returns Prediction if done, else updated REPLHistory."""
Expand All @@ -614,7 +608,7 @@ def _execute_iteration(
code = action.code
result = f"[Error] {e}"
return self._process_execution_result(action, code, result, history, output_field_names)
result = self._execute_code(repl, code, input_args)
result = self._execute_code(repl, code)
return self._process_execution_result(action, code, result, history, output_field_names)

# =========================================================================
Expand All @@ -640,12 +634,12 @@ def forward(self, **input_args) -> Prediction:
variables = self._build_variables(**input_args)

with self._interpreter_context(execution_tools) as repl:
regular_args = self._prepare_serializable_vars(input_args, repl)
self._initialize_inputs(input_args, repl)
history: REPLHistory = REPLHistory(max_output_chars=self.max_output_chars)

for iteration in range(self.max_iters):
result: Prediction | REPLHistory = self._execute_iteration(
repl, variables, history, iteration, regular_args, output_field_names
repl, variables, history, iteration, output_field_names
)
if isinstance(result, Prediction):
return result
Expand Down Expand Up @@ -681,7 +675,6 @@ async def _aexecute_iteration(
variables: list[REPLVariable],
history: REPLHistory,
iteration: int,
input_args: dict[str, Any],
output_field_names: list[str],
) -> Prediction | REPLHistory:
"""Async version: Execute one iteration."""
Expand All @@ -703,7 +696,7 @@ async def _aexecute_iteration(
code = pred.code
result = f"[Error] {e}"
return self._process_execution_result(pred, code, result, history, output_field_names)
result = self._execute_code(repl, code, input_args)
result = self._execute_code(repl, code)
return self._process_execution_result(pred, code, result, history, output_field_names)

async def aforward(self, **input_args) -> Prediction:
Expand All @@ -725,12 +718,12 @@ async def aforward(self, **input_args) -> Prediction:
variables = self._build_variables(**input_args)

with self._interpreter_context(execution_tools) as repl:
regular_args = self._prepare_serializable_vars(input_args, repl)
self._initialize_inputs(input_args, repl)
history = REPLHistory(max_output_chars=self.max_output_chars)

for iteration in range(self.max_iters):
result = await self._aexecute_iteration(
repl, variables, history, iteration, regular_args, output_field_names
repl, variables, history, iteration, output_field_names
)
if isinstance(result, Prediction):
return result
Expand Down
78 changes: 45 additions & 33 deletions tests/predict/test_rlm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -511,7 +511,7 @@ class TestRLMCallMethod:

def test_call_is_alias_for_forward(self):
"""Test that __call__ is an alias for forward()."""
mock = MockInterpreter(responses=[FinalOutput({"answer": "42"})])
mock = MockInterpreter(responses=["", FinalOutput({"answer": "42"})])
rlm = RLM("query -> answer", max_iters=3, interpreter=mock)
rlm.generate_action = make_mock_predictor([
{"reasoning": "Return answer", "code": 'SUBMIT("42")'},
Expand All @@ -527,6 +527,7 @@ class TestRLMMaxIterationsFallback:
def test_max_iters_triggers_extract(self):
"""Test that reaching max_iters uses extract fallback."""
mock = MockInterpreter(responses=[
"",
"exploring...",
"still exploring...",
"more exploring...",
Expand Down Expand Up @@ -556,6 +557,7 @@ def failing_tool() -> str:
raise RuntimeError("Tool failed!")

mock = MockInterpreter(responses=[
"",
CodeInterpreterError("RuntimeError: Tool failed!"),
FinalOutput({"answer": "recovered"}),
])
Expand All @@ -571,6 +573,7 @@ def failing_tool() -> str:
def test_runtime_error_history_uses_stripped_code(self):
"""Runtime execution failures should preserve stripped code in history."""
mock = MockInterpreter(responses=[
"",
CodeInterpreterError("NameError: name 'x' is not defined"),
FinalOutput({"answer": "recovered"}),
])
Expand All @@ -588,6 +591,7 @@ def test_runtime_error_history_uses_stripped_code(self):
def test_syntax_error_from_execute_is_recoverable(self):
"""SyntaxError from interpreter.execute should be surfaced as an iteration error."""
mock = MockInterpreter(responses=[
"",
SyntaxError("invalid syntax"),
FinalOutput({"answer": "recovered"}),
])
Expand All @@ -604,6 +608,7 @@ def test_syntax_error_from_execute_is_recoverable(self):
def test_syntax_error_from_strip_code_fences_is_recoverable(self):
"""SyntaxError raised by _strip_code_fences (e.g. non-Python fence tag) should be recoverable."""
mock = MockInterpreter(responses=[
"",
FinalOutput({"answer": "recovered"}),
])
rlm = RLM("query -> answer", max_iters=5, interpreter=mock)
Expand Down Expand Up @@ -855,7 +860,7 @@ class TestRLMAsyncMock:
@pytest.mark.asyncio
async def test_aforward_basic(self):
"""Test aforward() returns Prediction with expected output (MockInterpreter)."""
mock = MockInterpreter(responses=[FinalOutput({"answer": "42"})])
mock = MockInterpreter(responses=["", FinalOutput({"answer": "42"})])
rlm = RLM("query -> answer", max_iters=3, interpreter=mock)
rlm.generate_action = make_mock_predictor([
{"reasoning": "Return answer", "code": 'SUBMIT("42")'},
Expand All @@ -867,7 +872,7 @@ async def test_aforward_basic(self):
@pytest.mark.asyncio
async def test_aforward_int_output_mock(self):
"""Test aforward() returns int when signature expects int (MockInterpreter)."""
mock = MockInterpreter(responses=[FinalOutput({"count": 42})])
mock = MockInterpreter(responses=["", FinalOutput({"count": 42})])
rlm = RLM("query -> count: int", max_iters=3, interpreter=mock)
rlm.generate_action = make_mock_predictor([
{"reasoning": "Return count", "code": "SUBMIT(42)"},
Expand All @@ -881,6 +886,7 @@ async def test_aforward_int_output_mock(self):
async def test_aforward_multi_iteration_mock(self):
"""Test aforward() handles multiple iterations before SUBMIT (MockInterpreter)."""
mock = MockInterpreter(responses=[
"",
"explored data",
FinalOutput({"answer": "done"}),
])
Expand All @@ -906,7 +912,7 @@ class TestRLMTypeCoercionMock:
])
def test_type_coercion(self, output_field, output_type, final_value, code, expected):
"""Test RLM type coercion for various types (MockInterpreter)."""
mock = MockInterpreter(responses=[FinalOutput({output_field: final_value})])
mock = MockInterpreter(responses=["", FinalOutput({output_field: final_value})])
rlm = RLM(f"query -> {output_field}: {output_type}", max_iters=3, interpreter=mock)
rlm.generate_action = make_mock_predictor([
{"reasoning": "Return value", "code": code},
Expand All @@ -918,6 +924,7 @@ def test_type_coercion(self, output_field, output_type, final_value, code, expec
def test_type_error_retries(self):
"""Test RLM retries when type validation fails (MockInterpreter)."""
mock = MockInterpreter(responses=[
"",
FinalOutput({"answer": "maybe"}), # Invalid for Literal
FinalOutput({"answer": "yes"}), # Valid
])
Expand Down Expand Up @@ -1108,6 +1115,17 @@ def test_with_input_variables_e2e(self):

assert result.total == 15

def test_input_mutations_persist_between_iterations(self):
"""Input variables share the same persistent namespace as generated variables."""
with dummy_lm_context([
{"reasoning": "Extend the input", "code": "numbers.append(4)\nprint(numbers)"},
{"reasoning": "Return the updated sum", "code": "SUBMIT(sum(numbers))"},
]):
rlm = RLM("numbers: list[int] -> total: int", max_iters=3)
result = rlm.forward(numbers=[1, 2, 3])

assert result.total == 10

def test_with_tool_e2e(self):
"""Test RLM calling a host-side tool through the sandbox."""
def lookup(key: str) -> str:
Expand Down Expand Up @@ -1259,41 +1277,32 @@ def test_regular_values_unchanged(self):
assert "plain text" in variables[0].preview


class TestPrepareSerializableVars:
"""Tests for _prepare_serializable_vars with MockInterpreter."""
class TestInitializeInputs:
"""Tests for one-time interpreter input initialization."""

def test_separates_serializable_from_regular(self):
"""Serializable values are injected; regular values are returned."""
mock = MockInterpreter(responses=["", FinalOutput({"answer": "42"})])
def test_initializes_serializable_and_regular_inputs_once(self):
mock = MockInterpreter(responses=["", ""])
rlm = RLM("data, query -> answer", max_iters=3, interpreter=mock)

stub = _StubSerializable("payload")

# Manually call _prepare_serializable_vars
rlm._inject_execution_context(mock, rlm._prepare_execution_tools())
regular = rlm._prepare_serializable_vars({"data": stub, "query": "hello"}, mock)

# Regular args should only contain non-serializable values
assert "query" in regular
assert regular["query"] == "hello"
assert "data" not in regular
rlm._initialize_inputs({"data": stub, "query": "hello"}, mock)

# MockInterpreter should have received an execute call for the setup
assert mock.call_count == 1
assert mock.call_count == 2
code, variables = mock.call_history[0]
assert "import json" in code
assert "_raw_data" in variables
assert mock.call_history[1] == ("pass", {"query": "hello"})

def test_no_serializable_returns_all(self):
"""When no SandboxSerializable values exist, all args are returned."""
mock = MockInterpreter(responses=[FinalOutput({"answer": "42"})])
def test_regular_inputs_are_initialized_once(self):
mock = MockInterpreter(responses=["", ""])
rlm = RLM("query -> answer", max_iters=3, interpreter=mock)

rlm._inject_execution_context(mock, rlm._prepare_execution_tools())
regular = rlm._prepare_serializable_vars({"query": "hello"}, mock)
rlm._initialize_inputs({"query": "hello"}, mock)

assert regular == {"query": "hello"}
assert mock.call_count == 0
assert mock.call_history == [("pass", {"query": "hello"})]

def test_binary_payload_uses_base64_transport(self):
"""Non-UTF8 bytes should be transported via base64 and decoded in sandbox code."""
Expand All @@ -1302,12 +1311,13 @@ def test_binary_payload_uses_base64_transport(self):

payload = _BinarySerializable()
rlm._inject_execution_context(mock, rlm._prepare_execution_tools())
rlm._prepare_serializable_vars({"data": payload, "query": "hello"}, mock)
rlm._initialize_inputs({"data": payload, "query": "hello"}, mock)

assert mock.call_count == 1
assert mock.call_count == 2
code, variables = mock.call_history[0]
assert "_raw_data = base64.b64decode(_raw_data_base64)" in code
assert variables["_raw_data_base64"] == base64.b64encode(b"\xff\xfe\xfd").decode("ascii")
assert mock.call_history[1] == ("pass", {"query": "hello"})

def test_large_payload_not_inlined_in_code(self):
"""Large payloads should ride in the variables kwarg, not the code string.
Expand All @@ -1316,7 +1326,7 @@ def test_large_payload_not_inlined_in_code(self):
subsequent prompt and could blow past sandbox limits. The transport
contract is: code stays small, payload travels as a named variable.
"""
mock = MockInterpreter(responses=[""])
mock = MockInterpreter(responses=["", ""])
rlm = RLM("data, query -> answer", interpreter=mock)

large_text = "x" * (2 * 1024 * 1024) # 2 MB UTF-8 payload
Expand All @@ -1335,19 +1345,21 @@ def rlm_preview(self, max_chars: int = 500) -> str:
return f"LargeText({len(large_text)} chars)"

rlm._inject_execution_context(mock, rlm._prepare_execution_tools())
rlm._prepare_serializable_vars({"data": _LargeText(), "query": "hi"}, mock)
rlm._initialize_inputs({"data": _LargeText(), "query": "hi"}, mock)

assert mock.call_count == 1
assert mock.call_count == 2
code, variables = mock.call_history[0]
# Payload must be in variables, not the code string.
assert variables["_raw_data"] == large_text
assert mock.call_history[1] == ("pass", {"query": "hi"})
assert large_text not in code
assert len(code) < 1000

def test_forward_with_serializable(self):
"""Full forward() pass with a SandboxSerializable input."""
mock = MockInterpreter(responses=[
"", # setup execution for _prepare_serializable_vars
"", # serializable setup
"", # regular input initialization
FinalOutput({"answer": "done"}),
])
rlm = RLM("data, query -> answer", max_iters=3, interpreter=mock)
Expand All @@ -1359,8 +1371,8 @@ def test_forward_with_serializable(self):
result = rlm.forward(data=stub, query="test")
assert result.answer == "done"

# First call should be the serializable setup, second should be the iteration
assert mock.call_count == 2
# Setup and ordinary inputs run once before model-generated code.
assert mock.call_count == 3


@pytest.mark.deno
Expand All @@ -1387,7 +1399,7 @@ def rlm_preview(self, max_chars: int = 500) -> str:
with PythonInterpreter(tools={}) as interp:
rlm = RLM("data -> answer", interpreter=interp)
rlm._inject_execution_context(interp, rlm._prepare_execution_tools())
rlm._prepare_serializable_vars({"data": _LargeText()}, interp)
rlm._initialize_inputs({"data": _LargeText()}, interp)
result = interp.execute("print(len(data)); print(data[:6])")

assert str(len(large_text)) in result
Expand Down
Loading