From 98c032229e28c38faa85a4add771d88e959d50b4 Mon Sep 17 00:00:00 2001 From: Andrey Buzin Date: Fri, 8 May 2026 10:32:42 -0700 Subject: [PATCH 1/8] Add provider metadata field on all datamodel units --- src/ai/types/events.py | 3 ++- src/ai/types/messages.py | 22 +++++++++++++++------- src/ai/types/metadata.py | 5 +++++ 3 files changed, 22 insertions(+), 8 deletions(-) create mode 100644 src/ai/types/metadata.py diff --git a/src/ai/types/events.py b/src/ai/types/events.py index e3bf0f62..e5a451bb 100644 --- a/src/ai/types/events.py +++ b/src/ai/types/events.py @@ -4,7 +4,7 @@ import pydantic -from . import messages +from . import messages, metadata from . import usage as usage_ # we're using pydantic because events are crossing @@ -34,6 +34,7 @@ class BaseEvent(pydantic.BaseModel): message: messages.Message = _DUMMY_MESSAGE usage: usage_.Usage | None = None + provider_metadata: pydantic.SerializeAsAny[metadata.ProviderMetadata] | None = None replay: bool = pydantic.Field(default=False, exclude=True, repr=False) model_config = pydantic.ConfigDict(frozen=True) diff --git a/src/ai/types/messages.py b/src/ai/types/messages.py index cfb11e0a..1d62769f 100644 --- a/src/ai/types/messages.py +++ b/src/ai/types/messages.py @@ -4,7 +4,7 @@ import pydantic -from . import media +from . import media, metadata from . import usage as usage_ @@ -17,6 +17,7 @@ def generate_id(prefix: str | None = None) -> str: class TextPart(pydantic.BaseModel): id: str = pydantic.Field(default_factory=generate_id) text: str + provider_metadata: pydantic.SerializeAsAny[metadata.ProviderMetadata] | None = None kind: Literal["text"] = "text" @@ -26,6 +27,7 @@ class ToolCallPart(pydantic.BaseModel): tool_call_id: str tool_name: str tool_args: str + provider_metadata: pydantic.SerializeAsAny[metadata.ProviderMetadata] | None = None kind: Literal["tool_call"] = "tool_call" @@ -41,6 +43,7 @@ class ToolResultPart(pydantic.BaseModel): tool_name: str result: Any = None is_error: bool = False + provider_metadata: pydantic.SerializeAsAny[metadata.ProviderMetadata] | None = None kind: Literal["tool_result"] = "tool_result" model_config = pydantic.ConfigDict(frozen=True) @@ -57,7 +60,8 @@ class BuiltinToolCallPart(pydantic.BaseModel): tool_call_id: str tool_name: str tool_args: str = "" - provider_name: str | None = None + provider_name: str | None = None # TODO replace with provider_metadata + provider_metadata: pydantic.SerializeAsAny[metadata.ProviderMetadata] | None = None kind: Literal["builtin_tool_call"] = "builtin_tool_call" @@ -71,7 +75,10 @@ class BuiltinToolReturnPart(pydantic.BaseModel): result: Any = None is_error: bool = False provider_name: str | None = None - provider_details: dict[str, Any] | None = None + provider_details: dict[str, Any] | None = ( + None # TODO replace with provider_metadata + ) + provider_metadata: pydantic.SerializeAsAny[metadata.ProviderMetadata] | None = None kind: Literal["builtin_tool_return"] = "builtin_tool_return" model_config = pydantic.ConfigDict(frozen=True) @@ -82,7 +89,8 @@ class ReasoningPart(pydantic.BaseModel): text: str # Anthropic's thinking blocks include a signature for cache/verification. # This must be preserved and sent back in multi-turn conversations. - signature: str | None = None + signature: str | None = None # TODO replace with provider_metadata + provider_metadata: pydantic.SerializeAsAny[metadata.ProviderMetadata] | None = None kind: Literal["reasoning"] = "reasoning" @@ -136,6 +144,7 @@ class StructuredOutputPart(pydantic.BaseModel): data: dict[str, Any] output_type_name: str kind: Literal["structured_output"] = "structured_output" + provider_metadata: pydantic.SerializeAsAny[metadata.ProviderMetadata] | None = None _hydrated: Any = pydantic.PrivateAttr(default=None) @@ -169,6 +178,7 @@ class FilePart(pydantic.BaseModel): media_type: str # IANA media type, e.g. "image/png", "audio/wav" filename: str | None = None kind: Literal["file"] = "file" + provider_metadata: pydantic.SerializeAsAny[metadata.ProviderMetadata] | None = None @classmethod def from_url(cls, url: str, *, media_type: str | None = None) -> Self: @@ -228,6 +238,7 @@ class Message(pydantic.BaseModel): id: str = pydantic.Field(default_factory=generate_id) turn_id: str | None = None usage: usage_.Usage | None = None + provider_metadata: pydantic.SerializeAsAny[metadata.ProviderMetadata] | None = None # Set on the seeded message that ``models.stream`` returns when # short-circuiting an existing assistant turn (resume-after-approval @@ -282,6 +293,3 @@ def output(self) -> Any: if isinstance(part, StructuredOutputPart): return part.value return None - - -Usage = usage_.Usage diff --git a/src/ai/types/metadata.py b/src/ai/types/metadata.py new file mode 100644 index 00000000..a533ea4f --- /dev/null +++ b/src/ai/types/metadata.py @@ -0,0 +1,5 @@ +import pydantic + + +class ProviderMetadata(pydantic.BaseModel): + model_config = pydantic.ConfigDict(frozen=True, populate_by_name=True) From 5dd7d623ef06bcb7c1eadb6079ebcae24946a504 Mon Sep 17 00:00:00 2001 From: Andrey Buzin Date: Fri, 8 May 2026 10:53:07 -0700 Subject: [PATCH 2/8] Merge metadata from events into parts as latest wins --- src/ai/models/core/api.py | 84 +++++++++++++++++++++++++++++---------- 1 file changed, 64 insertions(+), 20 deletions(-) diff --git a/src/ai/models/core/api.py b/src/ai/models/core/api.py index 789f0b53..e227c4ca 100644 --- a/src/ai/models/core/api.py +++ b/src/ai/models/core/api.py @@ -195,48 +195,76 @@ def _aggregate_event(self, event: types.events.Event) -> dict[str, Any]: self._message.usage = event.usage match event: - case types.events.TextStart(block_id=bid): - tp = types.messages.TextPart(id=bid, text="") + case types.events.TextStart(block_id=bid, provider_metadata=pm): + tp = types.messages.TextPart(id=bid, text="", provider_metadata=pm) self._message.parts.append(tp) self._parts[bid] = tp - case types.events.TextDelta(block_id=bid, chunk=c): + case types.events.TextDelta(block_id=bid, chunk=c, provider_metadata=pm): existing_text = self._parts.get(bid) if isinstance(existing_text, types.messages.TextPart): existing_text.text += c - case types.events.ReasoningStart(block_id=bid): - rp = types.messages.ReasoningPart(id=bid, text="") + if pm is not None: + existing_text.provider_metadata = pm + case types.events.TextEnd(block_id=bid, provider_metadata=pm): + existing_text = self._parts.get(bid) + if ( + isinstance(existing_text, types.messages.TextPart) + and pm is not None + ): + existing_text.provider_metadata = pm + case types.events.ReasoningStart(block_id=bid, provider_metadata=pm): + rp = types.messages.ReasoningPart(id=bid, text="", provider_metadata=pm) self._message.parts.append(rp) self._parts[bid] = rp - case types.events.ReasoningDelta(block_id=bid, chunk=c): + case types.events.ReasoningDelta( + block_id=bid, chunk=c, provider_metadata=pm + ): existing_reasoning = self._parts.get(bid) if isinstance(existing_reasoning, types.messages.ReasoningPart): existing_reasoning.text += c - case types.events.ReasoningEnd(block_id=bid, signature=sig): + if pm is not None: + existing_reasoning.provider_metadata = pm + case types.events.ReasoningEnd( + block_id=bid, signature=sig, provider_metadata=pm + ): existing_reasoning = self._parts.get(bid) - if ( - isinstance(existing_reasoning, types.messages.ReasoningPart) - and sig is not None - ): - existing_reasoning.signature = sig - case types.events.ToolStart(tool_call_id=tcid, tool_name=name): + if isinstance(existing_reasoning, types.messages.ReasoningPart): + if sig is not None: + existing_reasoning.signature = sig + if pm is not None: + existing_reasoning.provider_metadata = pm + case types.events.ToolStart( + tool_call_id=tcid, tool_name=name, provider_metadata=pm + ): tcp = types.messages.ToolCallPart( id=tcid, tool_call_id=tcid, tool_name=name, tool_args="", + provider_metadata=pm, ) self._message.parts.append(tcp) self._parts[tcid] = tcp - case types.events.ToolDelta(tool_call_id=tcid, chunk=c): + case types.events.ToolDelta( + tool_call_id=tcid, chunk=c, provider_metadata=pm + ): existing_tool = self._parts.get(tcid) if isinstance(existing_tool, types.messages.ToolCallPart): existing_tool.tool_args += c - case types.events.ToolEnd(tool_call_id=tcid): + if pm is not None: + existing_tool.provider_metadata = pm + + case types.events.ToolEnd(tool_call_id=tcid, provider_metadata=pm): existing_tool = self._parts.get(tcid) if isinstance(existing_tool, types.messages.ToolCallPart): updates["tool_call"] = existing_tool + if pm is not None: + existing_tool.provider_metadata = pm case types.events.BuiltinToolStart( - tool_call_id=tcid, tool_name=name, provider_name=pname + tool_call_id=tcid, + tool_name=name, + provider_name=pname, + provider_metadata=pm, ): btcp = types.messages.BuiltinToolCallPart( id=tcid, @@ -244,30 +272,46 @@ def _aggregate_event(self, event: types.events.Event) -> dict[str, Any]: tool_name=name, tool_args="", provider_name=pname, + provider_metadata=pm, ) self._message.parts.append(btcp) self._parts[tcid] = btcp - case types.events.BuiltinToolDelta(tool_call_id=tcid, chunk=c): + case types.events.BuiltinToolDelta( + tool_call_id=tcid, chunk=c, provider_metadata=pm + ): existing_btc = self._parts.get(tcid) if isinstance(existing_btc, types.messages.BuiltinToolCallPart): existing_btc.tool_args += c - case types.events.BuiltinToolEnd(tool_call_id=tcid): + if pm is not None: + existing_btc.provider_metadata = pm + case types.events.BuiltinToolEnd(tool_call_id=tcid, provider_metadata=pm): existing_btc = self._parts.get(tcid) if isinstance(existing_btc, types.messages.BuiltinToolCallPart): updates["tool_call"] = existing_btc - case types.events.BuiltinToolResult(result=res): + if pm is not None: + existing_btc.provider_metadata = pm + case types.events.BuiltinToolResult(result=res, provider_metadata=pm): self._message.parts.append(res) case types.events.FileEvent( - block_id=bid, media_type=mt, data=d, filename=fname + block_id=bid, + media_type=mt, + data=d, + filename=fname, + provider_metadata=pm, ): fp = types.messages.FilePart( id=bid or types.messages.generate_id(), data=d, media_type=mt, filename=fname, + provider_metadata=pm, ) self._message.parts.append(fp) self._parts[fp.id] = fp + + case types.events.StreamEnd(provider_metadata=pm): + if pm is not None: + self._message.provider_metadata = pm case _: pass From 1e43b1a82e857f7db2bf4ef3191f5c91ddc4a321 Mon Sep 17 00:00:00 2001 From: Andrey Buzin Date: Fri, 8 May 2026 10:57:28 -0700 Subject: [PATCH 3/8] Add a test for provider metadata --- tests/models/core/test_api.py | 97 +++++++++++++++++++++++++++++++++++ 1 file changed, 97 insertions(+) diff --git a/tests/models/core/test_api.py b/tests/models/core/test_api.py index 3a460887..5a45c969 100644 --- a/tests/models/core/test_api.py +++ b/tests/models/core/test_api.py @@ -10,6 +10,7 @@ from ai import models from ai.models.openai import openai from ai.types import events as events_ +from ai.types import metadata as metadata_ from ai.types import messages as messages_ from ...conftest import MOCK_MODEL, MOCK_PROVIDER, MockProvider, mock_llm, text_msg @@ -19,6 +20,17 @@ class _MockStreamParams(pydantic.BaseModel): value: str +class _TestProviderMetadata(metadata_.ProviderMetadata): + marker: str + + +def _provider_metadata_marker( + provider_metadata: metadata_.ProviderMetadata | None, +) -> str: + assert isinstance(provider_metadata, _TestProviderMetadata) + return provider_metadata.marker + + async def test_stream_aggregates_registered_adapter_events() -> None: mock = mock_llm([[text_msg("Hello world")]]) @@ -68,6 +80,91 @@ async def _tool_stream( assert stream.tool_calls == [tool_end.tool_call] +async def test_stream_accumulates_provider_metadata_latest_wins() -> None: + async def _metadata_stream() -> AsyncGenerator[events_.Event]: + yield events_.StreamStart() + yield events_.TextStart( + block_id="text", + provider_metadata=_TestProviderMetadata(marker="text-start"), + ) + yield events_.TextDelta(block_id="text", chunk="hello") + yield events_.TextDelta( + block_id="text", + chunk=" world", + provider_metadata=_TestProviderMetadata(marker="text-delta"), + ) + yield events_.TextEnd( + block_id="text", + provider_metadata=_TestProviderMetadata(marker="text-end"), + ) + yield events_.ReasoningStart( + block_id="reasoning", + provider_metadata=_TestProviderMetadata(marker="reasoning-start"), + ) + yield events_.ReasoningDelta( + block_id="reasoning", + chunk="thinking", + provider_metadata=_TestProviderMetadata(marker="reasoning-delta"), + ) + yield events_.ReasoningEnd( + block_id="reasoning", + signature="sig-1", + provider_metadata=_TestProviderMetadata(marker="reasoning-end"), + ) + yield events_.ToolStart( + tool_call_id="tc-1", + tool_name="weather", + provider_metadata=_TestProviderMetadata(marker="tool-start"), + ) + yield events_.ToolDelta(tool_call_id="tc-1", chunk='{"city"') + yield events_.ToolDelta( + tool_call_id="tc-1", + chunk=':"SF"}', + provider_metadata=_TestProviderMetadata(marker="tool-delta"), + ) + yield events_.ToolEnd( + tool_call_id="tc-1", + tool_call=messages_.DUMMY_TOOL_CALL, + provider_metadata=_TestProviderMetadata(marker="tool-end"), + ) + yield events_.FileEvent( + block_id="file", + media_type="image/png", + data="base64-data", + provider_metadata=_TestProviderMetadata(marker="file"), + ) + yield events_.StreamEnd( + provider_metadata=_TestProviderMetadata(marker="message"), + ) + + stream = models.Stream(_metadata_stream()) + async for _ in stream: + pass + + assert _provider_metadata_marker(stream.message.provider_metadata) == "message" + + text = stream.message.parts[0] + assert isinstance(text, messages_.TextPart) + assert text.text == "hello world" + assert _provider_metadata_marker(text.provider_metadata) == "text-end" + + reasoning = stream.message.parts[1] + assert isinstance(reasoning, messages_.ReasoningPart) + assert reasoning.text == "thinking" + assert reasoning.signature == "sig-1" + assert _provider_metadata_marker(reasoning.provider_metadata) == "reasoning-end" + + tool_call = stream.message.parts[2] + assert isinstance(tool_call, messages_.ToolCallPart) + assert tool_call.tool_args == '{"city":"SF"}' + assert _provider_metadata_marker(tool_call.provider_metadata) == "tool-end" + + file = stream.message.parts[3] + assert isinstance(file, messages_.FilePart) + assert file.data == "base64-data" + assert _provider_metadata_marker(file.provider_metadata) == "file" + + async def test_stream_uses_explicit_model_client() -> None: received_clients: list[models.Client] = [] From 2702d33cf82404f92a42fbed4a41fce907f0122a Mon Sep 17 00:00:00 2001 From: Andrey Buzin Date: Fri, 8 May 2026 11:21:57 -0700 Subject: [PATCH 4/8] Add anthropic-specific metadata types --- src/ai/models/anthropic/metadata.py | 167 ++++++++++++++++++++++++++++ src/ai/types/__init__.py | 3 +- 2 files changed, 169 insertions(+), 1 deletion(-) create mode 100644 src/ai/models/anthropic/metadata.py diff --git a/src/ai/models/anthropic/metadata.py b/src/ai/models/anthropic/metadata.py new file mode 100644 index 00000000..43b121d1 --- /dev/null +++ b/src/ai/models/anthropic/metadata.py @@ -0,0 +1,167 @@ +from __future__ import annotations + +from typing import Annotated, Any, Literal + +import pydantic + +from ... import types +from . import params + +_METADATA_CONFIG = pydantic.ConfigDict(frozen=True, populate_by_name=True) + + +class AnthropicCitations(pydantic.BaseModel): + model_config = _METADATA_CONFIG + + enabled: bool + + +class AnthropicCaller(pydantic.BaseModel): + model_config = _METADATA_CONFIG + + type: str + tool_id: str | None = pydantic.Field( + default=None, + validation_alias="toolId", + serialization_alias="toolId", + ) + + +class AnthropicUsageIteration(pydantic.BaseModel): + model_config = _METADATA_CONFIG + + type: Literal["compaction", "message"] + input_tokens: int = pydantic.Field( + validation_alias="inputTokens", + serialization_alias="inputTokens", + ) + output_tokens: int = pydantic.Field( + validation_alias="outputTokens", + serialization_alias="outputTokens", + ) + + +class AnthropicContainerSkill(pydantic.BaseModel): + model_config = _METADATA_CONFIG + + type: Literal["anthropic", "custom"] + skill_id: str = pydantic.Field( + validation_alias="skillId", + serialization_alias="skillId", + ) + version: str + + +class AnthropicContainerMetadata(pydantic.BaseModel): + model_config = _METADATA_CONFIG + + expires_at: str = pydantic.Field( + validation_alias="expiresAt", + serialization_alias="expiresAt", + ) + id: str + skills: list[AnthropicContainerSkill] | None = None + + +class AnthropicClearToolUsesEdit(pydantic.BaseModel): + model_config = _METADATA_CONFIG + + type: Literal["clear_tool_uses_20250919"] + cleared_tool_uses: int = pydantic.Field( + validation_alias="clearedToolUses", + serialization_alias="clearedToolUses", + ) + cleared_input_tokens: int = pydantic.Field( + validation_alias="clearedInputTokens", + serialization_alias="clearedInputTokens", + ) + + +class AnthropicClearThinkingEdit(pydantic.BaseModel): + model_config = _METADATA_CONFIG + + type: Literal["clear_thinking_20251015"] + cleared_thinking_turns: int = pydantic.Field( + validation_alias="clearedThinkingTurns", + serialization_alias="clearedThinkingTurns", + ) + cleared_input_tokens: int = pydantic.Field( + validation_alias="clearedInputTokens", + serialization_alias="clearedInputTokens", + ) + + +class AnthropicCompactEdit(pydantic.BaseModel): + model_config = _METADATA_CONFIG + + type: Literal["compact_20260112"] + + +type AnthropicContextManagementEdit = Annotated[ + AnthropicClearToolUsesEdit | AnthropicClearThinkingEdit | AnthropicCompactEdit, + pydantic.Field(discriminator="type"), +] + + +class AnthropicContextManagementMetadata(pydantic.BaseModel): + model_config = _METADATA_CONFIG + + applied_edits: list[AnthropicContextManagementEdit] = pydantic.Field( + validation_alias="appliedEdits", + serialization_alias="appliedEdits", + ) + + +class AnthropicProviderMetadata(types.metadata.ProviderMetadata): + """Anthropic-specific metadata for messages, parts, and stream events.""" + + provider: Literal["anthropic"] = "anthropic" + + # Input/replay options. + cache_control: params.AnthropicCacheControl | None = pydantic.Field( + default=None, + validation_alias="cacheControl", + serialization_alias="cacheControl", + ) + citations: AnthropicCitations | None = None + title: str | None = None + context: str | None = None + + # Special part markers. + type: Literal["compaction", "mcp-tool-use"] | None = None + server_name: str | None = pydantic.Field( + default=None, + validation_alias="serverName", + serialization_alias="serverName", + ) + + # Reasoning replay metadata. + signature: str | None = None + redacted_data: str | None = pydantic.Field( + default=None, + validation_alias="redactedData", + serialization_alias="redactedData", + ) + + # Tool replay metadata. + result_type: str | None = pydantic.Field( + default=None, + validation_alias="resultType", + serialization_alias="resultType", + ) + caller: AnthropicCaller | dict[str, Any] | None = None + + # Response/message metadata. + usage: dict[str, Any] | None = None + stop_sequence: str | None = pydantic.Field( + default=None, + validation_alias="stopSequence", + serialization_alias="stopSequence", + ) + iterations: list[AnthropicUsageIteration] | None = None + container: AnthropicContainerMetadata | None = None + context_management: AnthropicContextManagementMetadata | None = pydantic.Field( + default=None, + validation_alias="contextManagement", + serialization_alias="contextManagement", + ) diff --git a/src/ai/types/__init__.py b/src/ai/types/__init__.py index ff62fc5b..937a3a1a 100644 --- a/src/ai/types/__init__.py +++ b/src/ai/types/__init__.py @@ -1,9 +1,10 @@ -from . import events, media, messages, proto, tools, usage +from . import events, media, messages, metadata, proto, tools, usage __all__ = [ "events", "media", "messages", + "metadata", "proto", "tools", "usage", From b4bcfbf59d9b087c03532506881f1cc25d5a0483 Mon Sep 17 00:00:00 2001 From: Andrey Buzin Date: Fri, 8 May 2026 11:34:12 -0700 Subject: [PATCH 5/8] Namespace imports in anthropic adapter --- src/ai/models/anthropic/adapter.py | 37 ++++++++++++++---------------- 1 file changed, 17 insertions(+), 20 deletions(-) diff --git a/src/ai/models/anthropic/adapter.py b/src/ai/models/anthropic/adapter.py index 1d743e93..cdc78f31 100644 --- a/src/ai/models/anthropic/adapter.py +++ b/src/ai/models/anthropic/adapter.py @@ -13,16 +13,9 @@ from ... import types from ...types import events -from ...types import messages as messages_ from .. import core +from . import params from . import tools as anthropic_tools -from .params import ( - AnthropicContainer, - AnthropicCustomSkill, - AnthropicDisabledThinking, - AnthropicParams, - AnthropicProviderSkill, -) PROVIDER_NAME = "anthropic" @@ -343,19 +336,23 @@ def _make_client( ) -def _coerce_anthropic_params(value: Any) -> AnthropicParams: +def _coerce_anthropic_params(value: Any) -> params.AnthropicParams: if value is None: - return AnthropicParams() - if isinstance(value, AnthropicParams): + return params.AnthropicParams() + if isinstance(value, params.AnthropicParams): return value if isinstance(value, Sequence) and not isinstance(value, str | bytes | bytearray): items = list(value) - if len(items) == 1 and isinstance(items[0], AnthropicParams): + if len(items) == 1 and isinstance(items[0], params.AnthropicParams): return items[0] - raise TypeError(f"anthropic stream params must be {AnthropicParams.__name__}") + raise TypeError( + f"anthropic stream params must be {params.AnthropicParams.__name__}" + ) -def _container_to_wire(container: AnthropicContainer) -> str | dict[str, Any] | None: +def _container_to_wire( + container: params.AnthropicContainer, +) -> str | dict[str, Any] | None: if not container.skills: return container.id @@ -367,9 +364,9 @@ def _container_to_wire(container: AnthropicContainer) -> str | dict[str, Any] | def _skill_to_wire( - skill: AnthropicProviderSkill | AnthropicCustomSkill, + skill: params.AnthropicProviderSkill | params.AnthropicCustomSkill, ) -> dict[str, Any]: - if isinstance(skill, AnthropicProviderSkill): + if isinstance(skill, params.AnthropicProviderSkill): result: dict[str, Any] = { "type": "anthropic", "skill_id": skill.skill_id, @@ -466,7 +463,7 @@ async def stream( api_kwargs["tools"] = wire_tools if anthropic_params.thinking is not None: - if not isinstance(anthropic_params.thinking, AnthropicDisabledThinking): + if not isinstance(anthropic_params.thinking, params.AnthropicDisabledThinking): api_kwargs["thinking"] = anthropic_params.thinking.model_dump( exclude_none=True, ) @@ -623,14 +620,14 @@ async def stream( if tool_id: yield events.ToolEnd( tool_call_id=tool_id, - tool_call=messages_.DUMMY_TOOL_CALL, + tool_call=types.messages.DUMMY_TOOL_CALL, ) elif block_type == "server_tool_use": tool_id = tool_ids.get(idx) if tool_id: yield events.BuiltinToolEnd( tool_call_id=tool_id, - tool_call=messages_.BuiltinToolCallPart( + tool_call=types.messages.BuiltinToolCallPart( tool_call_id=tool_id, tool_name=tool_names.get(idx, ""), provider_name=PROVIDER_NAME, @@ -662,7 +659,7 @@ async def stream( break yield events.BuiltinToolResult( tool_call_id=tool_use_id, - result=messages_.BuiltinToolReturnPart( + result=types.messages.BuiltinToolReturnPart( tool_call_id=tool_use_id, tool_name=tool_name, result=content_payload, From 5ef7846828a2d5325296d8360c706dd36023d994 Mon Sep 17 00:00:00 2001 From: Andrey Buzin Date: Fri, 8 May 2026 12:52:28 -0700 Subject: [PATCH 6/8] Add a missing case in metadata merge, allow extras in metadata model --- src/ai/models/core/api.py | 2 ++ src/ai/types/metadata.py | 6 +++++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/src/ai/models/core/api.py b/src/ai/models/core/api.py index e227c4ca..f0de0c03 100644 --- a/src/ai/models/core/api.py +++ b/src/ai/models/core/api.py @@ -291,6 +291,8 @@ def _aggregate_event(self, event: types.events.Event) -> dict[str, Any]: if pm is not None: existing_btc.provider_metadata = pm case types.events.BuiltinToolResult(result=res, provider_metadata=pm): + if pm is not None: + res = res.model_copy(update={"provider_metadata": pm}) self._message.parts.append(res) case types.events.FileEvent( block_id=bid, diff --git a/src/ai/types/metadata.py b/src/ai/types/metadata.py index a533ea4f..d0089f69 100644 --- a/src/ai/types/metadata.py +++ b/src/ai/types/metadata.py @@ -2,4 +2,8 @@ class ProviderMetadata(pydantic.BaseModel): - model_config = pydantic.ConfigDict(frozen=True, populate_by_name=True) + model_config = pydantic.ConfigDict( + frozen=True, + populate_by_name=True, + extra="allow", + ) From 77646f2d5899a63744dafa87d2e0e393967a78b8 Mon Sep 17 00:00:00 2001 From: Andrey Buzin Date: Fri, 8 May 2026 13:35:08 -0700 Subject: [PATCH 7/8] Fold existing metadata-related fields into metadata --- src/ai/models/anthropic/adapter.py | 59 ++++++++++++++++++++------ src/ai/models/core/api.py | 16 +++---- src/ai/types/builders.py | 9 +++- src/ai/types/events.py | 2 - src/ai/types/messages.py | 8 ---- tests/conftest.py | 5 ++- tests/models/ai_gateway/test_stream.py | 2 - tests/models/anthropic/test_adapter.py | 52 +++++++++++++++++++++-- tests/models/anthropic/test_stream.py | 32 +++++++++++++- tests/models/core/test_api.py | 4 +- tests/models/openai/test_adapter.py | 1 - 11 files changed, 141 insertions(+), 49 deletions(-) diff --git a/src/ai/models/anthropic/adapter.py b/src/ai/models/anthropic/adapter.py index cdc78f31..3d3475c0 100644 --- a/src/ai/models/anthropic/adapter.py +++ b/src/ai/models/anthropic/adapter.py @@ -14,6 +14,7 @@ from ... import types from ...types import events from .. import core +from . import metadata as anthropic_metadata from . import params from . import tools as anthropic_tools @@ -31,6 +32,7 @@ } ) + # --------------------------------------------------------------------------- # Message / tool conversion — internal types → Anthropic wire format # --------------------------------------------------------------------------- @@ -194,8 +196,22 @@ async def _messages_to_anthropic( for part in msg.parts: match part: case types.messages.ReasoningPart( - text=text, signature=signature + text=text, + provider_metadata=provider_metadata, ): + part_metadata = ( + provider_metadata + if isinstance( + provider_metadata, + anthropic_metadata.AnthropicProviderMetadata, + ) + else None + ) + signature = ( + part_metadata.signature + if part_metadata is not None + else None + ) if send_reasoning and signature: content.append( { @@ -232,12 +248,21 @@ async def _messages_to_anthropic( ) case types.messages.BuiltinToolReturnPart(): # Result block type comes from the original wire - # event ("web_search_tool_result", etc.); stored - # in provider_details when emitted. - details = part.provider_details or {} - wire_type = details.get( - "result_type", - f"{part.tool_name}_tool_result", + # event ("web_search_tool_result", etc.); stored in + # provider metadata when emitted. + part_metadata = ( + part.provider_metadata + if isinstance( + part.provider_metadata, + anthropic_metadata.AnthropicProviderMetadata, + ) + else None + ) + wire_type = ( + part_metadata.result_type + if part_metadata is not None + and part_metadata.result_type is not None + else f"{part.tool_name}_tool_result" ) content.append( { @@ -566,7 +591,7 @@ async def stream( yield events.BuiltinToolStart( tool_call_id=block.id, tool_name=block.name, - provider_name=PROVIDER_NAME, + provider_metadata=anthropic_metadata.AnthropicProviderMetadata(), ) # Result blocks (web_search_tool_result etc.) arrive # complete; we emit on stop so we have full content. @@ -611,9 +636,16 @@ async def stream( if block_type == "text": yield events.TextEnd(block_id=str(idx)) elif block_type == "thinking": + signature = signature_buffer.get(idx) yield events.ReasoningEnd( block_id=str(idx), - signature=signature_buffer.get(idx), + provider_metadata=( + anthropic_metadata.AnthropicProviderMetadata( + signature=signature + ) + if signature is not None + else None + ), ) elif block_type == "tool_use": tool_id = tool_ids.get(idx) @@ -630,7 +662,7 @@ async def stream( tool_call=types.messages.BuiltinToolCallPart( tool_call_id=tool_id, tool_name=tool_names.get(idx, ""), - provider_name=PROVIDER_NAME, + provider_metadata=anthropic_metadata.AnthropicProviderMetadata(), ), ) elif block_type in _TOOL_RESULT_BLOCK_TYPES: @@ -663,10 +695,9 @@ async def stream( tool_call_id=tool_use_id, tool_name=tool_name, result=content_payload, - provider_name=PROVIDER_NAME, - provider_details={ - "result_type": block_type or "", - }, + provider_metadata=anthropic_metadata.AnthropicProviderMetadata( + result_type=block_type or "" + ), ), ) diff --git a/src/ai/models/core/api.py b/src/ai/models/core/api.py index f0de0c03..2f41206b 100644 --- a/src/ai/models/core/api.py +++ b/src/ai/models/core/api.py @@ -224,15 +224,13 @@ def _aggregate_event(self, event: types.events.Event) -> dict[str, Any]: existing_reasoning.text += c if pm is not None: existing_reasoning.provider_metadata = pm - case types.events.ReasoningEnd( - block_id=bid, signature=sig, provider_metadata=pm - ): + case types.events.ReasoningEnd(block_id=bid, provider_metadata=pm): existing_reasoning = self._parts.get(bid) - if isinstance(existing_reasoning, types.messages.ReasoningPart): - if sig is not None: - existing_reasoning.signature = sig - if pm is not None: - existing_reasoning.provider_metadata = pm + if ( + isinstance(existing_reasoning, types.messages.ReasoningPart) + and pm is not None + ): + existing_reasoning.provider_metadata = pm case types.events.ToolStart( tool_call_id=tcid, tool_name=name, provider_metadata=pm ): @@ -263,7 +261,6 @@ def _aggregate_event(self, event: types.events.Event) -> dict[str, Any]: case types.events.BuiltinToolStart( tool_call_id=tcid, tool_name=name, - provider_name=pname, provider_metadata=pm, ): btcp = types.messages.BuiltinToolCallPart( @@ -271,7 +268,6 @@ def _aggregate_event(self, event: types.events.Event) -> dict[str, Any]: tool_call_id=tcid, tool_name=name, tool_args="", - provider_name=pname, provider_metadata=pm, ) self._message.parts.append(btcp) diff --git a/src/ai/types/builders.py b/src/ai/types/builders.py index 35c178ad..9a901ec8 100644 --- a/src/ai/types/builders.py +++ b/src/ai/types/builders.py @@ -23,6 +23,7 @@ ToolCallPart, ToolResultPart, ) +from .metadata import ProviderMetadata _PART_TYPES = ( TextPart, @@ -91,12 +92,16 @@ def file_part( return FilePart.from_bytes(data, media_type=media_type, filename=filename) -def thinking(text: str, *, signature: str | None = None) -> ReasoningPart: +def thinking( + text: str, + *, + provider_metadata: ProviderMetadata | None = None, +) -> ReasoningPart: """Create a :class:`ReasoningPart`. Useful for replaying conversation history that includes model reasoning. """ - return ReasoningPart(text=text, signature=signature) + return ReasoningPart(text=text, provider_metadata=provider_metadata) def _tool_results_from_messages(messages: list[Message]) -> list[ToolResultPart]: diff --git a/src/ai/types/events.py b/src/ai/types/events.py index e5a451bb..f192cb71 100644 --- a/src/ai/types/events.py +++ b/src/ai/types/events.py @@ -82,7 +82,6 @@ class ReasoningDelta(BaseEvent): class ReasoningEnd(BaseEvent): block_id: str = "" - signature: str | None = None kind: Literal["reasoning_end"] = "reasoning_end" @@ -111,7 +110,6 @@ class ToolEnd(BaseEvent): class BuiltinToolStart(BaseEvent): tool_call_id: str = "" tool_name: str = "" - provider_name: str | None = None kind: Literal["builtin_tool_start"] = "builtin_tool_start" diff --git a/src/ai/types/messages.py b/src/ai/types/messages.py index 1d62769f..a28d6c6e 100644 --- a/src/ai/types/messages.py +++ b/src/ai/types/messages.py @@ -60,7 +60,6 @@ class BuiltinToolCallPart(pydantic.BaseModel): tool_call_id: str tool_name: str tool_args: str = "" - provider_name: str | None = None # TODO replace with provider_metadata provider_metadata: pydantic.SerializeAsAny[metadata.ProviderMetadata] | None = None kind: Literal["builtin_tool_call"] = "builtin_tool_call" @@ -74,10 +73,6 @@ class BuiltinToolReturnPart(pydantic.BaseModel): tool_name: str result: Any = None is_error: bool = False - provider_name: str | None = None - provider_details: dict[str, Any] | None = ( - None # TODO replace with provider_metadata - ) provider_metadata: pydantic.SerializeAsAny[metadata.ProviderMetadata] | None = None kind: Literal["builtin_tool_return"] = "builtin_tool_return" @@ -87,9 +82,6 @@ class BuiltinToolReturnPart(pydantic.BaseModel): class ReasoningPart(pydantic.BaseModel): id: str = pydantic.Field(default_factory=generate_id) text: str - # Anthropic's thinking blocks include a signature for cache/verification. - # This must be preserved and sent back in multi-turn conversations. - signature: str | None = None # TODO replace with provider_metadata provider_metadata: pydantic.SerializeAsAny[metadata.ProviderMetadata] | None = None kind: Literal["reasoning"] = "reasoning" diff --git a/tests/conftest.py b/tests/conftest.py index aef80560..0a961da5 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -122,7 +122,10 @@ async def emit_events_for_messages( yield events_.ReasoningStart(block_id=bid) if part.text: yield events_.ReasoningDelta(block_id=bid, chunk=part.text) - yield events_.ReasoningEnd(block_id=bid, signature=part.signature) + yield events_.ReasoningEnd( + block_id=bid, + provider_metadata=part.provider_metadata, + ) elif isinstance(part, messages_.ToolCallPart): yield events_.ToolStart( diff --git a/tests/models/ai_gateway/test_stream.py b/tests/models/ai_gateway/test_stream.py index f289bec3..4bd8a712 100644 --- a/tests/models/ai_gateway/test_stream.py +++ b/tests/models/ai_gateway/test_stream.py @@ -507,13 +507,11 @@ def handler(req: httpx.Request) -> httpx.Response: tool_call_id="srvtoolu_1", tool_name="web_search", tool_args='{"q":"weather"}', - provider_name="anthropic", ) ret = messages.BuiltinToolReturnPart( tool_call_id="srvtoolu_1", tool_name="web_search", result=[{"title": "Forecast"}], - provider_name="anthropic", ) convo = [ user_msg("weather?"), diff --git a/tests/models/anthropic/test_adapter.py b/tests/models/anthropic/test_adapter.py index 9031dda9..2261ac77 100644 --- a/tests/models/anthropic/test_adapter.py +++ b/tests/models/anthropic/test_adapter.py @@ -16,6 +16,7 @@ import ai from ai import models from ai.models.anthropic import adapter, anthropic +from ai.models.anthropic import metadata as anthropic_metadata from ai.models.anthropic import params as anthropic_params from ai.types import messages @@ -161,7 +162,14 @@ async def test_send_reasoning_false_strips_thinking_blocks( _TEST_CLIENT, _MODEL, [ - ai.assistant_message(ai.thinking("hidden", signature="sig")), + ai.assistant_message( + ai.thinking( + "hidden", + provider_metadata=anthropic_metadata.AnthropicProviderMetadata( + signature="sig" + ), + ) + ), ai.user_message("Hi"), ], params=anthropic_params.AnthropicParams(send_reasoning=False), @@ -173,6 +181,41 @@ async def test_send_reasoning_false_strips_thinking_blocks( assert captured["messages"] == [{"role": "user", "content": "Hi"}] +async def test_reasoning_signature_round_trips_from_provider_metadata( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _, captured = _patch_client(monkeypatch) + + await _drain( + adapter.stream( + _TEST_CLIENT, + _MODEL, + [ + ai.assistant_message( + ai.thinking( + "hidden", + provider_metadata=anthropic_metadata.AnthropicProviderMetadata( + signature="sig" + ), + ) + ), + ai.user_message("Hi"), + ], + ) + ) + + assert captured["messages"][0] == { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "hidden", + "signature": "sig", + } + ], + } + + async def test_builtin_tool_parts_round_trip( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -183,14 +226,15 @@ async def test_builtin_tool_parts_round_trip( tool_call_id="srvtoolu_1", tool_name="web_search", tool_args='{"query":"weather"}', - provider_name="anthropic", + provider_metadata=anthropic_metadata.AnthropicProviderMetadata(), ) result = messages.BuiltinToolReturnPart( tool_call_id="srvtoolu_1", tool_name="web_search", result=[{"title": "Forecast", "url": "https://example.com"}], - provider_name="anthropic", - provider_details={"result_type": "web_search_tool_result"}, + provider_metadata=anthropic_metadata.AnthropicProviderMetadata( + result_type="web_search_tool_result" + ), ) convo = [ ai.user_message("What's the weather?"), diff --git a/tests/models/anthropic/test_stream.py b/tests/models/anthropic/test_stream.py index 0fb51f68..38d3c415 100644 --- a/tests/models/anthropic/test_stream.py +++ b/tests/models/anthropic/test_stream.py @@ -14,6 +14,7 @@ import ai from ai import models from ai.models.anthropic import adapter, anthropic +from ai.models.anthropic import metadata as anthropic_metadata from ai.types import events, messages from .conftest import ( @@ -54,7 +55,10 @@ async def test_server_tool_use_emits_builtin_events( assert calls[0].tool_call_id == "srvtoolu_1" assert calls[0].tool_name == "web_search" assert calls[0].tool_args == '{"query":"weather"}' - assert calls[0].provider_name == "anthropic" + assert isinstance( + calls[0].provider_metadata, + anthropic_metadata.AnthropicProviderMetadata, + ) async def test_tool_result_block_emits_builtin_result( @@ -88,7 +92,31 @@ async def test_tool_result_block_emits_builtin_result( assert ret.tool_call_id == "srvtoolu_1" assert ret.tool_name == "web_search" assert ret.result == payload - assert ret.provider_details == {"result_type": "web_search_tool_result"} + assert isinstance( + ret.provider_metadata, + anthropic_metadata.AnthropicProviderMetadata, + ) + assert ret.provider_metadata.result_type == "web_search_tool_result" + + +async def test_signature_delta_emits_anthropic_metadata( + monkeypatch: pytest.MonkeyPatch, +) -> None: + sdk_events = [ + block_start(0, "thinking"), + block_delta(0, "thinking_delta", thinking="hidden"), + block_delta(0, "signature_delta", signature="sig"), + block_stop(0), + ] + s = await _drain(FakeStream(sdk_events), monkeypatch) + + reasoning = s.message.parts[0] + assert isinstance(reasoning, messages.ReasoningPart) + assert isinstance( + reasoning.provider_metadata, + anthropic_metadata.AnthropicProviderMetadata, + ) + assert reasoning.provider_metadata.signature == "sig" async def test_event_kinds_in_order(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/tests/models/core/test_api.py b/tests/models/core/test_api.py index 5a45c969..5dcd983d 100644 --- a/tests/models/core/test_api.py +++ b/tests/models/core/test_api.py @@ -10,8 +10,8 @@ from ai import models from ai.models.openai import openai from ai.types import events as events_ -from ai.types import metadata as metadata_ from ai.types import messages as messages_ +from ai.types import metadata as metadata_ from ...conftest import MOCK_MODEL, MOCK_PROVIDER, MockProvider, mock_llm, text_msg @@ -108,7 +108,6 @@ async def _metadata_stream() -> AsyncGenerator[events_.Event]: ) yield events_.ReasoningEnd( block_id="reasoning", - signature="sig-1", provider_metadata=_TestProviderMetadata(marker="reasoning-end"), ) yield events_.ToolStart( @@ -151,7 +150,6 @@ async def _metadata_stream() -> AsyncGenerator[events_.Event]: reasoning = stream.message.parts[1] assert isinstance(reasoning, messages_.ReasoningPart) assert reasoning.text == "thinking" - assert reasoning.signature == "sig-1" assert _provider_metadata_marker(reasoning.provider_metadata) == "reasoning-end" tool_call = stream.message.parts[2] diff --git a/tests/models/openai/test_adapter.py b/tests/models/openai/test_adapter.py index 7bff8115..bc74249c 100644 --- a/tests/models/openai/test_adapter.py +++ b/tests/models/openai/test_adapter.py @@ -195,7 +195,6 @@ async def test_builtin_part_in_messages_raises( messages.BuiltinToolCallPart( tool_call_id="srvtoolu_1", tool_name="web_search", - provider_name="openai", ), ], ), From 95c72df60b2894b6ef1dc8a26fe01e1d0eb23a20 Mon Sep 17 00:00:00 2001 From: Andrey Buzin Date: Fri, 8 May 2026 13:52:43 -0700 Subject: [PATCH 8/8] Fix broken import in usage test --- tests/types/test_messages.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/types/test_messages.py b/tests/types/test_messages.py index dee0ee55..e9892eb6 100644 --- a/tests/types/test_messages.py +++ b/tests/types/test_messages.py @@ -5,7 +5,7 @@ import pydantic import pytest -from ai.types import messages +from ai.types import messages, usage class _Weather(pydantic.BaseModel): @@ -67,12 +67,12 @@ def test_structured_output_round_trip() -> None: def test_usage_add_merges_optional_fields() -> None: - a = messages.Usage( + a = usage.Usage( input_tokens=100, output_tokens=50, cache_read_tokens=20, ) - b = messages.Usage( + b = usage.Usage( input_tokens=200, output_tokens=80, reasoning_tokens=10,