models

- move finalize to aiter
- update models to the generator return type
parts manager
- disallow thinking after text
- delay emittion of thinking parts until there's content
tests
- swap out list calls for iteration
- add helper and consolidate tests to make them clearer

Author: dsfaccini

pydantic_ai_slim/pydantic_ai/_parts_manager.py

@@ -521,7 +521,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]:
+    def __aiter__(self) -> AsyncIterator[ModelResponseStreamEvent]:  # noqa: C901
         """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
@@ -580,6 +580,16 @@ def part_end_event(next_part: ModelResponsePart | None = None) -> PartEndEvent |
 
                     yield event
 
+                # Flush any buffered content and stream finalize events
+                for finalize_event in self._parts_manager.finalize():
+                    if isinstance(finalize_event, PartStartEvent):
+                        if last_start_event:
+                            end_event = part_end_event(finalize_event.part)
+                            if end_event:
+                                yield end_event
+                        last_start_event = finalize_event
+                    yield finalize_event
+
                 end_event = part_end_event()
                 if end_event:
                     yield end_event

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -734,19 +734,21 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                     ):
                         yield event_item
                 elif isinstance(current_block, BetaThinkingBlock):
-                    yield self._parts_manager.handle_thinking_delta(
+                    for e in self._parts_manager.handle_thinking_delta(
                         vendor_part_id=event.index,
                         content=current_block.thinking,
                         signature=current_block.signature,
                         provider_name=self.provider_name,
-                    )
+                    ):
+                        yield e
                 elif isinstance(current_block, BetaRedactedThinkingBlock):
-                    yield self._parts_manager.handle_thinking_delta(
+                    for e in self._parts_manager.handle_thinking_delta(
                         vendor_part_id=event.index,
                         id='redacted_thinking',
                         signature=current_block.data,
                         provider_name=self.provider_name,
-                    )
+                    ):
+                        yield e
                 elif isinstance(current_block, BetaToolUseBlock):
                     maybe_event = self._parts_manager.handle_tool_call_delta(
                         vendor_part_id=event.index,
@@ -807,17 +809,19 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                     ):
                         yield event_item
                 elif isinstance(event.delta, BetaThinkingDelta):
-                    yield self._parts_manager.handle_thinking_delta(
+                    for e in self._parts_manager.handle_thinking_delta(
                         vendor_part_id=event.index,
                         content=event.delta.thinking,
                         provider_name=self.provider_name,
-                    )
+                    ):
+                        yield e
                 elif isinstance(event.delta, BetaSignatureDelta):
-                    yield self._parts_manager.handle_thinking_delta(
+                    for e in self._parts_manager.handle_thinking_delta(
                         vendor_part_id=event.index,
                         signature=event.delta.signature,
                         provider_name=self.provider_name,
-                    )
+                    ):
+                        yield e
                 elif isinstance(event.delta, BetaInputJSONDelta):
                     maybe_event = self._parts_manager.handle_tool_call_delta(
                         vendor_part_id=event.index,

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -687,20 +687,22 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                     delta = content_block_delta['delta']
                     if 'reasoningContent' in delta:
                         if redacted_content := delta['reasoningContent'].get('redactedContent'):
-                            yield self._parts_manager.handle_thinking_delta(
+                            for e in self._parts_manager.handle_thinking_delta(
                                 vendor_part_id=index,
                                 id='redacted_content',
                                 signature=redacted_content.decode('utf-8'),
                                 provider_name=self.provider_name,
-                            )
+                            ):
+                                yield e
                         else:
                             signature = delta['reasoningContent'].get('signature')
-                            yield self._parts_manager.handle_thinking_delta(
+                            for e in self._parts_manager.handle_thinking_delta(
                                 vendor_part_id=index,
                                 content=delta['reasoningContent'].get('text'),
                                 signature=signature,
                                 provider_name=self.provider_name if signature else None,
-                            )
+                            ):
+                                yield e
                     if text := delta.get('text'):
                         for event in self._parts_manager.handle_text_delta(vendor_part_id=index, content=text):
                             yield event

pydantic_ai_slim/pydantic_ai/models/bedrock.py

@@ -284,7 +284,7 @@ class FunctionStreamedResponse(StreamedResponse):
     def __post_init__(self):
         self._usage += _estimate_usage([])
 
-    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:  # noqa: C901
         async for item in self._iter:
             if isinstance(item, str):
                 response_tokens = _estimate_string_tokens(item)
@@ -297,12 +297,13 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                         if delta.content:  # pragma: no branch
                             response_tokens = _estimate_string_tokens(delta.content)
                             self._usage += usage.RequestUsage(output_tokens=response_tokens)
-                        yield self._parts_manager.handle_thinking_delta(
+                        for e in self._parts_manager.handle_thinking_delta(
                             vendor_part_id=dtc_index,
                             content=delta.content,
                             signature=delta.signature,
                             provider_name='function' if delta.signature else None,
-                        )
+                        ):
+                            yield e
                     elif isinstance(delta, DeltaToolCall):
                         if delta.json_args:
                             response_tokens = _estimate_string_tokens(delta.json_args)

pydantic_ai_slim/pydantic_ai/models/function.py

@@ -668,15 +668,19 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
             for part in parts:
                 if part.thought_signature:
                     signature = base64.b64encode(part.thought_signature).decode('utf-8')
-                    yield self._parts_manager.handle_thinking_delta(
+                    for e in self._parts_manager.handle_thinking_delta(
                         vendor_part_id='thinking',
                         signature=signature,
                         provider_name=self.provider_name,
-                    )
+                    ):
+                        yield e
 
                 if part.text is not None:
                     if part.thought:
-                        yield self._parts_manager.handle_thinking_delta(vendor_part_id='thinking', content=part.text)
+                        for e in self._parts_manager.handle_thinking_delta(
+                            vendor_part_id='thinking', content=part.text
+                        ):
+                            yield e
                     else:
                         for event in self._parts_manager.handle_text_delta(vendor_part_id='content', content=part.text):
                             yield event

pydantic_ai_slim/pydantic_ai/models/google.py

@@ -547,9 +547,10 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                         reasoning = True
 
                     # NOTE: The `reasoning` field is only present if `groq_reasoning_format` is set to `parsed`.
-                    yield self._parts_manager.handle_thinking_delta(
+                    for e in self._parts_manager.handle_thinking_delta(
                         vendor_part_id=f'reasoning-{reasoning_index}', content=choice.delta.reasoning
-                    )
+                    ):
+                        yield e
                 else:
                     reasoning = False
 

pydantic_ai_slim/pydantic_ai/models/groq.py

@@ -1680,7 +1680,7 @@ class OpenAIStreamedResponse(StreamedResponse):
     _provider_name: str
     _provider_url: str
 
-    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:  # noqa: C901
         async for chunk in self._response:
             self._usage += _map_usage(chunk, self._provider_name, self._provider_url, self._model_name)
 
@@ -1706,23 +1706,25 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
             # The `reasoning_content` field is only present in DeepSeek models.
             # https://api-docs.deepseek.com/guides/reasoning_model
             if reasoning_content := getattr(choice.delta, 'reasoning_content', None):
-                yield self._parts_manager.handle_thinking_delta(
+                for e in self._parts_manager.handle_thinking_delta(
                     vendor_part_id='reasoning_content',
                     id='reasoning_content',
                     content=reasoning_content,
                     provider_name=self.provider_name,
-                )
+                ):
+                    yield e
 
             # The `reasoning` field is only present in gpt-oss via Ollama and OpenRouter.
             # - https://cookbook.openai.com/articles/gpt-oss/handle-raw-cot#chat-completions-api
             # - https://openrouter.ai/docs/use-cases/reasoning-tokens#basic-usage-with-reasoning-tokens
             if reasoning := getattr(choice.delta, 'reasoning', None):  # pragma: no cover
-                yield self._parts_manager.handle_thinking_delta(
+                for e in self._parts_manager.handle_thinking_delta(
                     vendor_part_id='reasoning',
                     id='reasoning',
                     content=reasoning,
                     provider_name=self.provider_name,
-                )
+                ):
+                    yield e
 
             # Handle the text part of the response
             content = choice.delta.content
@@ -1887,12 +1889,13 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                 if isinstance(chunk.item, responses.ResponseReasoningItem):
                     if signature := chunk.item.encrypted_content:  # pragma: no branch
                         # Add the signature to the part corresponding to the first summary item
-                        yield self._parts_manager.handle_thinking_delta(
+                        for e in self._parts_manager.handle_thinking_delta(
                             vendor_part_id=f'{chunk.item.id}-0',
                             id=chunk.item.id,
                             signature=signature,
                             provider_name=self.provider_name,
-                        )
+                        ):
+                            yield e
                 elif isinstance(chunk.item, responses.ResponseCodeInterpreterToolCall):
                     _, return_part, file_parts = _map_code_interpreter_tool_call(chunk.item, self.provider_name)
                     for i, file_part in enumerate(file_parts):
@@ -1925,11 +1928,12 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                     yield self._parts_manager.handle_part(vendor_part_id=f'{chunk.item.id}-return', part=return_part)
 
             elif isinstance(chunk, responses.ResponseReasoningSummaryPartAddedEvent):
-                yield self._parts_manager.handle_thinking_delta(
+                for e in self._parts_manager.handle_thinking_delta(
                     vendor_part_id=f'{chunk.item_id}-{chunk.summary_index}',
                     content=chunk.part.text,
                     id=chunk.item_id,
-                )
+                ):
+                    yield e
 
             elif isinstance(chunk, responses.ResponseReasoningSummaryPartDoneEvent):
                 pass  # there's nothing we need to do here
@@ -1938,11 +1942,12 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                 pass  # there's nothing we need to do here
 
             elif isinstance(chunk, responses.ResponseReasoningSummaryTextDeltaEvent):
-                yield self._parts_manager.handle_thinking_delta(
+                for e in self._parts_manager.handle_thinking_delta(
                     vendor_part_id=f'{chunk.item_id}-{chunk.summary_index}',
                     content=chunk.delta,
                     id=chunk.item_id,
-                )
+                ):
+                    yield e
 
             elif isinstance(chunk, responses.ResponseOutputTextAnnotationAddedEvent):
                 # TODO(Marcelo): We should support annotations in the future.

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -6,7 +6,7 @@
 from __future__ import annotations
 
 import io
-from collections.abc import AsyncIterable, AsyncIterator, Sequence
+from collections.abc import AsyncIterable, AsyncIterator, Iterator, Sequence
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from datetime import datetime, timezone
@@ -537,15 +537,18 @@ class OutlinesStreamedResponse(StreamedResponse):
     _provider_name: str
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
-        async for event in self._response:
-            event = self._parts_manager.handle_text_delta(
-                vendor_part_id='content',
-                content=event,
-                thinking_tags=self._model_profile.thinking_tags,
-                ignore_leading_whitespace=self._model_profile.ignore_streamed_leading_whitespace,
+        async for chunk in self._response:
+            events = cast(
+                Iterator[ModelResponseStreamEvent],
+                self._parts_manager.handle_text_delta(
+                    vendor_part_id='content',
+                    content=chunk,
+                    thinking_tags=self._model_profile.thinking_tags,
+                    ignore_leading_whitespace=self._model_profile.ignore_streamed_leading_whitespace,
+                ),
             )
-            if event is not None:  # pragma: no branch
-                yield event
+            for e in events:
+                yield e
 
     @property
     def model_name(self) -> str:

pydantic_ai_slim/pydantic_ai/models/outlines.py

@@ -2061,8 +2061,7 @@ async def test_groq_model_thinking_part_iter(allow_model_requests: None, groq_ap
 
     assert event_parts == snapshot(
         [
-            PartStartEvent(index=0, part=ThinkingPart(content='')),
-            PartDeltaEvent(index=0, delta=ThinkingPartDelta(content_delta='\n')),
+            PartStartEvent(index=0, part=ThinkingPart(content='\n')),
             PartDeltaEvent(index=0, delta=ThinkingPartDelta(content_delta='Okay')),
             PartDeltaEvent(index=0, delta=ThinkingPartDelta(content_delta=',')),
             PartDeltaEvent(index=0, delta=ThinkingPartDelta(content_delta=' so')),