From 51ba5693bf5db9278aaf5dc2b277d2c0c5c538ef Mon Sep 17 00:00:00 2001 From: Aditya Vardhan <76904033+adtyavrdhn@users.noreply.github.com> Date: Tue, 30 Jun 2026 22:36:40 +0530 Subject: [PATCH] =?UTF-8?q?Revert=20"fix(code=5Fmode):=20tighten=20`run=5F?= =?UTF-8?q?code`=20prompt=20against=20phantom=20imports=20and=E2=80=A6"=20?= =?UTF-8?q?(#317)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 3f4cbd0a66a0b17c9e4ff75154fddf49986ff9a6. --- pydantic_ai_harness/code_mode/__init__.py | 17 +------ pydantic_ai_harness/code_mode/_capability.py | 48 -------------------- pydantic_ai_harness/code_mode/_toolset.py | 46 +++---------------- tests/code_mode/test_code_mode.py | 45 +----------------- 4 files changed, 11 insertions(+), 145 deletions(-) diff --git a/pydantic_ai_harness/code_mode/__init__.py b/pydantic_ai_harness/code_mode/__init__.py index 36c6d32..234438c 100644 --- a/pydantic_ai_harness/code_mode/__init__.py +++ b/pydantic_ai_harness/code_mode/__init__.py @@ -1,19 +1,6 @@ """Code mode capability: route tool calls through a sandboxed Python environment.""" from pydantic_ai_harness.code_mode._capability import CodeMode -from pydantic_ai_harness.code_mode._toolset import ( - CodeModeMount, - CodeModeOS, - CodeModeOSCallback, - CodeModeToolset, - default_run_code_instructions, -) +from pydantic_ai_harness.code_mode._toolset import CodeModeMount, CodeModeOS, CodeModeOSCallback, CodeModeToolset -__all__ = [ - 'CodeMode', - 'CodeModeMount', - 'CodeModeOS', - 'CodeModeOSCallback', - 'CodeModeToolset', - 'default_run_code_instructions', -] +__all__ = ['CodeMode', 'CodeModeMount', 'CodeModeOS', 'CodeModeOSCallback', 'CodeModeToolset'] diff --git a/pydantic_ai_harness/code_mode/_capability.py b/pydantic_ai_harness/code_mode/_capability.py index 611e061..3371c09 100644 --- a/pydantic_ai_harness/code_mode/_capability.py +++ b/pydantic_ai_harness/code_mode/_capability.py @@ -10,7 +10,6 @@ from pydantic import TypeAdapter, ValidationError from pydantic_ai import AbstractToolset from pydantic_ai.capabilities import AbstractCapability, CapabilityOrdering from pydantic_ai.capabilities._tool_search import ToolSearch as _ToolSearch -from pydantic_ai.exceptions import UserError from pydantic_ai.messages import ModelResponse, NativeToolSearchReturnPart, SystemPromptPart from pydantic_ai.tools import AgentDepsT, RunContext, ToolDefinition, ToolSelector from typing_extensions import TypedDict @@ -120,53 +119,8 @@ class CodeMode(AbstractCapability[AgentDepsT]): keeps the system prompt shorter and is the better choice. """ - instructions: str | None = None - """Extra prose appended to the built-in `run_code` description. - - `None` (default) uses the built-in prose alone. When set, the text is appended after - the built-in prose and before the auto-generated tool catalog (TypedDict definitions + - function signatures), so the model still gets the Monty-tuned base guidance plus your - additions. - - ```python - agent = Agent(model, capabilities=[CodeMode( - instructions='Project notes: prefer the `search` tool before `fetch`.', - )]) - ``` - - To replace the built-in prose wholesale instead of appending, use - `dangerously_replace_instructions`. The two are mutually exclusive.""" - - dangerously_replace_instructions: str | None = None - """Fully replace the built-in `run_code` description prose (the part before the catalog). - - `None` (default) keeps the built-in prose. When set, your text replaces it entirely; the - auto-generated tool catalog is still appended afterward, so `run_code` stays callable. - - Dangerous because the built-in prose is tuned to Monty's sandbox semantics (importable - modules, the no-classes / no-third-party rules, the return-value contract). Replacing it - discards that guidance: keep your replacement accurate to the sandbox or `run_code` calls - will fail. To start from the real base and edit it, import - `default_run_code_instructions`. Mutually exclusive with `instructions`. - - ```python - from pydantic_ai_harness.code_mode import default_run_code_instructions - - agent = Agent(model, capabilities=[CodeMode( - dangerously_replace_instructions=default_run_code_instructions() + '\\n\\nProject notes: ...', - )]) - ```""" - _announced_tools: set[str] = field(default_factory=set[str], init=False, repr=False) - def __post_init__(self) -> None: - if self.instructions is not None and self.dangerously_replace_instructions is not None: - raise UserError( - '`instructions` and `dangerously_replace_instructions` are mutually exclusive: ' - '`instructions` appends to the built-in prose, `dangerously_replace_instructions` ' - 'replaces it. Set at most one.' - ) - def get_ordering(self) -> CapabilityOrdering: """CodeMode wraps around ToolSearch so that search_tools stays native.""" return CapabilityOrdering(position='outermost', wraps=[_ToolSearch]) @@ -186,8 +140,6 @@ class CodeMode(AbstractCapability[AgentDepsT]): dynamic_catalog=self.dynamic_catalog, os_access=self.os_access, mount=self.mount, - instructions=self.instructions, - dangerously_replace_instructions=self.dangerously_replace_instructions, ) async def after_tool_execute( diff --git a/pydantic_ai_harness/code_mode/_toolset.py b/pydantic_ai_harness/code_mode/_toolset.py index 3eb0954..d3b60a8 100644 --- a/pydantic_ai_harness/code_mode/_toolset.py +++ b/pydantic_ai_harness/code_mode/_toolset.py @@ -94,7 +94,7 @@ Write and run Python code in a sandboxed environment. The sandbox uses Monty, a subset of Python. Key restrictions: - **No classes**: class definitions are not supported - **No third-party libraries**: only the standard library modules listed below can be used -- **Importable standard library modules**: `sys`, `typing`, `asyncio`, `math`, `json`, `re`, `datetime`, `os`, `pathlib`. These must be imported before use, just like in regular Python. For example: `import asyncio` then `results = await asyncio.gather(tool_one(...), tool_two(...))`. No other module is importable -- anything outside this list is unavailable. The functions listed below are already in scope; call them by name, do not import them from any module.""" +- **Importable standard library modules**: `sys`, `typing`, `asyncio`, `math`, `json`, `re`, `datetime`, `os`, `pathlib`. These must be imported before use, just like in regular Python. For example: `import asyncio` then `results = await asyncio.gather(tool_one(...), tool_two(...))`.""" # Timing/OS restriction line, swapped depending on what host access the agent # configured. Three states, because `mount` and `os` enable different things: @@ -129,12 +129,7 @@ The last expression's value is automatically captured as the return value -- you structured data. Use `print()` only for supplementary logging or debug output. Returns the last expression's value directly. If `print()` was also called, returns \ -`{"output": "", "result": }`. - -Do not paste large data as a code literal -- a long string on one line causes an unterminated-literal \ -error; reference the variable holding a prior result instead. Tool results are already typed objects: \ -read their fields directly, do not `json.loads` them or index a string as a dict, and never pass a \ -truncation placeholder like `...` or `??` from a clamped result back as an argument.\ +`{"output": "", "result": }`.\ """ @@ -154,16 +149,6 @@ def _base_description(*, has_os: bool, has_mount: bool) -> str: return f'{_RUN_CODE_DESCRIPTION_HEAD}\n{restriction}\n{_RUN_CODE_DESCRIPTION_TAIL}' -def default_run_code_instructions(*, has_os: bool = False, has_mount: bool = False) -> str: - """The built-in `run_code` prose, so you can build on it for `CodeMode.instructions`. - - Pass `has_os`/`has_mount` matching the `CodeMode` config you use them with, so the - OS-access restriction line matches what the sandbox actually allows. The tool catalog - is appended by `CodeMode` itself and is not part of this string. - """ - return _base_description(has_os=has_os, has_mount=has_mount) - - def _functions_header(*, has_sync: bool, has_async: bool) -> str: """Build the functions-header paragraph for the `run_code` tool description.""" base = ( @@ -296,15 +281,6 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]): so Tool Search discoveries don't bust the tool-definitions cache prefix. """ - instructions: str | None = None - """Extra prose appended after the built-in `run_code` base prose (catalog still follows). - `None` uses the built-in prose alone. Set via `CodeMode.instructions`.""" - - dangerously_replace_instructions: str | None = None - """Fully replace the built-in `run_code` base prose (the tool catalog is always appended). - `None` keeps the built-in prose. Set via `CodeMode.dangerously_replace_instructions`. - Mutually exclusive with `instructions` (enforced by `CodeMode.__post_init__`).""" - # init=False so `replace()` in `for_run` produces a fresh instance with _repl=None, # giving each agent run isolated REPL state. Lazy-initialized on first call_tool. _repl: MontyRepl | None = field(default=None, init=False, repr=False) @@ -391,7 +367,7 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]): has_os = self.os_access is not None has_mount = self.mount is not None if self.dynamic_catalog: - description = self._resolved_base(has_os=has_os, has_mount=has_mount) + description = _base_description(has_os=has_os, has_mount=has_mount) self._last_catalog = self._render_catalog(callable_defs) else: description = self._build_description(callable_defs, has_os=has_os, has_mount=has_mount) @@ -664,19 +640,11 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]): callable_defs[safe_name] = td return callable_defs, sanitized_to_original - def _resolved_base(self, *, has_os: bool, has_mount: bool) -> str: - """The base prose: a full replacement, the built-in default, or default + appended text.""" - if self.dangerously_replace_instructions is not None: - return self.dangerously_replace_instructions - base = _base_description(has_os=has_os, has_mount=has_mount) - if self.instructions is None: - return base - return f'{base}\n\n{self.instructions}' - - def _build_description(self, callable_defs: dict[str, ToolDefinition], *, has_os: bool, has_mount: bool) -> str: + @staticmethod + def _build_description(callable_defs: dict[str, ToolDefinition], *, has_os: bool, has_mount: bool) -> str: """Render the `run_code` description: base prose + TypedDicts + function signatures.""" - base = self._resolved_base(has_os=has_os, has_mount=has_mount) - catalog = self._render_catalog(callable_defs) + base = _base_description(has_os=has_os, has_mount=has_mount) + catalog = CodeModeToolset._render_catalog(callable_defs) if not catalog: return base return base + '\n\n' + catalog diff --git a/tests/code_mode/test_code_mode.py b/tests/code_mode/test_code_mode.py index dc50c9c..cc6ad1a 100644 --- a/tests/code_mode/test_code_mode.py +++ b/tests/code_mode/test_code_mode.py @@ -19,7 +19,7 @@ from pydantic_ai import ( Tool, ToolDefinition, ) -from pydantic_ai.exceptions import ModelRetry, UserError +from pydantic_ai.exceptions import ModelRetry from pydantic_ai.messages import ToolCallPart from pydantic_ai.models.test import TestModel from pydantic_ai.tool_manager import ParallelExecutionMode @@ -31,7 +31,7 @@ from pydantic_monty import NOT_HANDLED, MountDir, OSAccess, OsFunction from typing_extensions import TypedDict from pydantic_ai_harness import CodeMode -from pydantic_ai_harness.code_mode import CodeModeToolset, default_run_code_instructions +from pydantic_ai_harness.code_mode import CodeModeToolset from pydantic_ai_harness.code_mode._toolset import ( # pyright: ignore[reportPrivateUsage] _SEARCH_TOOLS_MODIFIER, _TOOL_SEARCH_ADDENDUM, @@ -236,47 +236,6 @@ class TestCodeMode: # The base description must tell the model to await tool calls. assert 'await' in description - async def test_instructions_appends_to_base_prose(self) -> None: - """`instructions` appends to the built-in prose, with the catalog still after it.""" - toolset = _build_function_toolset(add) - wrapper = CodeMode[None](instructions='PROJECT NOTE').get_wrapper_toolset(toolset) - assert isinstance(wrapper, CodeModeToolset) - description = (await wrapper.get_tools(build_run_context(None)))['run_code'].tool_def.description - assert description is not None - assert 'Monty' in description # built-in prose kept - assert 'PROJECT NOTE' in description # the caller's addition - assert 'async def add(*, a: int, b: int) -> int' in description # catalog preserved - # Ordering: base prose, then the addition, then the catalog. - assert description.index('Monty') < description.index('PROJECT NOTE') < description.index('async def add') - - async def test_dangerously_replace_instructions_replaces_base_prose(self) -> None: - """`dangerously_replace_instructions` replaces the built-in prose; the catalog is still appended.""" - toolset = _build_function_toolset(add) - wrapper = CodeMode[None](dangerously_replace_instructions='CUSTOM RUN_CODE PROSE').get_wrapper_toolset(toolset) - assert isinstance(wrapper, CodeModeToolset) - description = (await wrapper.get_tools(build_run_context(None)))['run_code'].tool_def.description - assert description is not None - assert description.startswith('CUSTOM RUN_CODE PROSE') - assert 'async def add(*, a: int, b: int) -> int' in description # catalog preserved - assert 'Monty' not in description # built-in prose replaced - - async def test_dangerously_replace_builds_on_exported_default(self) -> None: - """`default_run_code_instructions()` is exported so a replacement can start from the real base.""" - toolset = _build_function_toolset(add) - custom = default_run_code_instructions() + '\n\nPROJECT NOTE' - wrapper = CodeMode[None](dangerously_replace_instructions=custom).get_wrapper_toolset(toolset) - assert isinstance(wrapper, CodeModeToolset) - description = (await wrapper.get_tools(build_run_context(None)))['run_code'].tool_def.description - assert description is not None - assert 'Monty' in description # the default prose was kept - assert 'PROJECT NOTE' in description # the caller's addition - assert 'async def add' in description # catalog preserved - - def test_instructions_and_replace_are_mutually_exclusive(self) -> None: - """Setting both `instructions` and `dangerously_replace_instructions` raises `UserError`.""" - with pytest.raises(UserError, match='mutually exclusive'): - CodeMode[None](instructions='a', dangerously_replace_instructions='b') - async def test_run_code_executes_call_through_monty(self) -> None: """End-to-end: `run_code` runs Python in Monty and dispatches to a sync wrapped tool.""" toolset = _build_function_toolset(add)