From 6c32588b6abcd0180d49fe548cd6c83d35a50f29 Mon Sep 17 00:00:00 2001 From: Thomas Steinacher Date: Thu, 18 Dec 2025 00:31:31 +0100 Subject: [PATCH 01/21] Operate on a deepcopy of `$defs` in `JsonSchemaTransformer` instead of the original schema (#3758) --- pydantic_ai_slim/pydantic_ai/_json_schema.py | 2 +- tests/test_json_schema.py | 38 ++++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/pydantic_ai_slim/pydantic_ai/_json_schema.py b/pydantic_ai_slim/pydantic_ai/_json_schema.py index a0360dc7fa..1fc0680dd2 100644 --- a/pydantic_ai_slim/pydantic_ai/_json_schema.py +++ b/pydantic_ai_slim/pydantic_ai/_json_schema.py @@ -45,7 +45,7 @@ def __init__( self.prefer_inlined_defs = prefer_inlined_defs self.simplify_nullable_unions = simplify_nullable_unions - self.defs: dict[str, JsonSchema] = self.schema.get('$defs', {}) + self.defs: dict[str, JsonSchema] = deepcopy(self.schema.get('$defs', {})) self.refs_stack: list[str] = [] self.recursive_refs = set[str]() diff --git a/tests/test_json_schema.py b/tests/test_json_schema.py index da2e141ebc..5a91655840 100644 --- a/tests/test_json_schema.py +++ b/tests/test_json_schema.py @@ -2,6 +2,7 @@ from __future__ import annotations as _annotations +from copy import deepcopy from typing import Any from pydantic_ai._json_schema import JsonSchemaTransformer @@ -49,3 +50,40 @@ def transform(self, schema: dict[str, Any]) -> dict[str, Any]: # Should keep anyOf since it's not nullable assert 'anyOf' in result3 assert len(result3['anyOf']) == 2 + + +def test_schema_defs_not_modified(): + """Test that the original schema $defs are not modified during transformation.""" + + # Create a concrete subclass for testing + class TestTransformer(JsonSchemaTransformer): + def transform(self, schema: dict[str, Any]) -> dict[str, Any]: + return schema + + # Create a schema with $defs that should not be modified + original_schema = { + 'type': 'object', + 'properties': {'value': {'$ref': '#/$defs/TestUnion'}}, + '$defs': { + 'TestUnion': { + 'anyOf': [ + {'type': 'string'}, + {'type': 'number'}, + ], + 'title': 'TestUnion', + } + }, + } + + # Keep a deepcopy to compare against later + original_schema_copy = deepcopy(original_schema) + + # Transform the schema + transformer = TestTransformer(original_schema) + result = transformer.walk() + + # Verify the original schema was not modified + assert original_schema == original_schema_copy + + # Verify the result is correct + assert result == original_schema_copy From 8b8a928db02c73642601ec83e6ee727ede75707c Mon Sep 17 00:00:00 2001 From: lif <1835304752@qq.com> Date: Fri, 19 Dec 2025 05:02:16 +0800 Subject: [PATCH 02/21] Ensure `type` field when converting `const` to `enum` in `GoogleJsonSchemaTransformer` (#3751) Co-authored-by: Claude Opus 4.5 --- .../pydantic_ai/profiles/google.py | 11 + tests/profiles/test_google.py | 208 ++++++++++++++++++ 2 files changed, 219 insertions(+) create mode 100644 tests/profiles/test_google.py diff --git a/pydantic_ai_slim/pydantic_ai/profiles/google.py b/pydantic_ai_slim/pydantic_ai/profiles/google.py index 6708fc41c3..76be665e39 100644 --- a/pydantic_ai_slim/pydantic_ai/profiles/google.py +++ b/pydantic_ai_slim/pydantic_ai/profiles/google.py @@ -44,6 +44,17 @@ def transform(self, schema: JsonSchema) -> JsonSchema: if (const := schema.pop('const', None)) is not None: # Gemini doesn't support const, but it does support enum with a single value schema['enum'] = [const] + # If type is not present, infer it from the const value for Gemini API compatibility + if 'type' not in schema: + if isinstance(const, str): + schema['type'] = 'string' + elif isinstance(const, bool): + # bool must be checked before int since bool is a subclass of int in Python + schema['type'] = 'boolean' + elif isinstance(const, int): + schema['type'] = 'integer' + elif isinstance(const, float): + schema['type'] = 'number' schema.pop('discriminator', None) schema.pop('examples', None) diff --git a/tests/profiles/test_google.py b/tests/profiles/test_google.py new file mode 100644 index 0000000000..de7ec4aa7a --- /dev/null +++ b/tests/profiles/test_google.py @@ -0,0 +1,208 @@ +"""Tests for Google JSON schema transformer. + +The GoogleJsonSchemaTransformer transforms JSON schemas for compatibility with Gemini API: +- Converts `const` to `enum` with inferred `type` field +- Removes unsupported fields like $schema, title, discriminator, examples +- Handles format fields by moving them to description +""" + +from __future__ import annotations as _annotations + +from typing import Literal + +from inline_snapshot import snapshot +from pydantic import BaseModel + +from pydantic_ai.profiles.google import GoogleJsonSchemaTransformer, google_model_profile + +# ============================================================================= +# Transformer Tests - const to enum conversion with type inference +# ============================================================================= + + +def test_const_string_infers_type(): + """When converting const to enum, type should be inferred for string values.""" + schema = {'const': 'hello'} + transformer = GoogleJsonSchemaTransformer(schema) + transformed = transformer.walk() + + assert transformed == snapshot({'enum': ['hello'], 'type': 'string'}) + + +def test_const_integer_infers_type(): + """When converting const to enum, type should be inferred for integer values.""" + schema = {'const': 42} + transformer = GoogleJsonSchemaTransformer(schema) + transformed = transformer.walk() + + assert transformed == snapshot({'enum': [42], 'type': 'integer'}) + + +def test_const_float_infers_type(): + """When converting const to enum, type should be inferred for float values.""" + schema = {'const': 3.14} + transformer = GoogleJsonSchemaTransformer(schema) + transformed = transformer.walk() + + assert transformed == snapshot({'enum': [3.14], 'type': 'number'}) + + +def test_const_boolean_infers_type(): + """When converting const to enum, type should be inferred for boolean values.""" + schema = {'const': True} + transformer = GoogleJsonSchemaTransformer(schema) + transformed = transformer.walk() + + assert transformed == snapshot({'enum': [True], 'type': 'boolean'}) + + +def test_const_false_boolean_infers_type(): + """When converting const to enum, type should be inferred for False boolean.""" + schema = {'const': False} + transformer = GoogleJsonSchemaTransformer(schema) + transformed = transformer.walk() + + assert transformed == snapshot({'enum': [False], 'type': 'boolean'}) + + +def test_const_preserves_existing_type(): + """When const has an existing type field, it should be preserved.""" + schema = {'const': 'hello', 'type': 'string'} + transformer = GoogleJsonSchemaTransformer(schema) + transformed = transformer.walk() + + assert transformed == snapshot({'enum': ['hello'], 'type': 'string'}) + + +def test_const_array_does_not_infer_type(): + """When const is an array, type cannot be inferred and should not be added.""" + schema = {'const': [1, 2, 3]} + transformer = GoogleJsonSchemaTransformer(schema) + transformed = transformer.walk() + + assert transformed == snapshot({'enum': [[1, 2, 3]]}) + + +def test_const_in_nested_object(): + """const should be properly converted in nested object properties.""" + + class TaggedModel(BaseModel): + tag: Literal['hello'] + value: str + + schema = TaggedModel.model_json_schema() + transformer = GoogleJsonSchemaTransformer(schema) + transformed = transformer.walk() + + # The tag property should have both enum and type + assert transformed['properties']['tag'] == snapshot({'enum': ['hello'], 'type': 'string'}) + + +# ============================================================================= +# Transformer Tests - field removal +# ============================================================================= + + +def test_removes_schema_field(): + """$schema field should be removed.""" + schema = {'$schema': 'http://json-schema.org/draft-07/schema#', 'type': 'string'} + transformer = GoogleJsonSchemaTransformer(schema) + transformed = transformer.walk() + + assert '$schema' not in transformed + assert transformed == snapshot({'type': 'string'}) + + +def test_removes_title_field(): + """title field should be removed.""" + schema = {'title': 'MyString', 'type': 'string'} + transformer = GoogleJsonSchemaTransformer(schema) + transformed = transformer.walk() + + assert 'title' not in transformed + assert transformed == snapshot({'type': 'string'}) + + +def test_removes_discriminator_field(): + """discriminator field should be removed.""" + schema = {'discriminator': {'propertyName': 'type'}, 'type': 'object'} + transformer = GoogleJsonSchemaTransformer(schema) + transformed = transformer.walk() + + assert 'discriminator' not in transformed + assert transformed == snapshot({'type': 'object'}) + + +def test_removes_examples_field(): + """examples field should be removed.""" + schema = {'examples': ['foo', 'bar'], 'type': 'string'} + transformer = GoogleJsonSchemaTransformer(schema) + transformed = transformer.walk() + + assert 'examples' not in transformed + assert transformed == snapshot({'type': 'string'}) + + +def test_removes_exclusive_min_max(): + """exclusiveMinimum and exclusiveMaximum should be removed.""" + schema = {'type': 'integer', 'exclusiveMinimum': 0, 'exclusiveMaximum': 100} + transformer = GoogleJsonSchemaTransformer(schema) + transformed = transformer.walk() + + assert 'exclusiveMinimum' not in transformed + assert 'exclusiveMaximum' not in transformed + assert transformed == snapshot({'type': 'integer'}) + + +# ============================================================================= +# Transformer Tests - format handling +# ============================================================================= + + +def test_format_moved_to_description(): + """format should be moved to description for string types.""" + schema = {'type': 'string', 'format': 'date-time'} + transformer = GoogleJsonSchemaTransformer(schema) + transformed = transformer.walk() + + assert 'format' not in transformed + assert transformed == snapshot({'type': 'string', 'description': 'Format: date-time'}) + + +def test_format_appended_to_existing_description(): + """format should be appended to existing description.""" + schema = {'type': 'string', 'format': 'email', 'description': 'User email address'} + transformer = GoogleJsonSchemaTransformer(schema) + transformed = transformer.walk() + + assert 'format' not in transformed + assert transformed == snapshot({'type': 'string', 'description': 'User email address (format: email)'}) + + +# ============================================================================= +# Model Profile Tests +# ============================================================================= + + +def test_model_profile_gemini_2(): + """Gemini 2.x models should have proper profile settings.""" + profile = google_model_profile('gemini-2.0-flash') + assert profile is not None + assert profile.json_schema_transformer == GoogleJsonSchemaTransformer + assert profile.supports_json_schema_output is True + + +def test_model_profile_gemini_3(): + """Gemini 3.x models should support native output with builtin tools.""" + profile = google_model_profile('gemini-3.0-pro') + assert profile is not None + assert profile.google_supports_native_output_with_builtin_tools is True # type: ignore + + +def test_model_profile_image_model(): + """Image models should have limited capabilities.""" + profile = google_model_profile('gemini-2.0-flash-image') + assert profile is not None + assert profile.supports_image_output is True + assert profile.supports_json_schema_output is False + assert profile.supports_tools is False From 4a5da9b4ae3105fdb955d0245f69c42654c67743 Mon Sep 17 00:00:00 2001 From: Peter Li Date: Thu, 18 Dec 2025 15:06:30 -0800 Subject: [PATCH 03/21] Bump `google-genai` to 1.56.0 (#3770) --- uv.lock | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/uv.lock b/uv.lock index ba4012112b..75bc6b08fb 100644 --- a/uv.lock +++ b/uv.lock @@ -1984,16 +1984,16 @@ wheels = [ [[package]] name = "google-auth" -version = "2.38.0" +version = "2.45.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cachetools" }, { name = "pyasn1-modules" }, { name = "rsa" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/c6/eb/d504ba1daf190af6b204a9d4714d457462b486043744901a6eeea711f913/google_auth-2.38.0.tar.gz", hash = "sha256:8285113607d3b80a3f1543b75962447ba8a09fe85783432a784fdeef6ac094c4", size = 270866, upload-time = "2025-01-23T01:05:29.119Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e5/00/3c794502a8b892c404b2dea5b3650eb21bfc7069612fbfd15c7f17c1cb0d/google_auth-2.45.0.tar.gz", hash = "sha256:90d3f41b6b72ea72dd9811e765699ee491ab24139f34ebf1ca2b9cc0c38708f3", size = 320708, upload-time = "2025-12-15T22:58:42.889Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/9d/47/603554949a37bca5b7f894d51896a9c534b9eab808e2520a748e081669d0/google_auth-2.38.0-py2.py3-none-any.whl", hash = "sha256:e7dae6694313f434a2727bf2906f27ad259bae090d7aa896590d86feec3d9d4a", size = 210770, upload-time = "2025-01-23T01:05:26.572Z" }, + { url = "https://files.pythonhosted.org/packages/c6/97/451d55e05487a5cd6279a01a7e34921858b16f7dc8aa38a2c684743cd2b3/google_auth-2.45.0-py2.py3-none-any.whl", hash = "sha256:82344e86dc00410ef5382d99be677c6043d72e502b625aa4f4afa0bdacca0f36", size = 233312, upload-time = "2025-12-15T22:58:40.777Z" }, ] [package.optional-dependencies] @@ -2003,7 +2003,7 @@ requests = [ [[package]] name = "google-genai" -version = "1.55.0" +version = "1.56.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, @@ -2017,9 +2017,9 @@ dependencies = [ { name = "typing-extensions" }, { name = "websockets" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/1d/7c/19b59750592702305ae211905985ec8ab56f34270af4a159fba5f0214846/google_genai-1.55.0.tar.gz", hash = "sha256:ae9f1318fedb05c7c1b671a4148724751201e8908a87568364a309804064d986", size = 477615, upload-time = "2025-12-11T02:49:28.624Z" } +sdist = { url = "https://files.pythonhosted.org/packages/70/ad/d3ac5a102135bd3f1e4b1475ca65d2bd4bcc22eb2e9348ac40fe3fadb1d6/google_genai-1.56.0.tar.gz", hash = "sha256:0491af33c375f099777ae207d9621f044e27091fafad4c50e617eba32165e82f", size = 340451, upload-time = "2025-12-17T12:35:05.412Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3e/86/a5a8e32b2d40b30b5fb20e7b8113fafd1e38befa4d1801abd5ce6991065a/google_genai-1.55.0-py3-none-any.whl", hash = "sha256:98c422762b5ff6e16b8d9a1e4938e8e0ad910392a5422e47f5301498d7f373a1", size = 703389, upload-time = "2025-12-11T02:49:27.105Z" }, + { url = "https://files.pythonhosted.org/packages/84/93/94bc7a89ef4e7ed3666add55cd859d1483a22737251df659bf1aa46e9405/google_genai-1.56.0-py3-none-any.whl", hash = "sha256:9e6b11e0c105ead229368cb5849a480e4d0185519f8d9f538d61ecfcf193b052", size = 426563, upload-time = "2025-12-17T12:35:03.717Z" }, ] [[package]] From 64a6fe4d059a7382042c32fa7079e508df23794f Mon Sep 17 00:00:00 2001 From: Ben Drucker Date: Thu, 18 Dec 2025 16:06:44 -0800 Subject: [PATCH 04/21] Add AI SDK data chunk ID and tool approval types (#3760) --- .../pydantic_ai/ui/vercel_ai/_event_stream.py | 23 ++++++++++- .../ui/vercel_ai/response_types.py | 21 ++++++++++ tests/test_vercel_ai.py | 40 +++++++++++++++++-- 3 files changed, 79 insertions(+), 5 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py index ca94c5c186..070166df98 100644 --- a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py +++ b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py @@ -13,6 +13,7 @@ BuiltinToolCallPart, BuiltinToolReturnPart, FilePart, + FinishReason as PydanticFinishReason, FunctionToolResultEvent, RetryPromptPart, TextPart, @@ -23,6 +24,7 @@ ToolCallPartDelta, ) from ...output import OutputDataT +from ...run import AgentRunResultEvent from ...tools import AgentDepsT from .. import UIEventStream from .request_types import RequestData @@ -32,6 +34,7 @@ ErrorChunk, FileChunk, FinishChunk, + FinishReason, FinishStepChunk, ReasoningDeltaChunk, ReasoningEndChunk, @@ -48,6 +51,15 @@ ToolOutputErrorChunk, ) +# Map Pydantic AI finish reasons to Vercel AI format +_FINISH_REASON_MAP: dict[PydanticFinishReason, FinishReason] = { + 'stop': 'stop', + 'length': 'length', + 'content_filter': 'content-filter', + 'tool_call': 'tool-calls', + 'error': 'error', +} + __all__ = ['VercelAIEventStream'] # See https://ai-sdk.dev/docs/ai-sdk-ui/stream-protocol#data-stream-protocol @@ -64,6 +76,7 @@ class VercelAIEventStream(UIEventStream[RequestData, BaseChunk, AgentDepsT, Outp """UI event stream transformer for the Vercel AI protocol.""" _step_started: bool = False + _finish_reason: FinishReason = None @property def response_headers(self) -> Mapping[str, str] | None: @@ -85,10 +98,18 @@ async def before_response(self) -> AsyncIterator[BaseChunk]: async def after_stream(self) -> AsyncIterator[BaseChunk]: yield FinishStepChunk() - yield FinishChunk() + yield FinishChunk(finish_reason=self._finish_reason) yield DoneChunk() + async def handle_run_result(self, event: AgentRunResultEvent) -> AsyncIterator[BaseChunk]: + pydantic_reason = event.result.response.finish_reason + if pydantic_reason: + self._finish_reason = _FINISH_REASON_MAP.get(pydantic_reason) + return + yield + async def on_error(self, error: Exception) -> AsyncIterator[BaseChunk]: + self._finish_reason = 'error' yield ErrorChunk(error_text=str(error)) async def handle_text_start(self, part: TextPart, follows_text: bool = False) -> AsyncIterator[BaseChunk]: diff --git a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/response_types.py b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/response_types.py index 1255503107..6a7b98a2dc 100644 --- a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/response_types.py +++ b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/response_types.py @@ -16,6 +16,9 @@ ProviderMetadata = dict[str, dict[str, JSONValue]] """Provider metadata.""" +FinishReason = Literal['stop', 'length', 'content-filter', 'tool-calls', 'error', 'other', 'unknown'] | None +"""Reason why the model finished generating.""" + class BaseChunk(CamelBaseModel, ABC): """Abstract base class for response SSE events.""" @@ -145,6 +148,21 @@ class ToolOutputErrorChunk(BaseChunk): dynamic: bool | None = None +class ToolApprovalRequestChunk(BaseChunk): + """Tool approval request chunk for human-in-the-loop approval.""" + + type: Literal['tool-approval-request'] = 'tool-approval-request' + approval_id: str + tool_call_id: str + + +class ToolOutputDeniedChunk(BaseChunk): + """Tool output denied chunk when user denies tool execution.""" + + type: Literal['tool-output-denied'] = 'tool-output-denied' + tool_call_id: str + + class SourceUrlChunk(BaseChunk): """Source URL chunk.""" @@ -178,7 +196,9 @@ class DataChunk(BaseChunk): """Data chunk with dynamic type.""" type: Annotated[str, Field(pattern=r'^data-')] + id: str | None = None data: Any + transient: bool | None = None class StartStepChunk(BaseChunk): @@ -205,6 +225,7 @@ class FinishChunk(BaseChunk): """Finish chunk.""" type: Literal['finish'] = 'finish' + finish_reason: FinishReason = None message_metadata: Any | None = None diff --git a/tests/test_vercel_ai.py b/tests/test_vercel_ai.py index 12a4ea3eaa..0e191d445d 100644 --- a/tests/test_vercel_ai.py +++ b/tests/test_vercel_ai.py @@ -1039,7 +1039,7 @@ def client_response\ {'type': 'text-delta', 'delta': ' bodies safely?', 'id': IsStr()}, {'type': 'text-end', 'id': IsStr()}, {'type': 'finish-step'}, - {'type': 'finish'}, + {'type': 'finish', 'finishReason': 'stop'}, '[DONE]', ] ) @@ -1488,7 +1488,7 @@ async def stream_function( {'type': 'tool-input-available', 'toolCallId': IsStr(), 'toolName': 'unknown_tool', 'input': {}}, {'type': 'error', 'errorText': 'Exceeded maximum retries (1) for output validation'}, {'type': 'finish-step'}, - {'type': 'finish'}, + {'type': 'finish', 'finishReason': 'error'}, '[DONE]', ] ) @@ -1531,7 +1531,7 @@ async def tool(query: str) -> str: }, {'type': 'error', 'errorText': 'Unknown tool'}, {'type': 'finish-step'}, - {'type': 'finish'}, + {'type': 'finish', 'finishReason': 'error'}, '[DONE]', ] ) @@ -1572,7 +1572,7 @@ def raise_error(run_result: AgentRunResult[Any]) -> None: {'type': 'text-end', 'id': IsStr()}, {'type': 'error', 'errorText': 'Faulty on_complete'}, {'type': 'finish-step'}, - {'type': 'finish'}, + {'type': 'finish', 'finishReason': 'error'}, '[DONE]', ] ) @@ -1619,6 +1619,38 @@ async def on_complete(run_result: AgentRunResult[Any]) -> AsyncIterator[BaseChun ) +async def test_data_chunk_with_id_and_transient(): + """Test DataChunk supports optional id and transient fields for AI SDK compatibility.""" + agent = Agent(model=TestModel()) + + request = SubmitMessage( + id='foo', + messages=[ + UIMessage( + id='bar', + role='user', + parts=[TextUIPart(text='Hello')], + ), + ], + ) + + async def on_complete(run_result: AgentRunResult[Any]) -> AsyncIterator[BaseChunk]: + # Yield a data chunk with id for reconciliation + yield DataChunk(type='data-task', id='task-123', data={'status': 'complete'}) + # Yield a transient data chunk (not persisted to history) + yield DataChunk(type='data-progress', data={'percent': 100}, transient=True) + + adapter = VercelAIAdapter(agent, request) + events = [ + '[DONE]' if '[DONE]' in event else json.loads(event.removeprefix('data: ')) + async for event in adapter.encode_stream(adapter.run_stream(on_complete=on_complete)) + ] + + # Verify the data chunks are present in the events with correct fields + assert {'type': 'data-task', 'id': 'task-123', 'data': {'status': 'complete'}} in events + assert {'type': 'data-progress', 'data': {'percent': 100}, 'transient': True} in events + + @pytest.mark.skipif(not starlette_import_successful, reason='Starlette is not installed') async def test_adapter_dispatch_request(): agent = Agent(model=TestModel()) From 116059102904b9ee1f52d74c089cb5dc1d843c8b Mon Sep 17 00:00:00 2001 From: bendrucker Date: Thu, 18 Dec 2025 21:03:02 -0800 Subject: [PATCH 05/21] Add tool approval integration for Vercel AI adapter --- docs/ui/vercel-ai.md | 58 +++++ .../pydantic_ai/ui/vercel_ai/_adapter.py | 51 ++++- .../pydantic_ai/ui/vercel_ai/_event_stream.py | 58 ++++- .../pydantic_ai/ui/vercel_ai/request_types.py | 36 +++- .../ui/vercel_ai/response_types.py | 4 +- tests/test_vercel_ai.py | 198 ++++++++++++++++++ 6 files changed, 395 insertions(+), 10 deletions(-) diff --git a/docs/ui/vercel-ai.md b/docs/ui/vercel-ai.md index 7d94aadbb9..b5e7e45c86 100644 --- a/docs/ui/vercel-ai.md +++ b/docs/ui/vercel-ai.md @@ -2,6 +2,9 @@ Pydantic AI natively supports the [Vercel AI Data Stream Protocol](https://ai-sdk.dev/docs/ai-sdk-ui/stream-protocol#data-stream-protocol) to receive agent run input from, and stream events to, a [Vercel AI Elements](https://ai-sdk.dev/elements) frontend. +!!! note "AI SDK Version Compatibility" + The base protocol is compatible with AI SDK v5 and later. However, [tool approval](#tool-approval) (human-in-the-loop) features require **AI SDK v6 or later**. + ## Usage The [`VercelAIAdapter`][pydantic_ai.ui.vercel_ai.VercelAIAdapter] class is responsible for transforming agent run input received from the frontend into arguments for [`Agent.run_stream_events()`](../agents.md#running-agents), running the agent, and then transforming Pydantic AI events into Vercel AI events. The event stream transformation is handled by the [`VercelAIEventStream`][pydantic_ai.ui.vercel_ai.VercelAIEventStream] class, but you typically won't use this directly. @@ -81,3 +84,58 @@ async def chat(request: Request) -> Response: sse_event_stream = adapter.encode_stream(event_stream) return StreamingResponse(sse_event_stream, media_type=accept) ``` + +## Tool Approval + +Pydantic AI supports [AI SDK's human-in-the-loop tool approval](https://ai-sdk.dev/cookbook/next/human-in-the-loop) workflow, allowing users to approve or deny tool executions before they run. + +!!! warning "Requires AI SDK v6" + Tool approval is an AI SDK v6 feature. The `tool-approval-request` and `tool-output-denied` stream chunks, along with the `approval` field on tool parts, are not available in AI SDK v5. + +### How It Works + +1. **Tool requests approval**: When an agent calls a tool with `requires_approval=True`, Pydantic AI emits a `tool-approval-request` chunk instead of executing the tool immediately. + +2. **User decides**: The AI SDK frontend displays an approval UI. The user can approve or deny the tool execution. + +3. **Response is sent**: The frontend sends the approval decision back to the server in a follow-up request. + +4. **Tool executes or is denied**: If approved, the tool runs normally. If denied, Pydantic AI emits a `tool-output-denied` chunk and informs the model that the tool was rejected. + +### Server-Side Setup + +To enable tool approval, define your tool with `requires_approval=True` and include [`DeferredToolRequests`][pydantic_ai.tools.DeferredToolRequests] in your agent's output types: + +```py +from pydantic_ai import Agent +from pydantic_ai.tools import DeferredToolRequests + +agent: Agent[None, str | DeferredToolRequests] = Agent( + 'openai:gpt-5', + output_type=[str, DeferredToolRequests], +) + +@agent.tool_plain(requires_approval=True) +def delete_file(path: str) -> str: + """Delete a file from the filesystem.""" + # This won't execute until the user approves + os.remove(path) + return f'Deleted {path}' +``` + +When processing a follow-up request with approval responses, extract and pass the deferred tool results: + +```py +@app.post('/chat') +async def chat(request: Request) -> Response: + adapter = await VercelAIAdapter.from_request(request, agent=agent) + return adapter.streaming_response( + adapter.run_stream(deferred_tool_results=adapter.deferred_tool_results) + ) +``` + +### Client-Side Setup + +On the frontend, use AI SDK v6's [`useChat`](https://v6.ai-sdk.dev/docs/ai-sdk-ui/use-chat) hook with the [`Confirmation`](https://ai-sdk.dev/elements/components/confirmation) component or the `addToolApprovalResponse` function to handle approval UI. + +See the [AI SDK Human-in-the-Loop Cookbook](https://ai-sdk.dev/cookbook/next/human-in-the-loop) for complete frontend examples. diff --git a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py index fa82b9255b..4cfd5580c0 100644 --- a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py +++ b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py @@ -35,7 +35,7 @@ VideoUrl, ) from ...output import OutputDataT -from ...tools import AgentDepsT +from ...tools import AgentDepsT, DeferredToolResults, ToolApproved, ToolDenied from .. import MessagesBuilder, UIAdapter, UIEventStream from ._event_stream import VercelAIEventStream from .request_types import ( @@ -51,6 +51,7 @@ SourceUrlUIPart, StepStartUIPart, TextUIPart, + ToolApprovalResponded, ToolInputAvailablePart, ToolOutputAvailablePart, ToolOutputErrorPart, @@ -87,6 +88,54 @@ def messages(self) -> list[ModelMessage]: """Pydantic AI messages from the Vercel AI run input.""" return self.load_messages(self.run_input.messages) + @cached_property + def deferred_tool_results(self) -> DeferredToolResults | None: + """Extract deferred tool results from tool parts with approval responses. + + When the Vercel AI SDK client responds to a tool-approval-request, it sends + the approval decision in the tool part's `approval` field. This method extracts + those responses and converts them to Pydantic AI's `DeferredToolResults` format. + + Returns: + DeferredToolResults if any tool parts have approval responses, None otherwise. + """ + return self.extract_deferred_tool_results(self.run_input.messages) + + @classmethod + def extract_deferred_tool_results(cls, messages: Sequence[UIMessage]) -> DeferredToolResults | None: + """Extract deferred tool results from UI messages. + + Args: + messages: The UI messages to scan for approval responses. + + Returns: + DeferredToolResults if any tool parts have approval responses, None otherwise. + """ + approvals: dict[str, bool | ToolApproved | ToolDenied] = {} + + for msg in messages: + if msg.role != 'assistant': + continue + + for part in msg.parts: + if not isinstance(part, ToolUIPart | DynamicToolUIPart): + continue + + approval = part.approval + if approval is None or not isinstance(approval, ToolApprovalResponded): + continue + + tool_call_id = part.tool_call_id + if approval.approved: + approvals[tool_call_id] = ToolApproved() + else: + approvals[tool_call_id] = ToolDenied(message=approval.reason or 'The tool call was denied.') + + if not approvals: + return None + + return DeferredToolResults(approvals=approvals) + @classmethod def load_messages(cls, messages: Sequence[UIMessage]) -> list[ModelMessage]: # noqa: C901 """Transform Vercel AI messages into Pydantic AI messages.""" diff --git a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py index 070166df98..564ac9fa3e 100644 --- a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py +++ b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py @@ -5,6 +5,7 @@ from collections.abc import AsyncIterator, Mapping from dataclasses import dataclass from typing import Any +from uuid import uuid4 from pydantic_core import to_json @@ -25,9 +26,15 @@ ) from ...output import OutputDataT from ...run import AgentRunResultEvent -from ...tools import AgentDepsT +from ...tools import AgentDepsT, DeferredToolRequests from .. import UIEventStream -from .request_types import RequestData +from .request_types import ( + DynamicToolUIPart, + RequestData, + ToolApprovalResponded, + ToolUIPart, + UIMessage, +) from .response_types import ( BaseChunk, DoneChunk, @@ -44,10 +51,12 @@ TextDeltaChunk, TextEndChunk, TextStartChunk, + ToolApprovalRequestChunk, ToolInputAvailableChunk, ToolInputDeltaChunk, ToolInputStartChunk, ToolOutputAvailableChunk, + ToolOutputDeniedChunk, ToolOutputErrorChunk, ) @@ -71,12 +80,35 @@ def _json_dumps(obj: Any) -> str: return to_json(obj).decode('utf-8') +def _extract_denied_tool_ids(messages: list[UIMessage]) -> set[str]: + """Extract tool_call_ids that were denied from UI messages.""" + denied_ids: set[str] = set() + for msg in messages: + if msg.role != 'assistant': + continue + for part in msg.parts: + if not isinstance(part, ToolUIPart | DynamicToolUIPart): + continue + approval = part.approval + if isinstance(approval, ToolApprovalResponded) and not approval.approved: + denied_ids.add(part.tool_call_id) + return denied_ids + + @dataclass class VercelAIEventStream(UIEventStream[RequestData, BaseChunk, AgentDepsT, OutputDataT]): """UI event stream transformer for the Vercel AI protocol.""" _step_started: bool = False _finish_reason: FinishReason = None + _denied_tool_ids: set[str] | None = None + + @property + def denied_tool_ids(self) -> set[str]: + """Get the set of tool_call_ids that were denied by the user.""" + if self._denied_tool_ids is None: + self._denied_tool_ids = _extract_denied_tool_ids(self.run_input.messages) + return self._denied_tool_ids @property def response_headers(self) -> Mapping[str, str] | None: @@ -105,8 +137,15 @@ async def handle_run_result(self, event: AgentRunResultEvent) -> AsyncIterator[B pydantic_reason = event.result.response.finish_reason if pydantic_reason: self._finish_reason = _FINISH_REASON_MAP.get(pydantic_reason) - return - yield + + # Emit tool approval requests for deferred approvals + output = event.result.output + if isinstance(output, DeferredToolRequests): + for tool_call in output.approvals: + yield ToolApprovalRequestChunk( + approval_id=str(uuid4()), + tool_call_id=tool_call.tool_call_id, + ) async def on_error(self, error: Exception) -> AsyncIterator[BaseChunk]: self._finish_reason = 'error' @@ -203,10 +242,15 @@ async def handle_file(self, part: FilePart) -> AsyncIterator[BaseChunk]: async def handle_function_tool_result(self, event: FunctionToolResultEvent) -> AsyncIterator[BaseChunk]: part = event.result - if isinstance(part, RetryPromptPart): - yield ToolOutputErrorChunk(tool_call_id=part.tool_call_id, error_text=part.model_response()) + tool_call_id = part.tool_call_id + + # Check if this tool was denied by the user + if tool_call_id in self.denied_tool_ids: + yield ToolOutputDeniedChunk(tool_call_id=tool_call_id) + elif isinstance(part, RetryPromptPart): + yield ToolOutputErrorChunk(tool_call_id=tool_call_id, error_text=part.model_response()) else: - yield ToolOutputAvailableChunk(tool_call_id=part.tool_call_id, output=self._tool_return_output(part)) + yield ToolOutputAvailableChunk(tool_call_id=tool_call_id, output=self._tool_return_output(part)) # ToolCallResultEvent.content may hold user parts (e.g. text, images) that Vercel AI does not currently have events for diff --git a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/request_types.py b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/request_types.py index 1fe9a593af..3203bb6162 100644 --- a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/request_types.py +++ b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/request_types.py @@ -1,7 +1,9 @@ """Vercel AI request types (UI messages). Converted to Python from: -https://github.com/vercel/ai/blob/ai%405.0.59/packages/ai/src/ui/ui-messages.ts +https://github.com/vercel/ai/blob/ai%406.0.0-beta.159/packages/ai/src/ui/ui-messages.ts + +Tool approval types (`ToolApprovalRequested`, `ToolApprovalResponded`) require AI SDK v6 or later. """ from abc import ABC @@ -110,6 +112,30 @@ class DataUIPart(BaseUIPart): data: Any +class ToolApprovalRequested(CamelBaseModel): + """Tool approval in requested state (awaiting user response).""" + + id: str + """The approval request ID.""" + + +class ToolApprovalResponded(CamelBaseModel): + """Tool approval in responded state (user has approved or denied).""" + + id: str + """The approval request ID.""" + + approved: bool + """Whether the user approved the tool call.""" + + reason: str | None = None + """Optional reason for the approval or denial.""" + + +ToolApproval = ToolApprovalRequested | ToolApprovalResponded +"""Union of tool approval states.""" + + # Tool part states as separate models class ToolInputStreamingPart(BaseUIPart): """Tool part in input-streaming state.""" @@ -119,6 +145,7 @@ class ToolInputStreamingPart(BaseUIPart): state: Literal['input-streaming'] = 'input-streaming' input: Any | None = None provider_executed: bool | None = None + approval: ToolApproval | None = None class ToolInputAvailablePart(BaseUIPart): @@ -130,6 +157,7 @@ class ToolInputAvailablePart(BaseUIPart): input: Any | None = None provider_executed: bool | None = None call_provider_metadata: ProviderMetadata | None = None + approval: ToolApproval | None = None class ToolOutputAvailablePart(BaseUIPart): @@ -143,6 +171,7 @@ class ToolOutputAvailablePart(BaseUIPart): provider_executed: bool | None = None call_provider_metadata: ProviderMetadata | None = None preliminary: bool | None = None + approval: ToolApproval | None = None class ToolOutputErrorPart(BaseUIPart): @@ -156,6 +185,7 @@ class ToolOutputErrorPart(BaseUIPart): error_text: str provider_executed: bool | None = None call_provider_metadata: ProviderMetadata | None = None + approval: ToolApproval | None = None ToolUIPart = ToolInputStreamingPart | ToolInputAvailablePart | ToolOutputAvailablePart | ToolOutputErrorPart @@ -171,6 +201,7 @@ class DynamicToolInputStreamingPart(BaseUIPart): tool_call_id: str state: Literal['input-streaming'] = 'input-streaming' input: Any | None = None + approval: ToolApproval | None = None class DynamicToolInputAvailablePart(BaseUIPart): @@ -182,6 +213,7 @@ class DynamicToolInputAvailablePart(BaseUIPart): state: Literal['input-available'] = 'input-available' input: Any call_provider_metadata: ProviderMetadata | None = None + approval: ToolApproval | None = None class DynamicToolOutputAvailablePart(BaseUIPart): @@ -195,6 +227,7 @@ class DynamicToolOutputAvailablePart(BaseUIPart): output: Any call_provider_metadata: ProviderMetadata | None = None preliminary: bool | None = None + approval: ToolApproval | None = None class DynamicToolOutputErrorPart(BaseUIPart): @@ -207,6 +240,7 @@ class DynamicToolOutputErrorPart(BaseUIPart): input: Any error_text: str call_provider_metadata: ProviderMetadata | None = None + approval: ToolApproval | None = None DynamicToolUIPart = ( diff --git a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/response_types.py b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/response_types.py index 6a7b98a2dc..1f333f22dd 100644 --- a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/response_types.py +++ b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/response_types.py @@ -1,7 +1,9 @@ """Vercel AI response types (SSE chunks). Converted to Python from: -https://github.com/vercel/ai/blob/ai%405.0.59/packages/ai/src/ui-message-stream/ui-message-chunks.ts +https://github.com/vercel/ai/blob/ai%406.0.0-beta.159/packages/ai/src/ui-message-stream/ui-message-chunks.ts + +Tool approval types (`ToolApprovalRequestChunk`, `ToolOutputDeniedChunk`) require AI SDK v6 or later. """ from abc import ABC diff --git a/tests/test_vercel_ai.py b/tests/test_vercel_ai.py index 0e191d445d..b76dc4aa33 100644 --- a/tests/test_vercel_ai.py +++ b/tests/test_vercel_ai.py @@ -1651,6 +1651,197 @@ async def on_complete(run_result: AgentRunResult[Any]) -> AsyncIterator[BaseChun assert {'type': 'data-progress', 'data': {'percent': 100}, 'transient': True} in events +async def test_tool_approval_request_emission(): + """Test that ToolApprovalRequestChunk is emitted when tools require approval.""" + from pydantic_ai.tools import DeferredToolRequests + + async def stream_function( + messages: list[ModelMessage], agent_info: AgentInfo + ) -> AsyncIterator[DeltaToolCalls | str]: + yield { + 0: DeltaToolCall( + name='delete_file', + json_args='{"path": "test.txt"}', + tool_call_id='delete_1', + ) + } + + agent: Agent[None, str | DeferredToolRequests] = Agent( + model=FunctionModel(stream_function=stream_function), output_type=[str, DeferredToolRequests] + ) + + @agent.tool_plain(requires_approval=True) + def delete_file(path: str) -> str: + return f'Deleted {path}' + + request = SubmitMessage( + id='foo', + messages=[ + UIMessage( + id='bar', + role='user', + parts=[TextUIPart(text='Delete test.txt')], + ), + ], + ) + + adapter = VercelAIAdapter(agent, request) + events: list[str | dict[str, Any]] = [ + '[DONE]' if '[DONE]' in event else json.loads(event.removeprefix('data: ')) + async for event in adapter.encode_stream(adapter.run_stream()) + ] + + # Verify tool-approval-request chunk is emitted with UUID approval_id + approval_event: dict[str, Any] | None = next( + (e for e in events if isinstance(e, dict) and e.get('type') == 'tool-approval-request'), + None, + ) + assert approval_event is not None + assert approval_event['toolCallId'] == 'delete_1' + assert 'approvalId' in approval_event + + +def test_extract_deferred_tool_results_approved(): + """Test that approved tool calls are correctly extracted from UI messages.""" + from pydantic_ai.tools import ToolApproved + from pydantic_ai.ui.vercel_ai.request_types import ( + DynamicToolInputAvailablePart, + ToolApprovalResponded, + ) + + messages = [ + UIMessage( + id='msg-1', + role='assistant', + parts=[ + DynamicToolInputAvailablePart( + tool_name='delete_file', + tool_call_id='delete_1', + input={'path': 'test.txt'}, + approval=ToolApprovalResponded(id='approval-123', approved=True), + ), + ], + ), + ] + + result = VercelAIAdapter.extract_deferred_tool_results(messages) + assert result is not None + assert 'delete_1' in result.approvals + assert isinstance(result.approvals['delete_1'], ToolApproved) + + +def test_extract_deferred_tool_results_denied(): + """Test that denied tool calls are correctly extracted from UI messages.""" + from pydantic_ai.tools import ToolDenied + from pydantic_ai.ui.vercel_ai.request_types import ( + DynamicToolInputAvailablePart, + ToolApprovalResponded, + ) + + messages = [ + UIMessage( + id='msg-1', + role='assistant', + parts=[ + DynamicToolInputAvailablePart( + tool_name='delete_file', + tool_call_id='delete_1', + input={'path': 'test.txt'}, + approval=ToolApprovalResponded(id='approval-123', approved=False, reason='User rejected deletion'), + ), + ], + ), + ] + + result = VercelAIAdapter.extract_deferred_tool_results(messages) + assert result is not None + assert 'delete_1' in result.approvals + denial = result.approvals['delete_1'] + assert isinstance(denial, ToolDenied) + assert denial.message == 'User rejected deletion' + + +def test_extract_deferred_tool_results_no_approvals(): + """Test that None is returned when no approval responses exist.""" + messages = [ + UIMessage( + id='msg-1', + role='user', + parts=[TextUIPart(text='Hello')], + ), + ] + + result = VercelAIAdapter.extract_deferred_tool_results(messages) + assert result is None + + +async def test_tool_output_denied_chunk_emission(): + """Test that ToolOutputDeniedChunk is emitted when a tool call is denied.""" + from pydantic_ai.tools import DeferredToolRequests + from pydantic_ai.ui.vercel_ai.request_types import ( + DynamicToolInputAvailablePart, + ToolApprovalResponded, + ) + + async def stream_function( + messages: list[ModelMessage], agent_info: AgentInfo + ) -> AsyncIterator[DeltaToolCalls | str]: + # Model acknowledges the denial + yield 'The file deletion was cancelled as requested.' + + agent: Agent[None, str | DeferredToolRequests] = Agent( + model=FunctionModel(stream_function=stream_function), output_type=[str, DeferredToolRequests] + ) + + @agent.tool_plain(requires_approval=True) + def delete_file(path: str) -> str: + return f'Deleted {path}' + + # Simulate a follow-up request where the user denied the tool + request = SubmitMessage( + id='foo', + messages=[ + UIMessage( + id='user-1', + role='user', + parts=[TextUIPart(text='Delete test.txt')], + ), + UIMessage( + id='assistant-1', + role='assistant', + parts=[ + DynamicToolInputAvailablePart( + tool_name='delete_file', + tool_call_id='delete_1', + input={'path': 'test.txt'}, + approval=ToolApprovalResponded( + id='approval-123', + approved=False, + reason='User cancelled the deletion', + ), + ), + ], + ), + ], + ) + + adapter = VercelAIAdapter(agent, request) + events: list[str | dict[str, Any]] = [ + '[DONE]' if '[DONE]' in event else json.loads(event.removeprefix('data: ')) + async for event in adapter.encode_stream( + adapter.run_stream(deferred_tool_results=adapter.deferred_tool_results) + ) + ] + + # Verify tool-output-denied chunk is emitted + denied_event: dict[str, Any] | None = next( + (e for e in events if isinstance(e, dict) and e.get('type') == 'tool-output-denied'), + None, + ) + assert denied_event is not None + assert denied_event['toolCallId'] == 'delete_1' + + @pytest.mark.skipif(not starlette_import_successful, reason='Starlette is not installed') async def test_adapter_dispatch_request(): agent = Agent(model=TestModel()) @@ -2089,6 +2280,7 @@ async def test_adapter_dump_messages_with_tools(): 'output': '{"results":["result1","result2"]}', 'call_provider_metadata': None, 'preliminary': None, + 'approval': None, }, ], }, @@ -2151,6 +2343,7 @@ async def test_adapter_dump_messages_with_builtin_tools(): 'provider_executed': True, 'call_provider_metadata': {'pydantic_ai': {'provider_name': 'openai'}}, 'preliminary': None, + 'approval': None, } ], }, @@ -2197,6 +2390,7 @@ async def test_adapter_dump_messages_with_builtin_tool_without_return(): 'input': '{"query":"orphan query"}', 'provider_executed': True, 'call_provider_metadata': {'pydantic_ai': {'provider_name': 'openai'}}, + 'approval': None, } ], }, @@ -2363,6 +2557,7 @@ async def test_adapter_dump_messages_with_retry(): Fix the errors and try again.\ """, 'call_provider_metadata': None, + 'approval': None, } ], }, @@ -2494,6 +2689,7 @@ async def test_adapter_dump_messages_text_with_interruption(): 'provider_executed': True, 'call_provider_metadata': {'pydantic_ai': {'provider_name': 'test'}}, 'preliminary': None, + 'approval': None, }, { 'type': 'text', @@ -2612,6 +2808,7 @@ async def test_adapter_dump_messages_tool_call_without_return(): 'state': 'input-available', 'input': '{"city":"New York"}', 'call_provider_metadata': None, + 'approval': None, } ], } @@ -2646,6 +2843,7 @@ async def test_adapter_dump_messages_assistant_starts_with_tool(): 'state': 'input-available', 'input': '{}', 'call_provider_metadata': None, + 'approval': None, }, { 'type': 'text', From 85a415b2c048b140c5d60cc0a40da97cc0bc7fb7 Mon Sep 17 00:00:00 2001 From: bendrucker Date: Thu, 18 Dec 2025 21:16:33 -0800 Subject: [PATCH 06/21] Add robustness improvements and tests for tool approval --- .../pydantic_ai/ui/vercel_ai/_adapter.py | 3 + .../pydantic_ai/ui/vercel_ai/_event_stream.py | 6 +- tests/test_vercel_ai.py | 152 +++++++++++++++++- 3 files changed, 158 insertions(+), 3 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py index 4cfd5580c0..ed5d28a02d 100644 --- a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py +++ b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py @@ -126,6 +126,9 @@ def extract_deferred_tool_results(cls, messages: Sequence[UIMessage]) -> Deferre continue tool_call_id = part.tool_call_id + if not tool_call_id: + continue + if approval.approved: approvals[tool_call_id] = ToolApproved() else: diff --git a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py index 564ac9fa3e..8cf18973ac 100644 --- a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py +++ b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py @@ -91,7 +91,9 @@ def _extract_denied_tool_ids(messages: list[UIMessage]) -> set[str]: continue approval = part.approval if isinstance(approval, ToolApprovalResponded) and not approval.approved: - denied_ids.add(part.tool_call_id) + tool_call_id = part.tool_call_id + if tool_call_id: + denied_ids.add(tool_call_id) return denied_ids @@ -136,7 +138,7 @@ async def after_stream(self) -> AsyncIterator[BaseChunk]: async def handle_run_result(self, event: AgentRunResultEvent) -> AsyncIterator[BaseChunk]: pydantic_reason = event.result.response.finish_reason if pydantic_reason: - self._finish_reason = _FINISH_REASON_MAP.get(pydantic_reason) + self._finish_reason = _FINISH_REASON_MAP.get(pydantic_reason, 'unknown') # Emit tool approval requests for deferred approvals output = event.result.output diff --git a/tests/test_vercel_ai.py b/tests/test_vercel_ai.py index b76dc4aa33..045b016307 100644 --- a/tests/test_vercel_ai.py +++ b/tests/test_vercel_ai.py @@ -3,6 +3,7 @@ import json from collections.abc import AsyncIterator, MutableMapping from typing import Any, cast +from uuid import UUID import pytest from inline_snapshot import snapshot @@ -1698,7 +1699,10 @@ def delete_file(path: str) -> str: ) assert approval_event is not None assert approval_event['toolCallId'] == 'delete_1' - assert 'approvalId' in approval_event + # Validate approval_id is a valid UUID + approval_id = approval_event.get('approvalId') + assert approval_id is not None + UUID(approval_id) # Raises ValueError if not a valid UUID def test_extract_deferred_tool_results_approved(): @@ -1775,6 +1779,152 @@ def test_extract_deferred_tool_results_no_approvals(): assert result is None +def test_extract_deferred_tool_results_with_tool_ui_part(): + """Test that ToolUIPart (builtin/provider-executed tools) is handled correctly.""" + from pydantic_ai.tools import ToolApproved, ToolDenied + from pydantic_ai.ui.vercel_ai.request_types import ToolApprovalResponded + + messages = [ + UIMessage( + id='msg-1', + role='assistant', + parts=[ + ToolInputAvailablePart( + type='tool-delete_file', + tool_call_id='delete_1', + input={'path': 'test.txt'}, + approval=ToolApprovalResponded(id='approval-123', approved=True), + ), + ToolOutputAvailablePart( + type='tool-read_file', + tool_call_id='read_1', + input={'path': 'other.txt'}, + output='file contents', + approval=ToolApprovalResponded(id='approval-456', approved=False, reason='Not allowed'), + ), + ], + ), + ] + + result = VercelAIAdapter.extract_deferred_tool_results(messages) + assert result is not None + assert 'delete_1' in result.approvals + assert 'read_1' in result.approvals + assert isinstance(result.approvals['delete_1'], ToolApproved) + denial = result.approvals['read_1'] + assert isinstance(denial, ToolDenied) + assert denial.message == 'Not allowed' + + +def test_extract_deferred_tool_results_denied_no_reason(): + """Test that denied tool calls use default message when no reason is provided.""" + from pydantic_ai.tools import ToolDenied + from pydantic_ai.ui.vercel_ai.request_types import ( + DynamicToolInputAvailablePart, + ToolApprovalResponded, + ) + + messages = [ + UIMessage( + id='msg-1', + role='assistant', + parts=[ + DynamicToolInputAvailablePart( + tool_name='delete_file', + tool_call_id='delete_1', + input={'path': 'test.txt'}, + approval=ToolApprovalResponded(id='approval-123', approved=False), + ), + ], + ), + ] + + result = VercelAIAdapter.extract_deferred_tool_results(messages) + assert result is not None + denial = result.approvals['delete_1'] + assert isinstance(denial, ToolDenied) + assert denial.message == 'The tool call was denied.' + + +def test_extract_deferred_tool_results_multiple_approvals(): + """Test that multiple tool approvals in a single message are all extracted.""" + from pydantic_ai.tools import ToolApproved, ToolDenied + from pydantic_ai.ui.vercel_ai.request_types import ( + DynamicToolInputAvailablePart, + ToolApprovalResponded, + ) + + messages = [ + UIMessage( + id='msg-1', + role='assistant', + parts=[ + DynamicToolInputAvailablePart( + tool_name='delete_file', + tool_call_id='tool_1', + input={'path': 'a.txt'}, + approval=ToolApprovalResponded(id='approval-1', approved=True), + ), + DynamicToolInputAvailablePart( + tool_name='delete_file', + tool_call_id='tool_2', + input={'path': 'b.txt'}, + approval=ToolApprovalResponded(id='approval-2', approved=False, reason='No'), + ), + DynamicToolInputAvailablePart( + tool_name='read_file', + tool_call_id='tool_3', + input={'path': 'c.txt'}, + approval=ToolApprovalResponded(id='approval-3', approved=True), + ), + ], + ), + ] + + result = VercelAIAdapter.extract_deferred_tool_results(messages) + assert result is not None + assert len(result.approvals) == 3 + assert isinstance(result.approvals['tool_1'], ToolApproved) + assert isinstance(result.approvals['tool_2'], ToolDenied) + assert isinstance(result.approvals['tool_3'], ToolApproved) + + +def test_extract_deferred_tool_results_ignores_pending_approval(): + """Test that ToolApprovalRequested (pending) is ignored, only ToolApprovalResponded is processed.""" + from pydantic_ai.ui.vercel_ai.request_types import ( + DynamicToolInputAvailablePart, + ToolApprovalRequested, + ToolApprovalResponded, + ) + + messages = [ + UIMessage( + id='msg-1', + role='assistant', + parts=[ + DynamicToolInputAvailablePart( + tool_name='delete_file', + tool_call_id='pending_tool', + input={'path': 'pending.txt'}, + approval=ToolApprovalRequested(id='approval-pending'), + ), + DynamicToolInputAvailablePart( + tool_name='delete_file', + tool_call_id='responded_tool', + input={'path': 'responded.txt'}, + approval=ToolApprovalResponded(id='approval-responded', approved=True), + ), + ], + ), + ] + + result = VercelAIAdapter.extract_deferred_tool_results(messages) + assert result is not None + # Only the responded approval should be extracted + assert 'pending_tool' not in result.approvals + assert 'responded_tool' in result.approvals + + async def test_tool_output_denied_chunk_emission(): """Test that ToolOutputDeniedChunk is emitted when a tool call is denied.""" from pydantic_ai.tools import DeferredToolRequests From 62f7262a680878870d7f66d85c9ffff6bb65a3bf Mon Sep 17 00:00:00 2001 From: bendrucker Date: Thu, 18 Dec 2025 21:53:31 -0800 Subject: [PATCH 07/21] Skip doc examples that are incomplete snippets --- docs/ui/vercel-ai.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/ui/vercel-ai.md b/docs/ui/vercel-ai.md index b5e7e45c86..88238a240b 100644 --- a/docs/ui/vercel-ai.md +++ b/docs/ui/vercel-ai.md @@ -106,7 +106,7 @@ Pydantic AI supports [AI SDK's human-in-the-loop tool approval](https://ai-sdk.d To enable tool approval, define your tool with `requires_approval=True` and include [`DeferredToolRequests`][pydantic_ai.tools.DeferredToolRequests] in your agent's output types: -```py +```py {test="skip"} from pydantic_ai import Agent from pydantic_ai.tools import DeferredToolRequests @@ -125,7 +125,7 @@ def delete_file(path: str) -> str: When processing a follow-up request with approval responses, extract and pass the deferred tool results: -```py +```py {test="skip"} @app.post('/chat') async def chat(request: Request) -> Response: adapter = await VercelAIAdapter.from_request(request, agent=agent) From f18cda68c21f11edb9f687eb77c3ee935f3af9de Mon Sep 17 00:00:00 2001 From: bendrucker Date: Fri, 19 Dec 2025 11:33:55 -0800 Subject: [PATCH 08/21] Also skip linting for incomplete doc snippets --- docs/ui/vercel-ai.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/ui/vercel-ai.md b/docs/ui/vercel-ai.md index 88238a240b..6c69061652 100644 --- a/docs/ui/vercel-ai.md +++ b/docs/ui/vercel-ai.md @@ -106,7 +106,7 @@ Pydantic AI supports [AI SDK's human-in-the-loop tool approval](https://ai-sdk.d To enable tool approval, define your tool with `requires_approval=True` and include [`DeferredToolRequests`][pydantic_ai.tools.DeferredToolRequests] in your agent's output types: -```py {test="skip"} +```py {test="skip" lint="skip"} from pydantic_ai import Agent from pydantic_ai.tools import DeferredToolRequests @@ -125,7 +125,7 @@ def delete_file(path: str) -> str: When processing a follow-up request with approval responses, extract and pass the deferred tool results: -```py {test="skip"} +```py {test="skip" lint="skip"} @app.post('/chat') async def chat(request: Request) -> Response: adapter = await VercelAIAdapter.from_request(request, agent=agent) From cf6dad7204887e77db046e3f8ec764b9e59c223e Mon Sep 17 00:00:00 2001 From: bendrucker Date: Fri, 19 Dec 2025 11:57:19 -0800 Subject: [PATCH 09/21] Add tests to cover edge cases for tool approval extraction --- tests/test_vercel_ai.py | 146 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 146 insertions(+) diff --git a/tests/test_vercel_ai.py b/tests/test_vercel_ai.py index 045b016307..9fafd4ecfb 100644 --- a/tests/test_vercel_ai.py +++ b/tests/test_vercel_ai.py @@ -1925,6 +1925,152 @@ def test_extract_deferred_tool_results_ignores_pending_approval(): assert 'responded_tool' in result.approvals +def test_extract_deferred_tool_results_skips_non_tool_parts(): + """Test that non-tool parts (like TextUIPart) are skipped when extracting approvals.""" + from pydantic_ai.tools import ToolApproved + from pydantic_ai.ui.vercel_ai.request_types import ( + DynamicToolInputAvailablePart, + ToolApprovalResponded, + ) + + messages = [ + UIMessage( + id='msg-1', + role='assistant', + parts=[ + TextUIPart(text='Some text response'), + DynamicToolInputAvailablePart( + tool_name='my_tool', + tool_call_id='tool_1', + input={'arg': 'value'}, + approval=ToolApprovalResponded(id='approval-1', approved=True), + ), + TextUIPart(text='More text'), + ], + ), + ] + + result = VercelAIAdapter.extract_deferred_tool_results(messages) + assert result is not None + assert len(result.approvals) == 1 + assert isinstance(result.approvals['tool_1'], ToolApproved) + + +def test_extract_deferred_tool_results_skips_empty_tool_call_id(): + """Test that tool parts with empty tool_call_id are skipped.""" + from pydantic_ai.ui.vercel_ai.request_types import ( + DynamicToolInputAvailablePart, + ToolApprovalResponded, + ) + + messages = [ + UIMessage( + id='msg-1', + role='assistant', + parts=[ + DynamicToolInputAvailablePart( + tool_name='my_tool', + tool_call_id='', + input={'arg': 'value'}, + approval=ToolApprovalResponded(id='approval-1', approved=True), + ), + ], + ), + ] + + result = VercelAIAdapter.extract_deferred_tool_results(messages) + # Should return None since there are no valid approvals + assert result is None + + +def test_extract_denied_tool_ids_skips_non_tool_parts(): + """Test that _extract_denied_tool_ids skips non-tool parts.""" + from pydantic_ai.ui.vercel_ai._event_stream import _extract_denied_tool_ids + from pydantic_ai.ui.vercel_ai.request_types import ( + DynamicToolInputAvailablePart, + ToolApprovalResponded, + ) + + messages = [ + UIMessage( + id='msg-1', + role='assistant', + parts=[ + TextUIPart(text='Some text'), + DynamicToolInputAvailablePart( + tool_name='my_tool', + tool_call_id='tool_1', + input={'arg': 'value'}, + approval=ToolApprovalResponded(id='approval-1', approved=False), + ), + ], + ), + ] + + denied_ids = _extract_denied_tool_ids(messages) + assert denied_ids == {'tool_1'} + + +def test_extract_denied_tool_ids_skips_empty_tool_call_id(): + """Test that _extract_denied_tool_ids skips parts with empty tool_call_id.""" + from pydantic_ai.ui.vercel_ai._event_stream import _extract_denied_tool_ids + from pydantic_ai.ui.vercel_ai.request_types import ( + DynamicToolInputAvailablePart, + ToolApprovalResponded, + ) + + messages = [ + UIMessage( + id='msg-1', + role='assistant', + parts=[ + DynamicToolInputAvailablePart( + tool_name='my_tool', + tool_call_id='', + input={'arg': 'value'}, + approval=ToolApprovalResponded(id='approval-1', approved=False), + ), + ], + ), + ] + + denied_ids = _extract_denied_tool_ids(messages) + assert denied_ids == set() + + +def test_extract_denied_tool_ids_skips_approved_tools(): + """Test that _extract_denied_tool_ids only extracts denied (not approved) tools.""" + from pydantic_ai.ui.vercel_ai._event_stream import _extract_denied_tool_ids + from pydantic_ai.ui.vercel_ai.request_types import ( + DynamicToolInputAvailablePart, + ToolApprovalResponded, + ) + + messages = [ + UIMessage( + id='msg-1', + role='assistant', + parts=[ + DynamicToolInputAvailablePart( + tool_name='my_tool', + tool_call_id='approved_tool', + input={'arg': 'value'}, + approval=ToolApprovalResponded(id='approval-1', approved=True), + ), + DynamicToolInputAvailablePart( + tool_name='my_tool', + tool_call_id='denied_tool', + input={'arg': 'value'}, + approval=ToolApprovalResponded(id='approval-2', approved=False), + ), + ], + ), + ] + + denied_ids = _extract_denied_tool_ids(messages) + assert denied_ids == {'denied_tool'} + + async def test_tool_output_denied_chunk_emission(): """Test that ToolOutputDeniedChunk is emitted when a tool call is denied.""" from pydantic_ai.tools import DeferredToolRequests From 2c8a0af026c28d9e639b98957ac952b917d6af9d Mon Sep 17 00:00:00 2001 From: bendrucker Date: Fri, 19 Dec 2025 12:12:55 -0800 Subject: [PATCH 10/21] Use public interface for denied_tool_ids tests --- tests/test_vercel_ai.py | 134 +++++++++++++++++++++------------------- 1 file changed, 70 insertions(+), 64 deletions(-) diff --git a/tests/test_vercel_ai.py b/tests/test_vercel_ai.py index 9fafd4ecfb..4e98dafa44 100644 --- a/tests/test_vercel_ai.py +++ b/tests/test_vercel_ai.py @@ -1983,92 +1983,98 @@ def test_extract_deferred_tool_results_skips_empty_tool_call_id(): assert result is None -def test_extract_denied_tool_ids_skips_non_tool_parts(): - """Test that _extract_denied_tool_ids skips non-tool parts.""" - from pydantic_ai.ui.vercel_ai._event_stream import _extract_denied_tool_ids +def test_denied_tool_ids_skips_non_tool_parts(): + """Test that denied_tool_ids property skips non-tool parts.""" from pydantic_ai.ui.vercel_ai.request_types import ( DynamicToolInputAvailablePart, ToolApprovalResponded, ) - messages = [ - UIMessage( - id='msg-1', - role='assistant', - parts=[ - TextUIPart(text='Some text'), - DynamicToolInputAvailablePart( - tool_name='my_tool', - tool_call_id='tool_1', - input={'arg': 'value'}, - approval=ToolApprovalResponded(id='approval-1', approved=False), - ), - ], - ), - ] + request = SubmitMessage( + id='req-1', + messages=[ + UIMessage( + id='msg-1', + role='assistant', + parts=[ + TextUIPart(text='Some text'), + DynamicToolInputAvailablePart( + tool_name='my_tool', + tool_call_id='tool_1', + input={'arg': 'value'}, + approval=ToolApprovalResponded(id='approval-1', approved=False), + ), + ], + ), + ], + ) - denied_ids = _extract_denied_tool_ids(messages) - assert denied_ids == {'tool_1'} + stream = VercelAIEventStream(run_input=request) + assert stream.denied_tool_ids == {'tool_1'} -def test_extract_denied_tool_ids_skips_empty_tool_call_id(): - """Test that _extract_denied_tool_ids skips parts with empty tool_call_id.""" - from pydantic_ai.ui.vercel_ai._event_stream import _extract_denied_tool_ids +def test_denied_tool_ids_skips_empty_tool_call_id(): + """Test that denied_tool_ids property skips parts with empty tool_call_id.""" from pydantic_ai.ui.vercel_ai.request_types import ( DynamicToolInputAvailablePart, ToolApprovalResponded, ) - messages = [ - UIMessage( - id='msg-1', - role='assistant', - parts=[ - DynamicToolInputAvailablePart( - tool_name='my_tool', - tool_call_id='', - input={'arg': 'value'}, - approval=ToolApprovalResponded(id='approval-1', approved=False), - ), - ], - ), - ] + request = SubmitMessage( + id='req-1', + messages=[ + UIMessage( + id='msg-1', + role='assistant', + parts=[ + DynamicToolInputAvailablePart( + tool_name='my_tool', + tool_call_id='', + input={'arg': 'value'}, + approval=ToolApprovalResponded(id='approval-1', approved=False), + ), + ], + ), + ], + ) - denied_ids = _extract_denied_tool_ids(messages) - assert denied_ids == set() + stream = VercelAIEventStream(run_input=request) + assert stream.denied_tool_ids == set() -def test_extract_denied_tool_ids_skips_approved_tools(): - """Test that _extract_denied_tool_ids only extracts denied (not approved) tools.""" - from pydantic_ai.ui.vercel_ai._event_stream import _extract_denied_tool_ids +def test_denied_tool_ids_skips_approved_tools(): + """Test that denied_tool_ids property only extracts denied (not approved) tools.""" from pydantic_ai.ui.vercel_ai.request_types import ( DynamicToolInputAvailablePart, ToolApprovalResponded, ) - messages = [ - UIMessage( - id='msg-1', - role='assistant', - parts=[ - DynamicToolInputAvailablePart( - tool_name='my_tool', - tool_call_id='approved_tool', - input={'arg': 'value'}, - approval=ToolApprovalResponded(id='approval-1', approved=True), - ), - DynamicToolInputAvailablePart( - tool_name='my_tool', - tool_call_id='denied_tool', - input={'arg': 'value'}, - approval=ToolApprovalResponded(id='approval-2', approved=False), - ), - ], - ), - ] + request = SubmitMessage( + id='req-1', + messages=[ + UIMessage( + id='msg-1', + role='assistant', + parts=[ + DynamicToolInputAvailablePart( + tool_name='my_tool', + tool_call_id='approved_tool', + input={'arg': 'value'}, + approval=ToolApprovalResponded(id='approval-1', approved=True), + ), + DynamicToolInputAvailablePart( + tool_name='my_tool', + tool_call_id='denied_tool', + input={'arg': 'value'}, + approval=ToolApprovalResponded(id='approval-2', approved=False), + ), + ], + ), + ], + ) - denied_ids = _extract_denied_tool_ids(messages) - assert denied_ids == {'denied_tool'} + stream = VercelAIEventStream(run_input=request) + assert stream.denied_tool_ids == {'denied_tool'} async def test_tool_output_denied_chunk_emission(): From f731bc78002431c1cd7caea889f7f08d64eea9dd Mon Sep 17 00:00:00 2001 From: bendrucker Date: Fri, 19 Dec 2025 12:29:20 -0800 Subject: [PATCH 11/21] Fix coverage: add pragma comments and caching test --- tests/test_vercel_ai.py | 38 ++++++++++++++++++++++++++++++++++++-- 1 file changed, 36 insertions(+), 2 deletions(-) diff --git a/tests/test_vercel_ai.py b/tests/test_vercel_ai.py index 4e98dafa44..25c0fc947f 100644 --- a/tests/test_vercel_ai.py +++ b/tests/test_vercel_ai.py @@ -1673,7 +1673,7 @@ async def stream_function( @agent.tool_plain(requires_approval=True) def delete_file(path: str) -> str: - return f'Deleted {path}' + return f'Deleted {path}' # pragma: no cover request = SubmitMessage( id='foo', @@ -2077,6 +2077,40 @@ def test_denied_tool_ids_skips_approved_tools(): assert stream.denied_tool_ids == {'denied_tool'} +def test_denied_tool_ids_caching(): + """Test that denied_tool_ids property caches the result.""" + from pydantic_ai.ui.vercel_ai.request_types import ( + DynamicToolInputAvailablePart, + ToolApprovalResponded, + ) + + request = SubmitMessage( + id='req-1', + messages=[ + UIMessage( + id='msg-1', + role='assistant', + parts=[ + DynamicToolInputAvailablePart( + tool_name='my_tool', + tool_call_id='tool_1', + input={'arg': 'value'}, + approval=ToolApprovalResponded(id='approval-1', approved=False), + ), + ], + ), + ], + ) + + stream = VercelAIEventStream(run_input=request) + # First call computes and caches + first_result = stream.denied_tool_ids + # Second call returns cached value + second_result = stream.denied_tool_ids + assert first_result is second_result + assert first_result == {'tool_1'} + + async def test_tool_output_denied_chunk_emission(): """Test that ToolOutputDeniedChunk is emitted when a tool call is denied.""" from pydantic_ai.tools import DeferredToolRequests @@ -2097,7 +2131,7 @@ async def stream_function( @agent.tool_plain(requires_approval=True) def delete_file(path: str) -> str: - return f'Deleted {path}' + return f'Deleted {path}' # pragma: no cover # Simulate a follow-up request where the user denied the tool request = SubmitMessage( From 296611a92d8f64aaa25631bbca2f38c97972c027 Mon Sep 17 00:00:00 2001 From: bendrucker Date: Fri, 19 Dec 2025 19:01:42 -0800 Subject: [PATCH 12/21] Address PR review: add tests and fix documentation - Add test for from_request() with tool_approval parameter - Add test verifying approval chunks not emitted when tool_approval=False - Add test for deferred_tool_results fallback from instance field - Add test for denied_tool_ids with ToolUIPart (builtin tools) - Fix docs link to point to VercelAIAdapter.from_request - Reword Tool Approval section to clarify AI SDK UI vs AI Elements --- docs/ui/vercel-ai.md | 58 +---- pydantic_ai_slim/pydantic_ai/ui/_adapter.py | 10 + .../pydantic_ai/ui/_event_stream.py | 3 + .../pydantic_ai/ui/vercel_ai/_adapter.py | 44 ++-- .../pydantic_ai/ui/vercel_ai/_event_stream.py | 8 +- tests/test_vercel_ai.py | 235 +++++++++++++++++- 6 files changed, 287 insertions(+), 71 deletions(-) diff --git a/docs/ui/vercel-ai.md b/docs/ui/vercel-ai.md index 6c69061652..3d82859e8f 100644 --- a/docs/ui/vercel-ai.md +++ b/docs/ui/vercel-ai.md @@ -1,9 +1,7 @@ # Vercel AI Data Stream Protocol -Pydantic AI natively supports the [Vercel AI Data Stream Protocol](https://ai-sdk.dev/docs/ai-sdk-ui/stream-protocol#data-stream-protocol) to receive agent run input from, and stream events to, a [Vercel AI Elements](https://ai-sdk.dev/elements) frontend. +Pydantic AI natively supports the [Vercel AI Data Stream Protocol](https://ai-sdk.dev/docs/ai-sdk-ui/stream-protocol#data-stream-protocol) to receive agent run input from, and stream events to, a frontend using [AI SDK UI](https://ai-sdk.dev/docs/ai-sdk-ui/overview) hooks like [`useChat`](https://ai-sdk.dev/docs/reference/ai-sdk-ui/use-chat). You can optionally use [AI Elements](https://ai-sdk.dev/elements) for pre-built UI components. -!!! note "AI SDK Version Compatibility" - The base protocol is compatible with AI SDK v5 and later. However, [tool approval](#tool-approval) (human-in-the-loop) features require **AI SDK v6 or later**. ## Usage @@ -39,7 +37,7 @@ async def chat(request: Request) -> Response: If you're using a web framework not based on Starlette (e.g. Django or Flask) or need fine-grained control over the input or output, you can create a `VercelAIAdapter` instance and directly use its methods, which can be chained to accomplish the same thing as the `VercelAIAdapter.dispatch_request()` class method shown above: 1. The [`VercelAIAdapter.build_run_input()`][pydantic_ai.ui.vercel_ai.VercelAIAdapter.build_run_input] class method takes the request body as bytes and returns a Vercel AI [`RequestData`][pydantic_ai.ui.vercel_ai.request_types.RequestData] run input object, which you can then pass to the [`VercelAIAdapter()`][pydantic_ai.ui.vercel_ai.VercelAIAdapter] constructor along with the agent. - - You can also use the [`VercelAIAdapter.from_request()`][pydantic_ai.ui.UIAdapter.from_request] class method to build an adapter directly from a Starlette/FastAPI request. + - You can also use the [`VercelAIAdapter.from_request()`][pydantic_ai.ui.vercel_ai.VercelAIAdapter.from_request] class method to build an adapter directly from a Starlette/FastAPI request. 2. The [`VercelAIAdapter.run_stream()`][pydantic_ai.ui.UIAdapter.run_stream] method runs the agent and returns a stream of Vercel AI events. It supports the same optional arguments as [`Agent.run_stream_events()`](../agents.md#running-agents) and an optional `on_complete` callback function that receives the completed [`AgentRunResult`][pydantic_ai.agent.AgentRunResult] and can optionally yield additional Vercel AI events. - You can also use [`VercelAIAdapter.run_stream_native()`][pydantic_ai.ui.UIAdapter.run_stream_native] to run the agent and return a stream of Pydantic AI events instead, which can then be transformed into Vercel AI events using [`VercelAIAdapter.transform_stream()`][pydantic_ai.ui.UIAdapter.transform_stream]. 3. The [`VercelAIAdapter.encode_stream()`][pydantic_ai.ui.UIAdapter.encode_stream] method encodes the stream of Vercel AI events as SSE (HTTP Server-Sent Events) strings, which you can then return as a streaming response. @@ -87,55 +85,21 @@ async def chat(request: Request) -> Response: ## Tool Approval -Pydantic AI supports [AI SDK's human-in-the-loop tool approval](https://ai-sdk.dev/cookbook/next/human-in-the-loop) workflow, allowing users to approve or deny tool executions before they run. +Pydantic AI supports human-in-the-loop tool approval workflows with AI SDK UI, allowing users to approve or deny tool executions before they run. See the [deferred tool calls documentation](../deferred-tools.md) for details on setting up tools that require approval. -!!! warning "Requires AI SDK v6" - Tool approval is an AI SDK v6 feature. The `tool-approval-request` and `tool-output-denied` stream chunks, along with the `approval` field on tool parts, are not available in AI SDK v5. - -### How It Works - -1. **Tool requests approval**: When an agent calls a tool with `requires_approval=True`, Pydantic AI emits a `tool-approval-request` chunk instead of executing the tool immediately. - -2. **User decides**: The AI SDK frontend displays an approval UI. The user can approve or deny the tool execution. - -3. **Response is sent**: The frontend sends the approval decision back to the server in a follow-up request. - -4. **Tool executes or is denied**: If approved, the tool runs normally. If denied, Pydantic AI emits a `tool-output-denied` chunk and informs the model that the tool was rejected. - -### Server-Side Setup - -To enable tool approval, define your tool with `requires_approval=True` and include [`DeferredToolRequests`][pydantic_ai.tools.DeferredToolRequests] in your agent's output types: - -```py {test="skip" lint="skip"} -from pydantic_ai import Agent -from pydantic_ai.tools import DeferredToolRequests - -agent: Agent[None, str | DeferredToolRequests] = Agent( - 'openai:gpt-5', - output_type=[str, DeferredToolRequests], -) - -@agent.tool_plain(requires_approval=True) -def delete_file(path: str) -> str: - """Delete a file from the filesystem.""" - # This won't execute until the user approves - os.remove(path) - return f'Deleted {path}' -``` - -When processing a follow-up request with approval responses, extract and pass the deferred tool results: +To enable tool approval streaming, pass `tool_approval=True` when creating the adapter: ```py {test="skip" lint="skip"} @app.post('/chat') async def chat(request: Request) -> Response: - adapter = await VercelAIAdapter.from_request(request, agent=agent) - return adapter.streaming_response( - adapter.run_stream(deferred_tool_results=adapter.deferred_tool_results) - ) + adapter = await VercelAIAdapter.from_request(request, agent=agent, tool_approval=True) + return adapter.streaming_response(adapter.run_stream()) ``` -### Client-Side Setup +When `tool_approval=True`, the adapter will: -On the frontend, use AI SDK v6's [`useChat`](https://v6.ai-sdk.dev/docs/ai-sdk-ui/use-chat) hook with the [`Confirmation`](https://ai-sdk.dev/elements/components/confirmation) component or the `addToolApprovalResponse` function to handle approval UI. +1. Emit `tool-approval-request` chunks when tools with `requires_approval=True` are called +2. Automatically extract approval responses from follow-up requests +3. Emit `tool-output-denied` chunks for rejected tools -See the [AI SDK Human-in-the-Loop Cookbook](https://ai-sdk.dev/cookbook/next/human-in-the-loop) for complete frontend examples. +On the frontend, AI SDK UI's [`useChat`](https://ai-sdk.dev/docs/reference/ai-sdk-ui/use-chat) hook handles the approval flow. You can use the [`Confirmation`](https://ai-sdk.dev/elements/components/confirmation) component from AI Elements for a pre-built approval UI, or build your own using the hook's `addToolResult` function. diff --git a/pydantic_ai_slim/pydantic_ai/ui/_adapter.py b/pydantic_ai_slim/pydantic_ai/ui/_adapter.py index 4ac609a07e..93310c34f3 100644 --- a/pydantic_ai_slim/pydantic_ai/ui/_adapter.py +++ b/pydantic_ai_slim/pydantic_ai/ui/_adapter.py @@ -125,6 +125,12 @@ class UIAdapter(ABC, Generic[RunInputT, MessageT, EventT, AgentDepsT, OutputData accept: str | None = None """The `Accept` header value of the request, used to determine how to encode the protocol-specific events for the streaming response.""" + tool_approval: bool = False + """Whether to enable tool approval streaming for human-in-the-loop workflows.""" + + deferred_tool_results: DeferredToolResults | None = None + """Deferred tool results extracted from the request, used for tool approval workflows.""" + @classmethod async def from_request( cls, request: Request, *, agent: AbstractAgent[AgentDepsT, OutputDataT] @@ -237,6 +243,10 @@ def run_stream_native( toolsets: Optional additional toolsets for this run. builtin_tools: Optional additional builtin tools to use for this run. """ + # Use instance field as fallback if not explicitly passed + if deferred_tool_results is None: + deferred_tool_results = self.deferred_tool_results + message_history = [*(message_history or []), *self.messages] toolset = self.toolset diff --git a/pydantic_ai_slim/pydantic_ai/ui/_event_stream.py b/pydantic_ai_slim/pydantic_ai/ui/_event_stream.py index 391cf06f2f..13b862147e 100644 --- a/pydantic_ai_slim/pydantic_ai/ui/_event_stream.py +++ b/pydantic_ai_slim/pydantic_ai/ui/_event_stream.py @@ -70,6 +70,9 @@ class UIEventStream(ABC, Generic[RunInputT, EventT, AgentDepsT, OutputDataT]): accept: str | None = None """The `Accept` header value of the request, used to determine how to encode the protocol-specific events for the streaming response.""" + tool_approval: bool = False + """Whether tool approval streaming is enabled for human-in-the-loop workflows.""" + message_id: str = field(default_factory=lambda: str(uuid4())) """The message ID to use for the next event.""" diff --git a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py index ed5d28a02d..37f3c5fad6 100644 --- a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py +++ b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py @@ -62,7 +62,9 @@ from .response_types import BaseChunk if TYPE_CHECKING: - pass + from starlette.requests import Request + + from ...agent import AbstractAgent __all__ = ['VercelAIAdapter'] @@ -79,28 +81,40 @@ def build_run_input(cls, body: bytes) -> RequestData: """Build a Vercel AI run input object from the request body.""" return request_data_ta.validate_json(body) + @classmethod + async def from_request( + cls, + request: Request, + *, + agent: AbstractAgent[AgentDepsT, OutputDataT], + tool_approval: bool = False, + ) -> VercelAIAdapter[AgentDepsT, OutputDataT]: + """Create a Vercel AI adapter from a request. + + Args: + request: The incoming Starlette/FastAPI request. + agent: The Pydantic AI agent to run. + tool_approval: Whether to enable tool approval streaming for human-in-the-loop workflows. + """ + run_input = cls.build_run_input(await request.body()) + deferred_tool_results = cls.extract_deferred_tool_results(run_input.messages) if tool_approval else None + return cls( + agent=agent, + run_input=run_input, + accept=request.headers.get('accept'), + tool_approval=tool_approval, + deferred_tool_results=deferred_tool_results, + ) + def build_event_stream(self) -> UIEventStream[RequestData, BaseChunk, AgentDepsT, OutputDataT]: """Build a Vercel AI event stream transformer.""" - return VercelAIEventStream(self.run_input, accept=self.accept) + return VercelAIEventStream(self.run_input, accept=self.accept, tool_approval=self.tool_approval) @cached_property def messages(self) -> list[ModelMessage]: """Pydantic AI messages from the Vercel AI run input.""" return self.load_messages(self.run_input.messages) - @cached_property - def deferred_tool_results(self) -> DeferredToolResults | None: - """Extract deferred tool results from tool parts with approval responses. - - When the Vercel AI SDK client responds to a tool-approval-request, it sends - the approval decision in the tool part's `approval` field. This method extracts - those responses and converts them to Pydantic AI's `DeferredToolResults` format. - - Returns: - DeferredToolResults if any tool parts have approval responses, None otherwise. - """ - return self.extract_deferred_tool_results(self.run_input.messages) - @classmethod def extract_deferred_tool_results(cls, messages: Sequence[UIMessage]) -> DeferredToolResults | None: """Extract deferred tool results from UI messages. diff --git a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py index 8cf18973ac..554ad5a512 100644 --- a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py +++ b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py @@ -140,9 +140,9 @@ async def handle_run_result(self, event: AgentRunResultEvent) -> AsyncIterator[B if pydantic_reason: self._finish_reason = _FINISH_REASON_MAP.get(pydantic_reason, 'unknown') - # Emit tool approval requests for deferred approvals + # Emit tool approval requests for deferred approvals (only when tool_approval is enabled) output = event.result.output - if isinstance(output, DeferredToolRequests): + if self.tool_approval and isinstance(output, DeferredToolRequests): for tool_call in output.approvals: yield ToolApprovalRequestChunk( approval_id=str(uuid4()), @@ -246,8 +246,8 @@ async def handle_function_tool_result(self, event: FunctionToolResultEvent) -> A part = event.result tool_call_id = part.tool_call_id - # Check if this tool was denied by the user - if tool_call_id in self.denied_tool_ids: + # Check if this tool was denied by the user (only when tool_approval is enabled) + if self.tool_approval and tool_call_id in self.denied_tool_ids: yield ToolOutputDeniedChunk(tool_call_id=tool_call_id) elif isinstance(part, RetryPromptPart): yield ToolOutputErrorChunk(tool_call_id=tool_call_id, error_text=part.model_response()) diff --git a/tests/test_vercel_ai.py b/tests/test_vercel_ai.py index 25c0fc947f..09ace703c9 100644 --- a/tests/test_vercel_ai.py +++ b/tests/test_vercel_ai.py @@ -1686,7 +1686,7 @@ def delete_file(path: str) -> str: ], ) - adapter = VercelAIAdapter(agent, request) + adapter = VercelAIAdapter(agent, request, tool_approval=True) events: list[str | dict[str, Any]] = [ '[DONE]' if '[DONE]' in event else json.loads(event.removeprefix('data: ')) async for event in adapter.encode_stream(adapter.run_stream()) @@ -1705,6 +1705,51 @@ def delete_file(path: str) -> str: UUID(approval_id) # Raises ValueError if not a valid UUID +async def test_tool_approval_false_does_not_emit_approval_chunks(): + """Test that ToolApprovalRequestChunk is NOT emitted when tool_approval=False.""" + from pydantic_ai.tools import DeferredToolRequests + + async def stream_function( + messages: list[ModelMessage], agent_info: AgentInfo + ) -> AsyncIterator[DeltaToolCalls | str]: + yield { + 0: DeltaToolCall( + name='delete_file', + json_args='{"path": "test.txt"}', + tool_call_id='delete_1', + ) + } + + agent: Agent[None, str | DeferredToolRequests] = Agent( + model=FunctionModel(stream_function=stream_function), output_type=[str, DeferredToolRequests] + ) + + @agent.tool_plain(requires_approval=True) + def delete_file(path: str) -> str: + return f'Deleted {path}' # pragma: no cover + + request = SubmitMessage( + id='foo', + messages=[ + UIMessage( + id='bar', + role='user', + parts=[TextUIPart(text='Delete test.txt')], + ), + ], + ) + + adapter = VercelAIAdapter(agent, request, tool_approval=False) + events: list[str | dict[str, Any]] = [ + '[DONE]' if '[DONE]' in event else json.loads(event.removeprefix('data: ')) + async for event in adapter.encode_stream(adapter.run_stream()) + ] + + # Verify tool-approval-request chunk is NOT emitted when tool_approval=False + approval_events = [e for e in events if isinstance(e, dict) and e.get('type') == 'tool-approval-request'] + assert len(approval_events) == 0 + + def test_extract_deferred_tool_results_approved(): """Test that approved tool calls are correctly extracted from UI messages.""" from pydantic_ai.tools import ToolApproved @@ -2111,6 +2156,33 @@ def test_denied_tool_ids_caching(): assert first_result == {'tool_1'} +def test_denied_tool_ids_with_builtin_tool_ui_part(): + """Test that denied_tool_ids works with ToolUIPart (provider-executed tools).""" + from pydantic_ai.ui.vercel_ai.request_types import ToolApprovalResponded + + request = SubmitMessage( + id='req-1', + messages=[ + UIMessage( + id='msg-1', + role='assistant', + parts=[ + ToolInputAvailablePart( + type='tool-web_search', + tool_call_id='builtin_tool_1', + input={'query': 'search term'}, + provider_executed=True, + approval=ToolApprovalResponded(id='approval-1', approved=False), + ), + ], + ), + ], + ) + + stream = VercelAIEventStream(run_input=request) + assert stream.denied_tool_ids == {'builtin_tool_1'} + + async def test_tool_output_denied_chunk_emission(): """Test that ToolOutputDeniedChunk is emitted when a tool call is denied.""" from pydantic_ai.tools import DeferredToolRequests @@ -2161,12 +2233,11 @@ def delete_file(path: str) -> str: ], ) - adapter = VercelAIAdapter(agent, request) + deferred_tool_results = VercelAIAdapter.extract_deferred_tool_results(request.messages) + adapter = VercelAIAdapter(agent, request, tool_approval=True, deferred_tool_results=deferred_tool_results) events: list[str | dict[str, Any]] = [ '[DONE]' if '[DONE]' in event else json.loads(event.removeprefix('data: ')) - async for event in adapter.encode_stream( - adapter.run_stream(deferred_tool_results=adapter.deferred_tool_results) - ) + async for event in adapter.encode_stream(adapter.run_stream()) ] # Verify tool-output-denied chunk is emitted @@ -3380,3 +3451,157 @@ async def test_adapter_dump_messages_with_cache_point(): } ] ) + + +def starlette_import_successful() -> bool: + """Check if starlette is installed.""" + try: + from starlette.requests import Request # noqa: F401 + + return True + except ImportError: + return False + + +@pytest.mark.skipif(not starlette_import_successful(), reason='Starlette not installed') +async def test_from_request_with_tool_approval_enabled(): + """Test that from_request correctly handles tool_approval parameter.""" + from unittest.mock import AsyncMock + + from starlette.requests import Request + + from pydantic_ai.ui.vercel_ai.request_types import ( + DynamicToolInputAvailablePart, + ToolApprovalResponded, + ) + + agent = Agent(TestModel()) + + # Case 1: tool_approval=True with approvals in messages + request_with_approvals = SubmitMessage( + trigger='submit-message', + id='test-id', + messages=[ + UIMessage( + id='msg-1', + role='assistant', + parts=[ + DynamicToolInputAvailablePart( + tool_call_id='delete_1', + tool_name='delete_file', + input={'file': 'test.txt'}, + approval=ToolApprovalResponded(id='approval-1', approved=True), + ), + ], + ), + ], + ) + + request_body_with_approvals = request_with_approvals.model_dump_json().encode() + mock_request_with_approvals = AsyncMock(spec=Request) + mock_request_with_approvals.body = AsyncMock(return_value=request_body_with_approvals) + mock_request_with_approvals.headers.get = lambda key: None + + adapter_with_approval = await VercelAIAdapter.from_request( + mock_request_with_approvals, + agent=agent, + tool_approval=True, + ) + + # Verify tool_approval is set + assert adapter_with_approval.tool_approval is True + # Verify deferred_tool_results is populated + assert adapter_with_approval.deferred_tool_results is not None + assert 'delete_1' in adapter_with_approval.deferred_tool_results.approvals + + # Case 2: tool_approval=False + mock_request_without_approval = AsyncMock(spec=Request) + mock_request_without_approval.body = AsyncMock(return_value=request_body_with_approvals) + mock_request_without_approval.headers.get = lambda key: None + + adapter_without_approval = await VercelAIAdapter.from_request( + mock_request_without_approval, + agent=agent, + tool_approval=False, + ) + + # Verify tool_approval is set + assert adapter_without_approval.tool_approval is False + # Verify deferred_tool_results is None when tool_approval=False + assert adapter_without_approval.deferred_tool_results is None + + # Case 3: tool_approval=True but no approvals in messages + request_no_approvals = SubmitMessage( + trigger='submit-message', + id='test-id-2', + messages=[ + UIMessage( + id='msg-2', + role='user', + parts=[TextUIPart(text='Hello')], + ), + ], + ) + + request_body_no_approvals = request_no_approvals.model_dump_json().encode() + mock_request_no_approvals = AsyncMock(spec=Request) + mock_request_no_approvals.body = AsyncMock(return_value=request_body_no_approvals) + mock_request_no_approvals.headers.get = lambda key: None + + adapter_no_approvals = await VercelAIAdapter.from_request( + mock_request_no_approvals, + agent=agent, + tool_approval=True, + ) + + # Verify tool_approval is set + assert adapter_no_approvals.tool_approval is True + # Verify deferred_tool_results is None when no approvals exist + assert adapter_no_approvals.deferred_tool_results is None + + +async def test_deferred_tool_results_fallback_from_instance(): + """Test that run_stream_native uses deferred_tool_results from instance when not explicitly passed.""" + from unittest.mock import patch + + from pydantic_ai.tools import DeferredToolResults, ToolApproved + + # Create test deferred_tool_results + test_deferred_results = DeferredToolResults(approvals={'tool_1': ToolApproved()}) + + # Create a basic request + request = SubmitMessage( + id='test-request', + messages=[ + UIMessage( + id='test-message', + role='user', + parts=[TextUIPart(text='Test message')], + ), + ], + ) + + # Create a simple agent + agent = Agent(model=TestModel()) + + # Create adapter with deferred_tool_results set on instance + adapter = VercelAIAdapter(agent, request, deferred_tool_results=test_deferred_results) + + # Mock the agent's run_stream_events to capture what deferred_tool_results is passed + captured_kwargs = {} + + async def mock_run_stream_events(**kwargs): + captured_kwargs.update(kwargs) + # Return empty async iterator + return + yield # pragma: no cover + + with patch.object(agent, 'run_stream_events', side_effect=mock_run_stream_events): + # Call run_stream_native WITHOUT passing deferred_tool_results + result = adapter.run_stream_native() + # Consume the iterator + async for _ in result: # pragma: no cover + pass + + # Verify that the instance's deferred_tool_results was used + assert captured_kwargs['deferred_tool_results'] is test_deferred_results From 2359f4e6d405ab2d89b4969d24ac5b42b5b19921 Mon Sep 17 00:00:00 2001 From: bendrucker Date: Fri, 19 Dec 2025 19:20:44 -0800 Subject: [PATCH 13/21] Address PR review comments for tool approval - Make _extract_deferred_tool_results private and inline into from_request - Make _denied_tool_ids a private cached_property - Simplify approval values to True/False (per reviewer suggestion) - Add AI SDK v6 requirement note to documentation and docstrings - Update tests for new API and simplified return values --- docs/ui/vercel-ai.md | 3 + .../pydantic_ai/ui/vercel_ai/_adapter.py | 49 ++++------ .../pydantic_ai/ui/vercel_ai/_event_stream.py | 40 +++----- .../ui/vercel_ai/response_types.py | 10 +- tests/test_vercel_ai.py | 92 +++++++++---------- 5 files changed, 85 insertions(+), 109 deletions(-) diff --git a/docs/ui/vercel-ai.md b/docs/ui/vercel-ai.md index 3d82859e8f..64d308ea50 100644 --- a/docs/ui/vercel-ai.md +++ b/docs/ui/vercel-ai.md @@ -85,6 +85,9 @@ async def chat(request: Request) -> Response: ## Tool Approval +!!! note + Tool approval requires AI SDK v6 or later on the frontend. + Pydantic AI supports human-in-the-loop tool approval workflows with AI SDK UI, allowing users to approve or deny tool executions before they run. See the [deferred tool calls documentation](../deferred-tools.md) for details on setting up tools that require approval. To enable tool approval streaming, pass `tool_approval=True` when creating the adapter: diff --git a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py index 37f3c5fad6..8bfd37e61a 100644 --- a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py +++ b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py @@ -35,7 +35,7 @@ VideoUrl, ) from ...output import OutputDataT -from ...tools import AgentDepsT, DeferredToolResults, ToolApproved, ToolDenied +from ...tools import AgentDepsT, DeferredToolApprovalResult, DeferredToolResults from .. import MessagesBuilder, UIAdapter, UIEventStream from ._event_stream import VercelAIEventStream from .request_types import ( @@ -97,7 +97,7 @@ async def from_request( tool_approval: Whether to enable tool approval streaming for human-in-the-loop workflows. """ run_input = cls.build_run_input(await request.body()) - deferred_tool_results = cls.extract_deferred_tool_results(run_input.messages) if tool_approval else None + deferred_tool_results = cls._extract_deferred_tool_results(run_input.messages) if tool_approval else None return cls( agent=agent, run_input=run_input, @@ -106,53 +106,36 @@ async def from_request( deferred_tool_results=deferred_tool_results, ) - def build_event_stream(self) -> UIEventStream[RequestData, BaseChunk, AgentDepsT, OutputDataT]: - """Build a Vercel AI event stream transformer.""" - return VercelAIEventStream(self.run_input, accept=self.accept, tool_approval=self.tool_approval) - - @cached_property - def messages(self) -> list[ModelMessage]: - """Pydantic AI messages from the Vercel AI run input.""" - return self.load_messages(self.run_input.messages) - @classmethod - def extract_deferred_tool_results(cls, messages: Sequence[UIMessage]) -> DeferredToolResults | None: - """Extract deferred tool results from UI messages. - - Args: - messages: The UI messages to scan for approval responses. - - Returns: - DeferredToolResults if any tool parts have approval responses, None otherwise. - """ - approvals: dict[str, bool | ToolApproved | ToolDenied] = {} - + def _extract_deferred_tool_results(cls, messages: Sequence[UIMessage]) -> DeferredToolResults | None: + """Extract deferred tool results from UI messages.""" + approvals: dict[str, bool | DeferredToolApprovalResult] = {} for msg in messages: if msg.role != 'assistant': continue - for part in msg.parts: if not isinstance(part, ToolUIPart | DynamicToolUIPart): continue - approval = part.approval - if approval is None or not isinstance(approval, ToolApprovalResponded): + if not isinstance(approval, ToolApprovalResponded): continue - tool_call_id = part.tool_call_id if not tool_call_id: continue - - if approval.approved: - approvals[tool_call_id] = ToolApproved() - else: - approvals[tool_call_id] = ToolDenied(message=approval.reason or 'The tool call was denied.') - + approvals[tool_call_id] = approval.approved if not approvals: return None - return DeferredToolResults(approvals=approvals) + def build_event_stream(self) -> UIEventStream[RequestData, BaseChunk, AgentDepsT, OutputDataT]: + """Build a Vercel AI event stream transformer.""" + return VercelAIEventStream(self.run_input, accept=self.accept, tool_approval=self.tool_approval) + + @cached_property + def messages(self) -> list[ModelMessage]: + """Pydantic AI messages from the Vercel AI run input.""" + return self.load_messages(self.run_input.messages) + @classmethod def load_messages(cls, messages: Sequence[UIMessage]) -> list[ModelMessage]: # noqa: C901 """Transform Vercel AI messages into Pydantic AI messages.""" diff --git a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py index 554ad5a512..6560c756ba 100644 --- a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py +++ b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py @@ -4,6 +4,7 @@ from collections.abc import AsyncIterator, Mapping from dataclasses import dataclass +from functools import cached_property from typing import Any from uuid import uuid4 @@ -33,7 +34,6 @@ RequestData, ToolApprovalResponded, ToolUIPart, - UIMessage, ) from .response_types import ( BaseChunk, @@ -80,37 +80,27 @@ def _json_dumps(obj: Any) -> str: return to_json(obj).decode('utf-8') -def _extract_denied_tool_ids(messages: list[UIMessage]) -> set[str]: - """Extract tool_call_ids that were denied from UI messages.""" - denied_ids: set[str] = set() - for msg in messages: - if msg.role != 'assistant': - continue - for part in msg.parts: - if not isinstance(part, ToolUIPart | DynamicToolUIPart): - continue - approval = part.approval - if isinstance(approval, ToolApprovalResponded) and not approval.approved: - tool_call_id = part.tool_call_id - if tool_call_id: - denied_ids.add(tool_call_id) - return denied_ids - - @dataclass class VercelAIEventStream(UIEventStream[RequestData, BaseChunk, AgentDepsT, OutputDataT]): """UI event stream transformer for the Vercel AI protocol.""" _step_started: bool = False _finish_reason: FinishReason = None - _denied_tool_ids: set[str] | None = None - @property - def denied_tool_ids(self) -> set[str]: + @cached_property + def _denied_tool_ids(self) -> set[str]: """Get the set of tool_call_ids that were denied by the user.""" - if self._denied_tool_ids is None: - self._denied_tool_ids = _extract_denied_tool_ids(self.run_input.messages) - return self._denied_tool_ids + denied_ids: set[str] = set() + for msg in self.run_input.messages: + if msg.role != 'assistant': + continue + for part in msg.parts: + if not isinstance(part, ToolUIPart | DynamicToolUIPart): + continue + tool_call_id = part.tool_call_id + if tool_call_id and isinstance(part.approval, ToolApprovalResponded) and not part.approval.approved: + denied_ids.add(tool_call_id) + return denied_ids @property def response_headers(self) -> Mapping[str, str] | None: @@ -247,7 +237,7 @@ async def handle_function_tool_result(self, event: FunctionToolResultEvent) -> A tool_call_id = part.tool_call_id # Check if this tool was denied by the user (only when tool_approval is enabled) - if self.tool_approval and tool_call_id in self.denied_tool_ids: + if self.tool_approval and tool_call_id in self._denied_tool_ids: yield ToolOutputDeniedChunk(tool_call_id=tool_call_id) elif isinstance(part, RetryPromptPart): yield ToolOutputErrorChunk(tool_call_id=tool_call_id, error_text=part.model_response()) diff --git a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/response_types.py b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/response_types.py index 1f333f22dd..68d92a30b3 100644 --- a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/response_types.py +++ b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/response_types.py @@ -151,7 +151,10 @@ class ToolOutputErrorChunk(BaseChunk): class ToolApprovalRequestChunk(BaseChunk): - """Tool approval request chunk for human-in-the-loop approval.""" + """Tool approval request chunk for human-in-the-loop approval. + + Requires AI SDK v6 or later. + """ type: Literal['tool-approval-request'] = 'tool-approval-request' approval_id: str @@ -159,7 +162,10 @@ class ToolApprovalRequestChunk(BaseChunk): class ToolOutputDeniedChunk(BaseChunk): - """Tool output denied chunk when user denies tool execution.""" + """Tool output denied chunk when user denies tool execution. + + Requires AI SDK v6 or later. + """ type: Literal['tool-output-denied'] = 'tool-output-denied' tool_call_id: str diff --git a/tests/test_vercel_ai.py b/tests/test_vercel_ai.py index 09ace703c9..00900b7644 100644 --- a/tests/test_vercel_ai.py +++ b/tests/test_vercel_ai.py @@ -1752,7 +1752,6 @@ def delete_file(path: str) -> str: def test_extract_deferred_tool_results_approved(): """Test that approved tool calls are correctly extracted from UI messages.""" - from pydantic_ai.tools import ToolApproved from pydantic_ai.ui.vercel_ai.request_types import ( DynamicToolInputAvailablePart, ToolApprovalResponded, @@ -1773,15 +1772,14 @@ def test_extract_deferred_tool_results_approved(): ), ] - result = VercelAIAdapter.extract_deferred_tool_results(messages) + result = VercelAIAdapter._extract_deferred_tool_results(messages) assert result is not None assert 'delete_1' in result.approvals - assert isinstance(result.approvals['delete_1'], ToolApproved) + assert result.approvals['delete_1'] is True def test_extract_deferred_tool_results_denied(): """Test that denied tool calls are correctly extracted from UI messages.""" - from pydantic_ai.tools import ToolDenied from pydantic_ai.ui.vercel_ai.request_types import ( DynamicToolInputAvailablePart, ToolApprovalResponded, @@ -1802,12 +1800,11 @@ def test_extract_deferred_tool_results_denied(): ), ] - result = VercelAIAdapter.extract_deferred_tool_results(messages) + result = VercelAIAdapter._extract_deferred_tool_results(messages) assert result is not None assert 'delete_1' in result.approvals - denial = result.approvals['delete_1'] - assert isinstance(denial, ToolDenied) - assert denial.message == 'User rejected deletion' + # Denied returns False + assert result.approvals['delete_1'] is False def test_extract_deferred_tool_results_no_approvals(): @@ -1820,13 +1817,12 @@ def test_extract_deferred_tool_results_no_approvals(): ), ] - result = VercelAIAdapter.extract_deferred_tool_results(messages) + result = VercelAIAdapter._extract_deferred_tool_results(messages) assert result is None def test_extract_deferred_tool_results_with_tool_ui_part(): """Test that ToolUIPart (builtin/provider-executed tools) is handled correctly.""" - from pydantic_ai.tools import ToolApproved, ToolDenied from pydantic_ai.ui.vercel_ai.request_types import ToolApprovalResponded messages = [ @@ -1851,19 +1847,16 @@ def test_extract_deferred_tool_results_with_tool_ui_part(): ), ] - result = VercelAIAdapter.extract_deferred_tool_results(messages) + result = VercelAIAdapter._extract_deferred_tool_results(messages) assert result is not None assert 'delete_1' in result.approvals assert 'read_1' in result.approvals - assert isinstance(result.approvals['delete_1'], ToolApproved) - denial = result.approvals['read_1'] - assert isinstance(denial, ToolDenied) - assert denial.message == 'Not allowed' + assert result.approvals['delete_1'] is True + assert result.approvals['read_1'] is False def test_extract_deferred_tool_results_denied_no_reason(): - """Test that denied tool calls use default message when no reason is provided.""" - from pydantic_ai.tools import ToolDenied + """Test that denied tool calls return False when no reason is provided.""" from pydantic_ai.ui.vercel_ai.request_types import ( DynamicToolInputAvailablePart, ToolApprovalResponded, @@ -1884,16 +1877,14 @@ def test_extract_deferred_tool_results_denied_no_reason(): ), ] - result = VercelAIAdapter.extract_deferred_tool_results(messages) + result = VercelAIAdapter._extract_deferred_tool_results(messages) assert result is not None - denial = result.approvals['delete_1'] - assert isinstance(denial, ToolDenied) - assert denial.message == 'The tool call was denied.' + # Denied without reason returns False + assert result.approvals['delete_1'] is False def test_extract_deferred_tool_results_multiple_approvals(): """Test that multiple tool approvals in a single message are all extracted.""" - from pydantic_ai.tools import ToolApproved, ToolDenied from pydantic_ai.ui.vercel_ai.request_types import ( DynamicToolInputAvailablePart, ToolApprovalResponded, @@ -1926,12 +1917,12 @@ def test_extract_deferred_tool_results_multiple_approvals(): ), ] - result = VercelAIAdapter.extract_deferred_tool_results(messages) + result = VercelAIAdapter._extract_deferred_tool_results(messages) assert result is not None assert len(result.approvals) == 3 - assert isinstance(result.approvals['tool_1'], ToolApproved) - assert isinstance(result.approvals['tool_2'], ToolDenied) - assert isinstance(result.approvals['tool_3'], ToolApproved) + assert result.approvals['tool_1'] is True + assert result.approvals['tool_2'] is False + assert result.approvals['tool_3'] is True def test_extract_deferred_tool_results_ignores_pending_approval(): @@ -1963,7 +1954,7 @@ def test_extract_deferred_tool_results_ignores_pending_approval(): ), ] - result = VercelAIAdapter.extract_deferred_tool_results(messages) + result = VercelAIAdapter._extract_deferred_tool_results(messages) assert result is not None # Only the responded approval should be extracted assert 'pending_tool' not in result.approvals @@ -1972,7 +1963,6 @@ def test_extract_deferred_tool_results_ignores_pending_approval(): def test_extract_deferred_tool_results_skips_non_tool_parts(): """Test that non-tool parts (like TextUIPart) are skipped when extracting approvals.""" - from pydantic_ai.tools import ToolApproved from pydantic_ai.ui.vercel_ai.request_types import ( DynamicToolInputAvailablePart, ToolApprovalResponded, @@ -1995,10 +1985,10 @@ def test_extract_deferred_tool_results_skips_non_tool_parts(): ), ] - result = VercelAIAdapter.extract_deferred_tool_results(messages) + result = VercelAIAdapter._extract_deferred_tool_results(messages) assert result is not None assert len(result.approvals) == 1 - assert isinstance(result.approvals['tool_1'], ToolApproved) + assert result.approvals['tool_1'] is True def test_extract_deferred_tool_results_skips_empty_tool_call_id(): @@ -2023,13 +2013,13 @@ def test_extract_deferred_tool_results_skips_empty_tool_call_id(): ), ] - result = VercelAIAdapter.extract_deferred_tool_results(messages) + result = VercelAIAdapter._extract_deferred_tool_results(messages) # Should return None since there are no valid approvals assert result is None def test_denied_tool_ids_skips_non_tool_parts(): - """Test that denied_tool_ids property skips non-tool parts.""" + """Test that _denied_tool_ids property skips non-tool parts.""" from pydantic_ai.ui.vercel_ai.request_types import ( DynamicToolInputAvailablePart, ToolApprovalResponded, @@ -2055,11 +2045,11 @@ def test_denied_tool_ids_skips_non_tool_parts(): ) stream = VercelAIEventStream(run_input=request) - assert stream.denied_tool_ids == {'tool_1'} + assert stream._denied_tool_ids == {'tool_1'} def test_denied_tool_ids_skips_empty_tool_call_id(): - """Test that denied_tool_ids property skips parts with empty tool_call_id.""" + """Test that _denied_tool_ids property skips parts with empty tool_call_id.""" from pydantic_ai.ui.vercel_ai.request_types import ( DynamicToolInputAvailablePart, ToolApprovalResponded, @@ -2084,11 +2074,11 @@ def test_denied_tool_ids_skips_empty_tool_call_id(): ) stream = VercelAIEventStream(run_input=request) - assert stream.denied_tool_ids == set() + assert stream._denied_tool_ids == set() def test_denied_tool_ids_skips_approved_tools(): - """Test that denied_tool_ids property only extracts denied (not approved) tools.""" + """Test that _denied_tool_ids property only extracts denied (not approved) tools.""" from pydantic_ai.ui.vercel_ai.request_types import ( DynamicToolInputAvailablePart, ToolApprovalResponded, @@ -2119,11 +2109,11 @@ def test_denied_tool_ids_skips_approved_tools(): ) stream = VercelAIEventStream(run_input=request) - assert stream.denied_tool_ids == {'denied_tool'} + assert stream._denied_tool_ids == {'denied_tool'} def test_denied_tool_ids_caching(): - """Test that denied_tool_ids property caches the result.""" + """Test that _denied_tool_ids property caches the result.""" from pydantic_ai.ui.vercel_ai.request_types import ( DynamicToolInputAvailablePart, ToolApprovalResponded, @@ -2149,15 +2139,15 @@ def test_denied_tool_ids_caching(): stream = VercelAIEventStream(run_input=request) # First call computes and caches - first_result = stream.denied_tool_ids + first_result = stream._denied_tool_ids # Second call returns cached value - second_result = stream.denied_tool_ids + second_result = stream._denied_tool_ids assert first_result is second_result assert first_result == {'tool_1'} def test_denied_tool_ids_with_builtin_tool_ui_part(): - """Test that denied_tool_ids works with ToolUIPart (provider-executed tools).""" + """Test that _denied_tool_ids works with ToolUIPart (provider-executed tools).""" from pydantic_ai.ui.vercel_ai.request_types import ToolApprovalResponded request = SubmitMessage( @@ -2180,7 +2170,7 @@ def test_denied_tool_ids_with_builtin_tool_ui_part(): ) stream = VercelAIEventStream(run_input=request) - assert stream.denied_tool_ids == {'builtin_tool_1'} + assert stream._denied_tool_ids == {'builtin_tool_1'} async def test_tool_output_denied_chunk_emission(): @@ -2233,7 +2223,7 @@ def delete_file(path: str) -> str: ], ) - deferred_tool_results = VercelAIAdapter.extract_deferred_tool_results(request.messages) + deferred_tool_results = VercelAIAdapter._extract_deferred_tool_results(request.messages) adapter = VercelAIAdapter(agent, request, tool_approval=True, deferred_tool_results=deferred_tool_results) events: list[str | dict[str, Any]] = [ '[DONE]' if '[DONE]' in event else json.loads(event.removeprefix('data: ')) @@ -3456,8 +3446,9 @@ async def test_adapter_dump_messages_with_cache_point(): def starlette_import_successful() -> bool: """Check if starlette is installed.""" try: - from starlette.requests import Request # noqa: F401 + from starlette.requests import Request + assert Request # Verify the import succeeded return True except ImportError: return False @@ -3475,6 +3466,9 @@ async def test_from_request_with_tool_approval_enabled(): ToolApprovalResponded, ) + def mock_header_get(key: str) -> str | None: + return None + agent = Agent(TestModel()) # Case 1: tool_approval=True with approvals in messages @@ -3500,7 +3494,7 @@ async def test_from_request_with_tool_approval_enabled(): request_body_with_approvals = request_with_approvals.model_dump_json().encode() mock_request_with_approvals = AsyncMock(spec=Request) mock_request_with_approvals.body = AsyncMock(return_value=request_body_with_approvals) - mock_request_with_approvals.headers.get = lambda key: None + mock_request_with_approvals.headers.get = mock_header_get adapter_with_approval = await VercelAIAdapter.from_request( mock_request_with_approvals, @@ -3517,7 +3511,7 @@ async def test_from_request_with_tool_approval_enabled(): # Case 2: tool_approval=False mock_request_without_approval = AsyncMock(spec=Request) mock_request_without_approval.body = AsyncMock(return_value=request_body_with_approvals) - mock_request_without_approval.headers.get = lambda key: None + mock_request_without_approval.headers.get = mock_header_get adapter_without_approval = await VercelAIAdapter.from_request( mock_request_without_approval, @@ -3546,7 +3540,7 @@ async def test_from_request_with_tool_approval_enabled(): request_body_no_approvals = request_no_approvals.model_dump_json().encode() mock_request_no_approvals = AsyncMock(spec=Request) mock_request_no_approvals.body = AsyncMock(return_value=request_body_no_approvals) - mock_request_no_approvals.headers.get = lambda key: None + mock_request_no_approvals.headers.get = mock_header_get adapter_no_approvals = await VercelAIAdapter.from_request( mock_request_no_approvals, @@ -3588,9 +3582,9 @@ async def test_deferred_tool_results_fallback_from_instance(): adapter = VercelAIAdapter(agent, request, deferred_tool_results=test_deferred_results) # Mock the agent's run_stream_events to capture what deferred_tool_results is passed - captured_kwargs = {} + captured_kwargs: dict[str, Any] = {} - async def mock_run_stream_events(**kwargs): + async def mock_run_stream_events(**kwargs: Any) -> None: captured_kwargs.update(kwargs) # Return empty async iterator return From fb5bc2ed064412cee1c8db74b08155d5c36b49c4 Mon Sep 17 00:00:00 2001 From: bendrucker Date: Fri, 19 Dec 2025 19:24:25 -0800 Subject: [PATCH 14/21] Use 'AI SDK UI v6' instead of 'AI SDK v6' for frontend requirement --- docs/ui/vercel-ai.md | 2 +- pydantic_ai_slim/pydantic_ai/ui/vercel_ai/response_types.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/ui/vercel-ai.md b/docs/ui/vercel-ai.md index 64d308ea50..f0da2110d3 100644 --- a/docs/ui/vercel-ai.md +++ b/docs/ui/vercel-ai.md @@ -86,7 +86,7 @@ async def chat(request: Request) -> Response: ## Tool Approval !!! note - Tool approval requires AI SDK v6 or later on the frontend. + Tool approval requires AI SDK UI v6 or later on the frontend. Pydantic AI supports human-in-the-loop tool approval workflows with AI SDK UI, allowing users to approve or deny tool executions before they run. See the [deferred tool calls documentation](../deferred-tools.md) for details on setting up tools that require approval. diff --git a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/response_types.py b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/response_types.py index 68d92a30b3..0757eb4848 100644 --- a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/response_types.py +++ b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/response_types.py @@ -3,7 +3,7 @@ Converted to Python from: https://github.com/vercel/ai/blob/ai%406.0.0-beta.159/packages/ai/src/ui-message-stream/ui-message-chunks.ts -Tool approval types (`ToolApprovalRequestChunk`, `ToolOutputDeniedChunk`) require AI SDK v6 or later. +Tool approval types (`ToolApprovalRequestChunk`, `ToolOutputDeniedChunk`) require AI SDK UI v6 or later. """ from abc import ABC @@ -153,7 +153,7 @@ class ToolOutputErrorChunk(BaseChunk): class ToolApprovalRequestChunk(BaseChunk): """Tool approval request chunk for human-in-the-loop approval. - Requires AI SDK v6 or later. + Requires AI SDK UI v6 or later. """ type: Literal['tool-approval-request'] = 'tool-approval-request' @@ -164,7 +164,7 @@ class ToolApprovalRequestChunk(BaseChunk): class ToolOutputDeniedChunk(BaseChunk): """Tool output denied chunk when user denies tool execution. - Requires AI SDK v6 or later. + Requires AI SDK UI v6 or later. """ type: Literal['tool-output-denied'] = 'tool-output-denied' From c077afbc1bea75612da0d3b38d9f339b627d0000 Mon Sep 17 00:00:00 2001 From: bendrucker Date: Fri, 19 Dec 2025 19:26:10 -0800 Subject: [PATCH 15/21] Remove unnecessary tool_call_id checks (always present) --- .../pydantic_ai/ui/vercel_ai/_adapter.py | 5 +- .../pydantic_ai/ui/vercel_ai/_event_stream.py | 5 +- tests/test_vercel_ai.py | 56 ------------------- 3 files changed, 3 insertions(+), 63 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py index 8bfd37e61a..1ce016671e 100644 --- a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py +++ b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py @@ -119,10 +119,7 @@ def _extract_deferred_tool_results(cls, messages: Sequence[UIMessage]) -> Deferr approval = part.approval if not isinstance(approval, ToolApprovalResponded): continue - tool_call_id = part.tool_call_id - if not tool_call_id: - continue - approvals[tool_call_id] = approval.approved + approvals[part.tool_call_id] = approval.approved if not approvals: return None return DeferredToolResults(approvals=approvals) diff --git a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py index 6560c756ba..bb854e29c6 100644 --- a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py +++ b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py @@ -97,9 +97,8 @@ def _denied_tool_ids(self) -> set[str]: for part in msg.parts: if not isinstance(part, ToolUIPart | DynamicToolUIPart): continue - tool_call_id = part.tool_call_id - if tool_call_id and isinstance(part.approval, ToolApprovalResponded) and not part.approval.approved: - denied_ids.add(tool_call_id) + if isinstance(part.approval, ToolApprovalResponded) and not part.approval.approved: + denied_ids.add(part.tool_call_id) return denied_ids @property diff --git a/tests/test_vercel_ai.py b/tests/test_vercel_ai.py index 00900b7644..46b5c462ab 100644 --- a/tests/test_vercel_ai.py +++ b/tests/test_vercel_ai.py @@ -1991,33 +1991,6 @@ def test_extract_deferred_tool_results_skips_non_tool_parts(): assert result.approvals['tool_1'] is True -def test_extract_deferred_tool_results_skips_empty_tool_call_id(): - """Test that tool parts with empty tool_call_id are skipped.""" - from pydantic_ai.ui.vercel_ai.request_types import ( - DynamicToolInputAvailablePart, - ToolApprovalResponded, - ) - - messages = [ - UIMessage( - id='msg-1', - role='assistant', - parts=[ - DynamicToolInputAvailablePart( - tool_name='my_tool', - tool_call_id='', - input={'arg': 'value'}, - approval=ToolApprovalResponded(id='approval-1', approved=True), - ), - ], - ), - ] - - result = VercelAIAdapter._extract_deferred_tool_results(messages) - # Should return None since there are no valid approvals - assert result is None - - def test_denied_tool_ids_skips_non_tool_parts(): """Test that _denied_tool_ids property skips non-tool parts.""" from pydantic_ai.ui.vercel_ai.request_types import ( @@ -2048,35 +2021,6 @@ def test_denied_tool_ids_skips_non_tool_parts(): assert stream._denied_tool_ids == {'tool_1'} -def test_denied_tool_ids_skips_empty_tool_call_id(): - """Test that _denied_tool_ids property skips parts with empty tool_call_id.""" - from pydantic_ai.ui.vercel_ai.request_types import ( - DynamicToolInputAvailablePart, - ToolApprovalResponded, - ) - - request = SubmitMessage( - id='req-1', - messages=[ - UIMessage( - id='msg-1', - role='assistant', - parts=[ - DynamicToolInputAvailablePart( - tool_name='my_tool', - tool_call_id='', - input={'arg': 'value'}, - approval=ToolApprovalResponded(id='approval-1', approved=False), - ), - ], - ), - ], - ) - - stream = VercelAIEventStream(run_input=request) - assert stream._denied_tool_ids == set() - - def test_denied_tool_ids_skips_approved_tools(): """Test that _denied_tool_ids property only extracts denied (not approved) tools.""" from pydantic_ai.ui.vercel_ai.request_types import ( From e2e967c8432d9527fc82a04632e0a8c2c8407387 Mon Sep 17 00:00:00 2001 From: bendrucker Date: Fri, 19 Dec 2025 19:35:53 -0800 Subject: [PATCH 16/21] Inline extraction into from_request and remove private method tests - Inline extract_deferred_tool_results logic into from_request() - Remove unit tests for private _extract_deferred_tool_results method - Remove unit tests for private _denied_tool_ids property - Update test_tool_output_denied_chunk_emission to use public interface --- .../pydantic_ai/ui/vercel_ai/_adapter.py | 37 +- tests/test_vercel_ai.py | 371 +----------------- 2 files changed, 20 insertions(+), 388 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py index 1ce016671e..c63efdf26d 100644 --- a/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py +++ b/pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py @@ -97,7 +97,24 @@ async def from_request( tool_approval: Whether to enable tool approval streaming for human-in-the-loop workflows. """ run_input = cls.build_run_input(await request.body()) - deferred_tool_results = cls._extract_deferred_tool_results(run_input.messages) if tool_approval else None + + # Extract deferred tool results from messages when tool_approval is enabled + deferred_tool_results: DeferredToolResults | None = None + if tool_approval: + approvals: dict[str, bool | DeferredToolApprovalResult] = {} + for msg in run_input.messages: + if msg.role != 'assistant': + continue + for part in msg.parts: + if not isinstance(part, ToolUIPart | DynamicToolUIPart): + continue + approval = part.approval + if not isinstance(approval, ToolApprovalResponded): + continue + approvals[part.tool_call_id] = approval.approved + if approvals: + deferred_tool_results = DeferredToolResults(approvals=approvals) + return cls( agent=agent, run_input=run_input, @@ -106,24 +123,6 @@ async def from_request( deferred_tool_results=deferred_tool_results, ) - @classmethod - def _extract_deferred_tool_results(cls, messages: Sequence[UIMessage]) -> DeferredToolResults | None: - """Extract deferred tool results from UI messages.""" - approvals: dict[str, bool | DeferredToolApprovalResult] = {} - for msg in messages: - if msg.role != 'assistant': - continue - for part in msg.parts: - if not isinstance(part, ToolUIPart | DynamicToolUIPart): - continue - approval = part.approval - if not isinstance(approval, ToolApprovalResponded): - continue - approvals[part.tool_call_id] = approval.approved - if not approvals: - return None - return DeferredToolResults(approvals=approvals) - def build_event_stream(self) -> UIEventStream[RequestData, BaseChunk, AgentDepsT, OutputDataT]: """Build a Vercel AI event stream transformer.""" return VercelAIEventStream(self.run_input, accept=self.accept, tool_approval=self.tool_approval) diff --git a/tests/test_vercel_ai.py b/tests/test_vercel_ai.py index 46b5c462ab..ba1e94e9fd 100644 --- a/tests/test_vercel_ai.py +++ b/tests/test_vercel_ai.py @@ -1750,376 +1750,9 @@ def delete_file(path: str) -> str: assert len(approval_events) == 0 -def test_extract_deferred_tool_results_approved(): - """Test that approved tool calls are correctly extracted from UI messages.""" - from pydantic_ai.ui.vercel_ai.request_types import ( - DynamicToolInputAvailablePart, - ToolApprovalResponded, - ) - - messages = [ - UIMessage( - id='msg-1', - role='assistant', - parts=[ - DynamicToolInputAvailablePart( - tool_name='delete_file', - tool_call_id='delete_1', - input={'path': 'test.txt'}, - approval=ToolApprovalResponded(id='approval-123', approved=True), - ), - ], - ), - ] - - result = VercelAIAdapter._extract_deferred_tool_results(messages) - assert result is not None - assert 'delete_1' in result.approvals - assert result.approvals['delete_1'] is True - - -def test_extract_deferred_tool_results_denied(): - """Test that denied tool calls are correctly extracted from UI messages.""" - from pydantic_ai.ui.vercel_ai.request_types import ( - DynamicToolInputAvailablePart, - ToolApprovalResponded, - ) - - messages = [ - UIMessage( - id='msg-1', - role='assistant', - parts=[ - DynamicToolInputAvailablePart( - tool_name='delete_file', - tool_call_id='delete_1', - input={'path': 'test.txt'}, - approval=ToolApprovalResponded(id='approval-123', approved=False, reason='User rejected deletion'), - ), - ], - ), - ] - - result = VercelAIAdapter._extract_deferred_tool_results(messages) - assert result is not None - assert 'delete_1' in result.approvals - # Denied returns False - assert result.approvals['delete_1'] is False - - -def test_extract_deferred_tool_results_no_approvals(): - """Test that None is returned when no approval responses exist.""" - messages = [ - UIMessage( - id='msg-1', - role='user', - parts=[TextUIPart(text='Hello')], - ), - ] - - result = VercelAIAdapter._extract_deferred_tool_results(messages) - assert result is None - - -def test_extract_deferred_tool_results_with_tool_ui_part(): - """Test that ToolUIPart (builtin/provider-executed tools) is handled correctly.""" - from pydantic_ai.ui.vercel_ai.request_types import ToolApprovalResponded - - messages = [ - UIMessage( - id='msg-1', - role='assistant', - parts=[ - ToolInputAvailablePart( - type='tool-delete_file', - tool_call_id='delete_1', - input={'path': 'test.txt'}, - approval=ToolApprovalResponded(id='approval-123', approved=True), - ), - ToolOutputAvailablePart( - type='tool-read_file', - tool_call_id='read_1', - input={'path': 'other.txt'}, - output='file contents', - approval=ToolApprovalResponded(id='approval-456', approved=False, reason='Not allowed'), - ), - ], - ), - ] - - result = VercelAIAdapter._extract_deferred_tool_results(messages) - assert result is not None - assert 'delete_1' in result.approvals - assert 'read_1' in result.approvals - assert result.approvals['delete_1'] is True - assert result.approvals['read_1'] is False - - -def test_extract_deferred_tool_results_denied_no_reason(): - """Test that denied tool calls return False when no reason is provided.""" - from pydantic_ai.ui.vercel_ai.request_types import ( - DynamicToolInputAvailablePart, - ToolApprovalResponded, - ) - - messages = [ - UIMessage( - id='msg-1', - role='assistant', - parts=[ - DynamicToolInputAvailablePart( - tool_name='delete_file', - tool_call_id='delete_1', - input={'path': 'test.txt'}, - approval=ToolApprovalResponded(id='approval-123', approved=False), - ), - ], - ), - ] - - result = VercelAIAdapter._extract_deferred_tool_results(messages) - assert result is not None - # Denied without reason returns False - assert result.approvals['delete_1'] is False - - -def test_extract_deferred_tool_results_multiple_approvals(): - """Test that multiple tool approvals in a single message are all extracted.""" - from pydantic_ai.ui.vercel_ai.request_types import ( - DynamicToolInputAvailablePart, - ToolApprovalResponded, - ) - - messages = [ - UIMessage( - id='msg-1', - role='assistant', - parts=[ - DynamicToolInputAvailablePart( - tool_name='delete_file', - tool_call_id='tool_1', - input={'path': 'a.txt'}, - approval=ToolApprovalResponded(id='approval-1', approved=True), - ), - DynamicToolInputAvailablePart( - tool_name='delete_file', - tool_call_id='tool_2', - input={'path': 'b.txt'}, - approval=ToolApprovalResponded(id='approval-2', approved=False, reason='No'), - ), - DynamicToolInputAvailablePart( - tool_name='read_file', - tool_call_id='tool_3', - input={'path': 'c.txt'}, - approval=ToolApprovalResponded(id='approval-3', approved=True), - ), - ], - ), - ] - - result = VercelAIAdapter._extract_deferred_tool_results(messages) - assert result is not None - assert len(result.approvals) == 3 - assert result.approvals['tool_1'] is True - assert result.approvals['tool_2'] is False - assert result.approvals['tool_3'] is True - - -def test_extract_deferred_tool_results_ignores_pending_approval(): - """Test that ToolApprovalRequested (pending) is ignored, only ToolApprovalResponded is processed.""" - from pydantic_ai.ui.vercel_ai.request_types import ( - DynamicToolInputAvailablePart, - ToolApprovalRequested, - ToolApprovalResponded, - ) - - messages = [ - UIMessage( - id='msg-1', - role='assistant', - parts=[ - DynamicToolInputAvailablePart( - tool_name='delete_file', - tool_call_id='pending_tool', - input={'path': 'pending.txt'}, - approval=ToolApprovalRequested(id='approval-pending'), - ), - DynamicToolInputAvailablePart( - tool_name='delete_file', - tool_call_id='responded_tool', - input={'path': 'responded.txt'}, - approval=ToolApprovalResponded(id='approval-responded', approved=True), - ), - ], - ), - ] - - result = VercelAIAdapter._extract_deferred_tool_results(messages) - assert result is not None - # Only the responded approval should be extracted - assert 'pending_tool' not in result.approvals - assert 'responded_tool' in result.approvals - - -def test_extract_deferred_tool_results_skips_non_tool_parts(): - """Test that non-tool parts (like TextUIPart) are skipped when extracting approvals.""" - from pydantic_ai.ui.vercel_ai.request_types import ( - DynamicToolInputAvailablePart, - ToolApprovalResponded, - ) - - messages = [ - UIMessage( - id='msg-1', - role='assistant', - parts=[ - TextUIPart(text='Some text response'), - DynamicToolInputAvailablePart( - tool_name='my_tool', - tool_call_id='tool_1', - input={'arg': 'value'}, - approval=ToolApprovalResponded(id='approval-1', approved=True), - ), - TextUIPart(text='More text'), - ], - ), - ] - - result = VercelAIAdapter._extract_deferred_tool_results(messages) - assert result is not None - assert len(result.approvals) == 1 - assert result.approvals['tool_1'] is True - - -def test_denied_tool_ids_skips_non_tool_parts(): - """Test that _denied_tool_ids property skips non-tool parts.""" - from pydantic_ai.ui.vercel_ai.request_types import ( - DynamicToolInputAvailablePart, - ToolApprovalResponded, - ) - - request = SubmitMessage( - id='req-1', - messages=[ - UIMessage( - id='msg-1', - role='assistant', - parts=[ - TextUIPart(text='Some text'), - DynamicToolInputAvailablePart( - tool_name='my_tool', - tool_call_id='tool_1', - input={'arg': 'value'}, - approval=ToolApprovalResponded(id='approval-1', approved=False), - ), - ], - ), - ], - ) - - stream = VercelAIEventStream(run_input=request) - assert stream._denied_tool_ids == {'tool_1'} - - -def test_denied_tool_ids_skips_approved_tools(): - """Test that _denied_tool_ids property only extracts denied (not approved) tools.""" - from pydantic_ai.ui.vercel_ai.request_types import ( - DynamicToolInputAvailablePart, - ToolApprovalResponded, - ) - - request = SubmitMessage( - id='req-1', - messages=[ - UIMessage( - id='msg-1', - role='assistant', - parts=[ - DynamicToolInputAvailablePart( - tool_name='my_tool', - tool_call_id='approved_tool', - input={'arg': 'value'}, - approval=ToolApprovalResponded(id='approval-1', approved=True), - ), - DynamicToolInputAvailablePart( - tool_name='my_tool', - tool_call_id='denied_tool', - input={'arg': 'value'}, - approval=ToolApprovalResponded(id='approval-2', approved=False), - ), - ], - ), - ], - ) - - stream = VercelAIEventStream(run_input=request) - assert stream._denied_tool_ids == {'denied_tool'} - - -def test_denied_tool_ids_caching(): - """Test that _denied_tool_ids property caches the result.""" - from pydantic_ai.ui.vercel_ai.request_types import ( - DynamicToolInputAvailablePart, - ToolApprovalResponded, - ) - - request = SubmitMessage( - id='req-1', - messages=[ - UIMessage( - id='msg-1', - role='assistant', - parts=[ - DynamicToolInputAvailablePart( - tool_name='my_tool', - tool_call_id='tool_1', - input={'arg': 'value'}, - approval=ToolApprovalResponded(id='approval-1', approved=False), - ), - ], - ), - ], - ) - - stream = VercelAIEventStream(run_input=request) - # First call computes and caches - first_result = stream._denied_tool_ids - # Second call returns cached value - second_result = stream._denied_tool_ids - assert first_result is second_result - assert first_result == {'tool_1'} - - -def test_denied_tool_ids_with_builtin_tool_ui_part(): - """Test that _denied_tool_ids works with ToolUIPart (provider-executed tools).""" - from pydantic_ai.ui.vercel_ai.request_types import ToolApprovalResponded - - request = SubmitMessage( - id='req-1', - messages=[ - UIMessage( - id='msg-1', - role='assistant', - parts=[ - ToolInputAvailablePart( - type='tool-web_search', - tool_call_id='builtin_tool_1', - input={'query': 'search term'}, - provider_executed=True, - approval=ToolApprovalResponded(id='approval-1', approved=False), - ), - ], - ), - ], - ) - - stream = VercelAIEventStream(run_input=request) - assert stream._denied_tool_ids == {'builtin_tool_1'} - - async def test_tool_output_denied_chunk_emission(): """Test that ToolOutputDeniedChunk is emitted when a tool call is denied.""" - from pydantic_ai.tools import DeferredToolRequests + from pydantic_ai.tools import DeferredToolRequests, DeferredToolResults from pydantic_ai.ui.vercel_ai.request_types import ( DynamicToolInputAvailablePart, ToolApprovalResponded, @@ -2167,7 +1800,7 @@ def delete_file(path: str) -> str: ], ) - deferred_tool_results = VercelAIAdapter._extract_deferred_tool_results(request.messages) + deferred_tool_results = DeferredToolResults(approvals={'delete_1': False}) adapter = VercelAIAdapter(agent, request, tool_approval=True, deferred_tool_results=deferred_tool_results) events: list[str | dict[str, Any]] = [ '[DONE]' if '[DONE]' in event else json.loads(event.removeprefix('data: ')) From 33eabc74895fcfa66248513dc0d8060aedfcb457 Mon Sep 17 00:00:00 2001 From: bendrucker Date: Fri, 19 Dec 2025 19:49:02 -0800 Subject: [PATCH 17/21] Consolidate tool approval tests to use from_request() - Update test_tool_output_denied_chunk_emission to use from_request() with explicit type binding to test the full public interface - Remove test_from_request_with_tool_approval_enabled (now redundant) - Remove test_deferred_tool_results_fallback_from_instance (tested internal plumbing rather than observable behavior) --- tests/test_vercel_ai.py | 175 ++++++---------------------------------- 1 file changed, 23 insertions(+), 152 deletions(-) diff --git a/tests/test_vercel_ai.py b/tests/test_vercel_ai.py index ba1e94e9fd..3f3a607525 100644 --- a/tests/test_vercel_ai.py +++ b/tests/test_vercel_ai.py @@ -1750,9 +1750,18 @@ def delete_file(path: str) -> str: assert len(approval_events) == 0 +@pytest.mark.skipif(not starlette_import_successful, reason='Starlette is not installed') async def test_tool_output_denied_chunk_emission(): - """Test that ToolOutputDeniedChunk is emitted when a tool call is denied.""" - from pydantic_ai.tools import DeferredToolRequests, DeferredToolResults + """Test that ToolOutputDeniedChunk is emitted when a tool call is denied. + + This test verifies the full public interface: from_request() extracts approval + data from messages, and the adapter emits tool-output-denied chunks for denied tools. + """ + from unittest.mock import AsyncMock + + from starlette.requests import Request + + from pydantic_ai.tools import DeferredToolRequests from pydantic_ai.ui.vercel_ai.request_types import ( DynamicToolInputAvailablePart, ToolApprovalResponded, @@ -1764,9 +1773,7 @@ async def stream_function( # Model acknowledges the denial yield 'The file deletion was cancelled as requested.' - agent: Agent[None, str | DeferredToolRequests] = Agent( - model=FunctionModel(stream_function=stream_function), output_type=[str, DeferredToolRequests] - ) + agent = Agent(model=FunctionModel(stream_function=stream_function), output_type=[str, DeferredToolRequests]) @agent.tool_plain(requires_approval=True) def delete_file(path: str) -> str: @@ -1800,8 +1807,17 @@ def delete_file(path: str) -> str: ], ) - deferred_tool_results = DeferredToolResults(approvals={'delete_1': False}) - adapter = VercelAIAdapter(agent, request, tool_approval=True, deferred_tool_results=deferred_tool_results) + def mock_header_get(key: str) -> str | None: + return None + + request_body = request.model_dump_json().encode() + mock_request = AsyncMock(spec=Request) + mock_request.body = AsyncMock(return_value=request_body) + mock_request.headers.get = mock_header_get + + adapter = await VercelAIAdapter[None, str | DeferredToolRequests].from_request( + mock_request, agent=agent, tool_approval=True + ) events: list[str | dict[str, Any]] = [ '[DONE]' if '[DONE]' in event else json.loads(event.removeprefix('data: ')) async for event in adapter.encode_stream(adapter.run_stream()) @@ -3031,148 +3047,3 @@ def starlette_import_successful() -> bool: return False -@pytest.mark.skipif(not starlette_import_successful(), reason='Starlette not installed') -async def test_from_request_with_tool_approval_enabled(): - """Test that from_request correctly handles tool_approval parameter.""" - from unittest.mock import AsyncMock - - from starlette.requests import Request - - from pydantic_ai.ui.vercel_ai.request_types import ( - DynamicToolInputAvailablePart, - ToolApprovalResponded, - ) - - def mock_header_get(key: str) -> str | None: - return None - - agent = Agent(TestModel()) - - # Case 1: tool_approval=True with approvals in messages - request_with_approvals = SubmitMessage( - trigger='submit-message', - id='test-id', - messages=[ - UIMessage( - id='msg-1', - role='assistant', - parts=[ - DynamicToolInputAvailablePart( - tool_call_id='delete_1', - tool_name='delete_file', - input={'file': 'test.txt'}, - approval=ToolApprovalResponded(id='approval-1', approved=True), - ), - ], - ), - ], - ) - - request_body_with_approvals = request_with_approvals.model_dump_json().encode() - mock_request_with_approvals = AsyncMock(spec=Request) - mock_request_with_approvals.body = AsyncMock(return_value=request_body_with_approvals) - mock_request_with_approvals.headers.get = mock_header_get - - adapter_with_approval = await VercelAIAdapter.from_request( - mock_request_with_approvals, - agent=agent, - tool_approval=True, - ) - - # Verify tool_approval is set - assert adapter_with_approval.tool_approval is True - # Verify deferred_tool_results is populated - assert adapter_with_approval.deferred_tool_results is not None - assert 'delete_1' in adapter_with_approval.deferred_tool_results.approvals - - # Case 2: tool_approval=False - mock_request_without_approval = AsyncMock(spec=Request) - mock_request_without_approval.body = AsyncMock(return_value=request_body_with_approvals) - mock_request_without_approval.headers.get = mock_header_get - - adapter_without_approval = await VercelAIAdapter.from_request( - mock_request_without_approval, - agent=agent, - tool_approval=False, - ) - - # Verify tool_approval is set - assert adapter_without_approval.tool_approval is False - # Verify deferred_tool_results is None when tool_approval=False - assert adapter_without_approval.deferred_tool_results is None - - # Case 3: tool_approval=True but no approvals in messages - request_no_approvals = SubmitMessage( - trigger='submit-message', - id='test-id-2', - messages=[ - UIMessage( - id='msg-2', - role='user', - parts=[TextUIPart(text='Hello')], - ), - ], - ) - - request_body_no_approvals = request_no_approvals.model_dump_json().encode() - mock_request_no_approvals = AsyncMock(spec=Request) - mock_request_no_approvals.body = AsyncMock(return_value=request_body_no_approvals) - mock_request_no_approvals.headers.get = mock_header_get - - adapter_no_approvals = await VercelAIAdapter.from_request( - mock_request_no_approvals, - agent=agent, - tool_approval=True, - ) - - # Verify tool_approval is set - assert adapter_no_approvals.tool_approval is True - # Verify deferred_tool_results is None when no approvals exist - assert adapter_no_approvals.deferred_tool_results is None - - -async def test_deferred_tool_results_fallback_from_instance(): - """Test that run_stream_native uses deferred_tool_results from instance when not explicitly passed.""" - from unittest.mock import patch - - from pydantic_ai.tools import DeferredToolResults, ToolApproved - - # Create test deferred_tool_results - test_deferred_results = DeferredToolResults(approvals={'tool_1': ToolApproved()}) - - # Create a basic request - request = SubmitMessage( - id='test-request', - messages=[ - UIMessage( - id='test-message', - role='user', - parts=[TextUIPart(text='Test message')], - ), - ], - ) - - # Create a simple agent - agent = Agent(model=TestModel()) - - # Create adapter with deferred_tool_results set on instance - adapter = VercelAIAdapter(agent, request, deferred_tool_results=test_deferred_results) - - # Mock the agent's run_stream_events to capture what deferred_tool_results is passed - captured_kwargs: dict[str, Any] = {} - - async def mock_run_stream_events(**kwargs: Any) -> None: - captured_kwargs.update(kwargs) - # Return empty async iterator - return - yield # pragma: no cover - - with patch.object(agent, 'run_stream_events', side_effect=mock_run_stream_events): - # Call run_stream_native WITHOUT passing deferred_tool_results - result = adapter.run_stream_native() - # Consume the iterator - async for _ in result: # pragma: no cover - pass - - # Verify that the instance's deferred_tool_results was used - assert captured_kwargs['deferred_tool_results'] is test_deferred_results From e885d43cbe51408abc8e3ad722e3182316675968 Mon Sep 17 00:00:00 2001 From: bendrucker Date: Fri, 19 Dec 2025 19:53:04 -0800 Subject: [PATCH 18/21] Remove trailing blank lines (ruff) --- tests/test_vercel_ai.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/test_vercel_ai.py b/tests/test_vercel_ai.py index 3f3a607525..13024798c1 100644 --- a/tests/test_vercel_ai.py +++ b/tests/test_vercel_ai.py @@ -3045,5 +3045,3 @@ def starlette_import_successful() -> bool: return True except ImportError: return False - - From 676530be1363970ffc7066119a1bf69842733891 Mon Sep 17 00:00:00 2001 From: bendrucker Date: Sat, 20 Dec 2025 01:49:31 -0800 Subject: [PATCH 19/21] Add edge case tests for tool approval extraction coverage --- tests/test_vercel_ai.py | 129 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 129 insertions(+) diff --git a/tests/test_vercel_ai.py b/tests/test_vercel_ai.py index 13024798c1..b3a2d7b8d6 100644 --- a/tests/test_vercel_ai.py +++ b/tests/test_vercel_ai.py @@ -1792,6 +1792,13 @@ def delete_file(path: str) -> str: id='assistant-1', role='assistant', parts=[ + TextUIPart(text='I will delete the file for you.'), + DynamicToolInputAvailablePart( + tool_name='delete_file', + tool_call_id='delete_approved', + input={'path': 'approved.txt'}, + approval=ToolApprovalResponded(id='approval-456', approved=True), + ), DynamicToolInputAvailablePart( tool_name='delete_file', tool_call_id='delete_1', @@ -1832,6 +1839,128 @@ def mock_header_get(key: str) -> str | None: assert denied_event['toolCallId'] == 'delete_1' +@pytest.mark.skipif(not starlette_import_successful, reason='Starlette is not installed') +async def test_tool_approval_extraction_with_edge_cases(): + """Test that approval extraction correctly skips non-tool parts and non-responded approvals.""" + from unittest.mock import AsyncMock + + from starlette.requests import Request + + from pydantic_ai.tools import DeferredToolRequests + from pydantic_ai.ui.vercel_ai.request_types import ( + DynamicToolInputAvailablePart, + ToolApprovalRequested, + ToolApprovalResponded, + ) + + agent = Agent(TestModel(), output_type=[str, DeferredToolRequests]) + + @agent.tool_plain(requires_approval=True) + def some_tool(x: str) -> str: + return x # pragma: no cover + + request = SubmitMessage( + id='foo', + messages=[ + UIMessage(id='user-1', role='user', parts=[TextUIPart(text='Test')]), + UIMessage( + id='assistant-1', + role='assistant', + parts=[ + TextUIPart(text='Here is my response.'), + DynamicToolInputAvailablePart( + tool_name='some_tool', + tool_call_id='pending_tool', + input={'x': 'pending'}, + approval=ToolApprovalRequested(id='pending-approval'), + ), + DynamicToolInputAvailablePart( + tool_name='some_tool', + tool_call_id='no_approval_tool', + input={'x': 'no_approval'}, + approval=None, + ), + DynamicToolInputAvailablePart( + tool_name='some_tool', + tool_call_id='approved_tool', + input={'x': 'approved'}, + approval=ToolApprovalResponded(id='approved-id', approved=True), + ), + ], + ), + ], + ) + + def mock_header_get(key: str) -> str | None: + return None + + request_body = request.model_dump_json().encode() + mock_request = AsyncMock(spec=Request) + mock_request.body = AsyncMock(return_value=request_body) + mock_request.headers.get = mock_header_get + + adapter = await VercelAIAdapter[None, str | DeferredToolRequests].from_request( + mock_request, agent=agent, tool_approval=True + ) + + # Verify that only the responded approval was extracted + assert adapter.deferred_tool_results is not None + assert adapter.deferred_tool_results.approvals == {'approved_tool': True} + + +@pytest.mark.skipif(not starlette_import_successful, reason='Starlette is not installed') +async def test_tool_approval_no_approvals_extracted(): + """Test that deferred_tool_results is None when no approvals are responded.""" + from unittest.mock import AsyncMock + + from starlette.requests import Request + + from pydantic_ai.tools import DeferredToolRequests + from pydantic_ai.ui.vercel_ai.request_types import ( + DynamicToolInputAvailablePart, + ToolApprovalRequested, + ) + + agent = Agent(TestModel(), output_type=[str, DeferredToolRequests]) + + @agent.tool_plain(requires_approval=True) + def some_tool(x: str) -> str: + return x # pragma: no cover + + request = SubmitMessage( + id='foo', + messages=[ + UIMessage(id='user-1', role='user', parts=[TextUIPart(text='Test')]), + UIMessage( + id='assistant-1', + role='assistant', + parts=[ + DynamicToolInputAvailablePart( + tool_name='some_tool', + tool_call_id='pending_tool', + input={'x': 'pending'}, + approval=ToolApprovalRequested(id='pending-approval'), + ), + ], + ), + ], + ) + + def mock_header_get(key: str) -> str | None: + return None + + request_body = request.model_dump_json().encode() + mock_request = AsyncMock(spec=Request) + mock_request.body = AsyncMock(return_value=request_body) + mock_request.headers.get = mock_header_get + + adapter = await VercelAIAdapter[None, str | DeferredToolRequests].from_request( + mock_request, agent=agent, tool_approval=True + ) + + assert adapter.deferred_tool_results is None + + @pytest.mark.skipif(not starlette_import_successful, reason='Starlette is not installed') async def test_adapter_dispatch_request(): agent = Agent(model=TestModel()) From 9eebd8095ad9a5c2c71763261c2629ae7a73440e Mon Sep 17 00:00:00 2001 From: bendrucker Date: Sat, 20 Dec 2025 02:55:16 -0800 Subject: [PATCH 20/21] Fix coverage: remove duplicate starlette check function --- tests/test_vercel_ai.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/tests/test_vercel_ai.py b/tests/test_vercel_ai.py index b3a2d7b8d6..e5079372a4 100644 --- a/tests/test_vercel_ai.py +++ b/tests/test_vercel_ai.py @@ -3163,14 +3163,3 @@ async def test_adapter_dump_messages_with_cache_point(): } ] ) - - -def starlette_import_successful() -> bool: - """Check if starlette is installed.""" - try: - from starlette.requests import Request - - assert Request # Verify the import succeeded - return True - except ImportError: - return False From a908ef76a3c2b912b7bfa0b7a7f89db7bfb85b35 Mon Sep 17 00:00:00 2001 From: bendrucker Date: Sat, 20 Dec 2025 03:04:39 -0800 Subject: [PATCH 21/21] Add test for explicit deferred_tool_results parameter --- tests/test_vercel_ai.py | 44 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/tests/test_vercel_ai.py b/tests/test_vercel_ai.py index e5079372a4..69d52cfa4b 100644 --- a/tests/test_vercel_ai.py +++ b/tests/test_vercel_ai.py @@ -1961,6 +1961,50 @@ def mock_header_get(key: str) -> str | None: assert adapter.deferred_tool_results is None +@pytest.mark.skipif(not starlette_import_successful, reason='Starlette is not installed') +async def test_run_stream_with_explicit_deferred_tool_results(): + """Test that run_stream accepts explicit deferred_tool_results parameter.""" + from unittest.mock import AsyncMock + + from starlette.requests import Request + + from pydantic_ai.tools import DeferredToolResults + + async def stream_function( + messages: list[ModelMessage], agent_info: AgentInfo + ) -> AsyncIterator[DeltaToolCalls | str]: + yield 'Done' + + agent = Agent(model=FunctionModel(stream_function=stream_function)) + + request = SubmitMessage( + id='foo', + messages=[ + UIMessage(id='user-1', role='user', parts=[TextUIPart(text='Test')]), + ], + ) + + def mock_header_get(key: str) -> str | None: + return None + + request_body = request.model_dump_json().encode() + mock_request = AsyncMock(spec=Request) + mock_request.body = AsyncMock(return_value=request_body) + mock_request.headers.get = mock_header_get + + adapter = await VercelAIAdapter[None, str].from_request(mock_request, agent=agent) + + # Pass deferred_tool_results explicitly (even though it's empty, it covers the else branch) + explicit_results = DeferredToolResults() + events: list[str | dict[str, Any]] = [ + '[DONE]' if '[DONE]' in event else json.loads(event.removeprefix('data: ')) + async for event in adapter.encode_stream(adapter.run_stream(deferred_tool_results=explicit_results)) + ] + + # Verify stream completed successfully + assert events[-1] == '[DONE]' + + @pytest.mark.skipif(not starlette_import_successful, reason='Starlette is not installed') async def test_adapter_dispatch_request(): agent = Agent(model=TestModel())