mirror of
https://github.com/hwchase17/langchain.git
synced 2026-07-12 02:55:21 +00:00
test(core): add stream lifecycle validator and provider coverage
New `langchain_tests.utils.stream_lifecycle.assert_valid_event_stream` helper enforces the protocol contract on any event stream: - single message-start / message-finish envelope - blocks do not interleave (each block finishes before the next starts) - sequential uint wire indices from 0 - accumulated deltas match the finish payload for deltaable types Applied at three levels: - core/test_compat_bridge: provider-style emission patterns exercised directly through chunks_to_events / message_to_events (openai chat completions int indices, openai responses/v1 string identifiers, anthropic-style per-chunk int indices, inline image, invalid tool call, empty stream) - openai partner: validator applied to stream_v2 against the existing responses-api mock and to a new chat-completions stream_v2 test - anthropic partner: new mock stream of RawMessageStartEvent + RawContentBlock* events threaded through _stream via `_create` patch; covers thinking + text + tool_use lifecycle with tool-use stop_reason Enabling thinking on the anthropic test flips coerce_content_to_string off so every block carries a proper integer index — the structured path the bridge actually exercises. Default-mode (no tools / thinking / docs) coerces text to a plain string and strips per-chunk indices; the bridge handles that branch by collapsing to positional-0 and it is a known separate code path, intentionally not covered here.
This commit is contained in:
@@ -3,6 +3,7 @@
|
||||
from typing import TYPE_CHECKING, cast
|
||||
|
||||
import pytest
|
||||
from langchain_tests.utils.stream_lifecycle import assert_valid_event_stream
|
||||
|
||||
from langchain_core.language_models._compat_bridge import (
|
||||
CompatBlock,
|
||||
@@ -677,3 +678,257 @@ async def test_amessage_to_events_matches_sync() -> None:
|
||||
sync_events = list(message_to_events(msg))
|
||||
async_events = [e async for e in amessage_to_events(msg)]
|
||||
assert async_events == sync_events
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Lifecycle validator: provider-style emission patterns
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _aimsg_chunk(blocks: list[CompatBlock], msg_id: str = "m") -> ChatGenerationChunk:
|
||||
"""Wrap a list of content blocks into a ChatGenerationChunk.
|
||||
|
||||
Matches what a provider's `_stream` would yield per SSE event.
|
||||
"""
|
||||
return ChatGenerationChunk(message=AIMessageChunk(content=blocks, id=msg_id))
|
||||
|
||||
|
||||
def test_lifecycle_validator_openai_chat_completions_style() -> None:
|
||||
"""Text + streaming tool call with int indices, all at index 0/1.
|
||||
|
||||
Mirrors OpenAI chat-completions API where each delta stays at the
|
||||
same integer index and a new tool call bumps the index.
|
||||
"""
|
||||
chunks = [
|
||||
_aimsg_chunk([{"type": "text", "text": "Hello", "index": 0}]),
|
||||
_aimsg_chunk([{"type": "text", "text": " there", "index": 0}]),
|
||||
]
|
||||
# Tool-call chunks go via the tool_call_chunks channel, not content.
|
||||
chunks.extend(
|
||||
[
|
||||
ChatGenerationChunk(
|
||||
message=AIMessageChunk(
|
||||
content="",
|
||||
id="m",
|
||||
tool_call_chunks=[
|
||||
{
|
||||
"type": "tool_call_chunk",
|
||||
"index": 1,
|
||||
"id": "tc1",
|
||||
"name": "lookup",
|
||||
"args": '{"q":',
|
||||
}
|
||||
],
|
||||
)
|
||||
),
|
||||
ChatGenerationChunk(
|
||||
message=AIMessageChunk(
|
||||
content="",
|
||||
id="m",
|
||||
tool_call_chunks=[
|
||||
{
|
||||
"type": "tool_call_chunk",
|
||||
"index": 1,
|
||||
"id": None,
|
||||
"name": None,
|
||||
"args": ' "pie"}',
|
||||
}
|
||||
],
|
||||
)
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
events = list(chunks_to_events(iter(chunks), message_id="m"))
|
||||
assert_valid_event_stream(events)
|
||||
|
||||
finishes = [e for e in events if e["event"] == "content-block-finish"]
|
||||
types = [f["content_block"]["type"] for f in finishes]
|
||||
assert types == ["text", "tool_call"]
|
||||
assert finishes[0]["content_block"]["text"] == "Hello there"
|
||||
assert finishes[1]["content_block"]["args"] == {"q": "pie"}
|
||||
|
||||
|
||||
def test_lifecycle_validator_openai_responses_style() -> None:
|
||||
"""Reasoning → text → reasoning → text with string block identifiers.
|
||||
|
||||
Mirrors OpenAI `responses/v1` output_version where each distinct
|
||||
block has a string index like `lc_rs_305f30`.
|
||||
"""
|
||||
chunks = [
|
||||
_aimsg_chunk([{"type": "reasoning", "reasoning": "hmm", "index": "rs_a"}]),
|
||||
_aimsg_chunk([{"type": "reasoning", "reasoning": " first", "index": "rs_a"}]),
|
||||
_aimsg_chunk([{"type": "text", "text": "Answer: ", "index": "txt_a"}]),
|
||||
_aimsg_chunk([{"type": "text", "text": "42", "index": "txt_a"}]),
|
||||
_aimsg_chunk([{"type": "reasoning", "reasoning": "actually", "index": "rs_b"}]),
|
||||
_aimsg_chunk([{"type": "text", "text": "42!", "index": "txt_b"}]),
|
||||
]
|
||||
|
||||
events = list(chunks_to_events(iter(chunks), message_id="m"))
|
||||
assert_valid_event_stream(events)
|
||||
|
||||
starts = [e for e in events if e["event"] == "content-block-start"]
|
||||
finishes = [e for e in events if e["event"] == "content-block-finish"]
|
||||
# Four distinct blocks: reasoning, text, reasoning, text
|
||||
assert [s["content_block"]["type"] for s in starts] == [
|
||||
"reasoning",
|
||||
"text",
|
||||
"reasoning",
|
||||
"text",
|
||||
]
|
||||
assert [s["index"] for s in starts] == [0, 1, 2, 3]
|
||||
assert [f["index"] for f in finishes] == [0, 1, 2, 3]
|
||||
assert finishes[0]["content_block"]["reasoning"] == "hmm first"
|
||||
assert finishes[1]["content_block"]["text"] == "Answer: 42"
|
||||
assert finishes[2]["content_block"]["reasoning"] == "actually"
|
||||
assert finishes[3]["content_block"]["text"] == "42!"
|
||||
|
||||
|
||||
def test_lifecycle_validator_anthropic_style_text_and_thinking() -> None:
|
||||
"""Interleaved text and thinking blocks with int indices.
|
||||
|
||||
Mirrors Anthropic's per-event structure: one block per chunk, each
|
||||
labeled with an int `index` from Anthropic's content_block_start /
|
||||
content_block_delta events.
|
||||
"""
|
||||
chunks = [
|
||||
_aimsg_chunk([{"type": "reasoning", "reasoning": "Let me think", "index": 0}]),
|
||||
_aimsg_chunk([{"type": "reasoning", "reasoning": " more", "index": 0}]),
|
||||
_aimsg_chunk([{"type": "text", "text": "The answer is", "index": 1}]),
|
||||
_aimsg_chunk([{"type": "text", "text": " 42.", "index": 1}]),
|
||||
]
|
||||
|
||||
events = list(chunks_to_events(iter(chunks), message_id="m"))
|
||||
assert_valid_event_stream(events)
|
||||
|
||||
finishes = [e for e in events if e["event"] == "content-block-finish"]
|
||||
assert [f["content_block"]["type"] for f in finishes] == ["reasoning", "text"]
|
||||
assert finishes[0]["content_block"]["reasoning"] == "Let me think more"
|
||||
assert finishes[1]["content_block"]["text"] == "The answer is 42."
|
||||
|
||||
|
||||
def test_lifecycle_validator_anthropic_style_tool_use_after_text() -> None:
|
||||
"""Text then tool_use (tool_call_chunk) — Anthropic tool-calling pattern."""
|
||||
chunks = [
|
||||
_aimsg_chunk([{"type": "text", "text": "Looking up...", "index": 0}]),
|
||||
ChatGenerationChunk(
|
||||
message=AIMessageChunk(
|
||||
content=[],
|
||||
id="m",
|
||||
tool_call_chunks=[
|
||||
{
|
||||
"type": "tool_call_chunk",
|
||||
"index": 1,
|
||||
"id": "toolu_1",
|
||||
"name": "search",
|
||||
"args": "",
|
||||
}
|
||||
],
|
||||
)
|
||||
),
|
||||
ChatGenerationChunk(
|
||||
message=AIMessageChunk(
|
||||
content=[],
|
||||
id="m",
|
||||
tool_call_chunks=[
|
||||
{
|
||||
"type": "tool_call_chunk",
|
||||
"index": 1,
|
||||
"id": None,
|
||||
"name": None,
|
||||
"args": '{"query": "42"}',
|
||||
}
|
||||
],
|
||||
)
|
||||
),
|
||||
]
|
||||
|
||||
events = list(chunks_to_events(iter(chunks), message_id="m"))
|
||||
assert_valid_event_stream(events)
|
||||
|
||||
finishes = [e for e in events if e["event"] == "content-block-finish"]
|
||||
assert [f["content_block"]["type"] for f in finishes] == ["text", "tool_call"]
|
||||
assert finishes[1]["content_block"]["args"] == {"query": "42"}
|
||||
assert finishes[1]["content_block"]["id"] == "toolu_1"
|
||||
|
||||
|
||||
def test_lifecycle_validator_inline_image_block() -> None:
|
||||
"""A self-contained image block gets start + finish with no delta."""
|
||||
chunks = [
|
||||
_aimsg_chunk(
|
||||
[
|
||||
{
|
||||
"type": "image",
|
||||
"id": "img1",
|
||||
"mime_type": "image/png",
|
||||
"data": "AAAA",
|
||||
"index": 0,
|
||||
}
|
||||
]
|
||||
),
|
||||
]
|
||||
events = list(chunks_to_events(iter(chunks), message_id="m"))
|
||||
assert_valid_event_stream(events)
|
||||
|
||||
starts = [e for e in events if e["event"] == "content-block-start"]
|
||||
deltas = [e for e in events if e["event"] == "content-block-delta"]
|
||||
finishes = [e for e in events if e["event"] == "content-block-finish"]
|
||||
assert [s["content_block"]["type"] for s in starts] == ["image"]
|
||||
# Self-contained block: no delta, and start has heavy fields stripped.
|
||||
assert deltas == []
|
||||
assert "data" not in starts[0]["content_block"]
|
||||
assert finishes[0]["content_block"]["data"] == "AAAA"
|
||||
|
||||
|
||||
def test_lifecycle_validator_invalid_tool_call_args() -> None:
|
||||
"""Malformed JSON args finalize to invalid_tool_call; lifecycle still valid."""
|
||||
chunks = [
|
||||
ChatGenerationChunk(
|
||||
message=AIMessageChunk(
|
||||
content="",
|
||||
id="m",
|
||||
tool_call_chunks=[
|
||||
{
|
||||
"type": "tool_call_chunk",
|
||||
"index": 0,
|
||||
"id": "bad1",
|
||||
"name": "noop",
|
||||
"args": "not json",
|
||||
}
|
||||
],
|
||||
)
|
||||
),
|
||||
]
|
||||
events = list(chunks_to_events(iter(chunks), message_id="m"))
|
||||
assert_valid_event_stream(events)
|
||||
|
||||
finishes = [e for e in events if e["event"] == "content-block-finish"]
|
||||
assert len(finishes) == 1
|
||||
assert finishes[0]["content_block"]["type"] == "invalid_tool_call"
|
||||
|
||||
|
||||
def test_lifecycle_validator_empty_stream() -> None:
|
||||
"""An empty chunk iterator produces no events (and still validates)."""
|
||||
assert_valid_event_stream(list(chunks_to_events(iter([]))))
|
||||
|
||||
|
||||
def test_lifecycle_validator_message_to_events_roundtrip() -> None:
|
||||
"""`message_to_events` also produces spec-conformant lifecycles."""
|
||||
msg = AIMessage(
|
||||
content=[
|
||||
{"type": "reasoning", "reasoning": "think"},
|
||||
{"type": "text", "text": "answer"},
|
||||
{
|
||||
"type": "image",
|
||||
"id": "img1",
|
||||
"mime_type": "image/png",
|
||||
"data": "X" * 256,
|
||||
},
|
||||
],
|
||||
id="msg-1",
|
||||
tool_calls=[
|
||||
{"id": "t1", "name": "search", "args": {"q": "pie"}, "type": "tool_call"},
|
||||
],
|
||||
)
|
||||
events = list(message_to_events(msg))
|
||||
assert_valid_event_stream(events)
|
||||
|
||||
@@ -2843,3 +2843,155 @@ def test_no_task_budget_no_beta() -> None:
|
||||
betas = payload.get("betas")
|
||||
if betas:
|
||||
assert "task-budgets-2026-03-13" not in betas
|
||||
|
||||
|
||||
def test_anthropic_stream_v2_lifecycle() -> None:
|
||||
"""Validate lifecycle events across a thinking + text + tool_use stream.
|
||||
|
||||
Anthropic emits raw `content_block_start` / `content_block_delta` /
|
||||
`content_block_stop` events with integer `index` fields, interleaved
|
||||
with `message_start` and `message_delta`. This test threads a
|
||||
realistic event sequence through `_stream` via a mocked raw client
|
||||
and asserts that `stream_v2` produces a spec-conformant event
|
||||
stream: paired start/finish per block, no interleaving, sequential
|
||||
`uint` wire indices.
|
||||
"""
|
||||
from unittest.mock import patch
|
||||
|
||||
from anthropic.types import (
|
||||
InputJSONDelta,
|
||||
RawContentBlockDeltaEvent,
|
||||
RawContentBlockStartEvent,
|
||||
RawContentBlockStopEvent,
|
||||
RawMessageDeltaEvent,
|
||||
RawMessageStartEvent,
|
||||
RawMessageStopEvent,
|
||||
TextDelta,
|
||||
ThinkingBlock,
|
||||
ThinkingDelta,
|
||||
ToolUseBlock,
|
||||
)
|
||||
from anthropic.types.raw_message_delta_event import Delta as RawMessageDelta
|
||||
from anthropic.types.raw_message_delta_event import (
|
||||
MessageDeltaUsage as RawMessageDeltaUsage,
|
||||
)
|
||||
from langchain_tests.utils.stream_lifecycle import assert_valid_event_stream
|
||||
|
||||
msg = Message(
|
||||
id="msg_1",
|
||||
content=[],
|
||||
model=MODEL_NAME,
|
||||
role="assistant",
|
||||
stop_reason=None,
|
||||
stop_sequence=None,
|
||||
usage=Usage(input_tokens=10, output_tokens=0),
|
||||
type="message",
|
||||
)
|
||||
|
||||
events = [
|
||||
RawMessageStartEvent(message=msg, type="message_start"),
|
||||
# thinking block (index=0)
|
||||
RawContentBlockStartEvent(
|
||||
content_block=ThinkingBlock(signature="", thinking="", type="thinking"),
|
||||
index=0,
|
||||
type="content_block_start",
|
||||
),
|
||||
RawContentBlockDeltaEvent(
|
||||
delta=ThinkingDelta(thinking="Let me ", type="thinking_delta"),
|
||||
index=0,
|
||||
type="content_block_delta",
|
||||
),
|
||||
RawContentBlockDeltaEvent(
|
||||
delta=ThinkingDelta(thinking="think.", type="thinking_delta"),
|
||||
index=0,
|
||||
type="content_block_delta",
|
||||
),
|
||||
RawContentBlockStopEvent(index=0, type="content_block_stop"),
|
||||
# text block (index=1)
|
||||
RawContentBlockStartEvent(
|
||||
content_block=TextBlock(text="", type="text"),
|
||||
index=1,
|
||||
type="content_block_start",
|
||||
),
|
||||
RawContentBlockDeltaEvent(
|
||||
delta=TextDelta(text="The answer ", type="text_delta"),
|
||||
index=1,
|
||||
type="content_block_delta",
|
||||
),
|
||||
RawContentBlockDeltaEvent(
|
||||
delta=TextDelta(text="is 42.", type="text_delta"),
|
||||
index=1,
|
||||
type="content_block_delta",
|
||||
),
|
||||
RawContentBlockStopEvent(index=1, type="content_block_stop"),
|
||||
# tool_use block (index=2)
|
||||
RawContentBlockStartEvent(
|
||||
content_block=ToolUseBlock(
|
||||
id="toolu_1",
|
||||
input={},
|
||||
name="search",
|
||||
type="tool_use",
|
||||
),
|
||||
index=2,
|
||||
type="content_block_start",
|
||||
),
|
||||
RawContentBlockDeltaEvent(
|
||||
delta=InputJSONDelta(partial_json='{"q":', type="input_json_delta"),
|
||||
index=2,
|
||||
type="content_block_delta",
|
||||
),
|
||||
RawContentBlockDeltaEvent(
|
||||
delta=InputJSONDelta(partial_json=' "weather"}', type="input_json_delta"),
|
||||
index=2,
|
||||
type="content_block_delta",
|
||||
),
|
||||
RawContentBlockStopEvent(index=2, type="content_block_stop"),
|
||||
# message_delta with final usage and stop_reason
|
||||
RawMessageDeltaEvent(
|
||||
delta=RawMessageDelta(stop_reason="tool_use", stop_sequence=None),
|
||||
type="message_delta",
|
||||
usage=RawMessageDeltaUsage(
|
||||
output_tokens=50,
|
||||
input_tokens=10,
|
||||
cache_read_input_tokens=0,
|
||||
cache_creation_input_tokens=0,
|
||||
),
|
||||
),
|
||||
RawMessageStopEvent(type="message_stop"),
|
||||
]
|
||||
|
||||
# Enable thinking so `coerce_content_to_string=False` in `_stream`,
|
||||
# which gives every content block an integer `index` field — the
|
||||
# structured path the protocol bridge actually exercises. Default
|
||||
# (no tools / thinking / documents) coerces text to a plain string,
|
||||
# which strips indices and is a separate code path not covered here.
|
||||
llm = ChatAnthropic(
|
||||
model=MODEL_NAME,
|
||||
thinking={"type": "enabled", "budget_tokens": 1024},
|
||||
)
|
||||
|
||||
def mock_create(_payload: Any) -> list:
|
||||
return events
|
||||
|
||||
with patch.object(llm, "_create", mock_create):
|
||||
stream_events = list(llm.stream_v2("Test query"))
|
||||
|
||||
assert_valid_event_stream(stream_events)
|
||||
|
||||
finishes = [e for e in stream_events if e["event"] == "content-block-finish"]
|
||||
types = [f["content_block"]["type"] for f in finishes]
|
||||
assert types == ["reasoning", "text", "tool_call"]
|
||||
|
||||
wire_indices = [f["index"] for f in finishes]
|
||||
assert wire_indices == [0, 1, 2]
|
||||
|
||||
# Content accumulation reaches content-block-finish intact.
|
||||
assert finishes[0]["content_block"]["reasoning"] == "Let me think."
|
||||
assert finishes[1]["content_block"]["text"] == "The answer is 42."
|
||||
assert finishes[2]["content_block"]["args"] == {"q": "weather"}
|
||||
assert finishes[2]["content_block"]["name"] == "search"
|
||||
|
||||
# message-finish carries the tool_use stop reason.
|
||||
message_finish = stream_events[-1]
|
||||
assert message_finish["event"] == "message-finish"
|
||||
assert message_finish["reason"] == "tool_use"
|
||||
|
||||
@@ -616,6 +616,28 @@ def test_openai_stream(mock_openai_completion: list) -> None:
|
||||
assert "stream_options" not in call_kwargs[-1]
|
||||
|
||||
|
||||
def test_openai_stream_v2_lifecycle(mock_openai_completion: list) -> None:
|
||||
"""`stream_v2` on chat completions emits a spec-conformant lifecycle."""
|
||||
from langchain_tests.utils.stream_lifecycle import assert_valid_event_stream
|
||||
|
||||
llm = ChatOpenAI(model="gpt-4o")
|
||||
mock_client = MagicMock()
|
||||
|
||||
def mock_create(*args: Any, **kwargs: Any) -> MockSyncContextManager:
|
||||
return MockSyncContextManager(mock_openai_completion)
|
||||
|
||||
mock_client.create = mock_create
|
||||
with patch.object(llm, "client", mock_client):
|
||||
events = list(llm.stream_v2("你的名字叫什么?只回答名字"))
|
||||
|
||||
assert_valid_event_stream(events)
|
||||
# At minimum, a text block with the accumulated answer.
|
||||
finishes = [e for e in events if e["event"] == "content-block-finish"]
|
||||
assert len(finishes) >= 1
|
||||
text_finishes = [f for f in finishes if f["content_block"]["type"] == "text"]
|
||||
assert len(text_finishes) == 1
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_completion() -> dict:
|
||||
return {
|
||||
|
||||
@@ -6,6 +6,7 @@ from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessageChunk, BaseMessageChunk
|
||||
from langchain_tests.utils.stream_lifecycle import assert_valid_event_stream
|
||||
from openai.types.responses import (
|
||||
ResponseCompletedEvent,
|
||||
ResponseContentPartAddedEvent,
|
||||
@@ -769,9 +770,7 @@ def test_responses_stream_v2_emits_reasoning_lifecycle() -> None:
|
||||
(`content-block-start` / `content-block-delta` / `content-block-finish`)
|
||||
for every reasoning block observed on the wire, not just text blocks.
|
||||
"""
|
||||
llm = ChatOpenAI(
|
||||
model="o4-mini", use_responses_api=True, output_version="v1"
|
||||
)
|
||||
llm = ChatOpenAI(model="o4-mini", use_responses_api=True, output_version="v1")
|
||||
mock_client = MagicMock()
|
||||
|
||||
def mock_create(*args: Any, **kwargs: Any) -> MockSyncContextManager:
|
||||
@@ -782,6 +781,8 @@ def test_responses_stream_v2_emits_reasoning_lifecycle() -> None:
|
||||
with patch.object(llm, "root_client", mock_client):
|
||||
events = list(llm.stream_v2("test"))
|
||||
|
||||
assert_valid_event_stream(events)
|
||||
|
||||
reasoning_starts = [
|
||||
e
|
||||
for e in events
|
||||
@@ -801,10 +802,14 @@ def test_responses_stream_v2_emits_reasoning_lifecycle() -> None:
|
||||
assert len(reasoning_starts) == 4, (
|
||||
f"expected 4 reasoning start events, got {len(reasoning_starts)}"
|
||||
)
|
||||
all_finish_types = [
|
||||
e["content_block"]["type"]
|
||||
for e in events
|
||||
if e["event"] == "content-block-finish"
|
||||
]
|
||||
assert len(reasoning_finishes) == 4, (
|
||||
f"expected 4 reasoning finish events, got {len(reasoning_finishes)}: "
|
||||
f"all finish events = "
|
||||
f"{[e['content_block']['type'] for e in events if e['event'] == 'content-block-finish']}"
|
||||
f"all finish events = {all_finish_types}"
|
||||
)
|
||||
|
||||
# Finish events must carry the accumulated reasoning text.
|
||||
|
||||
202
libs/standard-tests/langchain_tests/utils/stream_lifecycle.py
Normal file
202
libs/standard-tests/langchain_tests/utils/stream_lifecycle.py
Normal file
@@ -0,0 +1,202 @@
|
||||
"""Validator for LangChain content-block protocol event streams.
|
||||
|
||||
Checks that an event stream emitted by a chat model (via `stream_v2`,
|
||||
or by the compat bridge's `chunks_to_events` / `message_to_events`)
|
||||
conforms to the protocol lifecycle rules:
|
||||
|
||||
- `message-start` opens and `message-finish` closes the stream.
|
||||
- Content blocks do not interleave: each block runs
|
||||
`content-block-start` → optional `content-block-delta`s →
|
||||
`content-block-finish` before the next block begins.
|
||||
- Wire indices on content-block events are sequential `uint` values
|
||||
starting at 0.
|
||||
- For deltaable block types (`text`, `reasoning`, `tool_call_chunk`,
|
||||
`server_tool_call_chunk`), accumulated delta content matches the
|
||||
final payload delivered on `content-block-finish`.
|
||||
|
||||
The validator accepts any iterable of protocol event dicts. It raises
|
||||
`AssertionError` on the first violation with a descriptive message.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterable
|
||||
|
||||
|
||||
_DELTAABLE_TYPES = frozenset(
|
||||
{"text", "reasoning", "tool_call_chunk", "server_tool_call_chunk"}
|
||||
)
|
||||
|
||||
|
||||
def assert_valid_event_stream(events: Iterable[Any]) -> None:
|
||||
"""Assert that a stream of protocol events obeys the lifecycle contract.
|
||||
|
||||
Args:
|
||||
events: Iterable of protocol event dicts (as yielded by
|
||||
`stream_v2` or `chunks_to_events`).
|
||||
|
||||
Raises:
|
||||
AssertionError: On the first lifecycle violation found. The
|
||||
message identifies the event index and the specific rule
|
||||
that was broken.
|
||||
"""
|
||||
event_list = list(events)
|
||||
if not event_list:
|
||||
return
|
||||
|
||||
first = event_list[0]
|
||||
assert first["event"] == "message-start", (
|
||||
f"first event must be `message-start`, got {first['event']!r}"
|
||||
)
|
||||
message_start_positions = [
|
||||
i for i, e in enumerate(event_list) if e["event"] == "message-start"
|
||||
]
|
||||
assert message_start_positions == [0], (
|
||||
f"expected exactly one `message-start` at position 0, "
|
||||
f"got positions {message_start_positions}"
|
||||
)
|
||||
|
||||
message_finish_positions = [
|
||||
i for i, e in enumerate(event_list) if e["event"] == "message-finish"
|
||||
]
|
||||
assert len(message_finish_positions) <= 1, (
|
||||
f"expected at most one `message-finish`, got {len(message_finish_positions)}"
|
||||
)
|
||||
if message_finish_positions:
|
||||
assert message_finish_positions[0] == len(event_list) - 1, (
|
||||
"`message-finish` must be the final event"
|
||||
)
|
||||
|
||||
open_idx: int | None = None
|
||||
expected_next_idx = 0
|
||||
start_events: dict[int, dict[str, Any]] = {}
|
||||
finish_events: dict[int, dict[str, Any]] = {}
|
||||
delta_accum: dict[int, dict[str, Any]] = {}
|
||||
|
||||
for i, event in enumerate(event_list):
|
||||
ev = event["event"]
|
||||
if ev == "message-start":
|
||||
assert i == 0, f"duplicate `message-start` at event {i}"
|
||||
continue
|
||||
if ev == "message-finish":
|
||||
assert open_idx is None, (
|
||||
f"`message-finish` while block {open_idx} still open (event {i})"
|
||||
)
|
||||
continue
|
||||
if ev == "error":
|
||||
continue
|
||||
if ev == "content-block-start":
|
||||
idx = event["index"]
|
||||
assert isinstance(idx, int), (
|
||||
f"content-block-start wire index must be an int, "
|
||||
f"got {idx!r} at event {i}"
|
||||
)
|
||||
assert idx >= 0, (
|
||||
f"content-block-start wire index must be non-negative, "
|
||||
f"got {idx} at event {i}"
|
||||
)
|
||||
assert idx == expected_next_idx, (
|
||||
f"expected next wire index {expected_next_idx}, got {idx} at event {i}"
|
||||
)
|
||||
assert open_idx is None, (
|
||||
f"content-block-start at idx={idx} while block {open_idx} "
|
||||
f"still open (event {i}); blocks must not interleave"
|
||||
)
|
||||
open_idx = idx
|
||||
start_events[idx] = event["content_block"]
|
||||
delta_accum[idx] = {}
|
||||
expected_next_idx += 1
|
||||
elif ev == "content-block-delta":
|
||||
idx = event["index"]
|
||||
assert idx == open_idx, (
|
||||
f"content-block-delta at idx={idx} but currently-open block is "
|
||||
f"{open_idx} (event {i})"
|
||||
)
|
||||
block = event["content_block"]
|
||||
_accumulate_delta(delta_accum[idx], block)
|
||||
elif ev == "content-block-finish":
|
||||
idx = event["index"]
|
||||
assert idx == open_idx, (
|
||||
f"content-block-finish at idx={idx} but currently-open block is "
|
||||
f"{open_idx} (event {i})"
|
||||
)
|
||||
finish_events[idx] = event["content_block"]
|
||||
open_idx = None
|
||||
else:
|
||||
# Unknown event types are accepted; the CDDL allows extensions.
|
||||
continue
|
||||
|
||||
assert open_idx is None, (
|
||||
f"block {open_idx} still open at end of stream — no content-block-finish"
|
||||
)
|
||||
missing = set(start_events) - set(finish_events)
|
||||
assert not missing, (
|
||||
f"the following block indices have no content-block-finish event: "
|
||||
f"{sorted(missing)}"
|
||||
)
|
||||
|
||||
for idx, finish_block in finish_events.items():
|
||||
_assert_delta_matches_finish(idx, delta_accum[idx], finish_block)
|
||||
|
||||
|
||||
def _accumulate_delta(accum: dict[str, Any], block: dict[str, Any]) -> None:
|
||||
"""Fold a delta block into the running accumulator for its index."""
|
||||
btype = block.get("type")
|
||||
if btype not in _DELTAABLE_TYPES:
|
||||
return
|
||||
if btype == "text":
|
||||
accum["text"] = accum.get("text", "") + block.get("text", "")
|
||||
elif btype == "reasoning":
|
||||
accum["reasoning"] = accum.get("reasoning", "") + block.get("reasoning", "")
|
||||
else: # tool_call_chunk / server_tool_call_chunk
|
||||
accum["args"] = accum.get("args", "") + (block.get("args") or "")
|
||||
if block.get("id") is not None:
|
||||
accum["id"] = block["id"]
|
||||
if block.get("name") is not None:
|
||||
accum["name"] = block["name"]
|
||||
|
||||
|
||||
def _assert_delta_matches_finish(
|
||||
idx: int,
|
||||
accum: dict[str, Any],
|
||||
finish_block: dict[str, Any],
|
||||
) -> None:
|
||||
"""Assert accumulated delta content is reflected in the finish payload."""
|
||||
ftype = finish_block.get("type")
|
||||
if ftype == "text" and "text" in accum:
|
||||
assert finish_block.get("text", "") == accum["text"], (
|
||||
f"block {idx} text accumulation {accum['text']!r} does not match "
|
||||
f"finish text {finish_block.get('text', '')!r}"
|
||||
)
|
||||
elif ftype == "reasoning" and "reasoning" in accum:
|
||||
assert finish_block.get("reasoning", "") == accum["reasoning"], (
|
||||
f"block {idx} reasoning accumulation mismatch: "
|
||||
f"accumulated {accum['reasoning']!r}, finish "
|
||||
f"{finish_block.get('reasoning', '')!r}"
|
||||
)
|
||||
elif ftype == "tool_call" and "args" in accum:
|
||||
# tool_call_chunk args are concatenated partial-JSON strings that
|
||||
# parse to a dict on finish.
|
||||
try:
|
||||
parsed = json.loads(accum["args"]) if accum["args"] else {}
|
||||
except json.JSONDecodeError:
|
||||
# Finish upgrades malformed args to invalid_tool_call, not
|
||||
# tool_call — so a tool_call finish implies args parsed cleanly.
|
||||
parsed = None
|
||||
assert finish_block.get("args") == parsed, (
|
||||
f"block {idx} tool_call args mismatch: accumulated parse "
|
||||
f"{parsed!r}, finish {finish_block.get('args')!r}"
|
||||
)
|
||||
elif ftype == "server_tool_call" and "args" in accum:
|
||||
try:
|
||||
parsed = json.loads(accum["args"]) if accum["args"] else {}
|
||||
except json.JSONDecodeError:
|
||||
parsed = None
|
||||
assert finish_block.get("args") == parsed
|
||||
|
||||
|
||||
__all__ = ["assert_valid_event_stream"]
|
||||
Reference in New Issue
Block a user