From c9add013ab6dbf01b96114609bfc3d8726bd58af Mon Sep 17 00:00:00 2001 From: Shahmir Varqha Date: Wed, 25 Mar 2026 16:32:36 +0800 Subject: [PATCH 1/8] upgrade package --- frontend/package.json | 4 +- frontend/src/components/chat/chat-utils.ts | 15 ++--- frontend/src/core/ai/staged-cells.ts | 6 ++ marimo/_plugins/ui/_impl/chat/chat.py | 2 +- pnpm-lock.yaml | 76 ++++++++++------------ 5 files changed, 52 insertions(+), 51 deletions(-) diff --git a/frontend/package.json b/frontend/package.json index f00df7fec9f..fdce7083d84 100644 --- a/frontend/package.json +++ b/frontend/package.json @@ -21,7 +21,7 @@ "./unstable_internal/*": "./src/*" }, "dependencies": { - "@ai-sdk/react": "^2.0.125", + "@ai-sdk/react": "^3.0.131", "@anywidget/types": "^0.2.0", "@codemirror/autocomplete": "^6.20.1", "@codemirror/commands": "^6.10.2", @@ -114,7 +114,7 @@ "@xterm/addon-web-links": "^0.12.0", "@xterm/xterm": "^5.5.0", "@zed-industries/agent-client-protocol": "^0.4.5", - "ai": "^5.0.123", + "ai": "^6.0.129", "ansi_up": "^6.0.6", "class-variance-authority": "^0.7.1", "clsx": "^2.1.1", diff --git a/frontend/src/components/chat/chat-utils.ts b/frontend/src/components/chat/chat-utils.ts index eeb15ace196..2d2b1cb9bd0 100644 --- a/frontend/src/components/chat/chat-utils.ts +++ b/frontend/src/components/chat/chat-utils.ts @@ -1,7 +1,12 @@ /* Copyright 2026 Marimo. All rights reserved. */ import type { components } from "@marimo-team/marimo-api"; -import type { FileUIPart, ToolUIPart, UIMessage } from "ai"; +import type { + ChatAddToolOutputFunction, + FileUIPart, + ToolUIPart, + UIMessage, +} from "ai"; import { useState } from "react"; import useEvent from "react-use-event-hook"; import type { ProviderId } from "@/core/ai/ids/ids"; @@ -111,12 +116,6 @@ export async function buildCompletionRequestBody( }; } -interface AddToolOutput { - tool: string; - toolCallId: string; - output: unknown; -} - export async function handleToolCall({ invokeAiTool, addToolOutput, // Important that we don't await addToolOutput to prevent potential deadlocks @@ -124,7 +123,7 @@ export async function handleToolCall({ toolContext, }: { invokeAiTool: (request: InvokeAiToolRequest) => Promise; - addToolOutput: (output: AddToolOutput) => Promise; + addToolOutput: ChatAddToolOutputFunction; toolCall: { toolName: string; toolCallId: string; diff --git a/frontend/src/core/ai/staged-cells.ts b/frontend/src/core/ai/staged-cells.ts index d16e785ef76..bdf21d43ffb 100644 --- a/frontend/src/core/ai/staged-cells.ts +++ b/frontend/src/core/ai/staged-cells.ts @@ -161,6 +161,12 @@ export function useStagedCells(store: JotaiStore) { case "tool-output-error": Logger.error("Error", chunk.type, { chunk }); break; + case "tool-approval-request": + Logger.log("Tool approval request", { chunk }); + break; + case "tool-output-denied": + Logger.error("Tool output denied", { chunk }); + break; // These logs are not useful for debugging case "start": case "start-step": diff --git a/marimo/_plugins/ui/_impl/chat/chat.py b/marimo/_plugins/ui/_impl/chat/chat.py index e7a58cbcb21..2209df58749 100644 --- a/marimo/_plugins/ui/_impl/chat/chat.py +++ b/marimo/_plugins/ui/_impl/chat/chat.py @@ -36,7 +36,7 @@ ) # The version of the Vercel AI SDK we use -AI_SDK_VERSION: Final[Literal[5, 6]] = 5 +AI_SDK_VERSION: Final[Literal[5, 6]] = 6 DONE_CHUNK: Final[str] = "[DONE]" diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index a4a6c68e63b..8c03c880930 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -38,8 +38,8 @@ importers: frontend: dependencies: '@ai-sdk/react': - specifier: ^2.0.125 - version: 2.0.126(react@19.2.4)(zod@4.3.6) + specifier: ^3.0.131 + version: 3.0.131(react@19.2.4)(zod@4.3.6) '@anywidget/types': specifier: ^0.2.0 version: 0.2.0 @@ -317,8 +317,8 @@ importers: specifier: ^0.4.5 version: 0.4.5 ai: - specifier: ^5.0.123 - version: 5.0.124(zod@4.3.6) + specifier: ^6.0.129 + version: 6.0.129(zod@4.3.6) ansi_up: specifier: ^6.0.6 version: 6.0.6 @@ -802,31 +802,27 @@ packages: '@agentclientprotocol/sdk@0.4.9': resolution: {integrity: sha512-ExwH828LaTGoTTjxuw49l+fwOLA+Yx0+qkWn1TcHMOsY5mVI9CUfkj7ZDhv2klgZ7mJeT+lxX/Dn/KINv1AkNQ==} - '@ai-sdk/gateway@2.0.30': - resolution: {integrity: sha512-5Nrkj8B4MzkkOfjjA+Cs5pamkbkK4lI11bx80QV7TFcen/hWA8wEC+UVzwuM5H2zpekoNMjvl6GonHnR62XIZw==} + '@ai-sdk/gateway@3.0.75': + resolution: {integrity: sha512-7Hwa0VdH+l85NFS7zqZhRRaiwZMStDxEwUoTPxPNEH6V0Vgw9wi9OGopIsYdywmfSOPfSAsPL8XXPAuaSLGchw==} engines: {node: '>=18'} peerDependencies: zod: ^3.25.76 || ^4.1.8 - '@ai-sdk/provider-utils@3.0.20': - resolution: {integrity: sha512-iXHVe0apM2zUEzauqJwqmpC37A5rihrStAih5Ks+JE32iTe4LZ58y17UGBjpQQTCRw9YxMeo2UFLxLpBluyvLQ==} + '@ai-sdk/provider-utils@4.0.20': + resolution: {integrity: sha512-gpUIj9uDhIGbuo9afKEgQ074BWmhvK4THJAAeBjRnroTy2yQYo6rbtGD7pQDMZM8ouXPYmT/SCdkWVJ0KcpX8A==} engines: {node: '>=18'} peerDependencies: zod: ^3.25.76 || ^4.1.8 - '@ai-sdk/provider@2.0.1': - resolution: {integrity: sha512-KCUwswvsC5VsW2PWFqF8eJgSCu5Ysj7m1TxiHTVA6g7k360bk0RNQENT8KTMAYEs+8fWPD3Uu4dEmzGHc+jGng==} + '@ai-sdk/provider@3.0.8': + resolution: {integrity: sha512-oGMAgGoQdBXbZqNG0Ze56CHjDZ1IDYOwGYxYjO5KLSlz5HiNQ9udIXsPZ61VWaHGZ5XW/jyjmr6t2xz2jGVwbQ==} engines: {node: '>=18'} - '@ai-sdk/react@2.0.126': - resolution: {integrity: sha512-fvFphZbGZxOsxiUU9JKMvvs05ZHiyFCIKgMnup3TQJzx2TgvFDebh2wbCZSxsp4jPytucJ85c4Obt3ejBeQyeQ==} + '@ai-sdk/react@3.0.131': + resolution: {integrity: sha512-75CqZ84ELE2/xOYpNs3A0Ioxa8MVa3yJCp3n8Ru9dWEbwv3tL7/+5wZlcHI9bkoEAw1KeNbhBL80eD861mtFHw==} engines: {node: '>=18'} peerDependencies: react: ^19.2.4 - zod: ^3.25.76 || ^4.1.8 - peerDependenciesMeta: - zod: - optional: true '@alloc/quick-lru@5.2.0': resolution: {integrity: sha512-UrcABB+4bUrFABwbluTIBErXwvbsU/V7TZWfmbgJfbkwiBuziS9gxdODUyuiecfdGQ85jglMW6juS3+z5TsKLw==} @@ -3630,8 +3626,8 @@ packages: '@sourcegraph/vscode-ws-jsonrpc@0.0.3-fork': resolution: {integrity: sha512-EJLq/ni66glk3xYyOZtUIEbjTCw8kMI6RvO0YQtPd+4um2+aTSM1LfN4NrsiVrRkG7EG/U2OkFlKqT8mGo6w4Q==} - '@standard-schema/spec@1.0.0': - resolution: {integrity: sha512-m2bOd0f2RT9k8QJx1JN85cZYyH1RqFBdlwtkSlf4tBDYLCiiZnv1fIIwacK6cqwXavOydf0NPToMQgpKq+dVlA==} + '@standard-schema/spec@1.1.0': + resolution: {integrity: sha512-l2aFy5jALhniG5HgqrD6jXLi/rUWrKvqN/qJx6yoJsgKhblVd+iqqU4RCXavm/jPityDo5TCvKMnpjKnOriy0w==} '@standard-schema/utils@0.3.0': resolution: {integrity: sha512-e7Mew686owMaPJVNNLs55PUvgz371nKgwsc4vxE49zsODpJEnxgxRo2y/OKrqueavXgZNMDVj3DdHFlaSAeU8g==} @@ -4575,8 +4571,8 @@ packages: resolution: {integrity: sha512-MnA+YT8fwfJPgBx3m60MNqakm30XOkyIoH1y6huTQvC0PwZG7ki8NacLBcrPbNoo8vEZy7Jpuk7+jMO+CUovTQ==} engines: {node: '>= 14'} - ai@5.0.124: - resolution: {integrity: sha512-Li6Jw9F9qsvFJXZPBfxj38ddP2iURCnMs96f9Q3OeQzrDVcl1hvtwSEAuxA/qmfh6SDV2ERqFUOFzigvr0697g==} + ai@6.0.129: + resolution: {integrity: sha512-5nGckqbzwUBZD7wV9jsA8qaoYRwGpU9LVMtXD+ZrxSi2H6QNjpbrhsuuEBKS9xcMYevCviVNoFzpmSUWzn45Hw==} engines: {node: '>=18'} peerDependencies: zod: ^3.25.76 || ^4.1.8 @@ -10353,33 +10349,33 @@ snapshots: dependencies: zod: 3.25.76 - '@ai-sdk/gateway@2.0.30(zod@4.3.6)': + '@ai-sdk/gateway@3.0.75(zod@4.3.6)': dependencies: - '@ai-sdk/provider': 2.0.1 - '@ai-sdk/provider-utils': 3.0.20(zod@4.3.6) + '@ai-sdk/provider': 3.0.8 + '@ai-sdk/provider-utils': 4.0.20(zod@4.3.6) '@vercel/oidc': 3.1.0 zod: 4.3.6 - '@ai-sdk/provider-utils@3.0.20(zod@4.3.6)': + '@ai-sdk/provider-utils@4.0.20(zod@4.3.6)': dependencies: - '@ai-sdk/provider': 2.0.1 - '@standard-schema/spec': 1.0.0 + '@ai-sdk/provider': 3.0.8 + '@standard-schema/spec': 1.1.0 eventsource-parser: 3.0.6 zod: 4.3.6 - '@ai-sdk/provider@2.0.1': + '@ai-sdk/provider@3.0.8': dependencies: json-schema: 0.4.0 - '@ai-sdk/react@2.0.126(react@19.2.4)(zod@4.3.6)': + '@ai-sdk/react@3.0.131(react@19.2.4)(zod@4.3.6)': dependencies: - '@ai-sdk/provider-utils': 3.0.20(zod@4.3.6) - ai: 5.0.124(zod@4.3.6) + '@ai-sdk/provider-utils': 4.0.20(zod@4.3.6) + ai: 6.0.129(zod@4.3.6) react: 19.2.4 swr: 2.3.4(react@19.2.4) throttleit: 2.1.0 - optionalDependencies: - zod: 4.3.6 + transitivePeerDependencies: + - zod '@alloc/quick-lru@5.2.0': {} @@ -13993,7 +13989,7 @@ snapshots: dependencies: vscode-jsonrpc: 4.0.0 - '@standard-schema/spec@1.0.0': {} + '@standard-schema/spec@1.1.0': {} '@standard-schema/utils@0.3.0': {} @@ -14899,14 +14895,14 @@ snapshots: chai: 5.3.3 tinyrainbow: 2.0.0 - '@vitest/mocker@3.2.4(msw@2.12.7(@types/node@20.19.30)(typescript@5.9.3))(rolldown-vite@7.3.1(@types/node@24.11.0)(esbuild@0.27.3)(jiti@2.6.1)(terser@5.44.1)(yaml@2.8.2))': + '@vitest/mocker@3.2.4(msw@2.12.7(@types/node@20.19.30)(typescript@5.9.3))(rolldown-vite@7.3.1(@types/node@20.19.30)(esbuild@0.27.3)(jiti@2.6.1)(terser@5.44.1)(yaml@2.8.2))': dependencies: '@vitest/spy': 3.2.4 estree-walker: 3.0.3 magic-string: 0.30.21 optionalDependencies: msw: 2.12.7(@types/node@20.19.30)(typescript@5.9.3) - vite: rolldown-vite@7.3.1(@types/node@24.11.0)(esbuild@0.27.3)(jiti@2.6.1)(terser@5.44.1)(yaml@2.8.2) + vite: rolldown-vite@7.3.1(@types/node@20.19.30)(esbuild@0.27.3)(jiti@2.6.1)(terser@5.44.1)(yaml@2.8.2) '@vitest/mocker@3.2.4(msw@2.12.7(@types/node@24.11.0)(typescript@5.9.3))(rolldown-vite@7.3.1(@types/node@24.11.0)(esbuild@0.27.3)(jiti@2.6.1)(terser@5.44.1)(yaml@2.8.2))': dependencies: @@ -15093,11 +15089,11 @@ snapshots: agent-base@7.1.4: {} - ai@5.0.124(zod@4.3.6): + ai@6.0.129(zod@4.3.6): dependencies: - '@ai-sdk/gateway': 2.0.30(zod@4.3.6) - '@ai-sdk/provider': 2.0.1 - '@ai-sdk/provider-utils': 3.0.20(zod@4.3.6) + '@ai-sdk/gateway': 3.0.75(zod@4.3.6) + '@ai-sdk/provider': 3.0.8 + '@ai-sdk/provider-utils': 4.0.20(zod@4.3.6) '@opentelemetry/api': 1.9.0 zod: 4.3.6 @@ -21731,7 +21727,7 @@ snapshots: dependencies: '@types/chai': 5.2.3 '@vitest/expect': 3.2.4 - '@vitest/mocker': 3.2.4(msw@2.12.7(@types/node@20.19.30)(typescript@5.9.3))(rolldown-vite@7.3.1(@types/node@24.11.0)(esbuild@0.27.3)(jiti@2.6.1)(terser@5.44.1)(yaml@2.8.2)) + '@vitest/mocker': 3.2.4(msw@2.12.7(@types/node@20.19.30)(typescript@5.9.3))(rolldown-vite@7.3.1(@types/node@20.19.30)(esbuild@0.27.3)(jiti@2.6.1)(terser@5.44.1)(yaml@2.8.2)) '@vitest/pretty-format': 3.2.4 '@vitest/runner': 3.2.4 '@vitest/snapshot': 3.2.4 From 57e1544ea17bbee16e729ca0694951519ff7afad Mon Sep 17 00:00:00 2001 From: Shahmir Varqha Date: Wed, 25 Mar 2026 21:11:27 +0800 Subject: [PATCH 2/8] pnpm dedupe --- pnpm-lock.yaml | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 8c03c880930..55c97b9b6e1 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -8743,6 +8743,7 @@ packages: rolldown-vite@7.3.1: resolution: {integrity: sha512-LYzdNAjRHhF2yA4JUQm/QyARyi216N2rpJ0lJZb8E9FU2y5v6Vk+xq/U4XBOxMefpWixT5H3TslmAHm1rqIq2w==} engines: {node: ^20.19.0 || >=22.12.0} + deprecated: Use this package to migrate from Vite 7 to Vite 8. For the most recent updates, migrate to Vite 8 once you're ready. hasBin: true peerDependencies: '@types/node': ^20.19.0 || >=22.12.0 @@ -14895,14 +14896,14 @@ snapshots: chai: 5.3.3 tinyrainbow: 2.0.0 - '@vitest/mocker@3.2.4(msw@2.12.7(@types/node@20.19.30)(typescript@5.9.3))(rolldown-vite@7.3.1(@types/node@20.19.30)(esbuild@0.27.3)(jiti@2.6.1)(terser@5.44.1)(yaml@2.8.2))': + '@vitest/mocker@3.2.4(msw@2.12.7(@types/node@20.19.30)(typescript@5.9.3))(rolldown-vite@7.3.1(@types/node@24.11.0)(esbuild@0.27.3)(jiti@2.6.1)(terser@5.44.1)(yaml@2.8.2))': dependencies: '@vitest/spy': 3.2.4 estree-walker: 3.0.3 magic-string: 0.30.21 optionalDependencies: msw: 2.12.7(@types/node@20.19.30)(typescript@5.9.3) - vite: rolldown-vite@7.3.1(@types/node@20.19.30)(esbuild@0.27.3)(jiti@2.6.1)(terser@5.44.1)(yaml@2.8.2) + vite: rolldown-vite@7.3.1(@types/node@24.11.0)(esbuild@0.27.3)(jiti@2.6.1)(terser@5.44.1)(yaml@2.8.2) '@vitest/mocker@3.2.4(msw@2.12.7(@types/node@24.11.0)(typescript@5.9.3))(rolldown-vite@7.3.1(@types/node@24.11.0)(esbuild@0.27.3)(jiti@2.6.1)(terser@5.44.1)(yaml@2.8.2))': dependencies: @@ -21727,7 +21728,7 @@ snapshots: dependencies: '@types/chai': 5.2.3 '@vitest/expect': 3.2.4 - '@vitest/mocker': 3.2.4(msw@2.12.7(@types/node@20.19.30)(typescript@5.9.3))(rolldown-vite@7.3.1(@types/node@20.19.30)(esbuild@0.27.3)(jiti@2.6.1)(terser@5.44.1)(yaml@2.8.2)) + '@vitest/mocker': 3.2.4(msw@2.12.7(@types/node@20.19.30)(typescript@5.9.3))(rolldown-vite@7.3.1(@types/node@24.11.0)(esbuild@0.27.3)(jiti@2.6.1)(terser@5.44.1)(yaml@2.8.2)) '@vitest/pretty-format': 3.2.4 '@vitest/runner': 3.2.4 '@vitest/snapshot': 3.2.4 From 12f4ba199c4b879abb280bc0ec0d8bfd6fbe2867 Mon Sep 17 00:00:00 2001 From: Shahmir Varqha Date: Wed, 25 Mar 2026 21:59:28 +0800 Subject: [PATCH 3/8] fix test --- tests/_plugins/ui/_impl/chat/test_chat.py | 45 ----------------------- 1 file changed, 45 deletions(-) diff --git a/tests/_plugins/ui/_impl/chat/test_chat.py b/tests/_plugins/ui/_impl/chat/test_chat.py index 797cebd7357..622c66f08be 100644 --- a/tests/_plugins/ui/_impl/chat/test_chat.py +++ b/tests/_plugins/ui/_impl/chat/test_chat.py @@ -978,51 +978,6 @@ def on_send_chunk(chunk: dict): ] -@pytest.mark.skipif( - not DependencyManager.pydantic_ai.has(), - reason="Pydantic AI is not installed", -) -def test_serialize_pydantic_v5(): - """Test ChunkSerializer excludes providerMetadata from tool-input-start for SDK v5. - - The Vercel AI SDK v5 schema drifts from v6, so we need to use Pydantic's handling. - - Since pydantic-ai uses toolCallId, providerMetadata must be excluded. - See: https://github.com/pydantic/pydantic-ai/pull/4166 - """ - from pydantic_ai.ui.vercel_ai.response_types import ToolInputStartChunk - - sent_chunks: list[dict] = [] - - def on_send_chunk(chunk: dict): - sent_chunks.append(chunk) - - serializer = ChunkSerializer(on_send_chunk=on_send_chunk) - - # Create chunk with providerMetadata (like Google Gemini produces) - chunk = ToolInputStartChunk( - tool_call_id="tc_1", - tool_name="my_tool", - provider_metadata={ - "pydantic_ai": { - "id": "test_id", - "provider_name": "google-gla", - "provider_details": {"thought_signature": "encrypted_data"}, - } - }, - ) - serializer.handle_chunk(chunk) - - # providerMetadata should be excluded for SDK v5 compatibility - assert sent_chunks == [ - { - "type": "tool-input-start", - "toolCallId": "tc_1", - "toolName": "my_tool", - } - ] - - @pytest.mark.skipif( not DependencyManager.pydantic_ai.has(), reason="Pydantic AI is not installed", From fdbb37434129753bd2ceb91ddf9b3661d7aeb496 Mon Sep 17 00:00:00 2001 From: Shahmir Varqha Date: Wed, 25 Mar 2026 22:16:08 +0800 Subject: [PATCH 4/8] upgrade deps --- examples/ai/chat/pydantic-ai-chat.py | 10 +++++++++- pyproject.toml | 6 +++--- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/examples/ai/chat/pydantic-ai-chat.py b/examples/ai/chat/pydantic-ai-chat.py index 5c41f6b5e85..b57b615faad 100644 --- a/examples/ai/chat/pydantic-ai-chat.py +++ b/examples/ai/chat/pydantic-ai-chat.py @@ -1,6 +1,14 @@ +# /// script +# requires-python = ">=3.13" +# dependencies = [ +# "httpx==0.28.1", +# "marimo>=0.21.1", +# "pydantic==2.12.5", +# ] +# /// import marimo -__generated_with = "0.19.2" +__generated_with = "0.21.1" app = marimo.App(width="medium") with app.setup(hide_code=True): diff --git a/pyproject.toml b/pyproject.toml index d90dc62f5c0..8c4518dca12 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -131,7 +131,7 @@ dev = [ # For linting "ruff>=0.14.0,<0.15.2", # TODO: remove upper bound once we fix the ruff issues # For AI - "pydantic-ai-slim[openai]>=1.52.0", + "pydantic-ai-slim[openai]>=1.71.0", ] test = [ @@ -191,7 +191,7 @@ test-optional = [ "anywidget~=0.9.21", "ipython~=8.12.3", # testing gen ai - "pydantic-ai-slim[google,anthropic,bedrock,openai]>=1.52.0", + "pydantic-ai-slim[google,anthropic,bedrock,openai]>=1.71.0", # - google-auth uses cachetools, and cachetools<5.0.0 uses collections.MutableMapping (removed in Python 3.10) "cachetools>=5.0.0", "boto3>=1.38.46", @@ -228,7 +228,7 @@ typecheck = [ "sqlalchemy>=2.0.40", "obstore>=0.8.2", "fsspec>=2026.2.0", - "pydantic-ai-slim[google,anthropic,bedrock,openai]>=1.52.0", + "pydantic-ai-slim[google,anthropic,bedrock,openai]>=1.71.0", "loro>=1.5.0", "boto3-stubs>=1.38.46", "pandas-stubs>=1.5.3.230321", From 95a6667fd0abd051e025514a69417d39b9693051 Mon Sep 17 00:00:00 2001 From: Shahmir Varqha Date: Wed, 25 Mar 2026 23:17:21 +0800 Subject: [PATCH 5/8] upgrade main dep to 1.52.0 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 8c4518dca12..365dfafe01c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -101,7 +101,7 @@ recommended = [ "marimo[sql]", "marimo[sandbox]", # For `marimo edit --sandbox DIRECTORY` "altair>=5.4.0", # Plotting in datasource viewer - "pydantic-ai-slim[openai]>=1.39.0", # AI features + "pydantic-ai-slim[openai]>=1.52.0", # AI features "ruff", # Formatting "nbformat>=5.7.0", # Export as IPYNB ] From 9e97aeeeae810f25a089a863247d857b79edac3d Mon Sep 17 00:00:00 2001 From: Shahmir Varqha Date: Wed, 25 Mar 2026 23:30:16 +0800 Subject: [PATCH 6/8] enforce min 1.52.0 pydantic-ai --- marimo/_server/ai/providers.py | 103 +++++---------------------------- 1 file changed, 13 insertions(+), 90 deletions(-) diff --git a/marimo/_server/ai/providers.py b/marimo/_server/ai/providers.py index 3cad9597ff2..6294d4a6010 100644 --- a/marimo/_server/ai/providers.py +++ b/marimo/_server/ai/providers.py @@ -40,7 +40,6 @@ from openai import AsyncOpenAI from openai.types.shared.reasoning_effort import ReasoningEffort from pydantic_ai import Agent, DeferredToolRequests, FunctionToolset - from pydantic_ai.messages import ThinkingPart from pydantic_ai.models import Model from pydantic_ai.models.bedrock import BedrockConverseModel from pydantic_ai.models.google import GoogleModel @@ -54,9 +53,7 @@ ) from pydantic_ai.providers.google import GoogleProvider as PydanticGoogle from pydantic_ai.providers.openai import OpenAIProvider as PydanticOpenAI - from pydantic_ai.ui.vercel_ai import VercelAIAdapter from pydantic_ai.ui.vercel_ai.request_types import UIMessage, UIMessagePart - from pydantic_ai.ui.vercel_ai.response_types import BaseChunk from starlette.responses import StreamingResponse @@ -102,6 +99,10 @@ def __init__( source="server", ) + DependencyManager.pydantic_ai.require_at_version( + "for Vercel AI SDK support", min_version="1.52.0" + ) + self.model = model self.config = config self.provider = self.create_provider(config) @@ -132,12 +133,6 @@ def create_agent( output_type=output_type, ) - def get_vercel_adapter(self) -> type[VercelAIAdapter[Any, Any]]: - """Return the Vercel AI adapter for the given provider.""" - from pydantic_ai.ui.vercel_ai import VercelAIAdapter - - return VercelAIAdapter - def convert_messages( self, messages: list[ServerUIMessage] ) -> list[UIMessage]: @@ -153,6 +148,7 @@ async def stream_completion( stream_options: Optional[StreamOptions] = None, ) -> StreamingResponse: """Return a streaming response from the given messages. The response are AI SDK events.""" + from pydantic_ai.ui.vercel_ai import VercelAIAdapter from pydantic_ai.ui.vercel_ai.request_types import SubmitMessage tools = (self.config.tools or []) + additional_tools @@ -169,18 +165,12 @@ async def stream_completion( # TODO: Text only and format stream are not supported yet stream_options = stream_options or StreamOptions() - vercel_adapter = self.get_vercel_adapter() - if DependencyManager.pydantic_ai.has_at_version(min_version="1.52.0"): - adapter = vercel_adapter( - agent=agent, - run_input=run_input, - accept=stream_options.accept, - sdk_version=AI_SDK_VERSION, - ) - else: - adapter = vercel_adapter( - agent=agent, run_input=run_input, accept=stream_options.accept - ) + adapter = VercelAIAdapter( + agent=agent, + run_input=run_input, + accept=stream_options.accept, + sdk_version=AI_SDK_VERSION, + ) event_stream = adapter.run_stream() return adapter.streaming_response(event_stream) @@ -193,16 +183,16 @@ async def stream_text( additional_tools: list[ToolDefinition], ) -> AsyncGenerator[str]: """Return a stream of text from the given messages.""" + from pydantic_ai.ui.vercel_ai import VercelAIAdapter tools = (self.config.tools or []) + additional_tools agent = self.create_agent( max_tokens=max_tokens, tools=tools, system_prompt=system_prompt ) - vercel_adapter = self.get_vercel_adapter() async with agent.run_stream( user_prompt=user_prompt, - message_history=vercel_adapter.load_messages( + message_history=VercelAIAdapter.load_messages( self.convert_messages(messages) ), ) as result: @@ -800,73 +790,6 @@ def process_part(self, part: UIMessagePart) -> UIMessagePart: ) return part - def get_vercel_adapter( - self, - ) -> type[VercelAIAdapter[None, DeferredToolRequests | str]]: - """ - Return a custom adapter that includes thinking signatures in ReasoningEndChunk. - - pydantic_ai's VercelAIEventStream.handle_thinking_end doesn't pass the signature - from ThinkingPart to ReasoningEndChunk, which breaks Anthropic's extended thinking - on follow-up messages (Anthropic requires signatures on thinking blocks). - - This is a patch for pydantic-ai <1.47.0, which doesn't include the signature in the ReasoningEndChunk. - """ - if DependencyManager.pydantic_ai.has_at_version(min_version="1.47.0"): - return super().get_vercel_adapter() - - from pydantic_ai import DeferredToolRequests - from pydantic_ai.ui.vercel_ai import VercelAIAdapter - from pydantic_ai.ui.vercel_ai._event_stream import VercelAIEventStream - from pydantic_ai.ui.vercel_ai.response_types import ReasoningEndChunk - - AnthropicOutputType = DeferredToolRequests | str - - # Custom event stream that includes signature in ReasoningEndChunk - class AnthropicVercelAIEventStream( - VercelAIEventStream[None, AnthropicOutputType] - ): - async def handle_thinking_end( - self, part: ThinkingPart, followed_by_thinking: bool = False - ) -> AsyncIterator[BaseChunk]: - """Override to include signature in provider_metadata.""" - try: - provider_metadata = None - if part.signature: - pydantic_ai_meta: dict[str, Any] = { - "signature": part.signature - } - if part.provider_name: - pydantic_ai_meta["provider_name"] = ( - part.provider_name - ) - if part.id: - pydantic_ai_meta["id"] = part.id - provider_metadata = {"pydantic_ai": pydantic_ai_meta} - - yield ReasoningEndChunk( - id=self.message_id, provider_metadata=provider_metadata - ) - except Exception as e: - LOGGER.warning( - f"Error in AnthropicVercelAIEventStream.handle_thinking_end: {e}" - ) - async for chunk in super().handle_thinking_end( - part, followed_by_thinking - ): - yield chunk - - # Custom adapter that uses the custom event stream - class AnthropicVercelAIAdapter( - VercelAIAdapter[None, AnthropicOutputType] - ): - def build_event_stream(self) -> AnthropicVercelAIEventStream: - return AnthropicVercelAIEventStream( - self.run_input, accept=self.accept - ) - - return AnthropicVercelAIAdapter - class BedrockProvider(PydanticProvider["PydanticBedrock"]): def setup_credentials(self, config: AnyProviderConfig) -> None: From 7f705742a1935ae52941e1c79a4da8ee45b3181d Mon Sep 17 00:00:00 2001 From: Shahmir Varqha Date: Fri, 27 Mar 2026 10:36:57 +0800 Subject: [PATCH 7/8] require 1.52.0 in ui.chat --- marimo/_plugins/ui/_impl/chat/chat.py | 25 +++++++++++++------------ marimo/_server/ai/providers.py | 10 +++++----- 2 files changed, 18 insertions(+), 17 deletions(-) diff --git a/marimo/_plugins/ui/_impl/chat/chat.py b/marimo/_plugins/ui/_impl/chat/chat.py index 2209df58749..c07c5dc13f3 100644 --- a/marimo/_plugins/ui/_impl/chat/chat.py +++ b/marimo/_plugins/ui/_impl/chat/chat.py @@ -40,6 +40,13 @@ DONE_CHUNK: Final[str] = "[DONE]" +def require_vercel_ai_sdk_support() -> None: + """Only Pydantic AI >=1.52.0 supports AI SDK v6. So, we require it.""" + DependencyManager.pydantic_ai.require_at_version( + why="for Vercel AI SDK support", min_version="1.52.0" + ) + + @dataclass class SendMessageRequest: messages: list[ChatMessage] @@ -413,9 +420,8 @@ def _convert_value(self, value: dict[str, Any]) -> list[ChatMessage]: part_validator_class = None if DependencyManager.pydantic_ai.imported(): - from pydantic_ai.ui.vercel_ai.request_types import ( - UIMessagePart, - ) + require_vercel_ai_sdk_support() + from pydantic_ai.ui.vercel_ai.request_types import UIMessagePart # The frontend sends messages as ChatMessage parts so we use pydantic-ai to cast them # as Vercel UIMessagePart @@ -466,20 +472,15 @@ def handle_chunk(self, chunk: Any) -> None: # Handle Pydantic AI's Vercel AI SDK chunks if DependencyManager.pydantic_ai.imported(): + require_vercel_ai_sdk_support() from pydantic_ai.ui.vercel_ai.response_types import ( BaseChunk, ) if isinstance(chunk, BaseChunk): - try: - serialized = json.loads( - chunk.encode(sdk_version=AI_SDK_VERSION) - ) - except TypeError: - # Fallback for pydantic-ai < 1.52.0 which doesn't have sdk_version param - serialized = chunk.model_dump( - mode="json", by_alias=True, exclude_none=True - ) + serialized = json.loads( + chunk.encode(sdk_version=AI_SDK_VERSION) + ) self.on_send_chunk(serialized) return diff --git a/marimo/_server/ai/providers.py b/marimo/_server/ai/providers.py index 6294d4a6010..0f305cb6c45 100644 --- a/marimo/_server/ai/providers.py +++ b/marimo/_server/ai/providers.py @@ -25,7 +25,10 @@ generate_id, ) from marimo._dependencies.dependencies import Dependency, DependencyManager -from marimo._plugins.ui._impl.chat.chat import AI_SDK_VERSION +from marimo._plugins.ui._impl.chat.chat import ( + AI_SDK_VERSION, + require_vercel_ai_sdk_support, +) from marimo._server.ai.config import AnyProviderConfig from marimo._server.ai.ids import AiModelId from marimo._server.ai.tools.tool_manager import get_tool_manager @@ -98,10 +101,7 @@ def __init__( *(deps or []), source="server", ) - - DependencyManager.pydantic_ai.require_at_version( - "for Vercel AI SDK support", min_version="1.52.0" - ) + require_vercel_ai_sdk_support() self.model = model self.config = config From 308ce38c38746d462d0b6537883f0a8a0959d591 Mon Sep 17 00:00:00 2001 From: Shahmir Varqha Date: Fri, 27 Mar 2026 19:28:51 +0800 Subject: [PATCH 8/8] upgrade deps snapshot --- tests/_ai/llm/test_impl.py | 36 ------------------- .../optional-dependencies-recommended.txt | 2 +- 2 files changed, 1 insertion(+), 37 deletions(-) diff --git a/tests/_ai/llm/test_impl.py b/tests/_ai/llm/test_impl.py index 89e822ccb31..5f900441dfd 100644 --- a/tests/_ai/llm/test_impl.py +++ b/tests/_ai/llm/test_impl.py @@ -1571,42 +1571,6 @@ def test_pydantic_ai_serialize_vercel_ai_chunk(self) -> None: "input": {"query": "test"}, } - def test_pydantic_ai_serialize_vercel_ai_chunk_v5(self) -> None: - """Test that tool-input-start chunks exclude providerMetadata for SDK v5. - - The Vercel AI SDK v5 schema drifts from v6, so we need to use Pydantic's handling. - - For tool-input-start chunks, providerMetadata must be excluded. - See: https://github.com/pydantic/pydantic-ai/pull/4166 - """ - from pydantic_ai.ui.vercel_ai.response_types import ToolInputStartChunk - - mock_agent = MagicMock() - model = pydantic_ai(mock_agent) - - # Create chunk with providerMetadata (like Google Gemini produces) - chunk = ToolInputStartChunk( - tool_call_id="tc_1", - tool_name="my_tool", - provider_metadata={ - "pydantic_ai": { - "id": "test_id", - "provider_name": "google-gla", - "provider_details": { - "thought_signature": "encrypted_data" - }, - } - }, - ) - result = model._serialize_vercel_ai_chunk(chunk) - - # providerMetadata should be excluded for SDK v5 compatibility - assert result == { - "type": "tool-input-start", - "toolCallId": "tc_1", - "toolName": "my_tool", - } - def test_pydantic_ai_serialize_vercel_ai_chunk_done_type(self) -> None: """Test that 'done' type chunks are skipped.""" from pydantic_ai.ui.vercel_ai.response_types import DoneChunk diff --git a/tests/snapshots/optional-dependencies-recommended.txt b/tests/snapshots/optional-dependencies-recommended.txt index 1f4c4a57ed4..3367d4fca42 100644 --- a/tests/snapshots/optional-dependencies-recommended.txt +++ b/tests/snapshots/optional-dependencies-recommended.txt @@ -2,5 +2,5 @@ altair>=5.4.0 marimo[sandbox] marimo[sql] nbformat>=5.7.0 -pydantic-ai-slim[openai]>=1.39.0 +pydantic-ai-slim[openai]>=1.52.0 ruff \ No newline at end of file