coverage

Author: DouweM

pydantic_ai_slim/pydantic_ai/ui/adapter.py

@@ -190,16 +190,11 @@ async def process_stream(
         async for event in event_stream.handle_stream(stream):
             yield event
 
-        try:
-            result = event_stream.result
-            if on_complete is not None and result is not None:
+            if (result := event_stream.result) and on_complete is not None:
                 if _utils.is_async_callable(on_complete):
                     await on_complete(result)
                 else:
                     await _utils.run_in_executor(on_complete, result)
-        except Exception as e:  # TODO (DouweM): coverage
-            async for event in event_stream.on_error(e):
-                yield event
 
     async def run_stream(
         self,

pydantic_ai_slim/pydantic_ai/ui/event_stream.py

@@ -117,7 +117,7 @@ async def handle_stream(self, stream: AsyncIterator[SourceEvent]) -> AsyncIterat
                 if turn != next_turn:
                     if turn == 'request':
                         async for e in self.after_request():
-                            yield e  # TODO (DouweM): coverage
+                            yield e
                     elif turn == 'response':
                         async for e in self.after_response():
                             yield e  # TODO (DouweM): coverage
@@ -178,7 +178,7 @@ async def handle_event(self, event: SourceEvent) -> AsyncIterator[EventT]:  # no
             case FinalResultEvent():
                 self._final_result_event = event
                 async for e in self.handle_final_result(event):
-                    yield e  # TODO (DouweM): coverage
+                    yield e
             case FunctionToolCallEvent():
                 async for e in self.handle_function_tool_call(event):
                     yield e  # TODO (DouweM): coverage
@@ -204,7 +204,7 @@ async def handle_event(self, event: SourceEvent) -> AsyncIterator[EventT]:  # no
 
                 self.result = event.result
                 async for e in self.handle_run_result(event):
-                    yield e  # TODO (DouweM): coverage
+                    yield e
             case _:
                 pass
 
@@ -250,7 +250,7 @@ async def handle_part_delta(self, event: PartDeltaEvent) -> AsyncIterator[EventT
             case ThinkingPartDelta():
                 async for e in self.handle_thinking_delta(delta):
                     yield e
-            case ToolCallPartDelta():  # TODO (DouweM): coverage branch
+            case ToolCallPartDelta():  # pragma: no branch
                 async for e in self.handle_tool_call_delta(delta):
                     yield e
 
@@ -331,7 +331,7 @@ async def handle_thinking_delta(self, delta: ThinkingPartDelta) -> AsyncIterator
         Yields:
             Protocol-specific events.
         """
-        return
+        return  # TODO (DouweM): coverage
         yield  # Make this an async generator
 
     async def handle_thinking_end(

tests/test_ui.py

@@ -84,7 +84,7 @@ def event_stream(self) -> UIEventStream:
 
     @cached_property
     def messages(self) -> list[ModelMessage]:
-        return self.request.messages
+        return self.load_messages(self.request.messages)
 
     @cached_property
     def state(self) -> dict[str, Any] | None:
@@ -122,136 +122,75 @@ async def handle_part_end(self, event: PartEndEvent) -> AsyncIterator[str]:
 
     async def handle_text_start(self, part: TextPart, follows_text: bool = False) -> AsyncIterator[str]:
         yield f'<text>{part.content}'
-        async for e in super().handle_text_start(part, follows_text):
-            yield e
 
     async def handle_text_delta(self, delta: TextPartDelta) -> AsyncIterator[str]:
         yield delta.content_delta
-        async for e in super().handle_text_delta(delta):
-            yield e
 
     async def handle_text_end(self, part: TextPart, followed_by_text: bool = False) -> AsyncIterator[str]:
         yield '</text>'
-        async for e in super().handle_text_end(part, followed_by_text):
-            yield e
 
     async def handle_thinking_start(self, part: ThinkingPart, follows_thinking: bool = False) -> AsyncIterator[str]:
         yield f'<thinking>{part.content}'
-        async for e in super().handle_thinking_start(part, follows_thinking):
-            yield e
 
     async def handle_thinking_delta(self, delta: ThinkingPartDelta) -> AsyncIterator[str]:
         yield str(delta.content_delta)
-        async for e in super().handle_thinking_delta(delta):
-            yield e
 
     async def handle_thinking_end(self, part: ThinkingPart, followed_by_thinking: bool = False) -> AsyncIterator[str]:
         yield '</thinking>'
-        async for e in super().handle_thinking_end(part, followed_by_thinking):
-            yield e
 
     async def handle_tool_call_start(self, part: ToolCallPart) -> AsyncIterator[str]:
         yield f'<tool-call name={part.tool_name!r}>{part.args}'
-        async for e in super().handle_tool_call_start(part):
-            yield e
 
     async def handle_tool_call_delta(self, delta: ToolCallPartDelta) -> AsyncIterator[str]:
         yield str(delta.args_delta)
-        async for e in super().handle_tool_call_delta(delta):
-            yield e
 
     async def handle_tool_call_end(self, part: ToolCallPart) -> AsyncIterator[str]:
         yield f'</tool-call name={part.tool_name!r}>'
-        async for e in super().handle_tool_call_end(part):
-            yield e
 
     async def handle_builtin_tool_call_start(self, part: BuiltinToolCallPart) -> AsyncIterator[str]:
         yield f'<builtin-tool-call name={part.tool_name!r}>{part.args}'
-        async for e in super().handle_builtin_tool_call_start(part):
-            yield e
 
     async def handle_builtin_tool_call_end(self, part: BuiltinToolCallPart) -> AsyncIterator[str]:
         yield f'</builtin-tool-call name={part.tool_name!r}>'
-        async for e in super().handle_builtin_tool_call_end(part):
-            yield e
 
     async def handle_builtin_tool_return(self, part: BuiltinToolReturnPart) -> AsyncIterator[str]:
         yield f'<builtin-tool-return name={part.tool_name!r}>{part.content}</builtin-tool-return>'
-        async for e in super().handle_builtin_tool_return(part):
-            yield e
 
     async def handle_file(self, part: FilePart) -> AsyncIterator[str]:
         yield f'<file media_type={part.content.media_type!r} />'
-        async for e in super().handle_file(part):
-            yield e
 
     async def handle_final_result(self, event: FinalResultEvent) -> AsyncIterator[str]:
         yield '<final-result />'
-        async for e in super().handle_final_result(event):
-            yield e
 
     async def handle_function_tool_call(self, event: FunctionToolCallEvent) -> AsyncIterator[str]:
         yield f'<function-tool-call name={event.part.tool_name!r}>{event.part.args}</function-tool-call>'
-        async for e in super().handle_function_tool_call(event):
-            yield e
 
     async def handle_function_tool_result(self, event: FunctionToolResultEvent) -> AsyncIterator[str]:
         yield f'<function-tool-result name={event.result.tool_name!r}>{event.result.content}</function-tool-result>'
-        async for e in super().handle_function_tool_result(event):
-            yield e
 
     async def handle_run_result(self, event: AgentRunResultEvent) -> AsyncIterator[str]:
         yield f'<run-result>{event.result.output}</run-result>'
-        async for e in super().handle_run_result(event):
-            yield e
-
-    async def before_request(self) -> AsyncIterator[str]:
-        yield '<request>'
-        async for e in super().before_request():
-            yield e
 
-    async def after_request(self) -> AsyncIterator[str]:
-        yield '</request>'
-        async for e in super().after_request():
-            yield e
+    async def before_stream(self) -> AsyncIterator[str]:
+        yield '<stream>'
 
     async def before_response(self) -> AsyncIterator[str]:
         yield '<response>'
-        async for e in super().before_response():
-            yield e
 
     async def after_response(self) -> AsyncIterator[str]:
         yield '</response>'
-        async for e in super().after_response():
-            yield e
 
-    async def before_stream(self) -> AsyncIterator[str]:
-        yield '<stream>'
-        async for e in super().before_stream():
-            yield e
+    async def before_request(self) -> AsyncIterator[str]:
+        yield '<request>'
+
+    async def after_request(self) -> AsyncIterator[str]:
+        yield '</request>'
 
     async def after_stream(self) -> AsyncIterator[str]:
         yield '</stream>'
-        async for e in super().after_stream():
-            yield e
 
     async def on_error(self, error: Exception) -> AsyncIterator[str]:
-        yield f'on_error({error.__class__.__name__}({str(error)!r}))'
-        async for e in super().on_error(error):
-            yield e
-
-
-async def test_event_stream_text():
-    # text
-    #  - back to back
-    # thinking
-    #  - back to back
-    # tool call
-    # builtin tool call
-    # file
-    # error
-    # output tool
-    pass
+        yield f'<error type={error.__class__.__name__!r}>{str(error)}</error>'
 
 
 async def test_event_stream_builtin_tool_call():
@@ -313,3 +252,17 @@ async def stream_function(
             '</stream>',
         ]
     )
+
+
+async def test_event_stream_text():
+    # tool call
+    # error
+    # output tool
+    #  - dispatch_request
+    # text
+    #  - back to back
+    # thinking
+    #  - back to back
+    # tool call
+    # file
+    pass