coverage

Author: DouweM

pydantic_ai_slim/pydantic_ai/messages.py

@@ -523,7 +523,7 @@ def from_data_uri(cls, data_uri: str) -> Self:
         if not data_uri.startswith(prefix):
             raise ValueError('Data URI must start with "data:"')  # pragma: no cover
         media_type, data = data_uri[len(prefix) :].split(';base64,', 1)
-        return cls(data=base64.b64decode(data), media_type=media_type)
+        return cls.narrow_type(cls(data=base64.b64decode(data), media_type=media_type))
 
     @pydantic.computed_field
     @property

pydantic_ai_slim/pydantic_ai/ui/ag_ui/_event_stream.py

@@ -118,7 +118,7 @@ async def on_error(self, error: Exception) -> AsyncIterator[BaseEvent]:
     async def handle_text_start(self, part: TextPart, follows_text: bool = False) -> AsyncIterator[BaseEvent]:
         """Handle a TextPart at start."""
         if follows_text:
-            message_id = self.message_id  # TODO (DouweM): coverage
+            message_id = self.message_id
         else:
             message_id = self.new_message_id()
             yield TextMessageStartEvent(message_id=message_id)
@@ -133,7 +133,7 @@ async def handle_text_delta(self, delta: TextPartDelta) -> AsyncIterator[BaseEve
 
     async def handle_text_end(self, part: TextPart, followed_by_text: bool = False) -> AsyncIterator[BaseEvent]:
         """Handle a TextPart at end."""
-        if not followed_by_text:  # TODO (DouweM): coverage branch
+        if not followed_by_text:
             yield TextMessageEndEvent(message_id=self.message_id)
 
     async def handle_thinking_start(

pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py

@@ -6,7 +6,8 @@
 from dataclasses import dataclass
 from functools import cached_property
 
-from ...agent import AgentDepsT
+from typing_extensions import assert_never
+
 from ...messages import (
     AudioUrl,
     BinaryContent,
@@ -22,10 +23,12 @@
     ThinkingPart,
     ToolCallPart,
     ToolReturnPart,
+    UserContent,
     UserPromptPart,
     VideoUrl,
 )
 from ...output import OutputDataT
+from ...tools import AgentDepsT
 from ..adapter import UIAdapter
 from ..event_stream import UIEventStream
 from ..messages_builder import MessagesBuilder
@@ -36,6 +39,9 @@
     FileUIPart,
     ReasoningUIPart,
     RequestData,
+    SourceDocumentUIPart,
+    SourceUrlUIPart,
+    StepStartUIPart,
     TextUIPart,
     ToolOutputAvailablePart,
     ToolOutputErrorPart,
@@ -90,14 +96,16 @@ def load_messages(cls, messages: Sequence[UIMessage]) -> list[ModelMessage]:  #
         builder = MessagesBuilder()
 
         for msg in messages:
-            if msg.role in ('system', 'user'):
+            if msg.role == 'system':
+                for part in msg.parts:
+                    if isinstance(part, TextUIPart):  # pragma: no branch
+                        builder.add(SystemPromptPart(content=part.text))
+            elif msg.role == 'user':
+                user_prompt_content: str | list[UserContent] = []
                 for part in msg.parts:
                     if isinstance(part, TextUIPart):
-                        if msg.role == 'system':
-                            builder.add(SystemPromptPart(content=part.text))  # TODO (DouweM): coverage
-                        else:
-                            builder.add(UserPromptPart(content=part.text))
-                    elif isinstance(part, FileUIPart):  # TODO (DouweM): coverage
+                        user_prompt_content.append(part.text)
+                    elif isinstance(part, FileUIPart):
                         try:
                             file = BinaryContent.from_data_uri(part.url)
                         except ValueError:
@@ -111,28 +119,30 @@ def load_messages(cls, messages: Sequence[UIMessage]) -> list[ModelMessage]:  #
                                     file = AudioUrl(url=part.url, media_type=part.media_type)
                                 case _:
                                     file = DocumentUrl(url=part.url, media_type=part.media_type)
-                        builder.add(UserPromptPart(content=[file]))
+                        user_prompt_content.append(file)
 
-            elif msg.role == 'assistant':  # TODO (DouweM): coverage branch
+                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))  # TODO (DouweM): coverage
-                    elif isinstance(part, FileUIPart):  # TODO (DouweM): coverage
+                        builder.add(ThinkingPart(content=part.text))
+                    elif isinstance(part, FileUIPart):
                         try:
                             file = BinaryContent.from_data_uri(part.url)
-                        except ValueError as e:
+                        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, DataUIPart):
-                        # Not currently supported
-                        pass
-                    elif isinstance(part, ToolUIPart | DynamicToolUIPart):  # TODO (DouweM): coverage branch
-                        if isinstance(part, DynamicToolUIPart):  # TODO (DouweM): coverage
+                    elif isinstance(part, ToolUIPart | DynamicToolUIPart):
+                        if isinstance(part, DynamicToolUIPart):
                             tool_name = part.tool_name
                             builtin_tool = False
                         else:
@@ -142,15 +152,15 @@ def load_messages(cls, messages: Sequence[UIMessage]) -> list[ModelMessage]:  #
                         tool_call_id = part.tool_call_id
                         args = part.input
 
-                        if builtin_tool:  # TODO (DouweM): coverage
+                        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 = part.error_text
+                                    output = {'error_text': part.error_text, 'is_error': True}
 
                                 provider_name = (
                                     (part.call_provider_metadata or {}).get('pydantic_ai', {}).get('provider_name')
@@ -172,11 +182,27 @@ def load_messages(cls, messages: Sequence[UIMessage]) -> list[ModelMessage]:  #
                                 builder.add(
                                     ToolReturnPart(tool_name=tool_name, tool_call_id=tool_call_id, content=part.output)
                                 )
-                            elif part.state == 'output-error':  # TODO (DouweM): coverage
+                            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):
+                        # Contains custom data that shouldn't be sent to the model
+                        pass
+                    elif isinstance(part, SourceUrlUIPart):
+                        # TODO: Once we support citations: https://github.com/pydantic/pydantic-ai/issues/3126
+                        pass
+                    elif isinstance(part, SourceDocumentUIPart):
+                        # TODO: Once we support citations: https://github.com/pydantic/pydantic-ai/issues/3126
+                        pass
+                    elif isinstance(part, StepStartUIPart):
+                        # Nothing to do here
+                        pass
+                    else:
+                        assert_never(part)
+            else:
+                assert_never(msg.role)
 
         return builder.messages

pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py

@@ -80,19 +80,15 @@ async def before_stream(self) -> AsyncIterator[BaseChunk]:
 
     async def before_response(self) -> AsyncIterator[BaseChunk]:
         """Yield events before the request is processed."""
+        if self._step_started:
+            yield FinishStepChunk()
+
         self._step_started = True
         yield StartStepChunk()
 
-    async def after_request(self) -> AsyncIterator[BaseChunk]:
-        """Yield events after the response is processed."""
-        if self._step_started:  # TODO (DouweM): coverage
-            yield FinishStepChunk()
-            self._step_started = False
-
     async def after_stream(self) -> AsyncIterator[BaseChunk]:
         """Yield events after agent streaming completes."""
-        if self._step_started:  # TODO (DouweM): coverage branch
-            yield FinishStepChunk()
+        yield FinishStepChunk()
 
         yield FinishChunk()
         yield DoneChunk()
@@ -104,22 +100,22 @@ async def on_error(self, error: Exception) -> AsyncIterator[BaseChunk]:
     async def handle_text_start(self, part: TextPart, follows_text: bool = False) -> AsyncIterator[BaseChunk]:
         """Handle a TextPart at start."""
         if follows_text:
-            message_id = self.message_id  # TODO (DouweM): coverage
+            message_id = self.message_id
         else:
             message_id = self.new_message_id()
             yield TextStartChunk(id=message_id)
 
-        if part.content:  # TODO (DouweM): coverage branch
+        if part.content:
             yield TextDeltaChunk(id=message_id, delta=part.content)
 
     async def handle_text_delta(self, delta: TextPartDelta) -> AsyncIterator[BaseChunk]:
         """Handle a TextPartDelta."""
-        if delta.content_delta:  # TODO (DouweM): coverage branch
+        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]:
         """Handle a TextPart at end."""
-        if not followed_by_text:  # TODO (DouweM): coverage branch
+        if not followed_by_text:
             yield TextEndChunk(id=self.message_id)
 
     async def handle_thinking_start(
@@ -129,11 +125,11 @@ async def handle_thinking_start(
         message_id = self.new_message_id()
         yield ReasoningStartChunk(id=message_id)
         if part.content:
-            yield ReasoningDeltaChunk(id=message_id, delta=part.content)  # TODO (DouweM): coverage
+            yield ReasoningDeltaChunk(id=message_id, delta=part.content)
 
     async def handle_thinking_delta(self, delta: ThinkingPartDelta) -> AsyncIterator[BaseChunk]:
         """Handle a ThinkingPartDelta."""
-        if delta.content_delta:  # TODO (DouweM): coverage
+        if delta.content_delta:  # pragma: no branch
             yield ReasoningDeltaChunk(id=self.message_id, delta=delta.content_delta)
 
     async def handle_thinking_end(
@@ -164,9 +160,7 @@ async def _handle_tool_call_start(
             provider_executed=provider_executed,
         )
         if part.args:
-            yield ToolInputDeltaChunk(
-                tool_call_id=tool_call_id, input_text_delta=part.args_as_json_str()
-            )  # TODO (DouweM): coverage
+            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]:
         """Handle a ToolCallPartDelta."""
@@ -179,9 +173,7 @@ async def handle_tool_call_delta(self, delta: ToolCallPartDelta) -> AsyncIterato
 
     async def handle_tool_call_end(self, part: ToolCallPart) -> AsyncIterator[BaseChunk]:
         """Handle a ToolCallPart at end."""
-        yield ToolInputAvailableChunk(
-            tool_call_id=part.tool_call_id, tool_name=part.tool_name, input=part.args
-        )  # TODO (DouweM): coverage
+        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]:
         """Handle a BuiltinToolCallPart at end."""
@@ -204,11 +196,9 @@ async def handle_builtin_tool_return(self, part: BuiltinToolReturnPart) -> Async
     async def handle_file(self, part: FilePart) -> AsyncIterator[BaseChunk]:
         """Handle a FilePart."""
         file = part.content
-        yield FileChunk(url=file.data_uri, media_type=file.media_type)  # TODO (DouweM): coverage
+        yield FileChunk(url=file.data_uri, media_type=file.media_type)
 
-    async def handle_function_tool_result(
-        self, event: FunctionToolResultEvent
-    ) -> AsyncIterator[BaseChunk]:  # TODO (DouweM): coverage
+    async def handle_function_tool_result(self, event: FunctionToolResultEvent) -> AsyncIterator[BaseChunk]:
         """Handle a FunctionToolResultEvent, emitting tool result events."""
         result = event.result
         if isinstance(result, RetryPromptPart):