mirror of
https://github.com/pydantic/pydantic-ai-harness.git
synced 2026-07-21 02:45:34 +00:00
feat(subagents): add SubAgents capability for delegating to child agents (#267)
SubAgents exposes a single delegate_task(agent_name, task) tool that runs a named child agent in a fresh, isolated run, with the available agents listed in the system prompt as a cache-stable instruction. - deps are forwarded to every sub-agent; usage is shared by default (forward_usage) so usage limits apply across the whole agent tree - inherit_tools (opt-in) exposes the parent's tools to sub-agents, with the delegate tool filtered out so they can't recurse into delegation - shared_capabilities are applied to every sub-agent run - event_stream_handler is forwarded to every sub-agent run so sub-agent events surface to the caller - soft sub-agent failures (ModelRetry / UnexpectedModelBehavior) convert into a parent ModelRetry it can react to; hard limits propagate Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Bill Easton
@@ -0,0 +1,79 @@
|
||||
# SubAgents
|
||||
|
||||
Let an agent delegate self-contained tasks to named child agents.
|
||||
|
||||
## The problem
|
||||
|
||||
A single agent that does everything accumulates a large tool set and a long context. Splitting the work across specialized sub-agents keeps each context focused, but wiring up delegation by hand means writing a tool per agent, forwarding deps, threading usage limits, and telling the model what it can delegate to.
|
||||
|
||||
## 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.
|
||||
|
||||
```python
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai_harness.experimental.subagents import 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})],
|
||||
)
|
||||
|
||||
result = orchestrator.run_sync('Research the history of TLS and write a one-paragraph summary.')
|
||||
print(result.output)
|
||||
```
|
||||
|
||||
## The tool
|
||||
|
||||
| Tool | Purpose |
|
||||
|---|---|
|
||||
| `delegate_task(agent_name, task)` | Run the named sub-agent on a self-contained task and return its output. |
|
||||
|
||||
- The sub-agent runs with its own message history, so `task` must be self-contained.
|
||||
- An unknown `agent_name` raises `ModelRetry`, so the model can correct itself.
|
||||
- The result returned to the parent is `str(result.output)`.
|
||||
|
||||
## Deps, usage, tools, and capabilities
|
||||
|
||||
- **Deps are forwarded.** The parent run's `deps` are passed to each sub-agent, so sub-agents share the parent's `AgentDepsT` (enforced by the type signature -- every sub-agent is an `AbstractAgent[AgentDepsT, Any]`).
|
||||
- **Usage is shared by default.** The parent's `usage` is passed to each sub-agent run, so token usage aggregates and a parent `usage_limits` applies across the whole agent tree. Set `forward_usage=False` to give each sub-agent run its own accounting.
|
||||
- **Tools can be inherited.** With `inherit_tools=True`, the parent agent's tools are added to each sub-agent run (on top of the sub-agent's own). The delegate tool itself is filtered out, so a sub-agent can't recurse into further delegation. Off by default.
|
||||
- **Capabilities can be shared.** `shared_capabilities` are applied to every sub-agent run -- e.g. give all sub-agents a common guardrail, memory, or planning capability without rebuilding each `Agent`.
|
||||
- **Sub-agent events can be streamed.** Pass an `event_stream_handler` and it's forwarded to each sub-agent run, so the sub-agent's model-streaming and tool events surface to the caller (the handler receives the sub-agent's own `RunContext`).
|
||||
|
||||
## Failure handling
|
||||
|
||||
If a sub-agent run fails with a *soft* model error (`ModelRetry`, `UnexpectedModelBehavior`, e.g. it exhausted its own retries), the failure is converted into a `ModelRetry` for the parent -- so the parent's model sees `Sub-agent '<name>' failed: …` and can react. Hard errors (e.g. `UsageLimitExceeded`) propagate to stop the whole run.
|
||||
|
||||
## 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.
|
||||
|
||||
## Configuration
|
||||
|
||||
```python
|
||||
SubAgents(
|
||||
agents={}, # Mapping[str, AbstractAgent[AgentDepsT, Any]] -- name -> agent
|
||||
descriptions=None, # optional per-name description overrides for the prompt listing
|
||||
forward_usage=True, # share the parent's usage with sub-agent runs
|
||||
inherit_tools=False, # expose the parent's tools to sub-agents (delegate tool 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',
|
||||
)
|
||||
```
|
||||
|
||||
`SubAgents` is not serializable via the agent spec (it holds live `Agent` instances), so `get_serialization_name()` returns `None`.
|
||||
|
||||
## Notes
|
||||
|
||||
- Sub-agents can themselves have `SubAgents`, forming a tree. Share `usage` (the default) and set a `usage_limits` on the top-level run to bound the whole tree.
|
||||
- Delegations the model issues in parallel run as independent sub-agent runs.
|
||||
|
||||
## Further reading
|
||||
|
||||
- [Pydantic AI capabilities](https://ai.pydantic.dev/capabilities/)
|
||||
- [Multi-agent applications](https://ai.pydantic.dev/multi-agent-applications/)
|
||||
@@ -0,0 +1,9 @@
|
||||
"""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._toolset import SubAgentToolset
|
||||
|
||||
warn_experimental('subagents')
|
||||
|
||||
__all__ = ['SubAgentToolset', 'SubAgents']
|
||||
@@ -0,0 +1,113 @@
|
||||
"""Sub-agent capability: delegate self-contained tasks to named child agents."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from pydantic_ai.agent import AbstractAgent, EventStreamHandler
|
||||
from pydantic_ai.capabilities import AbstractCapability, AgentCapability
|
||||
from pydantic_ai.tools import AgentDepsT
|
||||
from pydantic_ai.toolsets import AgentToolset
|
||||
|
||||
from pydantic_ai_harness.experimental.subagents._toolset import SubAgentToolset
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pydantic_ai._instructions import AgentInstructions
|
||||
|
||||
|
||||
@dataclass
|
||||
class SubAgents(AbstractCapability[AgentDepsT]):
|
||||
"""Let an agent delegate self-contained tasks to named sub-agents.
|
||||
|
||||
Exposes a single `delegate_task(agent_name, task)` tool. Each delegation
|
||||
runs the chosen sub-agent in a fresh, isolated run (it never sees the parent
|
||||
conversation), and the available sub-agents are listed in the system prompt
|
||||
as a static, cache-stable instruction.
|
||||
|
||||
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
|
||||
parent's tools can be inherited (`inherit_tools`), extra capabilities can be
|
||||
applied to every sub-agent run (`shared_capabilities`), and sub-agent events
|
||||
can be streamed to a handler (`event_stream_handler`).
|
||||
|
||||
```python
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai_harness.experimental.subagents import 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})],
|
||||
)
|
||||
```
|
||||
"""
|
||||
|
||||
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)."""
|
||||
|
||||
forward_usage: bool = True
|
||||
"""If `True`, the parent run's `usage` is shared with each sub-agent run, so
|
||||
token usage aggregates and usage limits apply across the whole agent tree."""
|
||||
|
||||
inherit_tools: bool = False
|
||||
"""If `True`, the parent agent's tools are exposed to each sub-agent run (the
|
||||
delegate tool itself is filtered out, so sub-agents can't recurse into
|
||||
further delegation). Off by default to avoid silently widening sub-agent access."""
|
||||
|
||||
shared_capabilities: Sequence[AgentCapability[AgentDepsT]] = ()
|
||||
"""Capabilities applied to every sub-agent run, in addition to whatever each
|
||||
sub-agent already has."""
|
||||
|
||||
event_stream_handler: EventStreamHandler[AgentDepsT] | None = None
|
||||
"""If set, this handler is passed to each sub-agent run, so the sub-agent's
|
||||
model-streaming and tool events surface to the caller. The handler receives
|
||||
the sub-agent's own `RunContext` and event stream."""
|
||||
|
||||
tool_name: str = 'delegate_task'
|
||||
"""Name of the delegate tool exposed to the model."""
|
||||
|
||||
def get_instructions(self) -> AgentInstructions[AgentDepsT] | None:
|
||||
"""Static, cache-stable listing of the available sub-agents."""
|
||||
if not self.agents:
|
||||
return None
|
||||
overrides = self.descriptions or {}
|
||||
lines: list[str] = []
|
||||
for name, agent in self.agents.items():
|
||||
description = overrides.get(name) or agent.description
|
||||
lines.append(f'- {name}: {description}' if description else f'- {name}')
|
||||
listing = '\n'.join(lines)
|
||||
return (
|
||||
f'You can delegate self-contained tasks to these sub-agents using the `{self.tool_name}` '
|
||||
f'tool. Each runs in its own fresh context and does not see this conversation, so pass '
|
||||
f'everything it needs.\n\nAvailable sub-agents:\n{listing}'
|
||||
)
|
||||
|
||||
def get_toolset(self) -> AgentToolset[AgentDepsT] | None:
|
||||
"""Toolset providing the delegate tool, or `None` when no sub-agents are configured."""
|
||||
if not self.agents:
|
||||
return None
|
||||
return SubAgentToolset(
|
||||
agents=self.agents,
|
||||
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,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_serialization_name(cls) -> str | None:
|
||||
"""Not spec-serializable -- the capability holds live `Agent` instances."""
|
||||
return None
|
||||
@@ -0,0 +1,84 @@
|
||||
"""Sub-agent toolset: a single delegate tool that runs named child agents."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any
|
||||
|
||||
from pydantic_ai.agent import AbstractAgent, EventStreamHandler
|
||||
from pydantic_ai.capabilities import AgentCapability
|
||||
from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior
|
||||
from pydantic_ai.tools import AgentDepsT, RunContext
|
||||
from pydantic_ai.toolsets import AbstractToolset, FunctionToolset
|
||||
|
||||
|
||||
class SubAgentToolset(FunctionToolset[AgentDepsT]):
|
||||
"""Exposes one delegate tool that dispatches a task to a named sub-agent.
|
||||
|
||||
Each delegation runs the child agent in a fresh run with its own message
|
||||
history, so the sub-agent never sees the parent conversation. The parent's
|
||||
`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.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
agents: Mapping[str, AbstractAgent[AgentDepsT, Any]],
|
||||
forward_usage: bool,
|
||||
inherit_tools: bool,
|
||||
shared_capabilities: Sequence[AgentCapability[AgentDepsT]],
|
||||
event_stream_handler: EventStreamHandler[AgentDepsT] | None,
|
||||
tool_name: str,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self._agents: dict[str, AbstractAgent[AgentDepsT, Any]] = 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.add_function(self.delegate_task, name=tool_name)
|
||||
|
||||
def _inherited_toolsets(self, ctx: RunContext[AgentDepsT]) -> list[AbstractToolset[AgentDepsT]] | None:
|
||||
"""The parent's toolsets, with the delegate tool filtered out (no recursion)."""
|
||||
agent = ctx.agent
|
||||
if agent is None: # pragma: no cover - the running agent is always set during a run
|
||||
return None
|
||||
return [toolset.filtered(lambda _ctx, tool_def: tool_def.name != self._tool_name) for toolset in agent.toolsets]
|
||||
|
||||
async def delegate_task(self, ctx: RunContext[AgentDepsT], agent_name: str, task: str) -> str:
|
||||
"""Delegate a self-contained task to a named sub-agent and return its result.
|
||||
|
||||
The sub-agent runs in its own fresh context and does not see this
|
||||
conversation, so `task` must contain everything it needs.
|
||||
|
||||
Args:
|
||||
ctx: The run context (provides the parent's deps, usage, and tools).
|
||||
agent_name: Name of the sub-agent to run. Must be one of the agents
|
||||
listed in the instructions.
|
||||
task: The complete, self-contained instruction for the sub-agent.
|
||||
"""
|
||||
agent = self._agents.get(agent_name)
|
||||
if agent is None:
|
||||
available = ', '.join(sorted(self._agents))
|
||||
raise ModelRetry(f'Unknown sub-agent {agent_name!r}. Available sub-agents: {available}.')
|
||||
toolsets = self._inherited_toolsets(ctx) if self._inherit_tools else None
|
||||
capabilities = self._shared_capabilities or None
|
||||
usage = ctx.usage if self._forward_usage else None
|
||||
try:
|
||||
result = await agent.run(
|
||||
task,
|
||||
deps=ctx.deps,
|
||||
usage=usage,
|
||||
toolsets=toolsets,
|
||||
capabilities=capabilities,
|
||||
event_stream_handler=self._event_stream_handler,
|
||||
)
|
||||
except (ModelRetry, UnexpectedModelBehavior) as exc:
|
||||
# Soft sub-agent failures come back to the parent as a retry it can
|
||||
# react to; hard limits (e.g. UsageLimitExceeded) propagate.
|
||||
raise ModelRetry(f'Sub-agent {agent_name!r} failed: {exc}') from exc
|
||||
return str(result.output)
|
||||
@@ -0,0 +1,326 @@
|
||||
"""Tests for the SubAgents capability."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.capabilities import AbstractCapability
|
||||
from pydantic_ai.exceptions import UnexpectedModelBehavior, UsageLimitExceeded
|
||||
from pydantic_ai.messages import (
|
||||
AgentStreamEvent,
|
||||
ModelMessage,
|
||||
ModelResponse,
|
||||
RetryPromptPart,
|
||||
TextPart,
|
||||
ToolCallPart,
|
||||
ToolReturnPart,
|
||||
)
|
||||
from pydantic_ai.models.function import AgentInfo, FunctionModel
|
||||
from pydantic_ai.models.test import TestModel
|
||||
from pydantic_ai.tools import AgentDepsT, RunContext
|
||||
|
||||
from pydantic_ai_harness.experimental.subagents import SubAgents, SubAgentToolset
|
||||
|
||||
|
||||
@dataclass
|
||||
class _RecordingCapability(AbstractCapability[AgentDepsT]):
|
||||
"""Test capability whose dynamic instruction records each time it runs."""
|
||||
|
||||
log: list[str] = field(default_factory=list[str])
|
||||
|
||||
def get_instructions(self) -> Any:
|
||||
log = self.log
|
||||
|
||||
def _instructions(ctx: RunContext[AgentDepsT]) -> str:
|
||||
log.append('applied')
|
||||
return ''
|
||||
|
||||
return _instructions
|
||||
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def anyio_backend() -> str:
|
||||
"""Run async tests on the asyncio backend (matching upstream pydantic-ai)."""
|
||||
return 'asyncio'
|
||||
|
||||
|
||||
def _delegate_then_finish(agent_name: str, *, retries_before: int = 0) -> FunctionModel:
|
||||
"""A parent model that delegates to `agent_name` once, then replies with text.
|
||||
|
||||
`retries_before` extra delegations to a bogus agent happen first (to exercise
|
||||
the unknown-agent retry path).
|
||||
"""
|
||||
calls = {'n': 0}
|
||||
|
||||
def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
|
||||
calls['n'] += 1
|
||||
if calls['n'] <= retries_before:
|
||||
return ModelResponse(
|
||||
parts=[
|
||||
ToolCallPart('delegate_task', {'agent_name': 'ghost', 'task': 't'}, tool_call_id=f'b{calls["n"]}')
|
||||
]
|
||||
)
|
||||
if calls['n'] == retries_before + 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)
|
||||
|
||||
|
||||
class TestConstruction:
|
||||
def test_serialization_name_is_none(self) -> None:
|
||||
assert SubAgents.get_serialization_name() is None
|
||||
|
||||
def test_empty_agents_no_instructions(self) -> None:
|
||||
assert SubAgents[None]().get_instructions() is None
|
||||
|
||||
def test_empty_agents_no_toolset(self) -> None:
|
||||
assert SubAgents[None]().get_toolset() is None
|
||||
|
||||
|
||||
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()
|
||||
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()
|
||||
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()
|
||||
assert isinstance(instructions, str)
|
||||
assert '- plain' in instructions
|
||||
assert '- plain:' 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()
|
||||
assert isinstance(instructions, str)
|
||||
assert 'run_agent' in instructions
|
||||
|
||||
|
||||
class TestToolset:
|
||||
def test_get_toolset_exposes_delegate_tool(self) -> None:
|
||||
agent = Agent(TestModel(), name='x')
|
||||
toolset = SubAgents(agents={'x': 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()
|
||||
assert isinstance(toolset, SubAgentToolset)
|
||||
assert 'run_agent' in toolset.tools
|
||||
|
||||
|
||||
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})]
|
||||
)
|
||||
result = await parent.run('go')
|
||||
assert result.output == 'all done'
|
||||
returns = [
|
||||
part.content
|
||||
for message in result.all_messages()
|
||||
for part in message.parts
|
||||
if isinstance(part, ToolReturnPart) and part.tool_name == 'delegate_task'
|
||||
]
|
||||
assert returns == ['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})],
|
||||
)
|
||||
result = await parent.run('go')
|
||||
assert result.output == 'all done'
|
||||
# The bogus delegation produced a retry prompt naming the unknown agent and
|
||||
# listing the valid ones, sorted and comma-separated.
|
||||
retries = [
|
||||
part.content
|
||||
for message in result.all_messages()
|
||||
for part in message.parts
|
||||
if isinstance(part, RetryPromptPart) and part.tool_name == 'delegate_task'
|
||||
]
|
||||
assert any(
|
||||
"Unknown sub-agent 'ghost'" in str(r) and 'Available sub-agents: helper, worker' in str(r) for r in retries
|
||||
)
|
||||
|
||||
async def test_forwards_deps_and_shares_usage_by_default(self) -> None:
|
||||
captured: dict[str, Any] = {}
|
||||
parent_usage: dict[str, Any] = {}
|
||||
|
||||
worker = Agent(TestModel(custom_output_text='W'), name='worker', deps_type=str)
|
||||
|
||||
@worker.instructions
|
||||
def _capture(ctx: RunContext[str]) -> str: # pyright: ignore[reportUnusedFunction]
|
||||
captured['deps'] = ctx.deps
|
||||
captured['usage_is_parent'] = ctx.usage is parent_usage.get('usage')
|
||||
return ''
|
||||
|
||||
parent: Agent[str, str] = Agent(
|
||||
_delegate_then_finish('worker'),
|
||||
deps_type=str,
|
||||
capabilities=[SubAgents(agents={'worker': worker})],
|
||||
)
|
||||
|
||||
@parent.instructions
|
||||
def _remember_usage(ctx: RunContext[str]) -> str: # pyright: ignore[reportUnusedFunction]
|
||||
parent_usage['usage'] = ctx.usage
|
||||
return ''
|
||||
|
||||
result = await parent.run('go', deps='PARENT')
|
||||
assert result.output == 'all done'
|
||||
assert captured['deps'] == 'PARENT' # deps always forwarded
|
||||
assert captured['usage_is_parent'] is True # usage shared by default
|
||||
|
||||
async def test_inherit_tools_exposes_parent_tools_but_not_delegate(self) -> None:
|
||||
offered: list[str] = []
|
||||
|
||||
def worker_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
|
||||
if not offered:
|
||||
offered.extend(tool.name for tool in info.function_tools)
|
||||
# Call the inherited tool to prove it's actually usable, not just listed.
|
||||
return ModelResponse(parts=[ToolCallPart('parent_tool', {}, tool_call_id='p1')])
|
||||
return ModelResponse(parts=[TextPart('sub done')])
|
||||
|
||||
worker = Agent(FunctionModel(worker_fn), name='worker')
|
||||
parent: Agent[None, str] = Agent(
|
||||
_delegate_then_finish('worker'),
|
||||
capabilities=[SubAgents(agents={'worker': worker}, inherit_tools=True)],
|
||||
)
|
||||
|
||||
@parent.tool_plain
|
||||
def parent_tool() -> str: # pyright: ignore[reportUnusedFunction]
|
||||
return 'PT'
|
||||
|
||||
result = await parent.run('go')
|
||||
assert result.output == 'all done'
|
||||
assert 'parent_tool' in offered # the parent's tool is inherited by the sub-agent
|
||||
assert 'delegate_task' not in offered # the delegate tool is filtered out, so no recursion
|
||||
|
||||
async def test_shared_capabilities_applied_to_subagent(self) -> None:
|
||||
cap: _RecordingCapability[None] = _RecordingCapability()
|
||||
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])],
|
||||
)
|
||||
result = await parent.run('go')
|
||||
assert result.output == 'all done'
|
||||
assert cap.log == ['applied'] # the shared capability ran during the sub-agent run
|
||||
|
||||
async def test_event_stream_handler_forwarded_to_subagent(self) -> None:
|
||||
events: list[str] = []
|
||||
|
||||
async def handler(ctx: RunContext[None], stream: AsyncIterable[AgentStreamEvent]) -> None:
|
||||
async for event in stream:
|
||||
events.append(type(event).__name__)
|
||||
|
||||
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)],
|
||||
)
|
||||
result = await parent.run('go')
|
||||
assert result.output == 'all done'
|
||||
assert events # the sub-agent's run streamed events to the handler
|
||||
|
||||
async def test_hard_limit_propagates(self) -> None:
|
||||
def boom(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
|
||||
raise UsageLimitExceeded('limit hit')
|
||||
|
||||
limited = Agent(FunctionModel(boom), name='limited')
|
||||
parent: Agent[None, str] = Agent(
|
||||
_delegate_then_finish('limited'),
|
||||
capabilities=[SubAgents(agents={'limited': limited})],
|
||||
)
|
||||
# Hard limits are not converted to a retry -- they propagate to stop the run.
|
||||
with pytest.raises(UsageLimitExceeded):
|
||||
await parent.run('go')
|
||||
|
||||
async def test_soft_subagent_failure_becomes_model_retry(self) -> None:
|
||||
def boom(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
|
||||
raise UnexpectedModelBehavior('kaboom')
|
||||
|
||||
boomer = Agent(FunctionModel(boom), name='boomer')
|
||||
worker = Agent(TestModel(custom_output_text='OK'), name='worker')
|
||||
|
||||
calls = {'n': 0}
|
||||
|
||||
def parent_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
|
||||
calls['n'] += 1
|
||||
if calls['n'] == 1:
|
||||
return ModelResponse(
|
||||
parts=[ToolCallPart('delegate_task', {'agent_name': 'boomer', 'task': 't'}, tool_call_id='c1')]
|
||||
)
|
||||
if calls['n'] == 2:
|
||||
return ModelResponse(
|
||||
parts=[ToolCallPart('delegate_task', {'agent_name': 'worker', 'task': 't'}, tool_call_id='c2')]
|
||||
)
|
||||
return ModelResponse(parts=[TextPart('all done')])
|
||||
|
||||
parent: Agent[None, str] = Agent(
|
||||
FunctionModel(parent_fn),
|
||||
capabilities=[SubAgents(agents={'boomer': boomer, 'worker': worker})],
|
||||
)
|
||||
result = await parent.run('go')
|
||||
assert result.output == 'all done'
|
||||
retries = [
|
||||
part.content
|
||||
for message in result.all_messages()
|
||||
for part in message.parts
|
||||
if isinstance(part, RetryPromptPart) and part.tool_name == 'delegate_task'
|
||||
]
|
||||
assert any("Sub-agent 'boomer' failed" in str(r) for r in retries)
|
||||
|
||||
async def test_usage_not_shared_when_disabled(self) -> None:
|
||||
captured: dict[str, Any] = {}
|
||||
parent_usage: dict[str, Any] = {}
|
||||
|
||||
worker = Agent(TestModel(custom_output_text='W'), name='worker', deps_type=str)
|
||||
|
||||
@worker.instructions
|
||||
def _capture(ctx: RunContext[str]) -> str: # pyright: ignore[reportUnusedFunction]
|
||||
captured['deps'] = ctx.deps
|
||||
captured['usage_is_parent'] = ctx.usage is parent_usage.get('usage')
|
||||
return ''
|
||||
|
||||
parent: Agent[str, str] = Agent(
|
||||
_delegate_then_finish('worker'),
|
||||
deps_type=str,
|
||||
capabilities=[SubAgents(agents={'worker': worker}, forward_usage=False)],
|
||||
)
|
||||
|
||||
@parent.instructions
|
||||
def _remember_usage(ctx: RunContext[str]) -> str: # pyright: ignore[reportUnusedFunction]
|
||||
parent_usage['usage'] = ctx.usage
|
||||
return ''
|
||||
|
||||
result = await parent.run('go', deps='PARENT')
|
||||
assert result.output == 'all done'
|
||||
assert captured['deps'] == 'PARENT' # deps still forwarded
|
||||
assert captured['usage_is_parent'] is False # usage isolated
|
||||
Reference in New Issue
Block a user