diff --git a/dspy/predict/rlm.py b/dspy/predict/rlm.py index 8906d03a0e..e0b7287e8a 100644 --- a/dspy/predict/rlm.py +++ b/dspy/predict/rlm.py @@ -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(): @@ -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 @@ -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}" @@ -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.""" @@ -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) # ========================================================================= @@ -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 @@ -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.""" @@ -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: @@ -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 diff --git a/tests/predict/test_rlm.py b/tests/predict/test_rlm.py index e5a25497d0..17e9b10cb1 100644 --- a/tests/predict/test_rlm.py +++ b/tests/predict/test_rlm.py @@ -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")'}, @@ -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...", @@ -556,6 +557,7 @@ def failing_tool() -> str: raise RuntimeError("Tool failed!") mock = MockInterpreter(responses=[ + "", CodeInterpreterError("RuntimeError: Tool failed!"), FinalOutput({"answer": "recovered"}), ]) @@ -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"}), ]) @@ -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"}), ]) @@ -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) @@ -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")'}, @@ -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)"}, @@ -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"}), ]) @@ -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}, @@ -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 ]) @@ -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: @@ -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.""" @@ -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. @@ -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 @@ -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) @@ -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 @@ -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