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

feat(subagents): replace name-keyed mappings with a SubAgent list (#298)

Browse Source 
* feat(subagents): replace name-keyed mappings with a SubAgent list

Referencing sub-agents by dict keys across `agents`, `descriptions`, and
`limits` meant the same name was repeated in several places and a typo in one
mapping silently desynced from the others. A single `SubAgent` entry now pairs
an agent with its name/description overrides and per-delegate run controls, so
`SubAgents(agents=[SubAgent(...), ...])` keeps everything for one delegate in
one place.

A delegate's name defaults to the agent's own `name` (overridable via
`SubAgent(name=...)`); two delegates resolving to the same name, or an agent
with no name and no override, is rejected at construction. `SubAgentLimits` is
folded into `SubAgent` and removed.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HkzP9cq8KvuTDLLf7BJZFx

* refactor(subagents): keep toolset constructor kwarg named `agents`

Per review feedback, SubAgentToolset keeps its `agents` constructor
kwarg rather than renaming to `sub_agents`, matching the public
SubAgents(agents=...) API and the internal _by_name source.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Aditya Vardhan <adtyavrdhn@gmail.com>
This commit is contained in:
Marcelo Trylesinski
2026-06-23 17:35:38 +05:30
committed by GitHub
co-authored by Claude Aditya Vardhan
parent 50638595fd
commit 6232e0762b
5 changed files with 178 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pydantic_ai_harness/experimental/subagents/README.md
+26 -15
View File
@@ -6,7 +6,7 @@
> experimental path -- there is no top-level export:
>
> ```python
> from pydantic_ai_harness.experimental.subagents import SubAgents
> from pydantic_ai_harness.experimental.subagents import SubAgent, SubAgents
> ```
>
> Importing any experimental capability emits a `HarnessExperimentalWarning`. Silence **all**
@@ -27,24 +27,26 @@ A single agent that does everything accumulates a large tool set and a long cont
## The solution
`SubAgents` takes a name-to-agent mapping and exposes a single `delegate_task(agent_name, task)` tool. Each delegation runs the chosen sub-agent in its own run -- with its own message history, so it never sees the parent conversation -- and returns its output to the parent. The available sub-agents are listed in the system prompt as a static instruction, so the listing stays in the cached prefix.
`SubAgents` takes a sequence of `SubAgent` entries and exposes a single `delegate_task(agent_name, task)` tool. Each delegation runs the chosen sub-agent in its own run -- with its own message history, so it never sees the parent conversation -- and returns its output to the parent. The available sub-agents are listed in the system prompt as a static instruction, so the listing stays in the cached prefix.
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.subagents import SubAgents
from pydantic_ai_harness.experimental.subagents import SubAgent, SubAgents
researcher = Agent('anthropic:claude-sonnet-4-6', name='researcher', description='Researches a topic and reports findings')
writer = Agent('anthropic:claude-sonnet-4-6', name='writer', description='Turns notes into polished prose')
orchestrator = Agent(
'anthropic:claude-opus-4-7',
capabilities=[SubAgents(agents={'researcher': researcher, 'writer': writer})],
capabilities=[SubAgents(agents=[SubAgent(researcher), SubAgent(writer)])],
)
result = orchestrator.run_sync('Research the history of TLS and write a one-paragraph summary.')
print(result.output)
```
A delegate's name -- how the parent model refers to it, and how it is listed in the prompt -- is the agent's own `name`, or a `SubAgent(name=...)` override. Two delegates resolving to the same name is an error, and an agent with no name and no override is rejected.
## The tool
| Tool | Purpose |
@@ -65,22 +67,21 @@ print(result.output)
## Per-delegate run controls
`limits` maps a sub-agent name to a `SubAgentLimits`, giving one delegate its own budgets without touching the others. A name absent from `limits` runs with the `SubAgents` defaults.
Each `SubAgent` carries its own budgets, so one delegate's controls do not touch the others. A `SubAgent` with no controls set runs with the `SubAgents` defaults.
```python
from pydantic_ai.usage import UsageLimits
from pydantic_ai_harness.experimental.subagents import SubAgentLimits, SubAgents
from pydantic_ai_harness.experimental.subagents import SubAgent, SubAgents
# reproducer and librarian are Agent instances, as in the example above.
orchestrator = Agent(
'anthropic:claude-opus-4-7',
capabilities=[
SubAgents(
agents={'reproducer': reproducer, 'librarian': librarian},
limits={
'reproducer': SubAgentLimits(usage_limits=UsageLimits(request_limit=35), timeout_seconds=600, max_calls=1),
'librarian': SubAgentLimits(usage_limits=UsageLimits(request_limit=18), timeout_seconds=300, max_calls=2),
},
agents=[
SubAgent(reproducer, usage_limits=UsageLimits(request_limit=35), timeout_seconds=600, max_calls=1),
SubAgent(librarian, usage_limits=UsageLimits(request_limit=18), timeout_seconds=300, max_calls=2),
]
)
],
)
@@ -103,20 +104,30 @@ Hard errors propagate to stop the whole run. A `UsageLimitExceeded` from a child
## Discovery
The sub-agents are listed in the system prompt via `get_instructions`, using each agent's `description` (or a per-name `descriptions` override). A sub-agent with no description is listed by name alone.
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.
## Configuration
```python
SubAgents(
agents={}, # Mapping[str, AbstractAgent[AgentDepsT, Any]] -- name -> agent
descriptions=None, # optional per-name description overrides for the prompt listing
agents=(), # Sequence[SubAgent[AgentDepsT]] -- each pairs an agent with its run controls
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
event_stream_handler=None, # forwarded to each sub-agent run to stream its events
tool_name='delegate_task',
limits={}, # Mapping[str, SubAgentLimits] -- per-delegate run controls
)
```
```python
SubAgent(
agent, # AbstractAgent[AgentDepsT, Any] -- the child agent to run
name=None, # delegate name; defaults to the agent's own `name`
description=None, # prompt-listing description; defaults to the agent's own `description`
usage_limits=None, # per-delegation request/token budget (isolated accounting)
timeout_seconds=None, # per-delegation wall-clock budget
max_calls=None, # max delegations to this sub-agent per parent run
on_failure=None, # steering message for soft degradations of this delegate
)
```
pydantic_ai_harness/experimental/subagents/__init__.py
+2 -2
View File
@@ -2,8 +2,8 @@
from pydantic_ai_harness.experimental._warn import warn_experimental
from pydantic_ai_harness.experimental.subagents._capability import SubAgents
from pydantic_ai_harness.experimental.subagents._toolset import SubAgentLimits, SubAgentToolset
from pydantic_ai_harness.experimental.subagents._toolset import SubAgent, SubAgentToolset
warn_experimental('subagents')
__all__ = ['SubAgentLimits', 'SubAgentToolset', 'SubAgents']
__all__ = ['SubAgent', 'SubAgentToolset', 'SubAgents']
pydantic_ai_harness/experimental/subagents/_capability.py
+38 -39
View File
@@ -2,16 +2,16 @@
from __future__ import annotations
from collections.abc import Mapping, Sequence
from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
from pydantic_ai.agent import AbstractAgent, AgentRunResult, EventStreamHandler
from pydantic_ai.agent import AgentRunResult, EventStreamHandler
from pydantic_ai.capabilities import AbstractCapability, AgentCapability, WrapRunHandler
from pydantic_ai.tools import AgentDepsT, RunContext
from pydantic_ai.toolsets import AgentToolset
from pydantic_ai_harness.experimental.subagents._toolset import SubAgentLimits, SubAgentToolset
from pydantic_ai_harness.experimental.subagents._toolset import SubAgent, SubAgentToolset
if TYPE_CHECKING:
from pydantic_ai._instructions import AgentInstructions
@@ -26,6 +26,13 @@ class SubAgents(AbstractCapability[AgentDepsT]):
conversation), and the available sub-agents are listed in the system prompt
as a static, cache-stable instruction.
Sub-agents are passed as a sequence of `SubAgent` entries, each pairing an
agent with its per-delegate run controls (a `usage_limits` budget, a
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.
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
shared so usage limits apply across the whole agent tree. Optionally, the
@@ -33,34 +40,23 @@ class SubAgents(AbstractCapability[AgentDepsT]):
applied to every sub-agent run (`shared_capabilities`), and sub-agent events
can be streamed to a handler (`event_stream_handler`).
Per-delegate run controls (a `usage_limits` budget, a wall-clock
`timeout_seconds`, a per-run `max_calls` budget, and an `on_failure` steering
message) are configured via `limits`, keyed by sub-agent name. See
`SubAgentLimits`.
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.subagents import SubAgents
from pydantic_ai_harness.experimental.subagents import SubAgent, SubAgents
researcher = Agent('anthropic:claude-sonnet-4-6', name='researcher', description='Researches topics')
writer = Agent('anthropic:claude-sonnet-4-6', name='writer', description='Writes prose')
orchestrator = Agent(
'anthropic:claude-opus-4-7',
capabilities=[SubAgents(agents={'researcher': researcher, 'writer': writer})],
capabilities=[SubAgents(agents=[SubAgent(researcher), SubAgent(writer)])],
)
```
"""
agents: Mapping[str, AbstractAgent[AgentDepsT, Any]] = field(
default_factory=dict[str, 'AbstractAgent[AgentDepsT, Any]']
)
"""Mapping of sub-agent name to the agent that runs when it's delegated to."""
descriptions: Mapping[str, str] | None = None
"""Optional per-name description overrides for the system-prompt listing.
When a name is absent here, the agent's own `description` is used (if any)."""
agents: Sequence[SubAgent[AgentDepsT]] = ()
"""The sub-agents to expose, each a `SubAgent` pairing an agent with its
per-delegate run controls. See `SubAgent`."""
forward_usage: bool = True
"""If `True`, the parent run's `usage` is shared with each sub-agent run, so
@@ -83,26 +79,31 @@ class SubAgents(AbstractCapability[AgentDepsT]):
tool_name: str = 'delegate_task'
"""Name of the delegate tool exposed to the model."""
limits: Mapping[str, SubAgentLimits] = field(default_factory=dict[str, SubAgentLimits])
"""Per-delegate run controls, keyed by sub-agent name. A name absent here runs
with no per-delegate controls (the `SubAgents` defaults). See `SubAgentLimits`
for `usage_limits`, `timeout_seconds`, `max_calls`, and `on_failure`."""
_by_name: dict[str, SubAgent[AgentDepsT]] = field(
default_factory=dict[str, 'SubAgent[AgentDepsT]'], init=False, repr=False, compare=False
)
"""Sub-agents keyed by resolved name, built in `__post_init__` and passed to
the toolset. Insertion order matches `agents` for a stable prompt listing."""
_call_counts: dict[str, dict[str, int]] = field(
default_factory=dict[str, 'dict[str, int]'], init=False, repr=False, compare=False
)
"""Run-scoped delegation counts (run_id -> name -> count), shared with the
toolset and cleared per run in `wrap_run`. Backs `SubAgentLimits.max_calls`."""
toolset and cleared per run in `wrap_run`. Backs `SubAgent.max_calls`."""
def __post_init__(self) -> None:
# A limits entry for an unknown name would silently run unbudgeted, defeating
# the point of configuring a budget; fail fast on the likely typo instead.
unknown = set(self.limits) - set(self.agents)
if unknown:
raise ValueError(
f'`limits` names sub-agents not in `agents`: {", ".join(sorted(unknown))}. '
f'Available sub-agents: {", ".join(sorted(self.agents)) or "(none)"}.'
)
by_name: dict[str, SubAgent[AgentDepsT]] = {}
for sub_agent in self.agents:
name = sub_agent.resolved_name
if name is None:
raise ValueError('Sub-agent has no name: give its `Agent` a `name`, or set `SubAgent(name=...)`.')
if name in by_name:
raise ValueError(
f'Duplicate sub-agent name {name!r}. Each sub-agent needs a distinct name; '
f'set `SubAgent(name=...)` to disambiguate.'
)
by_name[name] = sub_agent
self._by_name = by_name
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."""
@@ -113,12 +114,11 @@ class SubAgents(AbstractCapability[AgentDepsT]):
def get_instructions(self) -> AgentInstructions[AgentDepsT] | None:
"""Static, cache-stable listing of the available sub-agents."""
if not self.agents:
if not self._by_name:
return None
overrides = self.descriptions or {}
lines: list[str] = []
for name, agent in self.agents.items():
description = overrides.get(name) or agent.description
for name, sub_agent in self._by_name.items():
description = sub_agent.description or sub_agent.agent.description
lines.append(f'- {name}: {description}' if description else f'- {name}')
listing = '\n'.join(lines)
return (
@@ -129,16 +129,15 @@ class SubAgents(AbstractCapability[AgentDepsT]):
def get_toolset(self) -> AgentToolset[AgentDepsT] | None:
"""Toolset providing the delegate tool, or `None` when no sub-agents are configured."""
if not self.agents:
if not self._by_name:
return None
return SubAgentToolset(
agents=self.agents,
agents=self._by_name,
forward_usage=self.forward_usage,
inherit_tools=self.inherit_tools,
shared_capabilities=self.shared_capabilities,
event_stream_handler=self.event_stream_handler,
tool_name=self.tool_name,
limits=self.limits,
call_counts=self._call_counts,
)
pydantic_ai_harness/experimental/subagents/_toolset.py
+45 -32
View File
@@ -5,7 +5,7 @@ from __future__ import annotations
import asyncio
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Any
from typing import Any, Generic
from pydantic_ai.agent import AbstractAgent, EventStreamHandler
from pydantic_ai.capabilities import AgentCapability
@@ -20,14 +20,29 @@ from pydantic_ai.usage import UsageLimits
@dataclass(frozen=True)
class SubAgentLimits:
"""Per-delegate run controls for one named sub-agent.
class SubAgent(Generic[AgentDepsT]):
"""One delegate: a child agent plus its per-delegate run controls.
Every field is optional; an unset field leaves the corresponding behaviour
at the `SubAgents` default. Pass a mapping of sub-agent name to
`SubAgentLimits` as `SubAgents(limits=...)`.
Pass a sequence of these as `SubAgents(agents=[...])`. The delegate's name --
how the parent model refers to it, and how it is listed in the system prompt --
is `name` when set, otherwise the agent's own `name`. An agent with neither is
rejected by `SubAgents`.
Every control below is optional; an unset field leaves the corresponding
behaviour at the `SubAgents` default.
"""
agent: AbstractAgent[AgentDepsT, Any]
"""The agent that runs when this delegate is invoked."""
name: str | None = None
"""Name the parent model uses to delegate to this agent. Defaults to the
agent's own `name` when unset."""
description: str | None = None
"""Description for the system-prompt listing. Defaults to the agent's own
`description` when unset; a delegate with neither is listed by name alone."""
usage_limits: UsageLimits | None = None
"""Request/token budget for one delegation. When set, the child runs with
its own usage accounting so the budget counts only the child's own requests
@@ -53,6 +68,11 @@ class SubAgentLimits:
failures soft: a child error returns this message as a normal tool result
instead of raising a parent `ModelRetry`."""
@property
def resolved_name(self) -> str | None:
"""The delegate's name: `name` if set, else the agent's own `name`."""
return self.name or self.agent.name
def _is_capability_contributed(toolset: AbstractToolset[AgentDepsT]) -> bool:
"""Whether `toolset`'s tree contains a `CapabilityOwnedToolset`."""
@@ -75,29 +95,27 @@ class SubAgentToolset(FunctionToolset[AgentDepsT]):
`deps` are forwarded; its `usage` is shared when enabled; its tools are
inherited when enabled; any `shared_capabilities` are applied to every
sub-agent run; and sub-agent events are streamed to `event_stream_handler`
when one is set. Per-delegate run controls come from `limits`.
when one is set. Per-delegate run controls come from each `SubAgent`.
"""
def __init__(
self,
*,
agents: Mapping[str, AbstractAgent[AgentDepsT, Any]],
agents: Mapping[str, SubAgent[AgentDepsT]],
forward_usage: bool,
inherit_tools: bool,
shared_capabilities: Sequence[AgentCapability[AgentDepsT]],
event_stream_handler: EventStreamHandler[AgentDepsT] | None,
tool_name: str,
limits: Mapping[str, SubAgentLimits],
call_counts: dict[str, dict[str, int]],
) -> None:
super().__init__()
self._agents: dict[str, AbstractAgent[AgentDepsT, Any]] = dict(agents)
self._agents: dict[str, SubAgent[AgentDepsT]] = dict(agents)
self._forward_usage = forward_usage
self._inherit_tools = inherit_tools
self._shared_capabilities = list(shared_capabilities)
self._event_stream_handler = event_stream_handler
self._tool_name = tool_name
self._limits = dict(limits)
# Run-scoped delegation counts, keyed by run_id then sub-agent name.
# Shared with the capability, which clears each run's entry in wrap_run.
self._call_counts = call_counts
@@ -152,38 +170,33 @@ class SubAgentToolset(FunctionToolset[AgentDepsT]):
listed in the instructions.
task: The complete, self-contained instruction for the sub-agent.
"""
agent = self._agents.get(agent_name)
if agent is None:
sub_agent = self._agents.get(agent_name)
if sub_agent is None:
available = ', '.join(sorted(self._agents))
raise ModelRetry(f'Unknown sub-agent {agent_name!r}. Available sub-agents: {available}.')
limits = self._limits.get(agent_name)
if (
limits is not None
and limits.max_calls is not None
and self._budget_exhausted(ctx, agent_name, limits.max_calls)
):
if sub_agent.max_calls is not None and self._budget_exhausted(ctx, agent_name, sub_agent.max_calls):
return self._steer(
limits,
sub_agent.on_failure,
f'Delegate budget for {agent_name!r} is exhausted for this run '
f'({limits.max_calls} call(s)). Synthesize from existing evidence and '
f'({sub_agent.max_calls} call(s)). Synthesize from existing evidence and '
f'choose the next action; do not delegate to {agent_name!r} again.',
)
toolsets = self._inherited_toolsets(ctx) if self._inherit_tools else None
capabilities = self._shared_capabilities or None
usage_limits: UsageLimits | None
if limits is not None and limits.usage_limits is not None:
if sub_agent.usage_limits is not None:
# Isolated accounting so the per-child budget counts only this child.
own_budget = True
usage = None
usage_limits = limits.usage_limits
usage_limits = sub_agent.usage_limits
else:
own_budget = False
usage = ctx.usage if self._forward_usage else None
usage_limits = None
run = agent.run(
run = sub_agent.agent.run(
task,
deps=ctx.deps,
usage=usage,
@@ -192,34 +205,34 @@ class SubAgentToolset(FunctionToolset[AgentDepsT]):
capabilities=capabilities,
event_stream_handler=self._event_stream_handler,
)
timeout = limits.timeout_seconds if limits is not None else None
timeout = sub_agent.timeout_seconds
try:
result = await (asyncio.wait_for(run, timeout) if timeout is not None else run)
except asyncio.TimeoutError:
return self._steer(
limits,
sub_agent.on_failure,
f'Sub-agent {agent_name!r} exceeded its {timeout}s time budget. '
f'Treat this as a recoverable observation and decide from existing evidence.',
)
except UsageLimitExceeded:
if own_budget:
return self._steer(
limits,
sub_agent.on_failure,
f'Sub-agent {agent_name!r} reached its usage budget. '
f'Treat this as a recoverable observation and decide from existing evidence.',
)
# A shared/parent usage limit means the whole tree is out of budget.
raise
except (ModelRetry, UnexpectedModelBehavior) as exc:
if limits is not None and limits.on_failure is not None:
return limits.on_failure
if sub_agent.on_failure is not None:
return sub_agent.on_failure
# Soft sub-agent failures come back to the parent as a retry it can react to.
raise ModelRetry(f'Sub-agent {agent_name!r} failed: {exc}') from exc
return str(result.output)
@staticmethod
def _steer(limits: SubAgentLimits | None, default: str) -> str:
def _steer(on_failure: str | None, default: str) -> str:
"""A soft steering message: the delegate's `on_failure` override, else `default`."""
if limits is not None and limits.on_failure is not None:
return limits.on_failure
if on_failure is not None:
return on_failure
return default
tests/experimental/subagents/test_subagents.py
+67 -66
View File
@@ -25,7 +25,7 @@ from pydantic_ai.models.test import TestModel
from pydantic_ai.tools import AgentDepsT, RunContext
from pydantic_ai.usage import UsageLimits
from pydantic_ai_harness.experimental.subagents import SubAgentLimits, SubAgents, SubAgentToolset
from pydantic_ai_harness.experimental.subagents import SubAgent, SubAgents, SubAgentToolset
@dataclass
@@ -121,30 +121,35 @@ class TestConstruction:
class TestInstructions:
def test_lists_agent_with_description(self) -> None:
agent = Agent(TestModel(), name='researcher', description='Researches topics')
instructions = SubAgents(agents={'researcher': agent}).get_instructions()
instructions = SubAgents(agents=[SubAgent(agent)]).get_instructions()
assert isinstance(instructions, str)
assert '- researcher: Researches topics' in instructions
assert 'delegate_task' in instructions
def test_description_override_wins(self) -> None:
agent = Agent(TestModel(), name='researcher', description='original')
instructions = SubAgents(
agents={'researcher': agent}, descriptions={'researcher': 'overridden'}
).get_instructions()
instructions = SubAgents(agents=[SubAgent(agent, description='overridden')]).get_instructions()
assert isinstance(instructions, str)
assert '- researcher: overridden' in instructions
assert 'original' not in instructions
def test_name_only_when_no_description(self) -> None:
agent = Agent(TestModel(), name='plain')
instructions = SubAgents(agents={'plain': agent}).get_instructions()
instructions = SubAgents(agents=[SubAgent(agent)]).get_instructions()
assert isinstance(instructions, str)
assert '- plain' in instructions
assert '- plain:' not in instructions
def test_name_override_wins(self) -> None:
agent = Agent(TestModel(), name='internal')
instructions = SubAgents(agents=[SubAgent(agent, name='public')]).get_instructions()
assert isinstance(instructions, str)
assert '- public' in instructions
assert 'internal' not in instructions
def test_custom_tool_name_in_instructions(self) -> None:
agent = Agent(TestModel(), name='x')
instructions = SubAgents(agents={'x': agent}, tool_name='run_agent').get_instructions()
instructions = SubAgents(agents=[SubAgent(agent)], tool_name='run_agent').get_instructions()
assert isinstance(instructions, str)
assert 'run_agent' in instructions
@@ -152,13 +157,13 @@ class TestInstructions:
class TestToolset:
def test_get_toolset_exposes_delegate_tool(self) -> None:
agent = Agent(TestModel(), name='x')
toolset = SubAgents(agents={'x': agent}).get_toolset()
toolset = SubAgents(agents=[SubAgent(agent)]).get_toolset()
assert isinstance(toolset, SubAgentToolset)
assert 'delegate_task' in toolset.tools
def test_custom_tool_name(self) -> None:
agent = Agent(TestModel(), name='x')
toolset = SubAgents(agents={'x': agent}, tool_name='run_agent').get_toolset()
toolset = SubAgents(agents=[SubAgent(agent)], tool_name='run_agent').get_toolset()
assert isinstance(toolset, SubAgentToolset)
assert 'run_agent' in toolset.tools
@@ -167,7 +172,7 @@ class TestDelegation:
async def test_delegates_and_returns_output(self) -> None:
worker = Agent(TestModel(custom_output_text='WORKER RESULT'), name='worker')
parent: Agent[None, str] = Agent(
_delegate_then_finish('worker'), capabilities=[SubAgents(agents={'worker': worker})]
_delegate_then_finish('worker'), capabilities=[SubAgents(agents=[SubAgent(worker)])]
)
result = await parent.run('go')
assert result.output == 'all done'
@@ -179,12 +184,22 @@ class TestDelegation:
]
assert returns == ['WORKER RESULT']
async def test_delegates_via_name_override(self) -> None:
worker = Agent(TestModel(custom_output_text='WORKER RESULT'), name='internal')
parent: Agent[None, str] = Agent(
_delegate_then_finish('public'),
capabilities=[SubAgents(agents=[SubAgent(worker, name='public')])],
)
result = await parent.run('go')
assert result.output == 'all done'
assert _delegate_returns(result) == ['WORKER RESULT']
async def test_unknown_agent_triggers_retry_then_succeeds(self) -> None:
worker = Agent(TestModel(custom_output_text='OK'), name='worker')
helper = Agent(TestModel(custom_output_text='OK'), name='helper')
parent: Agent[None, str] = Agent(
_delegate_then_finish('worker', retries_before=1),
capabilities=[SubAgents(agents={'worker': worker, 'helper': helper})],
capabilities=[SubAgents(agents=[SubAgent(worker), SubAgent(helper)])],
)
result = await parent.run('go')
assert result.output == 'all done'
@@ -215,7 +230,7 @@ class TestDelegation:
parent: Agent[str, str] = Agent(
_delegate_then_finish('worker'),
deps_type=str,
capabilities=[SubAgents(agents={'worker': worker})],
capabilities=[SubAgents(agents=[SubAgent(worker)])],
)
@parent.instructions
@@ -241,7 +256,7 @@ class TestDelegation:
worker = Agent(FunctionModel(worker_fn), name='worker')
parent: Agent[None, str] = Agent(
_delegate_then_finish('worker'),
capabilities=[SubAgents(agents={'worker': worker}, inherit_tools=True)],
capabilities=[SubAgents(agents=[SubAgent(worker)], inherit_tools=True)],
)
@parent.tool_plain
@@ -268,13 +283,12 @@ class TestDelegation:
worker = Agent(FunctionModel(worker_fn), name='worker')
toolset: SubAgentToolset[None] = SubAgentToolset(
agents={'worker': worker},
agents={'worker': SubAgent(worker)},
forward_usage=True,
inherit_tools=True,
shared_capabilities=[],
event_stream_handler=None,
tool_name='delegate_task',
limits={},
call_counts={},
)
parent: Agent[None, str] = Agent(_delegate_then_finish('worker'), toolsets=[toolset])
@@ -313,7 +327,7 @@ class TestDelegation:
worker = Agent(FunctionModel(worker_fn), name='worker')
parent: Agent[None, str] = Agent(
_delegate_then_finish('worker'),
capabilities=[SubAgents(agents={'worker': worker}, inherit_tools=True), _ToolCapability()],
capabilities=[SubAgents(agents=[SubAgent(worker)], inherit_tools=True), _ToolCapability()],
)
@parent.tool_plain
@@ -330,7 +344,7 @@ class TestDelegation:
worker = Agent(TestModel(custom_output_text='W'), name='worker')
parent: Agent[None, str] = Agent(
_delegate_then_finish('worker'),
capabilities=[SubAgents(agents={'worker': worker}, shared_capabilities=[cap])],
capabilities=[SubAgents(agents=[SubAgent(worker)], shared_capabilities=[cap])],
)
result = await parent.run('go')
assert result.output == 'all done'
@@ -346,7 +360,7 @@ class TestDelegation:
worker = Agent(TestModel(custom_output_text='W'), name='worker')
parent: Agent[None, str] = Agent(
_delegate_then_finish('worker'),
capabilities=[SubAgents(agents={'worker': worker}, event_stream_handler=handler)],
capabilities=[SubAgents(agents=[SubAgent(worker)], event_stream_handler=handler)],
)
result = await parent.run('go')
assert result.output == 'all done'
@@ -359,7 +373,7 @@ class TestDelegation:
limited = Agent(FunctionModel(boom), name='limited')
parent: Agent[None, str] = Agent(
_delegate_then_finish('limited'),
capabilities=[SubAgents(agents={'limited': limited})],
capabilities=[SubAgents(agents=[SubAgent(limited)])],
)
# Hard limits are not converted to a retry -- they propagate to stop the run.
with pytest.raises(UsageLimitExceeded):
@@ -388,7 +402,7 @@ class TestDelegation:
parent: Agent[None, str] = Agent(
FunctionModel(parent_fn),
capabilities=[SubAgents(agents={'boomer': boomer, 'worker': worker})],
capabilities=[SubAgents(agents=[SubAgent(boomer), SubAgent(worker)])],
)
result = await parent.run('go')
assert result.output == 'all done'
@@ -415,7 +429,7 @@ class TestDelegation:
parent: Agent[str, str] = Agent(
_delegate_then_finish('worker'),
deps_type=str,
capabilities=[SubAgents(agents={'worker': worker}, forward_usage=False)],
capabilities=[SubAgents(agents=[SubAgent(worker)], forward_usage=False)],
)
@parent.instructions
@@ -443,12 +457,7 @@ class TestRunControls:
parent: Agent[None, str] = Agent(
_delegate_then_finish('worker'),
capabilities=[
SubAgents(
agents={'worker': worker},
limits={'worker': SubAgentLimits(usage_limits=UsageLimits(request_limit=5))},
)
],
capabilities=[SubAgents(agents=[SubAgent(worker, usage_limits=UsageLimits(request_limit=5))])],
)
@parent.instructions
@@ -478,12 +487,7 @@ class TestRunControls:
parent: Agent[None, str] = Agent(
_delegate_then_finish('worker'),
capabilities=[
SubAgents(
agents={'worker': worker},
limits={'worker': SubAgentLimits(usage_limits=UsageLimits(request_limit=1))},
)
],
capabilities=[SubAgents(agents=[SubAgent(worker, usage_limits=UsageLimits(request_limit=1))])],
)
result = await parent.run('go')
assert result.output == 'all done'
@@ -498,7 +502,7 @@ class TestRunControls:
worker = Agent(TestModel(custom_output_text='W'), name='worker')
parent: Agent[None, str] = Agent(
_delegate_then_finish('worker'),
capabilities=[SubAgents(agents={'worker': worker})],
capabilities=[SubAgents(agents=[SubAgent(worker)])],
)
with pytest.raises(UsageLimitExceeded):
await parent.run('go', usage_limits=UsageLimits(request_limit=1))
@@ -511,12 +515,7 @@ class TestRunControls:
worker = Agent(FunctionModel(slow_fn), name='worker')
parent: Agent[None, str] = Agent(
_delegate_then_finish('worker'),
capabilities=[
SubAgents(
agents={'worker': worker},
limits={'worker': SubAgentLimits(timeout_seconds=0.01)},
)
],
capabilities=[SubAgents(agents=[SubAgent(worker, timeout_seconds=0.01)])],
)
result = await parent.run('go')
assert result.output == 'all done'
@@ -534,7 +533,7 @@ class TestRunControls:
worker = Agent(FunctionModel(worker_fn), name='worker')
parent: Agent[None, str] = Agent(
_delegate_n_then_finish('worker', 2),
capabilities=[SubAgents(agents={'worker': worker}, limits={'worker': SubAgentLimits(max_calls=1)})],
capabilities=[SubAgents(agents=[SubAgent(worker, max_calls=1)])],
)
result = await parent.run('go')
assert result.output == 'all done'
@@ -546,7 +545,7 @@ class TestRunControls:
async def test_call_counts_reset_between_runs(self) -> None:
worker = Agent(TestModel(custom_output_text='W'), name='worker')
capability = SubAgents(agents={'worker': worker}, limits={'worker': SubAgentLimits(max_calls=1)})
capability = SubAgents(agents=[SubAgent(worker, max_calls=1)])
# Two parents share the one capability (and its run-scoped count store); each
# gets a fresh delegate-once model so the second run actually delegates again.
first = await Agent(_delegate_then_finish('worker'), capabilities=[capability]).run('go')
@@ -564,12 +563,7 @@ class TestRunControls:
boomer = Agent(FunctionModel(boom), name='boomer')
parent: Agent[None, str] = Agent(
_delegate_then_finish('boomer'),
capabilities=[
SubAgents(
agents={'boomer': boomer},
limits={'boomer': SubAgentLimits(on_failure='steer: use existing evidence')},
)
],
capabilities=[SubAgents(agents=[SubAgent(boomer, on_failure='steer: use existing evidence')])],
)
result = await parent.run('go')
assert result.output == 'all done'
@@ -587,12 +581,7 @@ class TestRunControls:
worker = Agent(TestModel(custom_output_text='W'), name='worker')
parent: Agent[None, str] = Agent(
_delegate_n_then_finish('worker', 2),
capabilities=[
SubAgents(
agents={'worker': worker},
limits={'worker': SubAgentLimits(max_calls=1, on_failure='custom budget note')},
)
],
capabilities=[SubAgents(agents=[SubAgent(worker, max_calls=1, on_failure='custom budget note')])],
)
result = await parent.run('go')
assert result.output == 'all done'
@@ -600,11 +589,11 @@ class TestRunControls:
assert returns[1] == 'custom budget note'
async def test_limits_without_budget_run_normally(self) -> None:
# A limits entry with only an unrelated field set must not alter the happy path.
# A SubAgent with only an unrelated control set must not alter the happy path.
worker = Agent(TestModel(custom_output_text='W'), name='worker')
parent: Agent[None, str] = Agent(
_delegate_then_finish('worker'),
capabilities=[SubAgents(agents={'worker': worker}, limits={'worker': SubAgentLimits(timeout_seconds=30)})],
capabilities=[SubAgents(agents=[SubAgent(worker, timeout_seconds=30)])],
)
result = await parent.run('go')
assert _delegate_returns(result) == ['W']
@@ -635,7 +624,7 @@ class TestRunControls:
parent: Agent[None, str] = Agent(
FunctionModel(parent_fn),
capabilities=[SubAgents(agents={'worker': worker}, limits={'worker': SubAgentLimits(max_calls=1)})],
capabilities=[SubAgents(agents=[SubAgent(worker, max_calls=1)])],
)
result = await parent.run('go')
assert result.output == 'all done'
@@ -647,13 +636,25 @@ class TestRunControls:
assert any("Delegate budget for 'worker' is exhausted" in r for r in returns)
class TestLimitsValidation:
def test_limits_key_not_in_agents_raises(self) -> None:
worker = Agent(TestModel(), name='worker')
with pytest.raises(ValueError, match='names sub-agents not in `agents`: ghost'):
SubAgents(agents={'worker': worker}, limits={'ghost': SubAgentLimits(max_calls=1)})
class TestNameValidation:
def test_duplicate_name_raises(self) -> None:
first = Agent(TestModel(), name='dup')
second = Agent(TestModel(), name='dup')
with pytest.raises(ValueError, match="Duplicate sub-agent name 'dup'"):
SubAgents(agents=[SubAgent(first), SubAgent(second)])
def test_limits_subset_of_agents_is_accepted(self) -> None:
worker = Agent(TestModel(), name='worker')
capability = SubAgents(agents={'worker': worker}, limits={'worker': SubAgentLimits(max_calls=1)})
assert 'worker' in capability.limits
def test_duplicate_via_name_override_raises(self) -> None:
first = Agent(TestModel(), name='a')
second = Agent(TestModel(), name='b')
with pytest.raises(ValueError, match="Duplicate sub-agent name 'a'"):
SubAgents(agents=[SubAgent(first), SubAgent(second, name='a')])
def test_missing_name_raises(self) -> None:
nameless = Agent(TestModel())
with pytest.raises(ValueError, match='Sub-agent has no name'):
SubAgents(agents=[SubAgent(nameless)])
def test_name_override_satisfies_missing_agent_name(self) -> None:
nameless = Agent(TestModel())
capability = SubAgents(agents=[SubAgent(nameless, name='worker')])
assert 'worker' in capability._by_name