Fix streaming thinking tags split across multiple chunks

pydantic_ai_slim/pydantic_ai/messages.py

@@ -970,6 +970,11 @@ class TextPart:
     part_kind: Literal['text'] = 'text'
     """Part type identifier, this is available on all parts as a discriminator."""
 
+    potential_opening_tag_buffer: Annotated[str, pydantic.Field(exclude=True)] = field(
+        compare=False, default='', repr=False
+    )
+    """A buffer to accumulate a potential opening tag (like '<thi')."""
+
     def has_content(self) -> bool:
         """Return `True` if the text content is non-empty."""
         return bool(self.content)
@@ -1009,6 +1014,9 @@ class ThinkingPart:
     part_kind: Literal['thinking'] = 'thinking'
     """Part type identifier, this is available on all parts as a discriminator."""
 
+    closing_tag_buffer: Annotated[str, pydantic.Field(exclude=True)] = field(compare=False, default='', repr=False)
+    """A buffer to accumulate a potential closing tag (like '</th')."""
+
     def has_content(self) -> bool:
         """Return `True` if the thinking content is non-empty."""
         return bool(self.content)

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -7,6 +7,7 @@
 from __future__ import annotations as _annotations
 
 import base64
+import copy
 import warnings
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator, Callable, Iterator
@@ -568,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]:
+    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
@@ -615,7 +616,17 @@ def part_end_event(next_part: ModelResponsePart | None = None) -> PartEndEvent |
                         next_part_kind=next_part.part_kind if next_part else None,
                     )
 
-                async for event in iterator:
+                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:  # pragma: no cover - idk why this isn't covered
+                        yield event
+
+                async for event in chain_async_and_sync_iters(  # pragma: no cover - related to above
+                    iterator, self._parts_manager.final_flush()
+                ):
                     if isinstance(event, PartStartEvent):
                         if last_start_event:
                             end_event = part_end_event(event.part)
@@ -649,8 +660,14 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
 
     def get(self) -> ModelResponse:
         """Build a [`ModelResponse`][pydantic_ai.messages.ModelResponse] from the data received from the stream so far."""
+        # Flush any buffered content before building response
+        # clone parts manager to avoid modifying the ongoing stream state
+        cloned_manager = copy.deepcopy(self._parts_manager)
+        for _ in cloned_manager.final_flush():
+            pass
+
         return ModelResponse(
-            parts=self._parts_manager.get_parts(),
+            parts=cloned_manager.get_parts(),
             model_name=self.model_name,
             timestamp=self.timestamp,
             usage=self.usage(),

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -748,25 +748,26 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
             elif isinstance(event, BetaRawContentBlockStartEvent):
                 current_block = event.content_block
                 if isinstance(current_block, BetaTextBlock) and current_block.text:
-                    maybe_event = self._parts_manager.handle_text_delta(
+                    for event_item in self._parts_manager.handle_text_delta(
                         vendor_part_id=event.index, content=current_block.text
-                    )
-                    if maybe_event is not None:  # pragma: no branch
-                        yield maybe_event
+                    ):
+                        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,
@@ -822,23 +823,24 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
 
             elif isinstance(event, BetaRawContentBlockDeltaEvent):
                 if isinstance(event.delta, BetaTextDelta):
-                    maybe_event = self._parts_manager.handle_text_delta(
+                    for event_item in self._parts_manager.handle_text_delta(
                         vendor_part_id=event.index, content=event.delta.text
-                    )
-                    if maybe_event is not None:  # pragma: no branch
-                        yield maybe_event
+                    ):
+                        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/bedrock.py

@@ -687,24 +687,25 @@ 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'):
-                        maybe_event = self._parts_manager.handle_text_delta(vendor_part_id=index, content=text)
-                        if maybe_event is not None:  # pragma: no branch
-                            yield maybe_event
+                        for event in self._parts_manager.handle_text_delta(vendor_part_id=index, content=text):
+                            yield event
                     if 'toolUse' in delta:
                         tool_use = delta['toolUse']
                         maybe_event = self._parts_manager.handle_tool_call_delta(

pydantic_ai_slim/pydantic_ai/models/function.py

@@ -292,26 +292,26 @@ 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)
                 self._usage += usage.RequestUsage(output_tokens=response_tokens)
-                maybe_event = self._parts_manager.handle_text_delta(vendor_part_id='content', content=item)
-                if maybe_event is not None:  # pragma: no branch
-                    yield maybe_event
+                for event in self._parts_manager.handle_text_delta(vendor_part_id='content', content=item):
+                    yield event
             elif isinstance(item, dict) and item:
                 for dtc_index, delta in item.items():
                     if isinstance(delta, DeltaThinkingPart):
                         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/gemini.py

@@ -461,11 +461,10 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                 if 'text' in gemini_part:
                     # Using vendor_part_id=None means we can produce multiple text parts if their deltas are sprinkled
                     # amongst the tool call deltas
-                    maybe_event = self._parts_manager.handle_text_delta(
+                    for event in self._parts_manager.handle_text_delta(
                         vendor_part_id=None, content=gemini_part['text']
-                    )
-                    if maybe_event is not None:  # pragma: no branch
-                        yield maybe_event
+                    ):
+                        yield event
 
                 elif 'function_call' in gemini_part:
                     # Here, we assume all function_call parts are complete and don't have deltas.

pydantic_ai_slim/pydantic_ai/models/google.py

@@ -671,24 +671,25 @@ 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 len(part.text) > 0:
                         if part.thought:
-                            yield self._parts_manager.handle_thinking_delta(
+                            for event in self._parts_manager.handle_thinking_delta(
                                 vendor_part_id='thinking', content=part.text
-                            )
+                            ):
+                                yield event
                         else:
-                            maybe_event = self._parts_manager.handle_text_delta(
+                            for event in self._parts_manager.handle_text_delta(
                                 vendor_part_id='content', content=part.text
-                            )
-                            if maybe_event is not None:  # pragma: no branch
-                                yield maybe_event
+                            ):
+                                yield event
                 elif part.function_call:
                     maybe_event = self._parts_manager.handle_tool_call_delta(
                         vendor_part_id=uuid4(),

pydantic_ai_slim/pydantic_ai/models/groq.py

@@ -549,9 +549,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
 
@@ -574,14 +575,13 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                 # Handle the text part of the response
                 content = choice.delta.content
                 if content:
-                    maybe_event = self._parts_manager.handle_text_delta(
+                    for event in self._parts_manager.handle_text_delta(
                         vendor_part_id='content',
                         content=content,
                         thinking_tags=self._model_profile.thinking_tags,
                         ignore_leading_whitespace=self._model_profile.ignore_streamed_leading_whitespace,
-                    )
-                    if maybe_event is not None:  # pragma: no branch
-                        yield maybe_event
+                    ):
+                        yield event
 
                 # Handle the tool calls
                 for dtc in choice.delta.tool_calls or []:

pydantic_ai_slim/pydantic_ai/models/huggingface.py

@@ -483,14 +483,13 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
             # Handle the text part of the response
             content = choice.delta.content
             if content:
-                maybe_event = self._parts_manager.handle_text_delta(
+                for event in self._parts_manager.handle_text_delta(
                     vendor_part_id='content',
                     content=content,
                     thinking_tags=self._model_profile.thinking_tags,
                     ignore_leading_whitespace=self._model_profile.ignore_streamed_leading_whitespace,
-                )
-                if maybe_event is not None:  # pragma: no branch
-                    yield maybe_event
+                ):
+                    yield event
 
             for dtc in choice.delta.tool_calls or []:
                 maybe_event = self._parts_manager.handle_tool_call_delta(

pydantic_ai_slim/pydantic_ai/models/mistral.py

@@ -639,7 +639,8 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
             content = choice.delta.content
             text, thinking = _map_content(content)
             for thought in thinking:
-                self._parts_manager.handle_thinking_delta(vendor_part_id='thinking', content=thought)
+                for event in self._parts_manager.handle_thinking_delta(vendor_part_id='thinking', content=thought):
+                    yield event
             if text:
                 # Attempt to produce an output tool call from the received text
                 output_tools = {c.name: c for c in self.model_request_parameters.output_tools}
@@ -655,9 +656,8 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                             tool_call_id=maybe_tool_call_part.tool_call_id,
                         )
                 else:
-                    maybe_event = self._parts_manager.handle_text_delta(vendor_part_id='content', content=text)
-                    if maybe_event is not None:  # pragma: no branch
-                        yield maybe_event
+                    for event in self._parts_manager.handle_text_delta(vendor_part_id='content', content=text):
+                        yield event
 
             # Handle the explicit tool calls
             for index, dtc in enumerate(choice.delta.tool_calls or []):