Skip to content

Add spec parameter to agent.override() and agent.run() #4769

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
DouweM merged 16 commits into from
Mar 23, 2026
+699 −42
Merged

Add spec parameter to agent.override() and agent.run() #4769

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
3a9a537
Add `spec` parameter to `agent.override()` and `agent.run()`
DouweM Mar 21, 2026
fa70cc6
Merge branch 'capabilities' into agent-run-spec
DouweM Mar 21, 2026
94ee345
Add `spec` param to all durable agent method overloads and docstrings
DouweM Mar 21, 2026
65680fb
Fix CI: add missing spec to temporal iter overload, remove unused tes…
DouweM Mar 21, 2026
1504a8a
Address review feedback: fix bugs and strengthen tests
DouweM Mar 21, 2026
443a217
Extract shared _capabilities_from_spec() helper to deduplicate from_s…
DouweM Mar 21, 2026
1cda25d
Restore override replacement semantics for spec capabilities
DouweM Mar 22, 2026
76e9510
Fix PrefectAgent.iter() not forwarding builtin_tools, warn on descrip…
DouweM Mar 22, 2026
665ac93
Merge remote-tracking branch 'origin/capabilities' into agent-run-spec
DouweM Mar 22, 2026
f654390
Address review feedback round 2: extract _validate_spec, dynamic defa…
DouweM Mar 22, 2026
88fe894
Merge remote-tracking branch 'origin/capabilities' into agent-run-spec
DouweM Mar 23, 2026
ef7c0a9
Address review feedback round 2: extract _validate_spec, dynamic defa…
DouweM Mar 23, 2026
3994a6a
Fix pyright: use conftest IsDatetime/IsStr wrappers instead of dirty_…
DouweM Mar 23, 2026
af02274
Fix Prefect import formatting, clean up review feedback
DouweM Mar 23, 2026
2c536f3
Merge remote-tracking branch 'origin/capabilities' into agent-run-spec
DouweM Mar 23, 2026
a364a27
Merge remote-tracking branch 'origin/capabilities' into agent-run-spec
DouweM Mar 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
259 changes: 223 additions & 36 deletions pydantic_ai_slim/pydantic_ai/agent/__init__.py
View file
Open in desktop

Large diffs are not rendered by default.

34 changes: 34 additions & 0 deletions pydantic_ai_slim/pydantic_ai/agent/abstract.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@
from starlette.routing import BaseRoute, Route
from starlette.types import ExceptionHandler, Lifespan

from pydantic_ai.agent.spec import AgentSpec
from pydantic_ai.ui.ag_ui.app import AGUIApp


Expand Down Expand Up @@ -177,6 +178,7 @@ async def run(
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AgentBuiltinTool[AgentDepsT]] | None = None,
event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> AgentRunResult[OutputDataT]: ...

@overload
Expand All @@ -198,6 +200,7 @@ async def run(
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AgentBuiltinTool[AgentDepsT]] | None = None,
event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> AgentRunResult[RunOutputDataT]: ...

async def run(
Expand All @@ -218,6 +221,7 @@ async def run(
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AgentBuiltinTool[AgentDepsT]] | None = None,
event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> AgentRunResult[Any]:
"""Run the agent with a user prompt in async mode.

Expand Down Expand Up @@ -256,6 +260,7 @@ async def main():
toolsets: Optional additional toolsets for this run.
event_stream_handler: Optional handler for events from the model's streaming response and the agent's execution of tools to use for this run.
builtin_tools: Optional additional builtin tools for this run.
spec: Optional agent spec to apply for this run. At run time, spec values are additive.

Returns:
The result of the run.
Expand All @@ -279,6 +284,7 @@ async def main():
metadata=metadata,
toolsets=toolsets,
builtin_tools=builtin_tools,
spec=spec,
) as agent_run:
# Drive via next() so wrap_node_run hooks fire for each node.
node = agent_run.next_node
Expand Down Expand Up @@ -314,6 +320,7 @@ def run_sync(
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AgentBuiltinTool[AgentDepsT]] | None = None,
event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> AgentRunResult[OutputDataT]: ...

@overload
Expand All @@ -335,6 +342,7 @@ def run_sync(
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AgentBuiltinTool[AgentDepsT]] | None = None,
event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> AgentRunResult[RunOutputDataT]: ...

def run_sync(
Expand All @@ -355,6 +363,7 @@ def run_sync(
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AgentBuiltinTool[AgentDepsT]] | None = None,
event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> AgentRunResult[Any]:
"""Synchronously run the agent with a user prompt.

Expand Down Expand Up @@ -392,6 +401,7 @@ def run_sync(
toolsets: Optional additional toolsets for this run.
event_stream_handler: Optional handler for events from the model's streaming response and the agent's execution of tools to use for this run.
builtin_tools: Optional additional builtin tools for this run.
spec: Optional agent spec to apply for this run. At run time, spec values are additive.

Returns:
The result of the run.
Expand All @@ -416,6 +426,7 @@ def run_sync(
toolsets=toolsets,
builtin_tools=builtin_tools,
event_stream_handler=event_stream_handler,
spec=spec,
)
)

Expand All @@ -438,6 +449,7 @@ def run_stream(
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AgentBuiltinTool[AgentDepsT]] | None = None,
event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> AbstractAsyncContextManager[result.StreamedRunResult[AgentDepsT, OutputDataT]]: ...

@overload
Expand All @@ -459,6 +471,7 @@ def run_stream(
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AgentBuiltinTool[AgentDepsT]] | None = None,
event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> AbstractAsyncContextManager[result.StreamedRunResult[AgentDepsT, RunOutputDataT]]: ...

@asynccontextmanager
Expand All @@ -480,6 +493,7 @@ async def run_stream( # noqa: C901
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AgentBuiltinTool[AgentDepsT]] | None = None,
event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> AsyncIterator[result.StreamedRunResult[AgentDepsT, Any]]:
"""Run the agent with a user prompt in async streaming mode.

Expand Down Expand Up @@ -527,6 +541,7 @@ async def main():
event_stream_handler: Optional handler for events from the model's streaming response and the agent's execution of tools to use for this run.
It will receive all the events up until the final result is found, which you can then read or stream from inside the context manager.
Note that it does _not_ receive any events after the final result is found.
spec: Optional agent spec to apply for this run. At run time, spec values are additive.

Returns:
The result of the run.
Expand Down Expand Up @@ -554,6 +569,7 @@ async def main():
infer_name=False,
toolsets=toolsets,
builtin_tools=builtin_tools,
spec=spec,
) as agent_run:
# Handle wrap_run short-circuit: result is already available
if agent_run.result is not None:
Expand Down Expand Up @@ -702,6 +718,7 @@ def run_stream_sync(
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AgentBuiltinTool[AgentDepsT]] | None = None,
event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> result.StreamedRunResultSync[AgentDepsT, OutputDataT]: ...

@overload
Expand All @@ -722,6 +739,7 @@ def run_stream_sync(
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AbstractBuiltinTool] | None = None,
event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> result.StreamedRunResultSync[AgentDepsT, RunOutputDataT]: ...

def run_stream_sync(
Expand All @@ -741,6 +759,7 @@ def run_stream_sync(
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AgentBuiltinTool[AgentDepsT]] | None = None,
event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> result.StreamedRunResultSync[AgentDepsT, Any]:
"""Run the agent with a user prompt in sync streaming mode.

Expand Down Expand Up @@ -790,6 +809,7 @@ def main():
event_stream_handler: Optional handler for events from the model's streaming response and the agent's execution of tools to use for this run.
It will receive all the events up until the final result is found, which you can then read or stream from inside the context manager.
Note that it does _not_ receive any events after the final result is found.
spec: Optional agent spec to apply for this run. At run time, spec values are additive.

Returns:
The result of the run.
Expand All @@ -813,6 +833,7 @@ async def _consume_stream():
toolsets=toolsets,
builtin_tools=builtin_tools,
event_stream_handler=event_stream_handler,
spec=spec,
) as stream_result:
yield stream_result

Expand All @@ -837,6 +858,7 @@ def run_stream_events(
infer_name: bool = True,
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AgentBuiltinTool[AgentDepsT]] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> AsyncIterator[_messages.AgentStreamEvent | AgentRunResultEvent[OutputDataT]]: ...

@overload
Expand All @@ -857,6 +879,7 @@ def run_stream_events(
infer_name: bool = True,
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AgentBuiltinTool[AgentDepsT]] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> AsyncIterator[_messages.AgentStreamEvent | AgentRunResultEvent[RunOutputDataT]]: ...

def run_stream_events(
Expand All @@ -876,6 +899,7 @@ def run_stream_events(
infer_name: bool = True,
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AgentBuiltinTool[AgentDepsT]] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> AsyncIterator[_messages.AgentStreamEvent | AgentRunResultEvent[Any]]:
"""Run the agent with a user prompt in async mode and stream events from the run.

Expand Down Expand Up @@ -930,6 +954,7 @@ async def main():
infer_name: Whether to try to infer the agent name from the call frame if it's not set.
toolsets: Optional additional toolsets for this run.
builtin_tools: Optional additional builtin tools for this run.
spec: Optional agent spec to apply for this run. At run time, spec values are additive.

Returns:
An async iterable of stream events `AgentStreamEvent` and finally a `AgentRunResultEvent` with the final
Expand All @@ -955,6 +980,7 @@ async def main():
metadata=metadata,
toolsets=toolsets,
builtin_tools=builtin_tools,
spec=spec,
)

async def _run_stream_events(
Expand All @@ -973,6 +999,7 @@ async def _run_stream_events(
metadata: AgentMetadata[AgentDepsT] | None = None,
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AgentBuiltinTool[AgentDepsT]] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> AsyncIterator[_messages.AgentStreamEvent | AgentRunResultEvent[Any]]:
send_stream, receive_stream = anyio.create_memory_object_stream[
_messages.AgentStreamEvent | AgentRunResultEvent[Any]
Expand Down Expand Up @@ -1002,6 +1029,7 @@ async def run_agent() -> AgentRunResult[Any]:
toolsets=toolsets,
builtin_tools=builtin_tools,
event_stream_handler=event_stream_handler,
spec=spec,
)

task = asyncio.create_task(run_agent())
Expand Down Expand Up @@ -1037,6 +1065,7 @@ def iter(
infer_name: bool = True,
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AgentBuiltinTool[AgentDepsT]] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> AbstractAsyncContextManager[AgentRun[AgentDepsT, OutputDataT]]: ...

@overload
Expand All @@ -1057,6 +1086,7 @@ def iter(
infer_name: bool = True,
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AgentBuiltinTool[AgentDepsT]] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> AbstractAsyncContextManager[AgentRun[AgentDepsT, RunOutputDataT]]: ...

@asynccontextmanager
Expand All @@ -1078,6 +1108,7 @@ async def iter(
infer_name: bool = True,
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AgentBuiltinTool[AgentDepsT]] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> AsyncIterator[AgentRun[AgentDepsT, Any]]:
"""A contextmanager which can be used to iterate over the agent graph's nodes as they are executed.

Expand Down Expand Up @@ -1159,6 +1190,7 @@ async def main():
infer_name: Whether to try to infer the agent name from the call frame if it's not set.
toolsets: Optional additional toolsets for this run.
builtin_tools: Optional additional builtin tools for this run.
spec: Optional agent spec to apply for this run. At run time, spec values are additive.

Returns:
The result of the run.
Expand All @@ -1178,6 +1210,7 @@ def override(
tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] | _utils.Unset = _utils.UNSET,
instructions: _instructions.AgentInstructions[AgentDepsT] | _utils.Unset = _utils.UNSET,
model_settings: AgentModelSettings[AgentDepsT] | _utils.Unset = _utils.UNSET,
spec: dict[str, Any] | AgentSpec | None = None,
) -> Iterator[None]:
"""Context manager to temporarily override agent name, dependencies, model, toolsets, tools, or instructions.

Expand All @@ -1193,6 +1226,7 @@ def override(
instructions: The instructions to use instead of the instructions registered with the agent.
model_settings: The model settings to use instead of the model settings passed to the agent constructor.
When set, any per-run `model_settings` argument is ignored.
spec: Optional agent spec providing defaults for override.
"""
raise NotImplementedError
yield
Expand Down
4 changes: 2 additions & 2 deletions pydantic_ai_slim/pydantic_ai/agent/spec.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ class AgentSpec(BaseModel):

# $schema is included to avoid validation fails from the `$schema` key, see `_add_json_schema` below for context
json_schema_path: str | None = Field(default=None, alias='$schema')
model: str
model: str | None = None
Copy link
Contributor

@github-actions github-actions bot Mar 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Making model optional changes the contract of AgentSpec — specs that were previously valid (always had a model) now allow None. While the from_spec() method has a runtime check (line 697-700), this is a breaking change for existing code that relies on spec.model being a str. Any serialized specs without a model field will now silently validate as None rather than failing at parse time.

Is this intentional? If model should only be optional at run/override time (not for from_spec()/from_file()), consider keeping model: str on AgentSpec and using a separate type or a model: str | None only in the _resolve_spec path.

Copy link
Contributor

@github-actions github-actions bot Mar 22, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Making model optional here changes the contract of AgentSpec globally — it's now valid to create a spec (including via from_file()) without a model. While the from_spec() method has a runtime check, this is a weakening of the type-level contract for all users of AgentSpec, not just the run/override use case.

Consider instead keeping model: str required in AgentSpec and introducing a separate type (e.g. PartialAgentSpec or just using dict[str, Any]) for the run/override case where model is optional. Alternatively, if AgentSpec should support partial specs in general, the from_spec() error message should be updated to reflect that this is now an expected validation at the from_spec() level rather than a "bug" in the spec, and the JSON schema should note that model is optional.

@DouweM — would appreciate your input on whether partial specs should be a first-class concept in AgentSpec or scoped to the run/override path.

name: str | None = None
description: TemplateStr[Any] | str | None = None
instructions: TemplateStr[Any] | str | list[TemplateStr[Any] | str] | None = None
Expand Down Expand Up @@ -194,7 +194,7 @@ def model_json_schema_with_capabilities(
# - extra='forbid' enables strict validation in the generated schema
# When adding or removing fields on AgentSpec, update this class to match.
class _AgentSpecSchema(BaseModel, extra='forbid'):
model: str
model: str | None = None
name: str | None = None
description: str | None = None
instructions: str | list[str] | None = None
Expand Down
13 changes: 12 additions & 1 deletion pydantic_ai_slim/pydantic_ai/agent/wrapper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

from collections.abc import AsyncIterator, Iterator, Sequence
from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager
from typing import Any, overload
from typing import TYPE_CHECKING, Any, overload

from .. import (
_instructions,
Expand All @@ -25,6 +25,9 @@
from ..toolsets import AbstractToolset
from .abstract import AbstractAgent, AgentMetadata, AgentModelSettings, EventStreamHandler, RunOutputDataT

if TYPE_CHECKING:
from .spec import AgentSpec


class WrapperAgent(AbstractAgent[AgentDepsT, OutputDataT]):
"""Agent which wraps another agent.
Expand Down Expand Up @@ -98,6 +101,7 @@ def iter(
infer_name: bool = True,
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AgentBuiltinTool[AgentDepsT]] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> AbstractAsyncContextManager[AgentRun[AgentDepsT, OutputDataT]]: ...

@overload
Expand All @@ -118,6 +122,7 @@ def iter(
infer_name: bool = True,
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AgentBuiltinTool[AgentDepsT]] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> AbstractAsyncContextManager[AgentRun[AgentDepsT, RunOutputDataT]]: ...

@asynccontextmanager
Expand All @@ -138,6 +143,7 @@ async def iter(
infer_name: bool = True,
toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
builtin_tools: Sequence[AgentBuiltinTool[AgentDepsT]] | None = None,
spec: dict[str, Any] | AgentSpec | None = None,
) -> AsyncIterator[AgentRun[AgentDepsT, Any]]:
"""A contextmanager which can be used to iterate over the agent graph's nodes as they are executed.

Expand Down Expand Up @@ -216,6 +222,7 @@ async def main():
infer_name: Whether to try to infer the agent name from the call frame if it's not set.
toolsets: Optional additional toolsets for this run.
builtin_tools: Optional additional builtin tools for this run.
spec: Optional agent spec to apply for this run.

Returns:
The result of the run.
Expand All @@ -235,6 +242,7 @@ async def main():
infer_name=infer_name,
toolsets=toolsets,
builtin_tools=builtin_tools,
spec=spec,
) as run:
yield run

Expand All @@ -249,6 +257,7 @@ def override(
tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] | _utils.Unset = _utils.UNSET,
instructions: _instructions.AgentInstructions[AgentDepsT] | _utils.Unset = _utils.UNSET,
model_settings: AgentModelSettings[AgentDepsT] | _utils.Unset = _utils.UNSET,
spec: dict[str, Any] | AgentSpec | None = None,
) -> Iterator[None]:
"""Context manager to temporarily override agent name, dependencies, model, toolsets, tools, or instructions.

Expand All @@ -264,6 +273,7 @@ def override(
instructions: The instructions to use instead of the instructions registered with the agent.
model_settings: The model settings to use instead of the model settings passed to the agent constructor.
When set, any per-run `model_settings` argument is ignored.
spec: Optional agent spec to apply as overrides.
"""
with self.wrapped.override(
name=name,
Expand All @@ -273,5 +283,6 @@ def override(
tools=tools,
instructions=instructions,
model_settings=model_settings,
spec=spec,
):
yield
Loading
Loading