diff --git a/pydantic_ai_harness/experimental/subagents/README.md b/pydantic_ai_harness/experimental/subagents/README.md index 881c420..71a5b97 100644 --- a/pydantic_ai_harness/experimental/subagents/README.md +++ b/pydantic_ai_harness/experimental/subagents/README.md @@ -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 ) ``` diff --git a/pydantic_ai_harness/experimental/subagents/__init__.py b/pydantic_ai_harness/experimental/subagents/__init__.py index 3a98a84..1e564cb 100644 --- a/pydantic_ai_harness/experimental/subagents/__init__.py +++ b/pydantic_ai_harness/experimental/subagents/__init__.py @@ -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'] diff --git a/pydantic_ai_harness/experimental/subagents/_capability.py b/pydantic_ai_harness/experimental/subagents/_capability.py index 7dc619d..8b5fa8c 100644 --- a/pydantic_ai_harness/experimental/subagents/_capability.py +++ b/pydantic_ai_harness/experimental/subagents/_capability.py @@ -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, ) diff --git a/pydantic_ai_harness/experimental/subagents/_toolset.py b/pydantic_ai_harness/experimental/subagents/_toolset.py index 57c082e..c0f8533 100644 --- a/pydantic_ai_harness/experimental/subagents/_toolset.py +++ b/pydantic_ai_harness/experimental/subagents/_toolset.py @@ -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 diff --git a/tests/experimental/subagents/test_subagents.py b/tests/experimental/subagents/test_subagents.py index 3af629c..6809b90 100644 --- a/tests/experimental/subagents/test_subagents.py +++ b/tests/experimental/subagents/test_subagents.py @@ -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