Add invocation-scoped interpreter binding

Co-authored-by: Isaac Miller <isaac@cmpnd.ai>
Amp-Thread-ID: https://ampcode.com/threads/T-019fcce1-bc8b-753e-886f-9b77c0ccb376

Author: ampagent

docs/docs/api/modules/RLM.md

@@ -199,7 +199,7 @@ rlm = dspy.RLM(
 )
 ```
 
-RLM creates and shuts down one interpreter from this factory per invocation. It adds invocation-scoped tools to the returned interpreter's mutable `tools` dictionary, so remote sandboxes need a `CodeInterpreter` adapter that supports that protocol. To reuse a caller-owned interpreter, pass it as the first positional argument when calling the module: `rlm(interpreter, context=data, query=query)`. RLM updates its tools and output metadata but does not shut down or restore it. Reuse is supported only for sequential calls to the same RLM instance; use the factory path for concurrency.
+RLM creates and shuts down one interpreter from this factory per invocation. It adds invocation-scoped tools to the returned interpreter's mutable `tools` dictionary, so remote sandboxes need a `CodeInterpreter` adapter that supports that protocol. To reuse a caller-owned interpreter, pass it as the first positional argument when calling the module: `rlm(interpreter, context=data, query=query)`. RLM temporarily replaces its tools and output metadata, restores them afterward, and does not shut it down. Interpreter namespace state created during the invocation persists. Reuse is supported only for sequential calls to the same RLM instance; use the factory path for concurrency.
 
 If the factory exposes an `execution_instructions` string, RLM adds it to the action predictor's task instructions,
 which DSPy adapters place in the system prompt. Optimizers such as GEPA may therefore adapt the execution guidance

docs/docs/diving-deeper/rlm.md

@@ -93,7 +93,7 @@ A list of plain functions or `dspy.Tool` objects. RLM normalizes each to a `Tool
 A zero-argument callable that returns a fresh `CodeInterpreter` for one invocation. RLM may call the factory concurrently, and it always shuts down the returned interpreter. A class such as `PythonInterpreter` is already a factory; use `functools.partial` or a callable provider object when construction needs configuration. RLM adds invocation-scoped tools to the returned interpreter's mutable `tools` dictionary, so remote sandboxes need a `CodeInterpreter` adapter that supports that protocol.
 
 **`__call__(interpreter, **inputs)` / `acall(interpreter, **inputs)`**
-An escape hatch for a caller-owned interpreter, supplied as the first positional argument. RLM mutates its `tools` dictionary and, when supported, its output-field metadata, but does not shut down or restore the instance. Reuse is supported only for sequential calls to the same RLM instance, so retained variables and tool registrations stay within one program and trust boundary. Use `interpreter_factory` for concurrent invocations. A `PythonInterpreter` override must also stay on the thread where it was first used.
+An escape hatch for a caller-owned interpreter, supplied as the first positional argument. RLM temporarily replaces its `tools` dictionary and, when supported, its output-field metadata, then restores both after the invocation without shutting down the instance. Interpreter namespace state created during the invocation persists. Reuse is supported only for sequential calls to the same RLM instance, so retained variables stay within one program and trust boundary. Use `interpreter_factory` for concurrent invocations. A `PythonInterpreter` override must also stay on the thread where it was first used.
 
 **`dspy.SandboxSerializable`**
 The base class for inputs that need custom loading. Implement `sandbox_setup`, `to_sandbox`, `sandbox_assignment`, and `rlm_preview`. It also defines a Pydantic schema hook, so a subclass can be a typed field in a signature, as in `data: DataFrame = dspy.InputField()`.

dspy/predict/flex/bridge.py

@@ -42,7 +42,7 @@
 from dspy import CodeInterpreterError
 from dspy.adapters.types.base_type import Type as _CustomType
 from dspy.adapters.utils import annotation_allows_none, parse_value
-from dspy.primitives.code_interpreter import _create_interpreter
+from dspy.primitives.code_interpreter import _bind_interpreter, _create_interpreter
 from dspy.signatures.signature import make_signature
 from dspy.utils.exceptions import LMError
 
@@ -256,11 +256,9 @@ def forward(self, inputs: dict[str, Any]) -> Any:
         invocation = _Invocation(self, originals)
         interp = _create_interpreter(self._factory)
         try:
-            interp.tools.update({CONSTRUCT_TOOL: invocation.construct, CALL_TOOL: invocation.call})
-            # User tools callable by name in the sandbox.
-            interp.tools.update(
-                {name: _restoring_entrypoint(fn, originals) for name, fn in self._tool_callables().items()}
-            )
+            tools = {CONSTRUCT_TOOL: invocation.construct, CALL_TOOL: invocation.call}
+            tools.update({name: _restoring_entrypoint(fn, originals) for name, fn in self._tool_callables().items()})
+            _bind_interpreter(interp, {**interp.tools, **tools})
             interp.execute(SHIM_SETUP)
             interp.execute(self._module_src)  # defines the class in the sandbox
             code = (

dspy/predict/rlm.py

@@ -31,6 +31,7 @@
     CodeExecutionError,
     CodeInterpreter,
     FinalOutput,
+    _bind_interpreter,
     _create_interpreter,
     _validate_interpreter,
     _validate_interpreter_factory,
@@ -510,12 +511,7 @@ def _inject_execution_context(self, interpreter: CodeInterpreter, execution_tool
         even for user-provided interpreters. Each forward() call gets fresh tools with a
         fresh call counter, so we must inject on every execution.
         """
-        interpreter.tools.update(execution_tools)
-        if hasattr(interpreter, "output_fields"):
-            interpreter.output_fields = self._get_output_fields_info()
-        # Reset registration flag to force re-registration with fresh tools
-        if hasattr(interpreter, "_tools_registered"):
-            interpreter._tools_registered = False
+        _bind_interpreter(interpreter, {**interpreter.tools, **execution_tools}, self._get_output_fields_info())
 
     @contextmanager
     def _interpreter_context(
@@ -526,8 +522,12 @@ def _interpreter_context(
         """Yield a caller-owned interpreter or manage a factory-created one."""
         if interpreter is not None:
             _validate_interpreter(interpreter)
-            self._inject_execution_context(interpreter, execution_tools)
-            yield interpreter
+            original = (dict(interpreter.tools), getattr(interpreter, "output_fields", None))
+            try:
+                self._inject_execution_context(interpreter, execution_tools)
+                yield interpreter
+            finally:
+                _bind_interpreter(interpreter, *original)
             return
 
         interpreter = _create_interpreter(self._interpreter_factory)

dspy/primitives/code_interpreter.py

@@ -177,3 +177,14 @@ def _validate_interpreter(interpreter: Any) -> None:
     """Validate a caller-owned interpreter."""
     if not isinstance(interpreter, CodeInterpreter):
         raise TypeError(f"interpreter must implement CodeInterpreter, not {type(interpreter).__name__}.")
+
+
+def _bind_interpreter(interpreter, tools, output_fields=None) -> None:
+    if callable(bind := getattr(interpreter, "bind", None)):
+        bind(tools=tools, output_fields=output_fields)
+    else:
+        interpreter.tools.clear()
+        interpreter.tools.update(tools)
+        for name, value in (("output_fields", output_fields), ("_tools_registered", False)):
+            if hasattr(interpreter, name):
+                setattr(interpreter, name, value)

dspy/primitives/python_interpreter.py

@@ -287,7 +287,7 @@ def __init__(
         self.tools = dict(tools) if tools else {}
         self.output_fields = output_fields
         self.callbacks = list(callbacks or [])
-        self._tools_registered = False
+        self._tools_registered = not bool(self.tools or self.output_fields)
         # TODO later on add enable_run (--allow-run) by proxying subprocess.run through Deno.run() to fix 'emscripten does not support processes' error
 
         self._uses_default_deno_command = not deno_command
@@ -348,6 +348,11 @@ def __init__(
         "extensions are unavailable. Host filesystem, environment, and network access require explicit permission."
     )
 
+    def bind(self, *, tools: dict[str, Callable], output_fields: list[dict] | None = None) -> None:
+        fields = None if output_fields is None else [dict(field) for field in output_fields]
+        self.tools, self.output_fields = dict(tools), fields
+        self._tools_registered = False
+
     def _check_session_active(self) -> None:
         if self._session_ended:
             raise CodeInterpreterError(
@@ -479,7 +484,7 @@ def _register_tools(self) -> None:
             return
 
         # Build registration params with typed tool signatures
-        params = {}
+        params = {"outputs": self.output_fields or []}
 
         if self.tools:
             tools_info = []
@@ -490,14 +495,6 @@ def _register_tools(self) -> None:
                 })
             params["tools"] = tools_info
 
-        if self.output_fields:
-            params["outputs"] = self.output_fields
-
-        # Skip if nothing to register
-        if not params:
-            self._tools_registered = True
-            return
-
         self._send_request("register", params, "registering tools/outputs")
         self._tools_registered = True
 

tests/predict/test_rlm.py

@@ -502,11 +502,16 @@ async def test_factory_creates_and_shuts_down_one_interpreter_per_call(self):
     @pytest.mark.asyncio
     async def test_caller_owned_interpreter_can_be_reused_across_sequential_calls(self):
         factory = MockInterpreterFactory()
+
+        def backend_tool():
+            return "backend"
+
         interpreter = MockInterpreter(
             responses=[
                 FinalOutput({"answer": "first"}),
                 FinalOutput({"answer": "second"}),
-            ]
+            ],
+            tools={"backend_tool": backend_tool},
         )
         rlm = RLM("query -> answer", max_iters=1, interpreter_factory=factory)
         rlm.generate_action = make_mock_predictor([
@@ -520,6 +525,7 @@ async def test_caller_owned_interpreter_can_be_reused_across_sequential_calls(se
             assert sync_result.answer == "first"
             assert async_result.answer == "second"
             assert factory.instances == []
+            assert interpreter.tools == {"backend_tool": backend_tool}
             assert interpreter.execute("print('still open')") == ""
         finally:
             interpreter.shutdown()

tests/primitives/test_python_interpreter.py

@@ -19,7 +19,7 @@
 from pydantic import BaseModel, ConfigDict
 
 import dspy.primitives.python_interpreter as python_interpreter
-from dspy.primitives.code_interpreter import CodeExecutionError, CodeInterpreterError, FinalOutput
+from dspy.primitives.code_interpreter import CodeExecutionError, CodeInterpreterError, FinalOutput, _bind_interpreter
 from dspy.primitives.python_interpreter import (
     LARGE_VAR_THRESHOLD,
     PythonInterpreter,
@@ -740,6 +740,35 @@ def test_submit_valid_runtime_edge_cases(field, code, expected):
         assert interpreter.execute(code) == FinalOutput(expected)
 
 
+def test_bind_replaces_tools_and_submit_shape_without_resetting_state(pooled_interpreter):
+    tools = {"tool": lambda: "ok"}
+    fields = [{"name": "answer", "type": "str"}]
+    pooled_interpreter.execute("state = 41")
+    pooled_interpreter.bind(tools=tools, output_fields=fields)
+    tools.clear()
+    fields[0]["name"] = "changed"
+
+    assert pooled_interpreter.execute("tool()") == "ok"
+    assert pooled_interpreter.execute("state + 1") == 42
+    assert pooled_interpreter.execute("SUBMIT(answer='done')") == FinalOutput({"answer": "done"})
+    with pytest.raises(TypeError):
+        pooled_interpreter.bind(tools={}, output_fields=[None])
+    assert pooled_interpreter.execute("tool()") == "ok"
+
+    pooled_interpreter.bind(tools={})
+    assert pooled_interpreter.execute("SUBMIT('default')") == FinalOutput({"output": "default"})
+
+
+def test_bind_fallback_invalidates_legacy_registration():
+    class LegacyPythonInterpreter(PythonInterpreter):
+        bind = None
+
+    with LegacyPythonInterpreter(tools={"old": lambda: "old"}) as interpreter:
+        assert interpreter.execute("old()") == "old"
+        _bind_interpreter(interpreter, {"new": lambda: "new"})
+        assert interpreter.execute("new()") == "new"
+
+
 def test_tool_positional_args(configure_pooled_interpreter):
     """Test that tools work with positional arguments."""