refactor parts manager and add parametrized cases

Author: dsfaccini

pydantic_ai_slim/pydantic_ai/_parts_manager.py

@@ -967,6 +967,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)
@@ -1006,6 +1011,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/messages.py

@@ -569,7 +569,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 - loop never started - flush_buffer() seems to be being called before
+                        yield event
+
+                async for event in chain_async_and_sync_iters(iterator, self._parts_manager.flush_buffer()):
                     if isinstance(event, PartStartEvent):
                         if last_start_event:
                             end_event = part_end_event(event.part)
@@ -581,16 +591,6 @@ 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():  # pragma: no cover
-                    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
@@ -616,7 +616,7 @@ def get(self) -> ModelResponse:
         # 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.finalize():
+        for _ in cloned_manager.flush_buffer():
             pass
 
         return ModelResponse(

pydantic_ai_slim/pydantic_ai/models/__init__.py

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