implement pydantic-ai provider for mo.ui.chat

[Light2Dark]

📝 Summary

🔍 Description of Changes

📋 Checklist

• I have read the contributor guidelines.
• For large changes, or changes that affect the public API: this change was discussed or approved through an issue, on Discord, or the community discussions (Please provide a link if applicable).
• I have added tests for the changes made.
• I have run the code and verified that it works as expected.

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Review	Updated (UTC)
marimo-docs	Ready	Preview, Comment	Jan 5, 2026 4:30pm

[mscolnick]

is there a way to avoid passing around isPydanticAI around and avoid all the conditional logic? since we cannot assume every provider will be pydantic ai, are we just forced to keep these two branches of logic around?

[mscolnick]

maybe instead of a union type, we always return Record<string, unknown>[].

for non-pydantic AI, we can just wrap it in {type: 'text", content: "part"} and also end it with is_final

[Light2Dark]

this is the key change that forces us to pass isPydantic to the frontend, we are not storing the chat history from the backend, but instead from the frontend.

I can rename to something generic like validateFrontend

[Copilot]

Pull request overview

This pull request implements pydantic-ai provider support for mo.ui.chat, enabling streaming of pydantic-ai events to the frontend. The implementation introduces a "frontend-managed" mode where the frontend (using the AI SDK's useChat hook) manages chat state instead of the backend, which is necessary for handling structured pydantic-ai responses including reasoning blocks and tool calls.

Key changes:

• Added pydantic_ai ChatModel class that streams Vercel AI-compatible events
• Introduced frontend-managed streaming mode that bypasses backend state management
• Extended ChatMessage type to include id, metadata, and made parts non-optional with default empty list

Reviewed changes

Copilot reviewed 18 out of 19 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
marimo/_ai/llm/_impl.py	Implements new pydantic_ai ChatModel class with Vercel AI event streaming
marimo/_ai/llm/__init__.py	Exports the new pydantic_ai provider
marimo/_ai/_types.py	Adds id and metadata fields to ChatMessage, makes parts default to empty list, adds create() class method for validated message construction
marimo/_plugins/ui/_impl/chat/chat.py	Implements frontend-managed streaming mode, refactors chat message sending logic
marimo/_plugins/ui/_impl/chat/utils.py	Updates message conversion to handle new id and metadata fields
frontend/src/plugins/impl/chat/chat-ui.tsx	Implements frontend stream controller for pydantic-ai chunks, handles both frontend and backend-managed modes
frontend/src/plugins/impl/chat/ChatPlugin.tsx	Updates schemas and type definitions for new message structure
frontend/src/plugins/impl/chat/types.ts	Extends ChatMessage to inherit from UIMessage
frontend/src/components/chat/chat-utils.ts	Adds id field to message building
packages/openapi/api.yaml	Updates ChatMessage schema with id and metadata fields, changes parts default
packages/openapi/src/api.ts	Generated TypeScript types reflecting schema changes
tests/_ai/test_ai_types.py	Comprehensive tests for ChatMessage.create() and part conversion
tests/_ai/llm/test_impl.py	Tests for pydantic_ai ChatModel implementation
tests/_plugins/ui/_impl/chat/test_chat.py	Tests for frontend-managed mode and message handling
tests/_plugins/ui/_impl/chat/test_chat_delta_streaming.py	Tests for delta-based streaming accumulation
tests/_ai/test_chat_convert.py	Updated tests to reflect parts default change
examples/ai/chat/pydantic-ai-chat.py	Example demonstrating pydantic-ai integration with structured outputs and reasoning
tests/snapshots/api.txt	Updated API snapshot with pydantic_ai export

Comments suppressed due to low confidence (1)

frontend/src/components/chat/chat-utils.ts:106

• The toChatMessage function creates a ChatMessage object but does not include the metadata field from the UIMessage. The ChatMessage type now includes an optional metadata field (as seen in the schema changes), but it's not being preserved when converting from UIMessage. This could lead to loss of metadata when messages are sent to the backend. Consider adding metadata: message.metadata to the returned object.

    return {
      id: message.id,
      role: message.role,
      content: stringifyTextParts(message.parts), // This is no longer used in the backend
      parts,
    };
  }

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The code uses dataclasses.is_dataclass and dataclasses.asdict to convert parts, but pydantic-ai's UIMessagePart types are Pydantic BaseModel instances, not dataclasses. The is_dataclass check will return False for Pydantic models, so they'll be passed through as-is. However, if marimo's own ChatPart types (TextPart, FilePart, etc.) are dataclasses, they will be converted to dicts using asdict. This mixing of Pydantic models and plain dicts in the same parts list could cause type validation errors in pydantic-ai's UIMessage. Consider checking for Pydantic BaseModel instances and using model_dump instead of asdict, or ensuring all parts are properly typed for pydantic-ai.