mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
test(e2e): verify streamed answers and tool continuation
This commit is contained in:
parent
91588221cd
commit
7a7770db0d
5 changed files with 218 additions and 30 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 ())
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue