Update coverage

Author: dmontagu

pydantic_ai_slim/pydantic_ai/agent.py

@@ -1293,7 +1293,7 @@ def next_node(
             return next_node
         if _agent_graph.is_agent_node(next_node):
             return next_node
-        raise exceptions.AgentRunError(f'Unexpected node type: {type(next_node)}')
+        raise exceptions.AgentRunError(f'Unexpected node type: {type(next_node)}')  # pragma: no cover
 
     @property
     def result(self) -> AgentRunResult[ResultDataT] | None:

tests/test_streaming.py

@@ -2,14 +2,18 @@
 
 import datetime
 import json
+import re
 from collections.abc import AsyncIterator
+from copy import deepcopy
 from datetime import timezone
+from typing import Union
 
 import pytest
 from inline_snapshot import snapshot
 from pydantic import BaseModel
 
 from pydantic_ai import Agent, UnexpectedModelBehavior, UserError, capture_run_messages
+from pydantic_ai.agent import AgentRun
 from pydantic_ai.messages import (
     ModelMessage,
     ModelRequest,
@@ -22,7 +26,8 @@
 )
 from pydantic_ai.models.function import AgentInfo, DeltaToolCall, DeltaToolCalls, FunctionModel
 from pydantic_ai.models.test import TestModel
-from pydantic_ai.result import Usage
+from pydantic_ai.result import AgentStream, FinalResult, Usage
+from pydantic_graph import End
 
 from .conftest import IsNow
 
@@ -739,3 +744,109 @@ async def test_custom_result_type_default_structured() -> None:
     async with agent.run_stream('test', result_type=str) as result:
         response = await result.get_data()
         assert response == snapshot('success (no tool calls)')
+
+
+async def test_iter_stream_output():
+    m = TestModel(custom_result_text='The cat sat on the mat.')
+
+    agent = Agent(m)
+
+    @agent.result_validator
+    def result_validator_simple(data: str) -> str:
+        # Make a substitution in the validated results
+        return re.sub('cat sat', 'bat sat', data)
+
+    run: AgentRun
+    stream: AgentStream
+    messages: list[str] = []
+
+    stream_usage: Usage | None = None
+    with agent.iter('Hello') as run:
+        async for node in run:
+            if agent.is_model_request_node(node):
+                async with node.stream(run.ctx) as stream:
+                    async for chunk in stream.stream_output(debounce_by=None):
+                        messages.append(chunk)
+                stream_usage = deepcopy(stream.usage())
+    assert run.next_node == End(data=FinalResult(data='The bat sat on the mat.', tool_name=None))
+    assert (
+        run.usage()
+        == stream_usage
+        == Usage(requests=1, request_tokens=51, response_tokens=7, total_tokens=58, details=None)
+    )
+
+    assert messages == [
+        '',
+        'The ',
+        'The cat ',
+        'The bat sat ',
+        'The bat sat on ',
+        'The bat sat on the ',
+        'The bat sat on the mat.',
+        'The bat sat on the mat.',
+    ]
+
+
+async def test_iter_stream_responses():
+    m = TestModel(custom_result_text='The cat sat on the mat.')
+
+    agent = Agent(m)
+
+    @agent.result_validator
+    def result_validator_simple(data: str) -> str:
+        # Make a substitution in the validated results
+        return re.sub('cat sat', 'bat sat', data)
+
+    run: AgentRun
+    stream: AgentStream
+    messages: list[ModelResponse] = []
+    with agent.iter('Hello') as run:
+        async for node in run:
+            if agent.is_model_request_node(node):
+                async with node.stream(run.ctx) as stream:
+                    async for chunk in stream.stream_responses(debounce_by=None):
+                        messages.append(chunk)
+
+    assert messages == [
+        ModelResponse(
+            parts=[TextPart(content=text, part_kind='text')],
+            model_name='test',
+            timestamp=IsNow(tz=timezone.utc),
+            kind='response',
+        )
+        for text in [
+            '',
+            '',
+            'The ',
+            'The cat ',
+            'The cat sat ',
+            'The cat sat on ',
+            'The cat sat on the ',
+            'The cat sat on the mat.',
+        ]
+    ]
+
+    # Note: as you can see above, the result validator is not applied to the streamed responses, just the final result:
+    assert run.result is not None
+    assert run.result.data == 'The bat sat on the mat.'
+
+
+async def test_stream_iter_structured_validator() -> None:
+    class NotResultType(BaseModel):
+        not_value: str
+
+    agent = Agent[None, ResultType | NotResultType]('test', result_type=Union[ResultType, NotResultType])  # pyright: ignore[reportArgumentType]
+
+    @agent.result_validator
+    def result_validator(data: ResultType | NotResultType) -> ResultType | NotResultType:
+        assert isinstance(data, ResultType)
+        return ResultType(value=data.value + ' (validated)')
+
+    outputs: list[ResultType] = []
+    with agent.iter('test') as run:
+        async for node in run:
+            if agent.is_model_request_node(node):
+                async with node.stream(run.ctx) as stream:
+                    async for output in stream.stream_output(debounce_by=None):
+                        outputs.append(output)
+    assert outputs == [ResultType(value='a (validated)'), ResultType(value='a (validated)')]