Make AgentStream.stream_output (available inside agent.iter) stream validated output data instead of raising validation errors (#2134)

Author: DouweM

examples/pydantic_ai_examples/stream_whales.py

@@ -11,7 +11,7 @@
 from typing import Annotated
 
 import logfire
-from pydantic import Field, ValidationError
+from pydantic import Field
 from rich.console import Console
 from rich.live import Live
 from rich.table import Table
@@ -51,20 +51,7 @@ async def main():
         ) as result:
             console.print('Response:', style='green')
 
-            async for message, last in result.stream_structured(debounce_by=0.01):
-                try:
-                    whales = await result.validate_structured_output(
-                        message, allow_partial=not last
-                    )
-                except ValidationError as exc:
-                    if all(
-                        e['type'] == 'missing' and e['loc'] == ('response',)
-                        for e in exc.errors()
-                    ):
-                        continue
-                    else:
-                        raise
-
+            async for whales in result.stream(debounce_by=0.01):
                 table = Table(
                     title='Species of Whale',
                     caption='Streaming Structured responses from GPT-4',

pydantic_ai_slim/pydantic_ai/result.py

@@ -59,7 +59,12 @@ async def stream_output(self, *, debounce_by: float | None = 0.1) -> AsyncIterat
         """Asynchronously stream the (validated) agent outputs."""
         async for response in self.stream_responses(debounce_by=debounce_by):
             if self._final_result_event is not None:
-                yield await self._validate_response(response, self._final_result_event.tool_name, allow_partial=True)
+                try:
+                    yield await self._validate_response(
+                        response, self._final_result_event.tool_name, allow_partial=True
+                    )
+                except ValidationError:
+                    pass
         if self._final_result_event is not None:  # pragma: no branch
             yield await self._validate_response(
                 self._raw_stream_response.get(), self._final_result_event.tool_name, allow_partial=False

tests/test_streaming.py

@@ -1056,7 +1056,7 @@ def call_final_result_with_bad_data(messages: list[ModelMessage], info: AgentInf
     )
 
 
-async def test_stream_prompted_output():
+async def test_stream_structured_output():
     class CityLocation(BaseModel):
         city: str
         country: str | None = None
@@ -1079,6 +1079,30 @@ class CityLocation(BaseModel):
         assert result.is_complete
 
 
+async def test_iter_stream_structured_output():
+    class CityLocation(BaseModel):
+        city: str
+        country: str | None = None
+
+    m = TestModel(custom_output_text='{"city": "Mexico City", "country": "Mexico"}')
+
+    agent = Agent(m, output_type=PromptedOutput(CityLocation))
+
+    async with agent.iter('') as run:
+        async for node in run:
+            if agent.is_model_request_node(node):
+                async with node.stream(run.ctx) as stream:
+                    assert [c async for c in stream.stream_output(debounce_by=None)] == snapshot(
+                        [
+                            CityLocation(city='Mexico '),
+                            CityLocation(city='Mexico City'),
+                            CityLocation(city='Mexico City'),
+                            CityLocation(city='Mexico City', country='Mexico'),
+                            CityLocation(city='Mexico City', country='Mexico'),
+                        ]
+                    )
+
+
 def test_function_tool_event_tool_call_id_properties():
     """Ensure that the `tool_call_id` property on function tool events mirrors the underlying part's ID."""
     # Prepare a ToolCallPart with a fixed ID