From 7d882e120ce68be5e3e62fe1497a9b6be91ef24d Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Mon, 8 Jun 2026 18:30:21 -0700 Subject: [PATCH] Make `MessageBundle` a special tool result type Generalize the mechanism used currently for ContentOuput for MessageBundle as well. This gives it an approximately principled way for it to be fed into the outbound ai-sdk ui streaming (SpecialToolResults are allowed to stay as pydantic models) and a *more* principled hook for the ui protocol to preserve the shape. --- .../docs/core-framework/streaming-tools.mdx | 5 +- src/ai/agents/__init__.py | 2 +- src/ai/agents/agent.py | 20 +-- src/ai/agents/ui/ai_sdk/id_utils.py | 13 +- src/ai/agents/ui/ai_sdk/inbound_messages.py | 54 +++----- src/ai/agents/ui/ai_sdk/outbound_messages.py | 23 ++++ src/ai/agents/ui/ai_sdk/outbound_stream.py | 19 +-- src/ai/types/builders.py | 15 +-- src/ai/types/messages.py | 54 ++++++-- tests/agents/test_generator_tools.py | 2 +- .../agents/ui/ai_sdk/test_inbound_messages.py | 126 ++++++++++++++---- .../agents/ui/ai_sdk/test_outbound_stream.py | 14 +- tests/providers/ai_gateway/test_protocol.py | 2 +- .../anthropic/test_multipart_tool_result.py | 2 +- .../openai/test_multipart_tool_result.py | 2 +- tests/types/test_messages.py | 6 +- 16 files changed, 231 insertions(+), 128 deletions(-) diff --git a/docs/ai-python/content/docs/core-framework/streaming-tools.mdx b/docs/ai-python/content/docs/core-framework/streaming-tools.mdx index b418ceb4..72634a84 100644 --- a/docs/ai-python/content/docs/core-framework/streaming-tools.mdx +++ b/docs/ai-python/content/docs/core-framework/streaming-tools.mdx @@ -70,12 +70,13 @@ messages while the model input is the final assistant text. ## Rich snapshots Aggregators can preserve more than text. `MessageAggregator` stores nested -messages from a subagent: +messages from a subagent. The rich snapshot type is +`ai.types.messages.MessageBundle`: ```python if isinstance(event, ai.events.ToolCallResult): result = event.results[0].result - if isinstance(result, ai.agents.MessageBundle): + if isinstance(result, ai.types.messages.MessageBundle): print(result.messages[-1].text) ``` diff --git a/src/ai/agents/__init__.py b/src/ai/agents/__init__.py index 3d785520..9a4fbd6d 100644 --- a/src/ai/agents/__init__.py +++ b/src/ai/agents/__init__.py @@ -1,3 +1,4 @@ +from ..types.messages import MessageBundle from . import mcp, ui from .agent import ( Agent, @@ -9,7 +10,6 @@ GatedToolCall, LastAggregator, MessageAggregator, - MessageBundle, SimpleAggregator, StreamingStatusTool, StreamingTextTool, diff --git a/src/ai/agents/agent.py b/src/ai/agents/agent.py index cb136d53..f6ecc064 100644 --- a/src/ai/agents/agent.py +++ b/src/ai/agents/agent.py @@ -22,7 +22,6 @@ Any, ClassVar, Generic, - Literal, Protocol, Self, cast, @@ -39,6 +38,7 @@ from .. import models, types, util from ..types import builders from ..types import events as events_ +from ..types.messages import MessageBundle from . import _middleware as middleware_ from . import hooks as hooks_ from . import runtime @@ -59,18 +59,6 @@ def _unwrap_singleton_group(exc: BaseException) -> BaseException: return exc -def _result_kind(value: Any) -> Literal["json", "content"]: - """Tag a successful tool return value for ``ToolResultPart.result_kind``. - - A :class:`ContentOutput` becomes ``"content"`` (expanded into provider - multimodal blocks); everything else is ``"json"`` (the encoder sends a - ``str`` raw and JSON-encodes anything else). - """ - if isinstance(value, types.messages.ContentOutput): - return "content" - return "json" - - def _error_tool_result( exc: BaseException, *, @@ -232,10 +220,6 @@ def snapshot(self) -> T | None: return self._val -class MessageBundle(pydantic.BaseModel): - messages: tuple[types.messages.Message, ...] - - class MessageAggregator( events_.Aggregator[events_.AgentEvent, MessageBundle, str] ): @@ -629,7 +613,7 @@ async def _real( tool_call_id=call.tool_call_id, tool_name=call.tool_name, result=result, - result_kind=_result_kind(result), + result_kind=types.messages.ToolResultPart.kind_for(result), ) part.set_model_input(model_input) return tool_result(part) diff --git a/src/ai/agents/ui/ai_sdk/id_utils.py b/src/ai/agents/ui/ai_sdk/id_utils.py index 484d8a86..be123317 100644 --- a/src/ai/agents/ui/ai_sdk/id_utils.py +++ b/src/ai/agents/ui/ai_sdk/id_utils.py @@ -87,15 +87,18 @@ def _restore_message_ids( def _tool_result_kinds( source_messages: list[messages_.Message], ) -> dict[str, str]: - """Collect ``{tool_call_id: result_kind}`` for content tool results.""" + """Collect ``{tool_call_id: subtype}`` for special tool results. + + The recorded value is the :class:`SpecialToolResult` discriminator so the + inbound side can rehydrate the typed result without shape-sniffing it. + """ kinds: dict[str, str] = {} for message in source_messages: for part in message.parts: - if ( - isinstance(part, messages_.ToolResultPart) - and part.result_kind == "content" + if isinstance(part, messages_.ToolResultPart) and isinstance( + part.result, messages_.SpecialToolResult ): - kinds[part.tool_call_id] = part.result_kind + kinds[part.tool_call_id] = part.result.type return kinds diff --git a/src/ai/agents/ui/ai_sdk/inbound_messages.py b/src/ai/agents/ui/ai_sdk/inbound_messages.py index 373e8919..74e312bf 100644 --- a/src/ai/agents/ui/ai_sdk/inbound_messages.py +++ b/src/ai/agents/ui/ai_sdk/inbound_messages.py @@ -11,7 +11,7 @@ from typing import Any from ....types import messages as messages_ -from ...agent import MessageBundle +from ....types.messages import MessageBundle from . import approvals, id_utils from . import ui_messages as ui_messages_ from .approvals import ApprovalResponse, extract_approvals @@ -57,27 +57,6 @@ def _error_result(error_text: str | None, output: Any) -> dict[str, Any] | None: return normalized -def _decode_wire_output(output: Any) -> Any: - """Reconstruct the internal snapshot type from a wire tool output. - - Hacky special case: when the wire output looks like a ``UIMessage`` - (the wire shape we emit for sub-agent / ``MessageAggregator`` tools), - decode it back to a ``MessageBundle``. Other shapes pass through - unchanged. This avoids requiring callers to thread the tool - registry into inbound parsing. - """ - if not isinstance(output, dict): - return output - if output.get("role") != "assistant" or "parts" not in output: - return output - try: - ui_msg = ui_messages_.UIMessage.model_validate(output) - except Exception: - return output - inner = list(_parse([ui_msg])) - return MessageBundle(messages=tuple(inner)) - - def _build_result_part( *, tool_call_id: str, @@ -89,10 +68,14 @@ def _build_result_part( """Reconstruct a tool result from its wire form. ``kind_hint`` comes from the adapter's ``toolResultKinds`` metadata - (see :mod:`id_utils`). When it marks the result as ``content``, the - ``output`` -- a list of dumped content parts -- is rehydrated into a - typed :class:`ContentOutput` so providers re-expand it into multimodal - blocks; otherwise behaviour matches a plain value round-trip. + (see :mod:`id_utils`) and names the :class:`SpecialToolResult` subtype: + + * ``"content"`` rehydrates a :class:`ContentOutput` from the dumped + content parts so providers re-expand it into multimodal blocks; + * ``"messages"`` rebuilds a :class:`MessageBundle` by parsing the + carried sub-agent UIMessage(s). + + Without a hint the output is treated as a plain value round-trip. """ result: Any result_kind: messages_.ResultKind @@ -101,14 +84,19 @@ def _build_result_part( result_kind = "error" elif kind_hint == "content": result = messages_.ContentOutput.model_validate({"value": output}) - result_kind = "content" + result_kind = "special" + elif kind_hint == "messages": + raw = output if isinstance(output, list) else [output] + ui_msgs = [ + m + if isinstance(m, ui_messages_.UIMessage) + else ui_messages_.UIMessage.model_validate(m) + for m in raw + ] + result = MessageBundle(messages=tuple(_parse(ui_msgs))) + result_kind = "special" else: - decoded = _decode_wire_output(output) - result = ( - decoded - if isinstance(decoded, MessageBundle) - else _normalize_tool_result(decoded) - ) + result = _normalize_tool_result(output) result_kind = "json" return messages_.ToolResultPart( tool_call_id=tool_call_id, diff --git a/src/ai/agents/ui/ai_sdk/outbound_messages.py b/src/ai/agents/ui/ai_sdk/outbound_messages.py index 98366cd5..48507bfd 100644 --- a/src/ai/agents/ui/ai_sdk/outbound_messages.py +++ b/src/ai/agents/ui/ai_sdk/outbound_messages.py @@ -108,6 +108,23 @@ def dedupe_tool_parts( return result +def bundle_to_wire_output(bundle: messages_.MessageBundle) -> Any: + """Serialize a sub-agent transcript to its UI tool ``output``. + + Follows the AI SDK sub-agent convention of a single ``UIMessage`` for the + common case (one bubble), and only falls back to a JSON list when the + transcript spans multiple bubbles. Returns ``None`` for an empty bundle + so streaming callers can skip emitting until there's something to show. + The inbound side accepts either shape (see ``_build_result_part``). + """ + dumped = [ + m.model_dump(mode="json") for m in to_ui_messages(list(bundle.messages)) + ] + if not dumped: + return None + return dumped[0] if len(dumped) == 1 else dumped + + def _output_view( part: messages_.ToolResultPart, ) -> tuple[str, dict[str, Any]]: @@ -117,6 +134,12 @@ def _output_view( return "output-available", { "output": [item.model_dump(mode="json") for item in result.value] } + if isinstance(result, messages_.MessageBundle): + # `None` (empty bundle) becomes `[]` so a completed result still + # round-trips to an (empty) MessageBundle rather than a null output. + return "output-available", { + "output": bundle_to_wire_output(result) or [] + } if part.is_error: text = result if isinstance(result, str) else json.dumps(result) return "output-error", {"error_text": text} diff --git a/src/ai/agents/ui/ai_sdk/outbound_stream.py b/src/ai/agents/ui/ai_sdk/outbound_stream.py index 215174e4..e2183681 100644 --- a/src/ai/agents/ui/ai_sdk/outbound_stream.py +++ b/src/ai/agents/ui/ai_sdk/outbound_stream.py @@ -11,7 +11,7 @@ from ....types import events as events_ from ....types import media from ....types import messages as messages_ -from ...agent import MessageBundle +from ....types.messages import MessageBundle from . import approvals, outbound_messages, ui_events from .tool_utils import normalize_tool_input @@ -35,17 +35,18 @@ def _tool_error_text(part: messages_.ToolResultPart) -> str: def _to_wire_output(snapshot: Any) -> Any: """Convert an aggregator snapshot to its UI wire representation. - For ``MessageBundle`` (sub-agent transcripts) this produces a single - ``UIMessage`` assistant bubble — the canonical AI SDK shape. Other - snapshot types pass through unchanged. + For ``MessageBundle`` (sub-agent transcripts) this follows the AI SDK + sub-agent convention -- a single ``UIMessage`` for the common one-bubble + case, a JSON list only when the transcript spans multiple bubbles -- and + is paired with a ``toolResultKinds`` ``"messages"`` hint so the inbound + side can rebuild the bundle. Other snapshot types pass through unchanged. - Returns ``None`` if the bundle has no assistant anchor yet (e.g. a - streaming sub-agent that has produced no messages); callers should - skip emitting in that case. + Returns ``None`` if the bundle has no messages yet (e.g. a streaming + sub-agent that has produced nothing); callers should skip emitting in + that case. """ if isinstance(snapshot, MessageBundle): - ui_msgs = outbound_messages.to_ui_messages(list(snapshot.messages)) - return ui_msgs[-1] if ui_msgs else None + return outbound_messages.bundle_to_wire_output(snapshot) return snapshot diff --git a/src/ai/types/builders.py b/src/ai/types/builders.py index 99eb0896..9a3a69cf 100644 --- a/src/ai/types/builders.py +++ b/src/ai/types/builders.py @@ -241,18 +241,15 @@ def tool_result_part( """Create a :class:`ToolResultPart`. ``result`` is stored as-is; ``result_kind`` is derived: ``"error"`` when - ``is_error`` is set, ``"content"`` for a :class:`ContentOutput`, else - ``"json"`` (a ``str`` is sent raw to the model, anything else is - JSON-encoded at the provider boundary). + ``is_error`` is set, ``"special"`` for a :class:`ContentOutput` or + :class:`MessageBundle`, else ``"json"`` (a ``str`` is sent raw to the + model, anything else is JSON-encoded at the provider boundary). >>> ai.tool_result_part("tc-1", result={"temp": 72}, tool_name="weather") """ - if is_error: - result_kind: ResultKind = "error" - elif isinstance(result, ContentOutput): - result_kind = "content" - else: - result_kind = "json" + result_kind: ResultKind = ( + "error" if is_error else ToolResultPart.kind_for(result) + ) return ToolResultPart( tool_call_id=tool_call_id, tool_name=tool_name, diff --git a/src/ai/types/messages.py b/src/ai/types/messages.py index 6c61b1c3..644e9380 100644 --- a/src/ai/types/messages.py +++ b/src/ai/types/messages.py @@ -102,7 +102,7 @@ def from_bytes( # --------------------------------------------------------------------------- # Multipart tool result -- a tool may return a mix of text and file/image # parts so the model sees actual media. Stored on ``ToolResultPart.result`` -# with ``result_kind="content"``; providers expand it into their multimodal +# with ``result_kind="special"``; providers expand it into their multimodal # wire format. # --------------------------------------------------------------------------- @@ -122,13 +122,31 @@ class ContentOutput(pydantic.BaseModel): model_config = pydantic.ConfigDict(frozen=True) +class MessageBundle(pydantic.BaseModel): + type: Literal["messages"] = "messages" + messages: tuple["Message", ...] + + +SpecialToolResult = ContentOutput | MessageBundle + +_SPECIAL_TOOL_RESULT_ADAPTER: pydantic.TypeAdapter[SpecialToolResult] = ( + pydantic.TypeAdapter( + Annotated[ + SpecialToolResult, + pydantic.Field(discriminator="type"), + ] + ) +) + + _MODEL_INPUT_UNSET: Any = object() -# Coarse tag for the shape of ``ToolResultPart.result``. ``"content"`` means -# a :class:`ContentOutput`; ``"error"`` flags an error result; ``"json"`` (the -# default) is any plain value. Providers decide text-vs-json at the wire -# boundary (a ``str`` is sent raw, everything else is JSON-encoded). -ResultKind = Literal["error", "json", "content"] +# Coarse tag for the shape of ``ToolResultPart.result``. +# ``"special"`` means a :class:`SpecialToolResult`; ``"error"`` flags +# an error result; ``"json"`` (the default) is any plain value. +# Providers decide text-vs-json at the wire boundary (a ``str`` is +# sent raw, everything else is JSON-encoded). +ResultKind = Literal["error", "json", "special"] class ToolResultPart(pydantic.BaseModel): @@ -162,24 +180,36 @@ class ToolResultPart(pydantic.BaseModel): @pydantic.model_validator(mode="before") @classmethod def _restore_content(cls, data: Any) -> Any: - """Rebuild a typed :class:`ContentOutput` after a JSON round-trip. + """Rebuild a typed :class:`SpecialToolResult` after a JSON round-trip. ``result`` is ``Any``, so pydantic restores a serialized - ``ContentOutput`` as a plain dict. When ``result_kind`` says the - result is content, coerce it back so providers (and the UI adapter) - can rely on ``isinstance(result, ContentOutput)``. + ``ContentOutput`` / ``MessageBundle`` as a plain dict. When + ``result_kind`` is ``"special"``, coerce it back to the typed result + so providers (and the UI adapter) can rely on ``isinstance`` checks. """ if ( isinstance(data, dict) - and data.get("result_kind") == "content" + and data.get("result_kind") == "special" and isinstance(data.get("result"), dict) ): data = { **data, - "result": ContentOutput.model_validate(data["result"]), + "result": _SPECIAL_TOOL_RESULT_ADAPTER.validate_python( + data["result"] + ), } return data + @staticmethod + def kind_for(result: Any) -> ResultKind: + """Derive ``result_kind`` for a non-error result value. + + A :data:`SpecialToolResult` is ``"special"``; anything else is + ``"json"``. Error results are tagged ``"error"`` by the + caller, independent of the value. + """ + return "special" if isinstance(result, SpecialToolResult) else "json" + @property def is_error(self) -> bool: """Whether this result represents an error to the model.""" diff --git a/tests/agents/test_generator_tools.py b/tests/agents/test_generator_tools.py index 64136208..f83f49e0 100644 --- a/tests/agents/test_generator_tools.py +++ b/tests/agents/test_generator_tools.py @@ -9,10 +9,10 @@ import ai from ai import models -from ai.agents.agent import MessageBundle from ai.types import events as agent_events_ from ai.types import events as events_ from ai.types import messages as messages_ +from ai.types.messages import MessageBundle from ..conftest import ( MOCK_MODEL, diff --git a/tests/agents/ui/ai_sdk/test_inbound_messages.py b/tests/agents/ui/ai_sdk/test_inbound_messages.py index a3342e05..463ed263 100644 --- a/tests/agents/ui/ai_sdk/test_inbound_messages.py +++ b/tests/agents/ui/ai_sdk/test_inbound_messages.py @@ -4,11 +4,11 @@ import pytest -from ai.agents.agent import MessageBundle from ai.agents.ui.ai_sdk import to_messages, to_ui_messages from ai.agents.ui.ai_sdk.inbound_messages import _normalize_ui_messages from ai.agents.ui.ai_sdk.ui_messages import UIMessage, UIToolPart from ai.types import messages as messages_ +from ai.types.messages import MessageBundle def _ui(role: str, *parts: dict[str, Any], id: str = "m1") -> UIMessage: @@ -152,43 +152,117 @@ def test_to_messages_rejects_empty_user() -> None: to_messages(ui) -def test_to_messages_decodes_subagent_tool_output() -> None: - """A sub-agent tool's wire UIMessage decodes back to MessageBundle. +def test_subagent_bundle_round_trips_via_metadata() -> None: + """A sub-agent MessageBundle survives outbound -> inbound. - ``result`` carries the rich MessageBundle so a subsequent UI render - gets the same shape we sent. ``model_input`` is left unset here — - populating it requires the tool registry, which lives in + The transcript rides the UI tool part's ``output`` as a list of + UIMessages; the ``toolResultKinds`` adapter metadata carries the + ``"messages"`` signal to rebuild the bundle. ``model_input`` is left + unset here -- populating it requires the tool registry, which lives in :meth:`Agent.run`. """ - # Wire shape: tool-_research_tool with output = UIMessage{parts=[text]}. - ui = [ - _ui("user", _text("research mars"), id="u1"), - _ui( - "assistant", - _tool( - "_research_tool", - "tc1", - "output-available", - input={"topic": "mars"}, - output={ - "id": "sub-1", - "role": "assistant", - "parts": [{"type": "text", "text": "Mars has two moons."}], - }, + turn = "turn-1" + bundle = MessageBundle( + messages=( + messages_.Message( + role="assistant", + parts=[messages_.TextPart(text="Mars has two moons.")], ), + ) + ) + internal = [ + messages_.Message( id="a1", + turn_id=turn, + role="assistant", + parts=[ + messages_.ToolCallPart( + tool_call_id="tc1", + tool_name="research", + tool_args="{}", + ) + ], + ), + messages_.Message( + id="t1", + turn_id=turn, + role="tool", + parts=[ + messages_.ToolResultPart( + tool_call_id="tc1", + tool_name="research", + result=bundle, + result_kind="special", + ) + ], ), ] - messages, _ = to_messages(ui) - # Find the tool message with the decoded result. - tool_msgs = [m for m in messages if m.role == "tool"] + ui = to_ui_messages(internal) + restored, _ = to_messages(ui) + + tool_msgs = [m for m in restored if m.role == "tool"] assert len(tool_msgs) == 1 result_part = tool_msgs[0].tool_results[0] + assert result_part.result_kind == "special" assert isinstance(result_part.result, MessageBundle) + assert len(result_part.result.messages) == 1 + inner = result_part.result.messages[0] + assert inner.role == "assistant" + assert inner.text == "Mars has two moons." assert not result_part.has_model_input +def test_multi_bubble_subagent_bundle_round_trips_as_list() -> None: + """A transcript spanning multiple bubbles round-trips via a list output.""" + turn = "turn-1" + bundle = MessageBundle( + messages=( + messages_.Message( + role="assistant", + turn_id="s1", + parts=[messages_.TextPart(text="first")], + ), + messages_.Message( + role="assistant", + turn_id="s2", + parts=[messages_.TextPart(text="second")], + ), + ) + ) + internal = [ + messages_.Message( + id="a1", + turn_id=turn, + role="assistant", + parts=[ + messages_.ToolCallPart( + tool_call_id="tc1", tool_name="research", tool_args="{}" + ) + ], + ), + messages_.Message( + id="t1", + turn_id=turn, + role="tool", + parts=[ + messages_.ToolResultPart( + tool_call_id="tc1", + tool_name="research", + result=bundle, + result_kind="special", + ) + ], + ), + ] + + restored, _ = to_messages(to_ui_messages(internal)) + + result_part = next(m for m in restored if m.role == "tool").tool_results[0] + assert isinstance(result_part.result, MessageBundle) + assert [m.text for m in result_part.result.messages] == ["first", "second"] + + def test_content_result_round_trips_via_metadata() -> None: """A content tool result survives outbound -> inbound as ContentOutput. @@ -221,7 +295,7 @@ def test_content_result_round_trips_via_metadata() -> None: result=messages_.ContentOutput( value=[messages_.TextPart(text="desc"), fp] ), - result_kind="content", + result_kind="special", ) ], ), @@ -233,7 +307,7 @@ def test_content_result_round_trips_via_metadata() -> None: tool_msgs = [m for m in restored if m.role == "tool"] assert len(tool_msgs) == 1 part = tool_msgs[0].tool_results[0] - assert part.result_kind == "content" + assert part.result_kind == "special" assert isinstance(part.result, messages_.ContentOutput) text_part, file_part = part.result.value assert isinstance(text_part, messages_.TextPart) diff --git a/tests/agents/ui/ai_sdk/test_outbound_stream.py b/tests/agents/ui/ai_sdk/test_outbound_stream.py index e496405c..3907e006 100644 --- a/tests/agents/ui/ai_sdk/test_outbound_stream.py +++ b/tests/agents/ui/ai_sdk/test_outbound_stream.py @@ -413,10 +413,12 @@ async def test_partial_tool_results_emit_preliminary_outputs() -> None: assert all(p.tool_call_id == "tc1" for p in prelim) -async def test_partial_message_bundle_becomes_ui_message() -> None: - """MessageAggregator's snapshot collapses to one UIMessage.""" - from ai.agents.ui.ai_sdk.ui_messages import UIMessage +async def test_partial_message_bundle_becomes_single_ui_message() -> None: + """A one-bubble MessageAggregator snapshot serializes to one UIMessage. + Matches the AI SDK sub-agent convention: a single ``UIMessage`` (not a + one-element list) for the common case. + """ inner_msg = messages_.Message( role="assistant", parts=[messages_.TextPart(text="hi from sub-agent")], @@ -440,9 +442,9 @@ async def test_partial_message_bundle_becomes_ui_message() -> None: for p in out if isinstance(p, ui_events.UIToolOutputAvailableEvent) and p.preliminary ] - assert isinstance(prelim.output, UIMessage) - assert prelim.output.role == "assistant" - assert prelim.output.parts[0].type == "text" + assert isinstance(prelim.output, dict) + assert prelim.output["role"] == "assistant" + assert prelim.output["parts"][0]["type"] == "text" async def test_partial_tool_result_without_factory_is_skipped() -> None: diff --git a/tests/providers/ai_gateway/test_protocol.py b/tests/providers/ai_gateway/test_protocol.py index 8cfccd4c..3422dc20 100644 --- a/tests/providers/ai_gateway/test_protocol.py +++ b/tests/providers/ai_gateway/test_protocol.py @@ -545,7 +545,7 @@ def test_content_multipart(self) -> None: messages.ContentOutput( value=[messages.TextPart(text="desc"), fp] ), - result_kind="content", + result_kind="special", ) ) assert result["type"] == "content" diff --git a/tests/providers/anthropic/test_multipart_tool_result.py b/tests/providers/anthropic/test_multipart_tool_result.py index e68fd455..d7cfa51e 100644 --- a/tests/providers/anthropic/test_multipart_tool_result.py +++ b/tests/providers/anthropic/test_multipart_tool_result.py @@ -94,7 +94,7 @@ async def test_tool_result_with_file_part(self) -> None: fp, ] ), - result_kind="content", + result_kind="special", ) ], ), diff --git a/tests/providers/openai/test_multipart_tool_result.py b/tests/providers/openai/test_multipart_tool_result.py index 40f6687d..d40cb074 100644 --- a/tests/providers/openai/test_multipart_tool_result.py +++ b/tests/providers/openai/test_multipart_tool_result.py @@ -98,7 +98,7 @@ async def test_tool_result_with_file_part(self) -> None: fp, ] ), - result_kind="content", + result_kind="special", ) ], ), diff --git a/tests/types/test_messages.py b/tests/types/test_messages.py index 03806b69..08dee497 100644 --- a/tests/types/test_messages.py +++ b/tests/types/test_messages.py @@ -130,7 +130,7 @@ def test_tool_result_content_output_with_file_part_round_trip() -> None: result=messages.ContentOutput( value=[messages.TextPart(text="label"), fp] ), - result_kind="content", + result_kind="special", ) j = trp.model_dump_json() restored = messages.ToolResultPart.model_validate_json(j) @@ -175,7 +175,7 @@ def test_tool_result_content_in_message_round_trip() -> None: result=messages.ContentOutput( value=[messages.TextPart(text="Read image"), fp] ), - result_kind="content", + result_kind="special", ) ], ) @@ -203,7 +203,7 @@ def test_tool_result_file_part_base64_valid_after_round_trip() -> None: result=messages.ContentOutput( value=[messages.TextPart(text="label"), fp] ), - result_kind="content", + result_kind="special", ) restored = messages.ToolResultPart.model_validate_json( trp.model_dump_json()