fix final_flush coverage

Author: dsfaccini

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -569,7 +569,7 @@ class StreamedResponse(ABC):
     _event_iterator: AsyncIterator[ModelResponseStreamEvent] | None = field(default=None, init=False)
     _usage: RequestUsage = field(default_factory=RequestUsage, init=False)
 
-    def __aiter__(self) -> AsyncIterator[ModelResponseStreamEvent]:  # noqa: C901
+    def __aiter__(self) -> AsyncIterator[ModelResponseStreamEvent]:
         """Stream the response as an async iterable of [`ModelResponseStreamEvent`][pydantic_ai.messages.ModelResponseStreamEvent]s.
 
         This proxies the `_event_iterator()` and emits all events, while also checking for matches
@@ -616,15 +616,7 @@ def part_end_event(next_part: ModelResponsePart | None = None) -> PartEndEvent |
                         next_part_kind=next_part.part_kind if next_part else None,
                     )
 
-                async def chain_async_and_sync_iters(
-                    iter1: AsyncIterator[ModelResponseStreamEvent], iter2: Iterator[ModelResponseStreamEvent]
-                ) -> AsyncIterator[ModelResponseStreamEvent]:
-                    async for event in iter1:
-                        yield event
-                    for event in iter2:
-                        yield event
-
-                async for event in chain_async_and_sync_iters(iterator, self._parts_manager.final_flush()):
+                async for event in iterator:
                     if isinstance(event, PartStartEvent):
                         if last_start_event:
                             end_event = part_end_event(event.part)
@@ -642,8 +634,7 @@ async def chain_async_and_sync_iters(
 
             self._event_iterator = iterator_with_part_end(
                 iterator_with_final_event(
-                    # TODO chain_async_and_sync_iters(iterator, self._parts_manager.final_flush())
-                    self._get_event_iterator()
+                    chain_async_and_sync_iters(self._get_event_iterator(), self._parts_manager.final_flush())
                 )
             )
         return self._event_iterator
@@ -704,6 +695,16 @@ def timestamp(self) -> datetime:
         raise NotImplementedError()
 
 
+async def chain_async_and_sync_iters(
+    iter1: AsyncIterator[ModelResponseStreamEvent], iter2: Iterator[ModelResponseStreamEvent]
+) -> AsyncIterator[ModelResponseStreamEvent]:
+    """Chain an async iterator with a sync iterator."""
+    async for event in iter1:
+        yield event
+    for event in iter2:
+        yield event
+
+
 ALLOW_MODEL_REQUESTS = True
 """Whether to allow requests to models.
 

pydantic_ai_slim/pydantic_ai/models/function.py

@@ -186,6 +186,7 @@ async def request_stream(
 
         yield FunctionStreamedResponse(
             model_request_parameters=model_request_parameters,
+            _model_profile=self.profile,
             _model_name=self._model_name,
             _iter=response_stream,
         )
@@ -286,6 +287,7 @@ class FunctionStreamedResponse(StreamedResponse):
     """Implementation of `StreamedResponse` for [FunctionModel][pydantic_ai.models.function.FunctionModel]."""
 
     _model_name: str
+    _model_profile: ModelProfile
     _iter: AsyncIterator[str | DeltaToolCalls | DeltaThinkingCalls | BuiltinToolCallsReturns]
     _timestamp: datetime = field(default_factory=_utils.now_utc)
 
@@ -297,7 +299,9 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
             if isinstance(item, str):
                 response_tokens = _estimate_string_tokens(item)
                 self._usage += usage.RequestUsage(output_tokens=response_tokens)
-                for event in self._parts_manager.handle_text_delta(vendor_part_id='content', content=item):
+                for event in self._parts_manager.handle_text_delta(
+                    vendor_part_id='content', content=item, thinking_tags=self._model_profile.thinking_tags
+                ):
                     yield event
             elif isinstance(item, dict) and item:
                 for dtc_index, delta in item.items():

tests/test_parts_manager_thinking_tags.py

@@ -376,6 +376,16 @@ class Case:
         ],
         ignore_leading_whitespace=True,
     ),
+    Case(
+        name='new_part_ignore_whitespace_mixed_with_full_opening',
+        chunks=['  <think>'],
+        expected_parts=[TextPart('<think>')],
+        expected_normal_events=[],
+        expected_flushed_events=[
+            PartStartEvent(index=0, part=TextPart('<think>')),
+        ],
+        ignore_leading_whitespace=True,
+    ),
 ]
 
 # Category 9: No Vendor ID (updates, new after thinking, closings as text)

tests/test_streaming.py

@@ -2110,7 +2110,10 @@ async def stream_with_incomplete_thinking(
     ) -> AsyncIterator[str]:
         yield '<thi'
 
-    agent = Agent(FunctionModel(stream_function=stream_with_incomplete_thinking))
+    function_model = FunctionModel(stream_function=stream_with_incomplete_thinking)
+    function_model.profile.thinking_tags = ('<think>', '</think>')
+
+    agent = Agent(function_model)
 
     events: list[Any] = []
     async for event in agent.run_stream_events('Hello'):