mirror of
https://github.com/pydantic/pydantic-ai-harness.git
synced 2026-07-21 02:45:34 +00:00
Add experimental CE-discovery capabilities: context, subagents disk-load, docs, authoring (#304)
RepoContext (instruction autoload + asset inventory), SubAgents disk-load, PyaiDocs (on-demand pyai docs), RuntimeAuthoring (runtime capability authoring). Reviewed + hardened; validated under pydantic-ai 1.107 and 2.0.0.
This commit is contained in:
David SF
@@ -0,0 +1,111 @@
|
||||
# RuntimeAuthoring
|
||||
|
||||
> [!WARNING]
|
||||
> **Experimental.** This capability lives under `pydantic_ai_harness.experimental` and may
|
||||
> change or be removed in any release, without a deprecation period. Import it from the
|
||||
> experimental path -- there is no top-level export:
|
||||
>
|
||||
> ```python
|
||||
> from pydantic_ai_harness.experimental.authoring import RuntimeAuthoring
|
||||
> ```
|
||||
>
|
||||
> Importing any experimental capability emits a `HarnessExperimentalWarning`. Silence **all**
|
||||
> harness experimental warnings with a single filter (no per-capability lines needed):
|
||||
>
|
||||
> ```python
|
||||
> import warnings
|
||||
> from pydantic_ai_harness.experimental import HarnessExperimentalWarning
|
||||
>
|
||||
> warnings.filterwarnings('ignore', category=HarnessExperimentalWarning)
|
||||
> ```
|
||||
|
||||
Let an agent author, validate, and persist real pydantic-ai capabilities at runtime.
|
||||
|
||||
## The problem
|
||||
|
||||
A coding agent often discovers, mid-task, that it wants a behavior its host does not
|
||||
yet have: a guardrail, an extra instruction, a tool, a request hook. The capability
|
||||
surface to express that already exists -- but only a developer can write a capability
|
||||
class, wire it into the agent, and restart. The agent itself cannot extend its own host
|
||||
while it runs.
|
||||
|
||||
## The solution
|
||||
|
||||
`RuntimeAuthoring` exposes three tools:
|
||||
|
||||
- `author_capability(name, code)` -- write `code` to `<directory>/<name>.py`, import it,
|
||||
and validate it. Validation requires exactly one
|
||||
`pydantic_ai.capabilities.AbstractCapability` subclass that constructs with no
|
||||
arguments; the side-effect-free static getters (`get_instructions`, `get_toolset`,
|
||||
`get_native_tools`, `get_model_settings`, `get_serialization_name`) are exercised. The
|
||||
async lifecycle hooks are not run -- they need a live `RunContext`.
|
||||
- `list_authored_capabilities()` -- list authored capabilities with status and any
|
||||
validation error.
|
||||
- `disable_authored_capability(name)` -- stop a capability from being injected.
|
||||
|
||||
A "hook" is not a standalone object in pydantic-ai -- it is a method on a capability. So
|
||||
authoring a hook means authoring a capability that overrides one lifecycle method.
|
||||
|
||||
```python
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai_harness.experimental.authoring import RuntimeAuthoring
|
||||
|
||||
authoring = RuntimeAuthoring(directory=Path('.authored'))
|
||||
agent = Agent('anthropic:claude-sonnet-4-6', capabilities=[authoring])
|
||||
```
|
||||
|
||||
## Activation boundary
|
||||
|
||||
A capability **cannot** be added to a live, already-executing run. pydantic-ai resolves
|
||||
the effective capability set once at the start of each run (the run's `root_capability`
|
||||
is fixed; there is no setter). So an authored capability is live on the **next**
|
||||
`agent.run(...)`, not the run that authored it. This mirrors Loopy's runtime personas,
|
||||
which are usable on the next delegate call rather than mid-execution -- but one notch
|
||||
coarser: a persona adds no tools or hooks (it rides a single generic `delegate` tool), so
|
||||
it is usable later in the same run, whereas a full capability contributes tools and hooks
|
||||
that only exist once the run's toolset and capability chain are assembled at run start.
|
||||
|
||||
### Integration contract
|
||||
|
||||
The orchestrator drives the loop, so it owns the one-line contract: thread the store's
|
||||
active capabilities into each run. With `agent.run(..., capabilities=...)`, the authored
|
||||
capability is live on the very next loop iteration -- no process restart.
|
||||
|
||||
```python
|
||||
history = None
|
||||
while not done:
|
||||
extra = authoring.store.load_active()
|
||||
result = await agent.run(next_prompt, message_history=history, capabilities=extra)
|
||||
history = result.all_messages()
|
||||
```
|
||||
|
||||
Because the capabilities also persist to disk (`<directory>/<name>.py` plus a
|
||||
`manifest.json` index), a fresh process picks them up by constructing a new
|
||||
`RuntimeAuthoring` over the same `directory` and calling `store.load_active()`.
|
||||
|
||||
`manifest.json` records each capability's name, module file, class name, status, and last
|
||||
validation error -- the surface a UI can read to show what the agent has authored.
|
||||
|
||||
## Trust boundary and the sandboxed alternative
|
||||
|
||||
Authoring executes arbitrary Python in-process at import, construction, and run time. That
|
||||
is the same trust boundary an agent that already runs shell commands and edits files
|
||||
operates under, which is the deliberate choice here.
|
||||
|
||||
The sandboxed alternative is the dormant `pa` Monty hook-slot registration system in the
|
||||
Loopy tree (`pa/slots.py`, `pa/registration_tools.py`, `pa/capability.py`
|
||||
`PaRegistrations`, `pa/registrations.py`, `pa/registration_runtime.py`,
|
||||
`pa/monty_bridge.py`). It wires type-checked, resource-limited, allowlist-gated
|
||||
`pydantic_monty` snippets into typed hook slots instead of importing native `.py`. It
|
||||
trades native power for sandboxing and is also next-run-only. `RuntimeAuthoring` chooses
|
||||
native authoring; this note records the alternative so the tradeoff is on file.
|
||||
|
||||
## Typing
|
||||
|
||||
Imported authored code is dynamic, but nothing typed `Any` crosses back into the harness:
|
||||
every value pulled from an authored module is narrowed with `isinstance`/`issubclass`
|
||||
before use, and loaded instances are typed `AbstractCapability[object]`. Because
|
||||
`AgentDepsT` is contravariant, an `AbstractCapability[object]` is accepted by any agent's
|
||||
`capabilities=` parameter.
|
||||
@@ -0,0 +1,23 @@
|
||||
"""Runtime capability authoring: let an agent write, validate, and register real capabilities."""
|
||||
|
||||
from pydantic_ai_harness.experimental._warn import warn_experimental
|
||||
from pydantic_ai_harness.experimental.authoring._capability import RuntimeAuthoring
|
||||
from pydantic_ai_harness.experimental.authoring._store import AuthoredCapability, CapabilityStore
|
||||
from pydantic_ai_harness.experimental.authoring._toolset import AuthoringToolset
|
||||
from pydantic_ai_harness.experimental.authoring._validate import (
|
||||
CapabilityValidationError,
|
||||
load_capability_instance,
|
||||
validate_capability_file,
|
||||
)
|
||||
|
||||
warn_experimental('authoring')
|
||||
|
||||
__all__ = [
|
||||
'AuthoredCapability',
|
||||
'AuthoringToolset',
|
||||
'CapabilityStore',
|
||||
'CapabilityValidationError',
|
||||
'RuntimeAuthoring',
|
||||
'load_capability_instance',
|
||||
'validate_capability_file',
|
||||
]
|
||||
@@ -0,0 +1,90 @@
|
||||
"""Runtime authoring capability: let an agent write and register real pydantic-ai capabilities."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from pydantic_ai.capabilities import AbstractCapability
|
||||
from pydantic_ai.tools import AgentDepsT
|
||||
from pydantic_ai.toolsets import AgentToolset
|
||||
|
||||
from pydantic_ai_harness.experimental.authoring._store import CapabilityStore
|
||||
from pydantic_ai_harness.experimental.authoring._toolset import AuthoringToolset
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pydantic_ai._instructions import AgentInstructions
|
||||
|
||||
|
||||
_DEFAULT_GUIDANCE = (
|
||||
'You can author new pydantic-ai capabilities at runtime with `author_capability(name, code)`. '
|
||||
'A capability is a subclass of `pydantic_ai.capabilities.AbstractCapability` that constructs with '
|
||||
'no arguments and overrides one or more lifecycle hooks (a single overridden hook is a valid '
|
||||
'capability). Authored capabilities are validated immediately but become active on the next agent '
|
||||
'run, not the current one. Use `list_authored_capabilities` and `disable_authored_capability` to '
|
||||
'manage them.'
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RuntimeAuthoring(AbstractCapability[AgentDepsT]):
|
||||
"""Let an agent author, validate, and persist real pydantic-ai capabilities at runtime.
|
||||
|
||||
Exposes `author_capability(name, code)`, `list_authored_capabilities`, and
|
||||
`disable_authored_capability`. Authoring writes a real `.py` to `directory`,
|
||||
imports it, and validates it (exactly one `AbstractCapability` subclass that
|
||||
constructs with no arguments and whose static getters run). Authored
|
||||
capabilities hold live code, so they are not spec-serializable and are
|
||||
persisted as source rather than as a spec.
|
||||
|
||||
Activation boundary: a capability cannot be added to a live, already-executing
|
||||
run -- pydantic-ai resolves the capability set once at the start of each run.
|
||||
The authored capability becomes usable on the next `agent.run(...)`. The
|
||||
integration contract is one line on the orchestrator side: thread the store's
|
||||
active capabilities into the next run.
|
||||
|
||||
```python
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai_harness.experimental.authoring import RuntimeAuthoring
|
||||
|
||||
authoring = RuntimeAuthoring(directory=Path('.authored'))
|
||||
agent = Agent('anthropic:claude-sonnet-4-6', capabilities=[authoring])
|
||||
|
||||
# Loop: each iteration injects whatever the agent has authored so far.
|
||||
result = await agent.run('build a logging capability', capabilities=authoring.store.load_active())
|
||||
```
|
||||
|
||||
This executes authored Python in-process -- the same trust boundary an agent
|
||||
that already runs shell commands and edits files operates under. The dormant
|
||||
`pa` Monty hook-slot registration system is the sandboxed alternative; see the
|
||||
capability README.
|
||||
"""
|
||||
|
||||
directory: Path
|
||||
"""Directory holding the authored `<name>.py` files and the `manifest.json` index."""
|
||||
|
||||
guidance: str | None = None
|
||||
"""Static system-prompt guidance on authoring. Cache-stable. Leave `None` for the
|
||||
default, or set `''` to omit guidance entirely."""
|
||||
|
||||
@property
|
||||
def store(self) -> CapabilityStore:
|
||||
"""The disk-backed store. Call `store.load_active()` to inject authored capabilities into the next run."""
|
||||
return CapabilityStore(self.directory)
|
||||
|
||||
def get_instructions(self) -> AgentInstructions[AgentDepsT] | None:
|
||||
"""Static, cache-stable guidance on the authoring tools."""
|
||||
guidance = _DEFAULT_GUIDANCE if self.guidance is None else self.guidance
|
||||
return guidance or None
|
||||
|
||||
def get_toolset(self) -> AgentToolset[AgentDepsT] | None:
|
||||
"""Toolset providing the authoring tools over this capability's store."""
|
||||
return AuthoringToolset[AgentDepsT](self.store)
|
||||
|
||||
@classmethod
|
||||
def get_serialization_name(cls) -> str | None:
|
||||
"""Not spec-serializable: the capability holds a live, disk-backed store."""
|
||||
return None
|
||||
@@ -0,0 +1,155 @@
|
||||
"""Disk-backed store for runtime-authored capabilities.
|
||||
|
||||
Each authored capability is one `<name>.py` file under `directory`; a sibling
|
||||
`manifest.json` indexes them (name, module file, class, status, last validation
|
||||
error) and is the surface a UI can read. The store mirrors Loopy's persona
|
||||
persistence: a read-modify-write upsert keyed by name, fail-soft loading that
|
||||
skips corrupt entries rather than raising.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
import tempfile
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, ValidationError
|
||||
from pydantic_ai.capabilities import AbstractCapability
|
||||
|
||||
from pydantic_ai_harness.experimental.authoring._validate import (
|
||||
CapabilityValidationError,
|
||||
load_capability_instance,
|
||||
validate_capability_file,
|
||||
)
|
||||
|
||||
_NAME_RE = re.compile(r'^[a-z][a-z0-9_]*$')
|
||||
|
||||
|
||||
class AuthoredCapability(BaseModel):
|
||||
"""Manifest entry for one authored capability."""
|
||||
|
||||
name: str
|
||||
module_file: str
|
||||
class_name: str
|
||||
status: Literal['active', 'disabled'] = 'active'
|
||||
last_error: str | None = None
|
||||
|
||||
|
||||
class _Manifest(BaseModel):
|
||||
capabilities: list[AuthoredCapability] = []
|
||||
|
||||
|
||||
@dataclass
|
||||
class CapabilityStore:
|
||||
"""Read/write index of authored capability `.py` files under `directory`."""
|
||||
|
||||
directory: Path
|
||||
|
||||
@property
|
||||
def _manifest_path(self) -> Path:
|
||||
return self.directory / 'manifest.json'
|
||||
|
||||
def _load_manifest(self) -> _Manifest:
|
||||
path = self._manifest_path
|
||||
if not path.exists():
|
||||
return _Manifest()
|
||||
try:
|
||||
return _Manifest.model_validate_json(path.read_text(encoding='utf-8'))
|
||||
except (OSError, ValidationError, ValueError):
|
||||
return _Manifest()
|
||||
|
||||
def _save_manifest(self, manifest: _Manifest) -> None:
|
||||
self.directory.mkdir(parents=True, exist_ok=True)
|
||||
# Write to a temp file in the same directory, then atomically replace the
|
||||
# manifest, so a crash mid-write never leaves a partial/corrupt file that
|
||||
# `_load_manifest` would read as "no capabilities".
|
||||
fd, tmp_name = tempfile.mkstemp(dir=self.directory, prefix='manifest.', suffix='.json.tmp')
|
||||
try:
|
||||
with os.fdopen(fd, 'w', encoding='utf-8') as tmp:
|
||||
tmp.write(manifest.model_dump_json(indent=2))
|
||||
os.replace(tmp_name, self._manifest_path)
|
||||
except BaseException:
|
||||
os.unlink(tmp_name)
|
||||
raise
|
||||
|
||||
def _upsert(self, record: AuthoredCapability) -> None:
|
||||
manifest = self._load_manifest()
|
||||
for index, existing in enumerate(manifest.capabilities):
|
||||
if existing.name == record.name:
|
||||
manifest.capabilities[index] = record
|
||||
break
|
||||
else:
|
||||
manifest.capabilities.append(record)
|
||||
self._save_manifest(manifest)
|
||||
|
||||
def write(self, name: str, code: str) -> AuthoredCapability:
|
||||
"""Write `code` to `<name>.py`, validate it, and upsert the manifest entry.
|
||||
|
||||
Raises `ValueError` for an invalid name (before writing anything). A code
|
||||
that imports but fails validation is still written (so it can be
|
||||
inspected) and recorded with `last_error` set; `load_active` skips it.
|
||||
"""
|
||||
if not _NAME_RE.match(name):
|
||||
raise ValueError(
|
||||
f'invalid capability name {name!r}: use lowercase letters, digits, and underscores, '
|
||||
f'starting with a letter'
|
||||
)
|
||||
self.directory.mkdir(parents=True, exist_ok=True)
|
||||
path = self.directory / f'{name}.py'
|
||||
path.write_text(code, encoding='utf-8')
|
||||
try:
|
||||
class_name = validate_capability_file(path)
|
||||
record = AuthoredCapability(name=name, module_file=path.name, class_name=class_name)
|
||||
except CapabilityValidationError as exc:
|
||||
record = AuthoredCapability(name=name, module_file=path.name, class_name='', last_error=str(exc))
|
||||
self._upsert(record)
|
||||
return record
|
||||
|
||||
def disable(self, name: str) -> bool:
|
||||
"""Mark the named capability disabled so `load_active` stops returning it. Returns whether it existed."""
|
||||
manifest = self._load_manifest()
|
||||
found = False
|
||||
for record in manifest.capabilities:
|
||||
if record.name == name:
|
||||
record.status = 'disabled'
|
||||
found = True
|
||||
if found:
|
||||
self._save_manifest(manifest)
|
||||
return found
|
||||
|
||||
def list_all(self) -> list[AuthoredCapability]:
|
||||
"""Return every manifest entry, in insertion order."""
|
||||
return self._load_manifest().capabilities
|
||||
|
||||
def load_active(self) -> list[AbstractCapability[object]]:
|
||||
"""Construct every active authored capability for per-run injection.
|
||||
|
||||
Re-imports and re-constructs each active entry. Entries that fail to load
|
||||
(corrupt source, construction error) are skipped, not raised, so one bad
|
||||
capability never blocks the rest. A load outcome that disagrees with the
|
||||
record's `last_error` is persisted back to the manifest: a newly broken
|
||||
entry records its error, a re-fixed entry clears it, so the manifest stays
|
||||
truthful about which capabilities are actually active.
|
||||
"""
|
||||
manifest = self._load_manifest()
|
||||
instances: list[AbstractCapability[object]] = []
|
||||
changed = False
|
||||
for record in manifest.capabilities:
|
||||
if record.status != 'active':
|
||||
continue
|
||||
try:
|
||||
instances.append(load_capability_instance(self.directory / record.module_file))
|
||||
except CapabilityValidationError as exc:
|
||||
if record.last_error != str(exc):
|
||||
record.last_error = str(exc)
|
||||
changed = True
|
||||
continue
|
||||
if record.last_error is not None:
|
||||
record.last_error = None
|
||||
changed = True
|
||||
if changed:
|
||||
self._save_manifest(manifest)
|
||||
return instances
|
||||
@@ -0,0 +1,70 @@
|
||||
"""Authoring toolset: tools for the model to author, list, and disable capabilities."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic_ai.exceptions import ModelRetry
|
||||
from pydantic_ai.tools import AgentDepsT
|
||||
from pydantic_ai.toolsets import FunctionToolset
|
||||
|
||||
from pydantic_ai_harness.experimental.authoring._store import CapabilityStore
|
||||
|
||||
|
||||
class AuthoringToolset(FunctionToolset[AgentDepsT]):
|
||||
"""Exposes `author_capability`, `list_authored_capabilities`, and `disable_authored_capability`."""
|
||||
|
||||
def __init__(self, store: CapabilityStore) -> None:
|
||||
super().__init__()
|
||||
self._store = store
|
||||
self.add_function(self.author_capability, name='author_capability')
|
||||
self.add_function(self.list_authored_capabilities, name='list_authored_capabilities')
|
||||
self.add_function(self.disable_authored_capability, name='disable_authored_capability')
|
||||
|
||||
async def author_capability(self, name: str, code: str) -> str:
|
||||
"""Author a new pydantic-ai capability from Python source.
|
||||
|
||||
`code` must define exactly one `pydantic_ai.capabilities.AbstractCapability`
|
||||
subclass that constructs with no arguments. The capability is written to
|
||||
disk, imported, and validated immediately. It becomes usable on the next
|
||||
agent run (the next orchestrator loop iteration), not the current run.
|
||||
|
||||
Args:
|
||||
name: Identifier for the capability. Lowercase letters, digits, and
|
||||
underscores, starting with a letter. Reusing a name replaces the
|
||||
previous capability of that name.
|
||||
code: Complete Python source defining one `AbstractCapability` subclass.
|
||||
"""
|
||||
try:
|
||||
record = self._store.write(name, code)
|
||||
except ValueError as exc:
|
||||
raise ModelRetry(str(exc)) from exc
|
||||
if record.last_error is not None:
|
||||
return (
|
||||
f'Capability {name!r} was written but failed validation: {record.last_error}\n'
|
||||
f'Fix the code and call author_capability again with the same name.'
|
||||
)
|
||||
return (
|
||||
f'Capability {name!r} ({record.class_name}) authored and validated. It becomes active on '
|
||||
f'the next agent run (the next orchestrator loop iteration), not the current run.'
|
||||
)
|
||||
|
||||
async def list_authored_capabilities(self) -> str:
|
||||
"""List the capabilities authored so far, with their status and any validation error."""
|
||||
records = self._store.list_all()
|
||||
if not records:
|
||||
return 'No capabilities authored yet.'
|
||||
lines: list[str] = []
|
||||
for record in records:
|
||||
suffix = f' -- ERROR: {record.last_error}' if record.last_error is not None else ''
|
||||
class_name = record.class_name or '?'
|
||||
lines.append(f'- {record.name} [{record.status}] {class_name}{suffix}')
|
||||
return '\n'.join(lines)
|
||||
|
||||
async def disable_authored_capability(self, name: str) -> str:
|
||||
"""Disable an authored capability so it is no longer injected on the next run.
|
||||
|
||||
Args:
|
||||
name: Name of the capability to disable.
|
||||
"""
|
||||
if self._store.disable(name):
|
||||
return f'Capability {name!r} disabled; it will not be injected on the next run.'
|
||||
return f'No authored capability named {name!r}.'
|
||||
@@ -0,0 +1,132 @@
|
||||
"""Load and validate runtime-authored capability code without leaking `Any`.
|
||||
|
||||
Authored code is imported dynamically, so the values it produces are untyped at
|
||||
the import boundary. Every value crossing back into the harness is narrowed with
|
||||
`isinstance`/`issubclass` before use, so nothing typed `Any` escapes this module.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import sys
|
||||
from collections.abc import Mapping
|
||||
from types import ModuleType
|
||||
from typing import TYPE_CHECKING, TypeGuard
|
||||
|
||||
from pydantic_ai._instructions import normalize_instructions
|
||||
from pydantic_ai.capabilities import AbstractCapability
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class CapabilityValidationError(Exception):
|
||||
"""Raised when authored code cannot be imported, has the wrong shape, or fails to construct."""
|
||||
|
||||
|
||||
def _is_capability_subclass(obj: object) -> TypeGuard[type[AbstractCapability[object]]]:
|
||||
"""Narrow an attribute of unknown type to a concrete `AbstractCapability` subclass."""
|
||||
return isinstance(obj, type) and issubclass(obj, AbstractCapability)
|
||||
|
||||
|
||||
def _check_model_settings_return(value: object) -> None:
|
||||
"""Reject a `get_model_settings` return that the runtime cannot merge.
|
||||
|
||||
The runtime merges the return with `merge_model_settings`, which needs a
|
||||
`ModelSettings` mapping, a callable, or None. Typed `object` so an authored
|
||||
override that violates its declared signature is narrowed here, not trusted.
|
||||
"""
|
||||
if value is not None and not callable(value) and not isinstance(value, Mapping):
|
||||
raise CapabilityValidationError(
|
||||
f'get_model_settings() returned {type(value).__name__}, '
|
||||
f'expected a ModelSettings mapping, a callable, or None'
|
||||
)
|
||||
|
||||
|
||||
def _load_module(path: Path) -> ModuleType:
|
||||
"""Import `path` as a fresh module, not registered in `sys.modules`.
|
||||
|
||||
A fresh module object each call, plus suppressing the on-disk bytecode cache,
|
||||
means re-authoring under the same name always re-executes the new source.
|
||||
"""
|
||||
spec = importlib.util.spec_from_file_location(path.stem, path)
|
||||
if spec is None or spec.loader is None:
|
||||
raise CapabilityValidationError(f'cannot create an import spec for {path.name}')
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
# Don't let `exec_module` cache a `.pyc`. Re-authoring under the same name
|
||||
# changes the `.py`, but a stale cache would make the loader return the
|
||||
# previous bytecode for an unchanged file path, serving the old source.
|
||||
old_dont_write_bytecode = sys.dont_write_bytecode
|
||||
sys.dont_write_bytecode = True
|
||||
try:
|
||||
spec.loader.exec_module(module)
|
||||
finally:
|
||||
sys.dont_write_bytecode = old_dont_write_bytecode
|
||||
return module
|
||||
|
||||
|
||||
def _find_capability_class(module: ModuleType) -> type[AbstractCapability[object]]:
|
||||
"""Return the single `AbstractCapability` subclass defined in `module`.
|
||||
|
||||
Only classes whose `__module__` is the authored module count, so capabilities
|
||||
imported by the authored code (the base class itself, helpers from harness)
|
||||
are ignored. Anything other than exactly one is a validation error.
|
||||
"""
|
||||
found: list[type[AbstractCapability[object]]] = []
|
||||
for attr_name in dir(module):
|
||||
obj: object = getattr(module, attr_name)
|
||||
if _is_capability_subclass(obj) and obj.__module__ == module.__name__:
|
||||
found.append(obj)
|
||||
if not found:
|
||||
raise CapabilityValidationError('no `AbstractCapability` subclass found in the authored code')
|
||||
if len(found) > 1:
|
||||
names = ', '.join(sorted(cls.__name__ for cls in found))
|
||||
raise CapabilityValidationError(
|
||||
f'expected exactly one `AbstractCapability` subclass, found {len(found)}: {names}'
|
||||
)
|
||||
return found[0]
|
||||
|
||||
|
||||
def validate_capability_file(path: Path) -> str:
|
||||
"""Import, construct, and exercise the static getters of the authored capability.
|
||||
|
||||
Calls only the side-effect-free static getters (`get_instructions`,
|
||||
`get_toolset`, `get_native_tools`, `get_model_settings`,
|
||||
`get_serialization_name`). The async lifecycle hooks need a live `RunContext`
|
||||
and are not invoked here. Returns the capability class name on success.
|
||||
"""
|
||||
try:
|
||||
module = _load_module(path)
|
||||
cls = _find_capability_class(module)
|
||||
instance = cls()
|
||||
# Run the getter returns through the same coercion/shape-checks the
|
||||
# runtime uses, so a wrong return type fails validation here instead of
|
||||
# crashing the next `agent.run`. `get_native_tools`/`get_toolset` returns
|
||||
# are exercised by `list(...)`; the instructions/model-settings returns
|
||||
# would otherwise slip through untyped.
|
||||
normalize_instructions(instance.get_instructions())
|
||||
instance.get_toolset()
|
||||
list(instance.get_native_tools())
|
||||
_check_model_settings_return(instance.get_model_settings())
|
||||
cls.get_serialization_name()
|
||||
except CapabilityValidationError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise CapabilityValidationError(f'{type(exc).__name__}: {exc}') from exc
|
||||
return cls.__name__
|
||||
|
||||
|
||||
def load_capability_instance(path: Path) -> AbstractCapability[object]:
|
||||
"""Import the authored module and construct its capability instance.
|
||||
|
||||
Returns an `AbstractCapability[object]`; `AgentDepsT` is contravariant, so the
|
||||
instance is accepted by any agent's `capabilities=` parameter.
|
||||
"""
|
||||
try:
|
||||
module = _load_module(path)
|
||||
cls = _find_capability_class(module)
|
||||
return cls()
|
||||
except CapabilityValidationError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise CapabilityValidationError(f'{type(exc).__name__}: {exc}') from exc
|
||||
@@ -0,0 +1,123 @@
|
||||
# RepoContext
|
||||
|
||||
> [!WARNING]
|
||||
> **Experimental.** This capability lives under `pydantic_ai_harness.experimental` and may
|
||||
> change or be removed in any release, without a deprecation period. Import it from the
|
||||
> experimental path -- there is no top-level export:
|
||||
>
|
||||
> ```python
|
||||
> from pydantic_ai_harness.experimental.context import RepoContext
|
||||
> ```
|
||||
>
|
||||
> Importing any experimental capability emits a `HarnessExperimentalWarning`. Silence **all**
|
||||
> harness experimental warnings with a single filter (no per-capability lines needed):
|
||||
>
|
||||
> ```python
|
||||
> import warnings
|
||||
> from pydantic_ai_harness.experimental import HarnessExperimentalWarning
|
||||
>
|
||||
> warnings.filterwarnings('ignore', category=HarnessExperimentalWarning)
|
||||
> ```
|
||||
|
||||
Discover and load a repo's accumulated coding-assistant context engineering (CE).
|
||||
|
||||
## The problem
|
||||
|
||||
A repo accumulates CE for whatever coding assistant worked in it: instruction
|
||||
files (`CLAUDE.md`/`AGENTS.md`) scattered across the tree, and assets under
|
||||
`.claude`/`.agents`/`.codex`/`.grok` (skills, sub-agents, hooks). An agent that
|
||||
loads only the top-level instruction file misses the ancestor context and has no
|
||||
idea the rest of the setup exists, so it can neither honor it nor translate it.
|
||||
|
||||
## The solution
|
||||
|
||||
`RepoContext` bundles three strategies, each independently toggleable.
|
||||
|
||||
```python
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai_harness.experimental.context import RepoContext
|
||||
|
||||
agent = Agent(
|
||||
'anthropic:claude-sonnet-4-6',
|
||||
capabilities=[RepoContext(workspace_dir=Path('.'), home_dir=Path.home())],
|
||||
)
|
||||
```
|
||||
|
||||
### 1. Walk-up instruction autoload (on by default)
|
||||
|
||||
Loads `CLAUDE.md`/`AGENTS.md` from `workspace_dir` and every ancestor up to
|
||||
`home_dir` (inclusive). Precedence is ancestor-first, workspace-last: broadest
|
||||
context first, most specific last. Files are deduped by resolved real path and by
|
||||
content hash, so a symlinked `AGENTS.md -> CLAUDE.md` or two ancestors sharing
|
||||
identical content load once.
|
||||
|
||||
### 2. Asset inventory (on by default)
|
||||
|
||||
Exposes one tool, `inventory_agent_context()`, that reports where the repo's CE
|
||||
assets live -- the `.claude`/`.agents`/`.codex`/`.grok` roots and, within each,
|
||||
the `skills/` (SKILL.md), `agents/` (`.md`), and `settings.json` (hooks) it
|
||||
contains. It returns a structured `AgentContextInventory`; it locates assets and
|
||||
does not parse them, leaving translation to the orchestrator.
|
||||
|
||||
### 3. Nested-on-traversal (off by default)
|
||||
|
||||
When the model lists or reads a directory, surface that directory's
|
||||
`CLAUDE.md`/`AGENTS.md`. This couples to the host's list/read tools, so it is
|
||||
opt-in and configurable:
|
||||
|
||||
```python
|
||||
RepoContext(
|
||||
workspace_dir=Path('.'),
|
||||
nested_traversal=True,
|
||||
traversal_tool_names=frozenset({'list_dir', 'read_file'}), # match your tools
|
||||
traversal_path_arg='path', # the path arg key
|
||||
nested_inject='pointer', # or 'contents'
|
||||
)
|
||||
```
|
||||
|
||||
`nested_inject='pointer'` (default) appends a one-line note pointing at the file;
|
||||
`'contents'` inlines the file body. Each directory is surfaced at most once per
|
||||
run.
|
||||
|
||||
## Cache cost
|
||||
|
||||
Injecting file contents into the system prompt costs prompt-cache stability: a
|
||||
changed prefix re-bills the whole cached region. `RepoContext` keeps the two
|
||||
cache-relevant paths separate:
|
||||
|
||||
- Strategy 1 reads its files **once at run start** and injects them as static
|
||||
system instructions, so the cached prefix stays byte-identical across turns.
|
||||
- Strategy 3 is volatile (it depends on which directory was just touched), so its
|
||||
note is appended to the **tool result** in the message tail -- never to the
|
||||
system prompt -- and cannot invalidate the cached prefix.
|
||||
|
||||
## Configuration
|
||||
|
||||
```python
|
||||
RepoContext(
|
||||
workspace_dir, # Path -- the deepest dir the agent works in (required)
|
||||
home_dir=None, # Path | None -- shallowest dir to stop walk-up at, inclusive
|
||||
filenames=('CLAUDE.md', 'AGENTS.md'),
|
||||
autoload_instructions=True, # Strategy 1
|
||||
expose_inventory_tool=True, # Strategy 2
|
||||
inventory_tool_name='inventory_agent_context',
|
||||
nested_traversal=False, # Strategy 3
|
||||
nested_inject='pointer', # 'pointer' | 'contents'
|
||||
traversal_tool_names=frozenset({'list_directory', 'read_file'}),
|
||||
traversal_path_arg='path',
|
||||
asset_roots=('.claude', '.agents', '.codex', '.grok'),
|
||||
)
|
||||
```
|
||||
|
||||
## Scope
|
||||
|
||||
`RepoContext` locates and loads CE; it does not parse skill/sub-agent frontmatter
|
||||
or hook bodies, and it does not rewrite or translate assets. Strategy 1 reads its
|
||||
files once per run, so mid-run edits to those files are not reloaded.
|
||||
|
||||
## Further reading
|
||||
|
||||
- [Pydantic AI capabilities](https://ai.pydantic.dev/capabilities/)
|
||||
- [Pydantic AI hooks](https://ai.pydantic.dev/capabilities/#lifecycle-hooks)
|
||||
@@ -0,0 +1,11 @@
|
||||
"""Context capability: discover and load a repo's accumulated context engineering."""
|
||||
|
||||
from pydantic_ai_harness.experimental._warn import warn_experimental
|
||||
from pydantic_ai_harness.experimental.context._capability import RepoContext
|
||||
from pydantic_ai_harness.experimental.context._inventory import AgentContextInventory, AssetRoot
|
||||
from pydantic_ai_harness.experimental.context._loader import ContextFile
|
||||
from pydantic_ai_harness.experimental.context._toolset import RepoContextToolset
|
||||
|
||||
warn_experimental('context')
|
||||
|
||||
__all__ = ['AgentContextInventory', 'AssetRoot', 'ContextFile', 'RepoContext', 'RepoContextToolset']
|
||||
@@ -0,0 +1,196 @@
|
||||
"""RepoContext: discover and load a repo's accumulated context engineering."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass, field, replace
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Literal
|
||||
|
||||
from pydantic_ai.capabilities import AbstractCapability
|
||||
from pydantic_ai.messages import ToolCallPart
|
||||
from pydantic_ai.tools import AgentDepsT, RunContext, ToolDefinition
|
||||
from pydantic_ai.toolsets import AgentToolset
|
||||
|
||||
from pydantic_ai_harness.experimental.context._loader import (
|
||||
ContextFile,
|
||||
discover_instruction_files,
|
||||
find_dir_context_file,
|
||||
render_context_file,
|
||||
render_context_files,
|
||||
)
|
||||
from pydantic_ai_harness.experimental.context._toolset import RepoContextToolset
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pydantic_ai._instructions import AgentInstructions
|
||||
|
||||
|
||||
_INVENTORY_HINT = (
|
||||
'Call `{tool_name}` to map where this repo keeps its coding-assistant setup '
|
||||
'(instruction dirs, skills, sub-agents, and hooks) so you can read and translate it.'
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RepoContext(AbstractCapability[AgentDepsT]):
|
||||
"""Discover and load a repo's accumulated coding-assistant context engineering.
|
||||
|
||||
Three strategies, each independently toggleable:
|
||||
|
||||
1. Walk-up instruction autoload (`autoload_instructions`, on by default):
|
||||
load `CLAUDE.md`/`AGENTS.md` from `workspace_dir` and every ancestor up
|
||||
to `home_dir`, deduped, ancestor-first. These are read once at run start
|
||||
and injected as **static system instructions** via `get_instructions`, so
|
||||
they stay in the cached prefix and never re-read per turn.
|
||||
|
||||
2. Asset inventory (`expose_inventory_tool`, on by default): a tool that
|
||||
reports where the repo's CE assets live (`.claude`/`.agents`/`.codex`/
|
||||
`.grok` and their `skills/`, `agents/`, `settings.json`). It locates
|
||||
assets; it does not parse them.
|
||||
|
||||
3. Nested-on-traversal (`nested_traversal`, off by default): when the model
|
||||
lists or reads a directory (via a tool named in `traversal_tool_names`),
|
||||
surface that directory's `CLAUDE.md`/`AGENTS.md`. The note is appended to
|
||||
the **tool result** (message tail), not to system instructions, so it
|
||||
does not invalidate the cached prefix. `nested_inject='pointer'` (default)
|
||||
appends a one-line pointer; `'contents'` inlines the file body.
|
||||
|
||||
Cache note: injecting file contents into the system prompt costs prompt-cache
|
||||
stability. Strategy 1 is safe because its files are static; the volatile
|
||||
Strategy 3 content rides in the message tail instead.
|
||||
|
||||
```python
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai_harness.experimental.context import RepoContext
|
||||
|
||||
agent = Agent(
|
||||
'anthropic:claude-sonnet-4-6',
|
||||
capabilities=[RepoContext(workspace_dir=Path('.'), home_dir=Path.home())],
|
||||
)
|
||||
```
|
||||
"""
|
||||
|
||||
workspace_dir: Path
|
||||
"""The deepest directory the agent works in. The walk-up and asset scan are
|
||||
anchored here."""
|
||||
|
||||
home_dir: Path | None = None
|
||||
"""The shallowest directory to stop the walk-up at, inclusive. `None` (the
|
||||
default) scans only `workspace_dir` -- no walk-up."""
|
||||
|
||||
filenames: Sequence[str] = ('CLAUDE.md', 'AGENTS.md')
|
||||
"""Instruction filenames to look for, in within-directory precedence order."""
|
||||
|
||||
autoload_instructions: bool = True
|
||||
"""Strategy 1: load instruction files into the system prompt."""
|
||||
|
||||
expose_inventory_tool: bool = True
|
||||
"""Strategy 2: expose the asset-inventory tool."""
|
||||
|
||||
inventory_tool_name: str = 'inventory_agent_context'
|
||||
"""Name of the inventory tool exposed to the model."""
|
||||
|
||||
nested_traversal: bool = False
|
||||
"""Strategy 3: surface a directory's instruction file when the model lists or
|
||||
reads that directory. Off by default -- it couples to the list/read tools."""
|
||||
|
||||
nested_inject: Literal['pointer', 'contents'] = 'pointer'
|
||||
"""For Strategy 3: append a one-line `pointer`, or inline the file `contents`."""
|
||||
|
||||
traversal_tool_names: frozenset[str] = frozenset({'list_directory', 'read_file'})
|
||||
"""Tool names that trigger Strategy 3. Override to match the host's list/read
|
||||
tools (e.g. `frozenset({'list_dir', 'read_file'})`)."""
|
||||
|
||||
traversal_path_arg: str = 'path'
|
||||
"""The tool argument key holding the listed/read path."""
|
||||
|
||||
asset_roots: Sequence[str] = ('.claude', '.agents', '.codex', '.grok')
|
||||
"""Root directories the inventory tool scans, relative to `workspace_dir`."""
|
||||
|
||||
_context_files: list[ContextFile] | None = field(default=None, init=False, repr=False, compare=False)
|
||||
"""Cached walk-up result for this run, computed lazily on first access."""
|
||||
|
||||
_seen_dirs: set[str] = field(default_factory=set[str], init=False, repr=False, compare=False)
|
||||
"""Run-scoped set of directories already surfaced by Strategy 3."""
|
||||
|
||||
async def for_run(self, ctx: RunContext[AgentDepsT]) -> RepoContext[AgentDepsT]:
|
||||
"""Return a fresh per-run instance with isolated traversal/cache state."""
|
||||
return replace(self)
|
||||
|
||||
def _files(self) -> list[ContextFile]:
|
||||
if self._context_files is None:
|
||||
self._context_files = discover_instruction_files(self.workspace_dir, self.home_dir, self.filenames)
|
||||
return self._context_files
|
||||
|
||||
def get_instructions(self) -> AgentInstructions[AgentDepsT] | None:
|
||||
"""Static, cache-stable instructions: loaded files plus the inventory hint."""
|
||||
parts: list[str] = []
|
||||
if self.autoload_instructions:
|
||||
files = self._files()
|
||||
if files:
|
||||
parts.append(render_context_files(files, relative_to=self.workspace_dir))
|
||||
if self.expose_inventory_tool:
|
||||
parts.append(_INVENTORY_HINT.format(tool_name=self.inventory_tool_name))
|
||||
return '\n\n'.join(parts) or None
|
||||
|
||||
def get_toolset(self) -> AgentToolset[AgentDepsT] | None:
|
||||
"""The asset-inventory toolset, or `None` when the tool is disabled."""
|
||||
if not self.expose_inventory_tool:
|
||||
return None
|
||||
return RepoContextToolset[AgentDepsT](self.workspace_dir, self.asset_roots, self.inventory_tool_name)
|
||||
|
||||
async def after_tool_execute(
|
||||
self,
|
||||
ctx: RunContext[AgentDepsT],
|
||||
*,
|
||||
call: ToolCallPart,
|
||||
tool_def: ToolDefinition,
|
||||
args: dict[str, Any],
|
||||
result: Any,
|
||||
) -> Any:
|
||||
"""Strategy 3: append a directory's instruction file to a list/read result."""
|
||||
if not self.nested_traversal or call.tool_name not in self.traversal_tool_names:
|
||||
return result
|
||||
raw_path = args.get(self.traversal_path_arg)
|
||||
if not isinstance(raw_path, str):
|
||||
return result
|
||||
directory = self._resolve_directory(raw_path)
|
||||
context_file = find_dir_context_file(directory, self.filenames)
|
||||
if context_file is None:
|
||||
return result
|
||||
key = str(directory.resolve())
|
||||
if key in self._seen_dirs:
|
||||
return result
|
||||
if not isinstance(result, str):
|
||||
return result
|
||||
self._seen_dirs.add(key)
|
||||
note = self._render_note(context_file)
|
||||
return f'{result}\n\n{note}'
|
||||
|
||||
def _resolve_directory(self, raw_path: str) -> Path:
|
||||
candidate = Path(raw_path)
|
||||
if not candidate.is_absolute():
|
||||
candidate = self.workspace_dir / candidate
|
||||
return candidate.parent if candidate.is_file() else candidate
|
||||
|
||||
def _render_note(self, context_file: ContextFile) -> str:
|
||||
label = self._label(context_file.path)
|
||||
if self.nested_inject == 'contents':
|
||||
return render_context_file(context_file, label=label)
|
||||
return (
|
||||
f'<repo-context>This directory has {context_file.path.name} ({label}). '
|
||||
f'Read it if relevant to your task.</repo-context>'
|
||||
)
|
||||
|
||||
def _label(self, path: Path) -> str:
|
||||
try:
|
||||
return path.resolve().relative_to(self.workspace_dir.resolve()).as_posix()
|
||||
except ValueError:
|
||||
return path.as_posix()
|
||||
|
||||
@classmethod
|
||||
def get_serialization_name(cls) -> str | None:
|
||||
"""Serialization name for agent-spec support."""
|
||||
return 'RepoContext'
|
||||
@@ -0,0 +1,59 @@
|
||||
"""Locate (not parse) a repo's coding-assistant CE assets."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
_ROOT_NOTES = {
|
||||
'.codex': 'Codex uses TOML config; assets are derived from the .claude/.agents setup.',
|
||||
'.grok': 'Grok setup is derived from the .claude/.agents setup.',
|
||||
}
|
||||
|
||||
|
||||
class AssetRoot(BaseModel):
|
||||
"""Where CE assets live under a single root directory (e.g. `.claude`)."""
|
||||
|
||||
root: str = Field(description='The root directory name, relative to the workspace, e.g. ".claude".')
|
||||
exists: bool = Field(description='Whether the root directory is present in the workspace.')
|
||||
skills: list[str] = Field(default_factory=list, description='Paths to SKILL.md files found under skills/.')
|
||||
agents: list[str] = Field(default_factory=list, description='Paths to agent .md files found under agents/.')
|
||||
settings: str | None = Field(default=None, description='Path to settings.json (hooks), if present.')
|
||||
notes: str | None = Field(default=None, description='Format or derivation notes for this root, if any.')
|
||||
|
||||
|
||||
class AgentContextInventory(BaseModel):
|
||||
"""A map of where a repo's CE assets live, for an orchestrator to read or translate."""
|
||||
|
||||
roots: list[AssetRoot] = Field(default_factory=list[AssetRoot], description='One entry per scanned root directory.')
|
||||
|
||||
|
||||
def _relposix(path: Path, workspace: Path) -> str:
|
||||
try:
|
||||
return path.resolve().relative_to(workspace).as_posix()
|
||||
except ValueError:
|
||||
return path.as_posix()
|
||||
|
||||
|
||||
def scan_assets(workspace_dir: Path, asset_roots: Sequence[str]) -> AgentContextInventory:
|
||||
"""Scan `asset_roots` under `workspace_dir`, locating skills, agents, and hooks.
|
||||
|
||||
This locates assets only; it does not open or parse SKILL.md, agent `.md`, or
|
||||
`settings.json` contents.
|
||||
"""
|
||||
workspace = workspace_dir.resolve()
|
||||
roots: list[AssetRoot] = []
|
||||
for name in asset_roots:
|
||||
directory = workspace / name
|
||||
notes = _ROOT_NOTES.get(name)
|
||||
if not directory.is_dir():
|
||||
roots.append(AssetRoot(root=name, exists=False, notes=notes))
|
||||
continue
|
||||
skills = sorted(_relposix(p, workspace) for p in directory.glob('skills/**/SKILL.md') if p.is_file())
|
||||
agents = sorted(_relposix(p, workspace) for p in directory.glob('agents/*.md') if p.is_file())
|
||||
settings_path = directory / 'settings.json'
|
||||
settings = _relposix(settings_path, workspace) if settings_path.is_file() else None
|
||||
roots.append(AssetRoot(root=name, exists=True, skills=skills, agents=agents, settings=settings, notes=notes))
|
||||
return AgentContextInventory(roots=roots)
|
||||
@@ -0,0 +1,112 @@
|
||||
"""Walk-up discovery, dedup, precedence, and rendering of instruction files."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ContextFile:
|
||||
"""A single instruction file discovered during walk-up."""
|
||||
|
||||
directory: Path
|
||||
"""The directory the file was found in."""
|
||||
|
||||
path: Path
|
||||
"""The file's path."""
|
||||
|
||||
content: str
|
||||
"""The file's text content."""
|
||||
|
||||
|
||||
def _walk_dirs(workspace_dir: Path, home_dir: Path | None) -> list[Path]:
|
||||
"""Directories to scan, ancestor-first (home first, workspace last).
|
||||
|
||||
Walk up from `workspace_dir` to `home_dir` inclusive. When `home_dir` is
|
||||
`None`, or is not an ancestor of `workspace_dir`, only `workspace_dir` is
|
||||
scanned.
|
||||
"""
|
||||
workspace = workspace_dir.resolve()
|
||||
if home_dir is None:
|
||||
return [workspace]
|
||||
home = home_dir.resolve()
|
||||
chain: list[Path] = [workspace]
|
||||
if workspace != home:
|
||||
for parent in workspace.parents:
|
||||
chain.append(parent)
|
||||
if parent == home:
|
||||
break
|
||||
else:
|
||||
return [workspace]
|
||||
return list(reversed(chain))
|
||||
|
||||
|
||||
def discover_instruction_files(
|
||||
workspace_dir: Path,
|
||||
home_dir: Path | None,
|
||||
filenames: Sequence[str],
|
||||
) -> list[ContextFile]:
|
||||
"""Collect instruction files from `home_dir` down to `workspace_dir`.
|
||||
|
||||
Precedence is ancestor-first, workspace-last: the broadest context comes
|
||||
first and the most specific (closest to the model's recency window) comes
|
||||
last. Within a directory, `filenames` are tried in order.
|
||||
|
||||
Files are deduped by resolved real path and by content hash, so a symlinked
|
||||
`AGENTS.md -> CLAUDE.md` or two ancestors sharing identical content load
|
||||
once. The first occurrence in precedence order wins.
|
||||
"""
|
||||
seen_paths: set[Path] = set()
|
||||
seen_hashes: set[str] = set()
|
||||
found: list[ContextFile] = []
|
||||
for directory in _walk_dirs(workspace_dir, home_dir):
|
||||
for filename in filenames:
|
||||
candidate = directory / filename
|
||||
if not candidate.is_file():
|
||||
continue
|
||||
real = candidate.resolve()
|
||||
if real in seen_paths:
|
||||
continue
|
||||
content = candidate.read_text(encoding='utf-8', errors='replace')
|
||||
digest = hashlib.sha256(content.encode('utf-8')).hexdigest()
|
||||
if digest in seen_hashes:
|
||||
continue
|
||||
seen_paths.add(real)
|
||||
seen_hashes.add(digest)
|
||||
found.append(ContextFile(directory=directory, path=candidate, content=content))
|
||||
return found
|
||||
|
||||
|
||||
def find_dir_context_file(directory: Path, filenames: Sequence[str]) -> ContextFile | None:
|
||||
"""Return the first existing instruction file in `directory`, or `None`."""
|
||||
for filename in filenames:
|
||||
candidate = directory / filename
|
||||
if candidate.is_file():
|
||||
return ContextFile(
|
||||
directory=directory,
|
||||
path=candidate,
|
||||
content=candidate.read_text(encoding='utf-8', errors='replace'),
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def render_context_file(file: ContextFile, *, label: str) -> str:
|
||||
"""Render one file as a labeled block."""
|
||||
return f'<context-file path="{label}">\n{file.content}\n</context-file>'
|
||||
|
||||
|
||||
def render_context_files(files: Sequence[ContextFile], *, relative_to: Path) -> str:
|
||||
"""Render discovered files as labeled blocks in precedence order."""
|
||||
blocks = [render_context_file(file, label=_label(file.path, relative_to)) for file in files]
|
||||
return '\n\n'.join(blocks)
|
||||
|
||||
|
||||
def _label(path: Path, relative_to: Path) -> str:
|
||||
"""A stable display label: relative to `relative_to` when possible."""
|
||||
try:
|
||||
return path.resolve().relative_to(relative_to.resolve()).as_posix()
|
||||
except ValueError:
|
||||
return path.as_posix()
|
||||
@@ -0,0 +1,31 @@
|
||||
"""Toolset exposing the asset-inventory tool."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic_ai.tools import AgentDepsT
|
||||
from pydantic_ai.toolsets import FunctionToolset
|
||||
|
||||
from pydantic_ai_harness.experimental.context._inventory import AgentContextInventory, scan_assets
|
||||
|
||||
|
||||
class RepoContextToolset(FunctionToolset[AgentDepsT]):
|
||||
"""Exposes a single tool that reports where the repo's CE assets live."""
|
||||
|
||||
def __init__(self, workspace_dir: Path, asset_roots: Sequence[str], tool_name: str) -> None:
|
||||
super().__init__()
|
||||
self._workspace_dir = workspace_dir
|
||||
self._asset_roots = asset_roots
|
||||
self.add_function(self.inventory_agent_context, name=tool_name)
|
||||
|
||||
async def inventory_agent_context(self) -> AgentContextInventory:
|
||||
"""Report where this repo's coding-assistant setup lives.
|
||||
|
||||
Returns the locations of instruction dirs (`.claude`, `.agents`,
|
||||
`.codex`, `.grok`) and, within each, the `skills/`, `agents/`, and
|
||||
`settings.json` (hooks) it contains. This locates assets so you can read
|
||||
and translate them; it does not parse their contents.
|
||||
"""
|
||||
return scan_assets(self._workspace_dir, self._asset_roots)
|
||||
@@ -0,0 +1,67 @@
|
||||
# PyaiDocs
|
||||
|
||||
> [!WARNING]
|
||||
> **Experimental.** This capability lives under `pydantic_ai_harness.experimental` and may
|
||||
> change or be removed in any release, without a deprecation period. Import it from the
|
||||
> experimental path -- there is no top-level export:
|
||||
>
|
||||
> ```python
|
||||
> from pydantic_ai_harness.experimental.docs import PyaiDocs
|
||||
> ```
|
||||
>
|
||||
> Importing any experimental capability emits a `HarnessExperimentalWarning`. Silence **all**
|
||||
> harness experimental warnings with a single filter (no per-capability lines needed):
|
||||
>
|
||||
> ```python
|
||||
> import warnings
|
||||
> from pydantic_ai_harness.experimental import HarnessExperimentalWarning
|
||||
>
|
||||
> warnings.filterwarnings('ignore', category=HarnessExperimentalWarning)
|
||||
> ```
|
||||
|
||||
Give an agent a tool that locates and returns Pydantic AI documentation on demand.
|
||||
|
||||
## The problem
|
||||
|
||||
An agent that authors Pydantic AI capabilities, hooks, tools, or toolsets needs the
|
||||
current docs for those APIs. Preloading the docs into the system prompt spends context
|
||||
the agent rarely needs in full and pins a snapshot that drifts from `main`.
|
||||
|
||||
## The solution
|
||||
|
||||
`PyaiDocs` exposes one tool, `read_pyai_docs(topic)`, that locates the requested page and
|
||||
returns it verbatim -- nothing is bundled into context up front. Each call resolves the
|
||||
topic from a configured local checkout first, then falls back to fetching the page from
|
||||
`pydantic/pydantic-ai:main`, so it works in any environment.
|
||||
|
||||
Topics: `capabilities`, `hooks`, `tools`, `tools-advanced`, `toolsets`, `agent`.
|
||||
|
||||
```python
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai_harness.experimental.docs import PyaiDocs
|
||||
|
||||
agent = Agent(
|
||||
'anthropic:claude-sonnet-4-6',
|
||||
capabilities=[PyaiDocs(local_docs_path=Path('~/pydantic/ai/base/docs').expanduser())],
|
||||
)
|
||||
```
|
||||
|
||||
## Resolution order
|
||||
|
||||
1. **Local checkout** -- when `local_docs_path` (or the `PYDANTIC_AI_HARNESS_DOCS_PATH`
|
||||
env var) is set and `{path}/{topic}.md` exists, that file is read and returned.
|
||||
2. **Remote fetch** -- otherwise the page is fetched from
|
||||
`https://raw.githubusercontent.com/pydantic/pydantic-ai/main/docs/{topic}.md`.
|
||||
3. **Neither resolves** -- a descriptive error naming the local path tried and the URL.
|
||||
|
||||
The capability never runs git. Keep the local checkout current yourself; the remote path
|
||||
always reads `main`, so it is the fresh fallback.
|
||||
|
||||
## Configuration
|
||||
|
||||
| Option | Default | Purpose |
|
||||
| --- | --- | --- |
|
||||
| `local_docs_path` | `None` | Local pyai docs checkout to read first. Falls back to the `PYDANTIC_AI_HARNESS_DOCS_PATH` env var, then to the remote source. |
|
||||
| `cache` | `True` | Memoize each returned doc in-process for the capability's lifetime, so a topic is read or fetched at most once. |
|
||||
@@ -0,0 +1,9 @@
|
||||
"""Docs capability: an on-demand tool that locates Pydantic AI documentation."""
|
||||
|
||||
from pydantic_ai_harness.experimental._warn import warn_experimental
|
||||
from pydantic_ai_harness.experimental.docs._capability import PyaiDocs
|
||||
from pydantic_ai_harness.experimental.docs._toolset import PyaiDocsToolset, PyaiDocsTopic
|
||||
|
||||
warn_experimental('docs')
|
||||
|
||||
__all__ = ['PyaiDocs', 'PyaiDocsToolset', 'PyaiDocsTopic']
|
||||
@@ -0,0 +1,97 @@
|
||||
"""Docs capability: an on-demand tool that locates Pydantic AI documentation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from pydantic_ai.capabilities import AbstractCapability
|
||||
from pydantic_ai.tools import AgentDepsT
|
||||
from pydantic_ai.toolsets import AgentToolset
|
||||
|
||||
from pydantic_ai_harness.experimental.docs._toolset import PyaiDocsToolset, PyaiDocsTopic
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pydantic_ai._instructions import AgentInstructions
|
||||
|
||||
_DOCS_PATH_ENV = 'PYDANTIC_AI_HARNESS_DOCS_PATH'
|
||||
"""Env var holding a local pyai docs checkout path, used when `local_docs_path` is unset."""
|
||||
|
||||
_INSTRUCTIONS = (
|
||||
'You have a `read_pyai_docs` tool that returns Pydantic AI documentation on demand. '
|
||||
'Topics: capabilities, hooks, tools, tools-advanced, toolsets, agent. Read the relevant '
|
||||
'topic before authoring or modifying a Pydantic AI capability, hook, tool, or toolset, '
|
||||
'rather than relying on memory.'
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class PyaiDocs(AbstractCapability[AgentDepsT]):
|
||||
"""Locate and return Pydantic AI documentation on demand.
|
||||
|
||||
Exposes a single `read_pyai_docs(topic)` tool. Docs are located and returned
|
||||
when asked for -- never bundled into context. Each call resolves the topic
|
||||
from a configured local checkout first, then falls back to fetching the page
|
||||
from `pydantic/pydantic-ai:main`, so it works in any environment.
|
||||
|
||||
The local checkout path comes from `local_docs_path`, or the
|
||||
`PYDANTIC_AI_HARNESS_DOCS_PATH` env var when that is unset; with neither set
|
||||
every call goes straight to the remote source. The capability never runs git
|
||||
-- keep the local checkout current yourself; the remote path always reads
|
||||
`main`.
|
||||
|
||||
```python
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai_harness.experimental.docs import PyaiDocs
|
||||
|
||||
agent = Agent(
|
||||
'anthropic:claude-sonnet-4-6',
|
||||
capabilities=[PyaiDocs(local_docs_path=Path('~/pydantic/ai/base/docs').expanduser())],
|
||||
)
|
||||
```
|
||||
"""
|
||||
|
||||
local_docs_path: Path | None = None
|
||||
"""Local pyai docs checkout to read first. When `None`, falls back to the
|
||||
`PYDANTIC_AI_HARNESS_DOCS_PATH` env var, then to the remote source."""
|
||||
|
||||
cache: bool = True
|
||||
"""If `True`, each returned doc is memoized in-process for the capability's
|
||||
lifetime, so a topic is read or fetched at most once."""
|
||||
|
||||
_cache: dict[PyaiDocsTopic, str] = field(
|
||||
default_factory=dict[PyaiDocsTopic, str], init=False, repr=False, compare=False
|
||||
)
|
||||
"""In-memory doc cache shared with the toolset, so memoized docs outlive a
|
||||
single `get_toolset` call."""
|
||||
|
||||
def _resolved_local_path(self) -> Path | None:
|
||||
"""The local checkout path: `local_docs_path`, else the env var, else `None`.
|
||||
|
||||
`~` is expanded so a raw `~/...` path resolves to the local checkout
|
||||
instead of silently falling through to the remote source.
|
||||
"""
|
||||
if self.local_docs_path is not None:
|
||||
return self.local_docs_path.expanduser()
|
||||
env_path = os.environ.get(_DOCS_PATH_ENV)
|
||||
return Path(env_path).expanduser() if env_path else None
|
||||
|
||||
def get_instructions(self) -> AgentInstructions[AgentDepsT] | None:
|
||||
"""Static, cache-stable guidance on using the docs tool."""
|
||||
return _INSTRUCTIONS
|
||||
|
||||
def get_toolset(self) -> AgentToolset[AgentDepsT] | None:
|
||||
"""Toolset providing `read_pyai_docs` over the resolved local path and shared cache."""
|
||||
return PyaiDocsToolset[AgentDepsT](
|
||||
local_docs_path=self._resolved_local_path(),
|
||||
cache=self._cache if self.cache else None,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_serialization_name(cls) -> str | None:
|
||||
"""Serialization name for agent-spec support."""
|
||||
return 'PyaiDocs'
|
||||
@@ -0,0 +1,101 @@
|
||||
"""Docs toolset: a single `read_pyai_docs` tool that locates one pyai doc on demand."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
|
||||
import httpx
|
||||
from pydantic_ai.tools import AgentDepsT
|
||||
from pydantic_ai.toolsets import FunctionToolset
|
||||
|
||||
_REMOTE_BASE = 'https://raw.githubusercontent.com/pydantic/pydantic-ai/main/docs'
|
||||
"""Raw-markdown base for the live fallback. Tracks `pydantic/pydantic-ai:main`,
|
||||
and each topic maps to `{base}/{topic}.md` -- byte-identical to a local checkout."""
|
||||
|
||||
_FETCH_TIMEOUT = 30.0
|
||||
"""Per-request timeout (seconds) for the remote markdown fetch."""
|
||||
|
||||
|
||||
class PyaiDocsTopic(str, Enum):
|
||||
"""A Pydantic AI documentation page that `read_pyai_docs` can return.
|
||||
|
||||
Each value is the on-disk / remote file stem, so `{value}.md` resolves both a
|
||||
local checkout file and a raw-markdown URL.
|
||||
"""
|
||||
|
||||
capabilities = 'capabilities'
|
||||
hooks = 'hooks'
|
||||
tools = 'tools'
|
||||
tools_advanced = 'tools-advanced'
|
||||
toolsets = 'toolsets'
|
||||
agent = 'agent'
|
||||
|
||||
|
||||
class PyaiDocsToolset(FunctionToolset[AgentDepsT]):
|
||||
"""Exposes one `read_pyai_docs` tool that locates and returns a pyai doc.
|
||||
|
||||
Resolution per call: a configured local checkout first (when the topic's
|
||||
`{stem}.md` exists there), otherwise a raw-markdown fetch from `main`. The
|
||||
full doc is returned verbatim. Results are memoized in the shared `cache`
|
||||
dict (when caching is enabled) so a topic is read or fetched at most once.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
local_docs_path: Path | None,
|
||||
cache: dict[PyaiDocsTopic, str] | None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self._local_docs_path = local_docs_path
|
||||
# Shared with the capability so memoized docs outlive a single get_toolset
|
||||
# call; `None` disables caching entirely.
|
||||
self._cache = cache
|
||||
self.add_function(self.read_pyai_docs, name='read_pyai_docs')
|
||||
|
||||
async def read_pyai_docs(self, topic: PyaiDocsTopic) -> str:
|
||||
"""Return the Pydantic AI documentation for `topic` as markdown.
|
||||
|
||||
Reads from a configured local docs checkout when available, otherwise
|
||||
fetches the page from the live docs source. Returns the full page; call
|
||||
it once per topic you need.
|
||||
|
||||
Args:
|
||||
topic: Which documentation page to return. One of `capabilities`,
|
||||
`hooks`, `tools`, `tools-advanced`, `toolsets`, `agent`.
|
||||
"""
|
||||
if self._cache is not None and topic in self._cache:
|
||||
return self._cache[topic]
|
||||
|
||||
markdown = self._read_local(topic)
|
||||
if markdown is None:
|
||||
markdown = await self._fetch_remote(topic)
|
||||
|
||||
if self._cache is not None:
|
||||
self._cache[topic] = markdown
|
||||
return markdown
|
||||
|
||||
def _read_local(self, topic: PyaiDocsTopic) -> str | None:
|
||||
"""Return the local checkout's markdown for `topic`, or `None` to fall back to remote."""
|
||||
if self._local_docs_path is None:
|
||||
return None
|
||||
path = self._local_docs_path.expanduser() / f'{topic.value}.md'
|
||||
if not path.is_file():
|
||||
return None
|
||||
return path.read_text(encoding='utf-8')
|
||||
|
||||
async def _fetch_remote(self, topic: PyaiDocsTopic) -> str:
|
||||
"""Fetch `topic`'s markdown from the live source, or raise a descriptive error."""
|
||||
url = f'{_REMOTE_BASE}/{topic.value}.md'
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=_FETCH_TIMEOUT) as client:
|
||||
response = await client.get(url)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPError as exc:
|
||||
local = 'no local checkout configured' if self._local_docs_path is None else str(self._local_docs_path)
|
||||
raise RuntimeError(
|
||||
f'Could not locate the {topic.value!r} docs. Local source: {local}. '
|
||||
f'Remote fetch from {url} failed: {exc}'
|
||||
) from exc
|
||||
return response.text
|
||||
@@ -106,11 +106,85 @@ Hard errors propagate to stop the whole run. A `UsageLimitExceeded` from a child
|
||||
|
||||
The sub-agents are listed in the system prompt via `get_instructions`, using each agent's `description` (or a `SubAgent(description=...)` override). A sub-agent with no description is listed by name alone.
|
||||
|
||||
## Loading sub-agents from disk
|
||||
|
||||
A repo's markdown agent definitions become delegates without writing any `Agent` code. By default every `*.md` file under the conventional folders is loaded as a sub-agent, alongside the explicitly-passed `agents`.
|
||||
|
||||
```python
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai_harness.experimental.subagents import SubAgents
|
||||
|
||||
orchestrator = Agent(
|
||||
'anthropic:claude-opus-4-7',
|
||||
capabilities=[SubAgents(inherit_tools=True)], # auto-loads ./.agents/agents/ and ~/.agents/agents/
|
||||
)
|
||||
```
|
||||
|
||||
`agent_folders` controls where definitions come from. It defaults to `'agents'`, the conventional layout:
|
||||
|
||||
- A folder-name `str` (the default `'agents'`): for the project root (cwd) then the home root, load from `<root>/.agents/<name>/`, falling back to `<root>/.claude/<name>/` when `<root>/.agents/` is absent.
|
||||
- A sequence of paths loads from exactly those folders, in order.
|
||||
- `None` disables disk loading, exposing only the explicitly-passed `agents`.
|
||||
|
||||
### Definition format
|
||||
|
||||
A definition is a markdown file with optional frontmatter:
|
||||
|
||||
```markdown
|
||||
---
|
||||
name: researcher
|
||||
description: Researches a topic and reports findings
|
||||
tools: Read, Grep
|
||||
---
|
||||
You research topics. Report your findings, each with a source.
|
||||
```
|
||||
|
||||
- `name` is the delegate name (how the parent refers to it and how it is listed). It falls back to the filename stem when absent.
|
||||
- `description` drives the prompt listing.
|
||||
- The markdown body becomes the agent's instructions.
|
||||
- `tools` (or `allowed-tools`) is a comma-separated string or a YAML block list. See "Tools" below.
|
||||
- `model` and `color` are ignored: the model is inherited from the parent (see below), and `color` has no pyai equivalent.
|
||||
|
||||
Frontmatter is read by a small, dependency-free parser limited to those keys (`pyyaml` is not a harness dependency). Full YAML frontmatter is not supported.
|
||||
|
||||
### Models and effort
|
||||
|
||||
Disk agents inherit the parent run's model by default. Per agent, the caller can override the model and set a thinking/effort level via `agent_overrides`, keyed by the agent's name:
|
||||
|
||||
```python
|
||||
from pydantic_ai_harness.experimental.subagents import AgentOverride, SubAgents
|
||||
|
||||
SubAgents(
|
||||
agent_folders='agents',
|
||||
agent_overrides={'researcher': AgentOverride(model='anthropic:claude-sonnet-4-6', effort='high')},
|
||||
)
|
||||
```
|
||||
|
||||
Every agent the capability builds runs at a minimum thinking-effort floor. `MINIMUM_EFFORT_FLOOR` and the `clamp_effort(level, floor=...)` helper are exported so an orchestrator can apply the same floor to its own agents (that orchestrator-side application is the caller's responsibility). `clamp_effort` maps `None`/`False` to the floor, leaves `True` (provider-default effort) unchanged, and raises a concrete level below the floor up to it. Effort is applied through pyai's `ModelSettings.thinking`.
|
||||
|
||||
### Tools
|
||||
|
||||
A disk agent gets no tools by default (`inherit_tools` is `False`); set `inherit_tools=True` to expose the parent's tools to it through the `inherit_tools` mechanism, in which case its `tools` frontmatter is ignored. To map the frontmatter tool names to specific toolsets instead, pass a `tool_resolver`: it receives each tool name (so it can honor entries like `Bash(git:*)`) and returns the toolsets that provide it, or `None` for an unknown name, which is skipped with a warning.
|
||||
|
||||
```python
|
||||
def resolve(tool_name: str):
|
||||
return TOOLSETS.get(tool_name) # -> Sequence[AgentToolset[object]] | None
|
||||
|
||||
SubAgents(agent_folders='agents', tool_resolver=resolve)
|
||||
```
|
||||
|
||||
### Precedence
|
||||
|
||||
When the same name appears in more than one source, the higher-precedence one wins and the others are skipped with a warning: explicitly-passed `agents` first, then the project folder, then the home folder (and, for an explicit path sequence, earlier paths before later ones). A duplicate name within the explicitly-passed `agents` list is still an error.
|
||||
|
||||
## Configuration
|
||||
|
||||
```python
|
||||
SubAgents(
|
||||
agents=(), # Sequence[SubAgent[AgentDepsT]] -- each pairs an agent with its run controls
|
||||
agent_folders='agents',# folder-name str (convention) | Sequence[Path] | None (disable)
|
||||
agent_overrides={}, # Mapping[str, AgentOverride] -- per-disk-agent model/effort override
|
||||
tool_resolver=None, # Callable[[str], Sequence[AgentToolset[object]] | None] -- disk-agent tool mapping
|
||||
forward_usage=True, # share the parent's usage with sub-agent runs
|
||||
inherit_tools=False, # expose the parent's own tools to sub-agents (capability tools excluded)
|
||||
shared_capabilities=(),# capabilities applied to every sub-agent run
|
||||
|
||||
@@ -1,9 +1,19 @@
|
||||
"""Sub-agent capability: delegate self-contained tasks to named child agents."""
|
||||
|
||||
from pydantic_ai_harness.experimental._warn import warn_experimental
|
||||
from pydantic_ai_harness.experimental.subagents._capability import SubAgents
|
||||
from pydantic_ai_harness.experimental.subagents._capability import SubAgents, ToolResolver
|
||||
from pydantic_ai_harness.experimental.subagents._disk import AgentOverride
|
||||
from pydantic_ai_harness.experimental.subagents._effort import MINIMUM_EFFORT_FLOOR, clamp_effort
|
||||
from pydantic_ai_harness.experimental.subagents._toolset import SubAgent, SubAgentToolset
|
||||
|
||||
warn_experimental('subagents')
|
||||
|
||||
__all__ = ['SubAgent', 'SubAgentToolset', 'SubAgents']
|
||||
__all__ = [
|
||||
'MINIMUM_EFFORT_FLOOR',
|
||||
'AgentOverride',
|
||||
'SubAgent',
|
||||
'SubAgentToolset',
|
||||
'SubAgents',
|
||||
'ToolResolver',
|
||||
'clamp_effort',
|
||||
]
|
||||
|
||||
@@ -2,20 +2,34 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
import warnings
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from pydantic_ai.agent import AgentRunResult, EventStreamHandler
|
||||
from pydantic_ai.agent import Agent, AgentRunResult, EventStreamHandler
|
||||
from pydantic_ai.capabilities import AbstractCapability, AgentCapability, WrapRunHandler
|
||||
from pydantic_ai.settings import ModelSettings
|
||||
from pydantic_ai.tools import AgentDepsT, RunContext
|
||||
from pydantic_ai.toolsets import AgentToolset
|
||||
|
||||
from pydantic_ai_harness.experimental.subagents._disk import (
|
||||
AgentOverride,
|
||||
ParsedAgent,
|
||||
parse_agent_markdown,
|
||||
resolve_folders,
|
||||
)
|
||||
from pydantic_ai_harness.experimental.subagents._effort import clamp_effort
|
||||
from pydantic_ai_harness.experimental.subagents._toolset import SubAgent, SubAgentToolset
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pydantic_ai._instructions import AgentInstructions
|
||||
|
||||
ToolResolver = Callable[[str], 'Sequence[AgentToolset[object]] | None']
|
||||
"""Maps one tool name from a disk definition's `tools` list to the toolsets that
|
||||
provide it, or `None` when the name is unknown (the loader warns and skips it)."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class SubAgents(AbstractCapability[AgentDepsT]):
|
||||
@@ -31,7 +45,17 @@ class SubAgents(AbstractCapability[AgentDepsT]):
|
||||
wall-clock `timeout_seconds`, a per-run `max_calls` budget, an `on_failure`
|
||||
steering message, and optional `name`/`description` overrides). A delegate's
|
||||
name is its `SubAgent.name`, or the agent's own `name` when unset; two
|
||||
delegates resolving to the same name is an error.
|
||||
explicitly-passed delegates resolving to the same name is an error.
|
||||
|
||||
Sub-agents are also loaded from disk by default: each markdown agent definition
|
||||
under `./.agents/agents/` and `~/.agents/agents/` (or the `.claude/` equivalent)
|
||||
becomes a delegate, built with the parent's model. Disk delegates get no tools
|
||||
by default (`inherit_tools` is `False`); set `inherit_tools=True` to expose the
|
||||
parent's tools, or pass a `tool_resolver` to map their frontmatter tool names.
|
||||
Disk delegates coexist with explicitly-passed ones; explicitly-passed agents take
|
||||
precedence, then the project folder, then the home folder. A disk delegate whose
|
||||
name is already taken is skipped with a warning. Configure or disable this with
|
||||
`agent_folders`; see also `agent_overrides` and `tool_resolver`.
|
||||
|
||||
The parent's `deps` are forwarded to each sub-agent (sub-agents therefore
|
||||
share the parent's `AgentDepsT`), and by default the parent's `usage` is
|
||||
@@ -56,7 +80,34 @@ class SubAgents(AbstractCapability[AgentDepsT]):
|
||||
|
||||
agents: Sequence[SubAgent[AgentDepsT]] = ()
|
||||
"""The sub-agents to expose, each a `SubAgent` pairing an agent with its
|
||||
per-delegate run controls. See `SubAgent`."""
|
||||
per-delegate run controls. See `SubAgent`. These take precedence over any
|
||||
disk-loaded agents of the same name."""
|
||||
|
||||
agent_folders: str | Sequence[Path] | None = 'agents'
|
||||
"""Where to load markdown agent definitions from, in addition to `agents`.
|
||||
Defaults to the conventional layout, so constructing the capability auto-loads
|
||||
a repo's agent files with no extra configuration.
|
||||
|
||||
- a folder-name `str` (the default `'agents'` is the conventional layout): for
|
||||
the project root (cwd) then the home root, load from `<root>/.agents/<name>/`,
|
||||
falling back to `<root>/.claude/<name>/` when `<root>/.agents/` is absent.
|
||||
- a sequence of paths: load from exactly those folders, in order.
|
||||
- `None`: disable disk loading entirely (only `agents` are exposed).
|
||||
|
||||
Missing folders are skipped. Within a folder every `*.md` file is a candidate."""
|
||||
|
||||
agent_overrides: Mapping[str, AgentOverride] = field(default_factory=dict[str, AgentOverride])
|
||||
"""Per-disk-agent overrides keyed by the agent's name. An entry can set the
|
||||
agent's `model` (otherwise the parent's model is inherited) and its `effort`
|
||||
(otherwise the minimum floor). Has no effect on explicitly-passed `agents`."""
|
||||
|
||||
tool_resolver: ToolResolver | None = None
|
||||
"""Optional override for how a disk agent gets its tools. When set, each tool
|
||||
name in a definition's `tools`/`allowed-tools` frontmatter is passed to this
|
||||
resolver and the returned toolsets are attached to that agent; an unknown name
|
||||
(resolver returns `None`) is skipped with a warning. When unset, the
|
||||
frontmatter tool list is ignored and disk agents inherit the parent's tools
|
||||
via `inherit_tools` (set `inherit_tools=True` to expose them)."""
|
||||
|
||||
forward_usage: bool = True
|
||||
"""If `True`, the parent run's `usage` is shared with each sub-agent run, so
|
||||
@@ -103,8 +154,81 @@ class SubAgents(AbstractCapability[AgentDepsT]):
|
||||
f'set `SubAgent(name=...)` to disambiguate.'
|
||||
)
|
||||
by_name[name] = sub_agent
|
||||
# Disk agents are lower precedence than explicit ones and than earlier
|
||||
# folders, so a name already taken is shadowed (a warning, not an error --
|
||||
# overriding a home agent from the project, or a disk agent from code, is
|
||||
# the intended path).
|
||||
for sub_agent in self._load_disk_agents():
|
||||
name = sub_agent.resolved_name
|
||||
if name is None: # pragma: no cover - disk agents always get a name (frontmatter or stem)
|
||||
continue
|
||||
if name in by_name:
|
||||
warnings.warn(
|
||||
f'Disk sub-agent {name!r} is shadowed by a higher-precedence definition; skipping it.',
|
||||
stacklevel=2,
|
||||
)
|
||||
continue
|
||||
by_name[name] = sub_agent
|
||||
self._by_name = by_name
|
||||
|
||||
def _load_disk_agents(self) -> list[SubAgent[AgentDepsT]]:
|
||||
"""Build a `SubAgent` for every markdown definition in `agent_folders`.
|
||||
|
||||
Folders are returned in precedence order (project before home); within a
|
||||
folder, files are loaded in sorted name order for a stable listing.
|
||||
"""
|
||||
if self.agent_folders is None:
|
||||
return []
|
||||
result: list[SubAgent[AgentDepsT]] = []
|
||||
for folder in resolve_folders(self.agent_folders, Path.cwd(), Path.home()):
|
||||
if not folder.is_dir():
|
||||
continue
|
||||
for path in sorted(folder.glob('*.md')):
|
||||
try:
|
||||
text = path.read_text(encoding='utf-8')
|
||||
except (OSError, UnicodeDecodeError) as exc:
|
||||
warnings.warn(f'Skipping unreadable disk sub-agent file {str(path)!r}: {exc}', stacklevel=2)
|
||||
continue
|
||||
parsed = parse_agent_markdown(text)
|
||||
result.append(self._build_disk_agent(parsed.name or path.stem, parsed))
|
||||
return result
|
||||
|
||||
def _build_disk_agent(self, name: str, parsed: ParsedAgent) -> SubAgent[AgentDepsT]:
|
||||
"""Build one disk-defined sub-agent: parent model + floored effort, tools resolved or inherited.
|
||||
|
||||
The agent is constructed with `deps_type=object` so the parent's deps (of
|
||||
any type) flow through unused at delegation; this also lets a disk
|
||||
`SubAgent[object]` sit in the parent's `SubAgent[AgentDepsT]` roster.
|
||||
"""
|
||||
override = self.agent_overrides.get(name)
|
||||
model = override.model if override is not None else None
|
||||
effort = override.effort if override is not None else None
|
||||
toolsets = self._resolve_disk_tools(parsed.tools) if self.tool_resolver is not None else None
|
||||
agent = Agent(
|
||||
model,
|
||||
deps_type=object,
|
||||
name=name,
|
||||
description=parsed.description,
|
||||
instructions=parsed.body or None,
|
||||
model_settings=ModelSettings(thinking=clamp_effort(effort)),
|
||||
toolsets=toolsets,
|
||||
)
|
||||
return SubAgent(agent)
|
||||
|
||||
def _resolve_disk_tools(self, tool_names: Sequence[str]) -> list[AgentToolset[object]]:
|
||||
"""Map a definition's tool names to toolsets via `tool_resolver`, warning on unknown names."""
|
||||
resolver = self.tool_resolver
|
||||
if resolver is None: # pragma: no cover - only called when tool_resolver is set
|
||||
return []
|
||||
toolsets: list[AgentToolset[object]] = []
|
||||
for tool_name in tool_names:
|
||||
resolved = resolver(tool_name)
|
||||
if resolved is None:
|
||||
warnings.warn(f'Unknown tool {tool_name!r} in disk sub-agent definition; skipping it.', stacklevel=2)
|
||||
continue
|
||||
toolsets.extend(resolved)
|
||||
return toolsets
|
||||
|
||||
async def wrap_run(self, ctx: RunContext[AgentDepsT], *, handler: WrapRunHandler) -> AgentRunResult[Any]:
|
||||
"""Run the parent agent, then drop this run's delegation counts so they don't accumulate."""
|
||||
try:
|
||||
|
||||
@@ -0,0 +1,160 @@
|
||||
"""Load sub-agent definitions from markdown files on disk.
|
||||
|
||||
A definition is a markdown file with optional YAML-style frontmatter:
|
||||
|
||||
```markdown
|
||||
---
|
||||
name: researcher
|
||||
description: Researches a topic and reports findings
|
||||
tools: Read, Grep
|
||||
---
|
||||
You research topics. Report findings with sources.
|
||||
```
|
||||
|
||||
The frontmatter is parsed by a small, dependency-free reader limited to the keys
|
||||
coding assistants write (`name`, `description`, `model`, `color`, and `tools` or
|
||||
`allowed-tools`); `pyyaml` is not a runtime dependency of harness. The body after
|
||||
the frontmatter is the agent's instructions. `model` and `color` are ignored: the
|
||||
model is inherited from the parent (overridable via `SubAgents.agent_overrides`),
|
||||
and `color` has no pyai equivalent.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic_ai.models import KnownModelName, Model
|
||||
from pydantic_ai.settings import ThinkingLevel
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AgentOverride:
|
||||
"""Per-agent override for a disk-loaded sub-agent, keyed by the agent's name.
|
||||
|
||||
Both fields are optional. An unset `model` inherits the parent run's model; an
|
||||
unset `effort` runs at the capability's minimum effort floor (see
|
||||
`clamp_effort`).
|
||||
"""
|
||||
|
||||
model: Model | KnownModelName | str | None = None
|
||||
"""Model to run this disk agent with, in place of inheriting the parent's."""
|
||||
|
||||
effort: ThinkingLevel | None = None
|
||||
"""Thinking/reasoning level for this disk agent. Raised to at least the floor."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ParsedAgent:
|
||||
"""One parsed agent definition: frontmatter fields plus the markdown body."""
|
||||
|
||||
name: str | None
|
||||
description: str | None
|
||||
tools: tuple[str, ...]
|
||||
body: str
|
||||
|
||||
|
||||
def _strip_quotes(value: str) -> str:
|
||||
"""Drop a single layer of matching single or double quotes from a scalar."""
|
||||
if len(value) >= 2 and value[0] == value[-1] and value[0] in ('"', "'"):
|
||||
return value[1:-1]
|
||||
return value
|
||||
|
||||
|
||||
def _parse_frontmatter(lines: Sequence[str]) -> dict[str, str | list[str]]:
|
||||
"""Parse `key: value` and block-list (`- item`) frontmatter lines.
|
||||
|
||||
Only the shape coding assistants emit is supported: scalar values and, for
|
||||
list keys, either a `key: a, b` inline form (handled by callers) or a block
|
||||
list of `- item` lines under a key with an empty value.
|
||||
"""
|
||||
result: dict[str, str | list[str]] = {}
|
||||
current_list_key: str | None = None
|
||||
for raw in lines:
|
||||
if not raw.strip():
|
||||
continue
|
||||
stripped = raw.lstrip()
|
||||
if current_list_key is not None and stripped.startswith('- '):
|
||||
item = stripped[2:].strip()
|
||||
existing = result[current_list_key]
|
||||
if isinstance(existing, list) and item:
|
||||
existing.append(item)
|
||||
continue
|
||||
if ':' not in raw:
|
||||
current_list_key = None
|
||||
continue
|
||||
key, _, value = raw.partition(':')
|
||||
key = key.strip()
|
||||
value = value.strip()
|
||||
if value:
|
||||
result[key] = _strip_quotes(value)
|
||||
current_list_key = None
|
||||
else:
|
||||
result[key] = []
|
||||
current_list_key = key
|
||||
return result
|
||||
|
||||
|
||||
def _parse_tools(fields: dict[str, str | list[str]]) -> tuple[str, ...]:
|
||||
"""Read the `tools` or `allowed-tools` key as a tuple of tool-name strings."""
|
||||
raw = fields.get('tools')
|
||||
if raw is None:
|
||||
raw = fields.get('allowed-tools')
|
||||
if raw is None:
|
||||
return ()
|
||||
if isinstance(raw, list):
|
||||
return tuple(item for item in raw if item)
|
||||
return tuple(name.strip() for name in raw.split(',') if name.strip())
|
||||
|
||||
|
||||
def parse_agent_markdown(text: str) -> ParsedAgent:
|
||||
"""Parse a markdown agent definition into frontmatter fields and a body."""
|
||||
lines = text.splitlines()
|
||||
if not lines or lines[0].strip() != '---':
|
||||
return ParsedAgent(None, None, (), text.strip())
|
||||
closing: int | None = None
|
||||
for index in range(1, len(lines)):
|
||||
if lines[index].strip() == '---':
|
||||
closing = index
|
||||
break
|
||||
if closing is None:
|
||||
return ParsedAgent(None, None, (), text.strip())
|
||||
fields = _parse_frontmatter(lines[1:closing])
|
||||
body = '\n'.join(lines[closing + 1 :]).strip()
|
||||
name = fields.get('name')
|
||||
description = fields.get('description')
|
||||
return ParsedAgent(
|
||||
name=name if isinstance(name, str) else None,
|
||||
description=description if isinstance(description, str) else None,
|
||||
tools=_parse_tools(fields),
|
||||
body=body,
|
||||
)
|
||||
|
||||
|
||||
def _convention_folder(root: Path, leaf: str) -> Path:
|
||||
"""`<root>/.agents/<leaf>` when `.agents/` exists, else the `.claude/` equivalent."""
|
||||
if (root / '.agents').is_dir():
|
||||
return root / '.agents' / leaf
|
||||
return root / '.claude' / leaf
|
||||
|
||||
|
||||
def resolve_folders(agent_folders: str | Sequence[Path], cwd: Path, home: Path) -> list[Path]:
|
||||
"""Resolve the configured disk source into a precedence-ordered list of folders.
|
||||
|
||||
- a `str`: the convention `<root>/.agents/<str>/` (or `.claude/<str>/`) for the
|
||||
project root (`cwd`) then the home root, project first.
|
||||
- a sequence of paths: those folders verbatim, in order.
|
||||
|
||||
Folders resolving to the same absolute path are deduped (keeping the first), so
|
||||
a project root equal to the home root does not scan and warn about every agent
|
||||
twice.
|
||||
"""
|
||||
if isinstance(agent_folders, str):
|
||||
folders = [_convention_folder(cwd, agent_folders), _convention_folder(home, agent_folders)]
|
||||
else:
|
||||
folders = list(agent_folders)
|
||||
seen: dict[Path, Path] = {}
|
||||
for folder in folders:
|
||||
seen.setdefault(folder.resolve(), folder)
|
||||
return list(seen.values())
|
||||
@@ -0,0 +1,33 @@
|
||||
"""Minimum thinking-effort floor shared by the sub-agent capability and its orchestrator."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic_ai.settings import ThinkingEffort, ThinkingLevel
|
||||
|
||||
MINIMUM_EFFORT_FLOOR: ThinkingEffort = 'low'
|
||||
"""Lowest thinking effort any agent this capability builds is allowed to run at.
|
||||
|
||||
`clamp_effort` raises anything below this to the floor. The constant is exported
|
||||
so an orchestrator that builds its own agents can apply the same floor to them
|
||||
(the orchestrator-side application is the caller's responsibility)."""
|
||||
|
||||
_EFFORT_RANK: dict[ThinkingEffort, int] = {'minimal': 0, 'low': 1, 'medium': 2, 'high': 3, 'xhigh': 4}
|
||||
"""Ordering of the concrete effort levels, low to high."""
|
||||
|
||||
|
||||
def clamp_effort(level: ThinkingLevel | None, floor: ThinkingEffort = MINIMUM_EFFORT_FLOOR) -> ThinkingLevel:
|
||||
"""Raise a thinking level to at least `floor`.
|
||||
|
||||
- `None` or `False` (thinking unset or disabled, both below any floor) become `floor`.
|
||||
- `True` (thinking on at the provider's default effort) is left as `True`: its
|
||||
magnitude is provider-defined and not comparable to a named level, so it is
|
||||
neither downgraded nor bumped.
|
||||
- A concrete effort below `floor` becomes `floor`; one at or above `floor` is unchanged.
|
||||
"""
|
||||
if level is None or level is False:
|
||||
return floor
|
||||
if level is True:
|
||||
return True
|
||||
if _EFFORT_RANK[level] < _EFFORT_RANK[floor]:
|
||||
return floor
|
||||
return level
|
||||
@@ -196,9 +196,13 @@ class SubAgentToolset(FunctionToolset[AgentDepsT]):
|
||||
usage = ctx.usage if self._forward_usage else None
|
||||
usage_limits = None
|
||||
|
||||
# A sub-agent with no model of its own (e.g. one loaded from disk) inherits
|
||||
# the parent run's model; one that brought its own keeps it.
|
||||
model = None if sub_agent.agent.model is not None else ctx.model
|
||||
run = sub_agent.agent.run(
|
||||
task,
|
||||
deps=ctx.deps,
|
||||
model=model,
|
||||
usage=usage,
|
||||
usage_limits=usage_limits,
|
||||
toolsets=toolsets,
|
||||
|
||||
@@ -0,0 +1,481 @@
|
||||
"""Tests for the RuntimeAuthoring capability."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import os
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.capabilities import AbstractCapability
|
||||
from pydantic_ai.exceptions import ModelRetry
|
||||
from pydantic_ai.messages import ToolReturnPart
|
||||
from pydantic_ai.models.test import TestModel
|
||||
|
||||
from pydantic_ai_harness.experimental.authoring import (
|
||||
AuthoringToolset,
|
||||
CapabilityStore,
|
||||
CapabilityValidationError,
|
||||
RuntimeAuthoring,
|
||||
)
|
||||
from pydantic_ai_harness.experimental.authoring._validate import (
|
||||
load_capability_instance,
|
||||
validate_capability_file,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def anyio_backend() -> str:
|
||||
"""Run async tests on the asyncio backend (matching upstream pydantic-ai)."""
|
||||
return 'asyncio'
|
||||
|
||||
|
||||
VALID_CODE = """
|
||||
from dataclasses import dataclass
|
||||
|
||||
from pydantic_ai.capabilities import AbstractCapability
|
||||
from pydantic_ai.toolsets import FunctionToolset
|
||||
|
||||
|
||||
class _MarkerToolset(FunctionToolset):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.add_function(self.marker, name='marker')
|
||||
|
||||
async def marker(self) -> str:
|
||||
return 'marker-result'
|
||||
|
||||
|
||||
@dataclass
|
||||
class MarkerCapability(AbstractCapability):
|
||||
def get_instructions(self):
|
||||
return 'MARKER_INSTRUCTION'
|
||||
|
||||
def get_toolset(self):
|
||||
return _MarkerToolset()
|
||||
"""
|
||||
|
||||
NO_SUBCLASS_CODE = 'value = 1\n'
|
||||
|
||||
MULTI_CODE = """
|
||||
from dataclasses import dataclass
|
||||
|
||||
from pydantic_ai.capabilities import AbstractCapability
|
||||
|
||||
|
||||
@dataclass
|
||||
class CapA(AbstractCapability):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class CapB(AbstractCapability):
|
||||
pass
|
||||
"""
|
||||
|
||||
CONSTRUCT_FAILS_CODE = """
|
||||
from dataclasses import dataclass
|
||||
|
||||
from pydantic_ai.capabilities import AbstractCapability
|
||||
|
||||
|
||||
@dataclass
|
||||
class NeedsArg(AbstractCapability):
|
||||
value: int
|
||||
"""
|
||||
|
||||
GETTER_FAILS_CODE = """
|
||||
from dataclasses import dataclass
|
||||
|
||||
from pydantic_ai.capabilities import AbstractCapability
|
||||
|
||||
|
||||
@dataclass
|
||||
class BadGetter(AbstractCapability):
|
||||
def get_toolset(self):
|
||||
raise ValueError('boom')
|
||||
"""
|
||||
|
||||
SYNTAX_ERROR_CODE = 'def (:\n'
|
||||
|
||||
|
||||
def _marker_code(class_name: str, instruction: str) -> str:
|
||||
"""A valid capability whose class name and static instructions are controllable."""
|
||||
return f"""
|
||||
from dataclasses import dataclass
|
||||
|
||||
from pydantic_ai.capabilities import AbstractCapability
|
||||
|
||||
|
||||
@dataclass
|
||||
class {class_name}(AbstractCapability):
|
||||
def get_instructions(self):
|
||||
return {instruction!r}
|
||||
"""
|
||||
|
||||
|
||||
BAD_INSTRUCTIONS_RETURN_CODE = """
|
||||
from dataclasses import dataclass
|
||||
|
||||
from pydantic_ai.capabilities import AbstractCapability
|
||||
|
||||
|
||||
@dataclass
|
||||
class BadInstructions(AbstractCapability):
|
||||
def get_instructions(self):
|
||||
return 12345
|
||||
"""
|
||||
|
||||
BAD_MODEL_SETTINGS_RETURN_CODE = """
|
||||
from dataclasses import dataclass
|
||||
|
||||
from pydantic_ai.capabilities import AbstractCapability
|
||||
|
||||
|
||||
@dataclass
|
||||
class BadModelSettings(AbstractCapability):
|
||||
def get_model_settings(self):
|
||||
return 'not-a-mapping'
|
||||
"""
|
||||
|
||||
|
||||
def _write(directory: Path, name: str, code: str) -> Path:
|
||||
path = directory / f'{name}.py'
|
||||
path.write_text(code, encoding='utf-8')
|
||||
return path
|
||||
|
||||
|
||||
def _return_none(*args: object, **kwargs: object) -> None:
|
||||
return None
|
||||
|
||||
|
||||
def _return_spec_without_loader(*args: object, **kwargs: object) -> SimpleNamespace:
|
||||
return SimpleNamespace(loader=None)
|
||||
|
||||
|
||||
class TestValidate:
|
||||
def test_valid_returns_class_name(self, tmp_path: Path) -> None:
|
||||
assert validate_capability_file(_write(tmp_path, 'm', VALID_CODE)) == 'MarkerCapability'
|
||||
|
||||
def test_no_subclass(self, tmp_path: Path) -> None:
|
||||
with pytest.raises(CapabilityValidationError, match='no `AbstractCapability` subclass'):
|
||||
validate_capability_file(_write(tmp_path, 'm', NO_SUBCLASS_CODE))
|
||||
|
||||
def test_multiple_subclasses(self, tmp_path: Path) -> None:
|
||||
with pytest.raises(CapabilityValidationError, match='found 2: CapA, CapB'):
|
||||
validate_capability_file(_write(tmp_path, 'm', MULTI_CODE))
|
||||
|
||||
def test_construct_failure_wrapped(self, tmp_path: Path) -> None:
|
||||
with pytest.raises(CapabilityValidationError, match='TypeError'):
|
||||
validate_capability_file(_write(tmp_path, 'm', CONSTRUCT_FAILS_CODE))
|
||||
|
||||
def test_getter_failure_wrapped(self, tmp_path: Path) -> None:
|
||||
with pytest.raises(CapabilityValidationError, match='ValueError: boom'):
|
||||
validate_capability_file(_write(tmp_path, 'm', GETTER_FAILS_CODE))
|
||||
|
||||
def test_bad_instructions_return_rejected(self, tmp_path: Path) -> None:
|
||||
# A non-instructions return is caught at author time, not at the next agent.run.
|
||||
with pytest.raises(CapabilityValidationError):
|
||||
validate_capability_file(_write(tmp_path, 'm', BAD_INSTRUCTIONS_RETURN_CODE))
|
||||
|
||||
def test_bad_model_settings_return_rejected(self, tmp_path: Path) -> None:
|
||||
with pytest.raises(CapabilityValidationError, match='ModelSettings mapping'):
|
||||
validate_capability_file(_write(tmp_path, 'm', BAD_MODEL_SETTINGS_RETURN_CODE))
|
||||
|
||||
def test_syntax_error_wrapped(self, tmp_path: Path) -> None:
|
||||
with pytest.raises(CapabilityValidationError):
|
||||
validate_capability_file(_write(tmp_path, 'm', SYNTAX_ERROR_CODE))
|
||||
|
||||
def test_load_module_no_spec(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
path = _write(tmp_path, 'm', VALID_CODE)
|
||||
monkeypatch.setattr(importlib.util, 'spec_from_file_location', _return_none)
|
||||
with pytest.raises(CapabilityValidationError, match='import spec'):
|
||||
validate_capability_file(path)
|
||||
|
||||
def test_load_module_no_loader(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
path = _write(tmp_path, 'm', VALID_CODE)
|
||||
monkeypatch.setattr(importlib.util, 'spec_from_file_location', _return_spec_without_loader)
|
||||
with pytest.raises(CapabilityValidationError, match='import spec'):
|
||||
validate_capability_file(path)
|
||||
|
||||
def test_load_instance_success(self, tmp_path: Path) -> None:
|
||||
instance = load_capability_instance(_write(tmp_path, 'm', VALID_CODE))
|
||||
assert isinstance(instance, AbstractCapability)
|
||||
|
||||
def test_load_instance_passthrough_validation_error(self, tmp_path: Path) -> None:
|
||||
with pytest.raises(CapabilityValidationError, match='no `AbstractCapability` subclass'):
|
||||
load_capability_instance(_write(tmp_path, 'm', NO_SUBCLASS_CODE))
|
||||
|
||||
def test_load_instance_generic_wrapped(self, tmp_path: Path) -> None:
|
||||
with pytest.raises(CapabilityValidationError, match='TypeError'):
|
||||
load_capability_instance(_write(tmp_path, 'm', CONSTRUCT_FAILS_CODE))
|
||||
|
||||
|
||||
class TestCapabilityStore:
|
||||
def test_write_valid(self, tmp_path: Path) -> None:
|
||||
store = CapabilityStore(tmp_path)
|
||||
record = store.write('marker', VALID_CODE)
|
||||
assert record.last_error is None
|
||||
assert record.class_name == 'MarkerCapability'
|
||||
assert (tmp_path / 'marker.py').exists()
|
||||
assert (tmp_path / 'manifest.json').exists()
|
||||
assert [r.name for r in store.list_all()] == ['marker']
|
||||
|
||||
def test_write_invalid_name_raises(self, tmp_path: Path) -> None:
|
||||
store = CapabilityStore(tmp_path)
|
||||
with pytest.raises(ValueError, match='invalid capability name'):
|
||||
store.write('Bad-Name', VALID_CODE)
|
||||
|
||||
def test_write_validation_failure_records_error(self, tmp_path: Path) -> None:
|
||||
store = CapabilityStore(tmp_path)
|
||||
record = store.write('bad', NO_SUBCLASS_CODE)
|
||||
assert record.last_error is not None
|
||||
assert record.class_name == ''
|
||||
assert (tmp_path / 'bad.py').exists()
|
||||
|
||||
def test_write_upsert_replaces(self, tmp_path: Path) -> None:
|
||||
store = CapabilityStore(tmp_path)
|
||||
store.write('marker', VALID_CODE)
|
||||
store.write('marker', NO_SUBCLASS_CODE)
|
||||
records = store.list_all()
|
||||
assert len(records) == 1
|
||||
assert records[0].last_error is not None
|
||||
|
||||
def test_reauthor_replaces_stale_source(self, tmp_path: Path) -> None:
|
||||
# Re-authoring the same name must serve the new source, not cached bytecode.
|
||||
store = CapabilityStore(tmp_path)
|
||||
store.write('marker', _marker_code('Foo', 'V1'))
|
||||
record = store.write('marker', _marker_code('Bar', 'V2'))
|
||||
assert record.class_name == 'Bar'
|
||||
active = store.load_active()
|
||||
assert len(active) == 1
|
||||
assert active[0].get_instructions() == 'V2'
|
||||
|
||||
def test_write_bad_getter_return_records_error(self, tmp_path: Path) -> None:
|
||||
# A wrong getter return type is rejected at author time, with last_error set.
|
||||
store = CapabilityStore(tmp_path)
|
||||
record = store.write('bad', BAD_INSTRUCTIONS_RETURN_CODE)
|
||||
assert record.last_error is not None
|
||||
assert record.class_name == ''
|
||||
|
||||
def test_write_append_distinct(self, tmp_path: Path) -> None:
|
||||
store = CapabilityStore(tmp_path)
|
||||
store.write('a', VALID_CODE)
|
||||
store.write('b', VALID_CODE)
|
||||
assert [r.name for r in store.list_all()] == ['a', 'b']
|
||||
|
||||
def test_write_upsert_later_entry(self, tmp_path: Path) -> None:
|
||||
store = CapabilityStore(tmp_path)
|
||||
store.write('a', VALID_CODE)
|
||||
store.write('b', VALID_CODE)
|
||||
# Re-authoring the second entry skips past the first before matching.
|
||||
store.write('b', NO_SUBCLASS_CODE)
|
||||
records = store.list_all()
|
||||
assert [r.name for r in records] == ['a', 'b']
|
||||
assert records[1].last_error is not None
|
||||
|
||||
def test_load_manifest_missing_returns_empty(self, tmp_path: Path) -> None:
|
||||
assert CapabilityStore(tmp_path / 'absent').list_all() == []
|
||||
|
||||
def test_load_manifest_corrupt_returns_empty(self, tmp_path: Path) -> None:
|
||||
(tmp_path / 'manifest.json').write_text('not json{', encoding='utf-8')
|
||||
assert CapabilityStore(tmp_path).list_all() == []
|
||||
|
||||
def test_load_active_returns_instances(self, tmp_path: Path) -> None:
|
||||
store = CapabilityStore(tmp_path)
|
||||
store.write('marker', VALID_CODE)
|
||||
active = store.load_active()
|
||||
assert len(active) == 1
|
||||
assert isinstance(active[0], AbstractCapability)
|
||||
|
||||
def test_load_active_skips_disabled(self, tmp_path: Path) -> None:
|
||||
store = CapabilityStore(tmp_path)
|
||||
store.write('marker', VALID_CODE)
|
||||
store.disable('marker')
|
||||
assert store.load_active() == []
|
||||
|
||||
def test_load_active_skips_broken(self, tmp_path: Path) -> None:
|
||||
store = CapabilityStore(tmp_path)
|
||||
store.write('marker', VALID_CODE)
|
||||
# The manifest entry stays active, but the source on disk goes bad.
|
||||
(tmp_path / 'marker.py').write_text(NO_SUBCLASS_CODE, encoding='utf-8')
|
||||
assert store.load_active() == []
|
||||
|
||||
def test_save_manifest_atomic_over_partial_prior_file(self, tmp_path: Path) -> None:
|
||||
# A prior interrupted save left a partial/corrupt manifest; the next save
|
||||
# must replace it atomically, leaving a valid manifest and no temp files.
|
||||
(tmp_path / 'manifest.json').write_text('{"capabilities": [', encoding='utf-8')
|
||||
store = CapabilityStore(tmp_path)
|
||||
store.write('marker', VALID_CODE)
|
||||
assert [r.name for r in store.list_all()] == ['marker']
|
||||
leftovers = [p.name for p in tmp_path.iterdir() if p.name.startswith('manifest.') and p.name.endswith('.tmp')]
|
||||
assert leftovers == []
|
||||
|
||||
def test_save_manifest_cleans_temp_on_failure(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# If the atomic replace fails, the temp file must not be left behind.
|
||||
store = CapabilityStore(tmp_path)
|
||||
store.write('marker', VALID_CODE)
|
||||
|
||||
def _boom(*args: object, **kwargs: object) -> None:
|
||||
raise OSError('replace failed')
|
||||
|
||||
monkeypatch.setattr(os, 'replace', _boom)
|
||||
with pytest.raises(OSError, match='replace failed'):
|
||||
store.write('second', VALID_CODE)
|
||||
leftovers = [p.name for p in tmp_path.iterdir() if p.name.startswith('manifest.') and p.name.endswith('.tmp')]
|
||||
assert leftovers == []
|
||||
|
||||
def test_load_active_persists_new_error(self, tmp_path: Path) -> None:
|
||||
# A capability that validated once but later fails to load gets its error
|
||||
# persisted to the manifest, so list_all reflects the real state.
|
||||
store = CapabilityStore(tmp_path)
|
||||
store.write('marker', VALID_CODE)
|
||||
assert store.list_all()[0].last_error is None
|
||||
(tmp_path / 'marker.py').write_text(NO_SUBCLASS_CODE, encoding='utf-8')
|
||||
assert store.load_active() == []
|
||||
record = store.list_all()[0]
|
||||
assert record.last_error is not None
|
||||
assert 'no `AbstractCapability` subclass' in record.last_error
|
||||
|
||||
def test_load_active_broken_twice_does_not_rewrite(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# A still-broken capability whose error is already persisted must not
|
||||
# trigger another manifest write on the next reload.
|
||||
store = CapabilityStore(tmp_path)
|
||||
store.write('marker', VALID_CODE)
|
||||
(tmp_path / 'marker.py').write_text(NO_SUBCLASS_CODE, encoding='utf-8')
|
||||
store.load_active()
|
||||
assert store.list_all()[0].last_error is not None
|
||||
|
||||
save_spy = Mock()
|
||||
monkeypatch.setattr(CapabilityStore, '_save_manifest', save_spy)
|
||||
assert store.load_active() == []
|
||||
save_spy.assert_not_called()
|
||||
|
||||
def test_load_active_clears_error_on_refix(self, tmp_path: Path) -> None:
|
||||
# Re-fixing a broken capability clears its persisted last_error on reload.
|
||||
store = CapabilityStore(tmp_path)
|
||||
store.write('marker', VALID_CODE)
|
||||
(tmp_path / 'marker.py').write_text(NO_SUBCLASS_CODE, encoding='utf-8')
|
||||
store.load_active()
|
||||
assert store.list_all()[0].last_error is not None
|
||||
(tmp_path / 'marker.py').write_text(VALID_CODE, encoding='utf-8')
|
||||
active = store.load_active()
|
||||
assert len(active) == 1
|
||||
assert store.list_all()[0].last_error is None
|
||||
|
||||
def test_load_active_healthy_does_not_rewrite_manifest(
|
||||
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
# A reload with no error transitions must not touch the manifest on disk.
|
||||
store = CapabilityStore(tmp_path)
|
||||
store.write('marker', VALID_CODE)
|
||||
|
||||
save_spy = Mock()
|
||||
monkeypatch.setattr(CapabilityStore, '_save_manifest', save_spy)
|
||||
active = store.load_active()
|
||||
assert len(active) == 1
|
||||
save_spy.assert_not_called()
|
||||
|
||||
def test_disable_found(self, tmp_path: Path) -> None:
|
||||
store = CapabilityStore(tmp_path)
|
||||
store.write('marker', VALID_CODE)
|
||||
assert store.disable('marker') is True
|
||||
assert store.list_all()[0].status == 'disabled'
|
||||
|
||||
def test_disable_not_found(self, tmp_path: Path) -> None:
|
||||
store = CapabilityStore(tmp_path)
|
||||
store.write('marker', VALID_CODE)
|
||||
# An existing non-matching entry is skipped, and the result is still False.
|
||||
assert store.disable('nope') is False
|
||||
|
||||
|
||||
class TestAuthoringToolset:
|
||||
async def test_author_success_message(self, tmp_path: Path) -> None:
|
||||
toolset = AuthoringToolset(CapabilityStore(tmp_path))
|
||||
result = await toolset.author_capability('marker', VALID_CODE)
|
||||
assert 'authored and validated' in result
|
||||
assert 'MarkerCapability' in result
|
||||
assert 'next agent run' in result
|
||||
|
||||
async def test_author_validation_failure_message(self, tmp_path: Path) -> None:
|
||||
toolset = AuthoringToolset(CapabilityStore(tmp_path))
|
||||
result = await toolset.author_capability('bad', NO_SUBCLASS_CODE)
|
||||
assert 'failed validation' in result
|
||||
|
||||
async def test_author_invalid_name_model_retry(self, tmp_path: Path) -> None:
|
||||
toolset = AuthoringToolset(CapabilityStore(tmp_path))
|
||||
with pytest.raises(ModelRetry, match='invalid capability name'):
|
||||
await toolset.author_capability('Bad', VALID_CODE)
|
||||
|
||||
async def test_list_empty(self, tmp_path: Path) -> None:
|
||||
toolset = AuthoringToolset(CapabilityStore(tmp_path))
|
||||
assert await toolset.list_authored_capabilities() == 'No capabilities authored yet.'
|
||||
|
||||
async def test_list_with_entries(self, tmp_path: Path) -> None:
|
||||
toolset = AuthoringToolset(CapabilityStore(tmp_path))
|
||||
await toolset.author_capability('marker', VALID_CODE)
|
||||
await toolset.author_capability('bad', NO_SUBCLASS_CODE)
|
||||
listing = await toolset.list_authored_capabilities()
|
||||
assert '- marker [active] MarkerCapability' in listing
|
||||
assert '- bad [active] ?' in listing
|
||||
assert 'ERROR:' in listing
|
||||
|
||||
async def test_disable_found_message(self, tmp_path: Path) -> None:
|
||||
toolset = AuthoringToolset(CapabilityStore(tmp_path))
|
||||
await toolset.author_capability('marker', VALID_CODE)
|
||||
result = await toolset.disable_authored_capability('marker')
|
||||
assert 'disabled' in result
|
||||
|
||||
async def test_disable_not_found_message(self, tmp_path: Path) -> None:
|
||||
toolset = AuthoringToolset(CapabilityStore(tmp_path))
|
||||
result = await toolset.disable_authored_capability('nope')
|
||||
assert 'No authored capability' in result
|
||||
|
||||
|
||||
class TestRuntimeAuthoringCapability:
|
||||
def test_get_instructions_default(self, tmp_path: Path) -> None:
|
||||
instructions = RuntimeAuthoring[None](directory=tmp_path).get_instructions()
|
||||
assert isinstance(instructions, str)
|
||||
assert 'author_capability' in instructions
|
||||
|
||||
def test_get_instructions_custom(self, tmp_path: Path) -> None:
|
||||
assert RuntimeAuthoring[None](directory=tmp_path, guidance='X').get_instructions() == 'X'
|
||||
|
||||
def test_get_instructions_empty_omitted(self, tmp_path: Path) -> None:
|
||||
assert RuntimeAuthoring[None](directory=tmp_path, guidance='').get_instructions() is None
|
||||
|
||||
def test_get_toolset_type(self, tmp_path: Path) -> None:
|
||||
assert isinstance(RuntimeAuthoring[None](directory=tmp_path).get_toolset(), AuthoringToolset)
|
||||
|
||||
def test_serialization_name_none(self) -> None:
|
||||
assert RuntimeAuthoring.get_serialization_name() is None
|
||||
|
||||
def test_store_property(self, tmp_path: Path) -> None:
|
||||
store = RuntimeAuthoring[None](directory=tmp_path).store
|
||||
assert isinstance(store, CapabilityStore)
|
||||
assert store.directory == tmp_path
|
||||
|
||||
|
||||
class TestEndToEnd:
|
||||
async def test_authored_capability_injected_and_runs(self, tmp_path: Path) -> None:
|
||||
store = CapabilityStore(tmp_path)
|
||||
store.write('marker', VALID_CODE)
|
||||
agent = Agent(TestModel(), capabilities=store.load_active())
|
||||
result = await agent.run('go')
|
||||
returns = [
|
||||
part.content
|
||||
for message in result.all_messages()
|
||||
for part in message.parts
|
||||
if isinstance(part, ToolReturnPart)
|
||||
]
|
||||
assert 'marker-result' in returns
|
||||
|
||||
async def test_runtime_authoring_tools_wired(self, tmp_path: Path) -> None:
|
||||
agent = Agent(TestModel(), capabilities=[RuntimeAuthoring(directory=tmp_path)])
|
||||
result = await agent.run('go')
|
||||
assert result.output is not None
|
||||
@@ -0,0 +1,334 @@
|
||||
"""Tests for the RepoContext capability."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from pydantic_ai import Agent, RunContext
|
||||
from pydantic_ai.messages import ToolCallPart
|
||||
from pydantic_ai.models.test import TestModel
|
||||
from pydantic_ai.tools import ToolDefinition
|
||||
from pydantic_ai.usage import RunUsage
|
||||
|
||||
from pydantic_ai_harness.experimental.context import (
|
||||
AgentContextInventory,
|
||||
ContextFile,
|
||||
RepoContext,
|
||||
RepoContextToolset,
|
||||
)
|
||||
from pydantic_ai_harness.experimental.context._inventory import scan_assets
|
||||
from pydantic_ai_harness.experimental.context._loader import (
|
||||
discover_instruction_files,
|
||||
find_dir_context_file,
|
||||
render_context_files,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def anyio_backend() -> str:
|
||||
return 'asyncio'
|
||||
|
||||
|
||||
def _run_context() -> RunContext[None]:
|
||||
return RunContext[None](
|
||||
deps=None,
|
||||
model=TestModel(),
|
||||
usage=RunUsage(),
|
||||
prompt=None,
|
||||
messages=[],
|
||||
run_step=0,
|
||||
)
|
||||
|
||||
|
||||
def _call(tool_name: str, **args: str) -> tuple[ToolCallPart, ToolDefinition, dict[str, str]]:
|
||||
return ToolCallPart(tool_name=tool_name, args=args), ToolDefinition(name=tool_name), args
|
||||
|
||||
|
||||
def _write(path: Path, content: str) -> Path:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(content, encoding='utf-8')
|
||||
return path
|
||||
|
||||
|
||||
class TestDiscoverInstructionFiles:
|
||||
def test_walk_up_ancestor_first(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / 'CLAUDE.md', 'root')
|
||||
workspace = tmp_path / 'a' / 'b'
|
||||
_write(workspace / 'CLAUDE.md', 'leaf')
|
||||
files = discover_instruction_files(workspace, tmp_path, ('CLAUDE.md',))
|
||||
assert [f.content for f in files] == ['root', 'leaf']
|
||||
|
||||
def test_home_none_only_workspace(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / 'CLAUDE.md', 'root')
|
||||
workspace = tmp_path / 'a'
|
||||
_write(workspace / 'CLAUDE.md', 'leaf')
|
||||
files = discover_instruction_files(workspace, None, ('CLAUDE.md',))
|
||||
assert [f.content for f in files] == ['leaf']
|
||||
|
||||
def test_home_equals_workspace(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / 'CLAUDE.md', 'only')
|
||||
files = discover_instruction_files(tmp_path, tmp_path, ('CLAUDE.md',))
|
||||
assert [f.content for f in files] == ['only']
|
||||
|
||||
def test_home_not_ancestor_falls_back_to_workspace(self, tmp_path: Path) -> None:
|
||||
workspace = tmp_path / 'a'
|
||||
_write(workspace / 'CLAUDE.md', 'leaf')
|
||||
unrelated = tmp_path / 'other'
|
||||
unrelated.mkdir()
|
||||
files = discover_instruction_files(workspace, unrelated, ('CLAUDE.md',))
|
||||
assert [f.content for f in files] == ['leaf']
|
||||
|
||||
def test_both_filenames_within_dir_order(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / 'CLAUDE.md', 'claude')
|
||||
_write(tmp_path / 'AGENTS.md', 'agents')
|
||||
files = discover_instruction_files(tmp_path, None, ('CLAUDE.md', 'AGENTS.md'))
|
||||
assert [f.content for f in files] == ['claude', 'agents']
|
||||
|
||||
@pytest.mark.skipif(sys.platform == 'win32', reason='symlinks need privileges on Windows')
|
||||
def test_symlink_deduped_by_realpath(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / 'CLAUDE.md', 'shared')
|
||||
(tmp_path / 'AGENTS.md').symlink_to(tmp_path / 'CLAUDE.md')
|
||||
files = discover_instruction_files(tmp_path, None, ('CLAUDE.md', 'AGENTS.md'))
|
||||
assert len(files) == 1
|
||||
|
||||
def test_identical_content_deduped_by_hash(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / 'CLAUDE.md', 'same')
|
||||
workspace = tmp_path / 'a'
|
||||
_write(workspace / 'CLAUDE.md', 'same')
|
||||
files = discover_instruction_files(workspace, tmp_path, ('CLAUDE.md',))
|
||||
assert [f.content for f in files] == ['same']
|
||||
|
||||
def test_missing_files_skipped(self, tmp_path: Path) -> None:
|
||||
files = discover_instruction_files(tmp_path, None, ('CLAUDE.md',))
|
||||
assert files == []
|
||||
|
||||
def test_non_utf8_file_does_not_crash(self, tmp_path: Path) -> None:
|
||||
(tmp_path / 'CLAUDE.md').write_bytes(b'caf\xe9 instructions')
|
||||
files = discover_instruction_files(tmp_path, None, ('CLAUDE.md',))
|
||||
assert len(files) == 1
|
||||
assert 'instructions' in files[0].content
|
||||
|
||||
|
||||
class TestFindDirContextFile:
|
||||
def test_first_existing_wins(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / 'AGENTS.md', 'agents')
|
||||
found = find_dir_context_file(tmp_path, ('CLAUDE.md', 'AGENTS.md'))
|
||||
assert found is not None
|
||||
assert found.content == 'agents'
|
||||
|
||||
def test_none_when_absent(self, tmp_path: Path) -> None:
|
||||
assert find_dir_context_file(tmp_path, ('CLAUDE.md',)) is None
|
||||
|
||||
def test_non_utf8_file_does_not_crash(self, tmp_path: Path) -> None:
|
||||
(tmp_path / 'CLAUDE.md').write_bytes(b'caf\xe9 instructions')
|
||||
found = find_dir_context_file(tmp_path, ('CLAUDE.md',))
|
||||
assert found is not None
|
||||
assert 'instructions' in found.content
|
||||
|
||||
|
||||
class TestRender:
|
||||
def test_label_outside_workspace_falls_back_to_posix(self, tmp_path: Path) -> None:
|
||||
outside = _write(tmp_path / 'outer' / 'CLAUDE.md', 'x')
|
||||
cf = ContextFile(directory=outside.parent, path=outside, content='x')
|
||||
rendered = render_context_files([cf], relative_to=tmp_path / 'inner')
|
||||
assert outside.as_posix() in rendered
|
||||
|
||||
|
||||
class TestInstructions:
|
||||
def test_includes_files_and_inventory_hint(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / 'CLAUDE.md', 'be nice')
|
||||
cap = RepoContext[None](workspace_dir=tmp_path)
|
||||
instructions = cap.get_instructions()
|
||||
assert isinstance(instructions, str)
|
||||
assert 'be nice' in instructions
|
||||
assert 'inventory_agent_context' in instructions
|
||||
|
||||
def test_none_when_all_disabled(self, tmp_path: Path) -> None:
|
||||
cap = RepoContext[None](workspace_dir=tmp_path, autoload_instructions=False, expose_inventory_tool=False)
|
||||
assert cap.get_instructions() is None
|
||||
|
||||
def test_autoload_off_keeps_inventory_hint(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / 'CLAUDE.md', 'ignored')
|
||||
cap = RepoContext[None](workspace_dir=tmp_path, autoload_instructions=False)
|
||||
instructions = cap.get_instructions()
|
||||
assert isinstance(instructions, str)
|
||||
assert 'ignored' not in instructions
|
||||
assert 'inventory_agent_context' in instructions
|
||||
|
||||
def test_no_files_no_inventory_is_none(self, tmp_path: Path) -> None:
|
||||
cap = RepoContext[None](workspace_dir=tmp_path, expose_inventory_tool=False)
|
||||
assert cap.get_instructions() is None
|
||||
|
||||
def test_files_cached_across_calls(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / 'CLAUDE.md', 'first')
|
||||
cap = RepoContext[None](workspace_dir=tmp_path)
|
||||
assert cap.get_instructions() is not None and 'first' in cap.get_instructions() # type: ignore[operator]
|
||||
_write(tmp_path / 'CLAUDE.md', 'second')
|
||||
# Read-once: the cached result is reused, so the edit is not picked up.
|
||||
assert 'second' not in cap.get_instructions() # type: ignore[operator]
|
||||
|
||||
|
||||
class TestToolset:
|
||||
def test_get_toolset_none_when_disabled(self, tmp_path: Path) -> None:
|
||||
assert RepoContext[None](workspace_dir=tmp_path, expose_inventory_tool=False).get_toolset() is None
|
||||
|
||||
def test_get_toolset_present(self, tmp_path: Path) -> None:
|
||||
assert isinstance(RepoContext[None](workspace_dir=tmp_path).get_toolset(), RepoContextToolset)
|
||||
|
||||
async def test_inventory_tool_runs_through_agent(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / '.claude' / 'skills' / 'foo' / 'SKILL.md', 'skill')
|
||||
agent = Agent(
|
||||
TestModel(call_tools=['inventory_agent_context']), capabilities=[RepoContext[None](workspace_dir=tmp_path)]
|
||||
)
|
||||
result = await agent.run('go')
|
||||
assert 'inventory_agent_context' in result.output
|
||||
|
||||
|
||||
class TestScanAssets:
|
||||
def test_full_shape(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / '.claude' / 'skills' / 'foo' / 'SKILL.md', 's')
|
||||
_write(tmp_path / '.claude' / 'agents' / 'bar.md', 'a')
|
||||
_write(tmp_path / '.claude' / 'settings.json', '{}')
|
||||
inv = scan_assets(tmp_path, ('.claude', '.agents', '.codex', '.grok'))
|
||||
by_root = {r.root: r for r in inv.roots}
|
||||
claude = by_root['.claude']
|
||||
assert claude.exists
|
||||
assert claude.skills == ['.claude/skills/foo/SKILL.md']
|
||||
assert claude.agents == ['.claude/agents/bar.md']
|
||||
assert claude.settings == '.claude/settings.json'
|
||||
assert by_root['.agents'].exists is False
|
||||
assert by_root['.codex'].notes is not None
|
||||
assert by_root['.grok'].notes is not None
|
||||
|
||||
def test_existing_root_without_settings(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / '.claude' / 'skills' / 'foo' / 'SKILL.md', 's')
|
||||
inv = scan_assets(tmp_path, ('.claude',))
|
||||
assert inv.roots[0].settings is None
|
||||
assert inv.roots[0].notes is None
|
||||
|
||||
def test_returns_model(self, tmp_path: Path) -> None:
|
||||
assert isinstance(scan_assets(tmp_path, ()), AgentContextInventory)
|
||||
|
||||
@pytest.mark.skipif(sys.platform == 'win32', reason='symlinks need privileges on Windows')
|
||||
def test_symlinked_asset_escaping_workspace(self, tmp_path: Path) -> None:
|
||||
workspace = tmp_path / 'ws'
|
||||
outside = _write(tmp_path / 'outside' / 'foo' / 'SKILL.md', 's')
|
||||
link = workspace / '.claude' / 'skills' / 'foo' / 'SKILL.md'
|
||||
link.parent.mkdir(parents=True)
|
||||
link.symlink_to(outside)
|
||||
inv = scan_assets(workspace, ('.claude',))
|
||||
claude = inv.roots[0]
|
||||
assert claude.exists
|
||||
assert len(claude.skills) == 1
|
||||
assert claude.skills[0].endswith('.claude/skills/foo/SKILL.md')
|
||||
|
||||
|
||||
class TestNestedTraversal:
|
||||
async def test_off_by_default_returns_result(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / 'sub' / 'CLAUDE.md', 'nested')
|
||||
cap = RepoContext[None](workspace_dir=tmp_path)
|
||||
call, tool_def, args = _call('list_directory', path='sub')
|
||||
out = await cap.after_tool_execute(_run_context(), call=call, tool_def=tool_def, args=args, result='listing')
|
||||
assert out == 'listing'
|
||||
|
||||
async def test_pointer_appended_on_first_traversal(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / 'sub' / 'CLAUDE.md', 'nested')
|
||||
cap = RepoContext[None](workspace_dir=tmp_path, nested_traversal=True)
|
||||
call, tool_def, args = _call('list_directory', path='sub')
|
||||
out = await cap.after_tool_execute(_run_context(), call=call, tool_def=tool_def, args=args, result='listing')
|
||||
assert out.startswith('listing')
|
||||
assert 'sub/CLAUDE.md' in out
|
||||
assert 'nested' not in out
|
||||
|
||||
async def test_second_traversal_no_reappend(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / 'sub' / 'CLAUDE.md', 'nested')
|
||||
cap = RepoContext[None](workspace_dir=tmp_path, nested_traversal=True)
|
||||
call, tool_def, args = _call('list_directory', path='sub')
|
||||
ctx = _run_context()
|
||||
first = await cap.after_tool_execute(ctx, call=call, tool_def=tool_def, args=args, result='one')
|
||||
second = await cap.after_tool_execute(ctx, call=call, tool_def=tool_def, args=args, result='two')
|
||||
assert 'CLAUDE.md' in first
|
||||
assert second == 'two'
|
||||
|
||||
async def test_tool_name_not_matched(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / 'sub' / 'CLAUDE.md', 'nested')
|
||||
cap = RepoContext[None](workspace_dir=tmp_path, nested_traversal=True)
|
||||
call, tool_def, args = _call('write_file', path='sub')
|
||||
out = await cap.after_tool_execute(_run_context(), call=call, tool_def=tool_def, args=args, result='r')
|
||||
assert out == 'r'
|
||||
|
||||
async def test_non_str_path_arg_ignored(self, tmp_path: Path) -> None:
|
||||
cap = RepoContext[None](workspace_dir=tmp_path, nested_traversal=True)
|
||||
call = ToolCallPart(tool_name='list_directory', args={'path': 123})
|
||||
out = await cap.after_tool_execute(
|
||||
_run_context(), call=call, tool_def=ToolDefinition(name='list_directory'), args={'path': 123}, result='r'
|
||||
)
|
||||
assert out == 'r'
|
||||
|
||||
async def test_dir_without_context_file_untouched(self, tmp_path: Path) -> None:
|
||||
(tmp_path / 'sub').mkdir()
|
||||
cap = RepoContext[None](workspace_dir=tmp_path, nested_traversal=True)
|
||||
call, tool_def, args = _call('list_directory', path='sub')
|
||||
out = await cap.after_tool_execute(_run_context(), call=call, tool_def=tool_def, args=args, result='r')
|
||||
assert out == 'r'
|
||||
|
||||
async def test_read_file_uses_parent_dir(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / 'sub' / 'CLAUDE.md', 'nested')
|
||||
target = _write(tmp_path / 'sub' / 'code.py', 'x = 1')
|
||||
cap = RepoContext[None](workspace_dir=tmp_path, nested_traversal=True)
|
||||
call, tool_def, args = _call('read_file', path=str(target))
|
||||
out = await cap.after_tool_execute(_run_context(), call=call, tool_def=tool_def, args=args, result='file body')
|
||||
assert 'CLAUDE.md' in out
|
||||
|
||||
async def test_contents_mode_inlines_body(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / 'sub' / 'CLAUDE.md', 'NESTED BODY')
|
||||
cap = RepoContext[None](workspace_dir=tmp_path, nested_traversal=True, nested_inject='contents')
|
||||
call, tool_def, args = _call('list_directory', path='sub')
|
||||
out = await cap.after_tool_execute(_run_context(), call=call, tool_def=tool_def, args=args, result='r')
|
||||
assert 'NESTED BODY' in out
|
||||
|
||||
async def test_label_falls_back_when_dir_outside_workspace(self, tmp_path: Path) -> None:
|
||||
workspace = tmp_path / 'ws'
|
||||
workspace.mkdir()
|
||||
outside = _write(tmp_path / 'outside' / 'CLAUDE.md', 'nested').parent
|
||||
cap = RepoContext[None](workspace_dir=workspace, nested_traversal=True)
|
||||
call, tool_def, args = _call('list_directory', path=str(outside))
|
||||
out = await cap.after_tool_execute(_run_context(), call=call, tool_def=tool_def, args=args, result='r')
|
||||
assert outside.resolve().as_posix() in out
|
||||
|
||||
async def test_non_str_result_returned_unchanged(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / 'sub' / 'CLAUDE.md', 'nested')
|
||||
cap = RepoContext[None](
|
||||
workspace_dir=tmp_path, nested_traversal=True, traversal_tool_names=frozenset({'list_dir', 'read_file'})
|
||||
)
|
||||
call, tool_def, args = _call('list_dir', path='sub')
|
||||
listing = [{'name': 'CLAUDE.md'}, {'name': 'code.py'}]
|
||||
out = await cap.after_tool_execute(_run_context(), call=call, tool_def=tool_def, args=args, result=listing)
|
||||
assert out is listing
|
||||
|
||||
async def test_string_result_still_gets_note(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / 'sub' / 'CLAUDE.md', 'nested')
|
||||
cap = RepoContext[None](workspace_dir=tmp_path, nested_traversal=True)
|
||||
call, tool_def, args = _call('list_directory', path='sub')
|
||||
out = await cap.after_tool_execute(_run_context(), call=call, tool_def=tool_def, args=args, result='listing')
|
||||
assert out.startswith('listing')
|
||||
assert 'CLAUDE.md' in out
|
||||
|
||||
|
||||
class TestForRunAndMisc:
|
||||
async def test_for_run_isolates_state(self, tmp_path: Path) -> None:
|
||||
_write(tmp_path / 'sub' / 'CLAUDE.md', 'nested')
|
||||
base = RepoContext[None](workspace_dir=tmp_path, nested_traversal=True)
|
||||
run_cap = await base.for_run(_run_context())
|
||||
call, tool_def, args = _call('list_directory', path='sub')
|
||||
await run_cap.after_tool_execute(_run_context(), call=call, tool_def=tool_def, args=args, result='r')
|
||||
fresh = await base.for_run(_run_context())
|
||||
out = await fresh.after_tool_execute(_run_context(), call=call, tool_def=tool_def, args=args, result='r2')
|
||||
assert 'CLAUDE.md' in out
|
||||
|
||||
def test_serialization_name(self) -> None:
|
||||
assert RepoContext.get_serialization_name() == 'RepoContext'
|
||||
@@ -0,0 +1,183 @@
|
||||
"""Tests for the PyaiDocs capability."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
from pathlib import Path
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pydantic import TypeAdapter
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart, ToolCallPart, ToolReturnPart
|
||||
from pydantic_ai.models.function import AgentInfo, FunctionModel
|
||||
|
||||
from pydantic_ai_harness.experimental import HarnessExperimentalWarning
|
||||
from pydantic_ai_harness.experimental.docs import PyaiDocs, PyaiDocsToolset, PyaiDocsTopic
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def anyio_backend() -> str:
|
||||
"""Run async tests on the asyncio backend (matching upstream pydantic-ai)."""
|
||||
return 'asyncio'
|
||||
|
||||
|
||||
class _FakeClient:
|
||||
"""Stand-in for `httpx.AsyncClient` that returns a canned response or raises."""
|
||||
|
||||
def __init__(self, *, text: str = '', status: int = 200, error: httpx.HTTPError | None = None) -> None:
|
||||
self._text = text
|
||||
self._status = status
|
||||
self._error = error
|
||||
|
||||
async def __aenter__(self) -> _FakeClient:
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc: object) -> bool:
|
||||
return False
|
||||
|
||||
async def get(self, url: str) -> httpx.Response:
|
||||
if self._error is not None:
|
||||
raise self._error
|
||||
return httpx.Response(self._status, text=self._text, request=httpx.Request('GET', url))
|
||||
|
||||
|
||||
def _install_fake_httpx(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
*,
|
||||
text: str = '',
|
||||
status: int = 200,
|
||||
error: httpx.HTTPError | None = None,
|
||||
) -> None:
|
||||
"""Replace `httpx.AsyncClient` with a factory yielding a `_FakeClient`."""
|
||||
|
||||
def factory(*args: object, **kwargs: object) -> _FakeClient:
|
||||
return _FakeClient(text=text, status=status, error=error)
|
||||
|
||||
monkeypatch.setattr(httpx, 'AsyncClient', factory)
|
||||
|
||||
|
||||
class TestPyaiDocsToolset:
|
||||
async def test_local_hit_is_cached(self, tmp_path: Path) -> None:
|
||||
(tmp_path / 'hooks.md').write_text('# Hooks local', encoding='utf-8')
|
||||
cache: dict[PyaiDocsTopic, str] = {}
|
||||
toolset = PyaiDocsToolset[None](local_docs_path=tmp_path, cache=cache)
|
||||
|
||||
assert await toolset.read_pyai_docs(PyaiDocsTopic.hooks) == '# Hooks local'
|
||||
assert cache[PyaiDocsTopic.hooks] == '# Hooks local'
|
||||
|
||||
# Second call serves from cache: removing the file does not change the result.
|
||||
(tmp_path / 'hooks.md').unlink()
|
||||
assert await toolset.read_pyai_docs(PyaiDocsTopic.hooks) == '# Hooks local'
|
||||
|
||||
async def test_remote_fallback_without_local_and_caching_disabled(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_install_fake_httpx(monkeypatch, text='# Capabilities remote')
|
||||
toolset = PyaiDocsToolset[None](local_docs_path=None, cache=None)
|
||||
|
||||
assert await toolset.read_pyai_docs(PyaiDocsTopic.capabilities) == '# Capabilities remote'
|
||||
|
||||
async def test_remote_fallback_when_local_file_missing(
|
||||
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
_install_fake_httpx(monkeypatch, text='# Agent remote')
|
||||
cache: dict[PyaiDocsTopic, str] = {}
|
||||
toolset = PyaiDocsToolset[None](local_docs_path=tmp_path, cache=cache)
|
||||
|
||||
assert await toolset.read_pyai_docs(PyaiDocsTopic.agent) == '# Agent remote'
|
||||
assert cache[PyaiDocsTopic.agent] == '# Agent remote'
|
||||
|
||||
async def test_remote_error_without_local_checkout(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_install_fake_httpx(monkeypatch, error=httpx.ConnectError('boom'))
|
||||
toolset = PyaiDocsToolset[None](local_docs_path=None, cache=None)
|
||||
|
||||
with pytest.raises(RuntimeError, match='no local checkout configured'):
|
||||
await toolset.read_pyai_docs(PyaiDocsTopic.tools)
|
||||
|
||||
async def test_remote_error_reports_local_path(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
||||
_install_fake_httpx(monkeypatch, status=404)
|
||||
toolset = PyaiDocsToolset[None](local_docs_path=tmp_path, cache=None)
|
||||
|
||||
with pytest.raises(RuntimeError, match=str(tmp_path)):
|
||||
await toolset.read_pyai_docs(PyaiDocsTopic.toolsets)
|
||||
|
||||
def test_tools_advanced_value_coerces_to_member(self) -> None:
|
||||
# The LLM passes the enum VALUE; pydantic coerces it back to the member.
|
||||
# `tools_advanced` is the one topic whose name != value, so lock the mapping.
|
||||
assert PyaiDocsTopic.tools_advanced.value == 'tools-advanced'
|
||||
assert TypeAdapter(PyaiDocsTopic).validate_python('tools-advanced') is PyaiDocsTopic.tools_advanced
|
||||
|
||||
async def test_tools_advanced_reads_hyphenated_file(self, tmp_path: Path) -> None:
|
||||
(tmp_path / 'tools-advanced.md').write_text('# Tools advanced local', encoding='utf-8')
|
||||
toolset = PyaiDocsToolset[None](local_docs_path=tmp_path, cache=None)
|
||||
|
||||
assert await toolset.read_pyai_docs(PyaiDocsTopic.tools_advanced) == '# Tools advanced local'
|
||||
|
||||
async def test_local_path_expands_user(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
||||
monkeypatch.setenv('HOME', str(tmp_path))
|
||||
(tmp_path / 'hooks.md').write_text('# Hooks home', encoding='utf-8')
|
||||
toolset = PyaiDocsToolset[None](local_docs_path=Path('~'), cache=None)
|
||||
|
||||
assert await toolset.read_pyai_docs(PyaiDocsTopic.hooks) == '# Hooks home'
|
||||
|
||||
|
||||
class TestPyaiDocsCapability:
|
||||
def test_resolved_path_prefers_constructor_arg(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
||||
monkeypatch.setenv('PYDANTIC_AI_HARNESS_DOCS_PATH', '/env/ignored')
|
||||
assert PyaiDocs[None](local_docs_path=tmp_path)._resolved_local_path() == tmp_path
|
||||
|
||||
def test_resolved_path_falls_back_to_env(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
||||
monkeypatch.setenv('PYDANTIC_AI_HARNESS_DOCS_PATH', str(tmp_path))
|
||||
assert PyaiDocs[None]()._resolved_local_path() == tmp_path
|
||||
|
||||
def test_resolved_path_none_when_unset(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv('PYDANTIC_AI_HARNESS_DOCS_PATH', raising=False)
|
||||
assert PyaiDocs[None]()._resolved_local_path() is None
|
||||
|
||||
def test_get_toolset_shares_cache_when_enabled(self) -> None:
|
||||
capability = PyaiDocs[None]()
|
||||
toolset = capability.get_toolset()
|
||||
assert isinstance(toolset, PyaiDocsToolset)
|
||||
assert toolset._cache is capability._cache
|
||||
|
||||
def test_get_toolset_disables_cache(self) -> None:
|
||||
toolset = PyaiDocs[None](cache=False).get_toolset()
|
||||
assert isinstance(toolset, PyaiDocsToolset)
|
||||
assert toolset._cache is None
|
||||
|
||||
def test_instructions_mention_the_tool(self) -> None:
|
||||
instructions = PyaiDocs[None]().get_instructions()
|
||||
assert isinstance(instructions, str)
|
||||
assert 'read_pyai_docs' in instructions
|
||||
|
||||
def test_serialization_name(self) -> None:
|
||||
assert PyaiDocs.get_serialization_name() == 'PyaiDocs'
|
||||
|
||||
|
||||
class TestThroughAgent:
|
||||
async def test_tool_returns_local_doc(self, tmp_path: Path) -> None:
|
||||
(tmp_path / 'capabilities.md').write_text('# Capabilities doc', encoding='utf-8')
|
||||
|
||||
def call_then_finish(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
|
||||
if len(messages) == 1:
|
||||
return ModelResponse(parts=[ToolCallPart('read_pyai_docs', {'topic': 'capabilities'})])
|
||||
return ModelResponse(parts=[TextPart('done')])
|
||||
|
||||
agent = Agent(FunctionModel(call_then_finish), capabilities=[PyaiDocs(local_docs_path=tmp_path)])
|
||||
result = await agent.run('go')
|
||||
|
||||
assert result.output == 'done'
|
||||
returns = [
|
||||
part.content
|
||||
for message in result.all_messages()
|
||||
for part in message.parts
|
||||
if isinstance(part, ToolReturnPart) and part.tool_name == 'read_pyai_docs'
|
||||
]
|
||||
assert returns == ['# Capabilities doc']
|
||||
|
||||
|
||||
def test_import_emits_experimental_warning() -> None:
|
||||
module = importlib.import_module('pydantic_ai_harness.experimental.docs')
|
||||
with pytest.warns(HarnessExperimentalWarning, match='docs'):
|
||||
importlib.reload(module)
|
||||
@@ -0,0 +1,26 @@
|
||||
"""Shared fixtures for the sub-agents tests."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def isolate_agent_dirs(tmp_path_factory: pytest.TempPathFactory, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Point the disk-loading convention roots at empty dirs for every test.
|
||||
|
||||
`SubAgents` auto-loads from `./.agents/agents/` and `~/.agents/agents/` (with
|
||||
a `.claude/` fallback) by default. Redirecting cwd and home to fresh empty
|
||||
directories keeps tests that build `SubAgents` from reading the developer's real
|
||||
agent files. The two roots are distinct so the default project + home pair does
|
||||
not resolve to one folder. Tests that exercise loading pass explicit folders, or
|
||||
populate the redirected home root themselves."""
|
||||
home = tmp_path_factory.mktemp('home_root')
|
||||
monkeypatch.chdir(tmp_path_factory.mktemp('project_root'))
|
||||
|
||||
def fake_home(cls: type[Path]) -> Path:
|
||||
return home
|
||||
|
||||
monkeypatch.setattr(Path, 'home', classmethod(fake_home))
|
||||
@@ -0,0 +1,330 @@
|
||||
"""Tests for disk-loaded sub-agents, effort floor, and model inheritance."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import warnings
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart, ToolCallPart, ToolReturnPart
|
||||
from pydantic_ai.models import Model
|
||||
from pydantic_ai.models.function import AgentInfo, FunctionModel
|
||||
from pydantic_ai.models.test import TestModel
|
||||
from pydantic_ai.tools import RunContext
|
||||
from pydantic_ai.toolsets import AgentToolset, FunctionToolset
|
||||
|
||||
from pydantic_ai_harness.experimental.subagents import (
|
||||
MINIMUM_EFFORT_FLOOR,
|
||||
AgentOverride,
|
||||
SubAgent,
|
||||
SubAgents,
|
||||
clamp_effort,
|
||||
)
|
||||
from pydantic_ai_harness.experimental.subagents._disk import (
|
||||
ParsedAgent,
|
||||
parse_agent_markdown,
|
||||
resolve_folders,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def anyio_backend() -> str:
|
||||
return 'asyncio'
|
||||
|
||||
|
||||
def _delegate_then_finish(agent_name: str) -> FunctionModel:
|
||||
"""A parent model that delegates to `agent_name` once, then replies with text."""
|
||||
calls = {'n': 0}
|
||||
|
||||
def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
|
||||
calls['n'] += 1
|
||||
if calls['n'] == 1:
|
||||
return ModelResponse(
|
||||
parts=[ToolCallPart('delegate_task', {'agent_name': agent_name, 'task': 'do it'}, tool_call_id='c1')]
|
||||
)
|
||||
return ModelResponse(parts=[TextPart('all done')])
|
||||
|
||||
return FunctionModel(model_fn)
|
||||
|
||||
|
||||
def _delegate_returns(result: Any) -> list[str]:
|
||||
return [
|
||||
str(part.content)
|
||||
for message in result.all_messages()
|
||||
for part in message.parts
|
||||
if isinstance(part, ToolReturnPart) and part.tool_name == 'delegate_task'
|
||||
]
|
||||
|
||||
|
||||
def _write_agent(folder: Path, filename: str, content: str) -> None:
|
||||
folder.mkdir(parents=True, exist_ok=True)
|
||||
(folder / filename).write_text(content, encoding='utf-8')
|
||||
|
||||
|
||||
class TestClampEffort:
|
||||
def test_none_becomes_floor(self) -> None:
|
||||
assert clamp_effort(None) == MINIMUM_EFFORT_FLOOR
|
||||
|
||||
def test_false_becomes_floor(self) -> None:
|
||||
assert clamp_effort(False) == MINIMUM_EFFORT_FLOOR
|
||||
|
||||
def test_true_unchanged(self) -> None:
|
||||
assert clamp_effort(True) is True
|
||||
|
||||
def test_below_floor_raised(self) -> None:
|
||||
assert clamp_effort('minimal') == 'low'
|
||||
|
||||
def test_at_floor_unchanged(self) -> None:
|
||||
assert clamp_effort('low') == 'low'
|
||||
|
||||
def test_above_floor_unchanged(self) -> None:
|
||||
assert clamp_effort('high') == 'high'
|
||||
|
||||
def test_custom_floor(self) -> None:
|
||||
assert clamp_effort('low', floor='high') == 'high'
|
||||
assert clamp_effort('xhigh', floor='high') == 'xhigh'
|
||||
|
||||
|
||||
class TestParseAgentMarkdown:
|
||||
def test_full_frontmatter(self) -> None:
|
||||
text = '---\nname: researcher\ndescription: Researches topics\ntools: Read, Grep\ncolor: blue\n---\nBody here.'
|
||||
parsed = parse_agent_markdown(text)
|
||||
assert parsed == ParsedAgent(
|
||||
name='researcher', description='Researches topics', tools=('Read', 'Grep'), body='Body here.'
|
||||
)
|
||||
|
||||
def test_no_frontmatter_uses_whole_text_as_body(self) -> None:
|
||||
parsed = parse_agent_markdown('Just a body, no frontmatter.\n')
|
||||
assert parsed == ParsedAgent(None, None, (), 'Just a body, no frontmatter.')
|
||||
|
||||
def test_empty_text(self) -> None:
|
||||
assert parse_agent_markdown('') == ParsedAgent(None, None, (), '')
|
||||
|
||||
def test_unclosed_frontmatter_is_all_body(self) -> None:
|
||||
parsed = parse_agent_markdown('---\nname: x\nstill no close')
|
||||
assert parsed == ParsedAgent(None, None, (), '---\nname: x\nstill no close')
|
||||
|
||||
def test_block_list_tools(self) -> None:
|
||||
# Includes a blank line (skipped) and a `- ` item with an empty value (dropped).
|
||||
text = '---\nname: a\n\ntools:\n - Read\n - Edit\n - \n---\nBody'
|
||||
parsed = parse_agent_markdown(text)
|
||||
assert parsed.tools == ('Read', 'Edit')
|
||||
|
||||
def test_allowed_tools_key(self) -> None:
|
||||
parsed = parse_agent_markdown('---\nname: a\nallowed-tools: Bash(git:*), Read\n---\nB')
|
||||
assert parsed.tools == ('Bash(git:*)', 'Read')
|
||||
|
||||
def test_quoted_scalar_values(self) -> None:
|
||||
parsed = parse_agent_markdown('---\nname: "quoted"\ndescription: \'single\'\n---\nB')
|
||||
assert parsed.name == 'quoted'
|
||||
assert parsed.description == 'single'
|
||||
|
||||
def test_non_scalar_name_falls_back_to_none(self) -> None:
|
||||
# `name:` with an empty value parses as a (empty) list, which is not a str.
|
||||
parsed = parse_agent_markdown('---\nname:\ndescription: d\n---\nB')
|
||||
assert parsed.name is None
|
||||
assert parsed.description == 'd'
|
||||
|
||||
def test_stray_dash_line_without_list_key_is_ignored(self) -> None:
|
||||
# A `- item` line with no preceding list key has no colon, so it resets and is skipped.
|
||||
parsed = parse_agent_markdown('---\n- orphan\nname: a\n---\nB')
|
||||
assert parsed.name == 'a'
|
||||
assert parsed.tools == ()
|
||||
|
||||
def test_no_tools_key(self) -> None:
|
||||
assert parse_agent_markdown('---\nname: a\n---\nB').tools == ()
|
||||
|
||||
|
||||
class TestResolveFolders:
|
||||
def test_str_convention_with_claude_fallback(self, tmp_path: Path) -> None:
|
||||
# Project root has `.agents/`; home root has neither, so it falls back to `.claude/`.
|
||||
project = tmp_path / 'project'
|
||||
project.mkdir()
|
||||
(project / '.agents').mkdir()
|
||||
home = tmp_path / 'home'
|
||||
home.mkdir()
|
||||
folders = resolve_folders('agents', project, home)
|
||||
assert folders == [project / '.agents' / 'agents', home / '.claude' / 'agents']
|
||||
|
||||
def test_str_overrides_leaf_name(self, tmp_path: Path) -> None:
|
||||
(tmp_path / '.agents').mkdir()
|
||||
folders = resolve_folders('reviewers', tmp_path, tmp_path)
|
||||
assert folders == [tmp_path / '.agents' / 'reviewers']
|
||||
|
||||
def test_sequence_used_verbatim(self, tmp_path: Path) -> None:
|
||||
paths = [tmp_path / 'a', tmp_path / 'b']
|
||||
assert resolve_folders(paths, tmp_path, tmp_path) == paths
|
||||
|
||||
def test_cwd_equal_home_dedupes_folder(self, tmp_path: Path) -> None:
|
||||
# When the project root equals the home root, the project and home convention
|
||||
# folders resolve to the same directory and are deduped to a single entry.
|
||||
(tmp_path / '.agents').mkdir()
|
||||
assert resolve_folders('agents', tmp_path, tmp_path) == [tmp_path / '.agents' / 'agents']
|
||||
|
||||
def test_duplicate_paths_in_sequence_deduped(self, tmp_path: Path) -> None:
|
||||
assert resolve_folders([tmp_path / 'a', tmp_path / 'a'], tmp_path, tmp_path) == [tmp_path / 'a']
|
||||
|
||||
|
||||
class TestDiskLoading:
|
||||
def test_auto_loads_from_convention_by_default(self) -> None:
|
||||
# The isolation fixture points the home root at an empty dir; populate its
|
||||
# conventional folder and the default `SubAgents()` picks it up with no config.
|
||||
_write_agent(Path.home() / '.agents' / 'agents', 'planner.md', '---\nname: planner\n---\nPlan.')
|
||||
cap: SubAgents[None] = SubAgents()
|
||||
assert 'planner' in cap._by_name
|
||||
|
||||
def test_cwd_equal_home_loads_once_without_shadow_warning(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# When the project root equals the home root, the project and home convention
|
||||
# folders point at the same dir; deduping them avoids loading each agent twice
|
||||
# and emitting a spurious "shadowed" warning.
|
||||
root = Path.cwd()
|
||||
|
||||
def fake_home(cls: type[Path]) -> Path:
|
||||
return root
|
||||
|
||||
monkeypatch.setattr(Path, 'home', classmethod(fake_home))
|
||||
_write_agent(root / '.agents' / 'agents', 'planner.md', '---\nname: planner\n---\nPlan.')
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter('error')
|
||||
cap: SubAgents[None] = SubAgents()
|
||||
assert 'planner' in cap._by_name
|
||||
|
||||
def test_none_disables_loading(self) -> None:
|
||||
cap: SubAgents[None] = SubAgents(agent_folders=None)
|
||||
assert cap._by_name == {}
|
||||
assert cap.get_toolset() is None
|
||||
|
||||
def test_loads_agents_from_folder(self, tmp_path: Path) -> None:
|
||||
_write_agent(tmp_path, 'researcher.md', '---\nname: researcher\ndescription: Researches\n---\nResearch well.')
|
||||
cap: SubAgents[None] = SubAgents(agent_folders=[tmp_path])
|
||||
assert 'researcher' in cap._by_name
|
||||
agent = cap._by_name['researcher'].agent
|
||||
assert agent.name == 'researcher'
|
||||
assert agent.model is None # inherits the parent model at delegation
|
||||
|
||||
def test_name_falls_back_to_filename_stem(self, tmp_path: Path) -> None:
|
||||
_write_agent(tmp_path, 'planner.md', 'No frontmatter, just a body.')
|
||||
cap: SubAgents[None] = SubAgents(agent_folders=[tmp_path])
|
||||
assert 'planner' in cap._by_name
|
||||
|
||||
def test_missing_folder_is_skipped(self, tmp_path: Path) -> None:
|
||||
cap: SubAgents[None] = SubAgents(agent_folders=[tmp_path / 'does-not-exist'])
|
||||
assert cap._by_name == {}
|
||||
|
||||
def test_undecodable_file_is_skipped_with_warning(self, tmp_path: Path) -> None:
|
||||
# A non-UTF-8 `.md` file must not abort loading: it is skipped with a warning
|
||||
# and every valid definition in the same folder still loads.
|
||||
tmp_path.mkdir(parents=True, exist_ok=True)
|
||||
(tmp_path / 'broken.md').write_bytes(b'---\nname: broken\n---\n\xff\xfe not utf-8')
|
||||
_write_agent(tmp_path, 'valid.md', '---\nname: valid\n---\nWork.')
|
||||
with pytest.warns(UserWarning, match='Skipping unreadable disk sub-agent file'):
|
||||
cap: SubAgents[None] = SubAgents(agent_folders=[tmp_path])
|
||||
assert 'valid' in cap._by_name
|
||||
assert 'broken' not in cap._by_name
|
||||
|
||||
def test_listing_uses_description(self, tmp_path: Path) -> None:
|
||||
_write_agent(tmp_path, 'r.md', '---\nname: r\ndescription: Researches\n---\nB')
|
||||
cap: SubAgents[None] = SubAgents(agent_folders=[tmp_path])
|
||||
instructions = cap.get_instructions()
|
||||
assert isinstance(instructions, str)
|
||||
assert '- r: Researches' in instructions
|
||||
|
||||
|
||||
class TestPrecedence:
|
||||
def test_explicit_shadows_disk(self, tmp_path: Path) -> None:
|
||||
_write_agent(tmp_path, 'worker.md', '---\nname: worker\ndescription: from disk\n---\nB')
|
||||
explicit = Agent(TestModel(), name='worker', description='from code')
|
||||
with pytest.warns(UserWarning, match="Disk sub-agent 'worker' is shadowed"):
|
||||
cap: SubAgents[None] = SubAgents(agents=[SubAgent(explicit)], agent_folders=[tmp_path])
|
||||
assert cap._by_name['worker'].agent is explicit
|
||||
|
||||
def test_earlier_folder_shadows_later(self, tmp_path: Path) -> None:
|
||||
project = tmp_path / 'project'
|
||||
home = tmp_path / 'home'
|
||||
_write_agent(project, 'worker.md', '---\nname: worker\ndescription: project\n---\nB')
|
||||
_write_agent(home, 'worker.md', '---\nname: worker\ndescription: home\n---\nB')
|
||||
with pytest.warns(UserWarning, match="Disk sub-agent 'worker' is shadowed"):
|
||||
cap: SubAgents[None] = SubAgents(agent_folders=[project, home])
|
||||
listing = cap.get_instructions()
|
||||
assert isinstance(listing, str)
|
||||
assert 'project' in listing
|
||||
assert 'home' not in listing
|
||||
|
||||
|
||||
class TestOverrides:
|
||||
def test_model_and_effort_override(self, tmp_path: Path) -> None:
|
||||
_write_agent(tmp_path, 'w.md', '---\nname: w\n---\nB')
|
||||
model = TestModel()
|
||||
cap: SubAgents[None] = SubAgents(
|
||||
agent_folders=[tmp_path],
|
||||
agent_overrides={'w': AgentOverride(model=model, effort='high')},
|
||||
)
|
||||
agent = cap._by_name['w'].agent
|
||||
assert agent.model is model
|
||||
|
||||
def test_effort_floored_without_override(self, tmp_path: Path) -> None:
|
||||
_write_agent(tmp_path, 'w.md', '---\nname: w\n---\nB')
|
||||
cap: SubAgents[None] = SubAgents(agent_folders=[tmp_path])
|
||||
# No override -> effort defaults to the floor on the built agent's settings.
|
||||
agent = cap._by_name['w'].agent
|
||||
assert isinstance(agent, Agent)
|
||||
settings = agent.model_settings
|
||||
assert isinstance(settings, dict)
|
||||
assert settings == {'thinking': MINIMUM_EFFORT_FLOOR}
|
||||
|
||||
|
||||
class TestToolResolver:
|
||||
def test_resolver_attaches_tools(self, tmp_path: Path) -> None:
|
||||
_write_agent(tmp_path, 'w.md', '---\nname: w\ntools: search\n---\nB')
|
||||
|
||||
toolset: FunctionToolset[object] = FunctionToolset()
|
||||
|
||||
def resolver(name: str) -> Sequence[AgentToolset[object]] | None:
|
||||
return [toolset] if name == 'search' else None
|
||||
|
||||
cap: SubAgents[None] = SubAgents(agent_folders=[tmp_path], tool_resolver=resolver)
|
||||
# The resolved toolset is attached to the built agent.
|
||||
assert toolset in cap._by_name['w'].agent.toolsets
|
||||
|
||||
def test_unknown_tool_warns_and_skips(self, tmp_path: Path) -> None:
|
||||
_write_agent(tmp_path, 'w.md', '---\nname: w\ntools: mystery\n---\nB')
|
||||
|
||||
def resolver(name: str) -> Sequence[AgentToolset[object]] | None:
|
||||
return None
|
||||
|
||||
with pytest.warns(UserWarning, match="Unknown tool 'mystery'"):
|
||||
SubAgents(agent_folders=[tmp_path], tool_resolver=resolver)
|
||||
|
||||
def test_no_resolver_ignores_frontmatter_tools(self, tmp_path: Path) -> None:
|
||||
_write_agent(tmp_path, 'w.md', '---\nname: w\ntools: Read, Edit\n---\nB')
|
||||
# Without a resolver, no warning and no own tools -- inheritance is the path.
|
||||
cap: SubAgents[None] = SubAgents(agent_folders=[tmp_path])
|
||||
assert 'w' in cap._by_name
|
||||
|
||||
|
||||
class TestModelInheritance:
|
||||
async def test_disk_agent_inherits_parent_model(self, tmp_path: Path) -> None:
|
||||
_write_agent(tmp_path, 'worker.md', '---\nname: worker\n---\nDo the work.')
|
||||
cap: SubAgents[None] = SubAgents(agent_folders=[tmp_path])
|
||||
disk_agent = cap._by_name['worker'].agent
|
||||
assert isinstance(disk_agent, Agent)
|
||||
|
||||
captured: dict[str, Model] = {}
|
||||
|
||||
@disk_agent.instructions
|
||||
def _capture(ctx: RunContext[object]) -> str: # pyright: ignore[reportUnusedFunction]
|
||||
captured['model'] = ctx.model
|
||||
return ''
|
||||
|
||||
parent_model = _delegate_then_finish('worker')
|
||||
parent: Agent[None, str] = Agent(parent_model, capabilities=[cap])
|
||||
result = await parent.run('go')
|
||||
assert result.output == 'all done'
|
||||
# The model-less disk agent ran on the parent's resolved model.
|
||||
assert captured['model'] is parent_model
|
||||
assert _delegate_returns(result) == ['all done']
|
||||
Reference in New Issue
Block a user