fix(pydantic): support pydantic-ai v2 in restate.ext.pydantic while keeping v1#209
fix(pydantic): support pydantic-ai v2 in restate.ext.pydantic while keeping v1#209selimacerbas wants to merge 1 commit into
Conversation
…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.
|
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? |
|
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? |
|
@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 |
Thank you @selimacerbas and @gvdongen for picking this up. @gvdongen , is there an ETA if possible? |
|
@tai-bo I'll have it released by Wednesday |
Problem
restate.ext.pydanticis written against pydantic-ai v1 and crashes on import underpydantic-ai-slim==2.0.0:The
pydantic_aioptional 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)
pydantic_ai.builtin_toolsdeleted →pydantic_ai.native_tools;AbstractBuiltinTool→AbstractNativeTool,tools.BuiltinToolFunc→tools.NativeToolFunctrythe v2 names,except ImportErrorfall 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.MCPServerrenamed topydantic_ai.mcp.MCPToolset— separate classes: 1.x concrete servers subclassMCPServer, 2.0 only hasMCPToolsetMCP_TOOLSET_CLASSES);isinstanceaccepts 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_kwargsshim removed in v2; native tools now go viacapabilities=[NativeTool(tool), …]RestateAgent.runtranslatesbuiltin_tools→capabilities=[NativeTool(t) …]when theNativeToolcapability is available (1.68+), and never forwardsbuiltin_toolson v2; oldest 1.x falls back to forwardingbuiltin_tools=.The
RestateAgent(event_stream_handler=…)public API is unchanged. The handler was onlyever passed at run time (still accepted in v2) and stored on
RestateModelWrapper, neverto a pydantic-ai constructor — so
RestateModelWrapper.request_streamcontinues to readself._event_stream_handlerexactly as before. (For reference, v2'scapabilities.ProcessEventStream(handler=…)is the new way to register a handler on theagent, and its docstring confirms the run-time
event_stream_handlerparameter stillobserves live events under durable-execution integrations — so no change is needed here.)
Repro
The break is iterative (
builtin_toolsmodule →BuiltinToolFunc→MCPServer), plus aruntime-only
TypeErroron the firstagent.run()from forwardingbuiltin_tools=to abase
run()that no longer accepts it.Tests
Verified on both
pydantic-ai-slim==1.107.0and==2.0.0against this branch:import restate.ext.pydantic(clean, no deprecation warnings on v1)RestateAgent(inner, event_stream_handler=handler)construction; handler flows intoRestateModelWrapper._event_stream_handlerand theevent_stream_handlerpropertyRestateAgent.run(builtin_tools=[WebSearchTool()])translates tocapabilities=[NativeTool(...)]and forwards no unknown kwargs (validated against thereal base
run()signature)MCPServerStdioon v1,MCPToolseton v2) is wrapped asRestateMCPServerruffclean;pyrightintroduces no new errors (removes the v2-incompatibility ones).There is currently no test fixture for
restate.ext.pydanticin the repo — happy tocontribute 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.