fix(pydantic): support pydantic-ai v2 in restate.ext.pydantic while keeping v1#209
Open
selimacerbas wants to merge 1 commit into
Open
fix(pydantic): support pydantic-ai v2 in restate.ext.pydantic while keeping v1#209selimacerbas wants to merge 1 commit into
selimacerbas 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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.