diff --git a/tests/integration/translation/README.md b/tests/integration/translation/README.md index 1b143fc1a1a..424772d8e43 100644 --- a/tests/integration/translation/README.md +++ b/tests/integration/translation/README.md @@ -4,6 +4,8 @@ These tests check one request through the proxy as literals on a `TranslationTes Some response fields are generated by LiteLLM on every run, such as a chat completion's `id` and `created` or a response's `id`, `created_at` and output message `id`. A case writes those as `unittest.mock.ANY`, so the field must still be present but any value passes. Use `ANY` only for a value LiteLLM generates, never for one copied from the provider reply +`StreamingTranslationTestCase` stores the stream request, exact provider request and provider and LiteLLM SSE frames. The runner compares parsed frames exactly and checks the final message they join into against `expected_litellm_response`, with streaming cases living in `/streaming/` + Deployments used by translation tests are shared by the whole suite and declared in `proxy_config.yaml` under `model_list`, with `model_name` equal to the litellm model string, `api_base: http://127.0.0.1:8191` and a synthetic key. A case names the deployment literally in its client request. The fake provider behind them is the `provider` fixture: one `wire_server` on port 8191 inside the pytest process, started the first time a test asks for it. A test queues its replies with `provider.expect(...)` and reads what the proxy sent with `provider.received()`. Tests that use it run one at a time. After each test the fixture fails if a queued reply was never requested or a received request was never read, and before each test it fails if a request arrived in between, naming the previous test. It refuses to start under pytest-xdist, so the `mcp` and `cost` groups cannot use it. Client requests carry `"cache": {"no-cache": True}` because the proxy caches responses in Redis Provider responses in translation cases are captured once from the real provider and stored verbatim. First run the new case against the fake provider; the body the proxy sends is the case's `expected_provider_request`. Send that exact body to the real provider endpoint with a key from the 1Password `Shared` vault (`/qa-keys`), and only store the case when the provider answers 2xx. Keep only the response body and drop every response header, since headers carry account identifiers such as the organization id and rate limits. Never print or save the request headers you sent. Before committing, check the body contains no key and no account identifier, such as an organization id, an AWS account id inside an ARN, a GCP project id or an Azure resource name, and that the prompt is synthetic. Paste the body as the case's `mock_provider_response` without shortening ids, token counts or signatures, and move long opaque values such as thinking signatures into module-level constants used by both `mock_provider_response` and `expected_litellm_response`. Do not commit the script used for the capture diff --git a/tests/integration/translation/case.py b/tests/integration/translation/case.py index b64edea21a1..4180001268c 100644 --- a/tests/integration/translation/case.py +++ b/tests/integration/translation/case.py @@ -1,5 +1,6 @@ from collections.abc import Mapping from dataclasses import dataclass +from typing import Literal from pydantic import JsonValue @@ -24,3 +25,28 @@ class TranslationTestCase: @property def id(self) -> str: return f"{self.litellm_request['model']}-{self.scenario}" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SSEFrame: + data: Mapping[str, JsonValue] | Literal["[DONE]"] + event: str | None = None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class StreamingTranslationTestCase: + """The stream request, provider request, provider SSE frames, exact LiteLLM frames, and the final message they join into.""" + + scenario: str + litellm_endpoint: str + litellm_request: Mapping[str, JsonValue] + expected_provider_endpoint: str + expected_provider_headers: Mapping[str, str] + expected_provider_request: Mapping[str, JsonValue] + mock_provider_frames: tuple[SSEFrame, ...] + expected_litellm_frames: tuple[SSEFrame, ...] + expected_litellm_response: Mapping[str, JsonValue] + + @property + def id(self) -> str: + return f"{self.litellm_request['model']}-{self.scenario}" diff --git a/tests/integration/translation/chat_completions/streaming/__init__.py b/tests/integration/translation/chat_completions/streaming/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/integration/translation/chat_completions/streaming/test_chat_completions_streaming_anthropic.py b/tests/integration/translation/chat_completions/streaming/test_chat_completions_streaming_anthropic.py new file mode 100644 index 00000000000..ca852a92f94 --- /dev/null +++ b/tests/integration/translation/chat_completions/streaming/test_chat_completions_streaming_anthropic.py @@ -0,0 +1,120 @@ +from typing import Final +from unittest.mock import ANY + +import pytest +from integration._support.client import Gateway +from integration._support.provider import SharedProvider +from integration.translation.case import SSEFrame, StreamingTranslationTestCase +from integration.translation.chat_completions.bases.anthropic import CLAUDE_SONNET_5_5_TEST_CASE +from integration.translation.runner import assert_streaming_translation + +CLAUDE_SONNET_5_5_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase( + scenario="streaming", + litellm_endpoint=CLAUDE_SONNET_5_5_TEST_CASE.litellm_endpoint, + litellm_request={**CLAUDE_SONNET_5_5_TEST_CASE.litellm_request, "stream": True}, + expected_provider_endpoint=CLAUDE_SONNET_5_5_TEST_CASE.expected_provider_endpoint, + expected_provider_headers=CLAUDE_SONNET_5_5_TEST_CASE.expected_provider_headers, + expected_provider_request={**CLAUDE_SONNET_5_5_TEST_CASE.expected_provider_request, "stream": True}, + mock_provider_frames=( + SSEFrame( + event="message_start", + data={ + "type": "message_start", + "message": { + "model": "claude-sonnet-5-5", + "id": "msg_011CfyYUVrHK4BM53hfSH5Ld", + "type": "message", + "role": "assistant", + "content": [], + "container": None, + "stop_reason": None, + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 23, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 6, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + }, + ), + SSEFrame( + event="content_block_start", + data={ + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + ), + SSEFrame(event="ping", data={"type": "ping"}), + SSEFrame( + event="content_block_delta", + data={ + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hello."}, + }, + ), + SSEFrame(event="content_block_stop", data={"type": "content_block_stop", "index": 0}), + SSEFrame( + event="message_delta", + data={ + "type": "message_delta", + "delta": { + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "container": None, + }, + "usage": { + "input_tokens": 23, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "output_tokens": 6, + "output_tokens_details": {"thinking_tokens": 0}, + }, + }, + ), + SSEFrame(event="message_stop", data={"type": "message_stop"}), + ), + expected_litellm_frames=( + SSEFrame( + data={ + "id": ANY, + "object": "chat.completion.chunk", + "created": ANY, + "model": "anthropic/claude-sonnet-5-5", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hello."}}], + } + ), + SSEFrame( + data={ + "id": ANY, + "object": "chat.completion.chunk", + "created": ANY, + "model": "anthropic/claude-sonnet-5-5", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + } + ), + SSEFrame(data="[DONE]"), + ), + expected_litellm_response={ + "id": ANY, + "object": "chat.completion", + "created": ANY, + "model": "anthropic/claude-sonnet-5-5", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello."}, "finish_reason": "stop"}], + }, +) + + +@pytest.mark.parametrize("case", [CLAUDE_SONNET_5_5_STREAMING_TEST_CASE], ids=lambda case: case.id) +def test_chat_completions_streaming_anthropic( + case: StreamingTranslationTestCase, gateway: Gateway, provider: SharedProvider +) -> None: + assert_streaming_translation(case, gateway, provider) diff --git a/tests/integration/translation/runner.py b/tests/integration/translation/runner.py index 8e6587c916d..f911f368efd 100644 --- a/tests/integration/translation/runner.py +++ b/tests/integration/translation/runner.py @@ -5,10 +5,13 @@ from typing import Final from pydantic import JsonValue -from integration._support.client import Gateway +from openai.lib.streaming.chat import ChatCompletionStreamState +from openai.types.chat import ChatCompletionChunk + +from integration._support.client import JSON_OBJECT, Gateway from integration._support.provider import SharedProvider -from integration._support.wire import Reply -from integration.translation.case import TranslationTestCase +from integration._support.wire import Reply, Request +from integration.translation.case import SSEFrame, StreamingTranslationTestCase, TranslationTestCase from litellm.responses.utils import ResponsesAPIRequestUtils TRANSPORT_HEADERS: Final = frozenset( @@ -35,14 +38,12 @@ def assert_translation(case: TranslationTestCase, gateway: Gateway, provider: Sh provider.expect(Reply(status=case.mock_provider_status_code, body=json.dumps(case.mock_provider_response).encode())) response: Final = gateway.request("POST", case.litellm_endpoint, case.litellm_request) received: Final = provider.received() - assert [(request.method, request.target) for request in received] == [("POST", case.expected_provider_endpoint)] - sent: Final = received[0] - assert { - name: value - for name, value in sent.headers.items() - if name not in TRANSPORT_HEADERS and not name.startswith("x-stainless-") - } == dict(case.expected_provider_headers) - assert json.loads(sent.body) == case.expected_provider_request + _assert_provider_request( + received, + case.expected_provider_endpoint, + case.expected_provider_headers, + case.expected_provider_request, + ) assert response.status_code == case.expected_litellm_status_code, response.text model_id: Final = response.headers.get("x-litellm-model-id") expected: Final = expected_response_for_deployment(case.expected_litellm_response, model_id) @@ -57,3 +58,80 @@ def assert_translation(case: TranslationTestCase, gateway: Gateway, provider: Sh for deployment in deployments ), "The encoded response ID must match the configured deployment" assert response.json() == expected + + +def _assert_provider_request( + received: tuple[Request, ...], + endpoint: str, + headers: Mapping[str, str], + request: Mapping[str, JsonValue], +) -> None: + assert [(sent_request.method, sent_request.target) for sent_request in received] == [("POST", endpoint)] + sent: Final = received[0] + assert { + name: value + for name, value in sent.headers.items() + if name not in TRANSPORT_HEADERS and not name.startswith("x-stainless-") + } == dict(headers) + assert json.loads(sent.body) == request + + +def encode_sse_frame(frame: SSEFrame) -> bytes: + serialized: Final = frame.data if isinstance(frame.data, str) else json.dumps(frame.data) + event_line: Final = f"event: {frame.event}\n" if frame.event is not None else "" + return f"{event_line}data: {serialized}\n\n".encode() + + +def _parse_sse_frame(block: str) -> SSEFrame | None: + lines: Final = tuple(line for line in block.splitlines() if line and not line.startswith(":")) + if not lines: + return None + event_line: Final = next((line for line in lines if line.startswith("event:")), None) + event: Final = event_line.removeprefix("event:").removeprefix(" ") if event_line is not None else None + data_lines: Final = tuple(line for line in lines if line.startswith("data:")) + if not data_lines: + return None + data: Final = "\n".join(line.removeprefix("data:").removeprefix(" ") for line in data_lines) + if data == "[DONE]": + return SSEFrame(event=event, data="[DONE]") + return SSEFrame(event=event, data=JSON_OBJECT.validate_python(json.loads(data))) + + +def parse_sse_frames(body: str) -> tuple[SSEFrame, ...]: + normalized: Final = body.replace("\r\n", "\n") + blocks: Final = tuple(_parse_sse_frame(block) for block in normalized.split("\n\n")) + return tuple(frame for frame in blocks if frame is not None) + + +def _join_stream_frames(endpoint: str, frames: tuple[SSEFrame, ...]) -> Mapping[str, JsonValue]: + match endpoint: + case "/v1/chat/completions": + state: Final = ChatCompletionStreamState() + for frame in frames: + if frame.data != "[DONE]": + state.handle_chunk(ChatCompletionChunk.model_validate(frame.data)) + return JSON_OBJECT.validate_python(state.get_final_completion().model_dump(mode="json", exclude_none=True)) + case _: + raise AssertionError(f"no stream accumulator for {endpoint}") + + +def assert_streaming_translation( + case: StreamingTranslationTestCase, gateway: Gateway, provider: SharedProvider +) -> None: + provider.expect( + Reply( + content_type="text/event-stream", + chunks=tuple(encode_sse_frame(frame) for frame in case.mock_provider_frames), + ) + ) + response: Final = gateway.request("POST", case.litellm_endpoint, case.litellm_request) + _assert_provider_request( + provider.received(), + case.expected_provider_endpoint, + case.expected_provider_headers, + case.expected_provider_request, + ) + assert response.status_code == 200, response.text + frames: Final = parse_sse_frames(response.text) + assert frames == case.expected_litellm_frames + assert _join_stream_frames(case.litellm_endpoint, frames) == case.expected_litellm_response