Skip to content

fix(pydantic): support pydantic-ai v2 in restate.ext.pydantic while keeping v1#209

Open
selimacerbas wants to merge 1 commit into
restatedev:mainfrom
selimacerbas:fix/pydantic-ai-v2-compat
Open

fix(pydantic): support pydantic-ai v2 in restate.ext.pydantic while keeping v1#209
selimacerbas wants to merge 1 commit into
restatedev:mainfrom
selimacerbas:fix/pydantic-ai-v2-compat

Conversation

@selimacerbas

Copy link
Copy Markdown
Contributor

Problem

restate.ext.pydantic is written against pydantic-ai v1 and crashes on import under
pydantic-ai-slim==2.0.0:

>>> from restate.ext.pydantic import RestateAgent
ModuleNotFoundError: No module named 'pydantic_ai.builtin_tools'

The pydantic_ai optional dependency is declared open-ended (pydantic-ai-slim>=1.68.0),
so 2.0.0 resolves and then fails at import time.

What changed (v1 1.68+ AND v2 2.0.0 both supported)

v1 → v2 change Handling (kept v1-compatible)
pydantic_ai.builtin_tools deleted → pydantic_ai.native_tools; AbstractBuiltinToolAbstractNativeTool, tools.BuiltinToolFunctools.NativeToolFunc try the v2 names, except ImportError fall back to the v1 module aliased to the same local names. (v1 1.68+ already exposes the new names, so no deprecation warnings on v1.)
pydantic_ai.mcp.MCPServer renamed to pydantic_ai.mcp.MCPToolsetseparate classes: 1.x concrete servers subclass MCPServer, 2.0 only has MCPToolset Detect against the tuple of MCP base classes present in the installed version (MCP_TOOLSET_CLASSES); isinstance accepts a tuple. A single aliased name would silently stop wrapping MCP servers on whichever line it didn't match.
Agent.run(builtin_tools=…) and its **_deprecated_kwargs shim removed in v2; native tools now go via capabilities=[NativeTool(tool), …] RestateAgent.run translates builtin_toolscapabilities=[NativeTool(t) …] when the NativeTool capability is available (1.68+), and never forwards builtin_tools on v2; oldest 1.x falls back to forwarding builtin_tools=.

The RestateAgent(event_stream_handler=…) public API is unchanged. The handler was only
ever passed at run time (still accepted in v2) and stored on RestateModelWrapper, never
to a pydantic-ai constructor — so RestateModelWrapper.request_stream continues to read
self._event_stream_handler exactly as before. (For reference, v2's
capabilities.ProcessEventStream(handler=…) is the new way to register a handler on the
agent, and its docstring confirms the run-time event_stream_handler parameter still
observes live events under durable-execution integrations — so no change is needed here.)

Repro

uv add restate-sdk 'pydantic-ai-slim[openai,mcp]==2.0.0'
python -c "from restate.ext.pydantic import RestateAgent"   # before: ModuleNotFoundError

The break is iterative (builtin_tools module → BuiltinToolFuncMCPServer), plus a
runtime-only TypeError on the first agent.run() from forwarding builtin_tools= to a
base run() that no longer accepts it.

Tests

Verified on both pydantic-ai-slim==1.107.0 and ==2.0.0 against this branch:

  • import restate.ext.pydantic (clean, no deprecation warnings on v1)
  • RestateAgent(inner, event_stream_handler=handler) construction; handler flows into
    RestateModelWrapper._event_stream_handler and the event_stream_handler property
  • RestateAgent.run(builtin_tools=[WebSearchTool()]) translates to
    capabilities=[NativeTool(...)] and forwards no unknown kwargs (validated against the
    real base run() signature)
  • A real MCP server (MCPServerStdio on v1, MCPToolset on v2) is wrapped as
    RestateMCPServer

ruff clean; pyright introduces no new errors (removes the v2-incompatibility ones).

There is currently no test fixture for restate.ext.pydantic in the repo — happy to
contribute the two smoke tests above + a pydantic-ai 2.x CI matrix entry as a follow-up if
that's useful.

Follow-up to #201.

…estatedev#201)

restate.ext.pydantic was written against pydantic-ai v1 and crashes on import
under pydantic-ai-slim 2.0.0:

    from restate.ext.pydantic import RestateAgent
    -> ModuleNotFoundError: No module named 'pydantic_ai.builtin_tools'

The v1 -> v2 changes that break it, and how each is handled (all kept v1-compatible):

  * pydantic_ai.builtin_tools deleted -> pydantic_ai.native_tools;
    AbstractBuiltinTool -> AbstractNativeTool, tools.BuiltinToolFunc ->
    tools.NativeToolFunc. Import the v2 names first, fall back to the v1 module.

  * pydantic_ai.mcp.MCPServer renamed to pydantic_ai.mcp.MCPToolset. The two are
    *separate* classes (1.x concrete servers subclass MCPServer; 2.0 only has
    MCPToolset), so detection matches against the tuple of MCP base classes
    present in the installed version rather than a single aliased name.

  * The run-time builtin_tools= argument (and its **_deprecated_kwargs shim) was
    removed from Agent.run in v2; native tools now go via
    capabilities=[NativeTool(tool), ...]. RestateAgent.run translates
    builtin_tools -> capabilities when the NativeTool capability is available
    (1.68+), falling back to forwarding builtin_tools= on the oldest 1.x.

The RestateAgent(event_stream_handler=...) public API is unchanged: the handler
was only ever passed at run time (still accepted in v2) and stored on
RestateModelWrapper, never to a pydantic-ai constructor.

Verified: import + RestateAgent construction + event_stream_handler flow +
run() kwarg translation + MCP-toolset wrapping all pass on both 1.107.0 and
2.0.0. ruff clean; pyright introduces no new errors.
@gvdongen

gvdongen commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

Hi! Thank you for this PR. I think we don't want to keep supporting compatibility with older versions of Pydantic AI. So I'll migrate the SDK to Pydantic AI 2.0 and drop support for 1.x from the next Restate SDK release onwards. Hope this works for you?

@selimacerbas

selimacerbas commented Jul 14, 2026

Copy link
Copy Markdown
Contributor Author

Hey @gvdongen, thanks for your answer, of course it works for me, do you want me complete this PR removing v1 support, along with tests etc?

@gvdongen

Copy link
Copy Markdown
Contributor

@selimacerbas I will take it from here because there is anyway some other stuff I wanted to do. But thanks a lot getting things started! I'll let you know once it's out

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants