Streaming extract `<think>` tags split up over multiple chunks

[phemmer]

Initial Checks

• I confirm that I'm using the latest version of Pydantic AI
• I confirm that I searched for my issue in https://github.com/pydantic/pydantic-ai/issues before opening this issue

Description

When pydantic-ai receives a streaming response from an OpenAI compatible model, it does not extract the <think></think> tags into a separate ThinkingPart. However it works fine when not streaming.

Example Code

@pytest.mark.asyncio
async def test_thinking_nostream():
    agent = Agent(
        model=OpenAIChatModel('', provider='litellm')
    )

    async def fake_completions_create(*args, **kwargs) -> ChatCompletion:
        return ChatCompletion(
            id="abcd",
            choices=[Choice(
                finish_reason="stop",
                index=0,
                message=ChatCompletionMessage(
                    content="<think>\nthinking content</think>\nNormal content.",
                    role="assistant",
                )
            )],
            created=int(datetime.datetime.now().timestamp()),
            model="abcd",
            object="chat.completion",
        )

    agent.model._completions_create = fake_completions_create
    result = await agent.run("asdf")

    msgs = result.new_messages()
    parts = msgs[-1].parts
    assert isinstance(parts[0], ThinkingPart)
    assert parts[0].content.strip() == "thinking content"
    assert isinstance(parts[1], TextPart)
    assert parts[1].content.strip() == "Normal content."

@pytest.mark.asyncio
async def test_thinking_stream(httpx_mock: HTTPXMock):
    agent = Agent(
        model=OpenAIChatModel('', provider='litellm')
    )

    chunks = [
        '''data: {"id":"asdf","created":1234567890,"model":"asdf","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant","content":"<"}}]}''',
        '''data: {"id":"asdf","created":1234567890,"model":"asdf","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant","content":"think>\\nthinking content"}}]}''',
        '''data: {"id":"asdf","created":1234567890,"model":"asdf","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant","content":"</think>\\nNormal content."}}]}''',
        '''data: {"id":"asdf","created":1234567890,"model":"asdf","object":"chat.completion.chunk","choices":[{"index":0,"finish_reason":"stop","delta":{}}]}''',
        '''data: {"id":"asdf","created":1234567890,"model":"asdf","object":"chat.completion.chunk","choices":[{"index":0,"delta":{}}],"usage":{"completion_tokens":123,"prompt_tokens":1234,"total_tokens":123,"completion_tokens_details":{"reasoning_tokens":0},"prompt_tokens_details":{"cached_tokens":123,"text_tokens":123}}}''',
        '''data: [DONE]'''
    ]
    httpx_mock.add_response(
        url="https://api.openai.com/v1/chat/completions",
        method="POST",
        stream=IteratorStream([(c + "\n\n").encode() for c in chunks]),
        headers={'Content-Type': 'text/event-stream'},
    )


    async def event_stream_handler(ctx, stream):
        pass

    result = await agent.run("asdf", event_stream_handler=event_stream_handler)

    msgs = result.new_messages()
    parts = msgs[-1].parts
    assert isinstance(parts[0], ThinkingPart)
    assert parts[0].content.strip() == "thinking content"
    assert isinstance(parts[1], TextPart)
    assert parts[1].content.strip() == "Normal content."

test_thinking_nostream passes
test_thinking_stream fails with:

>       assert isinstance(parts[0], ThinkingPart)
E       AssertionError: assert False
E        +  where False = isinstance(TextPart(content='<think>\nthinking content</think>\nNormal content.'), ThinkingPart)

^{(Yes, the tests don't mock out the response the same way. But it's immaterial to the <think> extraction, as that happens after. I wrote the test_thinking_nostream first, and when that didn't reproduce the issue, I had to go to streaming, and mocking that wasn't nearly as simple)}

Python, Pydantic AI & LLM client version

Python 3.13.3
pydantic-ai 1.0.10

[DouweM]

The expectation currently is that the think start and end tags will arrive in individual text chunks by themselves with no surrounding text, which I confirmed to be the case with Ollama and others. What model are you seeing output the tags with text immediately after it in the same chunk?

[phemmer]

Gemini through LiteLLM

[DouweM]

@phemmer OK thanks. Have you seen the <think> tag split up over multiple chunks like in your test, or is it always in a chunk in full, even if there's text before/after it?

[phemmer]

Yes, it's been split up. I removed some parts in the middle to simplify the sample, but it is still accurate. The example I was looking at had a single < in the first chunk, had the end of the tag and the content in another chunk, and the closing tag and normal content in another chunk.

[DouweM]

@phemmer All right, that means we'll need to do some buffering when we start receiving something that looks like it could become a thinking tag.

We parse the start tag and switch into thinking mode here:

pydantic-ai/pydantic_ai_slim/pydantic_ai/_parts_manager.py

Lines 131 to 134 in bfcccba

	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='')

And we continue in thinking mode or parse the end tag and switch back to text mode here:

pydantic-ai/pydantic_ai_slim/pydantic_ai/_parts_manager.py

Lines 118 to 125 in bfcccba

	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
	else:
	return self.handle_thinking_delta(vendor_part_id=vendor_part_id, content=content)

Would you be interested in looking into a fix?

[phemmer]

Just FYI, I did look at this, but have not implemented anything because of what would be involved to make it work correctly.

The core of the issue is that the code currently assumes that handle_text_delta can only return 1 or 0 events. To properly fix this, that assumption must be broken. It needs to be able to return an arbitrary number of events. Consider the situation where a single chunk contains <think>foo</think>\nbar. This would have to emit 2 events. One for the thinking content foo, and another for the text content bar. You could buffer bar until the next time that handle_text_delta is called, and merge it with the next event. But what if the next event isn't a text event? Then you would have to emit the text event, and buffer the next event again. Or what if there are no more events? The buffer would have to be explicitly drained, which is possible, but likely should be considered as an artifact of a bad design. Basically it can't be fixed without changing the function to either a generator or returning a list, which would be a fairly significant and intrusive change.

[dsfaccini]

yeahp you're right, the current PR as it is would drop the bar completely

I'll be happy to make a new one after the right solution is discussed

[DouweM]

@phemmer @dsfaccini I think returning a generator is the way to go. It's already different from the other methods in that it can return None (0 or 1), so making it return 0+ seems fine.