pydantic-ai-slim 1.8.0__py3-none-any.whl → 1.9.0__py3-none-any.whl

pydantic_ai/ui/vercel_ai/_adapter.py

@@ -0,0 +1,199 @@
+"""Vercel AI adapter for handling requests."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from functools import cached_property
+from typing import TYPE_CHECKING
+
+from pydantic import TypeAdapter
+from typing_extensions import assert_never
+
+from ...messages import (
+    AudioUrl,
+    BinaryContent,
+    BuiltinToolCallPart,
+    BuiltinToolReturnPart,
+    DocumentUrl,
+    FilePart,
+    ImageUrl,
+    ModelMessage,
+    RetryPromptPart,
+    SystemPromptPart,
+    TextPart,
+    ThinkingPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserContent,
+    UserPromptPart,
+    VideoUrl,
+)
+from ...output import OutputDataT
+from ...tools import AgentDepsT
+from .. import MessagesBuilder, UIAdapter, UIEventStream
+from ._event_stream import VercelAIEventStream
+from .request_types import (
+    DataUIPart,
+    DynamicToolUIPart,
+    FileUIPart,
+    ReasoningUIPart,
+    RequestData,
+    SourceDocumentUIPart,
+    SourceUrlUIPart,
+    StepStartUIPart,
+    TextUIPart,
+    ToolOutputAvailablePart,
+    ToolOutputErrorPart,
+    ToolUIPart,
+    UIMessage,
+)
+from .response_types import BaseChunk
+
+if TYPE_CHECKING:
+    pass
+
+
+__all__ = ['VercelAIAdapter']
+
+request_data_ta: TypeAdapter[RequestData] = TypeAdapter(RequestData)
+
+
+@dataclass
+class VercelAIAdapter(UIAdapter[RequestData, UIMessage, BaseChunk, AgentDepsT, OutputDataT]):
+    """UI adapter for the Vercel AI protocol."""
+
+    @classmethod
+    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)
+
+    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)
+
+    @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."""
+        builder = MessagesBuilder()
+
+        for msg in messages:
+            if msg.role == 'system':
+                for part in msg.parts:
+                    if isinstance(part, TextUIPart):
+                        builder.add(SystemPromptPart(content=part.text))
+                    else:  # pragma: no cover
+                        raise ValueError(f'Unsupported system message part type: {type(part)}')
+            elif msg.role == 'user':
+                user_prompt_content: str | list[UserContent] = []
+                for part in msg.parts:
+                    if isinstance(part, TextUIPart):
+                        user_prompt_content.append(part.text)
+                    elif isinstance(part, FileUIPart):
+                        try:
+                            file = BinaryContent.from_data_uri(part.url)
+                        except ValueError:
+                            media_type_prefix = part.media_type.split('/', 1)[0]
+                            match media_type_prefix:
+                                case 'image':
+                                    file = ImageUrl(url=part.url, media_type=part.media_type)
+                                case 'video':
+                                    file = VideoUrl(url=part.url, media_type=part.media_type)
+                                case 'audio':
+                                    file = AudioUrl(url=part.url, media_type=part.media_type)
+                                case _:
+                                    file = DocumentUrl(url=part.url, media_type=part.media_type)
+                        user_prompt_content.append(file)
+                    else:  # pragma: no cover
+                        raise ValueError(f'Unsupported user message part type: {type(part)}')
+
+                if user_prompt_content:  # pragma: no branch
+                    if len(user_prompt_content) == 1 and isinstance(user_prompt_content[0], str):
+                        user_prompt_content = user_prompt_content[0]
+                    builder.add(UserPromptPart(content=user_prompt_content))
+
+            elif msg.role == 'assistant':
+                for part in msg.parts:
+                    if isinstance(part, TextUIPart):
+                        builder.add(TextPart(content=part.text))
+                    elif isinstance(part, ReasoningUIPart):
+                        builder.add(ThinkingPart(content=part.text))
+                    elif isinstance(part, FileUIPart):
+                        try:
+                            file = BinaryContent.from_data_uri(part.url)
+                        except ValueError as e:  # pragma: no cover
+                            # We don't yet handle non-data-URI file URLs returned by assistants, as no Pydantic AI models do this.
+                            raise ValueError(
+                                'Vercel AI integration can currently only handle assistant file parts with data URIs.'
+                            ) from e
+                        builder.add(FilePart(content=file))
+                    elif isinstance(part, ToolUIPart | DynamicToolUIPart):
+                        if isinstance(part, DynamicToolUIPart):
+                            tool_name = part.tool_name
+                            builtin_tool = False
+                        else:
+                            tool_name = part.type.removeprefix('tool-')
+                            builtin_tool = part.provider_executed
+
+                        tool_call_id = part.tool_call_id
+                        args = part.input
+
+                        if builtin_tool:
+                            call_part = BuiltinToolCallPart(tool_name=tool_name, tool_call_id=tool_call_id, args=args)
+                            builder.add(call_part)
+
+                            if isinstance(part, ToolOutputAvailablePart | ToolOutputErrorPart):
+                                if part.state == 'output-available':
+                                    output = part.output
+                                else:
+                                    output = {'error_text': part.error_text, 'is_error': True}
+
+                                provider_name = (
+                                    (part.call_provider_metadata or {}).get('pydantic_ai', {}).get('provider_name')
+                                )
+                                call_part.provider_name = provider_name
+
+                                builder.add(
+                                    BuiltinToolReturnPart(
+                                        tool_name=tool_name,
+                                        tool_call_id=tool_call_id,
+                                        content=output,
+                                        provider_name=provider_name,
+                                    )
+                                )
+                        else:
+                            builder.add(ToolCallPart(tool_name=tool_name, tool_call_id=tool_call_id, args=args))
+
+                            if part.state == 'output-available':
+                                builder.add(
+                                    ToolReturnPart(tool_name=tool_name, tool_call_id=tool_call_id, content=part.output)
+                                )
+                            elif part.state == 'output-error':
+                                builder.add(
+                                    RetryPromptPart(
+                                        tool_name=tool_name, tool_call_id=tool_call_id, content=part.error_text
+                                    )
+                                )
+                    elif isinstance(part, DataUIPart):  # pragma: no cover
+                        # Contains custom data that shouldn't be sent to the model
+                        pass
+                    elif isinstance(part, SourceUrlUIPart):  # pragma: no cover
+                        # TODO: Once we support citations: https://github.com/pydantic/pydantic-ai/issues/3126
+                        pass
+                    elif isinstance(part, SourceDocumentUIPart):  # pragma: no cover
+                        # TODO: Once we support citations: https://github.com/pydantic/pydantic-ai/issues/3126
+                        pass
+                    elif isinstance(part, StepStartUIPart):  # pragma: no cover
+                        # Nothing to do here
+                        pass
+                    else:
+                        assert_never(part)
+            else:
+                assert_never(msg.role)
+
+        return builder.messages

pydantic_ai/ui/vercel_ai/_event_stream.py

@@ -0,0 +1,187 @@
+"""Vercel AI event stream implementation."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Mapping
+from dataclasses import dataclass
+from typing import Any
+
+from pydantic_core import to_json
+
+from ...messages import (
+    BuiltinToolCallPart,
+    BuiltinToolReturnPart,
+    FilePart,
+    FunctionToolResultEvent,
+    RetryPromptPart,
+    TextPart,
+    TextPartDelta,
+    ThinkingPart,
+    ThinkingPartDelta,
+    ToolCallPart,
+    ToolCallPartDelta,
+)
+from ...output import OutputDataT
+from ...tools import AgentDepsT
+from .. import UIEventStream
+from .request_types import RequestData
+from .response_types import (
+    BaseChunk,
+    DoneChunk,
+    ErrorChunk,
+    FileChunk,
+    FinishChunk,
+    FinishStepChunk,
+    ReasoningDeltaChunk,
+    ReasoningEndChunk,
+    ReasoningStartChunk,
+    StartChunk,
+    StartStepChunk,
+    TextDeltaChunk,
+    TextEndChunk,
+    TextStartChunk,
+    ToolInputAvailableChunk,
+    ToolInputDeltaChunk,
+    ToolInputStartChunk,
+    ToolOutputAvailableChunk,
+    ToolOutputErrorChunk,
+)
+
+__all__ = ['VercelAIEventStream']
+
+# See https://ai-sdk.dev/docs/ai-sdk-ui/stream-protocol#data-stream-protocol
+VERCEL_AI_DSP_HEADERS = {'x-vercel-ai-ui-message-stream': 'v1'}
+
+
+def _json_dumps(obj: Any) -> str:
+    """Dump an object to JSON string."""
+    return to_json(obj).decode('utf-8')
+
+
+@dataclass
+class VercelAIEventStream(UIEventStream[RequestData, BaseChunk, AgentDepsT, OutputDataT]):
+    """UI event stream transformer for the Vercel AI protocol."""
+
+    _step_started: bool = False
+
+    @property
+    def response_headers(self) -> Mapping[str, str] | None:
+        return VERCEL_AI_DSP_HEADERS
+
+    def encode_event(self, event: BaseChunk) -> str:
+        return f'data: {event.encode()}\n\n'
+
+    async def before_stream(self) -> AsyncIterator[BaseChunk]:
+        yield StartChunk()
+
+    async def before_response(self) -> AsyncIterator[BaseChunk]:
+        if self._step_started:
+            yield FinishStepChunk()
+
+        self._step_started = True
+        yield StartStepChunk()
+
+    async def after_stream(self) -> AsyncIterator[BaseChunk]:
+        yield FinishStepChunk()
+
+        yield FinishChunk()
+        yield DoneChunk()
+
+    async def on_error(self, error: Exception) -> AsyncIterator[BaseChunk]:
+        yield ErrorChunk(error_text=str(error))
+
+    async def handle_text_start(self, part: TextPart, follows_text: bool = False) -> AsyncIterator[BaseChunk]:
+        if follows_text:
+            message_id = self.message_id
+        else:
+            message_id = self.new_message_id()
+            yield TextStartChunk(id=message_id)
+
+        if part.content:
+            yield TextDeltaChunk(id=message_id, delta=part.content)
+
+    async def handle_text_delta(self, delta: TextPartDelta) -> AsyncIterator[BaseChunk]:
+        if delta.content_delta:  # pragma: no branch
+            yield TextDeltaChunk(id=self.message_id, delta=delta.content_delta)
+
+    async def handle_text_end(self, part: TextPart, followed_by_text: bool = False) -> AsyncIterator[BaseChunk]:
+        if not followed_by_text:
+            yield TextEndChunk(id=self.message_id)
+
+    async def handle_thinking_start(
+        self, part: ThinkingPart, follows_thinking: bool = False
+    ) -> AsyncIterator[BaseChunk]:
+        message_id = self.new_message_id()
+        yield ReasoningStartChunk(id=message_id)
+        if part.content:
+            yield ReasoningDeltaChunk(id=message_id, delta=part.content)
+
+    async def handle_thinking_delta(self, delta: ThinkingPartDelta) -> AsyncIterator[BaseChunk]:
+        if delta.content_delta:  # pragma: no branch
+            yield ReasoningDeltaChunk(id=self.message_id, delta=delta.content_delta)
+
+    async def handle_thinking_end(
+        self, part: ThinkingPart, followed_by_thinking: bool = False
+    ) -> AsyncIterator[BaseChunk]:
+        yield ReasoningEndChunk(id=self.message_id)
+
+    def handle_tool_call_start(self, part: ToolCallPart | BuiltinToolCallPart) -> AsyncIterator[BaseChunk]:
+        return self._handle_tool_call_start(part)
+
+    def handle_builtin_tool_call_start(self, part: BuiltinToolCallPart) -> AsyncIterator[BaseChunk]:
+        return self._handle_tool_call_start(part, provider_executed=True)
+
+    async def _handle_tool_call_start(
+        self,
+        part: ToolCallPart | BuiltinToolCallPart,
+        tool_call_id: str | None = None,
+        provider_executed: bool | None = None,
+    ) -> AsyncIterator[BaseChunk]:
+        tool_call_id = tool_call_id or part.tool_call_id
+        yield ToolInputStartChunk(
+            tool_call_id=tool_call_id,
+            tool_name=part.tool_name,
+            provider_executed=provider_executed,
+        )
+        if part.args:
+            yield ToolInputDeltaChunk(tool_call_id=tool_call_id, input_text_delta=part.args_as_json_str())
+
+    async def handle_tool_call_delta(self, delta: ToolCallPartDelta) -> AsyncIterator[BaseChunk]:
+        tool_call_id = delta.tool_call_id or ''
+        assert tool_call_id, '`ToolCallPartDelta.tool_call_id` must be set'
+        yield ToolInputDeltaChunk(
+            tool_call_id=tool_call_id,
+            input_text_delta=delta.args_delta if isinstance(delta.args_delta, str) else _json_dumps(delta.args_delta),
+        )
+
+    async def handle_tool_call_end(self, part: ToolCallPart) -> AsyncIterator[BaseChunk]:
+        yield ToolInputAvailableChunk(tool_call_id=part.tool_call_id, tool_name=part.tool_name, input=part.args)
+
+    async def handle_builtin_tool_call_end(self, part: BuiltinToolCallPart) -> AsyncIterator[BaseChunk]:
+        yield ToolInputAvailableChunk(
+            tool_call_id=part.tool_call_id,
+            tool_name=part.tool_name,
+            input=part.args,
+            provider_executed=True,
+            provider_metadata={'pydantic_ai': {'provider_name': part.provider_name}},
+        )
+
+    async def handle_builtin_tool_return(self, part: BuiltinToolReturnPart) -> AsyncIterator[BaseChunk]:
+        yield ToolOutputAvailableChunk(
+            tool_call_id=part.tool_call_id,
+            output=part.content,
+            provider_executed=True,
+        )
+
+    async def handle_file(self, part: FilePart) -> AsyncIterator[BaseChunk]:
+        file = part.content
+        yield FileChunk(url=file.data_uri, media_type=file.media_type)
+
+    async def handle_function_tool_result(self, event: FunctionToolResultEvent) -> AsyncIterator[BaseChunk]:
+        result = event.result
+        if isinstance(result, RetryPromptPart):
+            yield ToolOutputErrorChunk(tool_call_id=result.tool_call_id, error_text=result.model_response())
+        else:
+            yield ToolOutputAvailableChunk(tool_call_id=result.tool_call_id, output=result.content)
+
+        # ToolCallResultEvent.content may hold user parts (e.g. text, images) that Vercel AI does not currently have events for

pydantic_ai/ui/vercel_ai/_utils.py

@@ -0,0 +1,16 @@
+"""Utilities for Vercel AI protocol.
+
+Converted to Python from:
+https://github.com/vercel/ai/blob/ai%405.0.34/packages/ai/src/ui/ui-messages.ts
+"""
+
+from abc import ABC
+
+from pydantic import BaseModel, ConfigDict
+from pydantic.alias_generators import to_camel
+
+
+class CamelBaseModel(BaseModel, ABC):
+    """Base model with camelCase aliases."""
+
+    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True, extra='forbid')

pydantic_ai/ui/vercel_ai/request_types.py

@@ -0,0 +1,275 @@
+"""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
+"""
+
+from abc import ABC
+from typing import Annotated, Any, Literal
+
+from pydantic import Discriminator, Field
+
+from ._utils import CamelBaseModel
+
+# Technically this is recursive union of JSON types; for simplicity, we call it Any
+JSONValue = Any
+ProviderMetadata = dict[str, dict[str, JSONValue]]
+"""Provider metadata."""
+
+
+class BaseUIPart(CamelBaseModel, ABC):
+    """Abstract base class for all UI parts."""
+
+
+class TextUIPart(BaseUIPart):
+    """A text part of a message."""
+
+    type: Literal['text'] = 'text'
+
+    text: str
+    """The text content."""
+
+    state: Literal['streaming', 'done'] | None = None
+    """The state of the text part."""
+
+    provider_metadata: ProviderMetadata | None = None
+    """The provider metadata."""
+
+
+class ReasoningUIPart(BaseUIPart):
+    """A reasoning part of a message."""
+
+    type: Literal['reasoning'] = 'reasoning'
+
+    text: str
+    """The reasoning text."""
+
+    state: Literal['streaming', 'done'] | None = None
+    """The state of the reasoning part."""
+
+    provider_metadata: ProviderMetadata | None = None
+    """The provider metadata."""
+
+
+class SourceUrlUIPart(BaseUIPart):
+    """A source part of a message."""
+
+    type: Literal['source-url'] = 'source-url'
+    source_id: str
+    url: str
+    title: str | None = None
+    provider_metadata: ProviderMetadata | None = None
+
+
+class SourceDocumentUIPart(BaseUIPart):
+    """A document source part of a message."""
+
+    type: Literal['source-document'] = 'source-document'
+    source_id: str
+    media_type: str
+    title: str
+    filename: str | None = None
+    provider_metadata: ProviderMetadata | None = None
+
+
+class FileUIPart(BaseUIPart):
+    """A file part of a message."""
+
+    type: Literal['file'] = 'file'
+
+    media_type: str
+    """
+    IANA media type of the file.
+    @see https://www.iana.org/assignments/media-types/media-types.xhtml
+    """
+
+    filename: str | None = None
+    """Optional filename of the file."""
+
+    url: str
+    """
+    The URL of the file.
+    It can either be a URL to a hosted file or a [Data URL](https://developer.mozilla.org/en-US/docs/Web/HTTP/Basics_of_HTTP/Data_URLs).
+    """
+
+    provider_metadata: ProviderMetadata | None = None
+    """The provider metadata."""
+
+
+class StepStartUIPart(BaseUIPart):
+    """A step boundary part of a message."""
+
+    type: Literal['step-start'] = 'step-start'
+
+
+class DataUIPart(BaseUIPart):
+    """Data part with dynamic type based on data name."""
+
+    type: Annotated[str, Field(pattern=r'^data-')]
+    id: str | None = None
+    data: Any
+
+
+# Tool part states as separate models
+class ToolInputStreamingPart(BaseUIPart):
+    """Tool part in input-streaming state."""
+
+    type: Annotated[str, Field(pattern=r'^tool-')]
+    tool_call_id: str
+    state: Literal['input-streaming'] = 'input-streaming'
+    input: Any | None = None
+    provider_executed: bool | None = None
+
+
+class ToolInputAvailablePart(BaseUIPart):
+    """Tool part in input-available state."""
+
+    type: Annotated[str, Field(pattern=r'^tool-')]
+    tool_call_id: str
+    state: Literal['input-available'] = 'input-available'
+    input: Any | None = None
+    provider_executed: bool | None = None
+    call_provider_metadata: ProviderMetadata | None = None
+
+
+class ToolOutputAvailablePart(BaseUIPart):
+    """Tool part in output-available state."""
+
+    type: Annotated[str, Field(pattern=r'^tool-')]
+    tool_call_id: str
+    state: Literal['output-available'] = 'output-available'
+    input: Any | None = None
+    output: Any | None = None
+    provider_executed: bool | None = None
+    call_provider_metadata: ProviderMetadata | None = None
+    preliminary: bool | None = None
+
+
+class ToolOutputErrorPart(BaseUIPart):
+    """Tool part in output-error state."""
+
+    type: Annotated[str, Field(pattern=r'^tool-')]
+    tool_call_id: str
+    state: Literal['output-error'] = 'output-error'
+    input: Any | None = None
+    raw_input: Any | None = None
+    error_text: str
+    provider_executed: bool | None = None
+    call_provider_metadata: ProviderMetadata | None = None
+
+
+ToolUIPart = ToolInputStreamingPart | ToolInputAvailablePart | ToolOutputAvailablePart | ToolOutputErrorPart
+"""Union of all tool part types."""
+
+
+# Dynamic tool part states as separate models
+class DynamicToolInputStreamingPart(BaseUIPart):
+    """Dynamic tool part in input-streaming state."""
+
+    type: Literal['dynamic-tool'] = 'dynamic-tool'
+    tool_name: str
+    tool_call_id: str
+    state: Literal['input-streaming'] = 'input-streaming'
+    input: Any | None = None
+
+
+class DynamicToolInputAvailablePart(BaseUIPart):
+    """Dynamic tool part in input-available state."""
+
+    type: Literal['dynamic-tool'] = 'dynamic-tool'
+    tool_name: str
+    tool_call_id: str
+    state: Literal['input-available'] = 'input-available'
+    input: Any
+    call_provider_metadata: ProviderMetadata | None = None
+
+
+class DynamicToolOutputAvailablePart(BaseUIPart):
+    """Dynamic tool part in output-available state."""
+
+    type: Literal['dynamic-tool'] = 'dynamic-tool'
+    tool_name: str
+    tool_call_id: str
+    state: Literal['output-available'] = 'output-available'
+    input: Any
+    output: Any
+    call_provider_metadata: ProviderMetadata | None = None
+    preliminary: bool | None = None
+
+
+class DynamicToolOutputErrorPart(BaseUIPart):
+    """Dynamic tool part in output-error state."""
+
+    type: Literal['dynamic-tool'] = 'dynamic-tool'
+    tool_name: str
+    tool_call_id: str
+    state: Literal['output-error'] = 'output-error'
+    input: Any
+    error_text: str
+    call_provider_metadata: ProviderMetadata | None = None
+
+
+DynamicToolUIPart = (
+    DynamicToolInputStreamingPart
+    | DynamicToolInputAvailablePart
+    | DynamicToolOutputAvailablePart
+    | DynamicToolOutputErrorPart
+)
+"""Union of all dynamic tool part types."""
+
+
+UIMessagePart = (
+    TextUIPart
+    | ReasoningUIPart
+    | ToolUIPart
+    | DynamicToolUIPart
+    | SourceUrlUIPart
+    | SourceDocumentUIPart
+    | FileUIPart
+    | DataUIPart
+    | StepStartUIPart
+)
+"""Union of all message part types."""
+
+
+class UIMessage(CamelBaseModel):
+    """A message as displayed in the UI by Vercel AI Elements."""
+
+    id: str
+    """A unique identifier for the message."""
+
+    role: Literal['system', 'user', 'assistant']
+    """The role of the message."""
+
+    metadata: Any | None = None
+    """The metadata of the message."""
+
+    parts: list[UIMessagePart]
+    """
+    The parts of the message. Use this for rendering the message in the UI.
+    System messages should be avoided (set the system prompt on the server instead).
+    They can have text parts.
+    User messages can have text parts and file parts.
+    Assistant messages can have text, reasoning, tool invocation, and file parts.
+    """
+
+
+class SubmitMessage(CamelBaseModel, extra='allow'):
+    """Submit message request."""
+
+    trigger: Literal['submit-message'] = 'submit-message'
+    id: str
+    messages: list[UIMessage]
+
+
+class RegenerateMessage(CamelBaseModel, extra='allow'):
+    """Ask the agent to regenerate a message."""
+
+    trigger: Literal['regenerate-message']
+    id: str
+    messages: list[UIMessage]
+    message_id: str
+
+
+RequestData = Annotated[SubmitMessage | RegenerateMessage, Discriminator('trigger')]
+"""Union of all request data types."""