simonw--llm
Simon Willison
a2547d8183
ReasoningPart.token_count duplicated info already on response.token_details (reasoning_tokens), and the side-channel `response._reasoning_token_count` attribute with its set_usage ordering footgun was the wrong shape. Replaced with a clean StreamEvent.redacted=True marker that plugins yield like any other event. The framework hoists redacted reasoning Parts to the start of the assembled message so UIs render them before content, even though the opaque count typically arrives at the end of the stream. Also fix parallel tool calls emitted without tool_call_id (e.g. Gemini): a fresh tool_call_name now always allocates a new index instead of falling through to the prior tool-call group, so N parallel calls produce N distinct ToolCallParts instead of one with concatenated names and args. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
326 行
11 KiB
Python
326 行
11 KiB
Python
"""Tests for llm.serialization — the TypedDict spec for the JSON-safe
|
|
wire form of Message, Part, and Response.
|
|
|
|
Uses pydantic.TypeAdapter to verify that actual to_dict() output
|
|
conforms to the TypedDict annotations. pydantic is already a runtime
|
|
dependency.
|
|
"""
|
|
|
|
import json
|
|
import pytest
|
|
from pydantic import TypeAdapter
|
|
|
|
import llm
|
|
from llm.serialization import (
|
|
AttachmentPartDict,
|
|
MessageDict,
|
|
PartDict,
|
|
ResponseDict,
|
|
ReasoningPartDict,
|
|
TextPartDict,
|
|
ToolCallPartDict,
|
|
ToolResultPartDict,
|
|
)
|
|
|
|
# ---- required/optional keys ----------------------------------------
|
|
|
|
|
|
class TestRequiredOptionalKeys:
|
|
def test_message_dict_required_keys(self):
|
|
assert MessageDict.__required_keys__ == {"role", "parts"}
|
|
assert MessageDict.__optional_keys__ == {"provider_metadata"}
|
|
|
|
def test_text_part_dict_required_keys(self):
|
|
assert TextPartDict.__required_keys__ == {"type", "text"}
|
|
assert TextPartDict.__optional_keys__ == {"provider_metadata"}
|
|
|
|
def test_reasoning_part_dict_required_keys(self):
|
|
assert ReasoningPartDict.__required_keys__ == {"type", "text"}
|
|
assert ReasoningPartDict.__optional_keys__ == {
|
|
"redacted",
|
|
"provider_metadata",
|
|
}
|
|
|
|
def test_tool_call_part_dict_required_keys(self):
|
|
assert ToolCallPartDict.__required_keys__ == {"type", "name", "arguments"}
|
|
assert ToolCallPartDict.__optional_keys__ == {
|
|
"tool_call_id",
|
|
"server_executed",
|
|
"provider_metadata",
|
|
}
|
|
|
|
def test_tool_result_part_dict_required_keys(self):
|
|
assert ToolResultPartDict.__required_keys__ == {"type", "name", "output"}
|
|
assert ToolResultPartDict.__optional_keys__ == {
|
|
"tool_call_id",
|
|
"server_executed",
|
|
"exception",
|
|
"attachments",
|
|
"provider_metadata",
|
|
}
|
|
|
|
def test_attachment_part_dict_required_keys(self):
|
|
assert AttachmentPartDict.__required_keys__ == {"type"}
|
|
assert AttachmentPartDict.__optional_keys__ == {
|
|
"attachment",
|
|
"provider_metadata",
|
|
}
|
|
|
|
def test_response_dict_required_keys(self):
|
|
assert ResponseDict.__required_keys__ == {"model", "prompt", "messages"}
|
|
assert ResponseDict.__optional_keys__ == {"id", "usage", "datetime_utc"}
|
|
|
|
|
|
# ---- to_dict output conforms to the TypedDict ----------------------
|
|
|
|
|
|
class TestPartRoundTrip:
|
|
def _adapter(self, td):
|
|
return TypeAdapter(td)
|
|
|
|
def test_text_part_matches(self):
|
|
d = llm.TextPart(text="hello").to_dict()
|
|
self._adapter(TextPartDict).validate_python(d)
|
|
|
|
def test_text_part_with_provider_metadata_matches(self):
|
|
d = llm.TextPart(
|
|
text="hi", provider_metadata={"anthropic": {"cached": True}}
|
|
).to_dict()
|
|
self._adapter(TextPartDict).validate_python(d)
|
|
|
|
def test_reasoning_part_redacted_matches(self):
|
|
d = llm.ReasoningPart(text="", redacted=True).to_dict()
|
|
self._adapter(ReasoningPartDict).validate_python(d)
|
|
|
|
def test_reasoning_part_with_signature_matches(self):
|
|
d = llm.ReasoningPart(
|
|
text="thinking...",
|
|
provider_metadata={"anthropic": {"signature": "sig-abc"}},
|
|
).to_dict()
|
|
self._adapter(ReasoningPartDict).validate_python(d)
|
|
|
|
def test_tool_call_part_matches(self):
|
|
d = llm.ToolCallPart(
|
|
name="search", arguments={"q": "x"}, tool_call_id="c1"
|
|
).to_dict()
|
|
self._adapter(ToolCallPartDict).validate_python(d)
|
|
|
|
def test_tool_result_part_matches(self):
|
|
d = llm.ToolResultPart(
|
|
name="search", output="result", tool_call_id="c1"
|
|
).to_dict()
|
|
self._adapter(ToolResultPartDict).validate_python(d)
|
|
|
|
def test_attachment_part_with_url_matches(self):
|
|
att = llm.Attachment(type="image/jpeg", url="https://example.com/cat.jpg")
|
|
d = llm.AttachmentPart(attachment=att).to_dict()
|
|
self._adapter(AttachmentPartDict).validate_python(d)
|
|
|
|
def test_attachment_part_with_bytes_matches(self):
|
|
att = llm.Attachment(type="image/png", content=b"\x89PNG...")
|
|
d = llm.AttachmentPart(attachment=att).to_dict()
|
|
self._adapter(AttachmentPartDict).validate_python(d)
|
|
|
|
|
|
class TestPartDiscriminatedUnion:
|
|
def test_text_part_validates_as_part_dict(self):
|
|
d = llm.TextPart(text="hi").to_dict()
|
|
TypeAdapter(PartDict).validate_python(d)
|
|
|
|
def test_reasoning_part_validates_as_part_dict(self):
|
|
d = llm.ReasoningPart(text="thinking").to_dict()
|
|
TypeAdapter(PartDict).validate_python(d)
|
|
|
|
def test_tool_call_part_validates_as_part_dict(self):
|
|
d = llm.ToolCallPart(name="t", arguments={}, tool_call_id="c1").to_dict()
|
|
TypeAdapter(PartDict).validate_python(d)
|
|
|
|
def test_tool_result_part_validates_as_part_dict(self):
|
|
d = llm.ToolResultPart(name="t", output="out", tool_call_id="c1").to_dict()
|
|
TypeAdapter(PartDict).validate_python(d)
|
|
|
|
def test_attachment_part_validates_as_part_dict(self):
|
|
att = llm.Attachment(type="image/jpeg", url="http://x")
|
|
d = llm.AttachmentPart(attachment=att).to_dict()
|
|
TypeAdapter(PartDict).validate_python(d)
|
|
|
|
def test_unknown_type_rejected(self):
|
|
with pytest.raises(Exception):
|
|
TypeAdapter(PartDict).validate_python({"type": "nonsense", "text": "x"})
|
|
|
|
|
|
class TestMessageDictRoundTrip:
|
|
def test_user_message_matches(self):
|
|
d = llm.user("hi").to_dict()
|
|
TypeAdapter(MessageDict).validate_python(d)
|
|
|
|
def test_assistant_with_mixed_parts_matches(self):
|
|
m = llm.Message(
|
|
role="assistant",
|
|
parts=[
|
|
llm.ReasoningPart(
|
|
text="thinking",
|
|
provider_metadata={"anthropic": {"signature": "s"}},
|
|
),
|
|
llm.TextPart(text="answer"),
|
|
llm.ToolCallPart(
|
|
name="search",
|
|
arguments={"q": "x"},
|
|
tool_call_id="c1",
|
|
),
|
|
],
|
|
)
|
|
TypeAdapter(MessageDict).validate_python(m.to_dict())
|
|
|
|
def test_tool_role_message_with_results_matches(self):
|
|
m = llm.tool_message(
|
|
llm.ToolResultPart(name="s", output="r", tool_call_id="c1"),
|
|
)
|
|
TypeAdapter(MessageDict).validate_python(m.to_dict())
|
|
|
|
|
|
class TestResponseDictRoundTrip:
|
|
def test_mock_response_to_dict_matches(self, mock_model):
|
|
mock_model.enqueue(["answer"])
|
|
r = mock_model.prompt("q")
|
|
r.text()
|
|
|
|
d = r.to_dict()
|
|
TypeAdapter(ResponseDict).validate_python(d)
|
|
|
|
def test_response_with_reasoning_matches(self, mock_model):
|
|
mock_model.enqueue(
|
|
[
|
|
llm.StreamEvent(
|
|
type="reasoning",
|
|
chunk="thinking",
|
|
part_index=0,
|
|
provider_metadata={"anthropic": {"signature": "s"}},
|
|
),
|
|
llm.StreamEvent(type="text", chunk="answer", part_index=1),
|
|
]
|
|
)
|
|
r = mock_model.prompt("q")
|
|
r.text()
|
|
|
|
d = r.to_dict()
|
|
TypeAdapter(ResponseDict).validate_python(d)
|
|
|
|
def test_response_with_options_matches(self, mock_model):
|
|
mock_model.enqueue(["ok"])
|
|
r = mock_model.prompt("q", max_tokens=42)
|
|
r.text()
|
|
|
|
d = r.to_dict()
|
|
TypeAdapter(ResponseDict).validate_python(d)
|
|
assert d["prompt"].get("options") == {"max_tokens": 42}
|
|
|
|
|
|
# ---- Literal discriminators ----------------------------------------
|
|
|
|
|
|
class TestLiteralDiscriminators:
|
|
"""The `type` field on each PartDict is a Literal — that's how
|
|
Pydantic's discriminated unions work. Verify each literal."""
|
|
|
|
def test_text_part_literal_is_text(self):
|
|
import typing
|
|
|
|
hints = typing.get_type_hints(TextPartDict)
|
|
# Literal["text"] — check the args
|
|
assert typing.get_args(hints["type"]) == ("text",)
|
|
|
|
def test_reasoning_part_literal_is_reasoning(self):
|
|
import typing
|
|
|
|
hints = typing.get_type_hints(ReasoningPartDict)
|
|
assert typing.get_args(hints["type"]) == ("reasoning",)
|
|
|
|
def test_tool_call_part_literal_is_tool_call(self):
|
|
import typing
|
|
|
|
hints = typing.get_type_hints(ToolCallPartDict)
|
|
assert typing.get_args(hints["type"]) == ("tool_call",)
|
|
|
|
def test_tool_result_part_literal_is_tool_result(self):
|
|
import typing
|
|
|
|
hints = typing.get_type_hints(ToolResultPartDict)
|
|
assert typing.get_args(hints["type"]) == ("tool_result",)
|
|
|
|
def test_attachment_part_literal_is_attachment(self):
|
|
import typing
|
|
|
|
hints = typing.get_type_hints(AttachmentPartDict)
|
|
assert typing.get_args(hints["type"]) == ("attachment",)
|
|
|
|
|
|
# ---- to_dict / from_dict return-type annotations -------------------
|
|
|
|
|
|
class TestAnnotations:
|
|
"""Method signatures should advertise the specific TypedDicts."""
|
|
|
|
def test_text_part_to_dict_annotation(self):
|
|
import typing
|
|
|
|
hints = typing.get_type_hints(llm.TextPart.to_dict)
|
|
assert hints["return"] is TextPartDict
|
|
|
|
def test_reasoning_part_to_dict_annotation(self):
|
|
import typing
|
|
|
|
hints = typing.get_type_hints(llm.ReasoningPart.to_dict)
|
|
assert hints["return"] is ReasoningPartDict
|
|
|
|
def test_tool_call_part_to_dict_annotation(self):
|
|
import typing
|
|
|
|
hints = typing.get_type_hints(llm.ToolCallPart.to_dict)
|
|
assert hints["return"] is ToolCallPartDict
|
|
|
|
def test_tool_result_part_to_dict_annotation(self):
|
|
import typing
|
|
|
|
hints = typing.get_type_hints(llm.ToolResultPart.to_dict)
|
|
assert hints["return"] is ToolResultPartDict
|
|
|
|
def test_attachment_part_to_dict_annotation(self):
|
|
import typing
|
|
|
|
hints = typing.get_type_hints(llm.AttachmentPart.to_dict)
|
|
assert hints["return"] is AttachmentPartDict
|
|
|
|
def test_message_to_dict_annotation(self):
|
|
import typing
|
|
|
|
hints = typing.get_type_hints(llm.Message.to_dict)
|
|
assert hints["return"] is MessageDict
|
|
|
|
def test_message_from_dict_annotation(self):
|
|
import typing
|
|
|
|
hints = typing.get_type_hints(llm.Message.from_dict)
|
|
assert hints["d"] is MessageDict
|
|
|
|
def test_response_to_dict_annotation(self):
|
|
import typing
|
|
|
|
hints = typing.get_type_hints(llm.Response.to_dict)
|
|
assert hints["return"] is ResponseDict
|
|
|
|
|
|
# ---- End-to-end JSON round-trip validates against schema -----------
|
|
|
|
|
|
class TestEndToEnd:
|
|
def test_json_roundtrip_validates(self, mock_model):
|
|
mock_model.enqueue(["text answer"])
|
|
r = mock_model.prompt("q")
|
|
r.text()
|
|
|
|
payload = json.dumps(r.to_dict())
|
|
parsed = json.loads(payload)
|
|
# Parsed dict should still conform to ResponseDict.
|
|
TypeAdapter(ResponseDict).validate_python(parsed)
|