Add provider_metadata to Vercel UI adapter, event stream and certain models handling

[Light2Dark]

Fixes #3748 . Pulls provider details from reasoning end blocks, and passes them as provider_metadata.

Should fix anthropic issues: https://platform.claude.com/docs/en/build-with-claude/extended-thinking#extended-thinking-with-tool-use

Affected chunks (event-stream)

• TextStartChunk
• TextDeltaChunk
• TextEndChunk
• ReasoningStartChunk
• ReasoningDeltaChunk
• ReasoningEndChunk
• ToolInputAvailableChunk

Affected parts (adapter)

• assistant's TextUIPart
• assistant's BuiltInToolCallPart
• assistant's ToolCallPart
• assistant's FileUIPart

These messages have been modified to add provider_name

• TextPart
• BaseToolCallPart (only ToolCallPart affected)
• TextDeltaPart
• ToolCallPartDelta

[DouweM]

Hmm, in Pydantic AI we can have separate provider metadata on the call and return parts, but Vercel represents them as a single part. So we need to make sure they don't accidentally get merged, somehow, so that all the fields make it back to the correct part.

[Light2Dark]

Not sure I'm getting it right but I did move getting this details to the top, so both call and return parts should be the same.

[Light2Dark]

I might be confused, I can also remove this..

[DouweM]

The problem is that in Pydantic AI, (Builtin)ToolCallPart and (Builtin)ToolReturnPart both have their own provider_details, but Vercel AI combines both into one part (ToolOutputAvailablePart etc).

We need to make sure that when we translate the 2 Pydantic AI parts into 1 Vercel AI part and back into 2 Pydantic AI parts, they end up with the original provider_details, not all details on one of the 2 parts, or both parts duplcating the same details. So I think in this case we should store 2 dicts on the Vercel AI part, so that we can separately store and extract the details for the call and those for the result

[DouweM]

Can we make this a method on VercelAIAdapter like _get_pydantic_ai_meta? To match the load/dump_messages methods, I'd call them _load_part_metadata and _dump_part_metadata

[Light2Dark]

I created those 2 methods under VercelAIAdapter, _get_pydantic_ai_meta -> _load_part_metadata and _dump_part_metadata.

However, does it mean we should remove this form_provider_metadata function? I think it'd be good to reuse the same function to standardize. If so, I need to make _dump_part_metadata function public and use it in _event_stream.py

[DouweM]

Yeah I'd rather use the same _dump_part_metadata here. I'm OK with the event stream using private methods from the adapter; I'd rather not have this be user-facing public

[Light2Dark]

Thanks for the thorough review, I'll try to get it all right.

Yeah I'd rather use the same _dump_part_metadata here. I'm OK with the event stream using private methods from the adapter; I'd rather not have this be user-facing public

Just realized this would cause a circular import (adapter imports VercelAIEventStream). I could move it to _utils.py.

Could also keep this function in _event_stream.py. Idt it's user facing since not exported in __all__ declaration.

[Light2Dark]

I've decided to move both functions to _utils.py, and move around some code to _models.py to prevent circular import errors. Let me know if this is unwanted.

[Light2Dark]

I made this change because I think it's more accurate, but it does differ from the old behaviour. Can refer to 9748bbb for the change on the test too.

[DouweM]

See above; we should store the data of both the call and the return, we can't lose any piece of data (like the 'tool_type': 'web_search_preview')

[DouweM]

The problem is that in Pydantic AI, (Builtin)ToolCallPart and (Builtin)ToolReturnPart both have their own provider_details, but Vercel AI combines both into one part (ToolOutputAvailablePart etc).

We need to make sure that when we translate the 2 Pydantic AI parts into 1 Vercel AI part and back into 2 Pydantic AI parts, they end up with the original provider_details, not all details on one of the 2 parts, or both parts duplcating the same details. So I think in this case we should store 2 dicts on the Vercel AI part, so that we can separately store and extract the details for the call and those for the result

[DouweM]

Hmm ToolCallPart doesn't have provider_name, but in some places we only use ToolCallPart.provider_details if ModelResponse.provider_name has the expected value:

pydantic-ai/pydantic_ai_slim/pydantic_ai/models/google.py

Lines 914 to 918 in b5dee9b

	if (
	item.provider_details
	and (thought_signature := item.provider_details.get('thought_signature'))
	and m.provider_name == provider_name
	):

So to make this work properly with Google though signatures, we need to make sure the provider_name is also stored on the ModelResponse we build here from the Vercel AI data. Can you see how we would get that from Pydantic AI to Vercel AI format and back on the ModelResponse?

[Light2Dark]

I don't think we are passing this data in _message_builder to convert to pydantic messages.

@dataclass
class MessagesBuilder:
    """Helper class to build Pydantic AI messages from request/response parts."""

    messages: list[ModelMessage] = field(default_factory=list)

    def add(self, part: ModelRequestPart | ModelResponsePart) -> None:
        """Add a new part, creating a new request or response message if necessary."""
        last_message = self.messages[-1] if self.messages else None
        if isinstance(part, get_union_args(ModelRequestPart)):
            part = cast(ModelRequestPart, part)
            if isinstance(last_message, ModelRequest):
                last_message.parts = [*last_message.parts, part]
            else:
                self.messages.append(ModelRequest(parts=[part])) ## <-- We don't pass additional data here
        else:
            part = cast(ModelResponsePart, part)
            if isinstance(last_message, ModelResponse):
                last_message.parts = [*last_message.parts, part]
            else:
                self.messages.append(ModelResponse(parts=[part])) ## <-- And here

Could infer details from the part

    def add(self, part: ModelRequestPart | ModelResponsePart) -> None:
        """Add a new part, creating a new request or response message if necessary."""
        ...
        else:
            part = cast(ModelResponsePart, part)
            if isinstance(last_message, ModelResponse):
                last_message.parts = [*last_message.parts, part]
            else:
                self.messages.append(ModelResponse(parts=[part], provider_name=part.provider_name)) ## <-- here

and build the metadata as parts come in.

Another option is to modify the add function to accept more arguments and pass it to ModelResponse 🤔. Anyway, I'd prefer this to be another PR

[DouweM]

I think it's worth fixing this in this PR, because the goal here is to make encrypted thinking work reliably when data goes from Pydantic AI to Vercel AI and back to Pydantic AI, which means making all fields that matter to thinking survive that round-trip, which includes this ones :(

One way to do that would be to see if we can add extra fields to add, making sure it works even if there's no part on ModelResponse.parts that actually has the provider_name, so it will need to be read off ModelResponse.provider_name itself.

But it may be better to add the provider_name field everywhere we have currently have provider_details (ToolCallPart, TextPart etc), and then making sure it's always set when where set provider_details. That may touch more code but is a lot more clean, so please give that a try.

[Light2Dark]

@DouweM , I assume you mean modifying the part like

class ToolCallPart(BaseToolCallPart):
    """A tool call from a model."""

    _: KW_ONLY

    provider_name: str | None = None.  <---- NEW
    """The name of the provider that generated the response."""
    ...

I can make the changes for _event_stream.py, _adapter.py and for Vercel AI related changes. The surface area is not large. But doing the wiring for all the providers is imo going to be very large.

Example for the issue you mentioned

ToolCallPart doesn't have provider_name, but in some places we only use ToolCallPart.provider_details

I'm guessing we would modify google.py like this or something similar:

        if (
            item.provider_details
            and (thought_signature := item.provider_details.get('thought_signature'))
            and (m.provider_name == provider_name or item.provider_name == provider_name) ## <-- check item provider
        ):

[DouweM]

@Light2Dark Yep, adding provider_name to ToolCallPart and TextPart and any others that currently have provider_details.

But doing the wiring for all the providers is imo going to be very large.

I believe it's just one instance in OpenAIResponsesModel:

• pydantic-ai/pydantic_ai_slim/pydantic_ai/models/openai.py

  Line 1290 in fbe0dfe

  items.append(TextPart(content.text, id=item.id, provider_details=part_provider_details))

And a bunch in GoogleModel: (some of which may require changes to methods like handle_tool_call_delta, but that should be straightforward):

• pydantic-ai/pydantic_ai_slim/pydantic_ai/models/google.py

  Line 858 in fbe0dfe

  part=FilePart(content=BinaryContent.narrow_type(content), provider_details=provider_details),

• pydantic-ai/pydantic_ai_slim/pydantic_ai/models/google.py

  Line 842 in fbe0dfe

  provider_details=provider_details,

• pydantic-ai/pydantic_ai_slim/pydantic_ai/models/google.py

  Lines 827 to 834 in fbe0dfe

  	for event in self._parts_manager.handle_thinking_delta(
  	vendor_part_id=None, content=part.text, provider_details=provider_details
  	):
  	yield event
  	else:
  	for event in self._parts_manager.handle_text_delta(
  	vendor_part_id=None, content=part.text, provider_details=provider_details
  	):

• pydantic-ai/pydantic_ai_slim/pydantic_ai/models/google.py

  Line 1101 in fbe0dfe

  item.provider_details = {**(item.provider_details or {}), **provider_details}

Really appreciate you working on this, we're almost there!

[DouweM]

See above; we should store the data of both the call and the return, we can't lose any piece of data (like the 'tool_type': 'web_search_preview')

[DouweM]

Yeah I'd rather use the same _dump_part_metadata here. I'm OK with the event stream using private methods from the adapter; I'd rather not have this be user-facing public