diff --git a/Makefile b/Makefile index 1552187e906..a7a12a8635a 100644 --- a/Makefile +++ b/Makefile @@ -290,7 +290,7 @@ test-rust-extension: "$$temporary/venv/bin/python" -I -m mypy.stubtest \ --mypy-config-file tests/unit/rust_bridge/stubtest.ini \ litellm.rust_bridge._native && \ - LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \ + LITELLM_LOCAL_MODEL_COST_MAP=True \ "$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust rust-sqlx-prepare: diff --git a/tests/test_litellm_rust/AGENTS.md b/tests/test_litellm_rust/AGENTS.md index 7e8cd67165f..ea661bc36f1 100644 --- a/tests/test_litellm_rust/AGENTS.md +++ b/tests/test_litellm_rust/AGENTS.md @@ -1,8 +1,20 @@ -This directory holds only the tests that cannot be written in the Rust code +# Invariants + +- This directory tests the Rust path and nothing else. Treat the Python implementation of every route as already removed. Never call it on purpose, never run one request through both engines, and never compare Python and Rust outputs, headers, cache keys or logs +- A test that would still pass with Python serving the call does not belong here. Every test asserts the call was served natively with `support.routing.assert_served_natively` +- Select the Rust path with `support.routing.native_route(...)`, which prepends `RUST_REQUIRED` rules to `catalog.RULES` and restores them on exit. Each route folder's `conftest.py` applies it to every test through an autouse fixture. A test outside a route folder wraps its whole body in one `native_route(...)` block, and a child interpreter script assigns `catalog.RULES` once at its top. Never select a backend by reading or setting `LITELLM_RUST`, calling `litellm.rust(...)` or rebinding `catalog.RULES` any other way. The rollout in `litellm/rust_bridge/catalog.py` is a product decision, and a test here must pass no matter what it says, including for a route marked `PYTHON_ONLY` +- `catalog.RULES` is static at runtime and `native_route(...)` swaps a module global, so open it once per test and never per call. Two overlapping blocks, such as one per task inside an `asyncio.gather`, restore out of order and leave the wrong rules behind for later tests +- A test whose subject is the rollout switch itself (env, `litellm.rust(...)`, catalog rules) is a unit test of `litellm/rust_bridge/configuration.py` or `catalog.py` and lives in `tests/unit/rust_bridge/` +- A test lives here only when it cannot be written in Rust. That means it needs the Python public API, Python-owned state such as `litellm.cache` and callbacks, or the process boundary (fork, threads, event loops). Anything checkable from the Rust side goes in the crate's own tests +- The wire is the local recording server. No test here reaches a real provider +- Never parametrize one test across routes. Each route folder owns its tests, its `server` fixture, its request data and its `call` +- Call `_native` directly only for a case the public API cannot reach + +# Layout 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_dispatch.py` the call reaches the native route - `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 @@ -10,10 +22,6 @@ Each route folder (`ocr/`, `messages/`, `responses/`, `chat_completions/`) uses - `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 +- `test_cache.py` the configured cache seen from this route -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` +Route data and the route's `call` live in `/data.py`, and `/conftest.py` owns the `server` fixture and the autouse `native` fixture. `support/` holds only route-neutral helpers diff --git a/tests/test_litellm_rust/cache/test_python_cache.py b/tests/test_litellm_rust/cache/test_python_cache.py index 9a0a874ddbf..d855d212fec 100644 --- a/tests/test_litellm_rust/cache/test_python_cache.py +++ b/tests/test_litellm_rust/cache/test_python_cache.py @@ -23,11 +23,12 @@ from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.public_call import NativeCall from litellm.types.caching import CachingSupportedCallTypes +from tests.test_litellm_rust.cache.support import cache_key, collect, invoke, payload 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.routing import native_route pytestmark = pytest.mark.requires_rust_extension @@ -227,22 +228,23 @@ async def test_v2_messages_replays_a_completed_stream(recording_server: Recordin "stream": True, "callbacks": [recorder], } - first_stream: Final = await litellm.anthropic_messages(**parameters) - assert cache_key(first_stream) is None - first: Final = await collect(first_stream) - await recorder.wait_for_async("async_log_success_event") - second_stream: Final = await litellm.anthropic_messages(**parameters) - assert isinstance(cache_key(second_stream), str) - assert cache_key(second_stream) == get_hidden_params_dict(second_stream)["cache_key"] - second: Final = await collect(second_stream) - assert payload(first) == payload(second) - assert first == b"".join(recording_server.default_response.payloads()) - assert len(recording_server.requests) == 1 - await drain_logging() - successes: Final = await recorder.wait_for_async("async_log_success_event", count=2) - cached_log: Final = TypeAdapter(dict[str, object]).validate_python(successes[-1].kwargs) - assert cached_log["cache_hit"] is True - assert cached_log["response_cost"] == 0 + with native_route(Route.MESSAGES): + first_stream: Final = await litellm.anthropic_messages(**parameters) + assert cache_key(first_stream) is None + first: Final = await collect(first_stream) + await recorder.wait_for_async("async_log_success_event") + second_stream: Final = await litellm.anthropic_messages(**parameters) + assert isinstance(cache_key(second_stream), str) + assert cache_key(second_stream) == get_hidden_params_dict(second_stream)["cache_key"] + second: Final = await collect(second_stream) + assert payload(first) == payload(second) + assert first == b"".join(recording_server.default_response.payloads()) + assert len(recording_server.requests) == 1 + await drain_logging() + successes: Final = await recorder.wait_for_async("async_log_success_event", count=2) + cached_log: Final = TypeAdapter(dict[str, object]).validate_python(successes[-1].kwargs) + assert cached_log["cache_hit"] is True + assert cached_log["response_cost"] == 0 @pytest.mark.asyncio diff --git a/tests/test_litellm_rust/chat_completions/conftest.py b/tests/test_litellm_rust/chat_completions/conftest.py index 25c46c8ae42..04f3f64a91e 100644 --- a/tests/test_litellm_rust/chat_completions/conftest.py +++ b/tests/test_litellm_rust/chat_completions/conftest.py @@ -1,10 +1,20 @@ +from collections.abc import Generator + import pytest +from litellm.rust_bridge.catalog import Route from tests.test_litellm_rust.messages.data import MESSAGES_RESPONSE from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec +from tests.test_litellm_rust.support.routing import native_route @pytest.fixture def server(recording_server: RecordingServer) -> RecordingServer: recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) return recording_server + + +@pytest.fixture(autouse=True) +def native() -> Generator[None]: + with native_route(Route.CHAT_COMPLETIONS): + yield diff --git a/tests/test_litellm_rust/chat_completions/data.py b/tests/test_litellm_rust/chat_completions/data.py index 42f8e518f55..2b44cc76c33 100644 --- a/tests/test_litellm_rust/chat_completions/data.py +++ b/tests/test_litellm_rust/chat_completions/data.py @@ -1,10 +1,9 @@ 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 +from tests.test_litellm_rust.support.routing import call_public def _error(kind: str) -> dict[str, object]: @@ -17,8 +16,7 @@ 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, + return await call_public( litellm.completion, litellm.acompletion, REQUEST, diff --git a/tests/test_litellm_rust/chat_completions/test_requests.py b/tests/test_litellm_rust/chat_completions/test_requests.py index 3f2ba7c60e0..c8800501cd4 100644 --- a/tests/test_litellm_rust/chat_completions/test_requests.py +++ b/tests/test_litellm_rust/chat_completions/test_requests.py @@ -5,11 +5,10 @@ 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 +from tests.test_litellm_rust.support.routing import assert_served_natively pytestmark = pytest.mark.requires_rust_extension @@ -27,12 +26,11 @@ async def test_positional_parameters_reach_the_wire(server: RecordingServer, asy ) 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) - ) + 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 diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py index d9406af4c59..c1d452e838e 100644 --- a/tests/test_litellm_rust/conftest.py +++ b/tests/test_litellm_rust/conftest.py @@ -1,8 +1,8 @@ import asyncio -import os from collections.abc import AsyncIterator, Generator from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack +from importlib import import_module from typing import Final import pytest @@ -12,10 +12,7 @@ import litellm from litellm import utils from litellm.litellm_core_utils import litellm_logging, thread_pool_executor from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.rust_bridge.configuration import ( - _CONFIGURATION, - _parse_env_bool, -) +from litellm.rust_bridge.configuration import _CONFIGURATION from tests.test_litellm_rust.support.callback_recorder import drain_logging from tests.test_litellm_rust.support.clickhouse import clickhouse_url as clickhouse_url from tests.test_litellm_rust.support.isolation import isolated_callback_registries, rebound @@ -49,17 +46,13 @@ def recording_server() -> Generator[RecordingServer]: def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: - if not _parse_env_bool(os.environ.get("LITELLM_RUST")): - skip: Final = pytest.mark.skip(reason="requires LITELLM_RUST=1 and a compiled Rust extension") + try: + import_module("litellm.rust_bridge._native") + except ImportError: + skip: Final = pytest.mark.skip(reason="requires a compiled litellm.rust_bridge._native extension") for item in items: if "test_litellm_rust" in item.path.parts: item.add_marker(skip) - return - - try: - from litellm.rust_bridge import _native # noqa: F401 # validates the installed extension - except ImportError as error: - raise pytest.UsageError("LITELLM_RUST=1 requires a compiled litellm.rust_bridge._native extension") from error @pytest.fixture diff --git a/tests/test_litellm_rust/messages/conftest.py b/tests/test_litellm_rust/messages/conftest.py index 25c46c8ae42..18a303dcd69 100644 --- a/tests/test_litellm_rust/messages/conftest.py +++ b/tests/test_litellm_rust/messages/conftest.py @@ -1,10 +1,20 @@ +from collections.abc import Generator + import pytest +from litellm.rust_bridge.catalog import Route from tests.test_litellm_rust.messages.data import MESSAGES_RESPONSE from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec +from tests.test_litellm_rust.support.routing import native_route @pytest.fixture def server(recording_server: RecordingServer) -> RecordingServer: recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) return recording_server + + +@pytest.fixture(autouse=True) +def native() -> Generator[None]: + with native_route(Route.MESSAGES): + yield diff --git a/tests/test_litellm_rust/messages/data.py b/tests/test_litellm_rust/messages/data.py index b7d039eb352..46127c9237a 100644 --- a/tests/test_litellm_rust/messages/data.py +++ b/tests/test_litellm_rust/messages/data.py @@ -1,10 +1,9 @@ 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 +from tests.test_litellm_rust.support.routing import call_public MESSAGES_MODEL: Final = "anthropic/claude-sonnet-5" MESSAGES: Final = ({"role": "user", "content": "Hello"},) @@ -50,8 +49,7 @@ RATE_LIMITED: Final = _error("rate_limit_error") async def call(server: RecordingServer, *, asynchronous: bool = True, **kwargs: object) -> object: - return await call_natively( - Route.MESSAGES, + return await call_public( litellm.anthropic.messages.create, litellm.anthropic.messages.acreate, REQUEST, diff --git a/tests/test_litellm_rust/messages/test_requests.py b/tests/test_litellm_rust/messages/test_requests.py index 6388c0eb7db..92031e72aee 100644 --- a/tests/test_litellm_rust/messages/test_requests.py +++ b/tests/test_litellm_rust/messages/test_requests.py @@ -3,10 +3,9 @@ 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 +from tests.test_litellm_rust.support.routing import assert_served_natively pytestmark = pytest.mark.requires_rust_extension @@ -141,8 +140,7 @@ async def test_optional_positional_body_parameters_reach_the_wire(server: Record 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) + 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"] diff --git a/tests/test_litellm_rust/ocr/conftest.py b/tests/test_litellm_rust/ocr/conftest.py index 6d3834516d7..b85741a5250 100644 --- a/tests/test_litellm_rust/ocr/conftest.py +++ b/tests/test_litellm_rust/ocr/conftest.py @@ -1,9 +1,11 @@ -from typing import Final +from collections.abc import Generator import pytest +from litellm.rust_bridge.catalog import Route from tests.test_litellm_rust.ocr.data import OCR_RESPONSE from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec +from tests.test_litellm_rust.support.routing import native_route @pytest.fixture @@ -12,8 +14,7 @@ def server(recording_server: RecordingServer) -> RecordingServer: 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 +@pytest.fixture(autouse=True) +def native() -> Generator[None]: + with native_route(Route.OCR): + yield diff --git a/tests/test_litellm_rust/ocr/test_dispatch.py b/tests/test_litellm_rust/ocr/test_dispatch.py index bb75782fb67..db60b6aea84 100644 --- a/tests/test_litellm_rust/ocr/test_dispatch.py +++ b/tests/test_litellm_rust/ocr/test_dispatch.py @@ -3,7 +3,6 @@ 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.routing import assert_served_natively @@ -11,21 +10,6 @@ from tests.test_litellm_rust.support.routing import assert_served_natively pytestmark = pytest.mark.requires_rust_extension -@pytest.mark.parametrize("enabled", [False, True, 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=server.base_url, - ) - - assert isinstance(response, OCRResponse) - assert response.pages[0].markdown == "native OCR response" - assert_served_natively(server, response) - - @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.parametrize("caching", [None, False, True]) @@ -47,17 +31,15 @@ async def test_ocr_does_not_depend_on_chat_cache( 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 + server: RecordingServer, monkeypatch: pytest.MonkeyPatch, 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.catalog import Route, RouteContext from litellm.rust_bridge.host import settings from litellm.rust_bridge.ocr.entrypoints import NATIVE_OCR from litellm.rust_bridge.public_call import NativeCall @@ -93,7 +75,6 @@ def test_native_projection_errors_never_select_python( 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 == [] diff --git a/tests/test_litellm_rust/ocr/test_failures.py b/tests/test_litellm_rust/ocr/test_failures.py index 41373fb6e77..b5fbfedf404 100644 --- a/tests/test_litellm_rust/ocr/test_failures.py +++ b/tests/test_litellm_rust/ocr/test_failures.py @@ -26,7 +26,6 @@ pytestmark = pytest.mark.requires_rust_extension @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) @@ -48,7 +47,6 @@ async def test_ocr_contract_upstream_status( @pytest.mark.parametrize("preserved", ["body", "headers"]) async def test_ocr_contract_provider_error_details( server: RecordingServer, - ocr_backend: bool, asynchronous: bool, preserved: str, ) -> None: @@ -71,7 +69,6 @@ async def test_ocr_contract_provider_error_details( @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 @@ -95,7 +92,6 @@ async def test_ocr_contract_invalid_response_format( ) async def test_ocr_contract_malformed_document_is_actionable( server: RecordingServer, - ocr_backend: bool, asynchronous: bool, document: JsonValue, field: str, @@ -115,7 +111,6 @@ async def test_ocr_contract_malformed_document_is_actionable( @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, diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py index c7b7a6d4059..a5733b9b966 100644 --- a/tests/test_litellm_rust/ocr/test_requests.py +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -20,7 +20,6 @@ pytestmark = pytest.mark.requires_rust_extension @pytest.mark.parametrize("model", ["mistral/mistral-ocr-latest", "azure_ai/mistral-ocr-latest", "reducto/parse-v3"]) async def test_ocr_contract_native_format_supported( server: RecordingServer, - ocr_backend: bool, asynchronous: bool, model: str, ) -> None: @@ -43,15 +42,13 @@ async def test_ocr_contract_native_format_supported( assert response.pages[0].markdown == "native OCR response" assert response.get_provider_native_response() == payload assert len(server.requests) == 1 - if ocr_backend: - assert_served_natively(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( server: RecordingServer, - ocr_backend: bool, asynchronous: bool, ) -> None: server.default_response = ResponseSpec(body={"result": {"chunks": [{"content": "future model response"}]}}) diff --git a/tests/test_litellm_rust/responses/conftest.py b/tests/test_litellm_rust/responses/conftest.py index c2ea94473dd..73b398a4329 100644 --- a/tests/test_litellm_rust/responses/conftest.py +++ b/tests/test_litellm_rust/responses/conftest.py @@ -1,10 +1,20 @@ +from collections.abc import Generator + import pytest +from litellm.rust_bridge.catalog import Route from tests.test_litellm_rust.responses.data import RESPONSES_RESPONSE from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec +from tests.test_litellm_rust.support.routing import native_route @pytest.fixture def server(recording_server: RecordingServer) -> RecordingServer: recording_server.default_response = ResponseSpec(body=RESPONSES_RESPONSE) return recording_server + + +@pytest.fixture(autouse=True) +def native() -> Generator[None]: + with native_route(Route.RESPONSES): + yield diff --git a/tests/test_litellm_rust/responses/data.py b/tests/test_litellm_rust/responses/data.py index ed489f949a5..604a7de7442 100644 --- a/tests/test_litellm_rust/responses/data.py +++ b/tests/test_litellm_rust/responses/data.py @@ -1,10 +1,9 @@ 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 +from tests.test_litellm_rust.support.routing import call_public RESPONSES_MODEL: Final = "openai/gpt-6.1-sol" RESPONSES_TEXT: Final = "Hello from native Responses" @@ -59,8 +58,7 @@ RATE_LIMITED: Final = _error("rate_limit_exceeded") async def call(server: RecordingServer, *, asynchronous: bool = True, **kwargs: object) -> object: - return await call_natively( - Route.RESPONSES, + return await call_public( litellm.responses, litellm.aresponses, REQUEST, diff --git a/tests/test_litellm_rust/responses/test_requests.py b/tests/test_litellm_rust/responses/test_requests.py index db04ccb5c65..8a227408af6 100644 --- a/tests/test_litellm_rust/responses/test_requests.py +++ b/tests/test_litellm_rust/responses/test_requests.py @@ -6,11 +6,10 @@ 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 +from tests.test_litellm_rust.support.routing import assert_served_natively pytestmark = pytest.mark.requires_rust_extension @@ -23,12 +22,11 @@ async def test_positional_parameters_reach_the_wire(server: RecordingServer, asy 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) - ) + 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 diff --git a/tests/test_litellm_rust/support/routing.py b/tests/test_litellm_rust/support/routing.py index 56bb1276012..4250589848d 100644 --- a/tests/test_litellm_rust/support/routing.py +++ b/tests/test_litellm_rust/support/routing.py @@ -11,20 +11,18 @@ 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"] +__all__ = ["CALL_ID", "assert_served_natively", "call_public", "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)): + with rebound(catalog, "RULES", (*(RouteRule(route, Rollout.RUST_REQUIRED) for route in routes), *catalog.RULES)): yield -async def call_natively( - route: Route, +async def call_public( create: Callable[..., object], acreate: Callable[..., Awaitable[object]], request: Mapping[str, object], @@ -40,9 +38,6 @@ async def call_natively( "litellm_call_id": CALL_ID, **kwargs, } - with native_route(route): - if asynchronous: - return await acreate(**arguments) - return await asyncio.to_thread(create, **arguments) - - + if asynchronous: + return await acreate(**arguments) + return await asyncio.to_thread(create, **arguments) diff --git a/tests/test_litellm_rust/test_fork_guard.py b/tests/test_litellm_rust/test_fork_guard.py index 8e1251de8c2..7ced3f8bdff 100644 --- a/tests/test_litellm_rust/test_fork_guard.py +++ b/tests/test_litellm_rust/test_fork_guard.py @@ -94,8 +94,6 @@ _SDK_CONTRACT = textwrap.dedent( "api_base": f"http://127.0.0.1:{server.server_port}", "num_retries": 0, } - litellm.rust(True) - SERVED, REFUSED, OTHER = 0, 3, 4 def outcome(asynchronous): @@ -138,7 +136,6 @@ def test_sdk_call_in_a_child_forked_after_native_use_raises_instead_of_hanging() env = { **os.environ, "OBJC_DISABLE_INITIALIZE_FORK_SAFETY": "YES", - "LITELLM_RUST": "1", "LITELLM_LOCAL_MODEL_COST_MAP": "True", } @@ -162,8 +159,8 @@ from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer from litellm.utils import claude_json_str catalog.RULES = ( - RouteRule(Route.TOKENIZER, Rollout.RUST_OPT_IN), - RouteRule(Route.TOKEN_COUNTER, Rollout.RUST_OPT_IN), + RouteRule(Route.TOKENIZER, Rollout.RUST_REQUIRED), + RouteRule(Route.TOKEN_COUNTER, Rollout.RUST_REQUIRED), *catalog.RULES, ) litellm.anthropic_models = {*litellm.anthropic_models, "tokenizer-fork-fixture"} @@ -199,7 +196,9 @@ if pid == 0: lambda: _native.TokenCounter.from_tokenizer(encoding), lambda: _native.TokenCounter.from_tokenizer(encoding, fast=True), lambda: _native.Tokenizer.from_tiktoken("cl100k_base"), - lambda: asyncio.run(count_input_tokens({"prompt": "hello"}, b'{"prompt": "hello"}', ("counter-fork-fixture",))), + lambda: asyncio.run( + count_input_tokens({"prompt": "hello"}, b'{"prompt": "hello"}', ("counter-fork-fixture",)) + ), ): try: call() diff --git a/tests/unit/rust_bridge/test_configuration.py b/tests/unit/rust_bridge/test_configuration.py index 38fdfd0f476..8116cf98076 100644 --- a/tests/unit/rust_bridge/test_configuration.py +++ b/tests/unit/rust_bridge/test_configuration.py @@ -30,7 +30,15 @@ Decision: Final = configuration.Decision ("rollout", "process", "environment", "expected"), ( (Rollout.PYTHON_ONLY, True, True, Decision.PYTHON), + (Rollout.RUST_REQUIRED, None, None, Decision.RUST_REQUIRED), + (Rollout.RUST_REQUIRED, None, False, Decision.RUST_REQUIRED), + (Rollout.RUST_REQUIRED, None, True, Decision.RUST_REQUIRED), + (Rollout.RUST_REQUIRED, False, None, Decision.RUST_REQUIRED), (Rollout.RUST_REQUIRED, False, False, Decision.RUST_REQUIRED), + (Rollout.RUST_REQUIRED, False, True, Decision.RUST_REQUIRED), + (Rollout.RUST_REQUIRED, True, None, Decision.RUST_REQUIRED), + (Rollout.RUST_REQUIRED, True, False, Decision.RUST_REQUIRED), + (Rollout.RUST_REQUIRED, True, True, Decision.RUST_REQUIRED), (Rollout.RUST_OPT_IN, None, None, Decision.PYTHON), (Rollout.RUST_OPT_IN, None, True, Decision.RUST_WITH_FALLBACK), (Rollout.RUST_OPT_IN, True, None, Decision.RUST_WITH_FALLBACK),