Rework PydanticAI adapter: upstream parity + granular checkpoints#156
Draft
Rework PydanticAI adapter: upstream parity + granular checkpoints#156
Conversation
Contributor
Site Preview
|
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.
Summary
Reworks the PydanticAI adapter to mirror the upstream
pydantic_ai.durable_exec.kitarustructure and adds Kitaru-specific durability controls on top._agent.py,_model.py,_toolset.py,_function_toolset.py,_mcp_server.py,_hitl.py,_policy.py,_events.py,_tracking.py,_otel.py,_kitaru_internal.py,_utils.py,_run_context.py. A narrow ruff ignore set(
UP006/007/008/035/045,E501,I001,SIM105) keeps vendored style stable against upstream.KitaruAgent(granular_checkpoints=True, model_checkpoint_config=..., tool_checkpoint_config=..., tool_checkpoint_config_by_name=..., mcp_checkpoint_config=...)opens a checkpoint per model / tool / MCP call instead of one per turn. Per-tool overrides acceptFalseto opt a tool out entirely. Turn mode (default) keeps the aggregating run artifact.turn_checkpoint_configkwarg so users can tune retries / type on the auto-opened turn checkpoint.persist_message_history=True. Opt-in conversation memory — the adapter remembersresult.all_messages()on the instance and auto-injects it asmessage_historyon the next call when the caller doesn't pass one. One instance = one conversation.event_stream_handleris supplied, granular mode transparently falls back to a turn checkpoint — per-call checkpointing can't drainan async stream inside a sync ZenML step.
cache_toolshonored to skip redundanttools/listround-trips on checkpoint re-entry; MCP calls are taggedtoolset_kind='mcp'and default totype='mcp_call'in the dashboard.requires_approval=True,ApprovalRequired, andCallDeferredall route throughkitaru.wait()with no decorator;hitl_tool(...)remains the explicit marker.opentelemetryis installed; no-ops otherwise.PYDANTIC_AI_WRAPPEDreportsgranular_checkpointsandpersist_message_history; run-completion events carry method + error class.pydantic-ai-slimsurfaces asKitaruFeatureNotAvailableErrorinstead of anImportErrordeep in the module graph.src/kitaru/adapters/pydantic_ai/README.mddocuments concepts, install, usage patterns, checkpoint modes, streaming, capture policy,conversation memory, and troubleshooting — shaped for later upstreaming into PydanticAI docs.
Known limits and follow-ups
Kitaru-core (not overcomable in this adapter):
Adapter follow-up (tracked, not in this PR):
runtime='isolated'— currently rejected withKitaruUsageError. Lifting the limit requires making every adapter wrapper (KitaruModel,KitaruToolset,KitaruMCPServer) reconstructible from serializable construction args on the far side of the process boundary and wiringKitaruRunContextthrough the granulardispatcher. The pydantic payloads already serialize via
TypeAdapter, so most of the work is on the wrapper-identity side.By design:
@kitaru.flow— serializing an arbitrary agent closure isn't worth the machinery when a one-line decoratordoes the job.
Test plan
just checkpasses (format, lint, typecheck, typos, yaml, actionlint, links)just testpasses — includingtests/test_pydantic_ai_adapter.pycoveringreject_isolated_runtime, per-toolFalseopt-out,with_default_typeimmutability, the
_use_granularstreaming fallback, andpersist_message_historyload/store semanticsexamples/pydantic_ai_agent/pydantic_ai_adapter.pyruns end-to-end with and withoutgranular_checkpoints=Truerun_sync()calls on apersist_message_history=Trueinstance see each other's turns without the caller threadingmessage_historyruntime='isolated'in any checkpoint config raisesKitaruUsageErrorwith the guidance messagecache_tools=Trueskipstools/liston checkpoint re-entryrequires_approval=True,ApprovalRequired,CallDeferred, andhitl_toolall suspend and resume throughkitaru.wait()event_stream_handlerwithgranular_checkpoints=Truerecords a turn checkpoint (fallback);run_stream()/iter()require explicit@kitaru.checkpointin both modes