test(e2e): verify streamed answers and tool continuation

This commit is contained in:
Yuneng Jiang 2026-09-14 22:46:55 -07:00
parent 91588221cd
commit 7a7770db0d
No known key found for this signature in database
5 changed files with 218 additions and 30 deletions

View file

@ -170,6 +170,7 @@ class StreamingResponse(BaseModel):
# the consumed body is elided, so this is the only place they surface.
stream_error: str | None = None
stream_done: bool = False
stream_done_positions: tuple[int, ...] = ()
@property
def ok(self) -> bool:
@ -647,6 +648,7 @@ def streaming_outcome(
stream_events=[payload for payload, _ in events],
stream_event_arrivals=[arrived for _, arrived in events],
stream_done=any(payload == _SSE_DONE for payload, _ in payloads),
stream_done_positions=tuple(index for index, (payload, _) in enumerate(payloads) if payload == _SSE_DONE),
stream_error=next(
(line.decode(errors="replace")[:300] for line, _ in stamped if _is_stream_error_line(line)),
None,

View file

@ -1,51 +1,90 @@
"""Vendor §12.3: chat completions streaming SSE contract (LIT-4778).
Asserts a streamed /chat/completions response is SSE, carries content chunks,
and terminates with the OpenAI [DONE] sentinel.
"""
from __future__ import annotations
from typing import Final
import pytest
from e2e_config import unique_marker
from e2e_config import provider_edge_base, unique_marker
from e2e_http import require_successful_call
from lifecycle import ResourceManager
from models import ChatBody, ChatMessage, LiteLLMParamsBody
from models import ChatBody, ChatMessage, ChatStreamOptions, LiteLLMParamsBody, Usage
from proxy_client import ProxyClient
from pydantic import BaseModel
pytestmark = pytest.mark.e2e
pytestmark = [pytest.mark.e2e, pytest.mark.replayable]
class _Delta(BaseModel):
content: str | None = None
class _Choice(BaseModel):
index: int
delta: _Delta
finish_reason: str | None = None
class _Chunk(BaseModel):
choices: tuple[_Choice, ...]
usage: Usage | None = None
class TestChatStreamContract:
@pytest.mark.covers("llm.chat_completions.openai.basic.stream.works")
def test_chat_stream_is_sse_and_ends_with_done(self, proxy: ProxyClient, resources: ResourceManager) -> None:
model = f"e2e-chat-stream-{unique_marker()}"
model_id = proxy.create_model(
model: Final = f"e2e-chat-stream-{unique_marker()}"
base: Final = provider_edge_base("openai")
model_id: Final = proxy.create_model(
model,
LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY"),
LiteLLMParamsBody(
model="openai/gpt-5.6",
api_key="os.environ/OPENAI_API_KEY",
api_base=f"{base}/v1" if base else None,
),
)
resources.defer(lambda: proxy.delete_model(model_id))
key = resources.key()
result = proxy.chat_stream(
key: Final = resources.key()
expected: Final = "The amber kite crosses the quiet lake."
result: Final = proxy.chat_stream(
key,
ChatBody(
model=model,
messages=[
ChatMessage(
role="user",
content=f"Reply with the single word ok. {unique_marker()}",
role="user", content=f"Repeat exactly this sentence, with no additional text: {expected}"
)
],
stream=True,
max_completion_tokens=32,
temperature=0.0,
stream_options=ChatStreamOptions(include_usage=True),
max_completion_tokens=256,
reasoning_effort="none",
),
)
require_successful_call(result)
assert result.is_streaming, f"expected SSE content-type, got {result.content_type!r}"
assert result.stream_events, "stream returned no data events"
assert result.stream_done, (
f"stream must terminate with [DONE]; "
f"chunks={result.chunks} done={result.stream_done} events={len(result.stream_events)}"
assert not result.stream_error, f"stream errored: {result.stream_error}"
assert result.stream_done, "stream must terminate with [DONE]"
assert result.stream_done_positions == (len(result.stream_events),), "[DONE] must occur once after all events"
chunks: Final = tuple(_Chunk.model_validate_json(event) for event in result.stream_events)
text_positions: Final = tuple(
i for i, chunk in enumerate(chunks) if any(c.delta.content for c in chunk.choices)
)
terminal_positions: Final = tuple(
i for i, chunk in enumerate(chunks) if any(c.finish_reason is not None for c in chunk.choices)
)
assert text_positions, "stream completed without meaningful text"
assert len(terminal_positions) == 1, "expected exactly one terminal choice"
assert text_positions[0] < terminal_positions[0], "meaningful text must arrive before termination"
assert text_positions[-1] <= terminal_positions[0], "text arrived after termination"
assert all(c.index == 0 for chunk in chunks for c in chunk.choices)
assert tuple(c.finish_reason for c in chunks[terminal_positions[0]].choices) == ("stop",)
text: Final = "".join(c.delta.content or "" for chunk in chunks for c in chunk.choices)
assert text.strip() == expected, f"streamed answer was altered or incomplete: {text!r}"
usage_positions: Final = tuple(i for i, chunk in enumerate(chunks) if chunk.usage is not None)
assert usage_positions == (len(chunks) - 1,), "expected one final usage chunk"
assert terminal_positions[0] < usage_positions[0], "usage must follow the terminal choice"
usage: Final = chunks[-1].usage
assert usage is not None
assert usage.prompt_tokens is not None and usage.prompt_tokens > 0
assert usage.completion_tokens is not None and usage.completion_tokens > 0
assert usage.total_tokens == usage.prompt_tokens + usage.completion_tokens

View file

@ -21,7 +21,12 @@ from e2e_http import assert_client_error, require_successful_call, unwrap
from endpoints_client import EndpointsClient, MessagesResult
from lifecycle import ResourceManager
from models import (
AnthropicAssistantTurn,
AnthropicContentBlock,
AnthropicCustomTool,
AnthropicToolChoice,
AnthropicToolResultBlock,
AnthropicToolResultTurn,
AnthropicMessagesBody,
ChatMessage,
JsonSchemaProperty,
@ -29,7 +34,7 @@ from models import (
SpendLogRow,
ToolInputSchema,
)
from pydantic import BaseModel
from pydantic import BaseModel, ConfigDict
pytestmark = [pytest.mark.e2e, pytest.mark.replayable]
@ -284,8 +289,139 @@ class TestAnthropicMessages:
result = endpoints_client.proxy.transport.send(
"/v1/messages",
headers=endpoints_client.proxy.transport.bearer(key),
json=_OptionalMessagesBody(
messages=[ChatMessage(role="user", content="hi")], max_tokens=50
),
json=_OptionalMessagesBody(messages=[ChatMessage(role="user", content="hi")], max_tokens=50),
)
assert_client_error(result, "messages missing model")
class _BridgeDelta(BaseModel):
type: str | None = None
partial_json: str | None = None
stop_reason: str | None = None
class _BridgeEvent(BaseModel):
type: str
index: int | None = None
content_block: AnthropicContentBlock | None = None
delta: _BridgeDelta | None = None
class _ParcelInput(BaseModel):
model_config = ConfigDict(extra="forbid", strict=True)
parcel: str
shelf: int
def _tool_from_stream(events: tuple[_BridgeEvent, ...]) -> AnthropicContentBlock:
starts: Final = tuple(
event
for event in events
if event.type == "content_block_start"
and event.content_block is not None
and event.content_block.type == "tool_use"
)
assert len(starts) == 1, "expected exactly one tool call"
start: Final = starts[0]
block: Final = start.content_block
assert block is not None and block.id and start.index is not None
fragments: Final = tuple(
event
for event in events
if event.type == "content_block_delta" and event.delta is not None and event.delta.type == "input_json_delta"
)
assert fragments, "tool stream contained no argument fragments"
assert all(event.index == start.index for event in fragments), "tool fragments changed index"
positions: Final = tuple(i for i, event in enumerate(events) if event in fragments)
stops: Final = tuple(
i for i, event in enumerate(events) if event.type == "content_block_stop" and event.index == start.index
)
assert len(stops) == 1 and events.index(start) < positions[0] <= positions[-1] < stops[0]
assert tuple(
event.delta.stop_reason for event in events if event.type == "message_delta" and event.delta is not None
) == ("tool_use",)
terminal_positions: Final = tuple(i for i, event in enumerate(events) if event.type == "message_delta")
assert len(terminal_positions) == 1 and stops[0] < terminal_positions[0] < len(events) - 1
assert tuple(i for i, event in enumerate(events) if event.type == "message_stop") == (len(events) - 1,), (
"tool stream did not terminate exactly once"
)
arguments: Final = _ParcelInput.model_validate_json(
"".join(event.delta.partial_json or "" for event in fragments if event.delta is not None)
)
return AnthropicContentBlock(type="tool_use", id=block.id, name=block.name, input=arguments.model_dump())
def _parcel_result(tool: AnthropicContentBlock, result: AnthropicToolResultBlock) -> AnthropicToolResultTurn:
assert tool.id and result.tool_use_id == tool.id, "tool result ID does not match the emitted call"
return AnthropicToolResultTurn(content=[result])
def _request_tool(
client: EndpointsClient, key: str, request: AnthropicMessagesBody, stream: bool
) -> AnthropicContentBlock:
if stream:
response: Final = client.proxy.messages_stream(key, request)
require_successful_call(response)
assert response.is_streaming and not response.stream_error
return _tool_from_stream(tuple(_BridgeEvent.model_validate_json(event) for event in response.stream_events))
response_body: Final = unwrap(client.proxy.messages(key, request))
blocks: Final = tuple(block for block in response_body.content or () if block.type == "tool_use")
assert len(blocks) == 1
return blocks[0]
class TestOpenAIMessagesToolContinuation:
@pytest.mark.parametrize("stream", [True, False], ids=["stream", "nonstream"])
def test_required_tool_arguments_and_correlated_result(
self, endpoints_client: EndpointsClient, resources: ResourceManager, stream: bool
) -> None:
model: Final = f"e2e-bridge-tool-{unique_marker()}"
base: Final = provider_edge_base("openai")
model_id: Final = endpoints_client.create_model(
model,
LiteLLMParamsBody(
model="openai/gpt-5.6", api_key="os.environ/OPENAI_API_KEY", api_base=f"{base}/v1" if base else None
),
)
resources.defer(lambda: endpoints_client.delete_model(model_id))
key: Final = resources.key(models=[model])
tool: Final = AnthropicCustomTool(
name="locate_parcel",
description="Look up the receipt for a parcel on a shelf. Return the receipt verbatim.",
input_schema=ToolInputSchema(
properties={"parcel": JsonSchemaProperty(type="string"), "shelf": JsonSchemaProperty(type="integer")},
required=["parcel", "shelf"],
),
)
question: Final = ChatMessage(
role="user",
content="Call locate_parcel with parcel exactly amber-kite and shelf exactly 7. After the tool result, reply with only the receipt returned by the tool.",
)
request: Final = AnthropicMessagesBody(
model=model,
max_tokens=2048,
messages=[question],
tools=[tool],
tool_choice=AnthropicToolChoice(type="tool", name=tool.name),
stream=stream,
)
emitted: Final = _request_tool(endpoints_client, key, request, stream)
assert emitted.id and emitted.name == "locate_parcel"
assert emitted.input == {"parcel": "amber-kite", "shelf": 7}, "required tool arguments were lost or changed"
receipt: Final = f"receipt-{unique_marker()}"
result_turn: Final = _parcel_result(emitted, AnthropicToolResultBlock(tool_use_id=emitted.id, content=receipt))
continuation: Final = unwrap(
endpoints_client.proxy.messages(
key,
AnthropicMessagesBody(
model=model,
max_tokens=2048,
tools=[tool],
tool_choice=AnthropicToolChoice(type="none"),
messages=[question, AnthropicAssistantTurn(content=[emitted]), result_turn],
),
)
)
answer: Final = "".join(block.text or "" for block in continuation.content or ())
assert answer.strip() == receipt, "continuation did not consume the correlated tool result"
assert all(block.type != "tool_use" for block in continuation.content or ())

View file

@ -283,10 +283,15 @@ class ChatToolResultTurn(BaseModel):
type ChatTurn = ChatMessage | ChatAssistantTurn | ChatToolResultTurn
class ChatStreamOptions(BaseModel):
include_usage: bool
class ChatBody(BaseModel):
model: str
messages: Sequence[ChatTurn]
stream: bool = False
stream_options: ChatStreamOptions | None = None
max_tokens: int | None = None
max_completion_tokens: int | None = None
temperature: float | None = None
@ -488,12 +493,18 @@ class AnthropicToolResultTurn(BaseModel):
type AnthropicMessage = ChatMessage | AnthropicAssistantTurn | AnthropicToolResultTurn
class AnthropicToolChoice(BaseModel):
type: Literal["auto", "any", "tool", "none"]
name: str | None = None
class AnthropicMessagesBody(BaseModel):
model: str
messages: list[AnthropicMessage]
max_tokens: int
stream: bool | None = None
tools: list[AnthropicTool] | None = None
tool_choice: AnthropicToolChoice | None = None
guardrails: list[str] | None = None
cache: dict[str, bool] | None = {"no-cache": True}

View file

@ -8,6 +8,8 @@ Without the fix, the AnthropicStreamWrapper silently dropped these
arguments, causing tool_use blocks to arrive with empty input {}.
"""
import json
from typing import List
from unittest.mock import MagicMock
@ -139,9 +141,7 @@ async def test_async_stream_emits_input_json_delta_for_bundled_tool_args():
# Verify the delta carries the tool arguments
delta_event = events[input_json_delta_idx]
assert delta_event["delta"][
"partial_json"
], "input_json_delta should have non-empty partial_json"
assert json.loads(delta_event["delta"]["partial_json"]) == {"location": "Boston"}
@pytest.mark.asyncio
@ -300,7 +300,7 @@ def test_sync_stream_emits_input_json_delta_for_bundled_tool_args():
assert (
input_json_delta_idx == tool_start_idx + 1
), "input_json_delta should immediately follow the tool_use content_block_start"
assert events[input_json_delta_idx]["delta"]["partial_json"]
assert json.loads(events[input_json_delta_idx]["delta"]["partial_json"]) == {"location": "Boston"}
def test_sync_stream_no_extra_delta_when_tool_args_empty():