Handle thinking tags split across multiple chunks in streaming responses

dsfaccini · dsfaccini · commit 5f5d0993fb55 · 2025-10-22T15:24:22.000-05:00
diff --git a/pydantic_ai_slim/pydantic_ai/_parts_manager.py b/pydantic_ai_slim/pydantic_ai/_parts_manager.py
diff --git a/pydantic_ai_slim/pydantic_ai/models/anthropic.py b/pydantic_ai_slim/pydantic_ai/models/anthropic.py
@@ -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 text_event 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 text_event
                 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 text_event 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 text_event
                 elif isinstance(event.delta, BetaThinkingDelta):
                     yield self._parts_manager.handle_thinking_delta(
                         vendor_part_id=event.index,
diff --git a/pydantic_ai_slim/pydantic_ai/models/bedrock.py b/pydantic_ai_slim/pydantic_ai/models/bedrock.py
@@ -702,9 +702,8 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                                 provider_name=self.provider_name if signature else None,
                             )
                     if 'text' in delta:
-                        maybe_event = self._parts_manager.handle_text_delta(vendor_part_id=index, content=delta['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=delta['text']):
+                            yield event
                     if 'toolUse' in delta:
                         tool_use = delta['toolUse']
                         maybe_event = self._parts_manager.handle_tool_call_delta(
diff --git a/pydantic_ai_slim/pydantic_ai/models/function.py b/pydantic_ai_slim/pydantic_ai/models/function.py
@@ -289,9 +289,8 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
             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):
diff --git a/pydantic_ai_slim/pydantic_ai/models/gemini.py b/pydantic_ai_slim/pydantic_ai/models/gemini.py
@@ -454,11 +454,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.
diff --git a/pydantic_ai_slim/pydantic_ai/models/google.py b/pydantic_ai_slim/pydantic_ai/models/google.py
@@ -668,9 +668,8 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                     if part.thought:
                         yield self._parts_manager.handle_thinking_delta(vendor_part_id='thinking', content=part.text)
                     else:
-                        maybe_event = 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
+                        for event in self._parts_manager.handle_text_delta(vendor_part_id='content', content=part.text):
+                            yield event
                 elif part.function_call:
                     maybe_event = self._parts_manager.handle_tool_call_delta(
                         vendor_part_id=uuid4(),
diff --git a/pydantic_ai_slim/pydantic_ai/models/groq.py b/pydantic_ai_slim/pydantic_ai/models/groq.py
@@ -564,14 +564,13 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                 # Handle the text part of the response
                 content = choice.delta.content
                 if content is not None:
-                    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 []:
diff --git a/pydantic_ai_slim/pydantic_ai/models/huggingface.py b/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 is not None:
-                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(
diff --git a/pydantic_ai_slim/pydantic_ai/models/mistral.py b/pydantic_ai_slim/pydantic_ai/models/mistral.py
@@ -653,9 +653,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 []):
diff --git a/pydantic_ai_slim/pydantic_ai/models/openai.py b/pydantic_ai_slim/pydantic_ai/models/openai.py
@@ -1619,17 +1619,16 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
             # Handle the text part of the response
             content = choice.delta.content
             if content is not None:
-                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
-                    if isinstance(maybe_event, PartStartEvent) and isinstance(maybe_event.part, ThinkingPart):
-                        maybe_event.part.id = 'content'
-                        maybe_event.part.provider_name = self.provider_name
-                    yield maybe_event
+                ):
+                    if isinstance(event, PartStartEvent) and isinstance(event.part, ThinkingPart):
+                        event.part.id = 'content'
+                        event.part.provider_name = self.provider_name
+                    yield event
 
             # The `reasoning_content` field is only present in DeepSeek models.
             # https://api-docs.deepseek.com/guides/reasoning_model
@@ -1835,11 +1834,10 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                 pass  # there's nothing we need to do here
 
             elif isinstance(chunk, responses.ResponseTextDeltaEvent):
-                maybe_event = self._parts_manager.handle_text_delta(
+                for event in self._parts_manager.handle_text_delta(
                     vendor_part_id=chunk.item_id, content=chunk.delta, id=chunk.item_id
-                )
-                if maybe_event is not None:  # pragma: no branch
-                    yield maybe_event
+                ):
+                    yield event
 
             elif isinstance(chunk, responses.ResponseTextDoneEvent):
                 pass  # there's nothing we need to do here
diff --git a/pydantic_ai_slim/pydantic_ai/models/test.py b/pydantic_ai_slim/pydantic_ai/models/test.py
@@ -310,14 +310,12 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                     mid = len(text) // 2
                     words = [text[:mid], text[mid:]]
                 self._usage += _get_string_usage('')
-                maybe_event = self._parts_manager.handle_text_delta(vendor_part_id=i, content='')
-                if maybe_event is not None:  # pragma: no branch
-                    yield maybe_event
+                for event in self._parts_manager.handle_text_delta(vendor_part_id=i, content=''):
+                    yield event
                 for word in words:
                     self._usage += _get_string_usage(word)
-                    maybe_event = self._parts_manager.handle_text_delta(vendor_part_id=i, content=word)
-                    if maybe_event is not None:  # pragma: no branch
-                        yield maybe_event
+                    for event in self._parts_manager.handle_text_delta(vendor_part_id=i, content=word):
+                        yield event
             elif isinstance(part, ToolCallPart):
                 yield self._parts_manager.handle_tool_call_part(
                     vendor_part_id=i, tool_name=part.tool_name, args=part.args, tool_call_id=part.tool_call_id
diff --git a/tests/models/test_instrumented.py b/tests/models/test_instrumented.py
@@ -107,12 +107,10 @@ async def request_stream(
 class MyResponseStream(StreamedResponse):
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
         self._usage = RequestUsage(input_tokens=300, output_tokens=400)
-        maybe_event = self._parts_manager.handle_text_delta(vendor_part_id=0, content='text1')
-        if maybe_event is not None:  # pragma: no branch
-            yield maybe_event
-        maybe_event = self._parts_manager.handle_text_delta(vendor_part_id=0, content='text2')
-        if maybe_event is not None:  # pragma: no branch
-            yield maybe_event
+        for event in self._parts_manager.handle_text_delta(vendor_part_id=0, content='text1'):
+            yield event
+        for event in self._parts_manager.handle_text_delta(vendor_part_id=0, content='text2'):
+            yield event
 
     @property
     def model_name(self) -> str:
diff --git a/tests/models/test_openai.py b/tests/models/test_openai.py
@@ -1,7 +1,7 @@
 from __future__ import annotations as _annotations
 
 import json
-from collections.abc import Callable
+from collections.abc import AsyncIterable, Callable
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from enum import Enum
@@ -25,8 +25,10 @@
     ModelResponse,
     ModelRetry,
     RetryPromptPart,
+    RunContext,
     SystemPromptPart,
     TextPart,
+    TextPartDelta,
     ThinkingPart,
     ToolCallPart,
     ToolReturnPart,
@@ -666,6 +668,85 @@ async def test_stream_thinking_tags_split_across_chunks(allow_model_requests: No
     assert parts[1].content.strip() == 'Normal content.'
 
 
+async def test_stream_thinking_tags_with_surrounding_text_event_handler(allow_model_requests: None):
+    """Test that thinking tags with text before them are treated as regular text.
+
+    With the simplified start-only detection logic (per maintainer decision in
+    https://github.com/pydantic/pydantic-ai/pull/3206#issuecomment-3429785920),
+    thinking tags are ONLY recognized at position 0 of the response.
+
+    When content like 'foo<think>content</think>bar' arrives (even split across chunks),
+    it should all be treated as regular text because "foo" comes before the opening tag.
+
+    Real-world scenario from Gemini via LiteLLM:
+    - Chunk 1: "foo<thi"
+    - Chunk 2: "nk>content</think>"
+    - Chunk 3: "bar"
+
+    Expected: All content treated as text (no thinking tags detected)
+
+    Note: Event handlers receive events "up until the final result is found" (documented behavior),
+    so they will only see the first PartStartEvent(TextPart) which triggers FinalResultEvent.
+    Subsequent deltas are not seen by the event handler, but the final message parts will contain
+    all the text.
+    """
+    # Simulate the scenario with text before thinking tags
+    stream = [
+        text_chunk('foo<thi'),  # Text before tag with partial start tag
+        text_chunk('nk>content</think>'),  # What looks like thinking content
+        text_chunk('bar'),  # Text after
+        chunk([]),
+    ]
+    mock_client = MockOpenAI.create_mock_stream(stream)
+    m = OpenAIChatModel('gpt-4o', provider=OpenAIProvider(openai_client=mock_client))
+    agent = Agent(m)
+
+    # Collect all events from the stream
+    from pydantic_ai.messages import AgentStreamEvent
+
+    collected_events: list[AgentStreamEvent] = []
+
+    async def collect_events(ctx: RunContext[None], events: AsyncIterable[AgentStreamEvent]) -> None:
+        """Event handler that collects streaming events up until FinalResultEvent."""
+        async for event in events:
+            collected_events.append(event)
+
+    async with agent.run_stream('', event_stream_handler=collect_events) as result:
+        # Consume the stream
+        async for _ in result.stream_output(debounce_by=None):
+            pass
+
+    # Verify final parts: all treated as text (no thinking tags because of "foo" prefix)
+    msgs = result.new_messages()
+    parts = msgs[-1].parts
+    assert len(parts) == 1
+    assert isinstance(parts[0], TextPart)
+    assert parts[0].content == 'foo<think>content</think>bar'
+
+    # Verify event stream handler received events up to FinalResultEvent
+    # According to documented behavior, event handlers see events "up until the final result is found"
+    from pydantic_ai.messages import FinalResultEvent, PartDeltaEvent, PartStartEvent
+
+    part_events = [e for e in collected_events if isinstance(e, PartStartEvent | PartDeltaEvent)]
+    final_result_events = [e for e in collected_events if isinstance(e, FinalResultEvent)]
+
+    # Should have received at least the PartStartEvent(TextPart) and FinalResultEvent
+    # The PartStartEvent triggers FinalResultEvent (for output_type=str), which stops the event handler
+    assert len(part_events) >= 1, f'Expected at least 1 part event, but got {len(part_events)}: {part_events}'
+    assert len(final_result_events) == 1, f'Expected 1 FinalResultEvent, but got {len(final_result_events)}'
+
+    # All part events should be for text parts (no thinking events)
+    text_events = [
+        e
+        for e in part_events
+        if isinstance(e, PartStartEvent)
+        and isinstance(e.part, TextPart)
+        or isinstance(e, PartDeltaEvent)
+        and isinstance(e.delta, TextPartDelta)
+    ]
+    assert len(text_events) == len(part_events), f'All events should be text events: {part_events}'
+
+
 async def test_no_delta(allow_model_requests: None):
     stream = [
         chunk([]),
diff --git a/tests/test_parts_manager.py b/tests/test_parts_manager.py
