Skip to content

Toolsets #2024

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.

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

Draft
wants to merge 91 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from 87 commits
Commits
Show all changes
91 commits
Select commit Hold shift + click to select a range
e290951
WIP: Output modes
DouweM Jun 3, 2025
2056539
WIP: More output modes
DouweM Jun 3, 2025
bceba19
Merge remote-tracking branch 'origin/main' into output-modes
DouweM Jun 3, 2025
0cb25c4
Fix tests
DouweM Jun 3, 2025
933b74e
Remove syntax invalid before Python 3.12
DouweM Jun 3, 2025
7974df0
Fix tests
DouweM Jun 3, 2025
9cc19e2
Add TextOutput marker
DouweM Jun 9, 2025
bc6bb65
Merge remote-tracking branch 'origin/main' into output-modes
DouweM Jun 9, 2025
0e356a3
Add VCR recording of new test
DouweM Jun 9, 2025
81312dc
Implement additional output modes in GeminiModel and GoogleModel
DouweM Jun 10, 2025
52ef4d5
Fix prompted_json on OpenAIResponses
DouweM Jun 10, 2025
fe05956
Test output modes on Gemini and Anthropic
DouweM Jun 10, 2025
94421f3
Add VCR recordings of Gemini output mode tests
DouweM Jun 10, 2025
1902d00
Remove some old TODO comments
DouweM Jun 10, 2025
1f53c9b
Add missing VCR recording of Gemini output mode test
DouweM Jun 10, 2025
a4c2877
Add more missing VCR recordings
DouweM Jun 10, 2025
56e58f9
Fix OpenAI tools
DouweM Jun 10, 2025
a5234e1
Improve test coverage
DouweM Jun 10, 2025
40def08
Update unsupported output mode error message
DouweM Jun 10, 2025
837d305
Improve test coverage
DouweM Jun 10, 2025
3598bef
Merge branch 'main' into output-modes
DouweM Jun 10, 2025
5f71ba8
Test streaming with structured text output
DouweM Jun 10, 2025
cfc2749
Make TextOutputFunction Python 3.9 compatible
DouweM Jun 10, 2025
a137641
Properly merge JSON schemas accounting for defs
DouweM Jun 11, 2025
f495d46
Refactor output schemas and modes: more 'isinstance(output_schema, ..…
DouweM Jun 12, 2025
449ed0d
Merge branch 'main' into output-modes
DouweM Jun 12, 2025
e70d249
Clean up some variable names
DouweM Jun 12, 2025
4592b0b
Improve test coverage
DouweM Jun 12, 2025
db1c628
Merge branch 'main' into output-modes
DouweM Jun 13, 2025
f57d078
Combine JsonSchemaOutput and PromptedJsonOutput into StructuredTextOu…
DouweM Jun 13, 2025
5112455
Add missing cassettes
DouweM Jun 13, 2025
416cc7d
Can't use dataclass kw_only on 3.9
DouweM Jun 13, 2025
4b0e5cf
Improve test coverage
DouweM Jun 13, 2025
094920f
Improve test coverage
DouweM Jun 13, 2025
9f61706
Improve test coverage
DouweM Jun 13, 2025
9f51387
Remove unnecessary coverage ignores
DouweM Jun 13, 2025
9a1e628
Remove unnecessary coverage ignore
DouweM Jun 13, 2025
2b5fa81
Add docs
DouweM Jun 13, 2025
6c4662b
Fix docs refs
DouweM Jun 13, 2025
3ed3431
Fix nested list in docs
DouweM Jun 13, 2025
3d77818
Merge branch 'main' into output-modes
DouweM Jun 17, 2025
a86d7d4
Split StructuredTextOutput into ModelStructuredOutput and PromptedStr…
DouweM Jun 17, 2025
ce985a0
Merge branch 'main' into output-modes
DouweM Jun 17, 2025
71d1655
Fix WrapperModel.profile
DouweM Jun 17, 2025
8c04144
Update output modes docs
DouweM Jun 17, 2025
d78b5f7
Add examples to output mode marker docstrings
DouweM Jun 17, 2025
70d1197
Fix mypy type inference
DouweM Jun 17, 2025
2eb7fd1
Improve test coverage
DouweM Jun 17, 2025
25ccb54
Merge branch 'main' into output-modes
DouweM Jun 17, 2025
9e00c32
Import cast and RunContext in _function_schema
DouweM Jun 17, 2025
7de3c0d
Move RunContext and AgentDepsT into their own module to solve circula…
DouweM Jun 17, 2025
4029fac
Make _run_context module private, RunContext can be accessed through …
DouweM Jun 17, 2025
98bccf2
Merge branch 'main' into output-modes
DouweM Jun 19, 2025
8041cf3
Fix thinking part related tests
DouweM Jun 19, 2025
9bfed04
Implement Toolset
DouweM Jun 20, 2025
0f8da74
Make MCPServer a Toolset
DouweM Jun 20, 2025
8a29836
--no-edit
DouweM Jun 21, 2025
3d2012c
Add MappedToolset
DouweM Jun 21, 2025
901267d
Import Never from typing_extensions instead of typing
DouweM Jun 21, 2025
b9258d7
from __future__ import annotations
DouweM Jun 21, 2025
27ccbd1
Update client.md
DouweM Jun 21, 2025
3031e55
Pass only RunToolset to agent graph
DouweM Jun 21, 2025
ebd0b57
Make WrapperToolset abstract
DouweM Jun 21, 2025
867bf68
Introduce ToolDefinition.kind == 'pending'
DouweM Jun 21, 2025
c1115ae
Rename pending tools to deferred tools
DouweM Jun 21, 2025
6abd603
Merge branch 'main' into toolsets
DouweM Jun 24, 2025
a2f69df
Fix retries
DouweM Jun 24, 2025
0e0bf35
Remove duplicate cassettes
DouweM Jun 24, 2025
735df29
Merge branch 'main' into toolsets
DouweM Jun 26, 2025
8745a7a
Pass just one toolset into the run
DouweM Jun 26, 2025
05aa972
WIP
DouweM Jun 26, 2025
ad6e826
Fix streaming tool calls
DouweM Jun 27, 2025
84cd954
Stop double counting retries and reset on success
DouweM Jun 27, 2025
74a56ae
Fix retry error wrapping
DouweM Jun 27, 2025
0360e77
Make DeferredToolCalls work with streaming
DouweM Jun 30, 2025
6607b00
Merge branch 'main' into toolsets
DouweM Jun 30, 2025
8a3febb
Let toolsets be overridden in run/iter/run_stream/run_sync
DouweM Jun 30, 2025
2e200ac
Add DeferredToolset
DouweM Jun 30, 2025
1cb7f32
Add LangChainToolset
DouweM Jun 30, 2025
a6eba43
Add Agent.prepare_output_tools
DouweM Jun 30, 2025
0c96126
Require WrapperToolset subclasses to implement their own prepare_for_run
DouweM Jul 1, 2025
2348f45
Require DeferredToolCalls to be used with other output type
DouweM Jul 1, 2025
9dc684e
Merge branch 'main' into toolsets
DouweM Jul 1, 2025
f3124c0
Lots of cleanup
DouweM Jul 1, 2025
f660cc1
Some more tweaks
DouweM Jul 2, 2025
64dacbb
Merge branch 'main' into toolsets
DouweM Jul 2, 2025
5ca305e
Fix docs example
DouweM Jul 2, 2025
c5ef5f6
Address some feedback
DouweM Jul 2, 2025
badbe23
Merge branch 'main' into toolsets
DouweM Jul 2, 2025
acddb8d
Add sampling_model to Agent __init__, iter, run (etc), and override, …
DouweM Jul 2, 2025
89fc266
Turn RunContext.retries from a defaultdict into a dict again as the 0…
DouweM Jul 2, 2025
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
2 changes: 1 addition & 1 deletion docs/agents.md
Original file line number Diff line number Diff line change
Expand Up @@ -808,7 +808,7 @@ with capture_run_messages() as messages: # (2)!
result = agent.run_sync('Please get me the volume of a box with size 6.')
except UnexpectedModelBehavior as e:
print('An error occurred:', e)
#> An error occurred: Tool exceeded max retries count of 1
#> An error occurred: Tool 'calc_volume' exceeded max retries count of 1
print('cause:', repr(e.__cause__))
#> cause: ModelRetry('Please try again.')
print('messages:', messages)
Expand Down
39 changes: 20 additions & 19 deletions docs/mcp/client.md
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ Examples of both are shown below; [mcp-run-python](run-python.md) is used as the
[`MCPServerSSE`][pydantic_ai.mcp.MCPServerSSE] connects over HTTP using the [HTTP + Server Sent Events transport](https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#http-with-sse) to a server.

!!! note
[`MCPServerSSE`][pydantic_ai.mcp.MCPServerSSE] requires an MCP server to be running and accepting HTTP connections before calling [`agent.run_mcp_servers()`][pydantic_ai.Agent.run_mcp_servers]. Running the server is not managed by PydanticAI.
[`MCPServerSSE`][pydantic_ai.mcp.MCPServerSSE] requires an MCP server to be running and accepting HTTP connections before calling [`agent.run_toolsets()`][pydantic_ai.Agent.run_toolsets]. Running the server is not managed by PydanticAI.

The name "HTTP" is used since this implementation will be adapted in future to use the new
[Streamable HTTP](https://github.com/modelcontextprotocol/specification/pull/206) currently in development.
Expand All @@ -47,11 +47,11 @@ from pydantic_ai import Agent
from pydantic_ai.mcp import MCPServerSSE

server = MCPServerSSE(url='http://localhost:3001/sse') # (1)!
agent = Agent('openai:gpt-4o', mcp_servers=[server]) # (2)!
agent = Agent('openai:gpt-4o', toolsets=[server]) # (2)!


async def main():
async with agent.run_mcp_servers(): # (3)!
async with agent.run_toolsets(): # (3)!
result = await agent.run('How many days between 2000-01-01 and 2025-03-18?')
print(result.output)
#> There are 9,208 days between January 1, 2000, and March 18, 2025.
Expand Down Expand Up @@ -93,7 +93,7 @@ Will display as follows:
!!! note
[`MCPServerStreamableHTTP`][pydantic_ai.mcp.MCPServerStreamableHTTP] requires an MCP server to be
running and accepting HTTP connections before calling
[`agent.run_mcp_servers()`][pydantic_ai.Agent.run_mcp_servers]. Running the server is not
[`agent.run_toolsets()`][pydantic_ai.Agent.run_toolsets]. Running the server is not
managed by PydanticAI.

Before creating the Streamable HTTP client, we need to run a server that supports the Streamable HTTP transport.
Expand All @@ -118,10 +118,10 @@ from pydantic_ai import Agent
from pydantic_ai.mcp import MCPServerStreamableHTTP

server = MCPServerStreamableHTTP('http://localhost:8000/mcp') # (1)!
agent = Agent('openai:gpt-4o', mcp_servers=[server]) # (2)!
agent = Agent('openai:gpt-4o', toolsets=[server]) # (2)!

async def main():
async with agent.run_mcp_servers(): # (3)!
async with agent.run_toolsets(): # (3)!
result = await agent.run('How many days between 2000-01-01 and 2025-03-18?')
print(result.output)
#> There are 9,208 days between January 1, 2000, and March 18, 2025.
Expand All @@ -138,7 +138,7 @@ _(This example is complete, it can be run "as is" with Python 3.10+ — you'll n
The other transport offered by MCP is the [stdio transport](https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#stdio) where the server is run as a subprocess and communicates with the client over `stdin` and `stdout`. In this case, you'd use the [`MCPServerStdio`][pydantic_ai.mcp.MCPServerStdio] class.

!!! note
When using [`MCPServerStdio`][pydantic_ai.mcp.MCPServerStdio] servers, the [`agent.run_mcp_servers()`][pydantic_ai.Agent.run_mcp_servers] context manager is responsible for starting and stopping the server.
When using [`MCPServerStdio`][pydantic_ai.mcp.MCPServerStdio] servers, the [`agent.run_toolsets()`][pydantic_ai.Agent.run_toolsets] context manager is responsible for starting and stopping the server.

```python {title="mcp_stdio_client.py" py="3.10"}
from pydantic_ai import Agent
Expand All @@ -156,11 +156,11 @@ server = MCPServerStdio( # (1)!
'stdio',
]
)
agent = Agent('openai:gpt-4o', mcp_servers=[server])
agent = Agent('openai:gpt-4o', toolsets=[server])


async def main():
async with agent.run_mcp_servers():
async with agent.run_toolsets():
result = await agent.run('How many days between 2000-01-01 and 2025-03-18?')
print(result.output)
#> There are 9,208 days between January 1, 2000, and March 18, 2025.
Expand All @@ -180,31 +180,32 @@ call needs.
from typing import Any

from pydantic_ai import Agent
from pydantic_ai.mcp import CallToolFunc, MCPServerStdio, ToolResult
from pydantic_ai.mcp import MCPServerStdio, ToolResult
from pydantic_ai.models.test import TestModel
from pydantic_ai.tools import RunContext
from pydantic_ai.toolsets.processed import CallToolFunc


async def process_tool_call(
ctx: RunContext[int],
call_tool: CallToolFunc,
tool_name: str,
args: dict[str, Any],
name: str,
tool_args: dict[str, Any],
) -> ToolResult:
"""A tool call processor that passes along the deps."""
return await call_tool(tool_name, args, metadata={'deps': ctx.deps})
return await call_tool(name, tool_args, metadata={'deps': ctx.deps})


server = MCPServerStdio('python', ['mcp_server.py'], process_tool_call=process_tool_call)
agent = Agent(
model=TestModel(call_tools=['echo_deps']),
deps_type=int,
mcp_servers=[server]
toolsets=[server]
)


async def main():
async with agent.run_mcp_servers():
async with agent.run_toolsets():
result = await agent.run('Echo with deps set to 42', deps=42)
print(result.output)
#> {"echo_deps":{"echo":"This is an echo message","deps":42}}
Expand Down Expand Up @@ -242,7 +243,7 @@ calculator_server = MCPServerSSE(
# Both servers might have a tool named 'get_data', but they'll be exposed as:
# - 'weather_get_data'
# - 'calc_get_data'
agent = Agent('openai:gpt-4o', mcp_servers=[weather_server, calculator_server])
agent = Agent('openai:gpt-4o', toolsets=[weather_server, calculator_server])
```

### Example with Stdio Server
Expand Down Expand Up @@ -272,7 +273,7 @@ js_server = MCPServerStdio(
tool_prefix='js' # Tools will be prefixed with 'js_'
)

agent = Agent('openai:gpt-4o', mcp_servers=[python_server, js_server])
agent = Agent('openai:gpt-4o', toolsets=[python_server, js_server])
```

When the model interacts with these servers, it will see the prefixed tool names, but the prefixes will be automatically handled when making tool calls.
Expand Down Expand Up @@ -359,11 +360,11 @@ from pydantic_ai import Agent
from pydantic_ai.mcp import MCPServerStdio

server = MCPServerStdio(command='python', args=['generate_svg.py'])
agent = Agent('openai:gpt-4o', mcp_servers=[server])
agent = Agent('openai:gpt-4o', toolsets=[server])


async def main():
async with agent.run_mcp_servers():
async with agent.run_toolsets():
result = await agent.run('Create an image of a robot in a punk style.')
print(result.output)
#> Image file written to robot_punk.svg.
Expand Down
4 changes: 2 additions & 2 deletions docs/output.md
Original file line number Diff line number Diff line change
Expand Up @@ -200,8 +200,8 @@ async def hand_off_to_sql_agent(ctx: RunContext, query: str) -> list[Row]:
return output
except UnexpectedModelBehavior as e:
# Bubble up potentially retryable errors to the router agent
if (cause := e.__cause__) and hasattr(cause, 'tool_retry'):
raise ModelRetry(f'SQL agent failed: {cause.tool_retry.content}') from e
if (cause := e.__cause__) and isinstance(cause, ModelRetry):
raise ModelRetry(f'SQL agent failed: {cause.message}') from e
else:
raise

Expand Down
4 changes: 2 additions & 2 deletions mcp-run-python/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -52,11 +52,11 @@ server = MCPServerStdio('deno',
'jsr:@pydantic/mcp-run-python',
'stdio',
])
agent = Agent('claude-3-5-haiku-latest', mcp_servers=[server])
agent = Agent('claude-3-5-haiku-latest', toolsets=[server])


async def main():
async with agent.run_mcp_servers():
async with agent.run_toolsets():
result = await agent.run('How many days between 2000-01-01 and 2025-03-18?')
print(result.output)
#> There are 9,208 days between January 1, 2000, and March 18, 2025.w
Expand Down
Loading
Loading