diff --git a/tests/test_litellm/rust_bridge/native_route_wheel_test.py b/tests/test_litellm/rust_bridge/native_route_wheel_test.py index a7f50a82a99..927c315e634 100644 --- a/tests/test_litellm/rust_bridge/native_route_wheel_test.py +++ b/tests/test_litellm/rust_bridge/native_route_wheel_test.py @@ -10,6 +10,8 @@ import sys import tempfile import threading import zipfile +from collections.abc import Callable, Mapping +from dataclasses import dataclass from http.client import HTTPMessage from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path @@ -27,6 +29,123 @@ ANTHROPIC_RESPONSE: Final = ( ) +def assert_ocr_request(path: str, headers: HTTPMessage, body: dict[str, object]) -> None: + assert path == "/v1/ocr" + assert headers.get("authorization") == "Bearer sk-native" + assert body == { + "model": "mistral-ocr-latest", + "document": {"type": "document_url", "document_url": "https://example.com/document.pdf"}, + "include_image_base64": True, + } + + +def assert_transcription_request(path: str, headers: HTTPMessage, body: dict[str, object]) -> None: + assert path == "/model/mistral.voxtral-mini-3b-2507/converse" + assert headers.get("authorization", "").startswith("AWS4-HMAC-SHA256 ") + assert headers.get("x-amz-date") + messages: Final = body["messages"] + assert isinstance(messages, list) + assert messages[0]["content"][0]["audio"]["source"]["bytes"] == "AQI=" + assert "The audio language is en" in messages[0]["content"][1]["text"] + + +def assert_messages_request(path: str, headers: HTTPMessage, body: dict[str, object]) -> None: + assert path == "/v1/messages" + assert headers.get("x-api-key") == "sk-native" + assert body == { + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hello-from-messages"}], + } + + +def assert_chat_completions_request(path: str, headers: HTTPMessage, body: dict[str, object]) -> None: + assert path == "/v1/messages" + assert headers.get("x-api-key") == "sk-native" + assert body == { + "model": "claude-sonnet-4-5", + "max_tokens": 17, + "messages": [{"role": "user", "content": [{"type": "text", "text": "hello-from-chat"}]}], + } + + +@dataclass(frozen=True, slots=True) +class RouteCase: + name: str + request_kwargs: Mapping[str, object] + assert_request: Callable[[str, HTTPMessage, dict[str, object]], None] + success_value: Callable[[dict[str, object]], object] + expected_success: str + native_upstream_error: bool + + +ROUTE_CASES: Final = ( + RouteCase( + name="ocr", + request_kwargs={ + "model": "mistral-ocr-latest", + "document": {"type": "document_url", "document_url": "https://example.com/document.pdf"}, + "api_key": "sk-native", + "custom_llm_provider": "mistral", + "optional_params": {"include_image_base64": True}, + }, + assert_request=assert_ocr_request, + success_value=lambda response: response["pages"][0]["markdown"], + expected_success="native-ocr", + native_upstream_error=True, + ), + RouteCase( + name="transcription", + request_kwargs={ + "model": "mistral.voxtral-mini-3b-2507", + "audio": {"data": "AQI=", "format": "wav", "filename": "audio.wav"}, + "custom_llm_provider": "bedrock", + "optional_params": { + "aws_access_key_id": "native-access-key", + "aws_secret_access_key": "native-secret-key", + "aws_region_name": "us-east-1", + "language": "en", + }, + }, + assert_request=assert_transcription_request, + success_value=lambda response: response["text"], + expected_success="native-transcription", + native_upstream_error=False, + ), + RouteCase( + name="messages", + request_kwargs={ + "model": "claude-sonnet-4-5", + "body": { + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hello-from-messages"}], + }, + "api_key": "sk-native", + "custom_llm_provider": "anthropic", + }, + assert_request=assert_messages_request, + success_value=lambda response: response["content"][0]["text"], + expected_success="native-message", + native_upstream_error=False, + ), + RouteCase( + name="chat_completions", + request_kwargs={ + "model": "anthropic/claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hello-from-chat"}], + "optional_params": {"max_tokens": 17}, + "api_key": "sk-native", + }, + assert_request=assert_chat_completions_request, + success_value=lambda response: response["choices"][0]["message"]["content"], + expected_success="native-message", + native_upstream_error=True, + ), +) +ROUTE_CASE_BY_NAME: Final = {case.name: case for case in ROUTE_CASES} + + class NativeRouteHandler(BaseHTTPRequestHandler): protocol_version = "HTTP/1.1" @@ -73,35 +192,14 @@ def assert_native_request( headers: HTTPMessage, body: object, ) -> None: - if route not in {"ocr", "transcription", "messages", "chat_completions"}: + case: Final = ROUTE_CASE_BY_NAME.get(route or "") + if case is None: raise AssertionError(f"unexpected route marker: {route!r}") if outcome not in {"success", "429", "hang"}: raise AssertionError(f"unexpected outcome marker: {outcome!r}") if not isinstance(body, dict): raise TypeError(f"{route} sent {type(body).__name__}, expected a JSON object") - if route == "ocr": - assert path == "/v1/ocr" - assert headers.get("authorization") == "Bearer sk-native" - assert body["model"] == "mistral-ocr-latest" - assert body["document"]["document_url"] == "https://example.com/document.pdf" - assert body["include_image_base64"] is True - return - if route == "transcription": - assert path == "/model/mistral.voxtral-mini-3b-2507/converse" - assert headers.get("authorization", "").startswith("AWS4-HMAC-SHA256 ") - assert headers.get("x-amz-date") - assert body["messages"][0]["content"][0]["audio"]["source"]["bytes"] == "AQI=" - assert "The audio language is en" in body["messages"][0]["content"][1]["text"] - return - assert path == "/v1/messages" - assert headers.get("x-api-key") == "sk-native" - assert body["model"] == "claude-sonnet-4-5" - if route == "messages": - assert body["max_tokens"] == 16 - assert body["messages"][0]["content"] == "hello-from-messages" - return - assert body["max_tokens"] == 17 - assert body["messages"][0]["content"] == [{"type": "text", "text": "hello-from-chat"}] + case.assert_request(path, headers, body) def native_response(status: int, route: str | None) -> bytes: @@ -124,75 +222,29 @@ def load_native(native_path: Path) -> object: def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]: + case: Final = ROUTE_CASE_BY_NAME.get(route) + if case is None: + raise AssertionError(f"unknown route: {route}") common: Final = { "api_base": api_base, "extra_headers": {"x-test-outcome": outcome, "x-test-route": route}, "timeout_seconds": 3.0, } - if route == "ocr": - return common | { - "model": "mistral-ocr-latest", - "document": {"type": "document_url", "document_url": "https://example.com/document.pdf"}, - "api_key": "sk-native", - "custom_llm_provider": "mistral", - "optional_params": {"include_image_base64": True}, - } - if route == "transcription": - return common | { - "model": "mistral.voxtral-mini-3b-2507", - "audio": {"data": "AQI=", "format": "wav", "filename": "audio.wav"}, - "custom_llm_provider": "bedrock", - "optional_params": { - "aws_access_key_id": "native-access-key", - "aws_secret_access_key": "native-secret-key", - "aws_region_name": "us-east-1", - "language": "en", - }, - } - if route == "messages": - return common | { - "model": "claude-sonnet-4-5", - "body": { - "model": "claude-sonnet-4-5", - "max_tokens": 16, - "messages": [{"role": "user", "content": "hello-from-messages"}], - }, - "api_key": "sk-native", - "custom_llm_provider": "anthropic", - } - if route == "chat_completions": - return common | { - "model": "anthropic/claude-sonnet-4-5", - "messages": [{"role": "user", "content": "hello-from-chat"}], - "optional_params": {"max_tokens": 17}, - "api_key": "sk-native", - } - raise AssertionError(f"unknown route: {route}") + return common | dict(case.request_kwargs) def assert_success(route: str, response: object) -> None: + case: Final = ROUTE_CASE_BY_NAME[route] if not isinstance(response, dict): raise TypeError(f"{route} returned {type(response).__name__}, expected dict") - actual: Final = success_value(route, response) - expected: Final = ( - "native-ocr" if route == "ocr" else "native-transcription" if route == "transcription" else "native-message" - ) - if actual != expected: - raise AssertionError(f"{route} returned {actual!r}, expected {expected!r}") - - -def success_value(route: str, response: dict[object, object]) -> object: - if route == "ocr": - return response["pages"][0]["markdown"] - if route == "transcription": - return response["text"] - if route == "messages": - return response["content"][0]["text"] - return response["choices"][0]["message"]["content"] + actual: Final = case.success_value(response) + if actual != case.expected_success: + raise AssertionError(f"{route} returned {actual!r}, expected {case.expected_success!r}") def assert_rate_limit(native: object, route: str, error: BaseException) -> None: - if route in {"ocr", "chat_completions"}: + case: Final = ROUTE_CASE_BY_NAME[route] + if case.native_upstream_error: upstream_error: Final = native.RustUpstreamError if not isinstance(error, upstream_error) or error.args[0] != 429: raise AssertionError(f"{route} returned the wrong 429 error: {error!r}") @@ -202,37 +254,32 @@ def assert_rate_limit(native: object, route: str, error: BaseException) -> None: def exercise_sync(native: object, api_base: str) -> None: - for route in ("ocr", "transcription", "messages", "chat_completions"): - function: Final = getattr(native, route) - assert_success(route, function(**route_kwargs(route, api_base, "success"))) + for case in ROUTE_CASES: + function: Final = getattr(native, case.name) + assert_success(case.name, function(**route_kwargs(case.name, api_base, "success"))) try: - function(**route_kwargs(route, api_base, "429")) + function(**route_kwargs(case.name, api_base, "429")) except (RuntimeError, native.RustUpstreamError) as error: - assert_rate_limit(native, route, error) + assert_rate_limit(native, case.name, error) else: - raise AssertionError(f"{route} accepted a 429 response") + raise AssertionError(f"{case.name} accepted a 429 response") async def exercise_async(native: object, api_base: str) -> None: - for route in ("ocr", "transcription", "messages", "chat_completions"): - function: Final = getattr(native, f"a{route}") - assert_success(route, await function(**route_kwargs(route, api_base, "success"))) + for case in ROUTE_CASES: + function: Final = getattr(native, f"a{case.name}") + assert_success(case.name, await function(**route_kwargs(case.name, api_base, "success"))) try: - await function(**route_kwargs(route, api_base, "429")) + await function(**route_kwargs(case.name, api_base, "429")) except (RuntimeError, native.RustUpstreamError) as error: - assert_rate_limit(native, route, error) + assert_rate_limit(native, case.name, error) else: - raise AssertionError(f"a{route} accepted a 429 response") + raise AssertionError(f"a{case.name} accepted a 429 response") async def exercise_async_concurrency(native: object, api_base: str) -> None: responses: Final = await asyncio.wait_for( - asyncio.gather( - *( - native.amessages(**route_kwargs("messages", api_base, "success")) - for _ in range(32) - ) - ), + asyncio.gather(*(native.amessages(**route_kwargs("messages", api_base, "success")) for _ in range(32))), timeout=15, ) for response in responses: diff --git a/tests/test_litellm/rust_bridge/test_chat_completions.py b/tests/test_litellm/rust_bridge/test_chat_completions.py index 9754e2b0552..bfe17a0d812 100644 --- a/tests/test_litellm/rust_bridge/test_chat_completions.py +++ b/tests/test_litellm/rust_bridge/test_chat_completions.py @@ -163,11 +163,7 @@ def test_gate_forwards_enablement_to_preflight( if enablement == "environment": monkeypatch.setenv("LITELLM_RUST", "true") litellm_params: Final = ( - {"rust": True} - if enablement == "request-true" - else {"rust": False} - if enablement == "request-false" - else {} + {"rust": True} if enablement == "request-true" else {"rust": False} if enablement == "request-false" else {} ) assert accepts(litellm_params=litellm_params) is expected @@ -218,11 +214,14 @@ def test_gate_short_circuits_python_only_requests( bridge.set_rust_chat_completions(decline=gate) monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", bedrock_metadata_fields) - assert accepts( - custom_llm_provider=provider, - litellm_params=litellm_params, - stream=stream, - ) is False + assert ( + accepts( + custom_llm_provider=provider, + litellm_params=litellm_params, + stream=stream, + ) + is False + ) assert gate.calls == [] diff --git a/tests/test_litellm/rust_bridge/test_runtime.py b/tests/test_litellm/rust_bridge/test_runtime.py index fd0de57b757..16b8b5c0da3 100644 --- a/tests/test_litellm/rust_bridge/test_runtime.py +++ b/tests/test_litellm/rust_bridge/test_runtime.py @@ -288,11 +288,14 @@ def test_accepts_only_eligible_supported_requests( bridge: Final = runtime.RustBridge(route="preflight", load=load, enabled=enabled) - assert bridge.accepts( - check=check, - request_override=False if state == "disabled" else None, - eligible=state != "ineligible", - ) is expected + assert ( + bridge.accepts( + check=check, + request_override=False if state == "disabled" else None, + eligible=state != "ineligible", + ) + is expected + ) assert tuple(events) == expected_events diff --git a/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py b/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py index 2a5987b29bd..fdddb86d087 100644 --- a/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py +++ b/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py @@ -176,8 +176,7 @@ class _InvalidWheelCase: _InvalidWheelCase( metadata_tags=("cp312-cp312-linux_x86_64",), expected_error=( - "WHEEL tags do not match filename: expected cp310-abi3-linux_x86_64, " - "found cp312-cp312-linux_x86_64" + "WHEEL tags do not match filename: expected cp310-abi3-linux_x86_64, found cp312-cp312-linux_x86_64" ), ), id="metadata-tag-mismatched",