Refactor handle_text_delta() to generator pattern with split tag buffering

Convert handle_text_delta() from returning a single event to yielding multiple
events via a generator pattern. This enables proper handling of thinking tags
that may be split across multiple streaming chunks.

Key changes:
- Convert handle_text_delta() return type from ModelResponseStreamEvent | None
  to Generator[ModelResponseStreamEvent, None, None]
- Add _tag_buffer field to track partial content across chunks
- Implement _handle_text_delta_simple() for non-thinking-tag cases
- Implement _handle_text_delta_with_thinking_tags() with buffering logic
- Add _could_be_tag_start() helper to detect potential split tags
- Update all model implementations (10 files) to iterate over events
- Adapt test_handle_text_deltas_with_think_tags for generator API

Behavior:
- Complete thinking tags work at any position (maintains original behavior)
- Split thinking tags are buffered when starting at position 0 of chunk
- Split tags only work when vendor_part_id is not None (buffering requirement)

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

Author: dsfaccini

pydantic_ai_slim/pydantic_ai/_parts_manager.py

@@ -13,7 +13,7 @@
 
 from __future__ import annotations as _annotations
 
-from collections.abc import Hashable
+from collections.abc import Generator, Hashable
 from dataclasses import dataclass, field, replace
 from typing import Any
 
@@ -58,6 +58,8 @@ class ModelResponsePartsManager:
     """A list of parts (text or tool calls) that make up the current state of the model's response."""
     _vendor_id_to_part_index: dict[VendorId, int] = field(default_factory=dict, init=False)
     """Maps a vendor's "part" ID (if provided) to the index in `_parts` where that part resides."""
+    _tag_buffer: dict[VendorId, str] = field(default_factory=dict, init=False)
+    """Buffers partial content when thinking tags might be split across chunks."""
 
     def get_parts(self) -> list[ModelResponsePart]:
         """Return only model response parts that are complete (i.e., not ToolCallPartDelta's).
@@ -75,82 +77,159 @@ def handle_text_delta(
         id: str | None = None,
         thinking_tags: tuple[str, str] | None = None,
         ignore_leading_whitespace: bool = False,
-    ) -> ModelResponseStreamEvent | None:
+    ) -> Generator[ModelResponseStreamEvent, None, None]:
         """Handle incoming text content, creating or updating a TextPart in the manager as appropriate.
 
         When `vendor_part_id` is None, the latest part is updated if it exists and is a TextPart;
         otherwise, a new TextPart is created. When a non-None ID is specified, the TextPart corresponding
         to that vendor ID is either created or updated.
 
+        Thinking tags may be split across multiple chunks. When `thinking_tags` is provided and
+        `vendor_part_id` is not None, this method buffers content that could be the start of a
+        thinking tag appearing at the beginning of the current chunk.
+
         Args:
             vendor_part_id: The ID the vendor uses to identify this piece
                 of text. If None, a new part will be created unless the latest part is already
                 a TextPart.
             content: The text content to append to the appropriate TextPart.
             id: An optional id for the text part.
             thinking_tags: If provided, will handle content between the thinking tags as thinking parts.
+                Buffering for split tags requires a non-None vendor_part_id.
             ignore_leading_whitespace: If True, will ignore leading whitespace in the content.
 
-        Returns:
-            - A `PartStartEvent` if a new part was created.
-            - A `PartDeltaEvent` if an existing part was updated.
-            - `None` if no new event is emitted (e.g., the first text part was all whitespace).
+        Yields:
+            - `PartStartEvent` if a new part was created.
+            - `PartDeltaEvent` if an existing part was updated.
+            May yield multiple events from a single call if buffered content is flushed.
 
         Raises:
             UnexpectedModelBehavior: If attempting to apply text content to a part that is not a TextPart.
         """
+        if thinking_tags and vendor_part_id is not None:
+            yield from self._handle_text_delta_with_thinking_tags(
+                vendor_part_id=vendor_part_id,
+                content=content,
+                id=id,
+                thinking_tags=thinking_tags,
+                ignore_leading_whitespace=ignore_leading_whitespace,
+            )
+        else:
+            yield from self._handle_text_delta_simple(
+                vendor_part_id=vendor_part_id,
+                content=content,
+                id=id,
+                thinking_tags=thinking_tags,
+                ignore_leading_whitespace=ignore_leading_whitespace,
+            )
+
+    def _handle_text_delta_simple(
+        self,
+        *,
+        vendor_part_id: VendorId | None,
+        content: str,
+        id: str | None,
+        thinking_tags: tuple[str, str] | None,
+        ignore_leading_whitespace: bool,
+    ) -> Generator[ModelResponseStreamEvent, None, None]:
+        """Handle text delta without split tag buffering (original logic)."""
         existing_text_part_and_index: tuple[TextPart, int] | None = None
 
         if vendor_part_id is None:
-            # If the vendor_part_id is None, check if the latest part is a TextPart to update
             if self._parts:
                 part_index = len(self._parts) - 1
                 latest_part = self._parts[part_index]
                 if isinstance(latest_part, TextPart):
                     existing_text_part_and_index = latest_part, part_index
         else:
-            # Otherwise, attempt to look up an existing TextPart by vendor_part_id
             part_index = self._vendor_id_to_part_index.get(vendor_part_id)
             if part_index is not None:
                 existing_part = self._parts[part_index]
 
                 if thinking_tags and isinstance(existing_part, ThinkingPart):
-                    # We may be building a thinking part instead of a text part if we had previously seen a thinking tag
                     if content == thinking_tags[1]:
-                        # When we see the thinking end tag, we're done with the thinking part and the next text delta will need a new part
                         self._vendor_id_to_part_index.pop(vendor_part_id)
-                        return None
+                        return
                     else:
-                        return self.handle_thinking_delta(vendor_part_id=vendor_part_id, content=content)
+                        yield self.handle_thinking_delta(vendor_part_id=vendor_part_id, content=content)
+                        return
                 elif isinstance(existing_part, TextPart):
                     existing_text_part_and_index = existing_part, part_index
                 else:
                     raise UnexpectedModelBehavior(f'Cannot apply a text delta to {existing_part=}')
 
         if thinking_tags and content == thinking_tags[0]:
-            # When we see a thinking start tag (which is a single token), we'll build a new thinking part instead
             self._vendor_id_to_part_index.pop(vendor_part_id, None)
-            return self.handle_thinking_delta(vendor_part_id=vendor_part_id, content='')
+            yield self.handle_thinking_delta(vendor_part_id=vendor_part_id, content='')
+            return
 
         if existing_text_part_and_index is None:
-            # This is a workaround for models that emit `<think>\n</think>\n\n` or an empty text part ahead of tool calls (e.g. Ollama + Qwen3),
-            # which we don't want to end up treating as a final result when using `run_stream` with `str` a valid `output_type`.
             if ignore_leading_whitespace and (len(content) == 0 or content.isspace()):
-                return None
+                return
 
-            # There is no existing text part that should be updated, so create a new one
             new_part_index = len(self._parts)
             part = TextPart(content=content, id=id)
             if vendor_part_id is not None:
                 self._vendor_id_to_part_index[vendor_part_id] = new_part_index
             self._parts.append(part)
-            return PartStartEvent(index=new_part_index, part=part)
+            yield PartStartEvent(index=new_part_index, part=part)
         else:
-            # Update the existing TextPart with the new content delta
             existing_text_part, part_index = existing_text_part_and_index
             part_delta = TextPartDelta(content_delta=content)
             self._parts[part_index] = part_delta.apply(existing_text_part)
-            return PartDeltaEvent(index=part_index, delta=part_delta)
+            yield PartDeltaEvent(index=part_index, delta=part_delta)
+
+    def _handle_text_delta_with_thinking_tags(
+        self,
+        *,
+        vendor_part_id: VendorId,
+        content: str,
+        id: str | None,
+        thinking_tags: tuple[str, str],
+        ignore_leading_whitespace: bool,
+    ) -> Generator[ModelResponseStreamEvent, None, None]:
+        """Handle text delta with thinking tag detection and buffering for split tags."""
+        start_tag, end_tag = thinking_tags
+        buffered = self._tag_buffer.get(vendor_part_id, '')
+        combined_content = buffered + content
+
+        part_index = self._vendor_id_to_part_index.get(vendor_part_id)
+        existing_part = self._parts[part_index] if part_index is not None else None
+
+        if existing_part is not None and isinstance(existing_part, ThinkingPart):
+            if combined_content == end_tag:
+                self._vendor_id_to_part_index.pop(vendor_part_id)
+                self._tag_buffer.pop(vendor_part_id, None)
+                return
+            else:
+                self._tag_buffer.pop(vendor_part_id, None)
+                yield self.handle_thinking_delta(vendor_part_id=vendor_part_id, content=combined_content)
+                return
+
+        if combined_content == start_tag:
+            self._tag_buffer.pop(vendor_part_id, None)
+            self._vendor_id_to_part_index.pop(vendor_part_id, None)
+            yield self.handle_thinking_delta(vendor_part_id=vendor_part_id, content='')
+            return
+
+        if content.startswith(start_tag[0]) and self._could_be_tag_start(combined_content, start_tag):
+            self._tag_buffer[vendor_part_id] = combined_content
+            return
+
+        self._tag_buffer.pop(vendor_part_id, None)
+        yield from self._handle_text_delta_simple(
+            vendor_part_id=vendor_part_id,
+            content=combined_content,
+            id=id,
+            thinking_tags=thinking_tags,
+            ignore_leading_whitespace=ignore_leading_whitespace,
+        )
+
+    def _could_be_tag_start(self, content: str, tag: str) -> bool:
+        """Check if content could be the start of a tag."""
+        if len(content) >= len(tag):
+            return False
+        return tag.startswith(content)
 
     def handle_thinking_delta(
         self,

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -43,6 +43,7 @@
 )
 from ..output import OutputMode
 from ..profiles import DEFAULT_PROFILE, ModelProfile, ModelProfileSpec
+from ..providers import infer_provider
 from ..settings import ModelSettings, merge_model_settings
 from ..tools import ToolDefinition
 from ..usage import RequestUsage
@@ -637,41 +638,39 @@ def infer_model(model: Model | KnownModelName | str) -> Model:  # noqa: C901
         return TestModel()
 
     try:
-        provider, model_name = model.split(':', maxsplit=1)
+        provider_name, model_name = model.split(':', maxsplit=1)
     except ValueError:
-        provider = None
+        provider_name = None
         model_name = model
         if model_name.startswith(('gpt', 'o1', 'o3')):
-            provider = 'openai'
+            provider_name = 'openai'
         elif model_name.startswith('claude'):
-            provider = 'anthropic'
+            provider_name = 'anthropic'
         elif model_name.startswith('gemini'):
-            provider = 'google-gla'
+            provider_name = 'google-gla'
 
-        if provider is not None:
+        if provider_name is not None:
             warnings.warn(
-                f"Specifying a model name without a provider prefix is deprecated. Instead of {model_name!r}, use '{provider}:{model_name}'.",
+                f"Specifying a model name without a provider prefix is deprecated. Instead of {model_name!r}, use '{provider_name}:{model_name}'.",
                 DeprecationWarning,
             )
         else:
             raise UserError(f'Unknown model: {model}')
 
-    if provider == 'vertexai':  # pragma: no cover
+    if provider_name == 'vertexai':  # pragma: no cover
         warnings.warn(
             "The 'vertexai' provider name is deprecated. Use 'google-vertex' instead.",
             DeprecationWarning,
         )
-        provider = 'google-vertex'
+        provider_name = 'google-vertex'
 
-    if provider == 'gateway':
-        from ..providers.gateway import infer_model as infer_model_from_gateway
+    provider = infer_provider(provider_name)
 
-        return infer_model_from_gateway(model_name)
-    elif provider == 'cohere':
-        from .cohere import CohereModel
-
-        return CohereModel(model_name, provider=provider)
-    elif provider in (
+    model_kind = provider_name
+    if model_kind.startswith('gateway/'):
+        model_kind = provider_name.removeprefix('gateway/')
+    if model_kind in (
+        'openai',
         'azure',
         'deepseek',
         'cerebras',
@@ -681,43 +680,50 @@ def infer_model(model: Model | KnownModelName | str) -> Model:  # noqa: C901
         'heroku',
         'moonshotai',
         'ollama',
-        'openai',
-        'openai-chat',
         'openrouter',
         'together',
         'vercel',
         'litellm',
         'nebius',
         'ovhcloud',
     ):
+        model_kind = 'openai-chat'
+    elif model_kind in ('google-gla', 'google-vertex'):
+        model_kind = 'google'
+
+    if model_kind == 'openai-chat':
         from .openai import OpenAIChatModel
 
         return OpenAIChatModel(model_name, provider=provider)
-    elif provider == 'openai-responses':
+    elif model_kind == 'openai-responses':
         from .openai import OpenAIResponsesModel
 
-        return OpenAIResponsesModel(model_name, provider='openai')
-    elif provider in ('google-gla', 'google-vertex'):
+        return OpenAIResponsesModel(model_name, provider=provider)
+    elif model_kind == 'google':
         from .google import GoogleModel
 
         return GoogleModel(model_name, provider=provider)
-    elif provider == 'groq':
+    elif model_kind == 'groq':
         from .groq import GroqModel
 
         return GroqModel(model_name, provider=provider)
-    elif provider == 'mistral':
+    elif model_kind == 'cohere':
+        from .cohere import CohereModel
+
+        return CohereModel(model_name, provider=provider)
+    elif model_kind == 'mistral':
         from .mistral import MistralModel
 
         return MistralModel(model_name, provider=provider)
-    elif provider == 'anthropic':
+    elif model_kind == 'anthropic':
         from .anthropic import AnthropicModel
 
         return AnthropicModel(model_name, provider=provider)
-    elif provider == 'bedrock':
+    elif model_kind == 'bedrock':
         from .bedrock import BedrockConverseModel
 
         return BedrockConverseModel(model_name, provider=provider)
-    elif provider == 'huggingface':
+    elif model_kind == 'huggingface':
         from .huggingface import HuggingFaceModel
 
         return HuggingFaceModel(model_name, provider=provider)

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -162,7 +162,7 @@ def __init__(
         self,
         model_name: AnthropicModelName,
         *,
-        provider: Literal['anthropic'] | Provider[AsyncAnthropicClient] = 'anthropic',
+        provider: Literal['anthropic', 'gateway'] | Provider[AsyncAnthropicClient] = 'anthropic',
         profile: ModelProfileSpec | None = None,
         settings: ModelSettings | None = None,
     ):
@@ -179,7 +179,7 @@ def __init__(
         self._model_name = model_name
 
         if isinstance(provider, str):
-            provider = infer_provider(provider)
+            provider = infer_provider('gateway/anthropic' if provider == 'gateway' else provider)
         self._provider = provider
         self.client = provider.client
 
@@ -669,11 +669,10 @@ 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(
                         vendor_part_id=event.index,
@@ -715,11 +714,10 @@ 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(
                         vendor_part_id=event.index,