wehub-resource-sync
2107 lines
75 KiB
Python
2107 lines
75 KiB
Python
import json
|
|
import pytest
|
|
import llm
|
|
|
|
|
|
class TestTextPart:
|
|
def test_roundtrip(self):
|
|
part = llm.parts.TextPart(text="Hello world")
|
|
restored = llm.parts.Part.from_dict(part.to_dict())
|
|
assert restored == part
|
|
assert isinstance(restored, llm.parts.TextPart)
|
|
assert restored.text == "Hello world"
|
|
|
|
def test_to_dict_shape(self):
|
|
assert llm.parts.TextPart(text="hi").to_dict() == {"type": "text", "text": "hi"}
|
|
|
|
def test_with_provider_metadata(self):
|
|
part = llm.parts.TextPart(
|
|
text="hi", provider_metadata={"openai": {"flag": True}}
|
|
)
|
|
restored = llm.parts.Part.from_dict(part.to_dict())
|
|
assert restored == part
|
|
|
|
|
|
class TestReasoningPart:
|
|
def test_roundtrip_with_text(self):
|
|
part = llm.parts.ReasoningPart(text="Let me think...")
|
|
restored = llm.parts.Part.from_dict(part.to_dict())
|
|
assert restored == part
|
|
assert restored.text == "Let me think..."
|
|
assert restored.redacted is False
|
|
|
|
def test_roundtrip_redacted(self):
|
|
part = llm.parts.ReasoningPart(text="", redacted=True)
|
|
d = part.to_dict()
|
|
assert d["redacted"] is True
|
|
assert "token_count" not in d
|
|
restored = llm.parts.Part.from_dict(d)
|
|
assert restored == part
|
|
|
|
def test_no_token_count_field(self):
|
|
# token_count was removed: opaque token totals live on
|
|
# response.token_details, not on the Part.
|
|
with pytest.raises(TypeError):
|
|
llm.parts.ReasoningPart(text="", redacted=True, token_count=150)
|
|
|
|
|
|
class TestToolCallPart:
|
|
def test_roundtrip(self):
|
|
part = llm.parts.ToolCallPart(
|
|
name="search",
|
|
arguments={"query": "weather"},
|
|
tool_call_id="call_123",
|
|
)
|
|
restored = llm.parts.Part.from_dict(part.to_dict())
|
|
assert restored == part
|
|
assert restored.server_executed is False
|
|
|
|
def test_server_executed_flag_roundtrips(self):
|
|
part = llm.parts.ToolCallPart(
|
|
name="web_search",
|
|
arguments={"q": "x"},
|
|
tool_call_id="c1",
|
|
server_executed=True,
|
|
)
|
|
d = part.to_dict()
|
|
assert d["server_executed"] is True
|
|
restored = llm.parts.Part.from_dict(d)
|
|
assert restored.server_executed is True
|
|
|
|
|
|
class TestToolResultPart:
|
|
def test_roundtrip(self):
|
|
part = llm.parts.ToolResultPart(
|
|
name="search", output="72F sunny", tool_call_id="c1"
|
|
)
|
|
restored = llm.parts.Part.from_dict(part.to_dict())
|
|
assert restored == part
|
|
assert restored.exception is None
|
|
assert restored.attachments == []
|
|
|
|
def test_with_exception(self):
|
|
part = llm.parts.ToolResultPart(
|
|
name="t", output="", tool_call_id="c1", exception="boom"
|
|
)
|
|
restored = llm.parts.Part.from_dict(part.to_dict())
|
|
assert restored.exception == "boom"
|
|
|
|
|
|
class TestAttachmentPart:
|
|
def test_roundtrip_with_url(self):
|
|
att = llm.Attachment(url="http://example.com/cat.jpg")
|
|
part = llm.parts.AttachmentPart(attachment=att)
|
|
restored = llm.parts.Part.from_dict(part.to_dict())
|
|
assert isinstance(restored, llm.parts.AttachmentPart)
|
|
assert restored.attachment.url == "http://example.com/cat.jpg"
|
|
|
|
def test_roundtrip_with_path(self):
|
|
att = llm.Attachment(type="image/jpeg", path="/tmp/x.jpg")
|
|
part = llm.parts.AttachmentPart(attachment=att)
|
|
restored = llm.parts.Part.from_dict(part.to_dict())
|
|
assert restored.attachment.path == "/tmp/x.jpg"
|
|
assert restored.attachment.type == "image/jpeg"
|
|
|
|
def test_roundtrip_with_bytes_uses_base64(self):
|
|
raw = b"\x89PNG fake bytes"
|
|
att = llm.Attachment(type="image/png", content=raw)
|
|
part = llm.parts.AttachmentPart(attachment=att)
|
|
d = part.to_dict()
|
|
# Content must be a base64-encoded string in the dict form
|
|
assert isinstance(d["attachment"]["content"], str)
|
|
import base64
|
|
|
|
assert base64.b64decode(d["attachment"]["content"]) == raw
|
|
# And round-trip back to the original bytes
|
|
restored = llm.parts.Part.from_dict(d)
|
|
assert restored.attachment.content == raw
|
|
|
|
def test_json_serializable(self):
|
|
att = llm.Attachment(type="image/png", content=b"\x00\x01\x02")
|
|
part = llm.parts.AttachmentPart(attachment=att)
|
|
# Must survive json dumps/loads
|
|
restored = llm.parts.Part.from_dict(json.loads(json.dumps(part.to_dict())))
|
|
assert restored.attachment.content == b"\x00\x01\x02"
|
|
|
|
|
|
class TestUnknownPart:
|
|
def test_from_dict_unknown_type_raises(self):
|
|
with pytest.raises(ValueError):
|
|
llm.parts.Part.from_dict({"type": "nonsense"})
|
|
|
|
|
|
class TestRoleNotOnPart:
|
|
def test_text_part_has_no_role_attribute(self):
|
|
# Role lives on Message. Parts are content-only.
|
|
part = llm.parts.TextPart(text="hi")
|
|
assert not hasattr(part, "role")
|
|
|
|
def test_reasoning_part_has_no_role_attribute(self):
|
|
assert not hasattr(llm.parts.ReasoningPart(text=""), "role")
|
|
|
|
def test_tool_call_part_has_no_role_attribute(self):
|
|
assert not hasattr(
|
|
llm.parts.ToolCallPart(name="t", arguments={}, tool_call_id="c1"),
|
|
"role",
|
|
)
|
|
|
|
|
|
class TestMessage:
|
|
def test_roundtrip_simple_user_message(self):
|
|
m = llm.Message(role="user", parts=[llm.parts.TextPart(text="hi")])
|
|
restored = llm.Message.from_dict(m.to_dict())
|
|
assert restored == m
|
|
|
|
def test_roundtrip_with_provider_metadata(self):
|
|
m = llm.Message(
|
|
role="assistant",
|
|
parts=[llm.parts.TextPart(text="hi")],
|
|
provider_metadata={"anthropic": {"signature": "abc"}},
|
|
)
|
|
restored = llm.Message.from_dict(m.to_dict())
|
|
assert restored == m
|
|
|
|
def test_roundtrip_mixed_parts(self):
|
|
m = llm.Message(
|
|
role="assistant",
|
|
parts=[
|
|
llm.parts.ReasoningPart(text="Thinking"),
|
|
llm.parts.TextPart(text="Result"),
|
|
llm.parts.ToolCallPart(
|
|
name="search",
|
|
arguments={"q": "x"},
|
|
tool_call_id="c1",
|
|
),
|
|
],
|
|
)
|
|
restored = llm.Message.from_dict(m.to_dict())
|
|
assert restored == m
|
|
|
|
def test_empty_provider_metadata_omitted(self):
|
|
m = llm.Message(role="user", parts=[llm.parts.TextPart(text="x")])
|
|
d = m.to_dict()
|
|
assert "provider_metadata" not in d
|
|
|
|
def test_none_and_empty_provider_metadata_equivalent(self):
|
|
m_none = llm.Message(role="user", parts=[llm.parts.TextPart(text="x")])
|
|
m_empty = llm.Message(
|
|
role="user",
|
|
parts=[llm.parts.TextPart(text="x")],
|
|
provider_metadata={},
|
|
)
|
|
# Both serialize the same (empty metadata is omitted)
|
|
assert m_none.to_dict() == m_empty.to_dict()
|
|
|
|
|
|
class TestHelpers:
|
|
def test_user_with_string(self):
|
|
m = llm.user("hi")
|
|
assert m.role == "user"
|
|
assert m.parts == [llm.parts.TextPart(text="hi")]
|
|
|
|
def test_assistant_with_string(self):
|
|
m = llm.assistant("there")
|
|
assert m.role == "assistant"
|
|
assert m.parts == [llm.parts.TextPart(text="there")]
|
|
|
|
def test_system_with_string(self):
|
|
m = llm.system("be brief")
|
|
assert m.role == "system"
|
|
assert m.parts == [llm.parts.TextPart(text="be brief")]
|
|
|
|
def test_tool_message_with_part(self):
|
|
tr = llm.parts.ToolResultPart(name="t", output="r", tool_call_id="c1")
|
|
m = llm.tool_message(tr)
|
|
assert m.role == "tool"
|
|
assert m.parts == [tr]
|
|
|
|
def test_helper_accepts_attachment(self):
|
|
att = llm.Attachment(url="http://example.com/x.jpg")
|
|
m = llm.user("describe this", att)
|
|
assert m.parts == [
|
|
llm.parts.TextPart(text="describe this"),
|
|
llm.parts.AttachmentPart(attachment=att),
|
|
]
|
|
|
|
def test_helper_accepts_existing_part(self):
|
|
tp = llm.parts.TextPart(text="pre-built")
|
|
m = llm.user(tp)
|
|
assert m.parts == [tp]
|
|
|
|
def test_helper_flattens_one_level(self):
|
|
# Nested list gets flattened one level.
|
|
m = llm.user(["one", "two"], "three")
|
|
assert m.parts == [
|
|
llm.parts.TextPart(text="one"),
|
|
llm.parts.TextPart(text="two"),
|
|
llm.parts.TextPart(text="three"),
|
|
]
|
|
|
|
def test_helper_rejects_unknown_types(self):
|
|
with pytest.raises(TypeError):
|
|
llm.user(42)
|
|
|
|
def test_helper_with_provider_metadata(self):
|
|
m = llm.assistant("hi", provider_metadata={"openai": {"id": "x"}})
|
|
assert m.provider_metadata == {"openai": {"id": "x"}}
|
|
|
|
|
|
class TestStreamEvent:
|
|
def test_dataclass_defaults(self):
|
|
ev = llm.parts.StreamEvent(type="text", chunk="hi", part_index=0)
|
|
assert ev.type == "text"
|
|
assert ev.chunk == "hi"
|
|
assert ev.part_index == 0
|
|
assert ev.tool_call_id is None
|
|
assert ev.server_executed is False
|
|
assert ev.tool_name is None
|
|
assert ev.provider_metadata is None
|
|
assert ev.message_index == 0
|
|
|
|
def test_all_fields_accepted(self):
|
|
ev = llm.parts.StreamEvent(
|
|
type="tool_call_args",
|
|
chunk='{"q":',
|
|
part_index=2,
|
|
tool_call_id="c1",
|
|
server_executed=True,
|
|
tool_name="search",
|
|
provider_metadata={"openai": {"x": 1}},
|
|
message_index=1,
|
|
)
|
|
assert ev.tool_call_id == "c1"
|
|
assert ev.server_executed is True
|
|
assert ev.tool_name == "search"
|
|
assert ev.provider_metadata == {"openai": {"x": 1}}
|
|
assert ev.message_index == 1
|
|
|
|
|
|
# Backward compat for plain-str plugins: iterating a Response still
|
|
# yields text strings, response.text() still works, self._chunks is
|
|
# still populated.
|
|
|
|
|
|
class TestPlainStrPluginCompat:
|
|
"""A plugin that yields plain str must still work unchanged."""
|
|
|
|
def test_iter_yields_strings(self, mock_model):
|
|
mock_model.enqueue(["hello", " ", "world"])
|
|
response = mock_model.prompt("hi")
|
|
chunks = list(response)
|
|
assert chunks == ["hello", " ", "world"]
|
|
|
|
def test_text_returns_concatenation(self, mock_model):
|
|
mock_model.enqueue(["hello ", "world"])
|
|
response = mock_model.prompt("hi")
|
|
assert response.text() == "hello world"
|
|
|
|
def test_chunks_are_preserved(self, mock_model):
|
|
mock_model.enqueue(["a", "b", "c"])
|
|
response = mock_model.prompt("hi")
|
|
response.text()
|
|
assert response._chunks == ["a", "b", "c"]
|
|
|
|
|
|
class TestStreamEventsFromPlainStrPlugin:
|
|
"""When a plugin yields plain str, stream_events synthesizes text events."""
|
|
|
|
def test_stream_events_yields_text_events(self, mock_model):
|
|
mock_model.enqueue(["hel", "lo"])
|
|
response = mock_model.prompt("hi")
|
|
events = list(response.stream_events())
|
|
assert all(isinstance(e, llm.parts.StreamEvent) for e in events)
|
|
assert [e.type for e in events] == ["text", "text"]
|
|
assert [e.chunk for e in events] == ["hel", "lo"]
|
|
assert all(e.part_index == 0 for e in events)
|
|
|
|
def test_response_messages_is_single_assistant_text(self, mock_model):
|
|
mock_model.enqueue(["hello"])
|
|
response = mock_model.prompt("hi")
|
|
response.text()
|
|
messages = response.messages()
|
|
assert messages == [
|
|
llm.Message(role="assistant", parts=[llm.parts.TextPart(text="hello")])
|
|
]
|
|
|
|
def test_empty_response_has_empty_messages(self, mock_model):
|
|
mock_model.enqueue([])
|
|
response = mock_model.prompt("hi")
|
|
response.text()
|
|
assert response.messages() == []
|
|
|
|
|
|
class TestStreamEventsFromStreamEventPlugin:
|
|
"""When a plugin yields StreamEvents, they pass through unchanged
|
|
and iteration filters to text only."""
|
|
|
|
def test_iter_yields_only_text_chunks(self, mock_model):
|
|
events = [
|
|
llm.parts.StreamEvent(type="reasoning", chunk="think ", part_index=0),
|
|
llm.parts.StreamEvent(type="text", chunk="hel", part_index=1),
|
|
llm.parts.StreamEvent(type="text", chunk="lo", part_index=1),
|
|
]
|
|
mock_model.enqueue(events)
|
|
response = mock_model.prompt("hi")
|
|
chunks = list(response)
|
|
assert chunks == ["hel", "lo"]
|
|
|
|
def test_stream_events_yields_all_events(self, mock_model):
|
|
events = [
|
|
llm.parts.StreamEvent(type="reasoning", chunk="t", part_index=0),
|
|
llm.parts.StreamEvent(type="text", chunk="x", part_index=1),
|
|
]
|
|
mock_model.enqueue(events)
|
|
response = mock_model.prompt("hi")
|
|
got = list(response.stream_events())
|
|
assert [e.type for e in got] == ["reasoning", "text"]
|
|
|
|
def test_messages_assembles_reasoning_then_text(self, mock_model):
|
|
events = [
|
|
llm.parts.StreamEvent(type="reasoning", chunk="thinking", part_index=0),
|
|
llm.parts.StreamEvent(type="text", chunk="hello", part_index=1),
|
|
]
|
|
mock_model.enqueue(events)
|
|
response = mock_model.prompt("hi")
|
|
response.text()
|
|
assert response.messages() == [
|
|
llm.Message(
|
|
role="assistant",
|
|
parts=[
|
|
llm.parts.ReasoningPart(text="thinking"),
|
|
llm.parts.TextPart(text="hello"),
|
|
],
|
|
)
|
|
]
|
|
|
|
def test_tool_call_name_and_args_merge(self, mock_model):
|
|
events = [
|
|
llm.parts.StreamEvent(type="text", chunk="calling", part_index=0),
|
|
llm.parts.StreamEvent(
|
|
type="tool_call_name",
|
|
chunk="search",
|
|
part_index=1,
|
|
tool_call_id="c1",
|
|
),
|
|
llm.parts.StreamEvent(
|
|
type="tool_call_args",
|
|
chunk='{"q":',
|
|
part_index=1,
|
|
tool_call_id="c1",
|
|
),
|
|
llm.parts.StreamEvent(
|
|
type="tool_call_args",
|
|
chunk='"weather"}',
|
|
part_index=1,
|
|
tool_call_id="c1",
|
|
),
|
|
]
|
|
mock_model.enqueue(events)
|
|
response = mock_model.prompt("hi")
|
|
response.text()
|
|
msgs = response.messages()
|
|
assert len(msgs) == 1
|
|
parts = msgs[0].parts
|
|
assert parts == [
|
|
llm.parts.TextPart(text="calling"),
|
|
llm.parts.ToolCallPart(
|
|
name="search",
|
|
arguments={"q": "weather"},
|
|
tool_call_id="c1",
|
|
),
|
|
]
|
|
|
|
def test_tool_call_args_unparseable_json_falls_back(self, mock_model):
|
|
events = [
|
|
llm.parts.StreamEvent(
|
|
type="tool_call_name",
|
|
chunk="t",
|
|
part_index=0,
|
|
tool_call_id="c1",
|
|
),
|
|
llm.parts.StreamEvent(
|
|
type="tool_call_args",
|
|
chunk="not json",
|
|
part_index=0,
|
|
tool_call_id="c1",
|
|
),
|
|
]
|
|
mock_model.enqueue(events)
|
|
response = mock_model.prompt("hi")
|
|
response.text()
|
|
part = response.messages()[0].parts[0]
|
|
assert part.name == "t"
|
|
assert part.arguments == {"_raw": "not json"}
|
|
|
|
def test_family_mismatch_at_same_part_index_raises(self, mock_model):
|
|
events = [
|
|
llm.parts.StreamEvent(type="text", chunk="x", part_index=0),
|
|
llm.parts.StreamEvent(
|
|
type="tool_call_name",
|
|
chunk="t",
|
|
part_index=0,
|
|
tool_call_id="c1",
|
|
),
|
|
]
|
|
mock_model.enqueue(events)
|
|
response = mock_model.prompt("hi")
|
|
response.text()
|
|
with pytest.raises(ValueError, match="part_index"):
|
|
response.messages() # noqa: B018
|
|
|
|
def test_provider_metadata_merges_last_wins(self, mock_model):
|
|
events = [
|
|
llm.parts.StreamEvent(
|
|
type="reasoning",
|
|
chunk="think",
|
|
part_index=0,
|
|
provider_metadata={"anthropic": {"signature": "one"}},
|
|
),
|
|
llm.parts.StreamEvent(
|
|
type="reasoning",
|
|
chunk="",
|
|
part_index=0,
|
|
provider_metadata={"anthropic": {"signature": "final"}},
|
|
),
|
|
]
|
|
mock_model.enqueue(events)
|
|
response = mock_model.prompt("hi")
|
|
response.text()
|
|
part = response.messages()[0].parts[0]
|
|
assert part.provider_metadata == {"anthropic": {"signature": "final"}}
|
|
|
|
def test_redacted_reasoning_event_emits_marker_part(self, mock_model):
|
|
# A reasoning StreamEvent with redacted=True yields a
|
|
# ReasoningPart(text="", redacted=True) marker — opaque token
|
|
# totals live on response.token_details, not on the Part.
|
|
events = [
|
|
llm.parts.StreamEvent(type="reasoning", chunk="", redacted=True),
|
|
llm.parts.StreamEvent(type="text", chunk="hi"),
|
|
]
|
|
mock_model.enqueue(events)
|
|
response = mock_model.prompt("x")
|
|
response.text()
|
|
parts = response.messages()[0].parts
|
|
assert parts == [
|
|
llm.parts.ReasoningPart(text="", redacted=True),
|
|
llm.parts.TextPart(text="hi"),
|
|
]
|
|
|
|
def test_redacted_reasoning_hoisted_to_start_when_emitted_late(self, mock_model):
|
|
# Plugins typically learn opaque reasoning happened only when
|
|
# the final usage chunk arrives, so they emit the marker last.
|
|
# The framework hoists redacted reasoning Parts to the start of
|
|
# the assembled message so UIs can render them before content.
|
|
events = [
|
|
llm.parts.StreamEvent(type="text", chunk="hello"),
|
|
llm.parts.StreamEvent(type="reasoning", chunk="", redacted=True),
|
|
]
|
|
mock_model.enqueue(events)
|
|
response = mock_model.prompt("x")
|
|
response.text()
|
|
parts = response.messages()[0].parts
|
|
assert parts == [
|
|
llm.parts.ReasoningPart(text="", redacted=True),
|
|
llm.parts.TextPart(text="hello"),
|
|
]
|
|
|
|
def test_redacted_reasoning_event_default_redacted_is_false(self):
|
|
ev = llm.parts.StreamEvent(type="reasoning", chunk="thinking")
|
|
assert ev.redacted is False
|
|
|
|
|
|
class TestPartIndexAutoAllocation:
|
|
"""When part_index is None (the default), the framework groups
|
|
events into Parts using same-family adjacency for text/reasoning
|
|
and tool_call_id for tool calls."""
|
|
|
|
def test_streamevent_part_index_defaults_to_none(self):
|
|
ev = llm.parts.StreamEvent(type="text", chunk="hi")
|
|
assert ev.part_index is None
|
|
|
|
def test_consecutive_text_concatenates_into_one_part(self, mock_model):
|
|
events = [
|
|
llm.parts.StreamEvent(type="text", chunk="hello "),
|
|
llm.parts.StreamEvent(type="text", chunk="world"),
|
|
]
|
|
mock_model.enqueue(events)
|
|
response = mock_model.prompt("hi")
|
|
response.text()
|
|
assert response.messages()[0].parts == [llm.parts.TextPart(text="hello world")]
|
|
|
|
def test_text_then_reasoning_splits_into_two_parts(self, mock_model):
|
|
events = [
|
|
llm.parts.StreamEvent(type="text", chunk="hello"),
|
|
llm.parts.StreamEvent(type="reasoning", chunk="thinking"),
|
|
]
|
|
mock_model.enqueue(events)
|
|
response = mock_model.prompt("hi")
|
|
response.text()
|
|
assert response.messages()[0].parts == [
|
|
llm.parts.TextPart(text="hello"),
|
|
llm.parts.ReasoningPart(text="thinking"),
|
|
]
|
|
|
|
def test_text_tool_call_text_produces_three_parts(self, mock_model):
|
|
events = [
|
|
llm.parts.StreamEvent(type="text", chunk="before"),
|
|
llm.parts.StreamEvent(
|
|
type="tool_call_name",
|
|
chunk="search",
|
|
tool_call_id="c1",
|
|
),
|
|
llm.parts.StreamEvent(
|
|
type="tool_call_args",
|
|
chunk='{"q": "x"}',
|
|
tool_call_id="c1",
|
|
),
|
|
llm.parts.StreamEvent(type="text", chunk="after"),
|
|
]
|
|
mock_model.enqueue(events)
|
|
response = mock_model.prompt("hi")
|
|
response.text()
|
|
assert response.messages()[0].parts == [
|
|
llm.parts.TextPart(text="before"),
|
|
llm.parts.ToolCallPart(
|
|
name="search", arguments={"q": "x"}, tool_call_id="c1"
|
|
),
|
|
llm.parts.TextPart(text="after"),
|
|
]
|
|
|
|
def test_tool_call_groups_by_tool_call_id(self, mock_model):
|
|
events = [
|
|
llm.parts.StreamEvent(
|
|
type="tool_call_name",
|
|
chunk="search",
|
|
tool_call_id="c1",
|
|
),
|
|
llm.parts.StreamEvent(
|
|
type="tool_call_args",
|
|
chunk='{"q":',
|
|
tool_call_id="c1",
|
|
),
|
|
llm.parts.StreamEvent(
|
|
type="tool_call_args",
|
|
chunk='"weather"}',
|
|
tool_call_id="c1",
|
|
),
|
|
]
|
|
mock_model.enqueue(events)
|
|
response = mock_model.prompt("hi")
|
|
response.text()
|
|
assert response.messages()[0].parts == [
|
|
llm.parts.ToolCallPart(
|
|
name="search",
|
|
arguments={"q": "weather"},
|
|
tool_call_id="c1",
|
|
)
|
|
]
|
|
|
|
def test_parallel_tool_calls_interleaved_by_id(self, mock_model):
|
|
# Two tool calls whose args interleave on the wire — must
|
|
# still produce two distinct ToolCallParts grouped by id.
|
|
events = [
|
|
llm.parts.StreamEvent(
|
|
type="tool_call_name", chunk="search", tool_call_id="A"
|
|
),
|
|
llm.parts.StreamEvent(
|
|
type="tool_call_name", chunk="lookup", tool_call_id="B"
|
|
),
|
|
llm.parts.StreamEvent(
|
|
type="tool_call_args", chunk='{"q":"a"}', tool_call_id="A"
|
|
),
|
|
llm.parts.StreamEvent(
|
|
type="tool_call_args", chunk='{"k":"b"}', tool_call_id="B"
|
|
),
|
|
]
|
|
mock_model.enqueue(events)
|
|
response = mock_model.prompt("hi")
|
|
response.text()
|
|
parts = response.messages()[0].parts
|
|
assert parts == [
|
|
llm.parts.ToolCallPart(
|
|
name="search", arguments={"q": "a"}, tool_call_id="A"
|
|
),
|
|
llm.parts.ToolCallPart(
|
|
name="lookup", arguments={"k": "b"}, tool_call_id="B"
|
|
),
|
|
]
|
|
|
|
def test_tool_result_is_always_own_part(self, mock_model):
|
|
events = [
|
|
llm.parts.StreamEvent(
|
|
type="tool_call_name",
|
|
chunk="web_search",
|
|
tool_call_id="c1",
|
|
server_executed=True,
|
|
),
|
|
llm.parts.StreamEvent(
|
|
type="tool_call_args",
|
|
chunk='{"q":"x"}',
|
|
tool_call_id="c1",
|
|
server_executed=True,
|
|
),
|
|
llm.parts.StreamEvent(
|
|
type="tool_result",
|
|
chunk="results...",
|
|
tool_call_id="c1",
|
|
tool_name="web_search",
|
|
server_executed=True,
|
|
),
|
|
]
|
|
mock_model.enqueue(events)
|
|
response = mock_model.prompt("hi")
|
|
response.text()
|
|
parts = response.messages()[0].parts
|
|
assert parts == [
|
|
llm.parts.ToolCallPart(
|
|
name="web_search",
|
|
arguments={"q": "x"},
|
|
tool_call_id="c1",
|
|
server_executed=True,
|
|
),
|
|
llm.parts.ToolResultPart(
|
|
name="web_search",
|
|
output="results...",
|
|
tool_call_id="c1",
|
|
server_executed=True,
|
|
),
|
|
]
|
|
|
|
def test_two_reasoning_blocks_split_by_tool_call(self, mock_model):
|
|
# Some providers emit two thinking blocks separated by a tool
|
|
# call — those should yield two ReasoningParts, not one.
|
|
events = [
|
|
llm.parts.StreamEvent(type="reasoning", chunk="first"),
|
|
llm.parts.StreamEvent(type="tool_call_name", chunk="t", tool_call_id="c1"),
|
|
llm.parts.StreamEvent(type="tool_call_args", chunk="{}", tool_call_id="c1"),
|
|
llm.parts.StreamEvent(type="reasoning", chunk="second"),
|
|
]
|
|
mock_model.enqueue(events)
|
|
response = mock_model.prompt("hi")
|
|
response.text()
|
|
parts = response.messages()[0].parts
|
|
assert parts == [
|
|
llm.parts.ReasoningPart(text="first"),
|
|
llm.parts.ToolCallPart(name="t", arguments={}, tool_call_id="c1"),
|
|
llm.parts.ReasoningPart(text="second"),
|
|
]
|
|
|
|
def test_parallel_tool_calls_without_id_each_get_own_part(self, mock_model):
|
|
# Gemini emits multiple functionCall parts back-to-back without
|
|
# a tool_call_id. Each tool_call_name must allocate a fresh
|
|
# part — otherwise the N tool calls collapse into one with
|
|
# concatenated names and args.
|
|
events = [
|
|
llm.parts.StreamEvent(type="tool_call_name", chunk="store_fact"),
|
|
llm.parts.StreamEvent(type="tool_call_args", chunk='{"fact":"a"}'),
|
|
llm.parts.StreamEvent(type="tool_call_name", chunk="store_fact"),
|
|
llm.parts.StreamEvent(type="tool_call_args", chunk='{"fact":"b"}'),
|
|
llm.parts.StreamEvent(type="tool_call_name", chunk="store_fact"),
|
|
llm.parts.StreamEvent(type="tool_call_args", chunk='{"fact":"c"}'),
|
|
]
|
|
mock_model.enqueue(events)
|
|
response = mock_model.prompt("hi")
|
|
response.text()
|
|
parts = response.messages()[0].parts
|
|
assert parts == [
|
|
llm.parts.ToolCallPart(name="store_fact", arguments={"fact": "a"}),
|
|
llm.parts.ToolCallPart(name="store_fact", arguments={"fact": "b"}),
|
|
llm.parts.ToolCallPart(name="store_fact", arguments={"fact": "c"}),
|
|
]
|
|
|
|
def test_explicit_part_index_still_works(self, mock_model):
|
|
# Back-compat: plugins that pass explicit part_index should
|
|
# behave exactly as before.
|
|
events = [
|
|
llm.parts.StreamEvent(type="reasoning", chunk="t", part_index=0),
|
|
llm.parts.StreamEvent(type="text", chunk="hi", part_index=1),
|
|
]
|
|
mock_model.enqueue(events)
|
|
response = mock_model.prompt("hi")
|
|
response.text()
|
|
assert response.messages()[0].parts == [
|
|
llm.parts.ReasoningPart(text="t"),
|
|
llm.parts.TextPart(text="hi"),
|
|
]
|
|
|
|
def test_mix_explicit_zero_and_none_for_text_concatenates(self, mock_model):
|
|
# Forcing a single TextPart across non-adjacent text bursts:
|
|
# plugin pins explicit part_index=0 on the wraparound text
|
|
# events, and the tool call in between gets None (auto).
|
|
events = [
|
|
llm.parts.StreamEvent(type="text", chunk="before ", part_index=0),
|
|
llm.parts.StreamEvent(type="tool_call_name", chunk="t", tool_call_id="c1"),
|
|
llm.parts.StreamEvent(type="tool_call_args", chunk="{}", tool_call_id="c1"),
|
|
llm.parts.StreamEvent(type="text", chunk="after", part_index=0),
|
|
]
|
|
mock_model.enqueue(events)
|
|
response = mock_model.prompt("hi")
|
|
response.text()
|
|
parts = response.messages()[0].parts
|
|
assert parts == [
|
|
llm.parts.TextPart(text="before after"),
|
|
llm.parts.ToolCallPart(name="t", arguments={}, tool_call_id="c1"),
|
|
]
|
|
|
|
|
|
class TestStreamEventsLiveDuringStreaming:
|
|
"""Client code sees events arrive before the response is done"""
|
|
|
|
def test_events_arrive_before_done(self, mock_model):
|
|
events = [
|
|
llm.parts.StreamEvent(type="reasoning", chunk="t", part_index=0),
|
|
llm.parts.StreamEvent(type="text", chunk="hi", part_index=1),
|
|
]
|
|
mock_model.enqueue(events)
|
|
response = mock_model.prompt("x")
|
|
seen = []
|
|
for event in response.stream_events():
|
|
# Record the _done state at the moment we receive the event.
|
|
seen.append((event.type, response._done))
|
|
# Events arrived before _done was set.
|
|
assert [s[0] for s in seen] == ["reasoning", "text"]
|
|
assert all(not done for _type, done in seen)
|
|
# And after the generator is drained, the response is done.
|
|
assert response._done
|
|
|
|
def test_stream_events_after_done_replays(self, mock_model):
|
|
mock_model.enqueue(
|
|
[llm.parts.StreamEvent(type="text", chunk="hi", part_index=0)]
|
|
)
|
|
response = mock_model.prompt("x")
|
|
first = list(response.stream_events())
|
|
# Second call replays from the stored events.
|
|
second = list(response.stream_events())
|
|
assert len(first) == 1
|
|
assert [e.type for e in second] == ["text"]
|
|
assert [e.chunk for e in second] == ["hi"]
|
|
|
|
def test_plain_str_stream_events_after_done_replays(self, mock_model):
|
|
mock_model.enqueue(["hello"])
|
|
response = mock_model.prompt("x")
|
|
response.text()
|
|
events = list(response.stream_events())
|
|
assert len(events) == 1
|
|
assert events[0].type == "text"
|
|
assert events[0].chunk == "hello"
|
|
|
|
|
|
class TestAsyncStreamEvents:
|
|
@pytest.mark.asyncio
|
|
async def test_async_stream_events_live(self, async_mock_model):
|
|
events = [
|
|
llm.parts.StreamEvent(type="reasoning", chunk="r", part_index=0),
|
|
llm.parts.StreamEvent(type="text", chunk="t", part_index=1),
|
|
]
|
|
async_mock_model.enqueue(events)
|
|
response = async_mock_model.prompt("x")
|
|
seen_types = []
|
|
async for event in response.astream_events():
|
|
seen_types.append(event.type)
|
|
assert seen_types == ["reasoning", "text"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_iter_yields_only_text(self, async_mock_model):
|
|
events = [
|
|
llm.parts.StreamEvent(type="reasoning", chunk="r", part_index=0),
|
|
llm.parts.StreamEvent(type="text", chunk="hi", part_index=1),
|
|
]
|
|
async_mock_model.enqueue(events)
|
|
response = async_mock_model.prompt("x")
|
|
chunks = []
|
|
async for chunk in response:
|
|
chunks.append(chunk)
|
|
assert chunks == ["hi"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_messages_after_await(self, async_mock_model):
|
|
async_mock_model.enqueue(["hi"])
|
|
response = async_mock_model.prompt("x")
|
|
await response.text()
|
|
assert await response.messages() == [
|
|
llm.Message(role="assistant", parts=[llm.parts.TextPart(text="hi")])
|
|
]
|
|
|
|
|
|
class TestMessagesIsCallable:
|
|
"""response.messages() is a method (matching .text(), .json(),
|
|
.tool_calls()) — invocation forces execution if not yet done.
|
|
"""
|
|
|
|
def test_sync_messages_is_callable_and_returns_list(self, mock_model):
|
|
mock_model.enqueue(["hi"])
|
|
response = mock_model.prompt("x")
|
|
# No prior .text() or iteration — calling messages() forces
|
|
# execution and returns the assembled list.
|
|
assert response.messages() == [
|
|
llm.Message(role="assistant", parts=[llm.parts.TextPart(text="hi")])
|
|
]
|
|
|
|
def test_sync_messages_after_text_returns_same_list(self, mock_model):
|
|
mock_model.enqueue(["hi"])
|
|
response = mock_model.prompt("x")
|
|
response.text()
|
|
assert response.messages() == [
|
|
llm.Message(role="assistant", parts=[llm.parts.TextPart(text="hi")])
|
|
]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_messages_is_awaitable(self, async_mock_model):
|
|
async_mock_model.enqueue(["hi"])
|
|
response = async_mock_model.prompt("x")
|
|
# No prior await — `await response.messages()` forces it.
|
|
result = await response.messages()
|
|
assert result == [
|
|
llm.Message(role="assistant", parts=[llm.parts.TextPart(text="hi")])
|
|
]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_messages_after_text_returns_same_list(self, async_mock_model):
|
|
async_mock_model.enqueue(["hi"])
|
|
response = async_mock_model.prompt("x")
|
|
await response.text()
|
|
result = await response.messages()
|
|
assert result == [
|
|
llm.Message(role="assistant", parts=[llm.parts.TextPart(text="hi")])
|
|
]
|
|
|
|
|
|
class TestPromptMessagesSynthesis:
|
|
"""Prompt.messages constructs a Message list from legacy inputs when
|
|
messages= wasn't passed explicitly."""
|
|
|
|
def test_empty_prompt_yields_empty_messages(self, mock_model):
|
|
from llm.models import Prompt
|
|
|
|
p = Prompt(None, model=mock_model)
|
|
assert p.messages == []
|
|
|
|
def test_prompt_text_synthesizes_user_message(self, mock_model):
|
|
from llm.models import Prompt
|
|
|
|
p = Prompt("hi", model=mock_model)
|
|
assert p.messages == [
|
|
llm.Message(role="user", parts=[llm.parts.TextPart(text="hi")])
|
|
]
|
|
|
|
def test_system_and_prompt_synthesizes_two_messages(self, mock_model):
|
|
from llm.models import Prompt
|
|
|
|
p = Prompt("hi", model=mock_model, system="be brief")
|
|
assert p.messages == [
|
|
llm.Message(role="system", parts=[llm.parts.TextPart(text="be brief")]),
|
|
llm.Message(role="user", parts=[llm.parts.TextPart(text="hi")]),
|
|
]
|
|
|
|
def test_attachments_join_user_message(self, mock_model):
|
|
from llm.models import Prompt
|
|
|
|
att = llm.Attachment(url="http://example.com/a.jpg")
|
|
p = Prompt("look", model=mock_model, attachments=[att])
|
|
assert p.messages == [
|
|
llm.Message(
|
|
role="user",
|
|
parts=[
|
|
llm.parts.TextPart(text="look"),
|
|
llm.parts.AttachmentPart(attachment=att),
|
|
],
|
|
)
|
|
]
|
|
|
|
def test_tool_results_become_tool_role_message(self, mock_model):
|
|
from llm.models import Prompt
|
|
from llm import ToolResult
|
|
|
|
tr = ToolResult(name="t", output="ok", tool_call_id="c1")
|
|
p = Prompt(None, model=mock_model, tool_results=[tr])
|
|
assert p.messages == [
|
|
llm.Message(
|
|
role="tool",
|
|
parts=[
|
|
llm.parts.ToolResultPart(name="t", output="ok", tool_call_id="c1")
|
|
],
|
|
)
|
|
]
|
|
|
|
|
|
class TestPromptMessagesExplicit:
|
|
"""When messages= is passed, it's authoritative."""
|
|
|
|
def test_explicit_messages_returned_verbatim(self, mock_model):
|
|
from llm.models import Prompt
|
|
|
|
explicit = [
|
|
llm.system("x"),
|
|
llm.user("y"),
|
|
]
|
|
p = Prompt(None, model=mock_model, messages=explicit)
|
|
assert p.messages == explicit
|
|
|
|
def test_explicit_messages_ignores_prompt_kwarg(self, mock_model):
|
|
"""Explicit messages= is authoritative. A prompt= string passed
|
|
alongside is no longer auto-appended — the invariant is that
|
|
prompt.messages equals exactly what the model was sent."""
|
|
from llm.models import Prompt
|
|
|
|
explicit = [llm.system("x"), llm.user("prior"), llm.user("follow-up")]
|
|
p = Prompt("ignored text", model=mock_model, messages=explicit)
|
|
assert p.messages == explicit
|
|
|
|
def test_explicit_messages_independent_copy(self, mock_model):
|
|
"""Mutating the caller's list must not mutate Prompt.messages."""
|
|
from llm.models import Prompt
|
|
|
|
explicit = [llm.user("x")]
|
|
p = Prompt(None, model=mock_model, messages=explicit)
|
|
explicit.append(llm.user("later"))
|
|
assert p.messages == [llm.user("x")]
|
|
|
|
|
|
class TestModelPromptMessagesKwarg:
|
|
"""model.prompt / conversation.prompt / async counterparts accept
|
|
messages= and the list is observable on the resulting Prompt."""
|
|
|
|
def test_model_prompt_accepts_messages(self, mock_model):
|
|
mock_model.enqueue(["ok"])
|
|
response = mock_model.prompt(messages=[llm.user("hi")])
|
|
response.text()
|
|
assert response.prompt.messages == [llm.user("hi")]
|
|
|
|
def test_model_prompt_messages_with_system(self, mock_model):
|
|
mock_model.enqueue(["ok"])
|
|
response = mock_model.prompt(messages=[llm.system("be brief"), llm.user("hi")])
|
|
response.text()
|
|
assert response.prompt.messages == [
|
|
llm.system("be brief"),
|
|
llm.user("hi"),
|
|
]
|
|
|
|
def test_conversation_prompt_accepts_messages(self, mock_model):
|
|
mock_model.enqueue(["ok"])
|
|
conv = mock_model.conversation()
|
|
response = conv.prompt(messages=[llm.user("q")])
|
|
response.text()
|
|
assert response.prompt.messages == [llm.user("q")]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_model_prompt_accepts_messages(self, async_mock_model):
|
|
async_mock_model.enqueue(["ok"])
|
|
response = async_mock_model.prompt(messages=[llm.user("hi")])
|
|
await response.text()
|
|
assert response.prompt.messages == [llm.user("hi")]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_conversation_prompt_accepts_messages(self, async_mock_model):
|
|
async_mock_model.enqueue(["ok"])
|
|
conv = async_mock_model.conversation()
|
|
response = conv.prompt(messages=[llm.user("q")])
|
|
await response.text()
|
|
assert response.prompt.messages == [llm.user("q")]
|
|
|
|
|
|
# Invariant: response.prompt.messages == exactly what the model was
|
|
# sent for this turn, regardless of whether the caller used
|
|
# model.prompt(messages=[...]), conversation.prompt("text"), or
|
|
# response.reply("text").
|
|
|
|
|
|
class TestConversationFullChainInvariant:
|
|
def test_explicit_messages_is_authoritative_no_prompt_combine(self, mock_model):
|
|
"""Explicit messages= is the whole list. If prompt= is ALSO
|
|
passed, it's ignored for messages-building — the caller asked
|
|
for exact control."""
|
|
mock_model.enqueue(["ok"])
|
|
response = mock_model.prompt(
|
|
"this prompt argument is ignored",
|
|
messages=[llm.user("q")],
|
|
)
|
|
response.text()
|
|
assert response.prompt.messages == [llm.user("q")]
|
|
|
|
def test_conversation_second_turn_prompt_messages_has_full_chain(self, mock_model):
|
|
mock_model.enqueue(["a1"])
|
|
mock_model.enqueue(["a2"])
|
|
conv = mock_model.conversation()
|
|
|
|
r1 = conv.prompt("q1")
|
|
r1.text()
|
|
r2 = conv.prompt("q2")
|
|
r2.text()
|
|
|
|
# r2 was sent the full chain.
|
|
assert r2.prompt.messages == [
|
|
llm.user("q1"),
|
|
llm.assistant("a1"),
|
|
llm.user("q2"),
|
|
]
|
|
|
|
def test_conversation_third_turn_includes_everything_before(self, mock_model):
|
|
mock_model.enqueue(["a1"])
|
|
mock_model.enqueue(["a2"])
|
|
mock_model.enqueue(["a3"])
|
|
conv = mock_model.conversation()
|
|
r1 = conv.prompt("q1")
|
|
r1.text()
|
|
r2 = conv.prompt("q2")
|
|
r2.text()
|
|
r3 = conv.prompt("q3")
|
|
r3.text()
|
|
|
|
assert r3.prompt.messages == [
|
|
llm.user("q1"),
|
|
llm.assistant("a1"),
|
|
llm.user("q2"),
|
|
llm.assistant("a2"),
|
|
llm.user("q3"),
|
|
]
|
|
|
|
def test_conversation_first_turn_chain_is_single_user_message(self, mock_model):
|
|
mock_model.enqueue(["a1"])
|
|
conv = mock_model.conversation()
|
|
r1 = conv.prompt("q1")
|
|
r1.text()
|
|
assert r1.prompt.messages == [llm.user("q1")]
|
|
|
|
def test_conversation_preserves_reasoning_and_tool_call_parts(self, mock_model):
|
|
"""The chain carries reasoning and tool calls from prior turns,
|
|
not just the flat text — required for multi-turn extended
|
|
thinking (Claude) and tool-use round-trips."""
|
|
mock_model.enqueue(
|
|
[
|
|
llm.parts.StreamEvent(
|
|
type="reasoning", chunk="thinking...", part_index=0
|
|
),
|
|
llm.parts.StreamEvent(type="text", chunk="answer", part_index=1),
|
|
]
|
|
)
|
|
mock_model.enqueue(["follow-up answer"])
|
|
conv = mock_model.conversation()
|
|
r1 = conv.prompt("q1")
|
|
r1.text()
|
|
r2 = conv.prompt("q2")
|
|
r2.text()
|
|
|
|
assert r2.prompt.messages == [
|
|
llm.user("q1"),
|
|
llm.Message(
|
|
role="assistant",
|
|
parts=[
|
|
llm.parts.ReasoningPart(text="thinking..."),
|
|
llm.parts.TextPart(text="answer"),
|
|
],
|
|
),
|
|
llm.user("q2"),
|
|
]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_conversation_full_chain(self, async_mock_model):
|
|
async_mock_model.enqueue(["a1"])
|
|
async_mock_model.enqueue(["a2"])
|
|
conv = async_mock_model.conversation()
|
|
r1 = conv.prompt("q1")
|
|
await r1.text()
|
|
r2 = conv.prompt("q2")
|
|
await r2.text()
|
|
|
|
assert r2.prompt.messages == [
|
|
llm.user("q1"),
|
|
llm.assistant("a1"),
|
|
llm.user("q2"),
|
|
]
|
|
|
|
|
|
class TestSqliteRehydrateMessages:
|
|
"""After Response.from_row, response.messages() must still yield the
|
|
assistant turn as a TextPart (+ any tool calls). Otherwise
|
|
Conversation.prompt builds a broken chain for `llm -c`.
|
|
"""
|
|
|
|
def test_from_row_response_messages_synthesized_from_chunks(
|
|
self, mock_model, tmp_path
|
|
):
|
|
import sqlite_utils
|
|
from llm.migrations import migrate
|
|
|
|
mock_model.enqueue(["answer text"])
|
|
r1 = mock_model.prompt("q1")
|
|
r1.text()
|
|
|
|
db = sqlite_utils.Database(str(tmp_path / "logs.db"))
|
|
migrate(db)
|
|
r1.log_to_db(db)
|
|
|
|
# Rehydrate the response
|
|
row = next(db["responses"].rows)
|
|
rehydrated = llm.Response.from_row(db, row)
|
|
# _stream_events is empty (SQLite doesn't persist those), but
|
|
# _chunks carries the text. response.messages() must fall back
|
|
# to synthesizing a TextPart.
|
|
assert rehydrated._stream_events == []
|
|
assert rehydrated.messages() == [
|
|
llm.Message(
|
|
role="assistant", parts=[llm.parts.TextPart(text="answer text")]
|
|
)
|
|
]
|
|
|
|
def test_llm_dash_c_chain_preserves_prior_assistant_turn(
|
|
self, mock_model, tmp_path
|
|
):
|
|
"""End-to-end: a follow-up turn via load_conversation must send
|
|
[user(q1), assistant(a1), user(q2)] — not drop the assistant."""
|
|
import sqlite_utils
|
|
from llm.migrations import migrate
|
|
from llm.cli import load_conversation
|
|
|
|
mock_model.enqueue(["first answer"])
|
|
mock_model.enqueue(["second answer"])
|
|
r1 = mock_model.prompt("q1")
|
|
r1.text()
|
|
|
|
db_path = tmp_path / "logs.db"
|
|
db = sqlite_utils.Database(str(db_path))
|
|
migrate(db)
|
|
r1.log_to_db(db)
|
|
|
|
conv = load_conversation(None, database=str(db_path))
|
|
r2 = conv.prompt("q2")
|
|
r2.text()
|
|
|
|
assert r2.prompt.messages == [
|
|
llm.user("q1"),
|
|
llm.assistant("first answer"),
|
|
llm.user("q2"),
|
|
]
|
|
|
|
def test_llm_dash_c_after_logged_tool_chain_preserves_full_chain(
|
|
self, mock_model, tmp_path
|
|
):
|
|
"""A loaded tool-result response must carry the preceding
|
|
assistant tool_use. Otherwise Anthropic sees an orphan
|
|
tool_result at the start of the continued request."""
|
|
import sqlite_utils
|
|
from llm.cli import load_conversation
|
|
from llm.migrations import migrate
|
|
|
|
class ToolChainMock(type(mock_model)):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.calls = 0
|
|
|
|
def execute(self, prompt, stream, response, conversation):
|
|
self.calls += 1
|
|
if self.calls == 1:
|
|
response.add_tool_call(
|
|
llm.ToolCall(name="tick", arguments={}, tool_call_id="c1")
|
|
)
|
|
if False:
|
|
yield ""
|
|
else:
|
|
yield "final answer"
|
|
|
|
def tick() -> str:
|
|
return "tock"
|
|
|
|
m = ToolChainMock()
|
|
chain_response = m.chain("q1", tools=[tick])
|
|
chain_response.text()
|
|
|
|
db_path = tmp_path / "logs.db"
|
|
db = sqlite_utils.Database(str(db_path))
|
|
migrate(db)
|
|
chain_response.log_to_db(db)
|
|
|
|
conv = load_conversation(None, database=str(db_path))
|
|
r3 = conv.prompt("q2")
|
|
|
|
assert [m.role for m in r3.prompt.messages] == [
|
|
"user",
|
|
"assistant",
|
|
"tool",
|
|
"assistant",
|
|
"user",
|
|
]
|
|
assert isinstance(r3.prompt.messages[1].parts[0], llm.parts.ToolCallPart)
|
|
assert isinstance(r3.prompt.messages[2].parts[0], llm.parts.ToolResultPart)
|
|
assert r3.prompt.messages[2].parts[0].tool_call_id == "c1"
|
|
|
|
|
|
class TestAddToolCallWithStreamEvents:
|
|
"""A plugin may yield StreamEvents *and* call response.add_tool_call().
|
|
The Part list must include the tool call regardless of whether
|
|
_stream_events is empty or populated; otherwise persistence drops the
|
|
tool call and the next turn ships an orphan tool_result.
|
|
"""
|
|
|
|
def test_text_yield_plus_add_tool_call_emits_both_parts(self, mock_model):
|
|
class TextAndAddToolCallMock(type(mock_model)):
|
|
def execute(self, prompt, stream, response, conversation):
|
|
yield "answer"
|
|
response.add_tool_call(
|
|
llm.ToolCall(
|
|
name="search",
|
|
arguments={"q": "weather"},
|
|
tool_call_id="c1",
|
|
)
|
|
)
|
|
|
|
m = TextAndAddToolCallMock()
|
|
response = m.prompt("hi")
|
|
response.text()
|
|
parts = response.messages()[0].parts
|
|
assert llm.parts.TextPart(text="answer") in parts
|
|
tool_call_parts = [p for p in parts if isinstance(p, llm.parts.ToolCallPart)]
|
|
assert tool_call_parts == [
|
|
llm.parts.ToolCallPart(
|
|
name="search",
|
|
arguments={"q": "weather"},
|
|
tool_call_id="c1",
|
|
)
|
|
]
|
|
|
|
def test_stream_event_tool_call_plus_matching_add_tool_call_dedups(
|
|
self, mock_model
|
|
):
|
|
class DualApiMock(type(mock_model)):
|
|
def execute(self, prompt, stream, response, conversation):
|
|
yield llm.parts.StreamEvent(
|
|
type="tool_call_name", chunk="search", tool_call_id="c1"
|
|
)
|
|
yield llm.parts.StreamEvent(
|
|
type="tool_call_args",
|
|
chunk='{"q":"weather"}',
|
|
tool_call_id="c1",
|
|
)
|
|
response.add_tool_call(
|
|
llm.ToolCall(
|
|
name="search",
|
|
arguments={"q": "weather"},
|
|
tool_call_id="c1",
|
|
)
|
|
)
|
|
|
|
m = DualApiMock()
|
|
response = m.prompt("hi")
|
|
response.text()
|
|
tool_call_parts = [
|
|
p
|
|
for p in response.messages()[0].parts
|
|
if isinstance(p, llm.parts.ToolCallPart)
|
|
]
|
|
assert tool_call_parts == [
|
|
llm.parts.ToolCallPart(
|
|
name="search",
|
|
arguments={"q": "weather"},
|
|
tool_call_id="c1",
|
|
)
|
|
]
|
|
|
|
|
|
class TestResponseReply:
|
|
def test_reply_builds_next_turn_from_this_response(self, mock_model):
|
|
mock_model.enqueue(["a1"])
|
|
mock_model.enqueue(["a2"])
|
|
r1 = mock_model.prompt("q1")
|
|
r1.text()
|
|
|
|
r2 = r1.reply("q2")
|
|
r2.text()
|
|
assert r2.prompt.messages == [
|
|
llm.user("q1"),
|
|
llm.assistant("a1"),
|
|
llm.user("q2"),
|
|
]
|
|
|
|
def test_reply_chains(self, mock_model):
|
|
mock_model.enqueue(["a1"])
|
|
mock_model.enqueue(["a2"])
|
|
mock_model.enqueue(["a3"])
|
|
r1 = mock_model.prompt("q1")
|
|
r1.text()
|
|
r2 = r1.reply("q2")
|
|
r2.text()
|
|
r3 = r2.reply("q3")
|
|
r3.text()
|
|
assert r3.prompt.messages == [
|
|
llm.user("q1"),
|
|
llm.assistant("a1"),
|
|
llm.user("q2"),
|
|
llm.assistant("a2"),
|
|
llm.user("q3"),
|
|
]
|
|
|
|
def test_reply_no_prompt_reuses_messages_kwarg(self, mock_model):
|
|
"""Passing messages= to reply() appends those onto the chain
|
|
in place of a new user string."""
|
|
mock_model.enqueue(["a1"])
|
|
mock_model.enqueue(["a2"])
|
|
r1 = mock_model.prompt("q1")
|
|
r1.text()
|
|
r2 = r1.reply(messages=[llm.user("alt")])
|
|
r2.text()
|
|
assert r2.prompt.messages == [
|
|
llm.user("q1"),
|
|
llm.assistant("a1"),
|
|
llm.user("alt"),
|
|
]
|
|
|
|
def test_reply_from_conversation_response_extends_chain(self, mock_model):
|
|
mock_model.enqueue(["a1"])
|
|
mock_model.enqueue(["a2"])
|
|
conv = mock_model.conversation()
|
|
r1 = conv.prompt("q1")
|
|
r1.text()
|
|
r2 = r1.reply("q2")
|
|
r2.text()
|
|
assert r2.prompt.messages == [
|
|
llm.user("q1"),
|
|
llm.assistant("a1"),
|
|
llm.user("q2"),
|
|
]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_reply(self, async_mock_model):
|
|
async_mock_model.enqueue(["a1"])
|
|
async_mock_model.enqueue(["a2"])
|
|
r1 = async_mock_model.prompt("q1")
|
|
await r1.text()
|
|
r2 = await r1.reply("q2")
|
|
await r2.text()
|
|
assert r2.prompt.messages == [
|
|
llm.user("q1"),
|
|
llm.assistant("a1"),
|
|
llm.user("q2"),
|
|
]
|
|
|
|
def test_reply_with_tool_results_appends_tool_message(self, mock_model):
|
|
# model.prompt(...) makes tool calls, the
|
|
# caller runs them, then reply(tool_results=...) sends the
|
|
# results back in one call. The chain should grow by a
|
|
# role="tool" message containing ToolResultParts.
|
|
from llm.parts import (
|
|
Message,
|
|
ToolCallPart,
|
|
ToolResultPart,
|
|
)
|
|
|
|
# First-turn assistant message has a tool call.
|
|
first_assistant = Message(
|
|
role="assistant",
|
|
parts=[ToolCallPart(name="echo", arguments={"x": 1}, tool_call_id="c1")],
|
|
)
|
|
|
|
class ToolCallMock(type(mock_model)):
|
|
supports_tools = True
|
|
|
|
def execute(self, prompt, stream, response, conversation):
|
|
# Yield the assistant turn's parts as StreamEvents so
|
|
# response.messages() contains the tool call.
|
|
yield llm.parts.StreamEvent(
|
|
type="tool_call_name",
|
|
chunk="echo",
|
|
tool_call_id="c1",
|
|
)
|
|
yield llm.parts.StreamEvent(
|
|
type="tool_call_args",
|
|
chunk='{"x": 1}',
|
|
tool_call_id="c1",
|
|
)
|
|
|
|
m = ToolCallMock()
|
|
r1 = m.prompt("call echo")
|
|
r1.text()
|
|
|
|
tool_results = [llm.ToolResult(name="echo", output="ok", tool_call_id="c1")]
|
|
# The bug we're fixing: this previously silently dropped the
|
|
# tool_results because reply() forwards via messages= and the
|
|
# Prompt synthesis path is bypassed.
|
|
m.enqueue(["follow-up text"])
|
|
r2 = r1.reply(tool_results=tool_results)
|
|
r2.text()
|
|
assert r2.prompt.messages == [
|
|
llm.user("call echo"),
|
|
first_assistant,
|
|
Message(
|
|
role="tool",
|
|
parts=[ToolResultPart(name="echo", output="ok", tool_call_id="c1")],
|
|
),
|
|
]
|
|
|
|
def test_reply_with_tool_results_and_prompt(self, mock_model):
|
|
from llm.parts import ToolResultPart
|
|
|
|
class ToolCallMock(type(mock_model)):
|
|
supports_tools = True
|
|
|
|
def execute(self, prompt, stream, response, conversation):
|
|
yield llm.parts.StreamEvent(
|
|
type="tool_call_name",
|
|
chunk="echo",
|
|
tool_call_id="c1",
|
|
)
|
|
yield llm.parts.StreamEvent(
|
|
type="tool_call_args",
|
|
chunk='{"x": 1}',
|
|
tool_call_id="c1",
|
|
)
|
|
|
|
m = ToolCallMock()
|
|
r1 = m.prompt("call echo")
|
|
r1.text()
|
|
m.enqueue(["follow-up"])
|
|
r2 = r1.reply(
|
|
"now summarise",
|
|
tool_results=[llm.ToolResult(name="echo", output="ok", tool_call_id="c1")],
|
|
)
|
|
r2.text()
|
|
roles = [m.role for m in r2.prompt.messages]
|
|
assert roles == ["user", "assistant", "tool", "user"]
|
|
# tool message goes BEFORE the new user prompt.
|
|
tool_msg = r2.prompt.messages[2]
|
|
assert tool_msg.parts == [
|
|
ToolResultPart(name="echo", output="ok", tool_call_id="c1")
|
|
]
|
|
assert r2.prompt.messages[3] == llm.user("now summarise")
|
|
|
|
def test_reply_auto_executes_tool_calls_when_none_passed(self, mock_model):
|
|
# Zero-arg sugar: response.reply() with tool calls present
|
|
# auto-executes them and threads results back into the chain.
|
|
from llm.parts import ToolResultPart
|
|
|
|
executed = []
|
|
|
|
def echo(x: int) -> str:
|
|
executed.append(x)
|
|
return f"echo:{x}"
|
|
|
|
class ToolCallMock(type(mock_model)):
|
|
supports_tools = True
|
|
|
|
def execute(self, prompt, stream, response, conversation):
|
|
response.add_tool_call(
|
|
llm.ToolCall(name="echo", arguments={"x": 42}, tool_call_id="c1")
|
|
)
|
|
yield llm.parts.StreamEvent(
|
|
type="tool_call_name", chunk="echo", tool_call_id="c1"
|
|
)
|
|
yield llm.parts.StreamEvent(
|
|
type="tool_call_args",
|
|
chunk='{"x": 42}',
|
|
tool_call_id="c1",
|
|
)
|
|
|
|
m = ToolCallMock()
|
|
r1 = m.prompt("call echo", tools=[echo])
|
|
r1.text()
|
|
|
|
m.enqueue(["follow-up"])
|
|
# No tool_results passed — sugar kicks in and auto-executes.
|
|
r2 = r1.reply()
|
|
r2.text()
|
|
|
|
assert executed == [42]
|
|
# The tool message landed in the chain.
|
|
roles = [msg.role for msg in r2.prompt.messages]
|
|
assert roles == ["user", "assistant", "tool"]
|
|
tool_msg = r2.prompt.messages[2]
|
|
assert tool_msg.parts == [
|
|
ToolResultPart(name="echo", output="echo:42", tool_call_id="c1")
|
|
]
|
|
|
|
def test_reply_auto_execute_with_prompt(self, mock_model):
|
|
# reply("more text") with tool calls present also auto-executes
|
|
# so the user prompt can land after the tool results.
|
|
executed = []
|
|
|
|
def echo(x: int) -> str:
|
|
executed.append(x)
|
|
return "out"
|
|
|
|
class ToolCallMock(type(mock_model)):
|
|
supports_tools = True
|
|
|
|
def execute(self, prompt, stream, response, conversation):
|
|
response.add_tool_call(
|
|
llm.ToolCall(name="echo", arguments={"x": 1}, tool_call_id="c1")
|
|
)
|
|
yield llm.parts.StreamEvent(
|
|
type="tool_call_name", chunk="echo", tool_call_id="c1"
|
|
)
|
|
yield llm.parts.StreamEvent(
|
|
type="tool_call_args",
|
|
chunk='{"x": 1}',
|
|
tool_call_id="c1",
|
|
)
|
|
|
|
m = ToolCallMock()
|
|
r1 = m.prompt("call echo", tools=[echo])
|
|
r1.text()
|
|
m.enqueue(["follow-up"])
|
|
r2 = r1.reply("now summarise")
|
|
r2.text()
|
|
assert executed == [1]
|
|
roles = [msg.role for msg in r2.prompt.messages]
|
|
assert roles == ["user", "assistant", "tool", "user"]
|
|
|
|
def test_reply_explicit_tool_results_skips_auto_execute(self, mock_model):
|
|
# Passing tool_results= explicitly overrides the sugar — the
|
|
# tool function does NOT run (caller already ran it / wants
|
|
# custom results).
|
|
executed = []
|
|
|
|
def echo(x: int) -> str:
|
|
executed.append(x)
|
|
return "should not see"
|
|
|
|
class ToolCallMock(type(mock_model)):
|
|
supports_tools = True
|
|
|
|
def execute(self, prompt, stream, response, conversation):
|
|
yield llm.parts.StreamEvent(
|
|
type="tool_call_name", chunk="echo", tool_call_id="c1"
|
|
)
|
|
yield llm.parts.StreamEvent(
|
|
type="tool_call_args",
|
|
chunk='{"x": 1}',
|
|
tool_call_id="c1",
|
|
)
|
|
|
|
m = ToolCallMock()
|
|
r1 = m.prompt("call echo", tools=[echo])
|
|
r1.text()
|
|
m.enqueue(["follow-up"])
|
|
r2 = r1.reply(
|
|
tool_results=[
|
|
llm.ToolResult(name="echo", output="custom", tool_call_id="c1")
|
|
]
|
|
)
|
|
r2.text()
|
|
assert executed == [] # echo was NOT called
|
|
tool_msg = r2.prompt.messages[2]
|
|
assert tool_msg.parts[0].output == "custom"
|
|
|
|
def test_reply_no_tool_calls_no_tool_message(self, mock_model):
|
|
# reply() on a response without tool calls is unchanged — no
|
|
# tool message gets injected.
|
|
mock_model.enqueue(["a1"])
|
|
mock_model.enqueue(["a2"])
|
|
r1 = mock_model.prompt("q1")
|
|
r1.text()
|
|
r2 = r1.reply()
|
|
r2.text()
|
|
assert r2.prompt.messages == [
|
|
llm.user("q1"),
|
|
llm.assistant("a1"),
|
|
]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_reply_auto_executes_tool_calls(self, async_mock_model):
|
|
# Async reply() is a coroutine; with tool calls present the
|
|
# zero-arg sugar awaits execute_tool_calls() internally.
|
|
from llm.parts import ToolResultPart
|
|
|
|
executed = []
|
|
|
|
async def echo(x: int) -> str:
|
|
executed.append(x)
|
|
return f"echo:{x}"
|
|
|
|
class ToolCallMock(type(async_mock_model)):
|
|
supports_tools = True
|
|
|
|
async def execute(self, prompt, stream, response, conversation):
|
|
response.add_tool_call(
|
|
llm.ToolCall(name="echo", arguments={"x": 7}, tool_call_id="c1")
|
|
)
|
|
yield llm.parts.StreamEvent(
|
|
type="tool_call_name", chunk="echo", tool_call_id="c1"
|
|
)
|
|
yield llm.parts.StreamEvent(
|
|
type="tool_call_args",
|
|
chunk='{"x": 7}',
|
|
tool_call_id="c1",
|
|
)
|
|
|
|
m = ToolCallMock()
|
|
r1 = m.prompt("call echo", tools=[echo])
|
|
await r1.text()
|
|
m.enqueue(["follow-up"])
|
|
r2 = await r1.reply()
|
|
await r2.text()
|
|
assert executed == [7]
|
|
tool_msg = r2.prompt.messages[2]
|
|
assert tool_msg.parts == [
|
|
ToolResultPart(name="echo", output="echo:7", tool_call_id="c1")
|
|
]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_reply_with_tool_results(self, async_mock_model):
|
|
from llm.parts import (
|
|
Message,
|
|
ToolCallPart,
|
|
ToolResultPart,
|
|
)
|
|
|
|
class ToolCallMock(type(async_mock_model)):
|
|
supports_tools = True
|
|
|
|
async def execute(self, prompt, stream, response, conversation):
|
|
yield llm.parts.StreamEvent(
|
|
type="tool_call_name",
|
|
chunk="echo",
|
|
tool_call_id="c1",
|
|
)
|
|
yield llm.parts.StreamEvent(
|
|
type="tool_call_args",
|
|
chunk='{"x": 1}',
|
|
tool_call_id="c1",
|
|
)
|
|
|
|
m = ToolCallMock()
|
|
r1 = m.prompt("call echo")
|
|
await r1.text()
|
|
m.enqueue(["follow-up"])
|
|
r2 = await r1.reply(
|
|
tool_results=[llm.ToolResult(name="echo", output="ok", tool_call_id="c1")]
|
|
)
|
|
await r2.text()
|
|
assert r2.prompt.messages == [
|
|
llm.user("call echo"),
|
|
Message(
|
|
role="assistant",
|
|
parts=[
|
|
ToolCallPart(name="echo", arguments={"x": 1}, tool_call_id="c1")
|
|
],
|
|
),
|
|
Message(
|
|
role="tool",
|
|
parts=[ToolResultPart(name="echo", output="ok", tool_call_id="c1")],
|
|
),
|
|
]
|
|
|
|
|
|
# chain() propagates system across tool-result turns
|
|
|
|
|
|
class TestChainPropagatesSystem:
|
|
"""On a tool-result turn within a chain loop, the Prompt must
|
|
carry forward the original system= and system_fragments= so
|
|
adapters that read prompt.system (OpenAI and other
|
|
stateless-per-turn providers) see it on every call."""
|
|
|
|
def assert_system(self, prompt, *expected):
|
|
assert prompt.messages[0].role == "system"
|
|
for e in expected:
|
|
assert e in prompt.system
|
|
assert e in prompt.messages[0].parts[0].text
|
|
|
|
def test_sync_chain_tool_result_turn_preserves_system(self, mock_model):
|
|
# First turn: fake a tool call so the chain iterates.
|
|
tool_call = llm.ToolCall(tool_call_id="c1", name="tick", arguments={})
|
|
|
|
class ChainMock(type(mock_model)):
|
|
def execute(self, prompt, stream, response, conversation):
|
|
if not self._queue:
|
|
yield "done"
|
|
return
|
|
msgs = self._queue.pop(0)
|
|
for m in msgs:
|
|
yield m
|
|
if not response._tool_calls:
|
|
response.add_tool_call(tool_call)
|
|
|
|
def tick() -> str:
|
|
"Tick"
|
|
return "tock"
|
|
|
|
m = ChainMock()
|
|
m.enqueue(["tool-turn"]) # first response; chain will loop
|
|
m.enqueue(["final"]) # second response, after tool results
|
|
|
|
chain = m.chain("q", system="be brief", tools=[tick])
|
|
list(chain.responses())
|
|
# Second response was the tool-result turn.
|
|
self.assert_system(chain._responses[1].prompt, "be brief")
|
|
|
|
def test_sync_chain_tool_result_turn_preserves_system_fragments(self, mock_model):
|
|
tool_call = llm.ToolCall(tool_call_id="c1", name="tick", arguments={})
|
|
|
|
class ChainMock(type(mock_model)):
|
|
def execute(self, prompt, stream, response, conversation):
|
|
if not self._queue:
|
|
yield "done"
|
|
return
|
|
msgs = self._queue.pop(0)
|
|
for m in msgs:
|
|
yield m
|
|
if not response._tool_calls:
|
|
response.add_tool_call(tool_call)
|
|
|
|
def tick() -> str:
|
|
"Tick"
|
|
return "tock"
|
|
|
|
m = ChainMock()
|
|
m.enqueue(["tool-turn"])
|
|
m.enqueue(["final"])
|
|
|
|
chain = m.chain(
|
|
"q",
|
|
system="inline sys",
|
|
system_fragments=["fragment A", "fragment B"],
|
|
tools=[tick],
|
|
)
|
|
list(chain.responses())
|
|
self.assert_system(
|
|
chain._responses[1].prompt, "inline sys", "fragment A", "fragment B"
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_chain_tool_result_turn_preserves_system(
|
|
self, async_mock_model
|
|
):
|
|
tool_call = llm.ToolCall(tool_call_id="c1", name="tick", arguments={})
|
|
|
|
class AsyncChainMock(type(async_mock_model)):
|
|
supports_tools = True
|
|
|
|
async def execute(self, prompt, stream, response, conversation):
|
|
if not self._queue:
|
|
yield "done"
|
|
return
|
|
msgs = self._queue.pop(0)
|
|
for m in msgs:
|
|
yield m
|
|
if not response._tool_calls:
|
|
response.add_tool_call(tool_call)
|
|
|
|
def tick() -> str:
|
|
"Tick"
|
|
return "tock"
|
|
|
|
m = AsyncChainMock()
|
|
m.enqueue(["tool-turn"])
|
|
m.enqueue(["final"])
|
|
|
|
chain = m.chain("q", system="be brief", tools=[tick])
|
|
responses = []
|
|
async for r in chain.responses():
|
|
responses.append(r)
|
|
self.assert_system(responses[1].prompt, "be brief")
|
|
|
|
def test_chain_includes_system_in_messages(self, mock_model):
|
|
chain = mock_model.chain("q", system="be brief")
|
|
self.assert_system(chain.prompt, "be brief")
|
|
|
|
|
|
# chain() accepts messages= (parity with prompt())
|
|
|
|
|
|
class TestChainMessagesKwarg:
|
|
def test_conversation_chain_accepts_messages(self, mock_model):
|
|
mock_model.enqueue(["ok"])
|
|
conv = mock_model.conversation()
|
|
chain = conv.chain(messages=[llm.user("explicit")])
|
|
chain.text()
|
|
r1 = chain._responses[0]
|
|
assert r1.prompt.messages == [llm.user("explicit")]
|
|
|
|
def test_model_chain_accepts_messages(self, mock_model):
|
|
mock_model.enqueue(["ok"])
|
|
chain = mock_model.chain(messages=[llm.user("explicit")])
|
|
chain.text()
|
|
r1 = chain._responses[0]
|
|
assert r1.prompt.messages == [llm.user("explicit")]
|
|
|
|
def test_chain_messages_is_authoritative_over_prompt_kwarg(self, mock_model):
|
|
"""Parity with prompt(): when both are passed, messages= wins
|
|
and the prompt= string is not folded into the chain."""
|
|
mock_model.enqueue(["ok"])
|
|
chain = mock_model.chain(
|
|
"ignored text",
|
|
messages=[llm.user("explicit")],
|
|
)
|
|
chain.text()
|
|
r1 = chain._responses[0]
|
|
assert r1.prompt.messages == [llm.user("explicit")]
|
|
|
|
def test_chain_with_messages_and_prior_conversation(self, mock_model):
|
|
"""Explicit messages= on chain() replaces history reconstruction;
|
|
the chain starts from that exact list."""
|
|
mock_model.enqueue(["first"])
|
|
mock_model.enqueue(["second"])
|
|
conv = mock_model.conversation()
|
|
r1 = conv.prompt("prior")
|
|
r1.text()
|
|
|
|
# Now start a chain with explicit messages= — prior turn is
|
|
# ignored (consistent with prompt() behavior).
|
|
chain = conv.chain(messages=[llm.user("fresh start")])
|
|
chain.text()
|
|
first_chain_response = chain._responses[0]
|
|
assert first_chain_response.prompt.messages == [llm.user("fresh start")]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_conversation_chain_accepts_messages(self, async_mock_model):
|
|
async_mock_model.enqueue(["ok"])
|
|
conv = async_mock_model.conversation()
|
|
chain = conv.chain(messages=[llm.user("explicit")])
|
|
await chain.text()
|
|
r1 = chain._responses[0]
|
|
assert r1.prompt.messages == [llm.user("explicit")]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_model_chain_accepts_messages(self, async_mock_model):
|
|
async_mock_model.enqueue(["ok"])
|
|
chain = async_mock_model.chain(messages=[llm.user("explicit")])
|
|
await chain.text()
|
|
r1 = chain._responses[0]
|
|
assert r1.prompt.messages == [llm.user("explicit")]
|
|
|
|
|
|
# Response.to_dict / Response.from_dict
|
|
|
|
|
|
class TestResponseToDictFromDict:
|
|
def test_to_dict_captures_chain_and_output(self, mock_model):
|
|
mock_model.enqueue(["hello"])
|
|
r = mock_model.prompt("hi")
|
|
r.text()
|
|
|
|
d = r.to_dict()
|
|
assert d["model"] == "mock"
|
|
assert d["prompt"]["messages"] == [llm.user("hi").to_dict()]
|
|
assert d["messages"] == [llm.assistant("hello").to_dict()]
|
|
|
|
def test_from_dict_rehydrates_with_messages(self, mock_model):
|
|
mock_model.enqueue(["hello"])
|
|
r = mock_model.prompt("hi")
|
|
r.text()
|
|
payload = json.dumps(r.to_dict())
|
|
|
|
restored = llm.Response.from_dict(json.loads(payload))
|
|
assert restored._done
|
|
assert restored.text() == "hello"
|
|
assert restored.messages() == [llm.assistant("hello")]
|
|
assert restored.prompt.messages == [llm.user("hi")]
|
|
|
|
def test_from_dict_then_reply_continues_conversation(self, mock_model):
|
|
mock_model.enqueue(["a1"])
|
|
mock_model.enqueue(["a2"])
|
|
r1 = mock_model.prompt("q1")
|
|
r1.text()
|
|
|
|
# Serialize across the process boundary
|
|
payload = json.dumps(r1.to_dict())
|
|
restored = llm.Response.from_dict(json.loads(payload))
|
|
|
|
# Continue from the restored response
|
|
r2 = restored.reply("q2")
|
|
r2.text()
|
|
assert r2.prompt.messages == [
|
|
llm.user("q1"),
|
|
llm.assistant("a1"),
|
|
llm.user("q2"),
|
|
]
|
|
|
|
def test_to_dict_preserves_reasoning_and_signatures(self, mock_model):
|
|
mock_model.enqueue(
|
|
[
|
|
llm.parts.StreamEvent(
|
|
type="reasoning",
|
|
chunk="thinking...",
|
|
part_index=0,
|
|
provider_metadata={"anthropic": {"signature": "sig-abc"}},
|
|
),
|
|
llm.parts.StreamEvent(type="text", chunk="answer", part_index=1),
|
|
]
|
|
)
|
|
r = mock_model.prompt("q")
|
|
r.text()
|
|
|
|
payload = json.dumps(r.to_dict())
|
|
restored = llm.Response.from_dict(json.loads(payload))
|
|
|
|
msgs = restored.messages()
|
|
assert msgs[0].role == "assistant"
|
|
assert isinstance(msgs[0].parts[0], llm.parts.ReasoningPart)
|
|
assert msgs[0].parts[0].text == "thinking..."
|
|
assert msgs[0].parts[0].provider_metadata == {
|
|
"anthropic": {"signature": "sig-abc"}
|
|
}
|
|
|
|
def test_from_dict_reply_includes_prior_reasoning_in_chain(self, mock_model):
|
|
"""a reply() after from_dict() sends the thinking signature
|
|
back to the model for multi-turn extended thinking."""
|
|
mock_model.enqueue(
|
|
[
|
|
llm.parts.StreamEvent(
|
|
type="reasoning",
|
|
chunk="thinking...",
|
|
part_index=0,
|
|
provider_metadata={"anthropic": {"signature": "sig-xyz"}},
|
|
),
|
|
llm.parts.StreamEvent(type="text", chunk="answer", part_index=1),
|
|
]
|
|
)
|
|
mock_model.enqueue(["a2"])
|
|
r1 = mock_model.prompt("q1")
|
|
r1.text()
|
|
|
|
payload = json.dumps(r1.to_dict())
|
|
restored = llm.Response.from_dict(json.loads(payload))
|
|
r2 = restored.reply("q2")
|
|
r2.text()
|
|
|
|
# The signature must be in the chain sent to the model.
|
|
chain = r2.prompt.messages
|
|
reasoning_parts = [
|
|
p for m in chain for p in m.parts if isinstance(p, llm.parts.ReasoningPart)
|
|
]
|
|
assert len(reasoning_parts) == 1
|
|
assert reasoning_parts[0].provider_metadata == {
|
|
"anthropic": {"signature": "sig-xyz"}
|
|
}
|
|
|
|
def test_to_dict_captures_options(self, mock_model):
|
|
mock_model.enqueue(["ok"])
|
|
r = mock_model.prompt("hi", max_tokens=42)
|
|
r.text()
|
|
|
|
d = r.to_dict()
|
|
assert d["prompt"]["options"] == {"max_tokens": 42}
|
|
|
|
def test_from_dict_options_restored(self, mock_model):
|
|
mock_model.enqueue(["ok"])
|
|
r = mock_model.prompt("hi", max_tokens=42)
|
|
r.text()
|
|
|
|
payload = json.dumps(r.to_dict())
|
|
restored = llm.Response.from_dict(json.loads(payload))
|
|
assert restored.prompt.options.max_tokens == 42
|
|
|
|
def test_message_from_dict_static_method_unchanged(self):
|
|
m = llm.assistant("hi")
|
|
assert llm.Message.from_dict(m.to_dict()) == m
|
|
|
|
|
|
class TestChainResponseStreamEvents:
|
|
def test_sync_chain_stream_events_yields_text_when_no_tools(self, mock_model):
|
|
# Chain with no tool calls is a single-response chain — its
|
|
# stream_events should concatenate from each underlying response.
|
|
mock_model.enqueue(
|
|
[llm.parts.StreamEvent(type="text", chunk="done", part_index=0)]
|
|
)
|
|
chain = mock_model.conversation().chain("q")
|
|
events = list(chain.stream_events())
|
|
assert [e.type for e in events] == ["text"]
|
|
assert [e.chunk for e in events] == ["done"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_chain_astream_events_yields(self, async_mock_model):
|
|
async_mock_model.enqueue(
|
|
[llm.parts.StreamEvent(type="text", chunk="done", part_index=0)]
|
|
)
|
|
chain = async_mock_model.conversation().chain("q")
|
|
events = []
|
|
async for event in chain.astream_events():
|
|
events.append(event)
|
|
assert [e.type for e in events] == ["text"]
|
|
|
|
|
|
# Client-side serialization round-trip
|
|
#
|
|
# A library user can persist a conversation by serializing response.messages
|
|
# to JSON and later re-inflate it as messages=[...] on a follow-up prompt.
|
|
# No SQLite involvement.
|
|
|
|
|
|
class TestClientSerializationRoundTrip:
|
|
def test_response_messages_json_roundtrip(self, mock_model):
|
|
mock_model.enqueue(["hello there"])
|
|
r = mock_model.prompt("hi")
|
|
r.text()
|
|
|
|
# Serialize via Message.to_dict / json.dumps
|
|
payload = json.dumps([m.to_dict() for m in r.messages()])
|
|
# Deserialize — no LLM state needed beyond the types.
|
|
restored = [llm.Message.from_dict(d) for d in json.loads(payload)]
|
|
|
|
assert restored == r.messages()
|
|
|
|
def test_rebuilt_messages_reach_plugin_via_prompt(self, mock_model):
|
|
"""Round-trip: serialize messages from turn 1, re-inflate, send
|
|
as messages= to turn 2. The plugin sees the full chain."""
|
|
# Turn 1
|
|
mock_model.enqueue(["turn 1 answer"])
|
|
r1 = mock_model.prompt("turn 1 question")
|
|
r1.text()
|
|
|
|
# Persist everything the client cares about.
|
|
history = [llm.user("turn 1 question").to_dict()] + [
|
|
m.to_dict() for m in r1.messages()
|
|
]
|
|
payload = json.dumps(history)
|
|
|
|
# Later — rebuild from the wire form and continue.
|
|
rebuilt = [llm.Message.from_dict(d) for d in json.loads(payload)]
|
|
mock_model.enqueue(["turn 2 answer"])
|
|
r2 = mock_model.prompt(messages=rebuilt + [llm.user("turn 2 question")])
|
|
r2.text()
|
|
|
|
# The plugin saw the full structured history on prompt.messages.
|
|
assert r2.prompt.messages == rebuilt + [llm.user("turn 2 question")]
|
|
assert r2.messages() == [llm.assistant("turn 2 answer")]
|
|
|
|
def test_roundtrip_preserves_tool_calls_and_results(self, mock_model):
|
|
"""Assistant messages with tool calls + subsequent tool role
|
|
messages survive json round-trip intact."""
|
|
messages = [
|
|
llm.user("what's the weather?"),
|
|
llm.assistant(
|
|
"let me check",
|
|
llm.parts.ToolCallPart(
|
|
name="get_weather",
|
|
arguments={"city": "Paris"},
|
|
tool_call_id="c1",
|
|
),
|
|
),
|
|
llm.tool_message(
|
|
llm.parts.ToolResultPart(
|
|
name="get_weather",
|
|
output="sunny",
|
|
tool_call_id="c1",
|
|
)
|
|
),
|
|
]
|
|
payload = json.dumps([m.to_dict() for m in messages])
|
|
restored = [llm.Message.from_dict(d) for d in json.loads(payload)]
|
|
assert restored == messages
|
|
|
|
def test_roundtrip_preserves_redacted_reasoning(self, mock_model):
|
|
"""The redacted=True marker on a ReasoningPart survives
|
|
round-trip — UIs use it to show that opaque reasoning happened
|
|
in this turn (the actual token count lives on response usage)."""
|
|
msg = llm.Message(
|
|
role="assistant",
|
|
parts=[
|
|
llm.parts.ReasoningPart(text="", redacted=True),
|
|
llm.parts.TextPart(text="result"),
|
|
],
|
|
)
|
|
restored = llm.Message.from_dict(json.loads(json.dumps(msg.to_dict())))
|
|
assert restored == msg
|
|
|
|
def test_roundtrip_preserves_provider_metadata(self, mock_model):
|
|
msg = llm.Message(
|
|
role="assistant",
|
|
parts=[
|
|
llm.parts.ReasoningPart(
|
|
text="thinking",
|
|
provider_metadata={"anthropic": {"signature": "abc"}},
|
|
),
|
|
llm.parts.TextPart(text="answer"),
|
|
],
|
|
)
|
|
restored = llm.Message.from_dict(json.loads(json.dumps(msg.to_dict())))
|
|
assert restored == msg
|