wehub-resource-sync
407 lines
14 KiB
Python
407 lines
14 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.parts.TextPart(text="hello").to_dict()
|
|
self._adapter(TextPartDict).validate_python(d)
|
|
|
|
def test_text_part_with_provider_metadata_matches(self):
|
|
d = llm.parts.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.parts.ReasoningPart(text="", redacted=True).to_dict()
|
|
self._adapter(ReasoningPartDict).validate_python(d)
|
|
|
|
def test_reasoning_part_with_signature_matches(self):
|
|
d = llm.parts.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.parts.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.parts.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.parts.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.parts.AttachmentPart(attachment=att).to_dict()
|
|
self._adapter(AttachmentPartDict).validate_python(d)
|
|
|
|
|
|
class TestPartDiscriminatedUnion:
|
|
def test_text_part_validates_as_part_dict(self):
|
|
d = llm.parts.TextPart(text="hi").to_dict()
|
|
TypeAdapter(PartDict).validate_python(d)
|
|
|
|
def test_reasoning_part_validates_as_part_dict(self):
|
|
d = llm.parts.ReasoningPart(text="thinking").to_dict()
|
|
TypeAdapter(PartDict).validate_python(d)
|
|
|
|
def test_tool_call_part_validates_as_part_dict(self):
|
|
d = llm.parts.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.parts.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.parts.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.parts.ReasoningPart(
|
|
text="thinking",
|
|
provider_metadata={"anthropic": {"signature": "s"}},
|
|
),
|
|
llm.parts.TextPart(text="answer"),
|
|
llm.parts.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.parts.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.parts.StreamEvent(
|
|
type="reasoning",
|
|
chunk="thinking",
|
|
part_index=0,
|
|
provider_metadata={"anthropic": {"signature": "s"}},
|
|
),
|
|
llm.parts.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.parts.TextPart.to_dict)
|
|
assert hints["return"] is TextPartDict
|
|
|
|
def test_reasoning_part_to_dict_annotation(self):
|
|
import typing
|
|
|
|
hints = typing.get_type_hints(llm.parts.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.parts.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.parts.ToolResultPart.to_dict)
|
|
assert hints["return"] is ToolResultPartDict
|
|
|
|
def test_attachment_part_to_dict_annotation(self):
|
|
import typing
|
|
|
|
hints = typing.get_type_hints(llm.parts.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)
|
|
|
|
|
|
# ---- to_dict() must not emit keys absent from the TypedDict --------
|
|
#
|
|
# pydantic's TypeAdapter on a TypedDict silently drops keys that aren't
|
|
# declared, so the round-trip tests above will not catch the case where
|
|
# .to_dict() starts emitting a brand-new key that nobody added to the
|
|
# TypedDict. These tests close that gap by asserting the set of keys
|
|
# .to_dict() returns is a subset of the union of required + optional
|
|
# keys declared on the corresponding TypedDict.
|
|
|
|
|
|
def _allowed(td):
|
|
return td.__required_keys__ | td.__optional_keys__
|
|
|
|
|
|
class TestNoUndeclaredKeys:
|
|
def test_text_part_keys(self):
|
|
d = llm.parts.TextPart(
|
|
text="hi",
|
|
provider_metadata={"k": "v"},
|
|
).to_dict()
|
|
assert set(d.keys()) <= _allowed(TextPartDict)
|
|
|
|
def test_reasoning_part_keys(self):
|
|
d = llm.parts.ReasoningPart(
|
|
text="t",
|
|
redacted=True,
|
|
provider_metadata={"k": "v"},
|
|
).to_dict()
|
|
assert set(d.keys()) <= _allowed(ReasoningPartDict)
|
|
|
|
def test_tool_call_part_keys(self):
|
|
d = llm.parts.ToolCallPart(
|
|
name="t",
|
|
arguments={"q": "x"},
|
|
tool_call_id="c1",
|
|
server_executed=True,
|
|
provider_metadata={"k": "v"},
|
|
).to_dict()
|
|
assert set(d.keys()) <= _allowed(ToolCallPartDict)
|
|
|
|
def test_tool_result_part_keys(self):
|
|
d = llm.parts.ToolResultPart(
|
|
name="t",
|
|
output="r",
|
|
tool_call_id="c1",
|
|
server_executed=True,
|
|
exception="boom",
|
|
attachments=[llm.Attachment(type="image/png", url="http://x/y.png")],
|
|
provider_metadata={"k": "v"},
|
|
).to_dict()
|
|
assert set(d.keys()) <= _allowed(ToolResultPartDict)
|
|
|
|
def test_attachment_part_keys(self):
|
|
d = llm.parts.AttachmentPart(
|
|
attachment=llm.Attachment(type="image/png", url="http://x/y.png"),
|
|
provider_metadata={"k": "v"},
|
|
).to_dict()
|
|
assert set(d.keys()) <= _allowed(AttachmentPartDict)
|
|
|
|
def test_message_keys(self):
|
|
d = llm.Message(
|
|
role="assistant",
|
|
parts=[llm.parts.TextPart(text="hi")],
|
|
provider_metadata={"k": "v"},
|
|
).to_dict()
|
|
assert set(d.keys()) <= _allowed(MessageDict)
|
|
|
|
def test_response_keys(self, mock_model):
|
|
mock_model.enqueue(["answer"])
|
|
r = mock_model.prompt("q", max_tokens=10)
|
|
r.text()
|
|
d = r.to_dict()
|
|
assert set(d.keys()) <= _allowed(ResponseDict)
|
|
# And the nested prompt sub-dict must conform too.
|
|
from llm.serialization import PromptDict
|
|
|
|
assert set(d["prompt"].keys()) <= _allowed(PromptDict)
|