From 2b6e023b3469b8cdbddff65bc974af81d156751c Mon Sep 17 00:00:00 2001 From: Amp Date: Fri, 14 Aug 2026 20:19:28 +0000 Subject: [PATCH] Add invocation-scoped interpreter binding Co-authored-by: Isaac Miller Amp-Thread-ID: https://ampcode.com/threads/T-019fcce1-bc8b-753e-886f-9b77c0ccb376 --- docs/docs/api/modules/RLM.md | 2 +- docs/docs/diving-deeper/rlm.md | 2 +- dspy/predict/flex/bridge.py | 10 +++---- dspy/predict/rlm.py | 16 +++++------ dspy/primitives/code_interpreter.py | 11 ++++++++ dspy/primitives/python_interpreter.py | 17 +++++------ tests/predict/test_rlm.py | 8 +++++- tests/primitives/test_python_interpreter.py | 31 ++++++++++++++++++++- 8 files changed, 69 insertions(+), 28 deletions(-) diff --git a/docs/docs/api/modules/RLM.md b/docs/docs/api/modules/RLM.md index 1998e98648..2d44b69daf 100644 --- a/docs/docs/api/modules/RLM.md +++ b/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 diff --git a/docs/docs/diving-deeper/rlm.md b/docs/docs/diving-deeper/rlm.md index e1ca0bcf54..ec65f6337b 100644 --- a/docs/docs/diving-deeper/rlm.md +++ b/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()`. diff --git a/dspy/predict/flex/bridge.py b/dspy/predict/flex/bridge.py index cc3a63b1fe..c8798b0c94 100644 --- a/dspy/predict/flex/bridge.py +++ b/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 = ( diff --git a/dspy/predict/rlm.py b/dspy/predict/rlm.py index eb547d1ad6..d20b7732b1 100644 --- a/dspy/predict/rlm.py +++ b/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) diff --git a/dspy/primitives/code_interpreter.py b/dspy/primitives/code_interpreter.py index 8a0393ebff..f3e4def3b5 100644 --- a/dspy/primitives/code_interpreter.py +++ b/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) diff --git a/dspy/primitives/python_interpreter.py b/dspy/primitives/python_interpreter.py index 815c80b2a2..7b0a01dccd 100644 --- a/dspy/primitives/python_interpreter.py +++ b/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 diff --git a/tests/predict/test_rlm.py b/tests/predict/test_rlm.py index 405325c92a..73e404d583 100644 --- a/tests/predict/test_rlm.py +++ b/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() diff --git a/tests/primitives/test_python_interpreter.py b/tests/primitives/test_python_interpreter.py index 7f15fca033..cbe70eee5a 100644 --- a/tests/primitives/test_python_interpreter.py +++ b/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."""