furyhawk/pydantic-ai-harness
1
0
Fork 0
mirror of https://github.com/pydantic/pydantic-ai-harness.git synced 2026-07-21 02:45:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Revert "fix(code_mode): tighten run_code prompt against phantom imports and…" (#317)

Browse Source 
This reverts commit 3f4cbd0a66.
This commit is contained in:
Aditya Vardhan
2026-06-30 22:36:40 +05:30
committed by GitHub
parent 326c3d648d
commit 51ba5693bf
4 changed files with 11 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pydantic_ai_harness/code_mode/__init__.py
+2 -15
View File
@@ -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']
pydantic_ai_harness/code_mode/_capability.py
-48
View File
@@ -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(
pydantic_ai_harness/code_mode/_toolset.py
+7 -39
View File
@@ -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": "<printed text>", "result": <last expression>}`.
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": "<printed text>", "result": <last expression>}`.\
"""
@@ -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
tests/code_mode/test_code_mode.py
+2 -43
View File
@@ -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)