From 0fc11653e23782f778b82509f3f4682298493efe Mon Sep 17 00:00:00 2001 From: yujonglee Date: Wed, 14 Oct 2026 17:07:20 -0700 Subject: [PATCH] test(rust-bridge): pin Messages and Responses callback contracts on the native path (#46443) * test(rust-bridge): pin Messages and Responses callback contracts on the native path Each test runs the native route and asserts what the Python path does today, so a known gap is a strict xfail that names its reason and the exception it fails with Messages and Responses share one Surface and hook recorder in support/callbacks.py, and each surface folder lists its own expected counts and gaps. Body-level hooks from the function body (pre-request hook, websearch short-circuit, Messages interceptors, planned and legacy agentic loops, streaming deployment hook) are xfail until native runs them Co-Authored-By: Claude Opus 5.5 * test(rust-bridge): give every route folder the same one-concern-per-file layout Route data, the route Surface and the server fixture move into each route folder, and support/ keeps only route-neutral helpers. test_inference.py splits into chat_completions/ and responses/, ocr/test_requests.py splits into requests, failures and settings, and Messages and Responses get streaming and body hook files. Messages request tests that duplicate inference-messages request.rs are removed, and drop_params read from litellm.drop_params is now covered Co-Authored-By: Claude Opus 5.5 * test(rust-bridge): expect the required-route decline error for unsupported Responses requests Co-Authored-By: Yujong Lee * test(rust-bridge): import route data from the route folders in the secrets unit tests Co-Authored-By: Yujong Lee * test(rust-bridge): import NoReturn in the OCR dispatch tests Co-Authored-By: Yujong Lee * test(rust-bridge): move route-neutral helpers out of support/callbacks.py Co-Authored-By: Yujong Lee * test(rust-bridge): cover deployment rejection and concurrent callbacks * test(rust-bridge): assert native serving via the x-litellm-rust marker in the route folder tests Co-Authored-By: Yujong Lee --------- Co-authored-by: Claude Opus 5.5 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm_rust/AGENTS.md | 18 + .../{support/cache.py => cache/support.py} | 4 +- .../cache/test_python_cache.py | 11 +- .../test_litellm_rust/cache/v2/test_cache.py | 8 +- .../chat_completions/__init__.py | 0 .../chat_completions/conftest.py | 10 + .../chat_completions/data.py | 28 + .../chat_completions/test_callbacks.py | 65 ++ .../chat_completions/test_failures.py | 69 ++ .../chat_completions/test_lifecycle.py | 37 + .../chat_completions/test_requests.py | 42 ++ .../chat_completions/test_settings.py | 82 +++ tests/test_litellm_rust/conftest.py | 4 +- tests/test_litellm_rust/messages/conftest.py | 10 + tests/test_litellm_rust/messages/data.py | 64 ++ .../messages/test_body_hooks.py | 137 ++++ .../messages/test_callbacks.py | 569 ++++++++++------ .../messages/test_request_shaping.py | 311 --------- .../messages/test_requests.py | 150 +++++ .../messages/test_streaming.py | 73 ++ tests/test_litellm_rust/ocr/conftest.py | 19 + tests/test_litellm_rust/ocr/data.py | 39 ++ tests/test_litellm_rust/ocr/test_callbacks.py | 181 +++-- tests/test_litellm_rust/ocr/test_cohere.py | 32 - tests/test_litellm_rust/ocr/test_dispatch.py | 108 ++- tests/test_litellm_rust/ocr/test_failures.py | 340 ++++++++++ .../test_litellm_rust/ocr/test_guardrails.py | 24 +- tests/test_litellm_rust/ocr/test_lifecycle.py | 124 ++-- tests/test_litellm_rust/ocr/test_requests.py | 636 +----------------- tests/test_litellm_rust/ocr/test_settings.py | 198 ++++++ tests/test_litellm_rust/responses/__init__.py | 0 tests/test_litellm_rust/responses/conftest.py | 10 + tests/test_litellm_rust/responses/data.py | 75 +++ .../responses/test_body_hooks.py | 114 ++++ .../responses/test_callbacks.py | 408 +++++++++++ .../responses/test_failures.py | 83 +++ .../responses/test_lifecycle.py | 46 ++ .../responses/test_requests.py | 56 ++ .../responses/test_settings.py | 83 +++ .../responses/test_streaming.py | 58 ++ tests/test_litellm_rust/support/body_hooks.py | 39 ++ .../support/callback_recorder.py | 61 ++ tests/test_litellm_rust/support/gaps.py | 5 + .../support/recording_server.py | 6 + tests/test_litellm_rust/support/requests.py | 97 --- tests/test_litellm_rust/support/results.py | 26 + tests/test_litellm_rust/support/routing.py | 48 ++ tests/test_litellm_rust/support/served.py | 15 +- tests/test_litellm_rust/test_inference.py | 403 ----------- .../unit/rust_bridge/messages/test_secrets.py | 2 +- tests/unit/rust_bridge/ocr/test_secrets.py | 2 +- 51 files changed, 3157 insertions(+), 1873 deletions(-) rename tests/test_litellm_rust/{support/cache.py => cache/support.py} (96%) create mode 100644 tests/test_litellm_rust/chat_completions/__init__.py create mode 100644 tests/test_litellm_rust/chat_completions/conftest.py create mode 100644 tests/test_litellm_rust/chat_completions/data.py create mode 100644 tests/test_litellm_rust/chat_completions/test_callbacks.py create mode 100644 tests/test_litellm_rust/chat_completions/test_failures.py create mode 100644 tests/test_litellm_rust/chat_completions/test_lifecycle.py create mode 100644 tests/test_litellm_rust/chat_completions/test_requests.py create mode 100644 tests/test_litellm_rust/chat_completions/test_settings.py create mode 100644 tests/test_litellm_rust/messages/conftest.py create mode 100644 tests/test_litellm_rust/messages/data.py create mode 100644 tests/test_litellm_rust/messages/test_body_hooks.py delete mode 100644 tests/test_litellm_rust/messages/test_request_shaping.py create mode 100644 tests/test_litellm_rust/messages/test_requests.py create mode 100644 tests/test_litellm_rust/messages/test_streaming.py create mode 100644 tests/test_litellm_rust/ocr/conftest.py create mode 100644 tests/test_litellm_rust/ocr/data.py delete mode 100644 tests/test_litellm_rust/ocr/test_cohere.py create mode 100644 tests/test_litellm_rust/ocr/test_failures.py create mode 100644 tests/test_litellm_rust/ocr/test_settings.py create mode 100644 tests/test_litellm_rust/responses/__init__.py create mode 100644 tests/test_litellm_rust/responses/conftest.py create mode 100644 tests/test_litellm_rust/responses/data.py create mode 100644 tests/test_litellm_rust/responses/test_body_hooks.py create mode 100644 tests/test_litellm_rust/responses/test_callbacks.py create mode 100644 tests/test_litellm_rust/responses/test_failures.py create mode 100644 tests/test_litellm_rust/responses/test_lifecycle.py create mode 100644 tests/test_litellm_rust/responses/test_requests.py create mode 100644 tests/test_litellm_rust/responses/test_settings.py create mode 100644 tests/test_litellm_rust/responses/test_streaming.py create mode 100644 tests/test_litellm_rust/support/body_hooks.py create mode 100644 tests/test_litellm_rust/support/gaps.py delete mode 100644 tests/test_litellm_rust/support/requests.py create mode 100644 tests/test_litellm_rust/support/results.py create mode 100644 tests/test_litellm_rust/support/routing.py delete mode 100644 tests/test_litellm_rust/test_inference.py diff --git a/tests/test_litellm_rust/AGENTS.md b/tests/test_litellm_rust/AGENTS.md index d65ffd613aa..7e8cd67165f 100644 --- a/tests/test_litellm_rust/AGENTS.md +++ b/tests/test_litellm_rust/AGENTS.md @@ -1 +1,19 @@ This directory holds only the tests that cannot be written in the Rust code + +Each route folder (`ocr/`, `messages/`, `responses/`, `chat_completions/`) uses the same file names, one concern per file + +- `test_dispatch.py` which path served the call +- `test_requests.py` Python inputs that reach the wire +- `test_failures.py` errors raised to the caller and rejection before provider IO +- `test_settings.py` globals, env, named credentials and their precedence +- `test_streaming.py` stream relay, early close and stream-time logging +- `test_callbacks.py` the logger hook contract +- `test_body_hooks.py` handler body hooks such as interceptors and agentic loops +- `test_lifecycle.py` cancellation, concurrency, context and nested calls + +A concern that exists for only one route gets its own file there, never a provider-named one + +- Route data and the route's `call` live in `/data.py`, and `/conftest.py` owns the `server` fixture +- `support/` holds only route-neutral helpers. Never parametrize one test across routes, except in feature folders such as `cache/` +- Call the public API inside `support.routing.native_route(...)`. Call `_native` directly only for a case the public API cannot reach +- Check the native path with `support.routing.assert_served_natively` diff --git a/tests/test_litellm_rust/support/cache.py b/tests/test_litellm_rust/cache/support.py similarity index 96% rename from tests/test_litellm_rust/support/cache.py rename to tests/test_litellm_rust/cache/support.py index 5deacc7869f..a30f6b68cc6 100644 --- a/tests/test_litellm_rust/support/cache.py +++ b/tests/test_litellm_rust/cache/support.py @@ -16,9 +16,9 @@ from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES from litellm.rust_bridge.public_call import NativeCall from litellm.rust_bridge.responses.entrypoints import NATIVE_ARESPONSES from litellm.types.utils import ModelResponse +from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE +from tests.test_litellm_rust.responses.data import RESPONSES_MODEL, RESPONSES_RESPONSE from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec -from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE -from tests.test_litellm_rust.test_inference import RESPONSES_MODEL, RESPONSES_RESPONSE def payload(value: object) -> object: diff --git a/tests/test_litellm_rust/cache/test_python_cache.py b/tests/test_litellm_rust/cache/test_python_cache.py index 139859bdf55..9a0a874ddbf 100644 --- a/tests/test_litellm_rust/cache/test_python_cache.py +++ b/tests/test_litellm_rust/cache/test_python_cache.py @@ -21,14 +21,13 @@ from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_di from litellm.rust_bridge import runtime from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule from litellm.rust_bridge.configuration import Rollout -from litellm.rust_bridge.dispatch import call_hook from litellm.rust_bridge.public_call import NativeCall from litellm.types.caching import CachingSupportedCallTypes -from tests.test_litellm_rust.support.cache import cache_key, collect, invoke, payload -from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE +from tests.test_litellm_rust.responses.data import RESPONSES_RESPONSE +from tests.test_litellm_rust.cache.support import cache_key, collect, invoke, payload +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging, request_body from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec -from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE -from tests.test_litellm_rust.test_inference import RESPONSES_RESPONSE pytestmark = pytest.mark.requires_rust_extension @@ -350,8 +349,6 @@ async def test_cache_lookup_uses_backend_request_callback_semantics( route: Literal["chat", "messages", "responses"], legacy: bool, ) -> None: - from tests.test_litellm_rust.support.requests import request_body - class Rewrite(RecordingLogger): temperature = 0.1 diff --git a/tests/test_litellm_rust/cache/v2/test_cache.py b/tests/test_litellm_rust/cache/v2/test_cache.py index 5f930bd4ba7..0ba0a17007f 100644 --- a/tests/test_litellm_rust/cache/v2/test_cache.py +++ b/tests/test_litellm_rust/cache/v2/test_cache.py @@ -13,11 +13,11 @@ from litellm.caching.caching_handler import ( ) from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.types.utils import ModelResponse -from tests.test_litellm_rust.support.cache import cache_key, invoke, payload +from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE +from tests.test_litellm_rust.responses.data import RESPONSES_MODEL, RESPONSES_RESPONSE, RESPONSES_TEXT +from tests.test_litellm_rust.cache.support import cache_key, invoke, payload from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec -from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE -from tests.test_litellm_rust.test_inference import RESPONSES_MODEL, RESPONSES_RESPONSE pytestmark = pytest.mark.requires_rust_extension @@ -153,7 +153,7 @@ async def test_v2_cache_does_not_force_native_responses_streaming( assert isinstance(response, AsyncIterator) chunks: Final = [chunk async for chunk in response] assert chunks[-1].type == "response.completed" - assert chunks[-1].response.output[0].content[0].text == "native response" + assert chunks[-1].response.output[0].content[0].text == RESPONSES_TEXT @pytest.mark.asyncio diff --git a/tests/test_litellm_rust/chat_completions/__init__.py b/tests/test_litellm_rust/chat_completions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm_rust/chat_completions/conftest.py b/tests/test_litellm_rust/chat_completions/conftest.py new file mode 100644 index 00000000000..25c46c8ae42 --- /dev/null +++ b/tests/test_litellm_rust/chat_completions/conftest.py @@ -0,0 +1,10 @@ +import pytest + +from tests.test_litellm_rust.messages.data import MESSAGES_RESPONSE +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec + + +@pytest.fixture +def server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) + return recording_server diff --git a/tests/test_litellm_rust/chat_completions/data.py b/tests/test_litellm_rust/chat_completions/data.py new file mode 100644 index 00000000000..42f8e518f55 --- /dev/null +++ b/tests/test_litellm_rust/chat_completions/data.py @@ -0,0 +1,28 @@ +from typing import Final + +import litellm +from litellm.rust_bridge.catalog import Route +from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL +from tests.test_litellm_rust.support.recording_server import RecordingServer +from tests.test_litellm_rust.support.routing import call_natively + + +def _error(kind: str) -> dict[str, object]: + return {"type": "error", "error": {"type": kind, "message": kind}} + + +REQUEST: Final = {"model": MESSAGES_MODEL, "messages": list(MESSAGES), "max_tokens": 32} +BAD_REQUEST: Final = _error("invalid_request_error") +RATE_LIMITED: Final = _error("rate_limit_error") + + +async def call(server: RecordingServer, *, asynchronous: bool = True, **kwargs: object) -> object: + return await call_natively( + Route.CHAT_COMPLETIONS, + litellm.completion, + litellm.acompletion, + REQUEST, + server, + asynchronous=asynchronous, + **kwargs, + ) diff --git a/tests/test_litellm_rust/chat_completions/test_callbacks.py b/tests/test_litellm_rust/chat_completions/test_callbacks.py new file mode 100644 index 00000000000..6743f9adbf0 --- /dev/null +++ b/tests/test_litellm_rust/chat_completions/test_callbacks.py @@ -0,0 +1,65 @@ +from typing import Final + +import pytest +from pydantic import TypeAdapter + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.utils import CallTypes, ModelResponse +from tests.test_litellm_rust.chat_completions.data import call +from tests.test_litellm_rust.messages.data import MESSAGES_TEXT +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, request_body +from tests.test_litellm_rust.support.recording_server import RecordingServer +from tests.test_litellm_rust.support.routing import assert_served_natively + +pytestmark = pytest.mark.requires_rust_extension + +_OBJECT: Final = TypeAdapter(dict[str, object]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_returns_a_model_response_and_logs_success_once(server: RecordingServer, asynchronous: bool) -> None: + recorder: Final = RecordingLogger() + + result: Final = await call(server, asynchronous=asynchronous, callbacks=[recorder], temperature=0.25) + + assert_served_natively(server) + sent: Final = server.requests[0] + body: Final = _OBJECT.validate_python(sent.body) + assert sent.path == "/v1/messages" + assert body["temperature"] == 0.25 + assert request_body(_OBJECT.validate_python(recorder.wait_for("log_pre_api_call")[0].kwargs)) == body + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == MESSAGES_TEXT + success: Final = await recorder.wait_for_async("async_log_success_event" if asynchronous else "log_success_event") + assert len(success) == 1 + assert success[0].response is result + + +@pytest.mark.asyncio +async def test_pre_api_call_edits_reach_the_provider(server: RecordingServer) -> None: + class Edit(CustomLogger): + def log_pre_api_call(self, model: object, messages: object, kwargs: dict[str, object]) -> None: + request_body(kwargs)["temperature"] = 0.75 + + await call(server, callbacks=[Edit()]) + + assert _OBJECT.validate_python(server.requests[0].body)["temperature"] == 0.75 + + +@pytest.mark.asyncio +async def test_pre_call_deployment_hook_can_delete_keywords(server: RecordingServer) -> None: + class DropKeywords(CustomLogger): + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: CallTypes | None + ) -> dict[str, object]: + return {name: value for name, value in kwargs.items() if name not in ("temperature", "extra_headers")} + + litellm.callbacks.append(DropKeywords()) + await call(server, temperature=0.25, extra_headers={"x-deleted": "remove this header"}) + + assert_served_natively(server) + sent: Final = server.requests[0] + assert "temperature" not in _OBJECT.validate_python(sent.body) + assert sent.headers.get("x-deleted") is None diff --git a/tests/test_litellm_rust/chat_completions/test_failures.py b/tests/test_litellm_rust/chat_completions/test_failures.py new file mode 100644 index 00000000000..68a81096a7c --- /dev/null +++ b/tests/test_litellm_rust/chat_completions/test_failures.py @@ -0,0 +1,69 @@ +from typing import Final + +import pytest +from pydantic import TypeAdapter + +import litellm +from litellm import RateLimitError +from tests.test_litellm_rust.chat_completions.data import RATE_LIMITED, call +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec + +pytestmark = pytest.mark.requires_rust_extension + +_OBJECT: Final = TypeAdapter(dict[str, object]) + + +def failure_signal(asynchronous: bool) -> str: + return "async_log_failure_event" if asynchronous else "log_failure_event" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_sdk_policy_rejection_precedes_resource_setup_and_is_logged_once( + server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool +) -> None: + monkeypatch.setattr(litellm, "max_budget", 1.0) + monkeypatch.setattr(litellm, "_current_cost", 2.0) + monkeypatch.setattr(litellm, "ssl_verify", object()) + recorder: Final = RecordingLogger() + server.expected_requests = 0 + + with pytest.raises(litellm.BudgetExceededError) as caught: + await call(server, asynchronous=asynchronous, callbacks=[recorder]) + + failure: Final = await recorder.wait_for_async(failure_signal(asynchronous)) + assert len(failure) == 1 + assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value + assert not server.requests + assert not any("success" in name for name in recorder.names) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_provider_failure_is_terminal_and_shared_with_callbacks( + server: RecordingServer, asynchronous: bool +) -> None: + recorder: Final = RecordingLogger() + server.enqueue(ResponseSpec(body=RATE_LIMITED, status=429)) + + with pytest.raises(RateLimitError) as caught: + await call(server, asynchronous=asynchronous, callbacks=[recorder]) + + assert caught.value.status_code == 429 + assert len(server.requests) == 1 + failure: Final = await recorder.wait_for_async(failure_signal(asynchronous)) + assert len(failure) == 1 + assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value + assert not any("success" in name for name in recorder.names) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_empty_messages_fail_before_provider_io(server: RecordingServer, asynchronous: bool) -> None: + server.expected_requests = 0 + + with pytest.raises(Exception, match="chat completions requires at least one message"): + await call(server, asynchronous=asynchronous, messages=[]) + + assert not server.requests diff --git a/tests/test_litellm_rust/chat_completions/test_lifecycle.py b/tests/test_litellm_rust/chat_completions/test_lifecycle.py new file mode 100644 index 00000000000..f335d6fa509 --- /dev/null +++ b/tests/test_litellm_rust/chat_completions/test_lifecycle.py @@ -0,0 +1,37 @@ +import asyncio +from typing import Final + +import pytest + +import litellm +from litellm.rust_bridge import _native +from litellm.rust_bridge.public_call import NativeCall +from tests.test_litellm_rust.chat_completions.data import REQUEST +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.recording_server import RecordingServer + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.mark.asyncio +async def test_unstarted_call_has_no_provider_or_callback_effects( + server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + recorder: Final = RecordingLogger() + server.expected_requests = 0 + monkeypatch.setattr(litellm, "ssl_verify", object()) + monkeypatch.setattr(litellm, "max_budget", 1.0) + monkeypatch.setattr(litellm, "_current_cost", 2.0) + kwargs: Final = { + **REQUEST, + "api_key": "test-key", + "api_base": server.base_url, + "callbacks": [recorder], + } + + pending: Final = _native.acompletion(NativeCall(args=(), kwargs=kwargs, base={})) + assert asyncio.iscoroutine(pending) + pending.close() + + assert not server.requests + assert not recorder.events diff --git a/tests/test_litellm_rust/chat_completions/test_requests.py b/tests/test_litellm_rust/chat_completions/test_requests.py new file mode 100644 index 00000000000..3f2ba7c60e0 --- /dev/null +++ b/tests/test_litellm_rust/chat_completions/test_requests.py @@ -0,0 +1,42 @@ +import asyncio +from typing import Final + +import pytest +from pydantic import TypeAdapter + +import litellm +from litellm.rust_bridge.catalog import Route +from litellm.types.utils import ModelResponse +from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL, MESSAGES_TEXT +from tests.test_litellm_rust.support.recording_server import RecordingServer +from tests.test_litellm_rust.support.routing import assert_served_natively, native_route + +pytestmark = pytest.mark.requires_rust_extension + +_OBJECT: Final = TypeAdapter(dict[str, object]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_positional_parameters_reach_the_wire(server: RecordingServer, asynchronous: bool) -> None: + temperature: Final = 0.35 + arguments: Final = ( + (MESSAGES_MODEL, list(MESSAGES), None, None, 12.0, temperature) + if asynchronous + else (MESSAGES_MODEL, list(MESSAGES), 12.0, temperature) + ) + supplied: Final = {"base_url": server.base_url, "api_key": "test-key", "max_tokens": 32} + + with native_route(Route.CHAT_COMPLETIONS): + result: Final = ( + await litellm.acompletion(*arguments, **supplied) + if asynchronous + else await asyncio.to_thread(litellm.completion, *arguments, **supplied) + ) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == MESSAGES_TEXT + assert_served_natively(server) + body: Final = _OBJECT.validate_python(server.requests[0].body) + assert body["temperature"] == temperature + assert body["max_tokens"] == supplied["max_tokens"] diff --git a/tests/test_litellm_rust/chat_completions/test_settings.py b/tests/test_litellm_rust/chat_completions/test_settings.py new file mode 100644 index 00000000000..df64ab00421 --- /dev/null +++ b/tests/test_litellm_rust/chat_completions/test_settings.py @@ -0,0 +1,82 @@ +from typing import Final + +import pytest +from pydantic import TypeAdapter + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.models.credentials import CredentialItem +from litellm.types.utils import CallTypes, ModelResponse +from tests.test_litellm_rust.chat_completions.data import call +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.recording_server import RecordingServer +from tests.test_litellm_rust.support.routing import assert_served_natively + +pytestmark = pytest.mark.requires_rust_extension + +_OBJECT: Final = TypeAdapter(dict[str, object]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("from_credentials", (False, True)) +async def test_resource_setup_uses_deployment_hook_arguments( + server: RecordingServer, monkeypatch: pytest.MonkeyPatch, from_credentials: bool +) -> None: + invalid_settings: Final = {"ssl_verify": object()} + credential: Final = CredentialItem( + credential_name="resource-settings", credential_info={}, credential_values=invalid_settings + ) + monkeypatch.setattr(litellm, "credential_list", [credential]) + + class Prepare(CustomLogger): + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: CallTypes | None + ) -> dict[str, object]: + return { + **kwargs, + **({"litellm_credential_name": credential.credential_name} if from_credentials else invalid_settings), + } + + litellm.callbacks.append(Prepare()) + recorder: Final = RecordingLogger() + server.expected_requests = 0 + + with pytest.raises(ValueError, match=r"request\.ssl_verify") as caught: + await call(server, callbacks=[recorder]) + + failure: Final = await recorder.wait_for_async("async_log_failure_event") + assert len(failure) == 1 + assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value + assert not server.requests + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +@pytest.mark.parametrize("source", ("explicit", "base_url", "global", "provider", "environment", "empty")) +async def test_connection_settings_reach_the_provider( + server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool, source: str +) -> None: + key: Final = "selected-key" + explicit: Final = source in ("explicit", "base_url") + monkeypatch.setattr(litellm, "api_key", key if source in ("global", "empty") else ("unused" if explicit else None)) + monkeypatch.setattr(litellm, "anthropic_key", key if source == "provider" else ("unused" if explicit else None)) + monkeypatch.setattr( + litellm, + "api_base", + None if source == "environment" else ("http://127.0.0.1:1" if explicit else server.base_url), + ) + monkeypatch.setenv("ANTHROPIC_API_KEY", key if source == "environment" else "unused") + for name in ("ANTHROPIC_BASE_URL", "ANTHROPIC_API_BASE"): + monkeypatch.setenv(name, server.base_url if source == "environment" else "http://127.0.0.1:1") + + result: Final = await call( + server, + asynchronous=asynchronous, + api_key=key if explicit else ("" if source == "empty" else None), + api_base=server.base_url if source == "explicit" else ("" if source == "empty" else None), + **({"base_url": server.base_url} if source == "base_url" else {}), + ) + + assert isinstance(result, ModelResponse) + assert_served_natively(server) + assert server.requests[0].headers["x-api-key"] == key diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py index 2e965fe2630..d9406af4c59 100644 --- a/tests/test_litellm_rust/conftest.py +++ b/tests/test_litellm_rust/conftest.py @@ -23,12 +23,12 @@ from tests.test_litellm_rust.support.recording_server import RecordingServer, re @pytest_asyncio.fixture(autouse=True, loop_scope="function") -async def isolate_ocr_test_state() -> AsyncIterator[None]: +async def isolate_test_state() -> AsyncIterator[None]: with ExitStack() as stack: stack.enter_context(isolated_callback_registries()) stack.enter_context(rebound(litellm, "cache", None)) # test-quality-ok: isolate process-global cache stack.enter_context(rebound(_CONFIGURATION, "override", None)) - executor: Final = ThreadPoolExecutor(thread_name_prefix="rust-ocr-test-logging") + executor: Final = ThreadPoolExecutor(thread_name_prefix="rust-test-logging") stack.enter_context(rebound(litellm_logging, "executor", executor)) stack.enter_context(rebound(utils, "executor", executor)) stack.enter_context(rebound(thread_pool_executor, "executor", executor)) diff --git a/tests/test_litellm_rust/messages/conftest.py b/tests/test_litellm_rust/messages/conftest.py new file mode 100644 index 00000000000..25c46c8ae42 --- /dev/null +++ b/tests/test_litellm_rust/messages/conftest.py @@ -0,0 +1,10 @@ +import pytest + +from tests.test_litellm_rust.messages.data import MESSAGES_RESPONSE +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec + + +@pytest.fixture +def server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) + return recording_server diff --git a/tests/test_litellm_rust/messages/data.py b/tests/test_litellm_rust/messages/data.py new file mode 100644 index 00000000000..b7d039eb352 --- /dev/null +++ b/tests/test_litellm_rust/messages/data.py @@ -0,0 +1,64 @@ +from typing import Final + +import litellm +from litellm.rust_bridge.catalog import Route +from tests.test_litellm_rust.support.gaps import gap +from tests.test_litellm_rust.support.recording_server import RecordingServer +from tests.test_litellm_rust.support.routing import call_natively + +MESSAGES_MODEL: Final = "anthropic/claude-sonnet-5" +MESSAGES: Final = ({"role": "user", "content": "Hello"},) +MESSAGES_TEXT: Final = "Hello from native Messages" +MESSAGES_RESPONSE: Final = { + "id": "msg_native", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [{"type": "text", "text": MESSAGES_TEXT}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 4}, +} +MESSAGES_EVENTS: Final = ( + ("message_start", {"type": "message_start", "message": {**MESSAGES_RESPONSE, "content": [], "stop_reason": None}}), + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + ( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": MESSAGES_TEXT}}, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 4}, + }, + ), + ("message_stop", {"type": "message_stop"}), +) + +REQUEST: Final = {"model": MESSAGES_MODEL, "messages": list(MESSAGES), "max_tokens": 64} + + +def _error(kind: str) -> dict[str, object]: + return {"type": "error", "error": {"type": kind, "message": kind}} + + +BAD_REQUEST: Final = _error("invalid_request_error") +RATE_LIMITED: Final = _error("rate_limit_error") + + +async def call(server: RecordingServer, *, asynchronous: bool = True, **kwargs: object) -> object: + return await call_natively( + Route.MESSAGES, + litellm.anthropic.messages.create, + litellm.anthropic.messages.acreate, + REQUEST, + server, + asynchronous=asynchronous, + **kwargs, + ) + + +BODY_HOOKS: Final = gap("native Messages does not run the handler body hooks") diff --git a/tests/test_litellm_rust/messages/test_body_hooks.py b/tests/test_litellm_rust/messages/test_body_hooks.py new file mode 100644 index 00000000000..188cd58e00e --- /dev/null +++ b/tests/test_litellm_rust/messages/test_body_hooks.py @@ -0,0 +1,137 @@ +from collections.abc import Iterator +from datetime import datetime +from typing import Final + +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.anthropic.pass_through.messages import interceptors +from litellm.llms.anthropic.pass_through.messages.interceptors.base import MessagesInterceptor +from tests.test_litellm_rust.messages.data import BODY_HOOKS, MESSAGES_RESPONSE, REQUEST, call +from tests.test_litellm_rust.support.callback_recorder import drain_logging +from tests.test_litellm_rust.support.body_hooks import LegacyAgenticLoop, PlannedAgenticLoop, agentic_hooks +from tests.test_litellm_rust.support.isolation import rebound +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed +from tests.test_litellm_rust.support.results import response_id +from tests.test_litellm_rust.support.routing import CALL_ID, assert_served_natively + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.fixture +def server(server: RecordingServer) -> Iterator[RecordingServer]: + yield from relaxed(server) + + +@pytest.mark.asyncio +@BODY_HOOKS +async def test_body_pre_request_hook_rewrites_the_provider_request(server: RecordingServer) -> None: + class Rewrite(CustomLogger): + async def async_pre_request_hook( + self, model: str, messages: list[dict[str, object]], kwargs: dict[str, object] + ) -> dict[str, object]: + return {**kwargs, "system": "body hook instruction"} + + with rebound(litellm, "callbacks", [Rewrite()]): + await call(server) + + assert_served_natively(server) + assert isinstance(server.requests[0].body, dict) + assert server.requests[0].body.get("system") == "body hook instruction" + + +@pytest.mark.asyncio +@BODY_HOOKS +async def test_body_websearch_short_circuit_answers_without_the_provider(server: RecordingServer) -> None: + class ShortCircuit(WebSearchInterceptionLogger): + async def try_short_circuit_search( + self, + model: str, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + custom_llm_provider: str | None, + kwargs: dict[str, object] | None = None, + ) -> object: + return {**MESSAGES_RESPONSE, "id": "short-circuited"} + + with rebound(litellm, "callbacks", [ShortCircuit()]): + response: Final = await call(server) + + assert response_id(response) == "short-circuited" + assert server.requests == [] + + +@pytest.mark.asyncio +@BODY_HOOKS +async def test_body_interceptor_replaces_the_provider_call(server: RecordingServer) -> None: + class Replace(MessagesInterceptor): + def can_handle(self, tools: list[dict[str, object]] | None, custom_llm_provider: str | None) -> bool: + return True + + async def handle(self, **kwargs: object) -> object: + return {**MESSAGES_RESPONSE, "id": "intercepted"} + + with rebound(interceptors, "_interceptors", [Replace()]): + response: Final = await call(server) + + assert response_id(response) == "intercepted" + assert server.requests == [] + + +@pytest.mark.asyncio +@BODY_HOOKS +async def test_body_planned_agentic_loop_sends_the_followup_and_replaces_the_response(server: RecordingServer) -> None: + loop: Final = PlannedAgenticLoop() + server.enqueue(ResponseSpec(body=MESSAGES_RESPONSE)) + server.default_response = ResponseSpec(body={**MESSAGES_RESPONSE, "id": "followup"}) + + with rebound(litellm, "callbacks", [loop]): + response: Final = await call(server) + await drain_logging() + + assert_served_natively(server, count=2) + assert isinstance(server.requests[1].body, dict) + assert server.requests[1].body["messages"] == [{"role": "user", "content": "followup"}] + assert response_id(response) == "agentic-replaced" + assert agentic_hooks(loop) == ( + "async_should_run_agentic_loop", + "async_build_agentic_loop_plan", + "async_should_run_agentic_loop", + "async_post_agentic_loop_response_hook", + ) + + +@pytest.mark.asyncio +@BODY_HOOKS +async def test_body_legacy_agentic_loop_replaces_the_response(server: RecordingServer) -> None: + loop: Final = LegacyAgenticLoop() + + with rebound(litellm, "callbacks", [loop]): + response: Final = await call(server) + + assert_served_natively(server) + assert response_id(response) == "agentic-replaced" + assert agentic_hooks(loop) == ("async_should_run_agentic_loop", "async_run_agentic_loop") + + +@pytest.mark.asyncio +async def test_body_logging_setup_drops_a_logger_fn_from_a_supplied_logging_object(server: RecordingServer) -> None: + observed: Final[list[object]] = [] + logging_obj: Final = Logging( + model=REQUEST["model"], + messages=[], + stream=False, + start_time=datetime.now(), + call_type="anthropic_messages", + litellm_call_id=CALL_ID, + function_id=CALL_ID, + ) + logging_obj.update_environment_variables(litellm_params={"logger_fn": observed.append}, optional_params={}) + + await call(server, litellm_logging_obj=logging_obj) + + assert_served_natively(server) + assert observed == [] diff --git a/tests/test_litellm_rust/messages/test_callbacks.py b/tests/test_litellm_rust/messages/test_callbacks.py index 6c317c3b9d9..e8d49453065 100644 --- a/tests/test_litellm_rust/messages/test_callbacks.py +++ b/tests/test_litellm_rust/messages/test_callbacks.py @@ -1,197 +1,138 @@ -from collections.abc import AsyncIterator, Iterator +import asyncio +from collections.abc import Iterator +from contextvars import ContextVar +from datetime import datetime from typing import Final import pytest +from pydantic import TypeAdapter import litellm +from litellm.integrations.argilla import ArgillaLogger from litellm.integrations.custom_logger import CustomLogger -from litellm.rust_bridge import catalog -from litellm.rust_bridge.catalog import Route, RouteRule -from litellm.rust_bridge.configuration import Rollout -from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +from litellm.types.integrations.argilla import ArgillaItem +from litellm.types.utils import ModelResponse, StandardLoggingPayload +from tests.test_litellm_rust.messages.data import BAD_REQUEST, MESSAGES_RESPONSE, RATE_LIMITED, call +from tests.test_litellm_rust.support.callback_recorder import HookRecorder, RecordingLogger, request_body +from tests.test_litellm_rust.support.gaps import gap from tests.test_litellm_rust.support.isolation import rebound -from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec -from tests.test_litellm_rust.support.requests import ( - MESSAGES, - MESSAGES_EVENTS, - MESSAGES_MODEL, - MESSAGES_RESPONSE, - request_body, -) -from tests.test_litellm_rust.support.served import assert_served_natively +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed +from tests.test_litellm_rust.support.results import response_id, with_id +from tests.test_litellm_rust.support.routing import assert_served_natively pytestmark = pytest.mark.requires_rust_extension -STREAM: Final = ResponseSpec(body=None, events=MESSAGES_EVENTS) - - -@pytest.fixture(autouse=True) -def opt_messages_into_rust() -> Iterator[None]: - with rebound(catalog, "RULES", (RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN), *catalog.RULES)): - yield +SYNC_HANDLERS_ON_ASYNC_CALLS: Final = gap("native async calls also run the sync success and failure handlers") +_OBJECT: Final = TypeAdapter(dict[str, object]) +_HEADERS: Final = TypeAdapter(dict[str, str]) @pytest.fixture -def messages_server(recording_server: RecordingServer) -> RecordingServer: - recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) - return recording_server - - -def arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]: - return { - "model": MESSAGES_MODEL, - "messages": [dict(message) for message in MESSAGES], - "max_tokens": 64, - "api_key": "test-key", - "api_base": server.base_url, - **kwargs, - } +def server(server: RecordingServer) -> Iterator[RecordingServer]: + yield from relaxed(server) @pytest.mark.asyncio -async def test_native_messages_callbacks_see_the_provider_request_and_the_public_response( - messages_server: RecordingServer, +@pytest.mark.parametrize( + "hook,expected", + ( + pytest.param("log_pre_api_call", 1), + pytest.param("log_post_api_call", 1), + pytest.param("async_log_success_event", 1), + pytest.param("async_logging_hook", 1), + pytest.param("redact_standard_logging_payload_from_model_call_details", 1, id="redact"), + pytest.param("async_pre_call_deployment_hook", 1), + pytest.param("async_post_call_success_deployment_hook", 1), + pytest.param("log_success_event", 0), + pytest.param("logging_hook", 0), + ), +) +async def test_success_hook_fires_the_expected_number_of_times( + server: RecordingServer, hook: str, expected: int ) -> None: + recorder: Final = HookRecorder() + + with rebound(litellm, "callbacks", [recorder]): + response: Final = await call(server) + await recorder.wait_for_async("async_log_success_event") + + assert_served_natively(server, response) + assert recorder.names.count(hook) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "hook,expected", + ( + pytest.param("async_log_failure_event", 1), + pytest.param("async_post_call_failure_deployment_hook", 1), + pytest.param("log_model_group_rate_limit_error", 1), + pytest.param("log_failure_event", 0, marks=SYNC_HANDLERS_ON_ASYNC_CALLS), + ), +) +async def test_failure_hook_fires_the_expected_number_of_times( + server: RecordingServer, hook: str, expected: int +) -> None: + server.default_response = ResponseSpec(body=RATE_LIMITED, status=429) + recorder: Final = HookRecorder() + + with rebound(litellm, "callbacks", [recorder]): + with pytest.raises(litellm.RateLimitError): + await call(server, metadata={"model_group_size": 1}) + await recorder.wait_for_async("async_log_failure_event") + + assert_served_natively(server) + assert recorder.names.count(hook) == expected + + +@pytest.mark.asyncio +async def test_failure_events_carry_the_error_raised_to_the_caller(server: RecordingServer) -> None: + server.default_response = ResponseSpec(body=BAD_REQUEST, status=400) recorder: Final = RecordingLogger() - response: Final = await litellm.anthropic.messages.acreate( - **arguments(messages_server, callbacks=[recorder], litellm_call_id="messages-success") - ) + with rebound(litellm, "callbacks", [recorder]): + with pytest.raises(litellm.BadRequestError) as raised: + await call(server) + events: Final = await recorder.wait_for_async("async_log_failure_event") - assert_served_natively(messages_server, response) - assert response["content"] == MESSAGES_RESPONSE["content"] - sent: Final = messages_server.requests[0] - assert sent.path == "/v1/messages" - assert sent.body == {"model": "claude-sonnet-5", "messages": list(MESSAGES), "max_tokens": 64, "stream": False} - pre_call: Final = recorder.wait_for("log_pre_api_call") - assert request_body(pre_call[0].kwargs) == sent.body - success: Final = await recorder.wait_for_async("async_log_success_event") - assert len(success) == 1 - assert success[0].call_type == "anthropic_messages" - assert success[0].kwargs["litellm_call_id"] == "messages-success" - assert success[0].response.choices[0].message.content == "Hello from native Messages" + assert_served_natively(server) + assert isinstance(events[0].kwargs, dict) + assert events[0].kwargs["exception"] is raised.value @pytest.mark.asyncio -async def test_native_messages_pre_call_body_edit_reaches_the_provider(messages_server: RecordingServer) -> None: +async def test_pre_api_call_sees_the_wire_body(server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + + response: Final = await call(server, callbacks=[recorder]) + + assert_served_natively(server, response) + pre_call: Final = recorder.wait_for("log_pre_api_call")[0].kwargs + assert isinstance(pre_call, dict) + assert request_body(pre_call) == server.requests[0].body + + +@pytest.mark.asyncio +async def test_pre_api_call_edits_reach_the_provider(server: RecordingServer) -> None: class Edit(CustomLogger): - def log_pre_api_call(self, model, messages, kwargs): + def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None: request_body(kwargs)["temperature"] = 0.25 - await litellm.anthropic.messages.acreate(**arguments(messages_server, callbacks=[Edit()])) + response: Final = await call(server, callbacks=[Edit()]) - assert messages_server.requests[0].body["temperature"] == 0.25 + assert_served_natively(server, response) + assert isinstance(server.requests[0].body, dict) + assert server.requests[0].body["temperature"] == 0.25 @pytest.mark.asyncio -async def test_native_messages_provider_error_reaches_caller_and_failure_callbacks_as_one_public_error( - messages_server: RecordingServer, -) -> None: - messages_server.enqueue( - ResponseSpec(body={"type": "error", "error": {"type": "invalid_request_error", "message": "bad"}}, status=400) - ) - observed: Final = [] - - class Observe(CustomLogger): - def log_failure_event(self, kwargs, response_obj, start_time, end_time): - observed.append(("sync", kwargs["exception"])) - - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - observed.append(("async", kwargs["exception"])) - - with pytest.raises(litellm.BadRequestError) as raised: - await litellm.anthropic.messages.acreate(**arguments(messages_server, callbacks=[Observe()])) - - assert len(messages_server.requests) == 1 - assert [phase for phase, _ in observed] == ["sync", "async"] - assert all(error is raised.value for _, error in observed) - - -def sse_payload() -> bytes: - return b"".join(STREAM.payloads()) - - -@pytest.mark.asyncio -async def test_native_messages_stream_relays_provider_events_and_logs_success_once_after_the_last_chunk( - messages_server: RecordingServer, -) -> None: - messages_server.enqueue(STREAM) +async def test_pre_api_call_sees_the_shaped_optional_params(server: RecordingServer) -> None: recorder: Final = RecordingLogger() - stream: Final = await litellm.anthropic.messages.acreate( - **arguments(messages_server, stream=True, callbacks=[recorder]) - ) - assert isinstance(stream, AsyncIterator) - first: Final = await anext(stream) - await drain_logging() - assert "async_log_success_event" not in recorder.names - rest: Final = [chunk async for chunk in stream] + await call(server, callbacks=[recorder], temperature=0.2, top_k=3, drop_params=True) - assert first + b"".join(rest) == sse_payload() - assert_served_natively(messages_server, stream) - assert messages_server.requests[0].body["stream"] is True - success: Final = await recorder.wait_for_async("async_log_success_event") - assert len(success) == 1 - assert success[0].kwargs["stream"] is True - assert success[0].kwargs["completion_start_time"] is not None - assert "log_failure_event" not in recorder.names - - -@pytest.mark.asyncio -async def test_native_messages_stream_closed_early_logs_success_once_for_what_was_delivered( - messages_server: RecordingServer, -) -> None: - messages_server.enqueue(STREAM) - recorder: Final = RecordingLogger() - - stream: Final = await litellm.anthropic.messages.acreate( - **arguments(messages_server, stream=True, callbacks=[recorder]) - ) - assert isinstance(stream, AsyncIterator) - await anext(stream) - await stream.aclose() - - success: Final = await recorder.wait_for_async("async_log_success_event") - assert len(success) == 1 - with pytest.raises(StopAsyncIteration): - await anext(stream) - - -def test_native_sync_messages_stream_relays_provider_events_and_logs_success_once( - messages_server: RecordingServer, -) -> None: - messages_server.enqueue(STREAM) - recorder: Final = RecordingLogger() - - stream: Final = litellm.anthropic.messages.create(**arguments(messages_server, stream=True, callbacks=[recorder])) - assert isinstance(stream, Iterator) - - assert b"".join(stream) == sse_payload() - assert_served_natively(messages_server, stream) - assert len(recorder.wait_for("async_log_success_event")) == 1 - - -def test_native_sync_messages_returns_the_provider_message(messages_server: RecordingServer) -> None: - recorder: Final = RecordingLogger() - - response: Final = litellm.anthropic.messages.create(**arguments(messages_server, callbacks=[recorder])) - - assert_served_natively(messages_server, response) - assert response["content"] == MESSAGES_RESPONSE["content"] - assert len(recorder.wait_for("log_success_event")) == 1 - - -@pytest.mark.asyncio -async def test_native_messages_pre_call_sees_the_shaped_optional_params( - messages_server: RecordingServer, -) -> None: - recorder: Final = RecordingLogger() - - await litellm.anthropic.messages.acreate( - **arguments(messages_server, callbacks=[recorder], temperature=0.2, top_k=3, drop_params=True) - ) - - sent: Final = messages_server.requests[0].body + sent: Final = server.requests[0].body + assert isinstance(sent, dict) assert not {"temperature", "top_k"} & sent.keys() pre_call: Final = recorder.wait_for("log_pre_api_call")[0].kwargs assert isinstance(pre_call, dict) @@ -202,65 +143,287 @@ async def test_native_messages_pre_call_sees_the_shaped_optional_params( @pytest.mark.asyncio -async def test_native_messages_failing_pre_call_logger_does_not_fail_the_call(messages_server: RecordingServer) -> None: - class Broken(CustomLogger): - def log_pre_api_call(self, model, messages, kwargs): - raise RuntimeError("logger exploded") +async def test_logging_hook_and_redaction_edits_reach_the_async_success_event(server: RecordingServer) -> None: + class Rewrite(HookRecorder): + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + return {**kwargs, "logging_hook_ran": True}, result - response: Final = await litellm.anthropic.messages.acreate(**arguments(messages_server, callbacks=[Broken()])) + def redact_standard_logging_payload_from_model_call_details( + self, model_call_details: dict[str, object] + ) -> dict[str, object]: + return {**model_call_details, "redacted": True} - assert_served_natively(messages_server, response) - assert response["content"] == MESSAGES_RESPONSE["content"] + recorder: Final = Rewrite() + + with rebound(litellm, "callbacks", [recorder]): + await call(server) + success: Final = await recorder.wait_for_async("async_log_success_event") + + assert isinstance(success[0].kwargs, dict) + assert success[0].kwargs["logging_hook_ran"] is True + assert success[0].kwargs["redacted"] is True @pytest.mark.asyncio -async def test_native_messages_stream_success_log_carries_usage_rebuilt_from_the_relayed_events( - messages_server: RecordingServer, +async def test_pre_call_deployment_hooks_compose_the_wire_request_in_registration_order( + server: RecordingServer, ) -> None: - messages_server.enqueue(STREAM) + class Rewrite(CustomLogger): + def __init__(self, marker: str) -> None: + super().__init__() + self.marker = marker + + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: object + ) -> dict[str, object]: + headers: Final = _HEADERS.validate_python(kwargs["extra_headers"]) + return { + **kwargs, + "extra_headers": {**headers, "x-deployment-hook": headers["x-deployment-hook"] + self.marker}, + } + + with rebound(litellm, "callbacks", [Rewrite("first"), Rewrite("-second")]): + response: Final = await call(server, extra_headers={"x-deployment-hook": "caller-", "x-kept": "unchanged"}) + + assert_served_natively(server, response) + assert server.requests[0].headers["x-deployment-hook"] == "caller-first-second" + assert server.requests[0].headers["x-kept"] == "unchanged" + + +@pytest.mark.asyncio +async def test_pre_call_deployment_rejection_logs_failure_without_provider_io(server: RecordingServer) -> None: + rejection: Final = litellm.BadRequestError( + message="deployment rejected", model="test-model", llm_provider="anthropic" + ) recorder: Final = RecordingLogger() - stream: Final = await litellm.anthropic.messages.acreate( - **arguments(messages_server, stream=True, callbacks=[recorder]) - ) - assert isinstance(stream, AsyncIterator) - async for _ in stream: - pass + class Reject(CustomLogger): + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: object + ) -> dict[str, object]: + raise rejection - success: Final = await recorder.wait_for_async("async_log_success_event") - usage: Final = success[0].response.usage - assert usage.completion_tokens == MESSAGES_EVENTS[4][1]["usage"]["output_tokens"] - assert usage.prompt_tokens == MESSAGES_RESPONSE["usage"]["input_tokens"] - assert success[0].response.choices[0].message.content == "Hello from native Messages" + with rebound(litellm, "callbacks", [Reject(), recorder]): + with pytest.raises(litellm.BadRequestError, match="deployment rejected") as raised: + await call(server) + failures: Final = await recorder.wait_for_async("async_log_failure_event") + + assert raised.value is rejection + assert server.requests == [] + assert len(failures) == 1 + assert _OBJECT.validate_python(failures[0].kwargs)["exception"] is raised.value + assert not {"log_pre_api_call", "log_success_event", "async_log_success_event"} & set(recorder.names) @pytest.mark.asyncio -async def test_native_messages_does_not_restore_keywords_deleted_by_a_pre_call_hook( - messages_server: RecordingServer, -) -> None: +async def test_pre_call_deployment_hook_can_delete_keywords(server: RecordingServer) -> None: class DropKeywords(CustomLogger): async def async_pre_call_deployment_hook( self, kwargs: dict[str, object], call_type: object ) -> dict[str, object]: return {name: value for name, value in kwargs.items() if name not in ("system", "extra_headers")} - call_arguments: Final = arguments( - messages_server, - system="remove this instruction", - extra_headers={"x-deleted": "remove this header"}, - headers={"x-kept": "keep this header"}, - stream=False, - ) with rebound(litellm, "callbacks", [DropKeywords()]): - response: Final = await litellm.anthropic.messages.acreate(**call_arguments) + response: Final = await call( + server, + system="remove this instruction", + extra_headers={"x-deleted": "remove this header"}, + headers={"x-kept": "keep this header"}, + ) - assert_served_natively(messages_server, response) - sent: Final = messages_server.requests[0] - assert sent.body == { - "model": MESSAGES_MODEL.removeprefix("anthropic/"), - "messages": call_arguments["messages"], - "max_tokens": call_arguments["max_tokens"], - "stream": call_arguments["stream"], - } - assert sent.headers.get("x-deleted") is None - assert sent.headers.get("x-kept") == "keep this header" + assert_served_natively(server, response) + sent: Final = server.requests[0] + assert isinstance(sent.body, dict) + assert "system" not in sent.body + assert "x-deleted" not in sent.headers + assert sent.headers["x-kept"] == "keep this header" + + +@pytest.mark.asyncio +async def test_post_call_success_deployment_hook_replaces_the_response(server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + + class Replace(CustomLogger): + async def async_post_call_success_deployment_hook( + self, request_data: dict[str, object], response: object, call_type: object + ) -> object: + return with_id(response, "replaced") + + with rebound(litellm, "callbacks", [Replace(), recorder]): + response: Final = await call(server) + success: Final = await recorder.wait_for_async("async_log_success_event") + + assert_served_natively(server, response) + assert response_id(response) == "replaced" + assert len(success) == 1 + assert isinstance(success[0].response, ModelResponse) + assert success[0].response.id == response_id(response) + + +@pytest.mark.asyncio +async def test_concurrent_calls_keep_callback_context_metadata_and_responses_isolated(server: RecordingServer) -> None: + context: Final = ContextVar("messages-callback-context", default="outside") + first_entered: Final = asyncio.Event() + second_entered: Final = asyncio.Event() + server.enqueue(ResponseSpec(body={**MESSAGES_RESPONSE, "id": "msg_first"})) + server.enqueue(ResponseSpec(body={**MESSAGES_RESPONSE, "id": "msg_second"})) + + class Observe(RecordingLogger): + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: object + ) -> dict[str, object]: + arrived, other = ( + (first_entered, second_entered) + if kwargs["litellm_call_id"] == "first" + else (second_entered, first_entered) + ) + arrived.set() + await asyncio.wait_for(other.wait(), timeout=5) + return kwargs + + def log_pre_api_call(self, model: object, messages: object, kwargs: dict[str, object]) -> None: + self._record("log_pre_api_call", {**kwargs, "caller_context": context.get()}) + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self._record("async_log_success_event", {**kwargs, "caller_context": context.get()}, response_obj) + + recorder: Final = Observe() + + async def invoke(marker: str) -> object: + token: Final = context.set(marker) + try: + return await call(server, litellm_call_id=marker, metadata={"marker": marker}) + finally: + context.reset(token) + + with rebound(litellm, "callbacks", [recorder]): + responses: Final = await asyncio.gather(invoke("first"), invoke("second")) + success: Final = await recorder.wait_for_async("async_log_success_event", count=2) + + assert_served_natively(server, *responses) + expected: Final = {"first": responses[0], "second": responses[1]} + pre_call: Final = recorder.wait_for("log_pre_api_call", count=2) + assert len(pre_call) == len(success) == 2 + assert {_OBJECT.validate_python(event.kwargs)["litellm_call_id"] for event in success} == expected.keys() + for event in (*pre_call, *success): + details: Final = _OBJECT.validate_python(event.kwargs) + call_id: Final = str(details["litellm_call_id"]) + params: Final = _OBJECT.validate_python(details["litellm_params"]) + assert details["caller_context"] == call_id + assert _OBJECT.validate_python(params["metadata"])["marker"] == call_id + for event in success: + logged_call_id: Final = str(_OBJECT.validate_python(event.kwargs)["litellm_call_id"]) + assert isinstance(event.response, ModelResponse) + assert event.response.id == response_id(expected[logged_call_id]) + assert context.get() == "outside" + + +@pytest.mark.asyncio +async def test_a_failing_logger_does_not_fail_the_call(server: RecordingServer) -> None: + class Broken(CustomLogger): + def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None: + raise RuntimeError("logger exploded") + + response: Final = await call(server, callbacks=[Broken()]) + + assert_served_natively(server, response) + assert response_id(response) == MESSAGES_RESPONSE["id"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "event,asynchronous,dynamic,expected", + ( + pytest.param("input", False, False, 1, id="global-input"), + pytest.param("success", True, False, 1, id="global-async-success"), + pytest.param("success", True, True, 1, id="request-async-success"), + pytest.param("success", False, False, 0, id="global-sync-success", marks=SYNC_HANDLERS_ON_ASYNC_CALLS), + pytest.param("success", False, True, 0, id="request-sync-success", marks=SYNC_HANDLERS_ON_ASYNC_CALLS), + pytest.param("failure", True, False, 1, id="global-async-failure"), + pytest.param("failure", False, False, 1, id="global-sync-failure"), + pytest.param("failure", False, True, 1, id="request-sync-failure"), + ), +) +async def test_function_callback_fires_the_expected_number_of_times( + server: RecordingServer, event: str, asynchronous: bool, dynamic: bool, expected: int +) -> None: + recorder: Final = RecordingLogger() + + def callback(kwargs: dict[str, object], *_: object) -> None: + recorder._record(event, kwargs) + + async def async_callback(kwargs: dict[str, object], *_: object) -> None: + recorder._record(event, kwargs) + + registered: Final = async_callback if asynchronous else callback + global_list: Final = f"_async_{event}_callback" if asynchronous else f"{event}_callback" + request_callbacks: Final = {f"{event}_callback": [registered]} if dynamic else {} + if event == "failure": + server.default_response = ResponseSpec(body=BAD_REQUEST, status=400) + + with rebound(litellm, "callbacks", [recorder]), rebound(litellm, global_list, [] if dynamic else [registered]): + try: + await call(server, **request_callbacks) + except litellm.BadRequestError: + assert event == "failure" + await recorder.wait_for_async(f"async_log_{'failure' if event == 'failure' else 'success'}_event") + + assert_served_natively(server) + assert recorder.names.count(event) == expected + + +@pytest.mark.asyncio +async def test_pre_call_rule_rejects_before_provider_io(server: RecordingServer) -> None: + checked: Final[list[str]] = [] + + def reject(text: str) -> bool: + checked.append(text) + return False + + with rebound(litellm, "pre_call_rules", [reject]): + with pytest.raises(litellm.APIResponseValidationError): + await call(server) + + assert checked == ["Hello"] + assert server.requests == [] + + +@pytest.mark.asyncio +async def test_dataset_hook_returning_none_drops_the_record_from_the_success_sink(server: RecordingServer) -> None: + datasets: Final = RecordingLogger() + + class Drop(CustomLogger): + async def async_dataset_hook( + self, logged_data: ArgillaItem, standard_logging_payload: StandardLoggingPayload | None + ) -> None: + datasets._record("async_dataset_hook", response=logged_data) + + server.enqueue(ResponseSpec(body={"items": [{"id": "callback-dataset"}]})) + with rebound(litellm, "argilla_transformation_object", {"prompt": "messages", "answer": "response"}): + sink: Final = ArgillaLogger( + argilla_api_key="test-key", + argilla_base_url=server.base_url, + argilla_dataset_name="callback-dataset", + flush_interval=3600, + ) + with rebound(litellm, "callbacks", [sink, Drop()]): + await call(server) + await datasets.wait_for_async("async_dataset_hook") + await sink.flush_queue() + + assert [request.method for request in server.requests] == ["GET", "POST"] + + +@pytest.mark.asyncio +async def test_sync_call_returns_the_provider_message(server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + + response: Final = await call(server, asynchronous=False, callbacks=[recorder]) + + assert_served_natively(server, response) + assert response_id(response) == MESSAGES_RESPONSE["id"] + assert len(await recorder.wait_for_async("log_success_event")) == 1 diff --git a/tests/test_litellm_rust/messages/test_request_shaping.py b/tests/test_litellm_rust/messages/test_request_shaping.py deleted file mode 100644 index 3ae17e0f614..00000000000 --- a/tests/test_litellm_rust/messages/test_request_shaping.py +++ /dev/null @@ -1,311 +0,0 @@ -"""The native Messages route shapes the wire request the way the Python handler does. - -Model capability expectations come from model_prices_and_context_window.json (Claude Sonnet 5 is an -adaptive-thinking model without sampling params; Claude Haiku 4.5 is a legacy-thinking model), read at -2026-09-24; the cost map is LiteLLM's own file. -""" - -from collections.abc import Iterator -from typing import Final - -import pytest - -import litellm -from litellm.rust_bridge import catalog -from litellm.rust_bridge.catalog import Route, RouteRule -from litellm.rust_bridge.configuration import Rollout -from tests.test_litellm_rust.support.isolation import rebound -from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec -from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_RESPONSE -from tests.test_litellm_rust.support.served import assert_served_natively - -pytestmark = pytest.mark.requires_rust_extension - -ADAPTIVE_MODEL: Final = "anthropic/claude-sonnet-5" -LEGACY_THINKING_MODEL: Final = "anthropic/claude-haiku-4-5" - - -@pytest.fixture(autouse=True) -def opt_messages_into_rust() -> Iterator[None]: - with rebound(catalog, "RULES", (RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN), *catalog.RULES)): - yield - - -@pytest.fixture -def messages_server(recording_server: RecordingServer) -> RecordingServer: - recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) - return recording_server - - -def arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]: - return { - "model": ADAPTIVE_MODEL, - "messages": [dict(message) for message in MESSAGES], - "max_tokens": 8192, - "api_key": "test-key", - "api_base": server.base_url, - **kwargs, - } - - -def sent(server: RecordingServer, response: object) -> tuple[dict[str, object], dict[str, str]]: - assert_served_natively(server, response) - request: Final = server.requests[0] - assert isinstance(request.body, dict) - return request.body, request.headers - - -@pytest.mark.asyncio -async def test_reasoning_effort_becomes_adaptive_thinking_and_effort_on_the_wire( - messages_server: RecordingServer, -) -> None: - response: Final = await litellm.anthropic.messages.acreate(**arguments(messages_server, reasoning_effort="high")) - - body, _ = sent(messages_server, response) - assert "reasoning_effort" not in body - assert body["thinking"] == {"type": "adaptive", "display": "summarized"} - assert body["output_config"] == {"effort": "high"} - - -@pytest.mark.asyncio -async def test_claude_code_adaptive_payload_is_downgraded_to_a_capped_budget_for_a_legacy_model( - messages_server: RecordingServer, -) -> None: - response: Final = await litellm.anthropic.messages.acreate( - **arguments( - messages_server, - model=LEGACY_THINKING_MODEL, - max_tokens=3000, - thinking={"type": "adaptive"}, - output_config={"effort": "high"}, - temperature=0, - ) - ) - - body, _ = sent(messages_server, response) - assert body["thinking"] == {"type": "enabled", "budget_tokens": 2999} - assert "output_config" not in body - assert "temperature" not in body - - -@pytest.mark.asyncio -async def test_removed_sampling_params_are_dropped_under_drop_params(messages_server: RecordingServer) -> None: - response: Final = await litellm.anthropic.messages.acreate( - **arguments(messages_server, temperature=0.2, top_p=0.9, top_k=5, drop_params=True) - ) - - body, _ = sent(messages_server, response) - assert not {"temperature", "top_p", "top_k"} & body.keys() - - -@pytest.mark.asyncio -async def test_removed_sampling_params_are_rejected_without_drop_params(messages_server: RecordingServer) -> None: - messages_server.expected_requests = 0 - - with pytest.raises(litellm.BadRequestError, match="does not support top_k=5"): - await litellm.anthropic.messages.acreate(**arguments(messages_server, top_k=5)) - - assert messages_server.requests == [] - - -@pytest.mark.asyncio -async def test_replayed_history_is_sanitized_before_it_reaches_the_provider( - messages_server: RecordingServer, -) -> None: - history: Final = [ - {"role": "user", "content": "run it"}, - { - "role": "assistant", - "content": [ - {"type": "text", "text": ""}, - { - "type": "tool_use", - "id": "functions.Bash:0", - "name": "Bash", - "input": {}, - "provider_specific_fields": {"x": 1}, - }, - ], - }, - {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions.Bash:0", "content": "ok"}]}, - ] - - response: Final = await litellm.anthropic.messages.acreate(**arguments(messages_server, messages=history)) - - body, _ = sent(messages_server, response) - assert body["messages"] == [ - {"role": "user", "content": "run it"}, - {"role": "assistant", "content": [{"type": "tool_use", "id": "functions_Bash_0", "name": "Bash", "input": {}}]}, - {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions_Bash_0", "content": "ok"}]}, - ] - - -@pytest.mark.asyncio -async def test_feature_betas_merge_into_the_forwarded_beta_header(messages_server: RecordingServer) -> None: - response: Final = await litellm.anthropic.messages.acreate( - **arguments( - messages_server, - output_format={"type": "json_schema", "schema": {"type": "object"}}, - extra_headers={"anthropic-beta": "web-search-2025-03-05"}, - ) - ) - - _, headers = sent(messages_server, response) - assert headers["anthropic-beta"] == "structured-outputs-2025-11-13,web-search-2025-03-05" - - -@pytest.mark.asyncio -async def test_oauth_token_authenticates_as_a_bearer_with_the_oauth_beta(messages_server: RecordingServer) -> None: - response: Final = await litellm.anthropic.messages.acreate( - **arguments(messages_server, api_key="sk-ant-oat01-token") - ) - - _, headers = sent(messages_server, response) - assert "x-api-key" not in headers - assert headers["authorization"] == "Bearer sk-ant-oat01-token" - assert headers["anthropic-beta"] == "oauth-2025-04-20" - assert headers["anthropic-dangerous-direct-browser-access"] == "true" - - -@pytest.mark.asyncio -async def test_metadata_is_reduced_to_the_fields_anthropic_accepts(messages_server: RecordingServer) -> None: - response: Final = await litellm.anthropic.messages.acreate( - **arguments(messages_server, metadata={"user_id": "u-1", "trace_id": "internal"}) - ) - - body, _ = sent(messages_server, response) - assert body["metadata"] == {"user_id": "u-1"} - - -@pytest.mark.asyncio -async def test_additional_drop_params_remove_nested_fields_from_the_wire(messages_server: RecordingServer) -> None: - tools: Final = [{"name": "lookup", "input_schema": {"type": "object"}, "input_examples": [{"q": "x"}]}] - - response: Final = await litellm.anthropic.messages.acreate( - **arguments(messages_server, tools=tools, additional_drop_params=["tools[*].input_examples"]) - ) - - body, _ = sent(messages_server, response) - assert body["tools"] == [{"name": "lookup", "input_schema": {"type": "object"}}] - - -@pytest.mark.asyncio -async def test_provider_specific_headers_scoped_to_anthropic_reach_the_wire(messages_server: RecordingServer) -> None: - response: Final = await litellm.anthropic.messages.acreate( - **arguments( - messages_server, - provider_specific_header=[ - {"custom_llm_provider": "anthropic, azure_ai", "extra_headers": {"x-scoped": "yes"}}, - {"custom_llm_provider": "openai", "extra_headers": {"x-other": "no"}}, - ], - ) - ) - - _, headers = sent(messages_server, response) - assert headers["x-scoped"] == "yes" - assert "x-other" not in headers - - -@pytest.mark.asyncio -async def test_scoped_headers_override_extra_headers_which_override_forwarded_headers( - messages_server: RecordingServer, -) -> None: - response: Final = await litellm.anthropic.messages.acreate( - **arguments( - messages_server, - headers={"x-priority": "forwarded", "x-forwarded-only": "kept"}, - extra_headers={"x-priority": "extra", "x-extra-only": "kept"}, - provider_specific_header={"custom_llm_provider": "anthropic", "extra_headers": {"x-priority": "scoped"}}, - ) - ) - - _, headers = sent(messages_server, response) - assert {name: headers.get(name) for name in ("x-priority", "x-forwarded-only", "x-extra-only")} == { - "x-priority": "scoped", - "x-forwarded-only": "kept", - "x-extra-only": "kept", - } - - -@pytest.mark.asyncio -async def test_non_string_metadata_user_id_is_rejected_before_the_provider_call( - messages_server: RecordingServer, -) -> None: - messages_server.expected_requests = 0 - - with pytest.raises(litellm.BadRequestError, match=r"metadata\.user_id must be a string"): - await litellm.anthropic.messages.acreate(**arguments(messages_server, metadata={"user_id": 123})) - - assert messages_server.requests == [] - - -@pytest.mark.asyncio -async def test_native_messages_observes_runtime_capabilities_and_separate_caller_settings( - messages_server: RecordingServer, monkeypatch: pytest.MonkeyPatch -) -> None: - from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES - from litellm.rust_bridge.public_call import NativeCall - - native: Final = NATIVE_AMESSAGES.load() - assert native is not None - model: Final = "claude-test-runtime-capabilities" - messages_server.expected_requests = 2 - request: Final = NativeCall( - args=(), - kwargs={"temperature": 0.2, "drop_params": True}, - base={ - "model": model, - "messages": MESSAGES, - "max_tokens": 16, - "stream": None, - "api_key": "test-key", - "api_base": messages_server.base_url, - "custom_llm_provider": "anthropic", - "temperature": 0.2, - "drop_params": True, - }, - ) - monkeypatch.setitem( - litellm.model_cost, - model, - { - "litellm_provider": "anthropic", - "mode": "chat", - "supports_sampling_params": True, - }, - ) - first: Final = await native(request) - monkeypatch.setitem( - litellm.model_cost, - model, - { - "litellm_provider": "anthropic", - "mode": "chat", - "supports_sampling_params": False, - }, - ) - second: Final = await native(request) - - assert isinstance(first, dict) - assert isinstance(second, dict) - assert first["id"] == second["id"] == MESSAGES_RESPONSE["id"] - assert len(messages_server.requests) == 2 - first_body: Final = messages_server.requests[0].body - second_body: Final = messages_server.requests[1].body - assert isinstance(first_body, dict) - assert isinstance(second_body, dict) - assert first_body["temperature"] == request.kwargs["temperature"] - assert second_body == {name: value for name, value in first_body.items() if name != "temperature"} - - -@pytest.mark.asyncio -async def test_native_messages_reads_optional_positional_body_parameters(messages_server: RecordingServer) -> None: - metadata: Final = {"user_id": "caller"} - args: Final = (16, MESSAGES, "anthropic/claude-test", metadata, None, False, "Be brief", 0.25) - kwargs: Final = {"api_key": "test-key", "api_base": messages_server.base_url} - response: Final = await litellm.anthropic.messages.acreate(*args, **kwargs) - - body, _ = sent(messages_server, response) - assert body["temperature"] == args[7] - assert body["system"] == args[6] - assert body["metadata"] == metadata diff --git a/tests/test_litellm_rust/messages/test_requests.py b/tests/test_litellm_rust/messages/test_requests.py new file mode 100644 index 00000000000..6388c0eb7db --- /dev/null +++ b/tests/test_litellm_rust/messages/test_requests.py @@ -0,0 +1,150 @@ +from typing import Final + +import pytest + +import litellm +from litellm.rust_bridge.catalog import Route +from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_RESPONSE, call +from tests.test_litellm_rust.support.recording_server import RecordingServer +from tests.test_litellm_rust.support.routing import assert_served_natively, native_route + +pytestmark = pytest.mark.requires_rust_extension + +CAPABILITY_MODEL: Final = "claude-test-runtime-capabilities" + + +def sent(server: RecordingServer) -> tuple[dict[str, object], dict[str, str]]: + assert_served_natively(server) + request: Final = server.requests[0] + assert isinstance(request.body, dict) + return request.body, request.headers + + +def capabilities(monkeypatch: pytest.MonkeyPatch, *, sampling: bool) -> None: + monkeypatch.setitem( + litellm.model_cost, + CAPABILITY_MODEL, + {"litellm_provider": "anthropic", "mode": "chat", "supports_sampling_params": sampling}, + ) + + +@pytest.mark.asyncio +async def test_replayed_tool_ids_are_sanitized_before_they_reach_the_provider(server: RecordingServer) -> None: + history: Final = [ + {"role": "user", "content": "run it"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": ""}, + { + "type": "tool_use", + "id": "functions.Bash:0", + "name": "Bash", + "input": {}, + "provider_specific_fields": {"x": 1}, + }, + ], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions.Bash:0", "content": "ok"}]}, + ] + + await call(server, messages=history) + + body, _ = sent(server) + assert body["messages"] == [ + {"role": "user", "content": "run it"}, + {"role": "assistant", "content": [{"type": "tool_use", "id": "functions_Bash_0", "name": "Bash", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions_Bash_0", "content": "ok"}]}, + ] + + +@pytest.mark.asyncio +async def test_provider_specific_headers_scoped_to_anthropic_reach_the_wire(server: RecordingServer) -> None: + await call( + server, + provider_specific_header=[ + {"custom_llm_provider": "anthropic, azure_ai", "extra_headers": {"x-scoped": "yes"}}, + {"custom_llm_provider": "openai", "extra_headers": {"x-other": "no"}}, + ], + ) + + _, headers = sent(server) + assert headers["x-scoped"] == "yes" + assert "x-other" not in headers + + +@pytest.mark.asyncio +async def test_scoped_headers_override_extra_headers_which_override_forwarded_headers(server: RecordingServer) -> None: + await call( + server, + headers={"x-priority": "forwarded", "x-forwarded-only": "kept"}, + extra_headers={"x-priority": "extra", "x-extra-only": "kept"}, + provider_specific_header={"custom_llm_provider": "anthropic", "extra_headers": {"x-priority": "scoped"}}, + ) + + _, headers = sent(server) + assert {name: headers.get(name) for name in ("x-priority", "x-forwarded-only", "x-extra-only")} == { + "x-priority": "scoped", + "x-forwarded-only": "kept", + "x-extra-only": "kept", + } + + +@pytest.mark.asyncio +async def test_model_capabilities_are_read_from_model_cost_at_call_time( + server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + server.expected_requests = 2 + model: Final = f"anthropic/{CAPABILITY_MODEL}" + + capabilities(monkeypatch, sampling=True) + await call(server, model=model, temperature=0.2, drop_params=True) + capabilities(monkeypatch, sampling=False) + await call(server, model=model, temperature=0.2, drop_params=True) + + assert_served_natively(server, count=2) + first, second = (request.body for request in server.requests) + assert isinstance(first, dict) + assert first["temperature"] == 0.2 + assert second == {name: value for name, value in first.items() if name != "temperature"} + + +@pytest.mark.asyncio +async def test_global_drop_params_drops_unsupported_params( + server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + capabilities(monkeypatch, sampling=False) + monkeypatch.setattr(litellm, "drop_params", True) + + await call(server, model=f"anthropic/{CAPABILITY_MODEL}", top_k=5) + + body, _ = sent(server) + assert "top_k" not in body + + +@pytest.mark.asyncio +async def test_unsupported_params_are_rejected_before_io_without_drop_params( + server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + capabilities(monkeypatch, sampling=False) + monkeypatch.setattr(litellm, "drop_params", False) + server.expected_requests = 0 + + with pytest.raises(litellm.BadRequestError, match="top_k"): + await call(server, model=f"anthropic/{CAPABILITY_MODEL}", top_k=5) + + assert server.requests == [] + + +@pytest.mark.asyncio +async def test_optional_positional_body_parameters_reach_the_wire(server: RecordingServer) -> None: + metadata: Final = {"user_id": "caller"} + args: Final = (16, list(MESSAGES), "anthropic/claude-test", metadata, None, False, "Be brief", 0.25) + + with native_route(Route.MESSAGES): + response: Final = await litellm.anthropic_messages(*args, api_key="test-key", api_base=server.base_url) + + assert isinstance(response, dict) + assert response["id"] == MESSAGES_RESPONSE["id"] + body, _ = sent(server) + assert (body["metadata"], body["system"], body["temperature"]) == (args[3], args[6], args[7]) diff --git a/tests/test_litellm_rust/messages/test_streaming.py b/tests/test_litellm_rust/messages/test_streaming.py new file mode 100644 index 00000000000..c700c093d1a --- /dev/null +++ b/tests/test_litellm_rust/messages/test_streaming.py @@ -0,0 +1,73 @@ +from collections.abc import AsyncIterator, Iterator +from typing import Final + +import pytest + +from tests.test_litellm_rust.messages.data import MESSAGES_EVENTS, MESSAGES_RESPONSE, call +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed +from tests.test_litellm_rust.support.results import consume +from tests.test_litellm_rust.support.routing import assert_served_natively + +pytestmark = pytest.mark.requires_rust_extension + +STREAM: Final = ResponseSpec(body=None, events=MESSAGES_EVENTS) + + +@pytest.fixture +def server(server: RecordingServer) -> Iterator[RecordingServer]: + yield from relaxed(server) + + +@pytest.mark.asyncio +async def test_stream_relays_provider_events_and_logs_success_once_after_the_last_chunk( + server: RecordingServer, +) -> None: + server.enqueue(STREAM) + recorder: Final = RecordingLogger() + + stream: Final = await call(server, stream=True, callbacks=[recorder]) + assert isinstance(stream, AsyncIterator) + first: Final = await anext(stream) + await drain_logging() + assert "async_log_success_event" not in recorder.names + rest: Final = [chunk async for chunk in stream] + + assert_served_natively(server) + assert first + b"".join(rest) == b"".join(STREAM.payloads()) + success: Final = await recorder.wait_for_async("async_log_success_event") + assert len(success) == 1 + assert isinstance(success[0].kwargs, dict) + assert success[0].kwargs["completion_start_time"] is not None + usage: Final = success[0].response.usage + assert usage.prompt_tokens == MESSAGES_RESPONSE["usage"]["input_tokens"] + assert usage.completion_tokens == MESSAGES_RESPONSE["usage"]["output_tokens"] + assert "log_failure_event" not in recorder.names + + +@pytest.mark.asyncio +async def test_sync_stream_relays_provider_events_and_logs_success_once(server: RecordingServer) -> None: + server.enqueue(STREAM) + recorder: Final = RecordingLogger() + + stream: Final = await call(server, asynchronous=False, stream=True, callbacks=[recorder]) + chunks: Final = await consume(stream) + + assert_served_natively(server) + assert b"".join(chunks) == b"".join(STREAM.payloads()) + assert len(await recorder.wait_for_async("async_log_success_event")) == 1 + + +@pytest.mark.asyncio +async def test_stream_closed_early_logs_success_once(server: RecordingServer) -> None: + server.enqueue(STREAM) + recorder: Final = RecordingLogger() + + stream: Final = await call(server, stream=True, callbacks=[recorder]) + assert isinstance(stream, AsyncIterator) + await anext(stream) + await stream.aclose() + + assert len(await recorder.wait_for_async("async_log_success_event")) == 1 + with pytest.raises(StopAsyncIteration): + await anext(stream) diff --git a/tests/test_litellm_rust/ocr/conftest.py b/tests/test_litellm_rust/ocr/conftest.py new file mode 100644 index 00000000000..6d3834516d7 --- /dev/null +++ b/tests/test_litellm_rust/ocr/conftest.py @@ -0,0 +1,19 @@ +from typing import Final + +import pytest + +from tests.test_litellm_rust.ocr.data import OCR_RESPONSE +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec + + +@pytest.fixture +def server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=OCR_RESPONSE) + return recording_server + + +@pytest.fixture(params=[False, True], ids=["python", "rust"]) +def ocr_backend(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> bool: + enabled: Final = bool(request.param) + monkeypatch.setenv("LITELLM_RUST", "1" if enabled else "0") + return enabled diff --git a/tests/test_litellm_rust/ocr/data.py b/tests/test_litellm_rust/ocr/data.py new file mode 100644 index 00000000000..0f039892245 --- /dev/null +++ b/tests/test_litellm_rust/ocr/data.py @@ -0,0 +1,39 @@ +from typing import Final + +import litellm +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from tests.test_litellm_rust.support.recording_server import RecordingServer + +OCR_DOCUMENT: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"} +OCR_MODEL: Final = "mistral/mistral-ocr-latest" +OCR_TEXT: Final = "native OCR response" +OCR_RESPONSE: Final = { + "pages": [{"index": 0, "markdown": OCR_TEXT, "images": [], "dimensions": None}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1, "doc_size_bytes": 3}, +} + + +def ocr_arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]: + return { + "model": OCR_MODEL, + "document": dict(OCR_DOCUMENT), + "api_key": "test-key", + "api_base": server.base_url, + **kwargs, + } + + +def call_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse: + response: Final = litellm.ocr(**ocr_arguments(server, **kwargs)) + if not isinstance(response, OCRResponse): + raise TypeError(f"Expected OCRResponse, got {type(response).__name__}") + return response + + +async def call_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse: + return await litellm.aocr(**ocr_arguments(server, **kwargs)) + + +async def call(server: RecordingServer, asynchronous: bool, **kwargs: object) -> OCRResponse: + return await call_aocr(server, **kwargs) if asynchronous else call_ocr(server, **kwargs) diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py index 9e3945f11b1..5f961288800 100644 --- a/tests/test_litellm_rust/ocr/test_callbacks.py +++ b/tests/test_litellm_rust/ocr/test_callbacks.py @@ -15,49 +15,38 @@ import litellm from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.llms.base_llm.ocr.transformation import OCRResponse -from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging -from tests.test_litellm_rust.support.isolation import isolated_callback_registries, rebound -from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec -from tests.test_litellm_rust.support.requests import ( +from tests.test_litellm_rust.ocr.data import ( OCR_DOCUMENT, OCR_RESPONSE, - call_native, - call_native_aocr, - call_native_ocr, + call, + call_aocr, + call_ocr, +) +from tests.test_litellm_rust.support.callback_recorder import ( + RecordingLogger, + drain_logging, request_body, request_headers, ) +from tests.test_litellm_rust.support.isolation import isolated_callback_registries, rebound +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec pytestmark = pytest.mark.requires_rust_extension -@pytest.fixture -def ocr_server(recording_server: RecordingServer) -> RecordingServer: - recording_server.default_response = ResponseSpec(body=OCR_RESPONSE) - return recording_server - - -def call_native_ocr_with_callbacks(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object): - return call_native_ocr(server, callbacks=callbacks, **kwargs) - - -async def call_native_aocr_with_callbacks(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object): - return await call_native_aocr(server, callbacks=callbacks, **kwargs) - - -def test_native_ocr_pre_call_callback_receives_transformed_provider_request(ocr_server: RecordingServer) -> None: +def test_native_ocr_pre_call_callback_receives_transformed_provider_request(server: RecordingServer) -> None: observations: Final = [] class Observe(CustomLogger): def log_pre_api_call(self, model, messages, kwargs): observations.append((model, copy.deepcopy(kwargs["additional_args"]))) - call_native_ocr_with_callbacks(ocr_server, [Observe()], pages=[0]) + call_ocr(server, callbacks=[Observe()], pages=[0]) assert len(observations) == 1 model, additional_args = observations[0] assert model == "mistral-ocr-latest" - assert additional_args["api_base"] == f"{ocr_server.base_url}/v1/ocr" + assert additional_args["api_base"] == f"{server.base_url}/v1/ocr" assert additional_args["complete_input_dict"] == { "model": "mistral-ocr-latest", "document": OCR_DOCUMENT, @@ -66,7 +55,7 @@ def test_native_ocr_pre_call_callback_receives_transformed_provider_request(ocr_ @pytest.mark.asyncio -async def test_native_ocr_does_not_restore_keywords_deleted_by_a_pre_call_hook(ocr_server: RecordingServer) -> None: +async def test_native_ocr_does_not_restore_keywords_deleted_by_a_pre_call_hook(server: RecordingServer) -> None: class DropKeywords(CustomLogger): async def async_pre_call_deployment_hook( self, kwargs: dict[str, object], call_type: object @@ -74,17 +63,17 @@ async def test_native_ocr_does_not_restore_keywords_deleted_by_a_pre_call_hook(o return {name: value for name, value in kwargs.items() if name not in ("pages", "extra_headers")} with rebound(litellm, "callbacks", [DropKeywords()]): - await call_native_aocr(ocr_server, pages=[0], extra_headers={"x-deleted": "remove this header"}) + await call_aocr(server, pages=[0], extra_headers={"x-deleted": "remove this header"}) - assert len(ocr_server.requests) == 1 - sent: Final = ocr_server.requests[0] + assert len(server.requests) == 1 + sent: Final = server.requests[0] assert sent.body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT} assert sent.headers.get("x-deleted") is None @pytest.mark.parametrize("raise_after_edit", [False, True], ids=["callback-returns", "callback-raises"]) def test_native_ocr_pre_call_body_edit_reaches_next_callback_and_provider( - ocr_server: RecordingServer, raise_after_edit: bool + server: RecordingServer, raise_after_edit: bool ) -> None: observed: Final = [] @@ -98,13 +87,13 @@ def test_native_ocr_pre_call_body_edit_reaches_next_callback_and_provider( def log_pre_api_call(self, model, messages, kwargs): observed.append(copy.deepcopy(request_body(kwargs))) - call_native_ocr_with_callbacks(ocr_server, [Edit(), Observe()], include_image_base64=False) + call_ocr(server, callbacks=[Edit(), Observe()], include_image_base64=False) assert observed[0]["include_image_base64"] is True - assert ocr_server.requests[0].body["include_image_base64"] is True + assert server.requests[0].body["include_image_base64"] is True -def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(ocr_server: RecordingServer) -> None: +def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(server: RecordingServer) -> None: observed: Final = [] class Edit(CustomLogger): @@ -115,16 +104,16 @@ def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(ocr_ def log_pre_api_call(self, model, messages, kwargs): observed.append(dict(request_headers(kwargs))) - call_native_ocr_with_callbacks(ocr_server, [Edit(), Observe()]) + call_ocr(server, callbacks=[Edit(), Observe()]) assert observed[0]["x-audit-tag"] == "reviewed" - assert ocr_server.requests[0].headers["x-audit-tag"] == "reviewed" + assert server.requests[0].headers["x-audit-tag"] == "reviewed" @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_and_provider_references( - ocr_server: RecordingServer, asynchronous: bool + server: RecordingServer, asynchronous: bool ) -> None: original: Final = dict(OCR_DOCUMENT) replacement_url: Final = "data:application/pdf;base64,ZGVm" @@ -144,22 +133,20 @@ async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_ "model": "mistral/mistral-ocr-latest", "document": original, "api_key": "test-key", - "api_base": ocr_server.base_url, + "api_base": server.base_url, "callbacks": [Retain(), Edit()], } - response: Final = ( - await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments) - ) + response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments) assert aliases == [True] assert retained[0]["document_url"] == replacement_url assert original["document_url"] == replacement_url - assert ocr_server.requests[0].body["document"]["document_url"] == replacement_url + assert server.requests[0].body["document"]["document_url"] == replacement_url assert response.pages[0].markdown == "native OCR response" def test_native_ocr_pre_call_body_rebinding_is_visible_to_callbacks_but_not_provider( - ocr_server: RecordingServer, + server: RecordingServer, ) -> None: observed: Final = [] @@ -171,13 +158,13 @@ def test_native_ocr_pre_call_body_rebinding_is_visible_to_callbacks_but_not_prov def log_pre_api_call(self, model, messages, kwargs): observed.append(request_body(kwargs)) - call_native_ocr_with_callbacks(ocr_server, [Rebind(), Observe()]) + call_ocr(server, callbacks=[Rebind(), Observe()]) assert observed == [{"replacement": True}] - assert ocr_server.requests[0].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT} + assert server.requests[0].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT} -def test_native_ocr_callback_retained_body_observes_later_callback_mutation(ocr_server: RecordingServer) -> None: +def test_native_ocr_callback_retained_body_observes_later_callback_mutation(server: RecordingServer) -> None: queued: Final = [] class QueuePayload(CustomLogger): @@ -188,12 +175,12 @@ def test_native_ocr_callback_retained_body_observes_later_callback_mutation(ocr_ def log_pre_api_call(self, model, messages, kwargs): request_body(kwargs)["queued-edit"] = True - call_native_ocr_with_callbacks(ocr_server, [QueuePayload(), Edit()]) + call_ocr(server, callbacks=[QueuePayload(), Edit()]) assert queued[0]["queued-edit"] is True -def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(ocr_server: RecordingServer) -> None: +def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(server: RecordingServer) -> None: token: Final = object() terminal_tokens: queue.SimpleQueue[object] = queue.SimpleQueue() finished: Final = threading.Event() @@ -206,7 +193,7 @@ def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(o terminal_tokens.put(kwargs["test-token"]) finished.set() - call_native_ocr_with_callbacks(ocr_server, [Stash()]) + call_ocr(server, callbacks=[Stash()]) assert finished.wait(10) assert terminal_tokens.get_nowait() is token @@ -214,13 +201,13 @@ def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(o @pytest.mark.asyncio async def test_native_aocr_success_callback_receives_call_id_metadata_and_response( - ocr_server: RecordingServer, + server: RecordingServer, ) -> None: recorder: Final = RecordingLogger() - await call_native_aocr_with_callbacks( - ocr_server, - [recorder], + await call_aocr( + server, + callbacks=[recorder], litellm_call_id="ocr-success", metadata={"source": "callback-test"}, ) @@ -235,9 +222,9 @@ async def test_native_aocr_success_callback_receives_call_id_metadata_and_respon @pytest.mark.asyncio async def test_native_aocr_failure_callbacks_receive_call_type_error_and_no_response( - ocr_server: RecordingServer, + server: RecordingServer, ) -> None: - ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500)) + server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500)) observations: Final = [] class Observe(CustomLogger): @@ -248,7 +235,7 @@ async def test_native_aocr_failure_callbacks_receive_call_type_error_and_no_resp observations.append(("async", kwargs["call_type"], kwargs["exception"], response_obj)) with pytest.raises(litellm.InternalServerError): - await call_native_aocr_with_callbacks(ocr_server, [Observe()]) + await call_aocr(server, callbacks=[Observe()]) assert [observation[0] for observation in observations] == ["sync", "async"] assert all(observation[1] == "aocr" for observation in observations) @@ -257,12 +244,12 @@ async def test_native_aocr_failure_callbacks_receive_call_type_error_and_no_resp @pytest.mark.asyncio -async def test_native_aocr_pre_call_callback_runs_on_caller_loop_and_thread(ocr_server: RecordingServer) -> None: +async def test_native_aocr_pre_call_callback_runs_on_caller_loop_and_thread(server: RecordingServer) -> None: caller_loop: Final = asyncio.get_running_loop() caller_thread: Final = threading.current_thread() recorder: Final = RecordingLogger() - await call_native_aocr_with_callbacks(ocr_server, [recorder]) + await call_aocr(server, callbacks=[recorder]) events: Final = await recorder.wait_for_async("log_pre_api_call") assert len(events) == 1 @@ -272,9 +259,9 @@ async def test_native_aocr_pre_call_callback_runs_on_caller_loop_and_thread(ocr_ @pytest.mark.asyncio async def test_native_aocr_failure_callbacks_receive_state_added_by_pre_call_callback( - ocr_server: RecordingServer, + server: RecordingServer, ) -> None: - ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500)) + server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500)) token: Final = object() observed: Final = [] @@ -289,20 +276,20 @@ async def test_native_aocr_failure_callbacks_receive_state_added_by_pre_call_cal observed.append(("async", kwargs["request-token"])) with pytest.raises(litellm.InternalServerError): - await call_native_aocr_with_callbacks(ocr_server, [TrackInFlightRequest()]) + await call_aocr(server, callbacks=[TrackInFlightRequest()]) assert [event for event, _ in observed] == ["sync", "async"] assert all(observed_token is token for _, observed_token in observed) def test_native_ocr_dispatches_each_callback_phase_once_when_logger_is_registered_multiple_times( - ocr_server: RecordingServer, + server: RecordingServer, ) -> None: recorder: Final = RecordingLogger() - call_native_ocr_with_callbacks( - ocr_server, - [recorder, recorder], + call_ocr( + server, + callbacks=[recorder, recorder], success_callback=[recorder], failure_callback=[recorder], ) @@ -340,21 +327,21 @@ class ApplyEdits(CustomLogger): @settings(max_examples=25, deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture]) @given(edits=st.dictionaries(st.from_regex(r"x_[a-z]{1,6}", fullmatch=True), JSON_VALUES, max_size=3)) def test_native_ocr_provider_receives_the_body_exactly_as_pre_call_callbacks_left_it( - ocr_server: RecordingServer, edits: dict[str, object] + server: RecordingServer, edits: dict[str, object] ) -> None: - ocr_server.expected_requests = None + server.expected_requests = None with isolated_callback_registries(): - call_native_ocr_with_callbacks(ocr_server, [ApplyEdits(MappingProxyType(edits))]) + call_ocr(server, callbacks=[ApplyEdits(MappingProxyType(edits))]) - assert ocr_server.requests[-1].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT, **edits} + assert server.requests[-1].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT, **edits} @pytest.mark.parametrize("hook", ["log_pre_api_call", "logging_hook", "log_success_event"]) -def test_native_ocr_sync_hooks_see_no_running_event_loop(ocr_server: RecordingServer, hook: str) -> None: +def test_native_ocr_sync_hooks_see_no_running_event_loop(server: RecordingServer, hook: str) -> None: recorder: Final = RecordingLogger() - call_native_ocr_with_callbacks(ocr_server, [recorder]) + call_ocr(server, callbacks=[recorder]) [event] = recorder.wait_for(hook) assert event.loop is None @@ -364,7 +351,7 @@ def test_native_ocr_sync_hooks_see_no_running_event_loop(ocr_server: RecordingSe @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) async def test_native_ocr_payload_a_callback_retains_outlives_the_call_intact( - ocr_server: RecordingServer, asynchronous: bool + server: RecordingServer, asynchronous: bool ) -> None: retained: Final = [] @@ -372,21 +359,21 @@ async def test_native_ocr_payload_a_callback_retains_outlives_the_call_intact( def log_pre_api_call(self, model, messages, kwargs): retained.append((kwargs, request_body(kwargs), request_headers(kwargs))) - await call_native(ocr_server, asynchronous, callbacks=[Retain()]) + await call(server, asynchronous, callbacks=[Retain()]) await drain_logging() gc.collect() [(details, body, headers)] = retained - assert body == ocr_server.requests[0].body + assert body == server.requests[0].body assert headers - assert all(ocr_server.requests[0].headers[name.lower()] == value for name, value in headers.items()) + assert all(server.requests[0].headers[name.lower()] == value for name, value in headers.items()) assert details["additional_args"]["complete_input_dict"] is body assert details["additional_args"]["headers"] is headers @pytest.mark.asyncio @pytest.mark.parametrize("family", ["sync", "async"]) -async def test_native_ocr_success_callbacks_share_one_logging_payload(ocr_server: RecordingServer, family: str) -> None: +async def test_native_ocr_success_callbacks_share_one_logging_payload(server: RecordingServer, family: str) -> None: queued: Final = [] finished: Final = threading.Event() @@ -421,7 +408,7 @@ async def test_native_ocr_success_callbacks_share_one_logging_payload(ocr_server async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): strip_payload(kwargs) - await call_native(ocr_server, family == "async", callbacks=[QueuePayload(), StripPayload()]) + await call(server, family == "async", callbacks=[QueuePayload(), StripPayload()]) await drain_logging() assert await asyncio.to_thread(finished.wait, 10) @@ -430,7 +417,7 @@ async def test_native_ocr_success_callbacks_share_one_logging_payload(ocr_server @pytest.mark.asyncio async def test_native_aocr_state_stashed_before_a_blocking_hook_raises_reaches_failure_callbacks( - ocr_server: RecordingServer, + server: RecordingServer, ) -> None: token: Final = object() observed: Final = [] @@ -455,7 +442,7 @@ async def test_native_aocr_state_stashed_before_a_blocking_hook_raises_reaches_f litellm.callbacks.append(Block()) with pytest.raises(Blocked) as raised: - await call_native_aocr(ocr_server) + await call_aocr(server) await drain_logging() assert observed == [("sync", token, raised.value), ("async", token, raised.value)] @@ -464,7 +451,7 @@ async def test_native_aocr_state_stashed_before_a_blocking_hook_raises_reaches_f @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) async def test_native_azure_ocr_resolves_token_before_pre_call_on_caller_context( - ocr_server: RecordingServer, + server: RecordingServer, isolated_azure_auth: None, asynchronous: bool, ) -> None: @@ -497,27 +484,25 @@ async def test_native_azure_ocr_resolves_token_before_pre_call_on_caller_context "azure_ad_token_provider": provider, "callbacks": [Edit()], } - response: Final = ( - await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments) - ) + response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments) assert response.pages[0].markdown == "native OCR response" assert observations == ["token", "pre_call"] - assert ocr_server.requests[0].headers["authorization"] == "Bearer edited" + assert server.requests[0].headers["authorization"] == "Bearer edited" @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) async def test_native_azure_ocr_token_provider_can_make_nested_native_ocr_call( - ocr_server: RecordingServer, + server: RecordingServer, isolated_azure_auth: None, asynchronous: bool, ) -> None: - ocr_server.expected_requests = 2 + server.expected_requests = 2 calls: Final = [] def provider() -> str: calls.append("token") - nested: Final = call_native_ocr(ocr_server) + nested: Final = call_ocr(server) assert nested.pages[0].markdown == "native OCR response" return "outer-token" @@ -526,12 +511,10 @@ async def test_native_azure_ocr_token_provider_can_make_nested_native_ocr_call( "api_key": None, "azure_ad_token_provider": provider, } - response: Final = ( - await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments) - ) + response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments) assert response.pages[0].markdown == "native OCR response" assert calls == ["token"] - assert [request.headers["authorization"] for request in ocr_server.requests] == [ + assert [request.headers["authorization"] for request in server.requests] == [ "Bearer test-key", "Bearer outer-token", ] @@ -539,10 +522,10 @@ async def test_native_azure_ocr_token_provider_can_make_nested_native_ocr_call( @pytest.mark.asyncio async def test_concurrent_native_azure_ocr_calls_isolate_token_results_and_error( - ocr_server: RecordingServer, + server: RecordingServer, isolated_azure_auth: None, ) -> None: - ocr_server.expected_requests = 2 + server.expected_requests = 2 async def request(token: str, fail: bool) -> object: def provider() -> str: @@ -550,8 +533,8 @@ async def test_concurrent_native_azure_ocr_calls_isolate_token_results_and_error raise ValueError(token) return token - return await call_native_aocr( - ocr_server, + return await call_aocr( + server, model="azure_ai/mistral-ocr-latest", api_key=None, azure_ad_token_provider=provider, @@ -567,7 +550,7 @@ async def test_concurrent_native_azure_ocr_calls_isolate_token_results_and_error assert isinstance(responses[1], litellm.APIConnectionError) assert "Failed to get Azure AD token: failed" in str(responses[1]) assert isinstance(responses[2], OCRResponse) - assert sorted(request.headers["authorization"] for request in ocr_server.requests) == [ + assert sorted(request.headers["authorization"] for request in server.requests) == [ "Bearer first", "Bearer second", ] @@ -575,7 +558,7 @@ async def test_concurrent_native_azure_ocr_calls_isolate_token_results_and_error @pytest.mark.asyncio async def test_native_azure_ocr_releases_token_provider_after_cancellation( - ocr_server: RecordingServer, + server: RecordingServer, isolated_azure_auth: None, ) -> None: import gc @@ -590,16 +573,16 @@ async def test_native_azure_ocr_releases_token_provider_after_cancellation( async def invoke() -> weakref.ReferenceType[Provider]: provider: Final = Provider() reference: Final = weakref.ref(provider) - ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.1)) + server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.1)) task: Final = asyncio.create_task( - call_native_aocr( - ocr_server, + call_aocr( + server, model="azure_ai/mistral-ocr-latest", api_key=None, azure_ad_token_provider=provider, ) ) - await ocr_server.wait_for_requests(1) + await server.wait_for_requests(1) assert reference() is provider task.cancel() with pytest.raises(asyncio.CancelledError): diff --git a/tests/test_litellm_rust/ocr/test_cohere.py b/tests/test_litellm_rust/ocr/test_cohere.py deleted file mode 100644 index 8474e971c6f..00000000000 --- a/tests/test_litellm_rust/ocr/test_cohere.py +++ /dev/null @@ -1,32 +0,0 @@ -from typing import Final - -import pytest - -import litellm -from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec - -pytestmark = pytest.mark.requires_rust_extension -MODELS: Final = ("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0") -IMAGE: Final = {"type": "image_url", "image_url": "data:image/png;base64,YWJj"} -BOX: Final = {"top_left_x": 0, "top_left_y": 0, "bottom_right_x": 32, "bottom_right_y": 32} -PAYLOAD: Final = { - "pages": [ - { - "index": 4, - "markdown": {"content": "receipt", "images": [{"id": "image", "bounding_box": BOX, "description": "scan"}]}, - }, - {"markdown": {"content": "page two"}}, - ], - "meta": {"billed_units": {"pages": 3}}, -} - - -@pytest.mark.asyncio -@pytest.mark.parametrize("model", MODELS) -async def test_public_cohere_health_check(recording_server: RecordingServer, model: str) -> None: - recording_server.enqueue(ResponseSpec(body=PAYLOAD)) - response: Final = await litellm.ahealth_check( - model_params={"model": model, "api_key": "test-key", "api_base": recording_server.base_url}, mode="ocr" - ) - assert "error" not in response - assert recording_server.requests[0].body["document"]["image_url"].startswith("data:image/png;base64,") diff --git a/tests/test_litellm_rust/ocr/test_dispatch.py b/tests/test_litellm_rust/ocr/test_dispatch.py index a29a45376d9..bb75782fb67 100644 --- a/tests/test_litellm_rust/ocr/test_dispatch.py +++ b/tests/test_litellm_rust/ocr/test_dispatch.py @@ -1,42 +1,36 @@ -from typing import Final +from typing import Final, NoReturn import pytest import litellm from litellm.llms.base_llm.ocr.transformation import OCRResponse +from tests.test_litellm_rust.ocr.data import OCR_DOCUMENT, OCR_MODEL from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec -from tests.test_litellm_rust.support.requests import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE -from tests.test_litellm_rust.support.served import assert_served_natively +from tests.test_litellm_rust.support.routing import assert_served_natively pytestmark = pytest.mark.requires_rust_extension -@pytest.fixture -def ocr_server(recording_server: RecordingServer) -> RecordingServer: - recording_server.default_response = ResponseSpec(body=OCR_RESPONSE) - return recording_server - - @pytest.mark.parametrize("enabled", [False, True, None]) -def test_public_ocr_uses_native_route_independently_of_flag(ocr_server: RecordingServer, enabled: bool | None) -> None: +def test_public_ocr_uses_native_route_independently_of_flag(server: RecordingServer, enabled: bool | None) -> None: litellm.rust(enabled) response: Final = litellm.ocr( model=OCR_MODEL, document=OCR_DOCUMENT, api_key="test-key", - api_base=ocr_server.base_url, + api_base=server.base_url, ) assert isinstance(response, OCRResponse) assert response.pages[0].markdown == "native OCR response" - assert_served_natively(ocr_server, response) + assert_served_natively(server, response) @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.parametrize("caching", [None, False, True]) async def test_ocr_does_not_depend_on_chat_cache( - ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool, caching: bool | None + server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool, caching: bool | None ) -> None: from litellm.caching.caching import Cache @@ -45,9 +39,93 @@ async def test_ocr_does_not_depend_on_chat_cache( "model": OCR_MODEL, "document": OCR_DOCUMENT, "api_key": "test-key", - "api_base": ocr_server.base_url, + "api_base": server.base_url, "caching": caching, } response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) assert response.pages[0].markdown == "native OCR response" - assert len(ocr_server.requests) == 1 + assert len(server.requests) == 1 + + +@pytest.mark.parametrize("required", [False, True]) +@pytest.mark.parametrize("failure", ["invalid", "live", "schema"]) +def test_native_projection_errors_never_select_python( + server: RecordingServer, monkeypatch: pytest.MonkeyPatch, required: bool, failure: str +) -> None: + import dataclasses + import ssl + + from litellm.rust_bridge import runtime + from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule + from litellm.rust_bridge.configuration import Rollout + from litellm.rust_bridge.host import settings + from litellm.rust_bridge.ocr.entrypoints import NATIVE_OCR + from litellm.rust_bridge.public_call import NativeCall + + server.expected_requests = 0 + snapshot: Final = dataclasses.replace(settings.http_settings(), user_agent=1) + if failure == "schema": + monkeypatch.setattr(settings, "http_settings", lambda: snapshot) + else: + monkeypatch.setattr( + litellm, "ssl_verify", ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) if failure == "live" else object() + ) + request: Final = NativeCall( + args=(), + kwargs={}, + base={ + "model": "mistral/mistral-ocr-latest", + "document": OCR_DOCUMENT, + "api_key": "test-key", + "api_base": server.base_url, + "timeout": None, + "custom_llm_provider": "mistral", + "extra_headers": None, + }, + ) + + def python_fallback() -> NoReturn: + pytest.fail("projection failures must not select Python") + + with pytest.raises(RuntimeError if failure == "schema" else ValueError, match="http_settings"): + runtime.run( + RouteContext(Route.OCR, provider="mistral"), + binding=NATIVE_OCR, + native=lambda native: native(request), + python=python_fallback, + rules=(RouteRule(Route.OCR, Rollout.RUST_REQUIRED if required else Rollout.RUST_OPT_OUT),), + ) + assert server.requests == [] + + +MODELS: Final = ("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0") + + +IMAGE: Final = {"type": "image_url", "image_url": "data:image/png;base64,YWJj"} + + +BOX: Final = {"top_left_x": 0, "top_left_y": 0, "bottom_right_x": 32, "bottom_right_y": 32} + + +PAYLOAD: Final = { + "pages": [ + { + "index": 4, + "markdown": {"content": "receipt", "images": [{"id": "image", "bounding_box": BOX, "description": "scan"}]}, + }, + {"markdown": {"content": "page two"}}, + ], + "meta": {"billed_units": {"pages": 3}}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", MODELS) +async def test_ocr_health_check_uses_the_native_route(recording_server: RecordingServer, model: str) -> None: + recording_server.enqueue(ResponseSpec(body=PAYLOAD)) + response: Final = await litellm.ahealth_check( + model_params={"model": model, "api_key": "test-key", "api_base": recording_server.base_url}, mode="ocr" + ) + assert "error" not in response + assert_served_natively(recording_server) + assert recording_server.requests[0].body["document"]["image_url"].startswith("data:image/png;base64,") diff --git a/tests/test_litellm_rust/ocr/test_failures.py b/tests/test_litellm_rust/ocr/test_failures.py new file mode 100644 index 00000000000..41373fb6e77 --- /dev/null +++ b/tests/test_litellm_rust/ocr/test_failures.py @@ -0,0 +1,340 @@ +import json +from collections.abc import Callable +from dataclasses import dataclass +from io import BytesIO +from pathlib import Path +from typing import Final + +import httpx +import pytest +from pydantic import JsonValue + +import litellm +from tests.test_litellm_rust.ocr.data import ( + OCR_RESPONSE, + call, + call_aocr, + call_ocr, +) +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_ocr_contract_upstream_status( + server: RecordingServer, + ocr_backend: bool, + asynchronous: bool, +) -> None: + upstream: Final = ResponseSpec(body={"detail": "invalid provider option"}, status=422) + server.enqueue(upstream) + arguments: Final = { + "model": "vertex_ai/mistral-ocr-latest", + "vertex_project": "test-project", + "vertex_location": "us-central1", + "num_retries": 0, + } + with pytest.raises(litellm.BadRequestError) as caught: + await call(server, asynchronous, **arguments) + assert caught.value.status_code == upstream.status + assert caught.value.response.status_code == upstream.status + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("preserved", ["body", "headers"]) +async def test_ocr_contract_provider_error_details( + server: RecordingServer, + ocr_backend: bool, + asynchronous: bool, + preserved: str, +) -> None: + payload: Final = {"message": "rate limited"} + headers: Final = {"Retry-After": "17", "X-Request-ID": "ocr-request-123", "X-Future-Header": "retained"} + server.enqueue(ResponseSpec(body=payload, status=429, headers=headers)) + with pytest.raises(litellm.RateLimitError) as caught: + await call(server, asynchronous, num_retries=0) + response: Final = caught.value.response + assert isinstance(response, httpx.Response) + if preserved == "body": + assert response.content == json.dumps(payload).encode() + else: + for name, value in headers.items(): + assert response.headers.get(name.lower()) == value + assert response.headers.get(name.upper()) == value + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_ocr_contract_invalid_response_format( + server: RecordingServer, + ocr_backend: bool, + asynchronous: bool, +) -> None: + server.expected_requests = 0 + with pytest.raises(litellm.UnsupportedParamsError) as caught: + await call(server, asynchronous, req_format="bogus", num_retries=0) + assert caught.value.status_code == 400 + for value in ("req_format", "bogus", "native", "litellm"): + assert value in str(caught.value) + assert server.requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + "document,field", + [ + ([], "document"), + ({"document_url": "https://example.com/a.pdf"}, "type"), + ({"type": "text"}, "type"), + ], +) +async def test_ocr_contract_malformed_document_is_actionable( + server: RecordingServer, + ocr_backend: bool, + asynchronous: bool, + document: JsonValue, + field: str, +) -> None: + server.expected_requests = None + with pytest.raises(litellm.BadRequestError) as caught: + await call(server, asynchronous, document=document, num_retries=0) + assert caught.value.status_code == 400 + assert field.lower() in str(caught.value).lower() + assert "NoneType: None" not in str(caught.value) + assert "indices must be" not in str(caught.value) + assert server.requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("option,value,field", [("pages", [-1], "pages"), ("features", [1], "features")]) +async def test_ocr_contract_azure_invalid_options_are_bad_requests( + server: RecordingServer, + ocr_backend: bool, + asynchronous: bool, + option: str, + value: JsonValue, + field: str, +) -> None: + server.expected_requests = 0 + arguments: Final = {"model": "azure_ai/doc-intelligence/prebuilt-read", option: value, "num_retries": 0} + with pytest.raises(litellm.BadRequestError) as caught: + await call(server, asynchronous, **arguments) + assert caught.value.status_code == 400 + assert field in str(caught.value) + assert server.requests == [] + + +def test_native_ocr_maps_provider_400_with_public_provider_details(server: RecordingServer) -> None: + server.enqueue(ResponseSpec(body={"message": "invalid OCR request"}, status=400)) + + with pytest.raises(litellm.BadRequestError) as caught: + call_ocr(server) + + assert caught.value.status_code == 400 + assert caught.value.model == "mistral-ocr-latest" + assert caught.value.llm_provider == "mistral" + assert "invalid OCR request" in str(caught.value) + + +class TokenAbort(BaseException): + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("failure", ["ordinary", "abort"], ids=["value-error", "base-exception"]) +async def test_native_azure_ocr_token_provider_failure_prevents_pre_call_callback_and_request( + server: RecordingServer, + isolated_azure_auth: None, + asynchronous: bool, + failure: str, +) -> None: + server.expected_requests = 0 + calls: Final = [] + recorder: Final = RecordingLogger() + original: Final = {"ordinary": ValueError("token unavailable"), "abort": TokenAbort("abort")} + + def token_provider() -> object: + calls.append("token") + raise original[failure] + + arguments: Final = { + "model": "azure_ai/mistral-ocr-latest", + "api_key": None, + "azure_ad_token_provider": token_provider, + "callbacks": [recorder], + } + expected: Final = TokenAbort if failure == "abort" else litellm.APIConnectionError + with pytest.raises(expected) as caught: + await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments) + assert calls == ["token"] + assert server.requests == [] + assert "log_pre_api_call" not in recorder.names + if failure == "ordinary": + assert "Failed to get Azure AD token: token unavailable" in str(caught.value) + assert isinstance(caught.value.__context__, RuntimeError) + assert caught.value.__context__.__cause__ is original[failure] + else: + assert caught.value is original[failure] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_native_file_preparation_preserves_reader_exception(server: RecordingServer, asynchronous: bool) -> None: + server.expected_requests = 0 + failure: Final = RuntimeError("reader failed") + + class Reader: + def read(self) -> bytes: + raise failure + + document: Final = {"type": "file", "file": Reader()} + with pytest.raises(litellm.APIConnectionError, match="reader failed") as caught: + await call_aocr(server, document=document) if asynchronous else call_ocr(server, document=document) + assert caught.value.__context__ is failure + + +COHERE_IMAGE: Final = {"type": "image_url", "image_url": "data:image/png;base64,YWJj"} + + +FILE_SIZE_LIMIT: Final = 50 * 1024 * 1024 + + +class IntReader: + def read(self) -> int: + return 1 + + +def oversized_file(tmp_path: Path) -> Path: + path: Final = tmp_path / "large.pdf" + with path.open("wb") as stream: + stream.truncate(FILE_SIZE_LIMIT + 1) + return path + + +def empty_token() -> str: + return "" + + +def unused_token() -> str: + raise AssertionError("the token provider must not run") + + +@dataclass(frozen=True, slots=True) +class PublicFailure: + arguments: Callable[[Path], dict[str, object]] + error: type[Exception] + match: str + provider_requests: int = 0 + response: ResponseSpec | None = None + cause: type[BaseException] | None = None + + +PUBLIC_FAILURES: Final = { + "unknown-req-format": PublicFailure( + lambda _: {"req_format": "raw"}, litellm.BadRequestError, "Invalid `req_format`" + ), + "empty-file": PublicFailure( + lambda _: {"document": {"type": "file", "file": BytesIO(b"")}}, litellm.BadRequestError, "File is empty" + ), + "oversized-file": PublicFailure( + lambda tmp_path: {"document": {"type": "file", "file": oversized_file(tmp_path)}}, + litellm.BadRequestError, + "exceeds the size limit", + ), + "missing-file": PublicFailure( + lambda tmp_path: {"document": {"type": "file", "file": tmp_path / "missing.pdf"}}, + litellm.APIConnectionError, + "File not found", + cause=FileNotFoundError, + ), + "reader-returns-non-bytes": PublicFailure( + lambda _: {"document": {"type": "file", "file": IntReader()}}, + litellm.APIConnectionError, + "bytes or str", + cause=TypeError, + ), + "cohere-non-image": PublicFailure( + lambda _: {"model": "cohere/parse-v5.0"}, litellm.BadRequestError, "only accepts `image_url`" + ), + "cohere-unknown-format": PublicFailure( + lambda _: {"model": "cohere/parse-v5.0", "document": COHERE_IMAGE, "output_format": "html"}, + litellm.BadRequestError, + "output_format", + ), + "azure-missing-api-base": PublicFailure( + lambda _: { + "model": "azure_ai/mistral-ocr-latest", + "api_key": None, + "api_base": None, + "azure_ad_token_provider": unused_token, + }, + litellm.APIConnectionError, + "Missing Azure AI API Base", + ), + "azure-empty-token": PublicFailure( + lambda _: { + "model": "azure_ai/mistral-ocr-latest", + "api_key": None, + "azure_ad_token": "static-token", + "azure_ad_token_provider": empty_token, + }, + litellm.APIConnectionError, + "Missing Azure AI credentials", + ), + "upstream-500": PublicFailure( + lambda _: {}, + litellm.InternalServerError, + "provider unavailable", + provider_requests=1, + response=ResponseSpec(body={"message": "provider unavailable"}, status=500), + ), + "invalid-provider-response": PublicFailure( + lambda _: {}, + litellm.APIConnectionError, + "pages", + provider_requests=1, + response=ResponseSpec(body={"pages": "invalid"}), + ), + "response-over-limit": PublicFailure( + lambda _: {"max_response_bytes": len(json.dumps(OCR_RESPONSE).encode()) - 1}, + litellm.APIConnectionError, + "OCR response exceeds the size limit", + provider_requests=1, + ), + "timeout": PublicFailure( + lambda _: {"timeout": 0.01}, + litellm.Timeout, + "", + provider_requests=1, + response=ResponseSpec(body=OCR_RESPONSE, delay=0.2), + ), +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("failure", PUBLIC_FAILURES.values(), ids=PUBLIC_FAILURES.keys()) +async def test_native_failures_raise_the_public_exception_class( + server: RecordingServer, + isolated_azure_auth: None, + tmp_path: Path, + asynchronous: bool, + failure: PublicFailure, +) -> None: + server.expected_requests = failure.provider_requests + if failure.response is not None: + server.enqueue(failure.response) + + with pytest.raises(failure.error, match=failure.match) as caught: + await call(server, asynchronous, **failure.arguments(tmp_path)) + + assert len(server.requests) == failure.provider_requests + if failure.cause is not None: + assert isinstance(caught.value.__context__, failure.cause) diff --git a/tests/test_litellm_rust/ocr/test_guardrails.py b/tests/test_litellm_rust/ocr/test_guardrails.py index de4590ba202..cfe221d2994 100644 --- a/tests/test_litellm_rust/ocr/test_guardrails.py +++ b/tests/test_litellm_rust/ocr/test_guardrails.py @@ -8,19 +8,13 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks from litellm.types.utils import CallTypes +from tests.test_litellm_rust.ocr.data import call_aocr from tests.test_litellm_rust.support.callback_recorder import RecordingLogger -from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec -from tests.test_litellm_rust.support.requests import OCR_RESPONSE, call_native, call_native_aocr +from tests.test_litellm_rust.support.recording_server import RecordingServer pytestmark = pytest.mark.requires_rust_extension -@pytest.fixture -def ocr_server(recording_server: RecordingServer) -> RecordingServer: - recording_server.default_response = ResponseSpec(body=OCR_RESPONSE) - return recording_server - - class ReplaceOCRMarkdown(CustomGuardrail): def __init__(self) -> None: super().__init__( @@ -36,7 +30,7 @@ class ReplaceOCRMarkdown(CustomGuardrail): @pytest.mark.asyncio async def test_native_aocr_post_call_content_filter_blocks_matching_markdown( - ocr_server: RecordingServer, + server: RecordingServer, ) -> None: guardrail: Final = ContentFilterGuardrail( guardrail_name="block-native-ocr-markdown", @@ -46,22 +40,22 @@ async def test_native_aocr_post_call_content_filter_blocks_matching_markdown( litellm.callbacks.append(guardrail) with pytest.raises(HTTPException, match="Content blocked") as blocked: - await call_native_aocr(ocr_server, guardrails=[guardrail.guardrail_name]) + await call_aocr(server, guardrails=[guardrail.guardrail_name]) assert blocked.value.status_code == 400 - assert len(ocr_server.requests) == 1 + assert len(server.requests) == 1 @pytest.mark.asyncio async def test_native_aocr_post_call_replacement_reaches_caller_and_success_callback( - ocr_server: RecordingServer, + server: RecordingServer, ) -> None: guardrail: Final = ReplaceOCRMarkdown() recorder: Final = RecordingLogger() litellm.callbacks.append(guardrail) - response: Final = await call_native_aocr( - ocr_server, + response: Final = await call_aocr( + server, callbacks=[recorder], guardrails=[guardrail.guardrail_name], ) @@ -71,4 +65,4 @@ async def test_native_aocr_post_call_replacement_reaches_caller_and_success_call assert response.pages[0].markdown == "Reviewed OCR" assert len(success_events) == 1 assert success_events[0].response.pages[0].markdown == "Reviewed OCR" - assert "guardrails" not in ocr_server.requests[0].body + assert "guardrails" not in server.requests[0].body diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index 40bcb992b22..c262454a678 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -15,51 +15,45 @@ from litellm._logging import trace_id_var from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from tests.test_litellm_rust.ocr.data import OCR_RESPONSE, call_aocr, call_ocr from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec -from tests.test_litellm_rust.support.requests import OCR_RESPONSE, call_aocr, call_ocr pytestmark = pytest.mark.requires_rust_extension -@pytest.fixture -def ocr_server(recording_server: RecordingServer) -> RecordingServer: - recording_server.default_response = ResponseSpec(body=OCR_RESPONSE) - return recording_server - - @pytest.mark.asyncio -async def test_proxy_metadata_remains_python_owned(ocr_server: RecordingServer) -> None: +async def test_proxy_metadata_remains_python_owned(server: RecordingServer) -> None: from litellm.proxy._types import UserAPIKeyAuth recorder: Final = RecordingLogger() auth: Final = UserAPIKeyAuth(user_id="ocr-user") response: Final = await call_aocr( - ocr_server, callbacks=[recorder], metadata={"user_api_key_auth": auth}, shared_session=object() + server, callbacks=[recorder], metadata={"user_api_key_auth": auth}, shared_session=object() ) events: Final = await recorder.wait_for_async("async_log_success_event") assert response.pages[0].markdown == "native OCR response" assert response._hidden_params["additional_headers"]["x-litellm-rust"] == "true" assert events[0].kwargs["litellm_params"]["metadata"]["user_api_key_auth"].user_id == "ocr-user" - assert "metadata" not in ocr_server.requests[0].body + assert "metadata" not in server.requests[0].body @pytest.mark.asyncio async def test_request_level_custom_pricing_reaches_logging_params_and_bills_the_call( - ocr_server: RecordingServer, + server: RecordingServer, ) -> None: recorder: Final = RecordingLogger() - response: Final = await call_aocr(ocr_server, callbacks=[recorder], ocr_cost_per_page=0.05) + response: Final = await call_aocr(server, callbacks=[recorder], ocr_cost_per_page=0.05) events: Final = await recorder.wait_for_async("async_log_success_event") assert response.usage_info is not None and response.usage_info.pages_processed == 1 assert events[0].kwargs["litellm_params"]["ocr_cost_per_page"] == 0.05 assert response._hidden_params["response_cost"] == pytest.approx(0.05) - assert "ocr_cost_per_page" not in ocr_server.requests[0].body + assert "ocr_cost_per_page" not in server.requests[0].body @pytest.mark.asyncio -async def test_response_replacement_finalized_before_dispatch_in_caller_task(ocr_server: RecordingServer) -> None: +async def test_response_replacement_finalized_before_dispatch_in_caller_task(server: RecordingServer) -> None: caller: Final = asyncio.current_task() context: Final = ContextVar("lifecycle-test", default="before") observations: Final = [] @@ -76,11 +70,11 @@ async def test_response_replacement_finalized_before_dispatch_in_caller_task(ocr return response.model_copy(update={"model": "replaced"}) litellm.callbacks.append(Replace()) - response: Final = await call_aocr(ocr_server, callbacks=[recorder], litellm_call_id="native-final") + response: Final = await call_aocr(server, callbacks=[recorder], litellm_call_id="native-final") events: Final = await recorder.wait_for_async("async_log_success_event") assert observations == [("pre", caller, "pre"), ("post", caller, "pre")] assert context.get() == "pre" - assert ocr_server.requests[0].body["pages"] == [2] + assert server.requests[0].body["pages"] == [2] assert response.model == "replaced" assert events[0].response is response assert response._hidden_params["litellm_call_id"] == "native-final" @@ -90,7 +84,7 @@ async def test_response_replacement_finalized_before_dispatch_in_caller_task(ocr @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True]) async def test_metadata_failure_dispatches_only_failure_and_releases_logger( - ocr_server: RecordingServer, asynchronous: bool + server: RecordingServer, asynchronous: bool ) -> None: failure: Final = RuntimeError("metadata failed") seen: Final = [] @@ -120,8 +114,8 @@ async def test_metadata_failure_dispatches_only_failure_and_releases_logger( ) reference: Final = weakref.ref(logger) with pytest.raises(RuntimeError) as caught: - await call_aocr(ocr_server, litellm_logging_obj=logger) if asynchronous else call_ocr( - ocr_server, litellm_logging_obj=logger + await call_aocr(server, litellm_logging_obj=logger) if asynchronous else call_ocr( + server, litellm_logging_obj=logger ) assert caught.value is failure failure.__traceback__ = None @@ -132,12 +126,12 @@ async def test_metadata_failure_dispatches_only_failure_and_releases_logger( gc.collect() assert seen == ([("sync", failure), ("async", failure)] if asynchronous else [("sync", failure)]) assert reference() is None - assert len(ocr_server.requests) == 1 + assert len(server.requests) == 1 @pytest.mark.asyncio -async def test_mapped_failure_identity_and_deployment_snapshot(ocr_server: RecordingServer) -> None: - ocr_server.enqueue(ResponseSpec(body={"message": "unavailable"}, status=500)) +async def test_mapped_failure_identity_and_deployment_snapshot(server: RecordingServer) -> None: + server.enqueue(ResponseSpec(body={"message": "unavailable"}, status=500)) recorder: Final = RecordingLogger() snapshots: Final = [] @@ -148,20 +142,20 @@ async def test_mapped_failure_identity_and_deployment_snapshot(ocr_server: Recor litellm.callbacks.append(Observe()) with pytest.raises(litellm.InternalServerError) as caught: - await call_aocr(ocr_server, callbacks=[recorder]) + await call_aocr(server, callbacks=[recorder]) failures: Final = tuple(event for event in recorder.events if "failure" in event.name) assert [event.name for event in failures] == ["log_failure_event", "async_log_failure_event"] assert all(event.kwargs["exception"] is caught.value for event in failures) assert caught.value.status_code == 500 assert snapshots[0] is not caught.value assert snapshots[0].status_code == 418 - assert len(ocr_server.requests) == 1 + assert len(server.requests) == 1 @pytest.mark.asyncio @pytest.mark.parametrize("phase", ["pre", "http"]) async def test_cancellation_cleans_up_in_caller_task_without_terminal_dispatch( - ocr_server: RecordingServer, phase: str + server: RecordingServer, phase: str ) -> None: entered: Final = asyncio.Event() recorder: Final = RecordingLogger() @@ -174,21 +168,21 @@ async def test_cancellation_cleans_up_in_caller_task_without_terminal_dispatch( litellm.callbacks.append(Pause()) if phase == "http": - ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.2)) + server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.2)) if phase == "pre": - ocr_server.expected_requests = 0 + server.expected_requests = 0 restored: Final = [] async def invoke(): trace_id_var.set("parent") try: - await call_aocr(ocr_server, callbacks=[recorder], litellm_trace_id="native-call") + await call_aocr(server, callbacks=[recorder], litellm_trace_id="native-call") finally: restored.append(trace_id_var.get()) task: Final = asyncio.create_task(invoke()) if phase == "http": - await ocr_server.wait_for_requests(1) + await server.wait_for_requests(1) else: await asyncio.wait_for(entered.wait(), 5) task.cancel() @@ -201,9 +195,7 @@ async def test_cancellation_cleans_up_in_caller_task_without_terminal_dispatch( @pytest.mark.asyncio @pytest.mark.parametrize("blocked", [False, True]) -async def test_deferred_logging_requires_release_and_runs_at_most_once( - ocr_server: RecordingServer, blocked: bool -) -> None: +async def test_deferred_logging_requires_release_and_runs_at_most_once(server: RecordingServer, blocked: bool) -> None: recorder: Final = RecordingLogger() logger: Final = Logging( model="mistral-ocr-latest", @@ -216,7 +208,7 @@ async def test_deferred_logging_requires_release_and_runs_at_most_once( dynamic_async_success_callbacks=[recorder], ) logger._defer_async_logging = True - response: Final = await call_aocr(ocr_server, litellm_logging_obj=logger) + response: Final = await call_aocr(server, litellm_logging_obj=logger) await drain_logging() assert "async_log_success_event" not in recorder.names ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, blocked) @@ -228,7 +220,7 @@ async def test_deferred_logging_requires_release_and_runs_at_most_once( assert events[0].response is response -def test_sync_success_uses_executor_and_copied_caller_context(ocr_server: RecordingServer) -> None: +def test_sync_success_uses_executor_and_copied_caller_context(server: RecordingServer) -> None: context: Final = ContextVar("sync-lifecycle", default="missing") context.set("caller") thread: Final = threading.current_thread() @@ -240,7 +232,7 @@ def test_sync_success_uses_executor_and_copied_caller_context(ocr_server: Record observations.append((threading.current_thread(), context.get(), response_obj)) finished.set() - response: Final = call_ocr(ocr_server, callbacks=[Observe()]) + response: Final = call_ocr(server, callbacks=[Observe()]) assert finished.wait(5) assert observations[0][0] is not thread assert observations[0][1] == "caller" @@ -248,44 +240,44 @@ def test_sync_success_uses_executor_and_copied_caller_context(ocr_server: Record @pytest.mark.asyncio -async def test_nested_native_calls_preserve_context_and_dispatch_each_outcome(ocr_server: RecordingServer) -> None: - ocr_server.expected_requests = 2 +async def test_nested_native_calls_preserve_context_and_dispatch_each_outcome(server: RecordingServer) -> None: + server.expected_requests = 2 recorder: Final = RecordingLogger() outcomes: Final = [] class Nested(CustomLogger): async def async_pre_call_deployment_hook(self, kwargs, call_type): if kwargs.get("litellm_call_id") == "outer": - outcomes.append(await call_aocr(ocr_server, callbacks=[recorder], litellm_call_id="inner")) + outcomes.append(await call_aocr(server, callbacks=[recorder], litellm_call_id="inner")) litellm.callbacks.append(Nested()) - outcomes.append(await call_aocr(ocr_server, callbacks=[recorder], litellm_call_id="outer")) + outcomes.append(await call_aocr(server, callbacks=[recorder], litellm_call_id="outer")) events: Final = await recorder.wait_for_async("async_log_success_event", count=2) assert [event.kwargs["litellm_call_id"] for event in events] == ["inner", "outer"] assert events[0].response is outcomes[0] assert events[1].response is outcomes[1] - assert len(ocr_server.requests) == 2 + assert len(server.requests) == 2 -def test_sync_pre_call_can_make_nested_native_request(ocr_server: RecordingServer) -> None: - ocr_server.expected_requests = 2 +def test_sync_pre_call_can_make_nested_native_request(server: RecordingServer) -> None: + server.expected_requests = 2 observed: Final = [] class Nested(CustomLogger): def log_pre_api_call(self, model, messages, kwargs): if kwargs["litellm_call_id"] == "outer-sync": - observed.append(call_ocr(ocr_server, litellm_call_id="inner-sync")) + observed.append(call_ocr(server, litellm_call_id="inner-sync")) - response: Final = call_ocr(ocr_server, callbacks=[Nested()], litellm_call_id="outer-sync") + response: Final = call_ocr(server, callbacks=[Nested()], litellm_call_id="outer-sync") assert observed[0].pages[0].markdown == response.pages[0].markdown - assert len(ocr_server.requests) == 2 + assert len(server.requests) == 2 -def test_unstarted_native_coroutine_releases_input_without_reading_file(ocr_server: RecordingServer) -> None: +def test_unstarted_native_coroutine_releases_input_without_reading_file(server: RecordingServer) -> None: from litellm.ocr.dispatch import _public_request from litellm.rust_bridge import _native - ocr_server.expected_requests = 0 + server.expected_requests = 0 effects: Final = [] class File: @@ -308,7 +300,7 @@ def test_unstarted_native_coroutine_releases_input_without_reading_file(ocr_serv @pytest.mark.asyncio -async def test_file_read_happens_after_deployment_hook_in_caller_task(ocr_server: RecordingServer) -> None: +async def test_file_read_happens_after_deployment_hook_in_caller_task(server: RecordingServer) -> None: effects: Final = [] caller: Final = asyncio.current_task() @@ -323,13 +315,13 @@ async def test_file_read_happens_after_deployment_hook_in_caller_task(ocr_server effects.append(("hook", asyncio.current_task())) litellm.callbacks.append(Deployment()) - await call_aocr(ocr_server, document={"type": "file", "file": File()}) + await call_aocr(server, document={"type": "file", "file": File()}) assert effects == [("hook", caller), ("read", caller)] @pytest.mark.asyncio -async def test_failure_callbacks_continue_within_both_families(ocr_server: RecordingServer) -> None: - ocr_server.enqueue(ResponseSpec(body={"message": "failed"}, status=500)) +async def test_failure_callbacks_continue_within_both_families(server: RecordingServer) -> None: + server.enqueue(ResponseSpec(body={"message": "failed"}, status=500)) observed: Final = [] class Broken(CustomLogger): @@ -349,7 +341,7 @@ async def test_failure_callbacks_continue_within_both_families(ocr_server: Recor observed.append(("following-async", kwargs["exception"])) with pytest.raises(litellm.InternalServerError) as caught: - await call_aocr(ocr_server, callbacks=[Broken(), Following()]) + await call_aocr(server, callbacks=[Broken(), Following()]) assert [name for name, _ in observed] == ["broken-sync", "following-sync", "broken-async", "following-async"] assert all(error is caught.value for _, error in observed) @@ -392,17 +384,17 @@ async def test_cancelling_native_transport_closes_connection_before_return() -> @pytest.mark.asyncio -async def test_response_limit_is_enforced_at_the_public_boundary(ocr_server: RecordingServer) -> None: +async def test_response_limit_is_enforced_at_the_public_boundary(server: RecordingServer) -> None: limit: Final = len(json.dumps(OCR_RESPONSE).encode()) - 1 with pytest.raises(litellm.APIConnectionError, match="OCR response exceeds the size limit"): - await call_aocr(ocr_server, max_response_bytes=limit) - assert len(ocr_server.requests) == 1 + await call_aocr(server, max_response_bytes=limit) + assert len(server.requests) == 1 @pytest.mark.asyncio @pytest.mark.parametrize("failure", [False, True]) async def test_empty_callbacks_run_deployment_hooks_and_defer_like_the_python_client_wrapper( - ocr_server: RecordingServer, + server: RecordingServer, monkeypatch: pytest.MonkeyPatch, failure: bool, created_loggers: list[Logging], @@ -439,14 +431,14 @@ async def test_empty_callbacks_run_deployment_hooks_and_defer_like_the_python_cl monkeypatch.setattr(litellm_logging, "executor", probe) monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", probe) if failure: - ocr_server.enqueue(ResponseSpec(body={"message": "provider failed"}, status=500)) + server.enqueue(ResponseSpec(body={"message": "provider failed"}, status=500)) trace_id_var.set("callback-free-parent") arguments: Final = {"litellm_trace_id": "callback-free-call", "litellm_call_id": "callback-free-id"} if failure: with pytest.raises(litellm.InternalServerError): - await call_aocr(ocr_server, **arguments) + await call_aocr(server, **arguments) else: - response: Final = await call_aocr(ocr_server, **arguments) + response: Final = await call_aocr(server, **arguments) assert response.pages[0].markdown == "native OCR response" assert response._hidden_params["litellm_call_id"] == "callback-free-id" assert response._hidden_params["response_cost"] is not None @@ -468,20 +460,18 @@ async def test_empty_callbacks_run_deployment_hooks_and_defer_like_the_python_cl @pytest.mark.parametrize( "registration", ["success_callback", "_async_success_callback", "failure_callback", "_async_failure_callback"] ) -async def test_terminal_registration_added_during_http_is_observed( - ocr_server: RecordingServer, registration: str -) -> None: +async def test_terminal_registration_added_during_http_is_observed(server: RecordingServer, registration: str) -> None: failure: Final = "failure" in registration observer: Final = RecordingLogger() - ocr_server.enqueue( + server.enqueue( ResponseSpec( body={"message": "provider failed"} if failure else OCR_RESPONSE, status=500 if failure else 200, delay=0.1 ) ) task: Final = asyncio.create_task( - asyncio.to_thread(call_ocr, ocr_server) if registration == "success_callback" else call_aocr(ocr_server) + asyncio.to_thread(call_ocr, server) if registration == "success_callback" else call_aocr(server) ) - await ocr_server.wait_for_requests(1) + await server.wait_for_requests(1) getattr(litellm, registration).append(observer) if failure: with pytest.raises(litellm.InternalServerError): @@ -523,7 +513,7 @@ def created_loggers(monkeypatch: pytest.MonkeyPatch) -> list[Logging]: @pytest.mark.asyncio @pytest.mark.parametrize("consumer", ["logger_fn", "raw_global", "request_debug"]) async def test_explicit_logging_consumers_keep_request_and_response_payloads( - ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, created_loggers: list[Logging], consumer: str + server: RecordingServer, monkeypatch: pytest.MonkeyPatch, created_loggers: list[Logging], consumer: str ) -> None: snapshots: Final[list[dict[str, object]]] = [] if consumer == "raw_global": @@ -533,7 +523,7 @@ async def test_explicit_logging_consumers_keep_request_and_response_payloads( "raw_global": {}, "request_debug": {"litellm_request_debug": True}, }[consumer] - response: Final = await call_aocr(ocr_server, **arguments) + response: Final = await call_aocr(server, **arguments) details: Final = created_loggers[0].model_call_details assert details["additional_args"]["complete_input_dict"]["model"] == "mistral-ocr-latest" assert json.loads(details["original_response"])["pages"][0]["markdown"] == response.pages[0].markdown diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py index 8c4f3a403d6..c7b7a6d4059 100644 --- a/tests/test_litellm_rust/ocr/test_requests.py +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -1,160 +1,36 @@ -import json -from collections.abc import Callable -from dataclasses import dataclass from io import BytesIO -from pathlib import Path -from typing import Final, NoReturn +from typing import Final -import httpx import pytest -from pydantic import JsonValue -import litellm -from tests.test_litellm_rust.support.callback_recorder import RecordingLogger -from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec -from tests.test_litellm_rust.support.requests import ( +from tests.test_litellm_rust.ocr.data import ( OCR_DOCUMENT, OCR_RESPONSE, - call_native, - call_native_aocr, - call_native_ocr, + call_aocr, + call_ocr, ) -from tests.test_litellm_rust.support.served import assert_served_natively +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec +from tests.test_litellm_rust.support.routing import assert_served_natively pytestmark = pytest.mark.requires_rust_extension -@pytest.fixture(params=[False, True], ids=["python", "rust"]) -def ocr_backend(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> bool: - enabled: Final = bool(request.param) - monkeypatch.setenv("LITELLM_RUST", "1" if enabled else "0") - return enabled - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -async def test_ocr_contract_upstream_status( - ocr_server: RecordingServer, - ocr_backend: bool, - asynchronous: bool, -) -> None: - upstream: Final = ResponseSpec(body={"detail": "invalid provider option"}, status=422) - ocr_server.enqueue(upstream) - arguments: Final = { - "model": "vertex_ai/mistral-ocr-latest", - "vertex_project": "test-project", - "vertex_location": "us-central1", - "num_retries": 0, - } - with pytest.raises(litellm.BadRequestError) as caught: - await call_native(ocr_server, asynchronous, **arguments) - assert caught.value.status_code == upstream.status - assert caught.value.response.status_code == upstream.status - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -@pytest.mark.parametrize("preserved", ["body", "headers"]) -async def test_ocr_contract_provider_error_details( - ocr_server: RecordingServer, - ocr_backend: bool, - asynchronous: bool, - preserved: str, -) -> None: - payload: Final = {"message": "rate limited"} - headers: Final = {"Retry-After": "17", "X-Request-ID": "ocr-request-123", "X-Future-Header": "retained"} - ocr_server.enqueue(ResponseSpec(body=payload, status=429, headers=headers)) - with pytest.raises(litellm.RateLimitError) as caught: - await call_native(ocr_server, asynchronous, num_retries=0) - response: Final = caught.value.response - assert isinstance(response, httpx.Response) - if preserved == "body": - assert response.content == json.dumps(payload).encode() - else: - for name, value in headers.items(): - assert response.headers.get(name.lower()) == value - assert response.headers.get(name.upper()) == value - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -async def test_ocr_contract_invalid_response_format( - ocr_server: RecordingServer, - ocr_backend: bool, - asynchronous: bool, -) -> None: - ocr_server.expected_requests = 0 - with pytest.raises(litellm.UnsupportedParamsError) as caught: - await call_native(ocr_server, asynchronous, req_format="bogus", num_retries=0) - assert caught.value.status_code == 400 - for value in ("req_format", "bogus", "native", "litellm"): - assert value in str(caught.value) - assert ocr_server.requests == [] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -@pytest.mark.parametrize( - "document,field", - [ - ([], "document"), - ({"document_url": "https://example.com/a.pdf"}, "type"), - ({"type": "text"}, "type"), - ], -) -async def test_ocr_contract_malformed_document_is_actionable( - ocr_server: RecordingServer, - ocr_backend: bool, - asynchronous: bool, - document: JsonValue, - field: str, -) -> None: - ocr_server.expected_requests = None - with pytest.raises(litellm.BadRequestError) as caught: - await call_native(ocr_server, asynchronous, document=document, num_retries=0) - assert caught.value.status_code == 400 - assert field.lower() in str(caught.value).lower() - assert "NoneType: None" not in str(caught.value) - assert "indices must be" not in str(caught.value) - assert ocr_server.requests == [] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -@pytest.mark.parametrize("option,value,field", [("pages", [-1], "pages"), ("features", [1], "features")]) -async def test_ocr_contract_azure_invalid_options_are_bad_requests( - ocr_server: RecordingServer, - ocr_backend: bool, - asynchronous: bool, - option: str, - value: JsonValue, - field: str, -) -> None: - ocr_server.expected_requests = 0 - arguments: Final = {"model": "azure_ai/doc-intelligence/prebuilt-read", option: value, "num_retries": 0} - with pytest.raises(litellm.BadRequestError) as caught: - await call_native(ocr_server, asynchronous, **arguments) - assert caught.value.status_code == 400 - assert field in str(caught.value) - assert ocr_server.requests == [] - - @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) @pytest.mark.parametrize("model", ["mistral/mistral-ocr-latest", "azure_ai/mistral-ocr-latest", "reducto/parse-v3"]) async def test_ocr_contract_native_format_supported( - ocr_server: RecordingServer, + server: RecordingServer, ocr_backend: bool, asynchronous: bool, model: str, ) -> None: - ocr_server.expected_requests = None + server.expected_requests = None payload: Final = ( {"result": {"chunks": [{"content": "native OCR response"}]}, "usage": {"num_pages": 1}} if model.startswith("reducto/") else OCR_RESPONSE ) - ocr_server.default_response = ResponseSpec(body=payload) + server.default_response = ResponseSpec(body=payload) arguments: Final = { "model": model, "req_format": "native", @@ -163,43 +39,39 @@ async def test_ocr_contract_native_format_supported( if model.startswith("reducto/") else OCR_DOCUMENT, } - response: Final = ( - await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments) - ) + response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments) assert response.pages[0].markdown == "native OCR response" assert response.get_provider_native_response() == payload - assert len(ocr_server.requests) == 1 + assert len(server.requests) == 1 if ocr_backend: - assert_served_natively(ocr_server, response) + assert_served_natively(server, response) @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) async def test_ocr_contract_unknown_reducto_model_reaches_provider( - ocr_server: RecordingServer, + server: RecordingServer, ocr_backend: bool, asynchronous: bool, ) -> None: - ocr_server.default_response = ResponseSpec(body={"result": {"chunks": [{"content": "future model response"}]}}) + server.default_response = ResponseSpec(body={"result": {"chunks": [{"content": "future model response"}]}}) arguments: Final = { "model": "reducto/future-parse-model", "document": {"type": "document_url", "document_url": "reducto://ready.pdf"}, "num_retries": 0, } - response: Final = ( - await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments) - ) + response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments) assert response.model == "future-parse-model" assert response.pages[0].markdown == "future model response" - assert len(ocr_server.requests) == 1 - assert ocr_server.requests[0].path == "/parse" - assert ocr_server.requests[0].body == {"input": "reducto://ready.pdf"} + assert len(server.requests) == 1 + assert server.requests[0].path == "/parse" + assert server.requests[0].body == {"input": "reducto://ready.pdf"} @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) async def test_native_azure_ocr_uses_token_provider_result_as_bearer_token( - ocr_server: RecordingServer, isolated_azure_auth: None, asynchronous: bool + server: RecordingServer, isolated_azure_auth: None, asynchronous: bool ) -> None: calls: Final = [] @@ -212,478 +84,24 @@ async def test_native_azure_ocr_uses_token_provider_result_as_bearer_token( "api_key": None, "azure_ad_token_provider": token_provider, } - response: Final = ( - await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments) - ) + response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments) assert calls == ["token"] assert response.pages[0].markdown == "native OCR response" - assert_served_natively(ocr_server, response) - assert ocr_server.requests[0].headers["authorization"] == "Bearer callback-token" + assert_served_natively(server, response) + assert server.requests[0].headers["authorization"] == "Bearer callback-token" -@pytest.fixture -def ocr_server(recording_server: RecordingServer) -> RecordingServer: - recording_server.default_response = ResponseSpec(body=OCR_RESPONSE) - return recording_server - - -def test_native_ocr_maps_provider_400_with_public_provider_details(ocr_server: RecordingServer) -> None: - ocr_server.enqueue(ResponseSpec(body={"message": "invalid OCR request"}, status=400)) - - with pytest.raises(litellm.BadRequestError) as caught: - call_native_ocr(ocr_server) - - assert caught.value.status_code == 400 - assert caught.value.model == "mistral-ocr-latest" - assert caught.value.llm_provider == "mistral" - assert "invalid OCR request" in str(caught.value) - - -def test_native_ocr_encodes_python_file_input_and_drops_unknown_arguments(ocr_server: RecordingServer) -> None: - response: Final = call_native_ocr( - ocr_server, +def test_native_ocr_encodes_python_file_input_and_drops_unknown_arguments(server: RecordingServer) -> None: + response: Final = call_ocr( + server, document={"type": "file", "file": BytesIO(b"abc"), "mime_type": "image/png"}, opaque_extension=object(), ) assert response.pages[0].markdown == "native OCR response" - assert_served_natively(ocr_server, response) - assert ocr_server.requests[0].body == { + assert_served_natively(server, response) + assert server.requests[0].body == { "model": "mistral-ocr-latest", "document": {"type": "image_url", "image_url": "data:image/png;base64,YWJj"}, } - - -class TokenAbort(BaseException): - pass - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -@pytest.mark.parametrize("failure", ["ordinary", "abort"], ids=["value-error", "base-exception"]) -async def test_native_azure_ocr_token_provider_failure_prevents_pre_call_callback_and_request( - ocr_server: RecordingServer, - isolated_azure_auth: None, - asynchronous: bool, - failure: str, -) -> None: - ocr_server.expected_requests = 0 - calls: Final = [] - recorder: Final = RecordingLogger() - original: Final = {"ordinary": ValueError("token unavailable"), "abort": TokenAbort("abort")} - - def token_provider() -> object: - calls.append("token") - raise original[failure] - - arguments: Final = { - "model": "azure_ai/mistral-ocr-latest", - "api_key": None, - "azure_ad_token_provider": token_provider, - "callbacks": [recorder], - } - expected: Final = TokenAbort if failure == "abort" else litellm.APIConnectionError - with pytest.raises(expected) as caught: - await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments) - assert calls == ["token"] - assert ocr_server.requests == [] - assert "log_pre_api_call" not in recorder.names - if failure == "ordinary": - assert "Failed to get Azure AD token: token unavailable" in str(caught.value) - assert isinstance(caught.value.__context__, RuntimeError) - assert caught.value.__context__.__cause__ is original[failure] - else: - assert caught.value is original[failure] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", [False, True]) -@pytest.mark.parametrize( - "override, expected_key", - [ - ({}, "credential-key"), - ({"api_key": "explicit-key"}, "explicit-key"), - ({"api_key": None}, "environment-key"), - ], - ids=["inherit", "explicit", "explicit-none"], -) -async def test_native_ocr_inherits_named_credentials_without_overwriting_arguments( - ocr_server: RecordingServer, - monkeypatch: pytest.MonkeyPatch, - asynchronous: bool, - override: dict[str, object], - expected_key: str, -) -> None: - from litellm.models.credentials import CredentialItem - - pages: Final = [0] - opaque: Final = object() - monkeypatch.setenv("MISTRAL_API_KEY", "environment-key") - monkeypatch.setattr( - litellm, - "credential_list", - [ - CredentialItem(credential_name="other", credential_info={}, credential_values={"api_key": "wrong-key"}), - CredentialItem( - credential_name="ocr-test", - credential_info={}, - credential_values={ - "api_key": "credential-key", - "api_base": ocr_server.base_url, - "pages": pages, - "opaque": opaque, - }, - ), - CredentialItem(credential_name="ocr-test", credential_info={}, credential_values={"api_key": "later-key"}), - ], - ) - - class Observer(RecordingLogger): - def log_pre_api_call(self, model, messages, kwargs): - super().log_pre_api_call(model, messages, kwargs) - pages.append(2) - - arguments: Final = { - "model": "mistral/mistral-ocr-latest", - "document": OCR_DOCUMENT, - "litellm_credential_name": "ocr-test", - "callbacks": [Observer()], - **override, - } - response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) - assert response.pages[0].markdown == "native OCR response" - assert ocr_server.requests[0].headers["authorization"] == f"Bearer {expected_key}" - assert ocr_server.requests[0].body["pages"] == [0, 2] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", [False, True]) -async def test_native_file_preparation_preserves_reader_exception( - ocr_server: RecordingServer, asynchronous: bool -) -> None: - ocr_server.expected_requests = 0 - failure: Final = RuntimeError("reader failed") - - class Reader: - def read(self) -> bytes: - raise failure - - document: Final = {"type": "file", "file": Reader()} - with pytest.raises(litellm.APIConnectionError, match="reader failed") as caught: - await call_native_aocr(ocr_server, document=document) if asynchronous else call_native_ocr( - ocr_server, document=document - ) - assert caught.value.__context__ is failure - - -COHERE_IMAGE: Final = {"type": "image_url", "image_url": "data:image/png;base64,YWJj"} -FILE_SIZE_LIMIT: Final = 50 * 1024 * 1024 - - -class IntReader: - def read(self) -> int: - return 1 - - -def oversized_file(tmp_path: Path) -> Path: - path: Final = tmp_path / "large.pdf" - with path.open("wb") as stream: - stream.truncate(FILE_SIZE_LIMIT + 1) - return path - - -def empty_token() -> str: - return "" - - -def unused_token() -> str: - raise AssertionError("the token provider must not run") - - -@dataclass(frozen=True, slots=True) -class PublicFailure: - arguments: Callable[[Path], dict[str, object]] - error: type[Exception] - match: str - provider_requests: int = 0 - response: ResponseSpec | None = None - cause: type[BaseException] | None = None - - -PUBLIC_FAILURES: Final = { - "unknown-req-format": PublicFailure( - lambda _: {"req_format": "raw"}, litellm.BadRequestError, "Invalid `req_format`" - ), - "empty-file": PublicFailure( - lambda _: {"document": {"type": "file", "file": BytesIO(b"")}}, litellm.BadRequestError, "File is empty" - ), - "oversized-file": PublicFailure( - lambda tmp_path: {"document": {"type": "file", "file": oversized_file(tmp_path)}}, - litellm.BadRequestError, - "exceeds the size limit", - ), - "missing-file": PublicFailure( - lambda tmp_path: {"document": {"type": "file", "file": tmp_path / "missing.pdf"}}, - litellm.APIConnectionError, - "File not found", - cause=FileNotFoundError, - ), - "reader-returns-non-bytes": PublicFailure( - lambda _: {"document": {"type": "file", "file": IntReader()}}, - litellm.APIConnectionError, - "bytes or str", - cause=TypeError, - ), - "cohere-non-image": PublicFailure( - lambda _: {"model": "cohere/parse-v5.0"}, litellm.BadRequestError, "only accepts `image_url`" - ), - "cohere-unknown-format": PublicFailure( - lambda _: {"model": "cohere/parse-v5.0", "document": COHERE_IMAGE, "output_format": "html"}, - litellm.BadRequestError, - "output_format", - ), - "azure-missing-api-base": PublicFailure( - lambda _: { - "model": "azure_ai/mistral-ocr-latest", - "api_key": None, - "api_base": None, - "azure_ad_token_provider": unused_token, - }, - litellm.APIConnectionError, - "Missing Azure AI API Base", - ), - "azure-empty-token": PublicFailure( - lambda _: { - "model": "azure_ai/mistral-ocr-latest", - "api_key": None, - "azure_ad_token": "static-token", - "azure_ad_token_provider": empty_token, - }, - litellm.APIConnectionError, - "Missing Azure AI credentials", - ), - "upstream-500": PublicFailure( - lambda _: {}, - litellm.InternalServerError, - "provider unavailable", - provider_requests=1, - response=ResponseSpec(body={"message": "provider unavailable"}, status=500), - ), - "invalid-provider-response": PublicFailure( - lambda _: {}, - litellm.APIConnectionError, - "pages", - provider_requests=1, - response=ResponseSpec(body={"pages": "invalid"}), - ), - "response-over-limit": PublicFailure( - lambda _: {"max_response_bytes": len(json.dumps(OCR_RESPONSE).encode()) - 1}, - litellm.APIConnectionError, - "OCR response exceeds the size limit", - provider_requests=1, - ), - "timeout": PublicFailure( - lambda _: {"timeout": 0.01}, - litellm.Timeout, - "", - provider_requests=1, - response=ResponseSpec(body=OCR_RESPONSE, delay=0.2), - ), -} - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -@pytest.mark.parametrize("failure", PUBLIC_FAILURES.values(), ids=PUBLIC_FAILURES.keys()) -async def test_native_failures_raise_the_public_exception_class( - ocr_server: RecordingServer, - isolated_azure_auth: None, - tmp_path: Path, - asynchronous: bool, - failure: PublicFailure, -) -> None: - ocr_server.expected_requests = failure.provider_requests - if failure.response is not None: - ocr_server.enqueue(failure.response) - - with pytest.raises(failure.error, match=failure.match) as caught: - await call_native(ocr_server, asynchronous, **failure.arguments(tmp_path)) - - assert len(ocr_server.requests) == failure.provider_requests - if failure.cause is not None: - assert isinstance(caught.value.__context__, failure.cause) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -@pytest.mark.parametrize( - "name,value", - [ - ("ssl_verify", object()), - ("ssl_certificate", 1), - ("ssl_certificate", ""), - ("user_url_allowed_hosts", ["example.test", 1]), - ], -) -async def test_native_settings_fail_before_provider_io( - ocr_server: RecordingServer, - monkeypatch: pytest.MonkeyPatch, - asynchronous: bool, - name: str, - value: object, -) -> None: - ocr_server.expected_requests = 0 - monkeypatch.setattr(litellm, name, value) - with pytest.raises(ValueError, match=r"http_settings|url_policy"): - await call_native(ocr_server, asynchronous, num_retries=0) - assert ocr_server.requests == [] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -@pytest.mark.parametrize( - "name,value", - [ - ("vertex_project", 1), - ("vertex_location", ["region"]), - ], -) -async def test_native_vertex_globals_fail_before_provider_io_only_for_vertex_calls( - ocr_server: RecordingServer, - monkeypatch: pytest.MonkeyPatch, - asynchronous: bool, - name: str, - value: object, -) -> None: - ocr_server.expected_requests = 1 - monkeypatch.setattr(litellm, name, value) - await call_native(ocr_server, asynchronous, num_retries=0) - with pytest.raises(ValueError, match=r"provider_defaults"): - await call_native(ocr_server, asynchronous, model="vertex_ai/mistral-ocr-latest", num_retries=0) - assert len(ocr_server.requests) == 1 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -async def test_native_ssl_context_is_terminal_configuration(ocr_server: RecordingServer, asynchronous: bool) -> None: - import ssl - - ocr_server.expected_requests = 0 - context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) - with pytest.raises(ValueError, match=r"request\.ssl_verify.*SSLContext"): - await call_native(ocr_server, asynchronous, ssl_verify=context, num_retries=0) - assert ocr_server.requests == [] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -async def test_native_settings_preserve_protocol_failures( - ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool -) -> None: - ocr_server.expected_requests = 0 - failure: Final = LookupError("settings truth test failed") - cause: Final = RuntimeError("settings cause") - - class RaisesBool: - def __bool__(self) -> bool: - raise failure from cause - - monkeypatch.setattr(litellm, "force_ipv4", RaisesBool()) - with pytest.raises(LookupError) as caught: - await call_native(ocr_server, asynchronous, num_retries=0) - assert caught.value is failure - assert caught.value.__cause__ is cause - assert caught.value.__traceback__ is not None - assert ocr_server.requests == [] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -async def test_native_settings_observe_mutation_between_calls( - ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool -) -> None: - monkeypatch.setattr(litellm, "force_ipv4", "yes") - monkeypatch.setattr(litellm, "http2", 1) - monkeypatch.setattr(litellm, "vertex_project", []) - monkeypatch.setattr(litellm, "vertex_location", 0) - monkeypatch.setattr(litellm, "user_url_allowed_hosts", "EXAMPLE.TEST.") - response: Final = await call_native(ocr_server, asynchronous, num_retries=0) - assert response.pages[0].markdown == "native OCR response" - assert_served_natively(ocr_server, response) - monkeypatch.setattr(litellm, "ssl_certificate", 1) - with pytest.raises(ValueError, match=r"http_settings\.ssl_certificate"): - await call_native(ocr_server, asynchronous, num_retries=0) - assert len(ocr_server.requests) == 1 - - -@pytest.mark.parametrize("required", [False, True]) -@pytest.mark.parametrize("failure", ["invalid", "live", "schema"]) -def test_native_projection_errors_never_select_python( - ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, required: bool, failure: str -) -> None: - import dataclasses - import ssl - - from litellm.rust_bridge import runtime - from litellm.rust_bridge.host import settings - from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule - from litellm.rust_bridge.configuration import Rollout - from litellm.rust_bridge.ocr.entrypoints import NATIVE_OCR - from litellm.rust_bridge.public_call import NativeCall - - ocr_server.expected_requests = 0 - snapshot: Final = dataclasses.replace(settings.http_settings(), user_agent=1) - if failure == "schema": - monkeypatch.setattr(settings, "http_settings", lambda: snapshot) - else: - monkeypatch.setattr( - litellm, "ssl_verify", ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) if failure == "live" else object() - ) - request: Final = NativeCall( - args=(), - kwargs={}, - base={ - "model": "mistral/mistral-ocr-latest", - "document": OCR_DOCUMENT, - "api_key": "test-key", - "api_base": ocr_server.base_url, - "timeout": None, - "custom_llm_provider": "mistral", - "extra_headers": None, - }, - ) - - def python_fallback() -> NoReturn: - pytest.fail("projection failures must not select Python") - - with pytest.raises(RuntimeError if failure == "schema" else ValueError, match="http_settings"): - runtime.run( - RouteContext(Route.OCR, provider="mistral"), - binding=NATIVE_OCR, - native=lambda native: native(request), - python=python_fallback, - rules=(RouteRule(Route.OCR, Rollout.RUST_REQUIRED if required else Rollout.RUST_OPT_OUT),), - ) - assert ocr_server.requests == [] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -@pytest.mark.parametrize("present", [False, True], ids=["missing", "invalid-pem"]) -async def test_native_client_certificate_is_validated_before_io( - ocr_server: RecordingServer, - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, - asynchronous: bool, - present: bool, -) -> None: - ocr_server.expected_requests = 0 - certificate: Final = tmp_path / "client.pem" - if present: - certificate.write_text("invalid certificate") - monkeypatch.setattr(litellm, "ssl_certificate", str(certificate)) - with pytest.raises(ValueError, match=r"http_settings\.ssl_certificate.*PEM") as caught: - await call_native(ocr_server, asynchronous, num_retries=0) - assert str(certificate) not in str(caught.value) - assert ocr_server.requests == [] diff --git a/tests/test_litellm_rust/ocr/test_settings.py b/tests/test_litellm_rust/ocr/test_settings.py new file mode 100644 index 00000000000..7505468db32 --- /dev/null +++ b/tests/test_litellm_rust/ocr/test_settings.py @@ -0,0 +1,198 @@ +from pathlib import Path +from typing import Final + +import pytest + +import litellm +from tests.test_litellm_rust.ocr.data import ( + OCR_DOCUMENT, + call, +) +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.recording_server import RecordingServer +from tests.test_litellm_rust.support.routing import assert_served_natively + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + "override, expected_key", + [ + ({}, "credential-key"), + ({"api_key": "explicit-key"}, "explicit-key"), + ({"api_key": None}, "environment-key"), + ], + ids=["inherit", "explicit", "explicit-none"], +) +async def test_native_ocr_inherits_named_credentials_without_overwriting_arguments( + server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + asynchronous: bool, + override: dict[str, object], + expected_key: str, +) -> None: + from litellm.models.credentials import CredentialItem + + pages: Final = [0] + opaque: Final = object() + monkeypatch.setenv("MISTRAL_API_KEY", "environment-key") + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem(credential_name="other", credential_info={}, credential_values={"api_key": "wrong-key"}), + CredentialItem( + credential_name="ocr-test", + credential_info={}, + credential_values={ + "api_key": "credential-key", + "api_base": server.base_url, + "pages": pages, + "opaque": opaque, + }, + ), + CredentialItem(credential_name="ocr-test", credential_info={}, credential_values={"api_key": "later-key"}), + ], + ) + + class Observer(RecordingLogger): + def log_pre_api_call(self, model, messages, kwargs): + super().log_pre_api_call(model, messages, kwargs) + pages.append(2) + + arguments: Final = { + "model": "mistral/mistral-ocr-latest", + "document": OCR_DOCUMENT, + "litellm_credential_name": "ocr-test", + "callbacks": [Observer()], + **override, + } + response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) + assert response.pages[0].markdown == "native OCR response" + assert server.requests[0].headers["authorization"] == f"Bearer {expected_key}" + assert server.requests[0].body["pages"] == [0, 2] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + "name,value", + [ + ("ssl_verify", object()), + ("ssl_certificate", 1), + ("ssl_certificate", ""), + ("user_url_allowed_hosts", ["example.test", 1]), + ], +) +async def test_native_settings_fail_before_provider_io( + server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + asynchronous: bool, + name: str, + value: object, +) -> None: + server.expected_requests = 0 + monkeypatch.setattr(litellm, name, value) + with pytest.raises(ValueError, match=r"http_settings|url_policy"): + await call(server, asynchronous, num_retries=0) + assert server.requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + "name,value", + [ + ("vertex_project", 1), + ("vertex_location", ["region"]), + ], +) +async def test_native_vertex_globals_fail_before_provider_io_only_for_vertex_calls( + server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + asynchronous: bool, + name: str, + value: object, +) -> None: + server.expected_requests = 1 + monkeypatch.setattr(litellm, name, value) + await call(server, asynchronous, num_retries=0) + with pytest.raises(ValueError, match=r"provider_defaults"): + await call(server, asynchronous, model="vertex_ai/mistral-ocr-latest", num_retries=0) + assert len(server.requests) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_native_ssl_context_is_terminal_configuration(server: RecordingServer, asynchronous: bool) -> None: + import ssl + + server.expected_requests = 0 + context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + with pytest.raises(ValueError, match=r"request\.ssl_verify.*SSLContext"): + await call(server, asynchronous, ssl_verify=context, num_retries=0) + assert server.requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_native_settings_preserve_protocol_failures( + server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool +) -> None: + server.expected_requests = 0 + failure: Final = LookupError("settings truth test failed") + cause: Final = RuntimeError("settings cause") + + class RaisesBool: + def __bool__(self) -> bool: + raise failure from cause + + monkeypatch.setattr(litellm, "force_ipv4", RaisesBool()) + with pytest.raises(LookupError) as caught: + await call(server, asynchronous, num_retries=0) + assert caught.value is failure + assert caught.value.__cause__ is cause + assert caught.value.__traceback__ is not None + assert server.requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_native_settings_observe_mutation_between_calls( + server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool +) -> None: + monkeypatch.setattr(litellm, "force_ipv4", "yes") + monkeypatch.setattr(litellm, "http2", 1) + monkeypatch.setattr(litellm, "vertex_project", []) + monkeypatch.setattr(litellm, "vertex_location", 0) + monkeypatch.setattr(litellm, "user_url_allowed_hosts", "EXAMPLE.TEST.") + response: Final = await call(server, asynchronous, num_retries=0) + assert response.pages[0].markdown == "native OCR response" + assert_served_natively(server) + monkeypatch.setattr(litellm, "ssl_certificate", 1) + with pytest.raises(ValueError, match=r"http_settings\.ssl_certificate"): + await call(server, asynchronous, num_retries=0) + assert len(server.requests) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("present", [False, True], ids=["missing", "invalid-pem"]) +async def test_native_client_certificate_is_validated_before_io( + server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + asynchronous: bool, + present: bool, +) -> None: + server.expected_requests = 0 + certificate: Final = tmp_path / "client.pem" + if present: + certificate.write_text("invalid certificate") + monkeypatch.setattr(litellm, "ssl_certificate", str(certificate)) + with pytest.raises(ValueError, match=r"http_settings\.ssl_certificate.*PEM") as caught: + await call(server, asynchronous, num_retries=0) + assert str(certificate) not in str(caught.value) + assert server.requests == [] diff --git a/tests/test_litellm_rust/responses/__init__.py b/tests/test_litellm_rust/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm_rust/responses/conftest.py b/tests/test_litellm_rust/responses/conftest.py new file mode 100644 index 00000000000..c2ea94473dd --- /dev/null +++ b/tests/test_litellm_rust/responses/conftest.py @@ -0,0 +1,10 @@ +import pytest + +from tests.test_litellm_rust.responses.data import RESPONSES_RESPONSE +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec + + +@pytest.fixture +def server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=RESPONSES_RESPONSE) + return recording_server diff --git a/tests/test_litellm_rust/responses/data.py b/tests/test_litellm_rust/responses/data.py new file mode 100644 index 00000000000..ed489f949a5 --- /dev/null +++ b/tests/test_litellm_rust/responses/data.py @@ -0,0 +1,75 @@ +from typing import Final + +import litellm +from litellm.rust_bridge.catalog import Route +from tests.test_litellm_rust.support.gaps import gap +from tests.test_litellm_rust.support.recording_server import RecordingServer +from tests.test_litellm_rust.support.routing import call_natively + +RESPONSES_MODEL: Final = "openai/gpt-6.1-sol" +RESPONSES_TEXT: Final = "Hello from native Responses" +RESPONSES_RESPONSE: Final = { + "id": "resp_native", + "object": "response", + "created_at": 1, + "status": "completed", + "model": RESPONSES_MODEL.removeprefix("openai/"), + "output": [ + { + "id": "msg_native", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": RESPONSES_TEXT, "annotations": []}], + } + ], + "usage": {"input_tokens": 5, "output_tokens": 4, "total_tokens": 9}, +} +RESPONSES_EVENTS: Final = ( + ( + "response.created", + { + "type": "response.created", + "sequence_number": 0, + "response": {**RESPONSES_RESPONSE, "status": "in_progress", "output": []}, + }, + ), + ( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_native", + "output_index": 0, + "content_index": 0, + "delta": RESPONSES_TEXT, + }, + ), + ("response.completed", {"type": "response.completed", "sequence_number": 2, "response": RESPONSES_RESPONSE}), +) + + +def _error(kind: str) -> dict[str, object]: + return {"error": {"type": kind, "message": kind, "param": None, "code": None}} + + +REQUEST: Final = {"model": RESPONSES_MODEL, "input": "Hello"} +BAD_REQUEST: Final = _error("invalid_request_error") +RATE_LIMITED: Final = _error("rate_limit_exceeded") + + +async def call(server: RecordingServer, *, asynchronous: bool = True, **kwargs: object) -> object: + return await call_natively( + Route.RESPONSES, + litellm.responses, + litellm.aresponses, + REQUEST, + server, + asynchronous=asynchronous, + **kwargs, + ) + + +SYNC_DEPLOYMENT_HOOK: Final = gap("native sync calls skip the async pre-call deployment hook") +STREAMING: Final = gap("native Responses streaming declines to Python", raises=RuntimeError) +BODY_HOOKS: Final = gap("native Responses does not run the handler body hooks") diff --git a/tests/test_litellm_rust/responses/test_body_hooks.py b/tests/test_litellm_rust/responses/test_body_hooks.py new file mode 100644 index 00000000000..59ba3413da9 --- /dev/null +++ b/tests/test_litellm_rust/responses/test_body_hooks.py @@ -0,0 +1,114 @@ +from collections.abc import Iterator +from datetime import datetime +from typing import Final + +import pytest +from pydantic import BaseModel + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging +from tests.test_litellm_rust.responses.data import ( + BODY_HOOKS, + REQUEST, + RESPONSES_EVENTS, + RESPONSES_RESPONSE, + STREAMING, + call, +) +from tests.test_litellm_rust.support.body_hooks import LegacyAgenticLoop, PlannedAgenticLoop, agentic_hooks +from tests.test_litellm_rust.support.callback_recorder import drain_logging +from tests.test_litellm_rust.support.isolation import rebound +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed +from tests.test_litellm_rust.support.results import consume, response_id +from tests.test_litellm_rust.support.routing import CALL_ID, assert_served_natively + +pytestmark = pytest.mark.requires_rust_extension + +STREAM: Final = ResponseSpec(body=None, events=RESPONSES_EVENTS) + + +@pytest.fixture +def server(server: RecordingServer) -> Iterator[RecordingServer]: + yield from relaxed(server) + + +@pytest.mark.asyncio +@STREAMING +async def test_body_streaming_deployment_hook_rewrites_chunks_before_delivery(server: RecordingServer) -> None: + seen: Final[list[str]] = [] + + class Rewrite(CustomLogger): + async def async_post_call_streaming_deployment_hook( + self, request_data: dict[str, object], response_chunk: BaseModel, call_type: object + ) -> BaseModel: + seen.append(getattr(response_chunk, "type", "")) + if getattr(response_chunk, "type", None) != "response.output_text.delta": + return response_chunk + return response_chunk.model_copy(update={"delta": "rewritten"}) + + server.enqueue(STREAM) + with rebound(litellm, "callbacks", [Rewrite()]): + stream: Final = await call(server, stream=True) + chunks: Final = await consume(stream) + + assert_served_natively(server) + assert seen == [event for event, _ in RESPONSES_EVENTS] + assert [getattr(chunk, "delta", None) for chunk in chunks] == [None, "rewritten", None] + + +@pytest.mark.asyncio +@BODY_HOOKS +async def test_body_planned_agentic_loop_sends_the_followup_and_replaces_the_response(server: RecordingServer) -> None: + loop: Final = PlannedAgenticLoop() + server.enqueue(ResponseSpec(body=RESPONSES_RESPONSE)) + server.default_response = ResponseSpec(body={**RESPONSES_RESPONSE, "id": "followup"}) + + with rebound(litellm, "callbacks", [loop]): + response: Final = await call(server) + await drain_logging() + + assert_served_natively(server, count=2) + assert isinstance(server.requests[1].body, dict) + assert server.requests[1].body["input"] == [{"role": "user", "content": "followup"}] + assert response_id(response) == "agentic-replaced" + assert agentic_hooks(loop) == ( + "async_should_run_agentic_loop", + "async_build_agentic_loop_plan", + "async_should_run_agentic_loop", + "async_post_agentic_loop_response_hook", + "async_agentic_loop_cleanup_hook", + ) + + +@pytest.mark.asyncio +@BODY_HOOKS +async def test_body_legacy_agentic_loop_replaces_the_response(server: RecordingServer) -> None: + loop: Final = LegacyAgenticLoop() + + with rebound(litellm, "callbacks", [loop]): + response: Final = await call(server) + + assert_served_natively(server) + assert response_id(response) == "agentic-replaced" + assert agentic_hooks(loop) == ("async_should_run_agentic_loop", "async_run_agentic_loop") + + +@pytest.mark.asyncio +async def test_body_logging_setup_drops_a_logger_fn_from_a_supplied_logging_object(server: RecordingServer) -> None: + observed: Final[list[object]] = [] + logging_obj: Final = Logging( + model=REQUEST["model"], + messages=[], + stream=False, + start_time=datetime.now(), + call_type="aresponses", + litellm_call_id=CALL_ID, + function_id=CALL_ID, + ) + logging_obj.update_environment_variables(litellm_params={"logger_fn": observed.append}, optional_params={}) + + await call(server, litellm_logging_obj=logging_obj) + + assert_served_natively(server) + assert observed == [] diff --git a/tests/test_litellm_rust/responses/test_callbacks.py b/tests/test_litellm_rust/responses/test_callbacks.py new file mode 100644 index 00000000000..a4abdd43529 --- /dev/null +++ b/tests/test_litellm_rust/responses/test_callbacks.py @@ -0,0 +1,408 @@ +import asyncio +from collections.abc import Iterator +from contextvars import ContextVar +from datetime import datetime +from typing import Final + +import pytest +from pydantic import TypeAdapter + +import litellm +from litellm.integrations.argilla import ArgillaLogger +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.integrations.argilla import ArgillaItem +from litellm.types.utils import StandardLoggingPayload +from tests.test_litellm_rust.responses.data import ( + BAD_REQUEST, + RATE_LIMITED, + RESPONSES_RESPONSE, + SYNC_DEPLOYMENT_HOOK, + call, +) +from tests.test_litellm_rust.support.callback_recorder import HookRecorder, RecordingLogger, request_body +from tests.test_litellm_rust.support.gaps import gap +from tests.test_litellm_rust.support.isolation import rebound +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed +from tests.test_litellm_rust.support.results import response_id, with_id +from tests.test_litellm_rust.support.routing import assert_served_natively + +pytestmark = pytest.mark.requires_rust_extension + +SYNC_HANDLERS_ON_ASYNC_CALLS: Final = gap("native async calls also run the sync success and failure handlers") +_OBJECT: Final = TypeAdapter(dict[str, object]) +_HEADERS: Final = TypeAdapter(dict[str, str]) + + +def success_signal(asynchronous: bool) -> str: + return "async_log_success_event" if asynchronous else "log_success_event" + + +@pytest.fixture +def server(server: RecordingServer) -> Iterator[RecordingServer]: + yield from relaxed(server) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "hook,asynchronous,expected", + ( + pytest.param("log_pre_api_call", True, 1, id="async-log_pre_api_call"), + pytest.param("log_post_api_call", True, 1, id="async-log_post_api_call"), + pytest.param("async_log_success_event", True, 1, id="async-async_log_success_event"), + pytest.param("async_logging_hook", True, 1, id="async-async_logging_hook"), + pytest.param("redact_standard_logging_payload_from_model_call_details", True, 1, id="async-redact"), + pytest.param("async_pre_call_deployment_hook", True, 1, id="async-async_pre_call_deployment_hook"), + pytest.param( + "async_post_call_success_deployment_hook", True, 1, id="async-async_post_call_success_deployment_hook" + ), + pytest.param("log_success_event", True, 0, id="async-log_success_event"), + pytest.param("logging_hook", True, 0, id="async-logging_hook"), + pytest.param("log_pre_api_call", False, 1, id="sync-log_pre_api_call"), + pytest.param("log_post_api_call", False, 1, id="sync-log_post_api_call"), + pytest.param("log_success_event", False, 1, id="sync-log_success_event"), + pytest.param("logging_hook", False, 1, id="sync-logging_hook"), + pytest.param( + "async_pre_call_deployment_hook", + False, + 1, + id="sync-async_pre_call_deployment_hook", + marks=SYNC_DEPLOYMENT_HOOK, + ), + pytest.param( + "async_post_call_success_deployment_hook", False, 0, id="sync-async_post_call_success_deployment_hook" + ), + ), +) +async def test_success_hook_fires_the_expected_number_of_times( + server: RecordingServer, hook: str, asynchronous: bool, expected: int +) -> None: + recorder: Final = HookRecorder() + + with rebound(litellm, "callbacks", [recorder]): + await call(server, asynchronous=asynchronous) + await recorder.wait_for_async(success_signal(asynchronous)) + + assert_served_natively(server) + assert recorder.names.count(hook) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "hook,asynchronous,expected", + ( + pytest.param("async_log_failure_event", True, 1, id="async-async_log_failure_event"), + pytest.param( + "async_post_call_failure_deployment_hook", True, 1, id="async-async_post_call_failure_deployment_hook" + ), + pytest.param("log_model_group_rate_limit_error", True, 1, id="async-log_model_group_rate_limit_error"), + pytest.param("log_failure_event", True, 0, id="async-log_failure_event", marks=SYNC_HANDLERS_ON_ASYNC_CALLS), + pytest.param("log_failure_event", False, 1, id="sync-log_failure_event"), + ), +) +async def test_failure_hook_fires_the_expected_number_of_times( + server: RecordingServer, hook: str, asynchronous: bool, expected: int +) -> None: + server.default_response = ResponseSpec(body=RATE_LIMITED, status=429) + recorder: Final = HookRecorder() + + with rebound(litellm, "callbacks", [recorder]): + with pytest.raises(litellm.RateLimitError): + await call(server, asynchronous=asynchronous, metadata={"model_group_size": 1}) + await recorder.wait_for_async("async_log_failure_event" if asynchronous else "log_failure_event") + + assert_served_natively(server) + assert recorder.names.count(hook) == expected + + +@pytest.mark.asyncio +async def test_failure_events_carry_the_error_raised_to_the_caller(server: RecordingServer) -> None: + server.default_response = ResponseSpec(body=BAD_REQUEST, status=400) + recorder: Final = RecordingLogger() + + with rebound(litellm, "callbacks", [recorder]): + with pytest.raises(litellm.BadRequestError) as raised: + await call(server) + events: Final = await recorder.wait_for_async("async_log_failure_event") + + assert_served_natively(server) + assert isinstance(events[0].kwargs, dict) + assert events[0].kwargs["exception"] is raised.value + + +@pytest.mark.asyncio +async def test_pre_api_call_sees_the_wire_body(server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + + await call(server, callbacks=[recorder]) + + assert_served_natively(server) + pre_call: Final = recorder.wait_for("log_pre_api_call")[0].kwargs + assert isinstance(pre_call, dict) + assert request_body(pre_call) == server.requests[0].body + + +@pytest.mark.asyncio +async def test_pre_api_call_edits_reach_the_provider(server: RecordingServer) -> None: + class Edit(CustomLogger): + def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None: + request_body(kwargs)["temperature"] = 0.25 + + await call(server, callbacks=[Edit()]) + + assert_served_natively(server) + assert isinstance(server.requests[0].body, dict) + assert server.requests[0].body["temperature"] == 0.25 + + +@pytest.mark.asyncio +async def test_logging_hook_and_redaction_edits_reach_the_async_success_event(server: RecordingServer) -> None: + class Rewrite(HookRecorder): + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + return {**kwargs, "logging_hook_ran": True}, result + + def redact_standard_logging_payload_from_model_call_details( + self, model_call_details: dict[str, object] + ) -> dict[str, object]: + return {**model_call_details, "redacted": True} + + recorder: Final = Rewrite() + + with rebound(litellm, "callbacks", [recorder]): + await call(server) + success: Final = await recorder.wait_for_async("async_log_success_event") + + assert isinstance(success[0].kwargs, dict) + assert success[0].kwargs["logging_hook_ran"] is True + assert success[0].kwargs["redacted"] is True + + +@pytest.mark.asyncio +async def test_pre_call_deployment_hooks_compose_the_wire_request_in_registration_order( + server: RecordingServer, +) -> None: + class Rewrite(CustomLogger): + def __init__(self, marker: str) -> None: + super().__init__() + self.marker = marker + + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: object + ) -> dict[str, object]: + headers: Final = _HEADERS.validate_python(kwargs["extra_headers"]) + return { + **kwargs, + "extra_headers": {**headers, "x-deployment-hook": headers["x-deployment-hook"] + self.marker}, + } + + with rebound(litellm, "callbacks", [Rewrite("first"), Rewrite("-second")]): + await call(server, extra_headers={"x-deployment-hook": "caller-", "x-kept": "unchanged"}) + + assert_served_natively(server) + assert server.requests[0].headers["x-deployment-hook"] == "caller-first-second" + assert server.requests[0].headers["x-kept"] == "unchanged" + + +@pytest.mark.asyncio +async def test_pre_call_deployment_rejection_logs_failure_without_provider_io(server: RecordingServer) -> None: + rejection: Final = litellm.BadRequestError(message="deployment rejected", model="test-model", llm_provider="openai") + recorder: Final = RecordingLogger() + + class Reject(CustomLogger): + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: object + ) -> dict[str, object]: + raise rejection + + with rebound(litellm, "callbacks", [Reject(), recorder]): + with pytest.raises(litellm.BadRequestError, match="deployment rejected") as raised: + await call(server) + failures: Final = await recorder.wait_for_async("async_log_failure_event") + + assert raised.value is rejection + assert server.requests == [] + assert len(failures) == 1 + assert _OBJECT.validate_python(failures[0].kwargs)["exception"] is raised.value + assert not {"log_pre_api_call", "log_success_event", "async_log_success_event"} & set(recorder.names) + + +@pytest.mark.asyncio +async def test_pre_call_deployment_hook_can_delete_keywords(server: RecordingServer) -> None: + class DropKeywords(CustomLogger): + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: object + ) -> dict[str, object]: + return {name: value for name, value in kwargs.items() if name not in ("instructions", "extra_headers")} + + with rebound(litellm, "callbacks", [DropKeywords()]): + await call( + server, + instructions="remove this instruction", + extra_headers={"x-deleted": "remove this header"}, + ) + + assert_served_natively(server) + sent: Final = server.requests[0] + assert isinstance(sent.body, dict) + assert "instructions" not in sent.body + assert "x-deleted" not in sent.headers + + +@pytest.mark.asyncio +async def test_post_call_success_deployment_hook_replaces_the_response(server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + + class Replace(CustomLogger): + async def async_post_call_success_deployment_hook( + self, request_data: dict[str, object], response: object, call_type: object + ) -> object: + return with_id(response, "replaced") + + with rebound(litellm, "callbacks", [Replace(), recorder]): + response: Final = await call(server) + success: Final = await recorder.wait_for_async("async_log_success_event") + + assert_served_natively(server) + assert response_id(response) == "replaced" + assert len(success) == 1 + assert response_id(success[0].response) == response_id(response) + + +@pytest.mark.asyncio +async def test_concurrent_calls_keep_callback_context_metadata_and_responses_isolated(server: RecordingServer) -> None: + context: Final = ContextVar("responses-callback-context", default="outside") + first_entered: Final = asyncio.Event() + second_entered: Final = asyncio.Event() + server.enqueue(ResponseSpec(body={**RESPONSES_RESPONSE, "id": "resp_first"})) + server.enqueue(ResponseSpec(body={**RESPONSES_RESPONSE, "id": "resp_second"})) + + class Observe(RecordingLogger): + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: object + ) -> dict[str, object]: + arrived, other = ( + (first_entered, second_entered) + if kwargs["litellm_call_id"] == "first" + else (second_entered, first_entered) + ) + arrived.set() + await asyncio.wait_for(other.wait(), timeout=5) + return kwargs + + def log_pre_api_call(self, model: object, messages: object, kwargs: dict[str, object]) -> None: + self._record("log_pre_api_call", {**kwargs, "caller_context": context.get()}) + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self._record("async_log_success_event", {**kwargs, "caller_context": context.get()}, response_obj) + + recorder: Final = Observe() + + async def invoke(marker: str) -> object: + token: Final = context.set(marker) + try: + return await call(server, litellm_call_id=marker, metadata={"marker": marker}) + finally: + context.reset(token) + + with rebound(litellm, "callbacks", [recorder]): + responses: Final = await asyncio.gather(invoke("first"), invoke("second")) + success: Final = await recorder.wait_for_async("async_log_success_event", count=2) + + assert_served_natively(server, count=2) + expected: Final = {"first": responses[0], "second": responses[1]} + pre_call: Final = recorder.wait_for("log_pre_api_call", count=2) + assert len(pre_call) == len(success) == 2 + assert {_OBJECT.validate_python(event.kwargs)["litellm_call_id"] for event in success} == expected.keys() + for event in (*pre_call, *success): + details: Final = _OBJECT.validate_python(event.kwargs) + call_id: Final = str(details["litellm_call_id"]) + params: Final = _OBJECT.validate_python(details["litellm_params"]) + assert details["caller_context"] == call_id + assert _OBJECT.validate_python(params["metadata"])["marker"] == call_id + for event in success: + logged_call_id: Final = str(_OBJECT.validate_python(event.kwargs)["litellm_call_id"]) + assert response_id(event.response) == response_id(expected[logged_call_id]) + assert context.get() == "outside" + + +@pytest.mark.asyncio +async def test_a_failing_logger_does_not_fail_the_call(server: RecordingServer) -> None: + class Broken(CustomLogger): + def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None: + raise RuntimeError("logger exploded") + + response: Final = await call(server, callbacks=[Broken()]) + + assert_served_natively(server) + assert response_id(response) == RESPONSES_RESPONSE["id"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "event,asynchronous,dynamic,expected", + ( + pytest.param("input", False, False, 1, id="global-input"), + pytest.param("success", True, False, 1, id="global-async-success"), + pytest.param("success", True, True, 1, id="request-async-success"), + pytest.param("success", False, False, 0, id="global-sync-success", marks=SYNC_HANDLERS_ON_ASYNC_CALLS), + pytest.param("success", False, True, 0, id="request-sync-success", marks=SYNC_HANDLERS_ON_ASYNC_CALLS), + pytest.param("failure", True, False, 1, id="global-async-failure"), + pytest.param("failure", False, False, 1, id="global-sync-failure"), + pytest.param("failure", False, True, 1, id="request-sync-failure"), + ), +) +async def test_function_callback_fires_the_expected_number_of_times( + server: RecordingServer, event: str, asynchronous: bool, dynamic: bool, expected: int +) -> None: + recorder: Final = RecordingLogger() + + def callback(kwargs: dict[str, object], *_: object) -> None: + recorder._record(event, kwargs) + + async def async_callback(kwargs: dict[str, object], *_: object) -> None: + recorder._record(event, kwargs) + + registered: Final = async_callback if asynchronous else callback + global_list: Final = f"_async_{event}_callback" if asynchronous else f"{event}_callback" + request_callbacks: Final = {f"{event}_callback": [registered]} if dynamic else {} + if event == "failure": + server.default_response = ResponseSpec(body=BAD_REQUEST, status=400) + + with rebound(litellm, "callbacks", [recorder]), rebound(litellm, global_list, [] if dynamic else [registered]): + try: + await call(server, **request_callbacks) + except litellm.BadRequestError: + assert event == "failure" + await recorder.wait_for_async(f"async_log_{'failure' if event == 'failure' else 'success'}_event") + + assert_served_natively(server) + assert recorder.names.count(event) == expected + + +@pytest.mark.asyncio +async def test_dataset_hook_returning_none_drops_the_record_from_the_success_sink(server: RecordingServer) -> None: + datasets: Final = RecordingLogger() + + class Drop(CustomLogger): + async def async_dataset_hook( + self, logged_data: ArgillaItem, standard_logging_payload: StandardLoggingPayload | None + ) -> None: + datasets._record("async_dataset_hook", response=logged_data) + + server.enqueue(ResponseSpec(body={"items": [{"id": "callback-dataset"}]})) + with rebound(litellm, "argilla_transformation_object", {"prompt": "messages", "answer": "response"}): + sink: Final = ArgillaLogger( + argilla_api_key="test-key", + argilla_base_url=server.base_url, + argilla_dataset_name="callback-dataset", + flush_interval=3600, + ) + with rebound(litellm, "callbacks", [sink, Drop()]): + await call(server) + await datasets.wait_for_async("async_dataset_hook") + await sink.flush_queue() + + assert [request.method for request in server.requests] == ["GET", "POST"] diff --git a/tests/test_litellm_rust/responses/test_failures.py b/tests/test_litellm_rust/responses/test_failures.py new file mode 100644 index 00000000000..6dad1bc7278 --- /dev/null +++ b/tests/test_litellm_rust/responses/test_failures.py @@ -0,0 +1,83 @@ +from typing import Final + +import pytest +from pydantic import TypeAdapter + +import litellm +from litellm import RateLimitError +from tests.test_litellm_rust.responses.data import RATE_LIMITED, call +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec + +pytestmark = pytest.mark.requires_rust_extension + +_OBJECT: Final = TypeAdapter(dict[str, object]) + + +def failure_signal(asynchronous: bool) -> str: + return "async_log_failure_event" if asynchronous else "log_failure_event" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_sdk_policy_rejection_precedes_resource_setup_and_is_logged_once( + server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool +) -> None: + monkeypatch.setattr(litellm, "max_budget", 1.0) + monkeypatch.setattr(litellm, "_current_cost", 2.0) + monkeypatch.setattr(litellm, "ssl_verify", object()) + recorder: Final = RecordingLogger() + server.expected_requests = 0 + + with pytest.raises(litellm.BudgetExceededError) as caught: + await call(server, asynchronous=asynchronous, callbacks=[recorder]) + + failure: Final = await recorder.wait_for_async(failure_signal(asynchronous)) + assert len(failure) == 1 + assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value + assert not server.requests + assert not any("success" in name for name in recorder.names) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_provider_failure_is_terminal_and_shared_with_callbacks( + server: RecordingServer, asynchronous: bool +) -> None: + recorder: Final = RecordingLogger() + server.enqueue(ResponseSpec(body=RATE_LIMITED, status=429)) + + with pytest.raises(RateLimitError) as caught: + await call(server, asynchronous=asynchronous, callbacks=[recorder]) + + assert caught.value.status_code == 429 + assert len(server.requests) == 1 + failure: Final = await recorder.wait_for_async(failure_signal(asynchronous)) + assert len(failure) == 1 + assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value + assert not any("success" in name for name in recorder.names) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "options", + ( + {"stream": True}, + {"extra_body": {"provider_option": True}}, + {"mock_response": "mock"}, + {"num_retries": 1}, + {"use_chat_completions_api": True}, + {"model_list": []}, + ), +) +async def test_unsupported_requests_fail_before_callbacks_and_provider_io( + server: RecordingServer, options: dict[str, object] +) -> None: + recorder: Final = RecordingLogger() + server.expected_requests = 0 + + with pytest.raises(RuntimeError, match="declined the request"): + await call(server, callbacks=[recorder], **options) + + assert not server.requests + assert not recorder.events diff --git a/tests/test_litellm_rust/responses/test_lifecycle.py b/tests/test_litellm_rust/responses/test_lifecycle.py new file mode 100644 index 00000000000..a8a4ce18f09 --- /dev/null +++ b/tests/test_litellm_rust/responses/test_lifecycle.py @@ -0,0 +1,46 @@ +import asyncio +from typing import Final + +import pytest + +import litellm +from litellm.rust_bridge import _native +from litellm.rust_bridge.public_call import NativeCall +from tests.test_litellm_rust.responses.data import REQUEST +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.recording_server import RecordingServer + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.mark.asyncio +async def test_unstarted_call_has_no_provider_or_callback_effects( + server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + recorder: Final = RecordingLogger() + server.expected_requests = 0 + monkeypatch.setattr(litellm, "ssl_verify", object()) + monkeypatch.setattr(litellm, "max_budget", 1.0) + monkeypatch.setattr(litellm, "_current_cost", 2.0) + kwargs: Final = { + **REQUEST, + "api_key": "test-key", + "api_base": server.base_url, + "callbacks": [recorder], + } + + base: Final = { + **REQUEST, + "stream": None, + "api_key": "test-key", + "api_base": server.base_url, + "custom_llm_provider": "openai", + "extra_headers": None, + } + + pending: Final = _native.aresponses(NativeCall(args=(), kwargs=kwargs, base=base)) + assert asyncio.iscoroutine(pending) + pending.close() + + assert not server.requests + assert not recorder.events diff --git a/tests/test_litellm_rust/responses/test_requests.py b/tests/test_litellm_rust/responses/test_requests.py new file mode 100644 index 00000000000..db04ccb5c65 --- /dev/null +++ b/tests/test_litellm_rust/responses/test_requests.py @@ -0,0 +1,56 @@ +import asyncio +from typing import Final + +import pytest +from pydantic import TypeAdapter + +import litellm +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.rust_bridge.catalog import Route +from litellm.types.llms.openai import ResponsesAPIResponse +from tests.test_litellm_rust.responses.data import RESPONSES_MODEL, RESPONSES_TEXT, call +from tests.test_litellm_rust.support.recording_server import RecordingServer +from tests.test_litellm_rust.support.routing import assert_served_natively, native_route + +pytestmark = pytest.mark.requires_rust_extension + +_OBJECT: Final = TypeAdapter(dict[str, object]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_positional_parameters_reach_the_wire(server: RecordingServer, asynchronous: bool) -> None: + arguments: Final = ("hello", RESPONSES_MODEL, None, "Be brief", 16) + supplied: Final = {"api_key": "test-key", "api_base": server.base_url} + + with native_route(Route.RESPONSES): + result: Final = ( + await litellm.aresponses(*arguments, **supplied) + if asynchronous + else await asyncio.to_thread(litellm.responses, *arguments, **supplied) + ) + + assert isinstance(result, ResponsesAPIResponse) + assert result.output_text == RESPONSES_TEXT + assert_served_natively(server) + body: Final = _OBJECT.validate_python(server.requests[0].body) + assert body["instructions"] == arguments[3] + assert body["max_output_tokens"] == arguments[4] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +@pytest.mark.parametrize("encoded", (False, True)) +async def test_continuation_ids_are_decoded_before_they_reach_the_wire( + server: RecordingServer, asynchronous: bool, encoded: bool +) -> None: + original: Final = "resp_upstream" + previous: Final = ( + ResponsesAPIRequestUtils._build_responses_api_response_id("openai", "deployment", original) + if encoded + else original + ) + + await call(server, asynchronous=asynchronous, previous_response_id=previous) + + assert _OBJECT.validate_python(server.requests[0].body)["previous_response_id"] == original diff --git a/tests/test_litellm_rust/responses/test_settings.py b/tests/test_litellm_rust/responses/test_settings.py new file mode 100644 index 00000000000..883a84b2a45 --- /dev/null +++ b/tests/test_litellm_rust/responses/test_settings.py @@ -0,0 +1,83 @@ +from typing import Final + +import pytest +from pydantic import TypeAdapter + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.models.credentials import CredentialItem +from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.utils import CallTypes +from tests.test_litellm_rust.responses.data import call +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.recording_server import RecordingServer +from tests.test_litellm_rust.support.routing import assert_served_natively + +pytestmark = pytest.mark.requires_rust_extension + +_OBJECT: Final = TypeAdapter(dict[str, object]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("from_credentials", (False, True)) +async def test_resource_setup_uses_deployment_hook_arguments( + server: RecordingServer, monkeypatch: pytest.MonkeyPatch, from_credentials: bool +) -> None: + invalid_settings: Final = {"ssl_verify": object()} + credential: Final = CredentialItem( + credential_name="resource-settings", credential_info={}, credential_values=invalid_settings + ) + monkeypatch.setattr(litellm, "credential_list", [credential]) + + class Prepare(CustomLogger): + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: CallTypes | None + ) -> dict[str, object]: + return { + **kwargs, + **({"litellm_credential_name": credential.credential_name} if from_credentials else invalid_settings), + } + + litellm.callbacks.append(Prepare()) + recorder: Final = RecordingLogger() + server.expected_requests = 0 + + with pytest.raises(ValueError, match=r"request\.ssl_verify") as caught: + await call(server, callbacks=[recorder]) + + failure: Final = await recorder.wait_for_async("async_log_failure_event") + assert len(failure) == 1 + assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value + assert not server.requests + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +@pytest.mark.parametrize("source", ("explicit", "base_url", "global", "provider", "environment", "empty")) +async def test_connection_settings_reach_the_provider( + server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool, source: str +) -> None: + key: Final = "selected-key" + explicit: Final = source in ("explicit", "base_url") + monkeypatch.setattr(litellm, "api_key", key if source in ("global", "empty") else ("unused" if explicit else None)) + monkeypatch.setattr(litellm, "openai_key", key if source == "provider" else ("unused" if explicit else None)) + monkeypatch.setattr( + litellm, + "api_base", + None if source == "environment" else ("http://127.0.0.1:1" if explicit else server.base_url), + ) + monkeypatch.setenv("OPENAI_API_KEY", key if source == "environment" else "unused") + for name in ("OPENAI_BASE_URL", "OPENAI_API_BASE"): + monkeypatch.setenv(name, server.base_url if source == "environment" else "http://127.0.0.1:1") + + result: Final = await call( + server, + asynchronous=asynchronous, + api_key=key if explicit else ("" if source == "empty" else None), + api_base=server.base_url if source == "explicit" else ("" if source == "empty" else None), + **({"base_url": server.base_url} if source == "base_url" else {}), + ) + + assert isinstance(result, ResponsesAPIResponse) + assert_served_natively(server) + assert server.requests[0].headers["authorization"] == f"Bearer {key}" diff --git a/tests/test_litellm_rust/responses/test_streaming.py b/tests/test_litellm_rust/responses/test_streaming.py new file mode 100644 index 00000000000..fe284bfa3fc --- /dev/null +++ b/tests/test_litellm_rust/responses/test_streaming.py @@ -0,0 +1,58 @@ +from collections.abc import AsyncIterator, Iterator +from typing import Final + +import pytest + +from tests.test_litellm_rust.responses.data import RESPONSES_EVENTS, RESPONSES_RESPONSE, STREAMING, call +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed +from tests.test_litellm_rust.support.results import consume, response_id +from tests.test_litellm_rust.support.routing import assert_served_natively + +pytestmark = pytest.mark.requires_rust_extension + +STREAM: Final = ResponseSpec(body=None, events=RESPONSES_EVENTS) + + +@pytest.fixture +def server(server: RecordingServer) -> Iterator[RecordingServer]: + yield from relaxed(server) + + +@pytest.mark.asyncio +@STREAMING +async def test_stream_relays_provider_events_and_logs_success_once_after_the_last_chunk( + server: RecordingServer, +) -> None: + server.enqueue(STREAM) + recorder: Final = RecordingLogger() + + stream: Final = await call(server, stream=True, callbacks=[recorder]) + assert isinstance(stream, AsyncIterator) + first: Final = await anext(stream) + await drain_logging() + assert "async_log_success_event" not in recorder.names + rest: Final = [chunk async for chunk in stream] + + assert_served_natively(server) + assert [getattr(chunk, "type", None) for chunk in (first, *rest)] == [event for event, _ in RESPONSES_EVENTS] + success: Final = await recorder.wait_for_async("async_log_success_event") + assert len(success) == 1 + assert isinstance(success[0].kwargs, dict) + assert success[0].kwargs["completion_start_time"] is not None + assert response_id(success[0].response) == RESPONSES_RESPONSE["id"] + assert "log_failure_event" not in recorder.names + + +@pytest.mark.asyncio +@STREAMING +async def test_sync_stream_relays_provider_events_and_logs_success_once(server: RecordingServer) -> None: + server.enqueue(STREAM) + recorder: Final = RecordingLogger() + + stream: Final = await call(server, asynchronous=False, stream=True, callbacks=[recorder]) + chunks: Final = await consume(stream) + + assert_served_natively(server) + assert [getattr(chunk, "type", None) for chunk in chunks] == [event for event, _ in RESPONSES_EVENTS] + assert len(await recorder.wait_for_async("log_success_event")) == 1 diff --git a/tests/test_litellm_rust/support/body_hooks.py b/tests/test_litellm_rust/support/body_hooks.py new file mode 100644 index 00000000000..1764b84da67 --- /dev/null +++ b/tests/test_litellm_rust/support/body_hooks.py @@ -0,0 +1,39 @@ +from litellm.types.integrations.custom_logger import AgenticLoopPlan, AgenticLoopRequestPatch +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.results import response_id, with_id + + +def agentic_hooks(recorder: RecordingLogger) -> tuple[str, ...]: + return tuple(name for name in recorder.names if "agentic" in name) + + +class PlannedAgenticLoop(RecordingLogger): + async def async_should_run_agentic_loop(self, response: object, **kwargs: object) -> tuple[bool, dict[str, object]]: + self._record("async_should_run_agentic_loop", response=response) + return response_id(response) != "followup", {} + + async def async_build_agentic_loop_plan(self, response: object, **kwargs: object) -> AgenticLoopPlan: + self._record("async_build_agentic_loop_plan", response=response) + return AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch.model_validate( + {"messages": [{"role": "user", "content": "followup"}]} + ), + ) + + async def async_post_agentic_loop_response_hook(self, response: object, **kwargs: object) -> object: + self._record("async_post_agentic_loop_response_hook", response=response) + return with_id(response, "agentic-replaced") + + async def async_agentic_loop_cleanup_hook(self, **kwargs: object) -> None: + self._record("async_agentic_loop_cleanup_hook") + + +class LegacyAgenticLoop(RecordingLogger): + async def async_should_run_agentic_loop(self, response: object, **kwargs: object) -> tuple[bool, dict[str, object]]: + self._record("async_should_run_agentic_loop", response=response) + return True, {} + + async def async_run_agentic_loop(self, response: object, **kwargs: object) -> object: + self._record("async_run_agentic_loop", response=response) + return with_id(response, "agentic-replaced") diff --git a/tests/test_litellm_rust/support/callback_recorder.py b/tests/test_litellm_rust/support/callback_recorder.py index d3749ccc095..c5e9fcd51a8 100644 --- a/tests/test_litellm_rust/support/callback_recorder.py +++ b/tests/test_litellm_rust/support/callback_recorder.py @@ -3,6 +3,7 @@ import copy import threading import time from dataclasses import dataclass +from datetime import datetime from typing import Final from litellm.integrations.custom_logger import CustomLogger @@ -99,3 +100,63 @@ class RecordingLogger(CustomLogger): def logging_hook(self, kwargs, result, call_type): self._record("logging_hook", kwargs, result) return kwargs, result + + +class HookRecorder(RecordingLogger): + def log_post_api_call( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self._record("log_post_api_call", kwargs, response_obj) + + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + self._record("async_logging_hook", kwargs, result) + return kwargs, result + + def redact_standard_logging_payload_from_model_call_details( + self, model_call_details: dict[str, object] + ) -> dict[str, object]: + self._record("redact_standard_logging_payload_from_model_call_details", model_call_details) + return model_call_details + + async def async_pre_call_deployment_hook(self, kwargs: dict[str, object], call_type: object) -> dict[str, object]: + self._record("async_pre_call_deployment_hook", kwargs) + return kwargs + + async def async_post_call_success_deployment_hook( + self, request_data: dict[str, object], response: object, call_type: object + ) -> object: + self._record("async_post_call_success_deployment_hook", request_data, response) + return response + + async def async_post_call_failure_deployment_hook( + self, + request_data: dict[str, object], + original_exception: Exception, + call_type: object, + *, + fallback_depth: int | None = None, + ) -> None: + self._record("async_post_call_failure_deployment_hook", request_data, original_exception) + + async def log_model_group_rate_limit_error( + self, exception: Exception, original_model_group: str | None, kwargs: dict[str, object] + ) -> None: + self._record("log_model_group_rate_limit_error", kwargs, exception) + + +def request_body(kwargs: dict[str, object]) -> dict[str, object]: + additional_args = kwargs["additional_args"] + assert isinstance(additional_args, dict) + body = additional_args["complete_input_dict"] + assert isinstance(body, dict) + return body + + +def request_headers(kwargs: dict[str, object]) -> dict[str, object]: + additional_args = kwargs["additional_args"] + assert isinstance(additional_args, dict) + headers = additional_args["headers"] + assert isinstance(headers, dict) + return headers diff --git a/tests/test_litellm_rust/support/gaps.py b/tests/test_litellm_rust/support/gaps.py new file mode 100644 index 00000000000..bd2cefec467 --- /dev/null +++ b/tests/test_litellm_rust/support/gaps.py @@ -0,0 +1,5 @@ +import pytest + + +def gap(reason: str, raises: type[Exception] = AssertionError) -> pytest.MarkDecorator: + return pytest.mark.xfail(reason=reason, raises=raises, strict=True) diff --git a/tests/test_litellm_rust/support/recording_server.py b/tests/test_litellm_rust/support/recording_server.py index 74ca2cda1c5..d7b82767b97 100644 --- a/tests/test_litellm_rust/support/recording_server.py +++ b/tests/test_litellm_rust/support/recording_server.py @@ -120,3 +120,9 @@ def recording_service() -> Iterator[RecordingServer]: if recording_server.expected_requests is not None: assert len(recording_server.requests) == recording_server.expected_requests assert recording_server.responses == [] + + +def relaxed(server: RecordingServer) -> Iterator[RecordingServer]: + server.expected_requests = None + yield server + server.responses.clear() diff --git a/tests/test_litellm_rust/support/requests.py b/tests/test_litellm_rust/support/requests.py deleted file mode 100644 index c9cf81b83ca..00000000000 --- a/tests/test_litellm_rust/support/requests.py +++ /dev/null @@ -1,97 +0,0 @@ -from typing import Final - -import litellm -from litellm.llms.base_llm.ocr.transformation import OCRResponse -from tests.test_litellm_rust.support.recording_server import RecordingServer - -OCR_DOCUMENT: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"} -OCR_MODEL: Final = "mistral/mistral-ocr-latest" -OCR_RESPONSE: Final = { - "pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}], - "model": "mistral-ocr-latest", - "usage_info": {"pages_processed": 1, "doc_size_bytes": 3}, -} - -MESSAGES_MODEL: Final = "anthropic/claude-sonnet-5" -MESSAGES: Final = ({"role": "user", "content": "Hello"},) -MESSAGES_RESPONSE: Final = { - "id": "msg_native", - "type": "message", - "role": "assistant", - "model": "claude-sonnet-5", - "content": [{"type": "text", "text": "Hello from native Messages"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 5, "output_tokens": 4}, -} -MESSAGES_EVENTS: Final = ( - ("message_start", {"type": "message_start", "message": {**MESSAGES_RESPONSE, "content": [], "stop_reason": None}}), - ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), - ( - "content_block_delta", - { - "type": "content_block_delta", - "index": 0, - "delta": {"type": "text_delta", "text": "Hello from native Messages"}, - }, - ), - ("content_block_stop", {"type": "content_block_stop", "index": 0}), - ( - "message_delta", - { - "type": "message_delta", - "delta": {"stop_reason": "end_turn", "stop_sequence": None}, - "usage": {"output_tokens": 4}, - }, - ), - ("message_stop", {"type": "message_stop"}), -) - - -def ocr_arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]: - return { - "model": OCR_MODEL, - "document": dict(OCR_DOCUMENT), - "api_key": "test-key", - "api_base": server.base_url, - **kwargs, - } - - -def call_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse: - response: Final = litellm.ocr(**ocr_arguments(server, **kwargs)) - if not isinstance(response, OCRResponse): - raise TypeError(f"Expected OCRResponse, got {type(response).__name__}") - return response - - -async def call_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse: - return await litellm.aocr(**ocr_arguments(server, **kwargs)) - - -def call_native_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse: - return call_ocr(server, **kwargs) - - -async def call_native_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse: - return await call_aocr(server, **kwargs) - - -async def call_native(server: RecordingServer, asynchronous: bool, **kwargs: object) -> OCRResponse: - return await call_native_aocr(server, **kwargs) if asynchronous else call_native_ocr(server, **kwargs) - - -def request_body(kwargs: dict[str, object]) -> dict[str, object]: - additional_args = kwargs["additional_args"] - assert isinstance(additional_args, dict) - body = additional_args["complete_input_dict"] - assert isinstance(body, dict) - return body - - -def request_headers(kwargs: dict[str, object]) -> dict[str, object]: - additional_args = kwargs["additional_args"] - assert isinstance(additional_args, dict) - headers = additional_args["headers"] - assert isinstance(headers, dict) - return headers diff --git a/tests/test_litellm_rust/support/results.py b/tests/test_litellm_rust/support/results.py new file mode 100644 index 00000000000..814cebe3efc --- /dev/null +++ b/tests/test_litellm_rust/support/results.py @@ -0,0 +1,26 @@ +from collections.abc import AsyncIterator, Iterator + +from pydantic import BaseModel + +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.llms.openai import ResponsesAPIResponse + + +def response_id(response: object) -> object: + if isinstance(response, ResponsesAPIResponse): + return ResponsesAPIRequestUtils.decode_responses_api_response_id(response.id).get("response_id") or response.id + return response.get("id") if isinstance(response, dict) else None + + +def with_id(response: object, new_id: str) -> object: + if isinstance(response, BaseModel): + return response.model_copy(update={"id": new_id}) + assert isinstance(response, dict) + return {**response, "id": new_id} + + +async def consume(stream: object) -> tuple[object, ...]: + if isinstance(stream, AsyncIterator): + return tuple([chunk async for chunk in stream]) + assert isinstance(stream, Iterator) + return tuple(stream) diff --git a/tests/test_litellm_rust/support/routing.py b/tests/test_litellm_rust/support/routing.py new file mode 100644 index 00000000000..56bb1276012 --- /dev/null +++ b/tests/test_litellm_rust/support/routing.py @@ -0,0 +1,48 @@ +import asyncio +import copy +from collections.abc import Awaitable, Callable, Generator, Mapping +from contextlib import contextmanager +from typing import Final + +from litellm.rust_bridge import catalog +from litellm.rust_bridge.catalog import Route, RouteRule +from litellm.rust_bridge.configuration import Rollout +from tests.test_litellm_rust.support.isolation import rebound +from tests.test_litellm_rust.support.recording_server import RecordingServer +from tests.test_litellm_rust.support.served import assert_served_natively + +__all__ = ["CALL_ID", "assert_served_natively", "call_natively", "native_route"] + +CALL_ID: Final = "callback-contract" + + +@contextmanager +def native_route(*routes: Route) -> Generator[None]: + rules: tuple[RouteRule, ...] = tuple(RouteRule(route, Rollout.RUST_REQUIRED) for route in routes) + with rebound(catalog, "RULES", (*rules, *catalog.RULES)): + yield + + +async def call_natively( + route: Route, + create: Callable[..., object], + acreate: Callable[..., Awaitable[object]], + request: Mapping[str, object], + server: RecordingServer, + *, + asynchronous: bool, + **kwargs: object, +) -> object: + arguments: Final = { + **copy.deepcopy(dict(request)), + "api_key": "test-key", + "api_base": server.base_url, + "litellm_call_id": CALL_ID, + **kwargs, + } + with native_route(route): + if asynchronous: + return await acreate(**arguments) + return await asyncio.to_thread(create, **arguments) + + diff --git a/tests/test_litellm_rust/support/served.py b/tests/test_litellm_rust/support/served.py index d3c568937f3..edf492ea7e6 100644 --- a/tests/test_litellm_rust/support/served.py +++ b/tests/test_litellm_rust/support/served.py @@ -4,8 +4,13 @@ from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_di from tests.test_litellm_rust.support.recording_server import RecordingServer -def assert_served_natively(server: RecordingServer, response: object) -> None: - assert len(server.requests) == 1 - additional_headers: Final = get_hidden_params_dict(response).get("additional_headers") - assert isinstance(additional_headers, dict) - assert additional_headers.get("x-litellm-rust") == "true" +def assert_served_natively(server: RecordingServer, *responses: object, count: int = 1) -> None: + assert len(server.requests) == (len(responses) or count) + if not responses: + assert not any( + request.headers.get("user-agent", "").startswith("python-httpx") for request in server.requests + ) + for response in responses: + additional_headers: Final = get_hidden_params_dict(response).get("additional_headers") + assert isinstance(additional_headers, dict) + assert additional_headers.get("x-litellm-rust") == "true" diff --git a/tests/test_litellm_rust/test_inference.py b/tests/test_litellm_rust/test_inference.py deleted file mode 100644 index 83a01e8ad4d..00000000000 --- a/tests/test_litellm_rust/test_inference.py +++ /dev/null @@ -1,403 +0,0 @@ -import asyncio -from collections.abc import Awaitable, Coroutine, Mapping -from typing import Final, Literal, TypeAlias - -import pytest -from pydantic import JsonValue, TypeAdapter - -import litellm -from litellm import RateLimitError -from litellm.chat_completions import dispatch as chat_dispatch -from litellm.integrations.custom_logger import CustomLogger -from litellm.models.credentials import CredentialItem -from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.rust_bridge import _native -from litellm.rust_bridge.public_call import NativeCall -from litellm.types.llms.openai import ResponsesAPIResponse -from litellm.types.utils import CallTypes, ModelResponse -from tests.test_litellm_rust.support.callback_recorder import RecordingLogger -from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec -from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE, request_body - -pytestmark = pytest.mark.requires_rust_extension -Route: TypeAlias = Literal["chat", "responses"] -_OBJECT: Final = TypeAdapter(dict[str, object]) -NativeResult: TypeAlias = ( - ModelResponse - | ResponsesAPIResponse - | Coroutine[object, object, ModelResponse] - | Coroutine[object, object, ResponsesAPIResponse] -) -RESPONSES_MODEL: Final = "openai/gpt-6-sol" -RESPONSES_RESPONSE: Final[dict[str, JsonValue]] = { - "id": "resp_native", - "object": "response", - "created_at": 1, - "model": RESPONSES_MODEL.removeprefix("openai/"), - "status": "completed", - "output": [ - { - "type": "message", - "id": "msg_native", - "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": "native response", "annotations": []}], - } - ], - "usage": {"input_tokens": 5, "output_tokens": 4, "total_tokens": 9}, -} - - -@pytest.fixture(params=("chat", "responses")) -def route(request: pytest.FixtureRequest) -> Route: - return TypeAdapter(Route).validate_python(request.param) - - -def native_call( - route: Route, asynchronous: bool, server: RecordingServer, options: Mapping[str, object] -) -> NativeResult: - server.default_response = ResponseSpec(body=MESSAGES_RESPONSE if route == "chat" else RESPONSES_RESPONSE) - if route == "chat": - kwargs: Final = { - "model": MESSAGES_MODEL, - "messages": list(MESSAGES), - "api_key": "test-key", - "api_base": server.base_url, - "max_tokens": 32, - **options, - } - request: Final = NativeCall(args=(), kwargs=kwargs, base={}) - return (_native.acompletion if asynchronous else _native.completion)(request) - response_kwargs: Final = { - "model": RESPONSES_MODEL, - "input": "hello", - "api_key": "test-key", - "api_base": server.base_url, - "max_output_tokens": 32, - **options, - } - response_request: Final = NativeCall( - args=(), - kwargs=response_kwargs, - base={ - "model": RESPONSES_MODEL, - "input": "hello", - "stream": None, - "api_key": "test-key", - "api_base": server.base_url, - "custom_llm_provider": "openai", - "extra_headers": None, - }, - ) - return (_native.aresponses if asynchronous else _native.responses)(response_request) - - -async def execute(route: Route, asynchronous: bool, server: RecordingServer, options: Mapping[str, object]) -> object: - if not asynchronous: - return await asyncio.to_thread(native_call, route, False, server, options) - result: Final = native_call(route, True, server, options) - assert isinstance(result, Awaitable) - return await result - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", (False, True)) -async def test_native_inference_returns_public_models_and_logs_once( - route: Route, - asynchronous: bool, - recording_server: RecordingServer, -) -> None: - recorder: Final = RecordingLogger() - result: Final = await execute(route, asynchronous, recording_server, {"callbacks": [recorder], "temperature": 0.25}) - assert len(recording_server.requests) == 1 - sent: Final = recording_server.requests[0] - body: Final = _OBJECT.validate_python(sent.body) - assert body["temperature"] == 0.25 - assert request_body(_OBJECT.validate_python(recorder.wait_for("log_pre_api_call")[0].kwargs)) == body - if route == "chat": - assert isinstance(result, ModelResponse) - assert result.choices[0].message.content == "Hello from native Messages" - assert sent.path == "/v1/messages" - else: - assert isinstance(result, ResponsesAPIResponse) - assert result.output_text == "native response" - assert sent.path == "/responses" - success: Final = await recorder.wait_for_async("async_log_success_event" if asynchronous else "log_success_event") - assert len(success) == 1 - if isinstance(result, ModelResponse): - assert success[0].response is result - else: - logged: Final = success[0].response - assert isinstance(logged, ResponsesAPIResponse) - assert isinstance(result, ResponsesAPIResponse) - assert logged.id == result.id - assert logged.output_text == result.output_text - - -@pytest.mark.asyncio -async def test_native_inference_pre_call_edits_reach_the_provider( - route: Route, recording_server: RecordingServer -) -> None: - class Edit(CustomLogger): - def log_pre_api_call(self, model: object, messages: object, kwargs: dict[str, object]) -> None: - request_body(kwargs)["temperature"] = 0.75 - - await execute(route, True, recording_server, {"callbacks": [Edit()]}) - assert _OBJECT.validate_python(recording_server.requests[0].body)["temperature"] == 0.75 - - -@pytest.mark.asyncio -async def test_native_inference_does_not_restore_keywords_deleted_by_a_pre_call_hook( - route: Route, recording_server: RecordingServer -) -> None: - class DropKeywords(CustomLogger): - async def async_pre_call_deployment_hook( - self, kwargs: dict[str, object], call_type: CallTypes | None - ) -> dict[str, object]: - return {name: value for name, value in kwargs.items() if name not in ("temperature", "extra_headers")} - - litellm.callbacks.append(DropKeywords()) - await execute( - route, True, recording_server, {"temperature": 0.25, "extra_headers": {"x-deleted": "remove this header"}} - ) - - assert len(recording_server.requests) == 1 - sent: Final = recording_server.requests[0] - assert "temperature" not in _OBJECT.validate_python(sent.body) - assert sent.headers.get("x-deleted") is None - - -@pytest.mark.asyncio -@pytest.mark.parametrize("from_credentials", (False, True)) -async def test_native_resource_setup_uses_deployment_hook_arguments( - route: Route, - recording_server: RecordingServer, - monkeypatch: pytest.MonkeyPatch, - from_credentials: bool, -) -> None: - invalid_settings: Final = {"ssl_verify": object()} - credential: Final = CredentialItem( - credential_name="resource-settings", credential_info={}, credential_values=invalid_settings - ) - monkeypatch.setattr(litellm, "credential_list", [credential]) - - class Prepare(CustomLogger): - async def async_pre_call_deployment_hook( - self, kwargs: dict[str, object], call_type: CallTypes | None - ) -> dict[str, object]: - return { - **kwargs, - **({"litellm_credential_name": credential.credential_name} if from_credentials else invalid_settings), - } - - litellm.callbacks.append(Prepare()) - recorder: Final = RecordingLogger() - recording_server.expected_requests = 0 - with pytest.raises(ValueError, match=r"request\.ssl_verify") as caught: - await execute(route, True, recording_server, {"callbacks": [recorder]}) - failure: Final = await recorder.wait_for_async("async_log_failure_event") - assert len(failure) == 1 - assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value - assert not recording_server.requests - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", (False, True)) -async def test_sdk_policy_rejection_precedes_resource_setup_and_is_logged_once( - route: Route, - asynchronous: bool, - recording_server: RecordingServer, - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setattr(litellm, "max_budget", 1.0) - monkeypatch.setattr(litellm, "_current_cost", 2.0) - monkeypatch.setattr(litellm, "ssl_verify", object()) - recorder: Final = RecordingLogger() - recording_server.expected_requests = 0 - with pytest.raises(litellm.BudgetExceededError) as caught: - await execute(route, asynchronous, recording_server, {"callbacks": [recorder]}) - failure: Final = await recorder.wait_for_async("async_log_failure_event" if asynchronous else "log_failure_event") - assert len(failure) == 1 - assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value - assert not recording_server.requests - assert not any("success" in name for name in recorder.names) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", (False, True)) -async def test_native_inference_provider_failure_is_terminal_and_shared_with_callbacks( - route: Route, - asynchronous: bool, - recording_server: RecordingServer, -) -> None: - recorder: Final = RecordingLogger() - recording_server.enqueue( - ResponseSpec(body={"error": {"message": "slow down", "type": "rate_limit_error"}}, status=429) - ) - with pytest.raises(RateLimitError) as caught: - await execute(route, asynchronous, recording_server, {"callbacks": [recorder]}) - assert getattr(caught.value, "status_code", None) == 429 - assert len(recording_server.requests) == 1 - failure: Final = await recorder.wait_for_async("async_log_failure_event" if asynchronous else "log_failure_event") - assert len(failure) == 1 - assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value - assert not any("success" in name for name in recorder.names) - - -@pytest.mark.asyncio -async def test_unstarted_native_inference_has_no_provider_or_callback_effects( - route: Route, - recording_server: RecordingServer, - monkeypatch: pytest.MonkeyPatch, -) -> None: - recorder: Final = RecordingLogger() - recording_server.expected_requests = 0 - monkeypatch.setattr(litellm, "ssl_verify", object()) - monkeypatch.setattr(litellm, "max_budget", 1.0) - monkeypatch.setattr(litellm, "_current_cost", 2.0) - pending: Final = native_call(route, True, recording_server, {"callbacks": [recorder]}) - assert asyncio.iscoroutine(pending) - pending.close() - assert not recording_server.requests - assert not recorder.events - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "options", - ( - {"stream": True}, - {"extra_body": {"provider_option": True}}, - {"mock_response": "mock"}, - {"num_retries": 1}, - {"use_chat_completions_api": True}, - {"model_list": []}, - ), -) -async def test_native_responses_declines_unsupported_requests_before_callbacks( - recording_server: RecordingServer, - options: Mapping[str, object], -) -> None: - recorder: Final = RecordingLogger() - recording_server.expected_requests = 0 - with pytest.raises(_native.RustBridgeDeclined): - native_call("responses", True, recording_server, {**options, "callbacks": [recorder]}) - assert not recording_server.requests - assert not recorder.events - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", (False, True)) -async def test_native_chat_validation_failure_is_terminal( - asynchronous: bool, recording_server: RecordingServer -) -> None: - recording_server.expected_requests = 0 - with pytest.raises(Exception, match="chat completions requires at least one message") as failure: - await execute("chat", asynchronous, recording_server, {"messages": []}) - assert not isinstance(failure.value, _native.RustBridgeDeclined) - assert not recording_server.requests - - -@pytest.mark.asyncio -async def test_native_projection_reads_positional_parameters(route: Route, recording_server: RecordingServer) -> None: - from litellm.chat_completions.dispatch import ( - _DISPATCH as chat_dispatch, # pyright: ignore[reportPrivateUsage] # exercise the request passed to the native boundary - ) - from litellm.responses.dispatch import ( - _DISPATCH as responses_dispatch, # pyright: ignore[reportPrivateUsage] # exercise the request passed to the native boundary - ) - - recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE if route == "chat" else RESPONSES_RESPONSE) - kwargs: Final = {"api_key": "test-key", "api_base": recording_server.base_url} - if route == "chat": - args: Final = (MESSAGES_MODEL, list(MESSAGES), 12.0, 0.25) - request: Final = chat_dispatch.request(args, kwargs) - assert request is not None - await asyncio.to_thread(_native.completion, request) - assert _OBJECT.validate_python(recording_server.requests[0].body)["temperature"] == 0.25 - else: - response_args: Final = ("hello", RESPONSES_MODEL, None, "Be brief", 16) - response_request: Final = responses_dispatch.request(response_args, kwargs) - assert response_request is not None - await asyncio.to_thread(_native.responses, response_request) - body: Final = _OBJECT.validate_python(recording_server.requests[0].body) - assert body["instructions"] == "Be brief" - assert body["max_output_tokens"] == 16 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", (False, True)) -@pytest.mark.parametrize("source", ("explicit", "base_url", "global", "provider", "environment", "empty")) -async def test_native_connection_settings_reach_the_provider( - route: Route, - asynchronous: bool, - source: str, - recording_server: RecordingServer, - monkeypatch: pytest.MonkeyPatch, -) -> None: - key: Final = "selected-key" - explicit: Final = source in ("explicit", "base_url") - monkeypatch.setattr(litellm, "api_key", key if source in ("global", "empty") else ("unused" if explicit else None)) - monkeypatch.setattr(litellm, "openai_key", key if source == "provider" else ("unused" if explicit else None)) - monkeypatch.setattr(litellm, "anthropic_key", key if source == "provider" else ("unused" if explicit else None)) - monkeypatch.setattr( - litellm, - "api_base", - None if source == "environment" else ("http://127.0.0.1:1" if explicit else recording_server.base_url), - ) - monkeypatch.setenv( - "OPENAI_API_KEY" if route == "responses" else "ANTHROPIC_API_KEY", key if source == "environment" else "unused" - ) - for name in ("OPENAI_BASE_URL", "OPENAI_API_BASE", "ANTHROPIC_BASE_URL", "ANTHROPIC_API_BASE"): - monkeypatch.setenv(name, recording_server.base_url if source == "environment" else "http://127.0.0.1:1") - result: Final = await execute( - route, - asynchronous, - recording_server, - { - "api_key": key if explicit else ("" if source == "empty" else None), - "api_base": recording_server.base_url if source == "explicit" else ("" if source == "empty" else None), - **({"base_url": recording_server.base_url} if source == "base_url" else {}), - }, - ) - assert isinstance(result, ModelResponse | ResponsesAPIResponse) - assert len(recording_server.requests) == 1 - headers: Final = recording_server.requests[0].headers - assert headers["x-api-key" if route == "chat" else "authorization"] == (key if route == "chat" else f"Bearer {key}") - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", (False, True)) -@pytest.mark.parametrize("encoded", (False, True)) -async def test_native_responses_decode_continuation_ids( - asynchronous: bool, encoded: bool, recording_server: RecordingServer -) -> None: - original: Final = "resp_upstream" - previous: Final = ( - ResponsesAPIRequestUtils._build_responses_api_response_id("openai", "deployment", original) - if encoded - else original - ) - await execute("responses", asynchronous, recording_server, {"previous_response_id": previous}) - assert _OBJECT.validate_python(recording_server.requests[0].body)["previous_response_id"] == original - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", (False, True)) -async def test_native_chat_uses_bound_positional_parameters( - asynchronous: bool, recording_server: RecordingServer -) -> None: - recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) - arguments: Final = (MESSAGES_MODEL, list(MESSAGES), 12.0, 0.35) - supplied: Final = {"base_url": recording_server.base_url, "api_key": "test-key", "max_tokens": 32} - call: Final = chat_dispatch._DISPATCH.request(arguments, supplied) # pyright: ignore[reportPrivateUsage] # exercise the native request produced by public binding - assert call is not None - result: Final = ( - await _native.acompletion(call) if asynchronous else await asyncio.to_thread(_native.completion, call) - ) - assert isinstance(result, ModelResponse) - assert result.choices[0].message.content == "Hello from native Messages" - assert len(recording_server.requests) == 1 - body: Final = _OBJECT.validate_python(recording_server.requests[0].body) - assert body["temperature"] == arguments[3] - assert body["max_tokens"] == supplied["max_tokens"] diff --git a/tests/unit/rust_bridge/messages/test_secrets.py b/tests/unit/rust_bridge/messages/test_secrets.py index 89fd1d1b771..64208182025 100644 --- a/tests/unit/rust_bridge/messages/test_secrets.py +++ b/tests/unit/rust_bridge/messages/test_secrets.py @@ -14,8 +14,8 @@ from litellm.rust_bridge.host import settings from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES, NATIVE_MESSAGES from litellm.rust_bridge.public_call import NativeCall from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem +from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE from tests.test_litellm_rust.support.recording_server import ResponseSpec, recording_service -from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE pytest.importorskip("litellm.rust_bridge._native") diff --git a/tests/unit/rust_bridge/ocr/test_secrets.py b/tests/unit/rust_bridge/ocr/test_secrets.py index a9c6072b357..33ded916888 100644 --- a/tests/unit/rust_bridge/ocr/test_secrets.py +++ b/tests/unit/rust_bridge/ocr/test_secrets.py @@ -16,8 +16,8 @@ from litellm.rust_bridge.host import settings from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR from litellm.rust_bridge.public_call import NativeCall from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem +from tests.test_litellm_rust.ocr.data import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, recording_service -from tests.test_litellm_rust.support.requests import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE native: Final = pytest.importorskip("litellm.rust_bridge._native")