- include incomplete closing tags in thinking part

- fix mistral's event iterator (wasn't iterating over thinking events)

Author: dsfaccini

pydantic_ai_slim/pydantic_ai/_parts_manager.py

@@ -74,11 +74,13 @@ def get_parts(self) -> list[ModelResponsePart]:
         return [p for p in self._parts if not isinstance(p, ToolCallPartDelta)]
 
     def finalize(self) -> Generator[ModelResponseStreamEvent, None, None]:
-        """Flush any buffered content as text parts.
+        """Flush any buffered content, appending to ThinkingParts or creating TextParts.
 
         This should be called when streaming is complete to ensure no content is lost.
-        Any content buffered in _thinking_tag_buffer that hasn't been processed will be
-        treated as regular text and emitted.
+        Any content buffered in _thinking_tag_buffer will be appended to its corresponding
+        ThinkingPart if one exists, otherwise it will be emitted as a TextPart.
+
+        The only possible buffered content to append to ThinkingParts are incomplete closing tags like `</th`
 
         Yields:
             ModelResponseStreamEvent for any buffered content that gets flushed.
@@ -102,19 +104,27 @@ def finalize(self) -> Generator[ModelResponseStreamEvent, None, None]:
                     yield PartStartEvent(index=part_index, part=text_part)
                     self._started_part_indices.add(part_index)
 
-        # flush any remaining buffered content (partial tags like '<thi')
+        # flush any remaining buffered content
         for vendor_part_id, buffered_content in list(self._thinking_tag_buffer.items()):
             if buffered_content:
-                # Remove the vendor_part_id mapping to avoid trying to update existing parts
-                # This ensures buffered partial tags create new TextParts
-                self._vendor_id_to_part_index.pop(vendor_part_id, None)
-                yield from self._handle_text_delta_simple(
-                    vendor_part_id=vendor_part_id,
-                    content=buffered_content,
-                    id=None,
-                    thinking_tags=None,
-                    ignore_leading_whitespace=False,
-                )
+                part_index = self._vendor_id_to_part_index.get(vendor_part_id)
+
+                # If buffered content belongs to a ThinkingPart, append it to the ThinkingPart
+                # (for orphaned buffers like '</th')
+                if part_index is not None and isinstance(self._parts[part_index], ThinkingPart):
+                    yield from self.handle_thinking_delta(vendor_part_id=vendor_part_id, content=buffered_content)
+                    self._vendor_id_to_part_index.pop(vendor_part_id)
+                else:
+                    # Otherwise flush as TextPart
+                    # (for orphaned buffers like '<thi')
+                    self._vendor_id_to_part_index.pop(vendor_part_id, None)
+                    yield from self._handle_text_delta_simple(
+                        vendor_part_id=vendor_part_id,
+                        content=buffered_content,
+                        id=None,
+                        thinking_tags=None,
+                        ignore_leading_whitespace=False,
+                    )
 
         self._thinking_tag_buffer.clear()
 

pydantic_ai_slim/pydantic_ai/models/mistral.py

@@ -637,7 +637,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}

tests/test_parts_manager_split_tags.py

@@ -276,11 +276,12 @@ def test_vendor_id_switch_during_thinking():
     assert parts[1] == snapshot(TextPart(content='different part', part_kind='text'))
 
 
-# this last one's a weird one because the closing tag gets buffered and then flushed (bc it doesn't close)
-# in accordance with the open question https://github.com/pydantic/pydantic-ai/pull/3206#discussion_r2483976551
-# if we auto-close <think> tags then this case will reach the user as `ThinkingPart(content='thinking foo</th')`
 def test_thinking_interrupted_by_incomplete_end_tag_and_vendor_switch():
-    """Test unclosed thinking tag followed by different vendor_part_id."""
+    """Test unclosed thinking tag followed by different vendor_part_id.
+
+    When a vendor_part_id switches and leaves a ThinkingPart with buffered partial end tag,
+    the buffered content is auto-closed by appending it to the ThinkingPart during finalize().
+    """
     manager = ModelResponsePartsManager()
     thinking_tags = ('<think>', '</think>')
 
@@ -299,11 +300,8 @@ def test_thinking_interrupted_by_incomplete_end_tag_and_vendor_switch():
         pass
 
     parts = manager.get_parts()
-    assert len(parts) == 3
+    assert len(parts) == 2
     assert isinstance(parts[0], ThinkingPart)
-    assert parts[0].content == 'thinking foo'
+    assert parts[0].content == 'thinking foo</th'
     assert isinstance(parts[1], TextPart)
     assert parts[1].content == 'new content'
-    # currently as it stands, the incomplete end tag gets flushed as text (which is even weirder from a UX perspective)
-    assert isinstance(parts[2], TextPart)
-    assert parts[2].content == '</th'