implement pydantic-ai provider for mo.ui.chat #7636
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
| return; | ||
| } | ||
|
|
||
| if (props.isPydanticAI) { |
There was a problem hiding this comment.
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?
| delete_chat_history: (req: {}) => Promise<null>; | ||
| delete_chat_message: (req: { index: number }) => Promise<null>; | ||
| send_prompt: (req: SendMessageRequest) => Promise<string>; | ||
| send_prompt: (req: SendMessageRequest) => Promise<string | object[]>; |
There was a problem hiding this comment.
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
| // For pydantic-ai, useChat will form the data structure, | ||
| // so we set the value directly from the frontend. | ||
| if (props.isPydanticAI) { | ||
| props.setValue(message.messages); | ||
| } |
There was a problem hiding this comment.
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
There was a problem hiding this comment.
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_aiChatModel class that streams Vercel AI-compatible events - Introduced frontend-managed streaming mode that bypasses backend state management
- Extended
ChatMessagetype to includeid,metadata, and madepartsnon-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.
marimo/_ai/llm/_impl.py
Outdated
| parts: list[UIMessagePart] = [] | ||
| if message.parts: | ||
| parts = [ | ||
| cast(UIMessagePart, dataclasses.asdict(part)) | ||
| if dataclasses.is_dataclass(part) | ||
| else part | ||
| for part in message.parts | ||
| ] |
There was a problem hiding this comment.
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.
## 📝 Summary <!-- Provide a concise summary of what this pull request is addressing. If this PR fixes any issues, list them here by number (e.g., Fixes #123). --> - Adds display for tools, reasoning, file attachments (standardized) - Fixes chatConfig in the UI. Eg. Claude reasoning models don't allow you to pass in `top_p` and `temperature`. So we should allow nulls. - Minor style fixes for the chatbot, buttons are too prominent. https://github.com/user-attachments/assets/669333ac-a034-4f54-86d0-df3df0f30c32 <img width="596" height="268" alt="image" src="https://github.com/user-attachments/assets/f59a15db-39ff-4ecf-a9ac-b687478a480a" /> ## 🔍 Description of Changes <!-- Detail the specific changes made in this pull request. Explain the problem addressed and how it was resolved. If applicable, provide before and after comparisons, screenshots, or any relevant details to help reviewers understand the changes easily. --> ## 📋 Checklist - [x] I have read the [contributor guidelines](https://github.com/marimo-team/marimo/blob/main/CONTRIBUTING.md). - [ ] For large changes, or changes that affect the public API: this change was discussed or approved through an issue, on [Discord](https://marimo.io/discord?ref=pr), or the community [discussions](https://github.com/marimo-team/marimo/discussions) (Please provide a link if applicable). - [ ] I have added tests for the changes made. - [x] I have run the code and verified that it works as expected.
| self._chat_history: list[ChatMessage] = [] | ||
| self._frontend_managed = False | ||
|
|
||
| if DependencyManager.pydantic_ai.has() and isinstance( |
There was a problem hiding this comment.
DependencyManager.pydantic_ai.imported?
There was a problem hiding this comment.
thanks, we don't need to check for dependency here. I've removed it
|
|
||
| def has(self, quiet: bool = False) -> bool: | ||
| """Return True if the dependency is installed.""" | ||
| """Return True if the dependency is installed. For a more performant check, use imported().""" |
|
🚀 Development release published. You may be able to view the changes at https://marimo.app?v=0.18.5-dev154 |
📝 Summary
Also adds various fixes for chatbot
🔍 Description of Changes
📋 Checklist