diff --git a/dspy/predict/code_act.py b/dspy/predict/code_act.py index f53023993e..4a89a4923f 100644 --- a/dspy/predict/code_act.py +++ b/dspy/predict/code_act.py @@ -1,3 +1,4 @@ +import functools import inspect import logging from typing import Callable @@ -6,7 +7,7 @@ from dspy.adapters.types.tool import Tool from dspy.predict.program_of_thought import ProgramOfThought from dspy.predict.react import ReAct -from dspy.primitives.code_interpreter import CodeInterpreter, _validate_interpreter_factory +from dspy.primitives.code_interpreter import CodeInterpreter, _execution_instructions, _validate_interpreter_factory from dspy.primitives.python_interpreter import PythonInterpreter from dspy.signatures.signature import Signature, ensure_signature @@ -87,7 +88,7 @@ def _build_instructions(self, signature, tools): "For each iteration, you will generate a code snippet that either solves the task or progresses towards the solution.\n" "Ensure any output you wish to extract from the code is printed to the console. The code should be enclosed in a fenced code block.\n" f"When all information for producing the outputs ({outputs}) are available to be extracted, mark `finished=True` besides the final Python code.\n" - "You have access to the Python Standard Library and the following functions:" + "You have access to the execution environment and the following approved functions:" ) for idx, tool in enumerate(tools.values()): @@ -95,6 +96,19 @@ def _build_instructions(self, signature, tools): return instructions + @staticmethod + def _make_interpreter_tool(tool: Tool): + if inspect.iscoroutinefunction(tool.func): + async def invoke(**kwargs): + return await tool.acall(**kwargs) + else: + def invoke(**kwargs): + return tool(**kwargs) + + functools.update_wrapper(invoke, tool.func) + invoke.__signature__ = inspect.signature(tool.func) + return invoke + def forward(self, interpreter: CodeInterpreter | None = None, /, **kwargs): """Run the program with a fresh interpreter or a caller-owned override. @@ -110,14 +124,23 @@ def forward(self, interpreter: CodeInterpreter | None = None, /, **kwargs): "To use a caller-owned interpreter, pass it as the first positional argument when calling the module." ) with self._interpreter_context(interpreter) as interpreter: - # Define the tool functions in the interpreter + # Register approved host functions instead of copying their source + # into the guest, which breaks closures and bypasses host policy. for tool in self.tools.values(): - interpreter.execute(inspect.getsource(tool.func)) + interpreter.tools[tool.name] = self._make_interpreter_tool(tool) + if hasattr(interpreter, "_tools_registered"): + interpreter._tools_registered = False trajectory = {} max_iters = kwargs.pop("max_iters", self.max_iters) for idx in range(max_iters): - code_data = self.codeact(trajectory=trajectory, **kwargs) + codeact_args = {"trajectory": trajectory, **kwargs} + if instructions := _execution_instructions(interpreter): + if codeact_signature := getattr(self.codeact, "signature", None): + codeact_args["signature"] = codeact_signature.append_instructions( + f"\n\nExecution environment:\n{instructions}" + ) + code_data = self.codeact(**codeact_args) output = None code, error = self._parse_code(code_data) diff --git a/dspy/predict/rlm.py b/dspy/predict/rlm.py index 09ec51f381..138a67a4fa 100644 --- a/dspy/predict/rlm.py +++ b/dspy/predict/rlm.py @@ -10,7 +10,6 @@ from __future__ import annotations -import base64 import contextvars import functools import inspect @@ -32,6 +31,7 @@ CodeInterpreter, FinalOutput, _create_interpreter, + _execution_instructions, _validate_interpreter, _validate_interpreter_factory, ) @@ -62,7 +62,6 @@ - `llm_query_batched(prompts)` - query multiple prompts concurrently (much faster for multiple queries) - `print()` - ALWAYS print to see results - `SUBMIT({final_output_names})` - submit final output when done -- Standard libraries: re, json, collections, math, etc. IMPORTANT: This is ITERATIVE. Each code block you write will execute, you'll see the output, then you decide what to do next. Do NOT try to solve everything in one step. @@ -457,12 +456,9 @@ def _prepare_serializable_vars( try: payload_vars[raw_var_name] = payload.decode("utf-8") except UnicodeDecodeError: - encoded_var_name = f"{raw_var_name}_base64" - payload_vars[encoded_var_name] = base64.b64encode(payload).decode("ascii") - code_lines.extend([ - "import base64", - f"{raw_var_name} = base64.b64decode({encoded_var_name})", - ]) + encoded_var_name = f"{raw_var_name}_hex" + payload_vars[encoded_var_name] = payload.hex() + code_lines.append(f"{raw_var_name} = bytes.fromhex({encoded_var_name})") else: payload_vars[raw_var_name] = str(payload) @@ -479,7 +475,7 @@ def _prepare_serializable_vars( def _make_interpreter_tool(self, tool: Tool) -> Callable: """Preserve function metadata while routing execution through Tool.""" - if inspect.iscoroutinefunction(tool.func) or inspect.iscoroutinefunction(getattr(tool.func, "__call__", None)): + if inspect.iscoroutinefunction(tool.func) or inspect.iscoroutinefunction(type(tool.func).__call__): async def invoke(**kwargs): return await tool.acall(**kwargs) else: @@ -668,11 +664,17 @@ def _execute_iteration( ) -> Prediction | REPLHistory: """Execute one iteration. Returns Prediction if done, else updated REPLHistory.""" variables_info = [variable.format() for variable in variables] - action = self.generate_action( - variables_info=variables_info, - repl_history=history, - iteration=f"{iteration + 1}/{self.max_iters}", - ) + action_args = { + "variables_info": variables_info, + "repl_history": history, + "iteration": f"{iteration + 1}/{self.max_iters}", + } + if instructions := _execution_instructions(repl): + if action_signature := getattr(self.generate_action, "signature", None): + action_args["signature"] = action_signature.append_instructions( + f"\n\nExecution environment:\n{instructions}" + ) + action = self.generate_action(**action_args) if self.verbose: logger.info( f"RLM iteration {iteration + 1}/{self.max_iters}\n" @@ -761,11 +763,17 @@ async def _aexecute_iteration( ) -> Prediction | REPLHistory: """Async version: Execute one iteration.""" variables_info = [variable.format() for variable in variables] - pred = await self.generate_action.acall( - variables_info=variables_info, - repl_history=history, - iteration=f"{iteration + 1}/{self.max_iters}", - ) + action_args = { + "variables_info": variables_info, + "repl_history": history, + "iteration": f"{iteration + 1}/{self.max_iters}", + } + if instructions := _execution_instructions(repl): + if action_signature := getattr(self.generate_action, "signature", None): + action_args["signature"] = action_signature.append_instructions( + f"\n\nExecution environment:\n{instructions}" + ) + pred = await self.generate_action.acall(**action_args) if self.verbose: logger.info( f"RLM iteration {iteration + 1}/{self.max_iters}\n" diff --git a/dspy/primitives/base_module.py b/dspy/primitives/base_module.py index ab2214a69b..de2d89e34c 100644 --- a/dspy/primitives/base_module.py +++ b/dspy/primitives/base_module.py @@ -157,11 +157,23 @@ def dump_state(self, json_mode=True): return {name: param.dump_state(json_mode=json_mode) for name, param in self.named_parameters()} def load_state(self, state, *, allow_unsafe_lm_state=False): + import inspect + from dspy.predict.predict import Predict + def _accepts_unsafe_state(param): + try: + parameters = inspect.signature(param.load_state).parameters + except (TypeError, ValueError): + return isinstance(param, Predict) + explicit = parameters.get("allow_unsafe_lm_state") + if explicit is not None: + return explicit.kind in {inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY} + return any(value.kind == inspect.Parameter.VAR_KEYWORD for value in parameters.values()) + def _apply(module): for name, param in module.named_parameters(): - if isinstance(param, Predict): + if _accepts_unsafe_state(param): param.load_state(state[name], allow_unsafe_lm_state=allow_unsafe_lm_state) else: param.load_state(state[name]) diff --git a/dspy/primitives/code_interpreter.py b/dspy/primitives/code_interpreter.py index ce43810aef..56db3984bd 100644 --- a/dspy/primitives/code_interpreter.py +++ b/dspy/primitives/code_interpreter.py @@ -169,6 +169,12 @@ def _create_interpreter(factory: Callable[[], CodeInterpreter]) -> CodeInterpret return interpreter +def _execution_instructions(interpreter: CodeInterpreter) -> str | None: + """Return optional generated-code guidance without expanding the required protocol.""" + instructions = getattr(interpreter, "execution_instructions", None) + return instructions.strip() if isinstance(instructions, str) and instructions.strip() else None + + def _validate_interpreter(interpreter: Any) -> None: """Validate a caller-owned interpreter.""" if not isinstance(interpreter, CodeInterpreter): diff --git a/dspy/primitives/python_interpreter.py b/dspy/primitives/python_interpreter.py index 66061b6347..3a46b65f03 100644 --- a/dspy/primitives/python_interpreter.py +++ b/dspy/primitives/python_interpreter.py @@ -157,6 +157,10 @@ def my_tool(question: str) -> str: ``` """ + execution_instructions = """This is CPython running in Pyodide. The Python standard library and +Pyodide-compatible packages are available, but subprocesses and native system interfaces are unavailable. Host +filesystem, environment, and network access exist only when explicitly enabled by the interpreter configuration.""" + def __init__( self, deno_command: list[str] | None = None, diff --git a/tests/predict/test_code_act.py b/tests/predict/test_code_act.py index 471b16a07a..faebb3cef1 100644 --- a/tests/predict/test_code_act.py +++ b/tests/predict/test_code_act.py @@ -1,4 +1,3 @@ -import threading from unittest.mock import Mock import pytest @@ -82,7 +81,7 @@ def test_codeact_support_multiple_fields(pooled_interpreter): assert res.maximum == "6" assert res.minimum == "2" assert res.trajectory == { - "code_output_0": '"{\'maximum\': 6.0, \'minimum\': 2.0}\\n"', + "code_output_0": '"{\'maximum\': 6, \'minimum\': 2}\\n"', "generated_code_0": "result = extract_maximum_minimum('2, 3, 5, 6')\nprint(result)", } @@ -144,12 +143,7 @@ def test_codeact_code_execution_failure(pooled_interpreter): def test_codeact_evaluate_creates_one_interpreter_per_example(): - tool_registration_barrier = threading.Barrier(4) - def execute(code, variables): - if code.startswith("def add"): - tool_registration_barrier.wait(timeout=30) - return "" return "2\n" factory = MockInterpreterFactory(execute_fn=execute) @@ -172,13 +166,14 @@ def execute(code, variables): assert len(factory.instances) == 4 assert len({id(interpreter) for interpreter in factory.instances}) == 4 for interpreter in factory.instances: - assert interpreter.call_count == 2 + assert interpreter.call_count == 1 + assert interpreter.tools["add"](a=1, b=1) == 2 with pytest.raises(CodeInterpreterError, match="shutdown"): interpreter.execute("print('closed')") def test_codeact_factory_creates_fresh_interpreter_per_sequential_call(): - factory = MockInterpreterFactory(responses=["", "2\n"]) + factory = MockInterpreterFactory(responses=["2\n"]) program = CodeAct(BasicQA, tools=[add], interpreter_factory=factory) program.codeact = StaticPredictor(generated_code="print(add(1, 1))", finished=True) program.extractor = StaticPredictor(answer="2") @@ -194,6 +189,37 @@ def test_codeact_factory_creates_fresh_interpreter_per_sequential_call(): interpreter.execute("print('closed')") +def test_codeact_uses_interpreter_instructions_without_changing_predictor_schema(): + seen = {} + + class CapturingPredictor: + def __init__(self, signature): + self.signature = signature + + def __call__(self, **kwargs): + seen.update(kwargs) + return dspy.Prediction(generated_code="print(add(1, 1))", finished=True) + + interpreter = MockInterpreter(responses=["2\n"]) + interpreter.execution_instructions = "Use the constrained test dialect." + program = CodeAct(BasicQA, tools=[add]) + program.codeact = CapturingPredictor(program.codeact.signature) + program.extractor = StaticPredictor(answer="2") + + assert program(interpreter, question="What is 1+1?").answer == "2" + assert "execution_instructions" not in seen + assert interpreter.execution_instructions in seen["signature"].instructions + + +@pytest.mark.asyncio +async def test_codeact_preserves_async_host_tool(): + async def async_add(a: int, b: int) -> int: + return a + b + + wrapped = CodeAct._make_interpreter_tool(dspy.Tool(async_add)) + assert await wrapped(a=20, b=22) == 42 + + def test_codeact_allows_interpreter_as_signature_input(): factory = MockInterpreterFactory(responses=["", "CPython\n"]) program = CodeAct("interpreter -> answer", tools=[add], interpreter_factory=factory) @@ -229,7 +255,7 @@ def test_codeact_does_not_shutdown_caller_owned_interpreter(): program = CodeAct(BasicQA, tools=[add], interpreter_factory=factory) program.codeact = StaticPredictor(generated_code="print(add(1, 1))", finished=True) program.extractor = StaticPredictor(answer="2") - interpreter = MockInterpreter(responses=["", "2\n"]) + interpreter = MockInterpreter(responses=["2\n"]) try: result = program(interpreter, question="What is 1+1?") @@ -256,7 +282,7 @@ def test_codeact_shuts_down_factory_interpreter_when_extractor_raises(): def test_codeact_propagates_terminal_interpreter_failure_and_shuts_down(): - factory = MockInterpreterFactory(responses=["", CodeInterpreterError("protocol corrupt")]) + factory = MockInterpreterFactory(responses=[CodeInterpreterError("protocol corrupt")]) program = CodeAct(BasicQA, tools=[add], interpreter_factory=factory) program.codeact = StaticPredictor(generated_code="print(add(1, 1))", finished=True) diff --git a/tests/predict/test_rlm.py b/tests/predict/test_rlm.py index 09c0ef8ac8..c4c2785ac0 100644 --- a/tests/predict/test_rlm.py +++ b/tests/predict/test_rlm.py @@ -6,7 +6,6 @@ - Integration tests (@pytest.mark.deno): PythonInterpreter with Deno """ -import base64 from contextlib import contextmanager import pytest @@ -556,6 +555,26 @@ def test_action_signature_has_iteration_field(self): rlm = RLM("context -> answer") action_sig = rlm.generate_action.signature assert "iteration" in action_sig.input_fields + assert "execution_instructions" not in action_sig.input_fields + + def test_action_receives_interpreter_execution_instructions(self): + seen = {} + + class CapturingPredictor: + def __call__(self, **kwargs): + seen.update(kwargs) + return Prediction(reasoning="done", code="SUBMIT(answer='ok')") + + mock = MockInterpreter(responses=[FinalOutput({"answer": "ok"})]) + mock.execution_instructions = "Use only the constrained test dialect." + rlm = RLM("context -> answer", max_iters=1) + predictor = CapturingPredictor() + predictor.signature = rlm.generate_action.signature + rlm.generate_action = predictor + + assert rlm.forward(mock, context="test").answer == "ok" + assert "execution_instructions" not in seen + assert mock.execution_instructions in seen["signature"].instructions def test_format_output(self): """Test output formatting.""" @@ -1620,8 +1639,8 @@ def test_no_serializable_returns_all(self): assert regular == {"query": "hello"} assert mock.call_count == 0 - def test_binary_payload_uses_base64_transport(self): - """Non-UTF8 bytes should be transported via base64 and decoded in sandbox code.""" + def test_binary_payload_uses_portable_hex_transport(self): + """Non-UTF8 bytes should use a transport supported by constrained interpreters.""" mock = MockInterpreter(responses=[""]) rlm = RLM("data, query -> answer") @@ -1631,8 +1650,8 @@ def test_binary_payload_uses_base64_transport(self): assert mock.call_count == 1 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 "_raw_data = bytes.fromhex(_raw_data_hex)" in code + assert variables["_raw_data_hex"] == "fffefd" def test_large_payload_not_inlined_in_code(self): """Large payloads should ride in the variables kwarg, not the code string. diff --git a/tests/primitives/test_base_module.py b/tests/primitives/test_base_module.py index fb538f906b..d6af5af684 100644 --- a/tests/primitives/test_base_module.py +++ b/tests/primitives/test_base_module.py @@ -9,10 +9,51 @@ from litellm.types.utils import Usage import dspy +from dspy.predict.parameter import Parameter from dspy.primitives.prediction import Prediction from dspy.utils.dummies import DummyLM +def test_load_state_forwards_safe_flag_to_custom_parameter_kwargs(): + class CustomParameter(Parameter): + def __init__(self): + self.received = None + + def dump_state(self, json_mode=True): + return {} + + def load_state(self, state, **kwargs): + self.received = kwargs + + class CustomModule(dspy.Module): + def __init__(self): + self.parameter = CustomParameter() + + module = CustomModule() + module.load_state({"parameter": {}}, allow_unsafe_lm_state=True) + assert module.parameter.received == {"allow_unsafe_lm_state": True} + + +def test_load_state_preserves_custom_parameter_without_safe_flag(): + class CustomParameter(Parameter): + def __init__(self): + self.loaded = False + + def dump_state(self, json_mode=True): + return {} + + def load_state(self, state): + self.loaded = True + + class CustomModule(dspy.Module): + def __init__(self): + self.parameter = CustomParameter() + + module = CustomModule() + module.load_state({"parameter": {}}, allow_unsafe_lm_state=True) + assert module.parameter.loaded is True + + def test_deepcopy_basic(): signature = dspy.Signature("q -> a") cot = dspy.ChainOfThought(signature)