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

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,
@@ -512,12 +513,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(
@@ -528,8 +524,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. 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

@@ -471,11 +471,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([
@@ -489,6 +494,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."""