diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index c6901411167..1e78ef0a78e 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -4,6 +4,7 @@ on: push: paths: - "litellm-rust/**" + - "tests/test_litellm_rust/**" - ".cargo/**" - "pyproject.toml" - "rust-toolchain.toml" @@ -20,6 +21,7 @@ on: - "litellm_**" paths: - "litellm-rust/**" + - "tests/test_litellm_rust/**" - ".cargo/**" - "pyproject.toml" - "rust-toolchain.toml" diff --git a/tests/test_litellm_rust/README.md b/tests/test_litellm_rust/README.md new file mode 100644 index 00000000000..772d6e82e00 --- /dev/null +++ b/tests/test_litellm_rust/README.md @@ -0,0 +1,11 @@ +# Rust bridge tests + +Put API-specific tests under `ocr/`, `messages/`, or `chat/`. Keep request construction, recording servers, callback recorders, and state isolation in `support/`. Integration tests and their fixtures belong in `integrations/`. + +Use `backend` with `"python"` and `"rust"` for parity tests. Keep callback mutation and lifecycle assertions next to the API surface that owns them. Shared route parametrization belongs in `integrations/routes.py` only when its observable contract is the same for every route. + +Run the suite with `LITELLM_RUST=1 uv run pytest tests/test_litellm_rust`. + +The newly ported suite is marked as a non-strict expected failure until the retained callback implementation from #40070 lands. Passing cases appear as XPASS so staging coverage remains visible + +`isolated_backend` restores the previous backend override and callback state on exit, including after exceptions or cancellation. Scopes can nest in the same task. Run backend comparisons sequentially: overlapping scopes in different tasks raise before changing process-global state. Use separate worker processes for parallel backend comparisons diff --git a/tests/test_litellm_rust/chat/__init__.py b/tests/test_litellm_rust/chat/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/test_litellm_rust/chat/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/test_litellm_rust/chat/test_callback_mutation.py b/tests/test_litellm_rust/chat/test_callback_mutation.py new file mode 100644 index 00000000000..a51fbd7cfb3 --- /dev/null +++ b/tests/test_litellm_rust/chat/test_callback_mutation.py @@ -0,0 +1,404 @@ +import hashlib +import hmac +import json +from typing import Final +from urllib.parse import urlsplit + +import pytest + +import litellm +from tests.test_litellm_rust.support.recording_server import RecordedRequest, RecordingServer, ResponseSpec + +pytestmark = pytest.mark.requires_rust_extension + + +def _verify_sigv4(request: RecordedRequest, secret_key: str) -> None: + authorization: Final = request.headers["authorization"] + algorithm, attributes_text = authorization.split(" ", 1) + attributes: Final = dict(item.split("=", 1) for item in attributes_text.split(", ")) + credential_scope: Final = attributes["Credential"].split("/", 1)[1] + signed_names: Final = attributes["SignedHeaders"].split(";") + canonical_headers: Final = "".join(f"{name}:{' '.join(request.headers[name].split())}\n" for name in signed_names) + parsed_path: Final = urlsplit(request.path) + canonical_request: Final = "\n".join( + ( + request.method, + parsed_path.path, + parsed_path.query, + canonical_headers, + attributes["SignedHeaders"], + hashlib.sha256(request.raw_body).hexdigest(), + ) + ) + amz_date: Final = request.headers["x-amz-date"] + string_to_sign: Final = "\n".join( + (algorithm, amz_date, credential_scope, hashlib.sha256(canonical_request.encode()).hexdigest()) + ) + date, region, service, terminator = credential_scope.split("/") + date_key: Final = hmac.new(f"AWS4{secret_key}".encode(), date.encode(), hashlib.sha256).digest() + region_key: Final = hmac.new(date_key, region.encode(), hashlib.sha256).digest() + service_key: Final = hmac.new(region_key, service.encode(), hashlib.sha256).digest() + signing_key: Final = hmac.new(service_key, terminator.encode(), hashlib.sha256).digest() + expected: Final = hmac.new(signing_key, string_to_sign.encode(), hashlib.sha256).hexdigest() + assert hmac.compare_digest(attributes["Signature"], expected) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("provider", ["anthropic", "bedrock"]) +@pytest.mark.parametrize("rebind_logging_view", [False, True]) +@pytest.mark.parametrize("status", [200, 429]) +@pytest.mark.parametrize("native", [False, True]) +async def test_chat_retains_callback_edits_through_public_dispatch( + recording_server: RecordingServer, + asynchronous: bool, + provider: str, + rebind_logging_view: bool, + status: int, + native: bool, +) -> None: + import threading + + litellm.rust(native) + + from tests.test_litellm_rust.support.callback_recorder import RecordingLogger + from tests.test_litellm_rust.support.requests import MESSAGES_RESPONSE + + recording_server.default_response = ResponseSpec( + status=status, + body=( + MESSAGES_RESPONSE + if provider == "anthropic" + else { + "output": {"message": {"role": "assistant", "content": [{"text": "native chat"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 5, "outputTokens": 4, "totalTokens": 9}, + "metrics": {"latencyMs": 1}, + } + ), + ) + caller_thread: Final = threading.current_thread() + observations: Final = [] + + class EditingLogger(RecordingLogger): + def log_pre_api_call(self, model, messages, kwargs): + super().log_pre_api_call(model, messages, kwargs) + body = kwargs["additional_args"]["complete_input_dict"] + headers = kwargs["additional_args"]["headers"] + if provider == "anthropic": + assert isinstance(body, dict) + body["messages"][0]["content"][0]["text"] = "edited by callback" + body["max_tokens"] = 32 + else: + assert isinstance(body, str) + assert json.loads(body)["messages"][0]["content"][0]["text"] == "original" + headers["x-retained-callback"] = "original" + observations.append((threading.current_thread(), body, headers)) + if rebind_logging_view: + kwargs["additional_args"]["complete_input_dict"] = {"replacement": True} + kwargs["additional_args"]["headers"] = {"x-retained-callback": "replacement"} + + recorder: Final = EditingLogger() + kwargs: Final = { + "model": "anthropic/claude-opus-5" if provider == "anthropic" else "bedrock/anthropic.claude-opus-5", + "messages": [{"role": "user", "content": "original"}], + "max_tokens": 64, + "api_key": "test-key", + "api_base": recording_server.base_url, + "callbacks": [recorder], + "num_retries": 0, + **( + {"aws_access_key_id": "test", "aws_secret_access_key": "test", "aws_region_name": "us-east-1"} + if provider == "bedrock" + else {} + ), + } + if status != 200: + with pytest.raises((litellm.APIError, litellm.RateLimitError)) as raised: + await litellm.acompletion(**kwargs) if asynchronous else litellm.completion(**kwargs) + failure_event: Final = "async_log_failure_event" if asynchronous else "log_failure_event" + failures: Final = await recorder.wait_for_async(failure_event) + assert raised.value.status_code == status + assert failures[0].kwargs["exception"] is raised.value + assert recorder.names.count(failure_event) == 1 + assert recorder.names.count("log_pre_api_call") == 1 + assert len(recording_server.requests) == 1 + return + + response: Final = await litellm.acompletion(**kwargs) if asynchronous else litellm.completion(**kwargs) + event_name: Final = "async_log_success_event" if asynchronous else "log_success_event" + events: Final = await recorder.wait_for_async(event_name) + + if native: + assert response._hidden_params["additional_headers"]["x-litellm-rust"] == "true" + assert recorder.names.count("log_pre_api_call") == 1 + assert recorder.names.count(event_name) == 1 + if asynchronous and provider == "anthropic" and not native: + assert observations[0][0] is not caller_thread + else: + assert observations[0][0] is caller_thread + assert events[0].response is response + assert recording_server.requests[0].body["messages"][0]["content"][0]["text"] == ( + "edited by callback" if provider == "anthropic" else "original" + ) + assert recording_server.requests[0].headers["x-retained-callback"] == "original" + body: Final = recording_server.requests[0].body + assert (body["max_tokens"] if provider == "anthropic" else body["inferenceConfig"]["maxTokens"]) == ( + 32 if provider == "anthropic" else 64 + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("status", [200, 429]) +@pytest.mark.parametrize("native", [False, True]) +async def test_bedrock_callbacks_share_state_without_replacing_signed_transport( + recording_server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + asynchronous: bool, + status: int, + native: bool, +) -> None: + from tests.test_litellm_rust.support.callback_recorder import RecordingLogger + + litellm.rust(native) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + recording_server.default_response = ResponseSpec( + status=status, + body={ + "output": {"message": {"role": "assistant", "content": [{"text": "native chat"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 5, "outputTokens": 4, "totalTokens": 9}, + "metrics": {"latencyMs": 1}, + }, + ) + observations: Final[list[tuple[int, int]]] = [] + replacement_body: Final = {"replacement": True} + rebound_headers: Final = {"x-callback-header": "rebound"} + + class FirstCallback(RecordingLogger): + def log_pre_api_call(self, model, messages, kwargs): + additional: Final = kwargs["additional_args"] + callback_body: Final = additional["complete_input_dict"] + original_headers: Final = additional["headers"] + observations.append((id(kwargs), id(additional))) + assert isinstance(callback_body, str) + assert json.loads(callback_body)["messages"][0]["content"][0]["text"] == "original" + kwargs["callback_marker"] = "visible" + original_headers["x-callback-header"] = "in-place" + additional["complete_input_dict"] = replacement_body + additional["headers"] = rebound_headers + raise RuntimeError("first callback failure is non-blocking") + + class SecondCallback(RecordingLogger): + def log_pre_api_call(self, model, messages, kwargs): + additional: Final = kwargs["additional_args"] + observations.append((id(kwargs), id(additional))) + assert kwargs["callback_marker"] == "visible" + assert additional["complete_input_dict"] is replacement_body + assert additional["headers"] is rebound_headers + + request: Final = { + "model": "bedrock/anthropic.claude-opus-5", + "messages": [{"role": "user", "content": "original"}], + "max_tokens": 64, + "api_base": recording_server.base_url, + "callbacks": [FirstCallback(), SecondCallback()], + "num_retries": 0, + "aws_access_key_id": "test-access", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-east-1", + } + if status == 200: + await litellm.acompletion(**request) if asynchronous else litellm.completion(**request) + else: + with pytest.raises((litellm.APIError, litellm.RateLimitError)): + await litellm.acompletion(**request) if asynchronous else litellm.completion(**request) + + assert observations[0] == observations[1] + recorded: Final = recording_server.requests[0] + assert recorded.body["messages"][0]["content"][0]["text"] == "original" + assert "replacement" not in recorded.body + assert recorded.headers["x-callback-header"] == "in-place" + assert recorded.headers["authorization"].startswith("AWS4-HMAC-SHA256 ") + _verify_sigv4(recorded, "test-secret") + + +@pytest.mark.asyncio +async def test_native_chat_pre_call_preserves_the_caller_task_and_context( + recording_server: RecordingServer, +) -> None: + import asyncio + from contextvars import ContextVar + + from litellm.integrations.custom_logger import CustomLogger + from tests.test_litellm_rust.support.requests import MESSAGES_RESPONSE + + recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) + marker: Final[ContextVar[str]] = ContextVar("native_chat_marker", default="missing") + marker.set("caller") + caller_task: Final = asyncio.current_task() + observations: Final = [] + + class Observe(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + observations.append((asyncio.current_task(), marker.get())) + marker.set("callback") + + response: Final = await litellm.acompletion( + model="anthropic/claude-opus-5", + messages=[{"role": "user", "content": "task context"}], + max_tokens=16, + api_key="test-key", + api_base=recording_server.base_url, + callbacks=[Observe()], + num_retries=0, + ) + + assert response._hidden_params["additional_headers"]["x-litellm-rust"] == "true" + assert observations == [(caller_task, "caller")] + assert marker.get() == "callback" + + +@pytest.mark.asyncio +async def test_native_chat_concurrent_calls_keep_callback_roots_isolated( + recording_server: RecordingServer, +) -> None: + import asyncio + import threading + + from litellm.integrations.custom_logger import CustomLogger + from tests.test_litellm_rust.support.callback_recorder import drain_logging + from tests.test_litellm_rust.support.requests import MESSAGES_RESPONSE + + call_ids: Final = tuple(f"chat-{index}" for index in range(4)) + recording_server.expected_requests = len(call_ids) + accepted: Final = tuple(threading.Event() for _ in call_ids) + release: Final = threading.Event() + for signal in accepted: + recording_server.enqueue(ResponseSpec(body=MESSAGES_RESPONSE, accepted=signal, release_before_response=release)) + tokens: Final = {call_id: object() for call_id in call_ids} + terminal_state: Final = [] + + class Correlate(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + body: Final = kwargs["additional_args"]["complete_input_dict"] + body["messages"][0]["content"][0]["text"] = f"callback-{kwargs['litellm_call_id']}" + kwargs["correlation-token"] = tokens[kwargs["litellm_call_id"]] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + terminal_state.append((kwargs["litellm_call_id"], kwargs["correlation-token"])) + + logger: Final = Correlate() + + async def invoke(call_id: str) -> None: + await litellm.acompletion( + model="anthropic/claude-opus-5", + messages=[{"role": "user", "content": f"original-{call_id}"}], + max_tokens=16, + api_key="test-key", + api_base=recording_server.base_url, + callbacks=[logger], + litellm_call_id=call_id, + num_retries=0, + ) + + tasks: Final = tuple(asyncio.create_task(invoke(call_id)) for call_id in call_ids) + try: + assert all(await asyncio.gather(*(asyncio.to_thread(signal.wait, 3) for signal in accepted))) + finally: + release.set() + await asyncio.gather(*tasks) + await drain_logging() + + sent: Final = {request.body["messages"][0]["content"][0]["text"] for request in recording_server.requests} + assert sent == {f"callback-{call_id}" for call_id in call_ids} + assert len(terminal_state) == len(call_ids) + assert all(token is tokens[call_id] for call_id, token in terminal_state) + + +@pytest.mark.asyncio +async def test_native_chat_callback_can_make_a_nested_native_call( + recording_server: RecordingServer, +) -> None: + import threading + + from litellm.integrations.custom_logger import CustomLogger + from tests.test_litellm_rust.support.requests import MESSAGES_RESPONSE + + recording_server.expected_requests = 2 + recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) + nested_responses: Final = [] + nested_started: Final = threading.Event() + + class NestedCall(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + if nested_started.is_set(): + return + nested_started.set() + nested_responses.append( + litellm.completion( + model="anthropic/claude-opus-5", + messages=[{"role": "user", "content": "nested"}], + max_tokens=16, + api_key="test-key", + api_base=recording_server.base_url, + num_retries=0, + ) + ) + + outer: Final = await litellm.acompletion( + model="anthropic/claude-opus-5", + messages=[{"role": "user", "content": "outer"}], + max_tokens=16, + api_key="test-key", + api_base=recording_server.base_url, + callbacks=[NestedCall()], + num_retries=0, + ) + + assert outer._hidden_params["additional_headers"]["x-litellm-rust"] == "true" + assert nested_responses[0]._hidden_params["additional_headers"]["x-litellm-rust"] == "true" + sent: Final = {request.body["messages"][0]["content"][0]["text"] for request in recording_server.requests} + assert sent == {"outer", "nested"} + + +@pytest.mark.asyncio +async def test_native_chat_cancellation_during_io_does_not_publish_a_terminal( + recording_server: RecordingServer, +) -> None: + import asyncio + + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from tests.test_litellm_rust.support.callback_recorder import RecordingLogger + from tests.test_litellm_rust.support.requests import MESSAGES_RESPONSE + + recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE, delay=0.5) + recorder: Final = RecordingLogger() + task: Final = asyncio.create_task( + litellm.acompletion( + model="anthropic/claude-opus-5", + messages=[{"role": "user", "content": "cancel"}], + max_tokens=16, + api_key="test-key", + api_base=recording_server.base_url, + callbacks=[recorder], + num_retries=0, + ) + ) + async with asyncio.timeout(10): + while not recording_server.requests: + await asyncio.sleep(0.01) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + + assert recorder.names.count("log_pre_api_call") == 1 + assert not { + "log_success_event", + "async_log_success_event", + "log_failure_event", + "async_log_failure_event", + }.intersection(recorder.names) diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py index 02d274bb405..ccf321a1c8b 100644 --- a/tests/test_litellm_rust/conftest.py +++ b/tests/test_litellm_rust/conftest.py @@ -1,24 +1,181 @@ +import asyncio import os +from collections.abc import AsyncIterator, Generator +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack, asynccontextmanager, contextmanager +from threading import Lock +from types import ModuleType +from typing import Final, Literal, cast import pytest +import pytest_asyncio +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +import litellm +from litellm import utils +from litellm.litellm_core_utils import litellm_logging +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.rust_bridge.configuration import ( # pyright: ignore[reportPrivateUsage] # preserve raw configuration state in test isolation + _CONFIGURATION, + _parse_env_bool, +) +from tests.test_litellm_rust.support.callback_recorder import drain_logging +from tests.test_litellm_rust.support.prometheus import isolated_prometheus_registry +from tests.test_litellm_rust.support.recording_server import RecordingServer, recording_service + +CALLBACK_ATTRIBUTES: Final = ( + "callbacks", + "input_callback", + "success_callback", + "failure_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", +) +Backend = Literal["python", "rust"] -def pytest_collection_modifyitems(items): - rust_enabled = os.environ.get("LITELLM_RUST", "").strip().lower() in { - "1", - "true", - "yes", - "on", - } - if not rust_enabled: +class _BackendScopes: + def __init__(self) -> None: + self._owner: asyncio.Task[object] | None = None + self._lock: Final = Lock() + + @contextmanager + def enter(self) -> Generator[None]: + task: Final = asyncio.current_task() + with self._lock: + previous: Final = self._owner + if previous is not None and previous is not task: + raise RuntimeError("isolated_backend scopes cannot overlap across tasks; run them sequentially") + self._owner = task + try: + yield + finally: + with self._lock: + self._owner = previous + + +_BACKEND_SCOPES: Final = _BackendScopes() + + +def _list_attribute(container: ModuleType, attribute: str) -> list[object]: + value: Final = getattr(container, attribute) + if not isinstance(value, list): + raise AssertionError(f"{container.__name__}.{attribute} is not a list") + return cast(list[object], value) + + +@contextmanager +def _isolated_list(container: ModuleType, attribute: str) -> Generator[None]: + source: Final = _list_attribute(container, attribute) + original: Final = list(source) + source.clear() # mutable-ok: test isolation mutates global registries by design + try: + yield + finally: + source.clear() + source.extend(original) + setattr(container, attribute, source) + + +@contextmanager +def _rebound(container: object, attribute: str, value: object) -> Generator[None]: + original: Final[object] = getattr(container, attribute) + setattr(container, attribute, value) + try: + yield + finally: + setattr(container, attribute, original) + + +@contextmanager +def _rust_mode(enabled: bool) -> Generator[None]: + with _rebound(_CONFIGURATION, "override", enabled): + yield + + +@asynccontextmanager +async def isolated_backend(backend: Backend) -> AsyncIterator[ExitStack]: + with _BACKEND_SCOPES.enter(): + async with _isolated_backend(backend) as stack: + yield stack + + +@asynccontextmanager +async def _isolated_backend(backend: Backend) -> AsyncIterator[ExitStack]: + with ExitStack() as stack: + for attribute in CALLBACK_ATTRIBUTES: + stack.enter_context(_isolated_list(litellm, attribute)) + stack.enter_context(_isolated_list(litellm_logging, "_in_memory_loggers")) # pyright: ignore[reportPrivateUsage] # string-name callback cache has no public accessor + stack.enter_context( + _rebound(utils, "callback_list", []) + ) # rebind-ok: legacy global registry mutated by set_callbacks + stack.enter_context( + _rebound(litellm, "cache", None) + ) # test-quality-ok: isolate the process-global cache from native extension tests + stack.enter_context(isolated_prometheus_registry()) + stack.enter_context(_rust_mode(backend == "rust")) + executor: Final = ThreadPoolExecutor(thread_name_prefix="rust-test-logging") + stack.enter_context(_rebound(utils, "executor", executor)) + try: + yield stack + finally: + try: + await drain_logging() + finally: + await asyncio.to_thread(executor.shutdown, wait=True) + await GLOBAL_LOGGING_WORKER.stop() + + +@pytest_asyncio.fixture(autouse=True, loop_scope="function") +async def isolate_rust_state() -> AsyncIterator[ExitStack]: + # pytest-asyncio runs fixture setup and the test body in different tasks. + async with _isolated_backend("rust") as stack: + yield stack + + +@pytest.fixture +def recording_server() -> Generator[RecordingServer]: + with recording_service() as server: + yield server + + +def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: + expected_failure = pytest.mark.xfail( + reason="requires the retained callback implementation from #40070", + strict=False, + ) + for item in items: + if "test_litellm_rust" in item.path.parts and item.path.name != "test_ocr.py": + item.add_marker(expected_failure) + + if not _parse_env_bool(os.environ.get("LITELLM_RUST")): skip = pytest.mark.skip(reason="requires LITELLM_RUST=1 and a compiled Rust extension") for item in items: - item.add_marker(skip) + 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 + raise pytest.UsageError("LITELLM_RUST=1 requires a compiled litellm.rust_bridge._native extension") from error + + +@pytest.fixture +def isolated_azure_auth(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ( + "AZURE_AI_API_KEY", + "AZURE_AI_API_BASE", + "AZURE_AD_TOKEN", + "AZURE_TENANT_ID", + "AZURE_CLIENT_ID", + "AZURE_CLIENT_SECRET", + "AZURE_USERNAME", + "AZURE_PASSWORD", + ): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", False) diff --git a/tests/test_litellm_rust/integrations/__init__.py b/tests/test_litellm_rust/integrations/__init__.py new file mode 100644 index 00000000000..c5f3569c9e0 --- /dev/null +++ b/tests/test_litellm_rust/integrations/__init__.py @@ -0,0 +1,75 @@ +from tests.test_litellm_rust.integrations.catalogue import ( + DISCOVERED_ONLY_GUARDRAIL_NAMES, + ENTERPRISE_LOGGER_NAMES, + GUARDRAIL_NAMES, + GUARDRAIL_OBLIGATIONS, + LOGGER_OBLIGATIONS, + OSS_LOGGER_NAMES, + REQUIRED_GUARDRAIL_BEHAVIOR, + REQUIRED_LOGGER_BEHAVIOR, +) +from tests.test_litellm_rust.integrations.harness import ( + AsyncBoundaryLogger, + GCSLiteralAIHarness, + GenericAPIExportHarness, + MutatingFailingLogger, + OtelHarness, + PurviewHarness, + RecordingAsyncClient, + RecordingVertexInstance, + ReviewGuardrail, + RunObservation, + gcs_literalai_harness, + metric_value, + purview_harness, + wait_for_audits, +) +from tests.test_litellm_rust.support.routes import ( + ALL_ROUTES, + ASYNC_ROUTES, + MESSAGES_ROUTE, + MESSAGES_STREAM, + NON_STREAM_ASYNC_ROUTES, + OCR_ASYNC, + OCR_SYNC, + Route, + provider_response, + route_id, + wait_for_callback, +) + +__all__ = ( + "ALL_ROUTES", + "ASYNC_ROUTES", + "DISCOVERED_ONLY_GUARDRAIL_NAMES", + "ENTERPRISE_LOGGER_NAMES", + "GUARDRAIL_NAMES", + "GUARDRAIL_OBLIGATIONS", + "LOGGER_OBLIGATIONS", + "MESSAGES_ROUTE", + "MESSAGES_STREAM", + "NON_STREAM_ASYNC_ROUTES", + "OCR_ASYNC", + "OCR_SYNC", + "OSS_LOGGER_NAMES", + "REQUIRED_GUARDRAIL_BEHAVIOR", + "REQUIRED_LOGGER_BEHAVIOR", + "AsyncBoundaryLogger", + "GCSLiteralAIHarness", + "GenericAPIExportHarness", + "MutatingFailingLogger", + "OtelHarness", + "PurviewHarness", + "RecordingAsyncClient", + "RecordingVertexInstance", + "ReviewGuardrail", + "Route", + "RunObservation", + "gcs_literalai_harness", + "metric_value", + "provider_response", + "purview_harness", + "route_id", + "wait_for_audits", + "wait_for_callback", +) diff --git a/tests/test_litellm_rust/integrations/catalogue.py b/tests/test_litellm_rust/integrations/catalogue.py new file mode 100644 index 00000000000..3bcd0c9e7bf --- /dev/null +++ b/tests/test_litellm_rust/integrations/catalogue.py @@ -0,0 +1,193 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Literal + +DependencyProfile = Literal["required", "enterprise"] + +OSS_LOGGER_NAMES: Final = frozenset( + { + "agentops", + "anthropic_cache_control_hook", + "argilla", + "arize", + "arize_phoenix", + "aws_sqs", + "azure_sentinel", + "azure_storage", + "bitbucket", + "braintrust", + "cloudzero", + "datadog", + "datadog_llm_observability", + "datadog_metrics", + "deepeval", + "dotprompt", + "dynamic_rate_limiter", + "dynamic_rate_limiter_v3", + "focus", + "galileo", + "gcs_bucket", + "gcs_pubsub", + "gitlab", + "humanloop", + "lago", + "langfuse", + "langfuse_otel", + "langsmith", + "langtrace", + "levo", + "literalai", + "litellm_agent", + "logfire", + "mavvrik", + "mlflow", + "newrelic", + "opentelemetry", + "openmeter", + "opik", + "otel", + "posthog", + "prometheus", + "s3_v2", + "vantage", + "vector_store_pre_call_hook", + "weave_otel", + } +) +ENTERPRISE_LOGGER_NAMES: Final = frozenset({"generic_api", "pagerduty", "resend_email", "sendgrid_email", "smtp_email"}) +GUARDRAIL_NAMES: Final = frozenset( + { + "aim", + "akto", + "alice", + "aporia", + "azure/prompt_shield", + "azure/text_moderations", + "bedrock", + "block_code_execution", + "cato_networks", + "cisco_ai_defense", + "compresr", + "crowdstrike_aidr", + "custom_code", + "deepkeep", + "dynamoai", + "enkryptai", + "generic_guardrail_api", + "grayswan", + "guardrails_ai", + "headroom", + "hiddenlayer", + "hide-secrets", + "ibm_guardrails", + "javelin", + "lakera", + "lakera_v2", + "lasso", + "litellm_content_filter", + "llm_as_a_judge", + "mcp_end_user_permission", + "mcp_jwt_signer", + "mcp_security", + "microsoft_purview", + "model_armor", + "noma", + "noma_v2", + "onyx", + "openai_moderation", + "ovalix", + "pangea", + "panw_prisma_airs", + "pillar", + "presidio", + "prompt_security", + "promptguard", + "qostodian_nexus", + "qualifire", + "repelloai", + "rubrik", + "semantic_guard", + "singulr", + "straiker", + "tool_permission", + "vigil_guard", + "xecguard", + "zscaler_ai_guard", + } +) +DISCOVERED_ONLY_GUARDRAIL_NAMES: Final = frozenset({"tool_policy"}) + + +@dataclass(frozen=True, slots=True) +class CoverageObligation: + stable_name: str + registration_names: tuple[str, ...] + dependency_profile: DependencyProfile + behavioral_cases: tuple[str, ...] + + +REQUIRED_LOGGER_BEHAVIOR: Final = frozenset({"generic_api", "gcs_bucket", "literalai", "prometheus", "opentelemetry"}) +REQUIRED_GUARDRAIL_BEHAVIOR: Final = frozenset( + { + "azure/text_moderations", + "crowdstrike_aidr", + "litellm_content_filter", + "microsoft_purview", + "rubrik", + } +) + +LOGGER_BEHAVIORAL_CASES: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType( + { + "generic_api": ("generic-api-success",), + "gcs_bucket": ("gcs-literalai-scheduling",), + "literalai": ("gcs-literalai-scheduling",), + "prometheus": ("prometheus-string-registration",), + "opentelemetry": ("otel-export",), + } +) +GUARDRAIL_BEHAVIORAL_CASES: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType( + { + "azure/text_moderations": ("azure-text-moderation",), + "crowdstrike_aidr": ("crowdstrike-redaction-native-chat",), + "rubrik": ("rubrik-block-native-chat",), + "microsoft_purview": ("purview-audit-native-chat",), + "litellm_content_filter": ("content-filter-block",), + } +) + + +def _logger_obligations() -> Mapping[str, CoverageObligation]: + from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry + + registry: Final = CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE + expected_names: Final = OSS_LOGGER_NAMES | ENTERPRISE_LOGGER_NAMES + grouped: Final[dict[type[object], tuple[str, ...]]] = { + implementation: tuple(sorted(name for name in expected_names if registry.get(name) is implementation)) + for implementation in frozenset(registry.values()) + } + obligations: Final = { + names[0]: CoverageObligation( + stable_name=names[0], + registration_names=names, + dependency_profile="enterprise" if set(names) & ENTERPRISE_LOGGER_NAMES else "required", + behavioral_cases=next( + (LOGGER_BEHAVIORAL_CASES[name] for name in names if name in LOGGER_BEHAVIORAL_CASES), () + ), + ) + for names in grouped.values() + if names + } + missing_optional: Final = ENTERPRISE_LOGGER_NAMES - registry.keys() + unavailable: Final = {name: CoverageObligation(name, (name,), "enterprise", ()) for name in missing_optional} + return MappingProxyType({**obligations, **unavailable}) + + +LOGGER_OBLIGATIONS: Final = _logger_obligations() +GUARDRAIL_OBLIGATIONS: Final = MappingProxyType( + { + name: CoverageObligation(name, (name,), "required", GUARDRAIL_BEHAVIORAL_CASES.get(name, ())) + for name in GUARDRAIL_NAMES | DISCOVERED_ONLY_GUARDRAIL_NAMES + } +) diff --git a/tests/test_litellm_rust/integrations/conftest.py b/tests/test_litellm_rust/integrations/conftest.py new file mode 100644 index 00000000000..1adf2bca5c2 --- /dev/null +++ b/tests/test_litellm_rust/integrations/conftest.py @@ -0,0 +1,43 @@ +from contextlib import ExitStack +from typing import Final + +import pytest +import pytest_asyncio +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger +from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig +from litellm.integrations.prometheus import PrometheusLogger +from tests.test_litellm_rust.integrations import GenericAPIExportHarness, OtelHarness, Route, provider_response +from tests.test_litellm_rust.support.recording_server import RecordingServer, recording_service + + +@pytest.fixture +def provider(recording_server: RecordingServer, route: Route) -> RecordingServer: + recording_server.default_response = provider_response(route) + return recording_server + + +@pytest.fixture +def otel(isolate_rust_state: ExitStack) -> OtelHarness: + exporter: Final = InMemorySpanExporter() + tracer_provider: Final = TracerProvider() + isolate_rust_state.callback(tracer_provider.shutdown) + tracer_provider.add_span_processor(SimpleSpanProcessor(exporter)) + logger: Final = OpenTelemetry(config=OpenTelemetryConfig(exporter=exporter), tracer_provider=tracer_provider) + return OtelHarness(logger=logger, exporter=exporter) + + +@pytest.fixture +def prometheus() -> PrometheusLogger: + return PrometheusLogger() + + +@pytest_asyncio.fixture +async def generic_api_export(isolate_rust_state: ExitStack) -> GenericAPIExportHarness: + server: Final = isolate_rust_state.enter_context(recording_service()) + server.expected_requests = None + logger: Final = GenericAPILogger(endpoint=f"{server.base_url}/logs", batch_size=1, log_format="single") + return GenericAPIExportHarness(server=server, logger=logger) diff --git a/tests/test_litellm_rust/integrations/harness.py b/tests/test_litellm_rust/integrations/harness.py new file mode 100644 index 00000000000..4c2cd854a85 --- /dev/null +++ b/tests/test_litellm_rust/integrations/harness.py @@ -0,0 +1,203 @@ +import asyncio +import threading +import time +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping +from contextlib import asynccontextmanager +from dataclasses import dataclass +from typing import Final + +import httpx +from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from prometheus_client import REGISTRY + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger +from litellm.integrations.gcs_bucket.gcs_bucket_base import IAM_AUTH_KEY +from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger +from litellm.integrations.literal_ai import LiteralAILogger +from litellm.integrations.opentelemetry import LITELLM_REQUEST_SPAN_NAME, OpenTelemetry +from litellm.proxy.guardrails.guardrail_hooks.microsoft_purview.purview_dlp import MicrosoftPurviewDLPGuardrail +from litellm.types.guardrails import GuardrailEventHooks +from tests.test_litellm_rust.support.callback_recorder import drain_logging +from tests.test_litellm_rust.support.recording_server import RecordedRequest, RecordingServer + + +@dataclass(frozen=True, slots=True) +class RunObservation: + call_type: str + model: str + response_cost: float + response_text: str + provider_body: Mapping[str, object] + rust_dispatch: bool + + +@dataclass(frozen=True, slots=True) +class RecordedPost: + url: str + headers: Mapping[str, str] + body: object + + +class RecordingAsyncClient: + def __init__( + self, responses: tuple[Mapping[str, object], ...] = (), blocked_url_fragment: str | None = None + ) -> None: + self.posts: list[RecordedPost] = [] + self._responses = list(responses) + self._blocked_url_fragment = blocked_url_fragment + self.accepted = threading.Event() + self.release = threading.Event() + self.release.set() + + async def post(self, url: str, **kwargs: object) -> httpx.Response: + data: Final = kwargs.get("data") + json_body: Final = kwargs.get("json") + body: Final = json_body if json_body is not None else data + headers_value: Final = kwargs.get("headers") + headers: Final = headers_value if isinstance(headers_value, Mapping) else {} + self.posts.append(RecordedPost(url=url, headers=headers, body=body)) + should_block: Final = self._blocked_url_fragment is None or self._blocked_url_fragment in url + if should_block: + self.accepted.set() + released: Final = await asyncio.to_thread(self.release.wait, 10) + if not released: + raise TimeoutError(f"Timed out releasing POST {url}") + payload: Final = self._responses.pop(0) if self._responses else {} + response_headers: Final = {"etag": '"test-scope"'} if "protectionScopes/compute" in url else {} + return httpx.Response(200, json=payload, headers=response_headers, request=httpx.Request("POST", url)) + + +class RecordingVertexInstance: + async def _ensure_access_token_async(self, **kwargs: object) -> tuple[str, str]: + return "test-access-token", "test-project" + + def _get_token_and_url(self, **kwargs: object) -> tuple[str, str]: + return "test-access-token", "" + + +class AsyncBoundaryLogger(CustomLogger): + def __init__(self, action: Callable[[], Awaitable[None]]) -> None: + self._action = action + + async def async_log_success_event( + self, kwargs: object, response_obj: object, start_time: object, end_time: object + ) -> None: + await self._action() + + +@dataclass(frozen=True, slots=True) +class OtelHarness: + logger: OpenTelemetry + exporter: InMemorySpanExporter + + def spans(self, name: str = LITELLM_REQUEST_SPAN_NAME) -> tuple[ReadableSpan, ...]: + return tuple(span for span in self.exporter.get_finished_spans() if span.name == name) + + async def wait_for_spans( + self, name: str = LITELLM_REQUEST_SPAN_NAME, count: int = 1, timeout: float = 10 + ) -> tuple[ReadableSpan, ...]: + deadline: Final = time.monotonic() + timeout + while len(self.spans(name)) < count: + if time.monotonic() >= deadline: + raise TimeoutError(f"Timed out waiting for {count} {name} spans; saw {self.spans(name)}") + await asyncio.sleep(0.01) + await drain_logging() + return self.spans(name) + + +@dataclass(frozen=True, slots=True) +class GCSLiteralAIHarness: + gcs: GCSBucketLogger + literal: LiteralAILogger + storage: RecordingAsyncClient + literal_sink: RecordingAsyncClient + + +@dataclass(frozen=True, slots=True) +class GenericAPIExportHarness: + server: RecordingServer + logger: GenericAPILogger + + @property + def exports(self) -> tuple[RecordedRequest, ...]: + return tuple(request for request in self.server.requests if request.path == "/logs") + + +@asynccontextmanager +async def gcs_literalai_harness() -> AsyncIterator[GCSLiteralAIHarness]: + storage: Final = RecordingAsyncClient() + literal_sink: Final = RecordingAsyncClient() + gcs: Final = GCSBucketLogger(bucket_name="composition-bucket") + gcs.async_httpx_client = storage + gcs.vertex_instances[IAM_AUTH_KEY] = RecordingVertexInstance() + literal: Final = LiteralAILogger(literalai_api_key="test-key") + literal.async_httpx_client = literal_sink + try: + yield GCSLiteralAIHarness(gcs=gcs, literal=literal, storage=storage, literal_sink=literal_sink) + finally: + await gcs.aclose() + + +@dataclass(frozen=True, slots=True) +class PurviewHarness: + graph: RecordingAsyncClient + guardrail: MicrosoftPurviewDLPGuardrail + + +def purview_harness(name: str) -> PurviewHarness: + graph: Final = RecordingAsyncClient( + responses=({"access_token": "test-token", "expires_in": 3600}, {}, {}, {}), + blocked_url_fragment="processContent", + ) + graph.release.clear() + guardrail: Final = MicrosoftPurviewDLPGuardrail( + guardrail_name=name, + tenant_id="test-tenant", + client_id="test-client", + client_secret="test-secret", + event_hook=GuardrailEventHooks.logging_only, + default_on=True, + ) + guardrail.async_handler = graph + return PurviewHarness(graph=graph, guardrail=guardrail) + + +async def wait_for_audits(graph: RecordingAsyncClient, count: int = 2, timeout: float = 3) -> tuple[RecordedPost, ...]: + async with asyncio.timeout(timeout): + while len(tuple(post for post in graph.posts if "processContent" in post.url)) < count: + await asyncio.sleep(0.01) + return tuple(post for post in graph.posts if "processContent" in post.url) + + +def metric_value(name: str, **labels: str) -> float: + for metric in REGISTRY.collect(): + for sample in metric.samples: + if sample.name == name and all(sample.labels.get(key) == value for key, value in labels.items()): + return sample.value + return 0.0 + + +class ReviewGuardrail(CustomGuardrail): + def __init__(self, review: Callable[[object], Awaitable[object]]) -> None: + super().__init__(guardrail_name="rust-review", event_hook=GuardrailEventHooks.post_call, default_on=True) + self._review = review + self.call_types: list[object] = [] + + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + self.call_types.append(call_type) + return await self._review(response) + + +class MutatingFailingLogger(CustomLogger): + def log_success_event(self, kwargs, response_obj, start_time, end_time): + payload: Final = kwargs["standard_logging_object"] + payload["metadata"]["composition_marker"] = "visible-before-failure" + raise RuntimeError("synthetic callback failure") + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + payload: Final = kwargs["standard_logging_object"] + payload["metadata"]["composition_marker"] = "visible-before-failure" + raise RuntimeError("synthetic callback failure") diff --git a/tests/test_litellm_rust/integrations/test_composition.py b/tests/test_litellm_rust/integrations/test_composition.py new file mode 100644 index 00000000000..a18c7de7f0e --- /dev/null +++ b/tests/test_litellm_rust/integrations/test_composition.py @@ -0,0 +1,548 @@ +import asyncio +import gc +import json +import threading +import weakref +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import datetime +from typing import Final + +import pytest +from opentelemetry.trace import StatusCode + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger +from litellm.integrations.prometheus import PrometheusLogger +from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.proxy.guardrails.guardrail_registry import guardrail_initializer_registry +from litellm.types.guardrails import SupportedGuardrailIntegrations +from litellm.types.utils import CallTypes +from tests.test_litellm_rust.conftest import Backend, isolated_backend +from tests.test_litellm_rust.integrations import ( + ASYNC_ROUTES, + DISCOVERED_ONLY_GUARDRAIL_NAMES, + ENTERPRISE_LOGGER_NAMES, + GUARDRAIL_NAMES, + GUARDRAIL_OBLIGATIONS, + LOGGER_OBLIGATIONS, + MESSAGES_ROUTE, + MESSAGES_STREAM, + OCR_ASYNC, + OCR_SYNC, + OSS_LOGGER_NAMES, + REQUIRED_GUARDRAIL_BEHAVIOR, + REQUIRED_LOGGER_BEHAVIOR, + AsyncBoundaryLogger, + MutatingFailingLogger, + OtelHarness, + Route, + RunObservation, + gcs_literalai_harness, + metric_value, + provider_response, + route_id, + wait_for_callback, +) +from tests.test_litellm_rust.support.callback_recorder import ( + LiveReferenceLogger, + RecordingLogger, + SecondaryLiveReferenceLogger, + drain_logging, +) +from tests.test_litellm_rust.support.provenance import has_rust_response_marker +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, recording_service +from tests.test_litellm_rust.support.requests import ( + CHAT_MESSAGES, + CHAT_MODEL, + CHAT_RESPONSE, + MESSAGES, + MESSAGES_EVENTS, +) + +pytestmark = pytest.mark.requires_rust_extension + +MESSAGES_TOOLS: Final = [ + { + "name": "get_weather", + "description": "Get current weather", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + } +] + + +async def observe_backend(route: Route, backend: Backend) -> RunObservation: + async with isolated_backend(backend): + with recording_service() as provider: + provider.default_response = provider_response(route) + recorder: Final = RecordingLogger() + response: Final = await route.invoke(provider, callbacks=[recorder]) + event: Final = (await wait_for_callback(route, recorder))[0] + payload: Final = event.kwargs["standard_logging_object"] + provider_body: Final = provider.requests[0].body + if not isinstance(provider_body, dict): + raise TypeError(f"Expected provider object body, got {type(provider_body).__name__}") + return RunObservation( + call_type=payload["call_type"], + model=payload["model"], + response_cost=payload["response_cost"], + response_text=route.response_text(response), + provider_body=provider_body, + rust_dispatch=has_rust_response_marker(response), + ) + + +def test_logger_catalogue_reconciles_with_production_registry() -> None: + actual_names: Final = frozenset(CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE) + catalogued_names: Final = frozenset( + name for obligation in LOGGER_OBLIGATIONS.values() for name in obligation.registration_names + ) + assert OSS_LOGGER_NAMES <= actual_names + assert actual_names <= OSS_LOGGER_NAMES | ENTERPRISE_LOGGER_NAMES + assert catalogued_names == OSS_LOGGER_NAMES | ENTERPRISE_LOGGER_NAMES + + +def test_guardrail_catalogue_reconciles_enum_and_runtime_discovery() -> None: + enum_names: Final = frozenset(integration.value for integration in SupportedGuardrailIntegrations) + initializer_names: Final = frozenset(guardrail_initializer_registry) + assert enum_names == GUARDRAIL_NAMES + assert initializer_names == GUARDRAIL_NAMES | DISCOVERED_ONLY_GUARDRAIL_NAMES + assert frozenset(GUARDRAIL_OBLIGATIONS) == initializer_names + + +def test_required_integrations_have_behavioral_cases() -> None: + logger_names: Final = frozenset( + name + for obligation in LOGGER_OBLIGATIONS.values() + if obligation.behavioral_cases + for name in obligation.registration_names + ) + guardrail_names: Final = frozenset( + name for name, obligation in GUARDRAIL_OBLIGATIONS.items() if obligation.behavioral_cases + ) + assert REQUIRED_LOGGER_BEHAVIOR <= logger_names + assert REQUIRED_GUARDRAIL_BEHAVIOR <= guardrail_names + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ASYNC_ROUTES, ids=lambda route: f"otel-prometheus-generic-api-{route.name}-success") +async def test_export_composition( + route: Route, recording_server: RecordingServer, otel: OtelHarness, prometheus: PrometheusLogger +) -> None: + recording_server.default_response = provider_response(route) + before: Final = metric_value("litellm_requests_metric_total", model=route.provider_model) + recorder: Final = RecordingLogger() + + with recording_service() as sink: + logger: Final = GenericAPILogger(endpoint=f"{sink.base_url}/logs", batch_size=1, log_format="single") + await route.invoke(recording_server, callbacks=[otel.logger, prometheus, logger, recorder]) + await wait_for_callback(route, recorder) + + exports: Final = tuple(request for request in sink.requests if request.path == "/logs") + spans: Final = await otel.wait_for_spans() + assert len(exports) == 1 + assert exports[0].body["call_type"] == route.call_type + assert exports[0].body["response_cost"] == pytest.approx(route.expected_cost) + assert len(spans) == 1 + assert spans[0].status.status_code is StatusCode.OK + assert metric_value("litellm_requests_metric_total", model=route.provider_model) == before + 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", (OCR_SYNC, OCR_ASYNC, MESSAGES_ROUTE), ids=route_id) +async def test_public_sdk_python_rust_composition_parity(route: Route) -> None: + python: Final = await observe_backend(route, "python") + rust: Final = await observe_backend(route, "rust") + assert python.call_type == rust.call_type == route.call_type + assert python.model == rust.model == route.provider_model + assert python.response_cost == pytest.approx(rust.response_cost) + assert python.response_text == rust.response_text == route.expected_text + python_body: Final = {**python.provider_body, "stream": python.provider_body.get("stream", False)} + rust_body: Final = {**rust.provider_body, "stream": rust.provider_body.get("stream", False)} + assert python_body == rust_body + assert python.rust_dispatch is False + assert rust.rust_dispatch is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ("python", "rust")) +async def test_terminal_callbacks_share_live_objects_within_one_run(backend: Backend) -> None: + async with isolated_backend(backend): + with recording_service() as provider: + provider.default_response = provider_response(MESSAGES_ROUTE) + first: Final = LiveReferenceLogger() + second: Final = SecondaryLiveReferenceLogger() + response: Final = await MESSAGES_ROUTE.invoke(provider, callbacks=[first, second]) + first_event: Final = (await first.wait_for_async())[0] + second_event: Final = (await second.wait_for_async())[0] + assert first_event.kwargs is second_event.kwargs + assert first_event.response is second_event.response + assert has_rust_response_marker(response) is (backend == "rust") + first.release() + second.release() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ("python", "rust")) +async def test_per_logger_redaction_preserves_the_shared_unredacted_payload(backend: Backend) -> None: + class Capture(CustomLogger): + def __init__(self, redact: bool) -> None: + super().__init__(turn_off_message_logging=redact) + self.kwargs: object = None + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.kwargs = kwargs + + async with isolated_backend(backend): + with recording_service() as provider: + provider.default_response = provider_response(MESSAGES_ROUTE) + redacted: Final = Capture(True) + plain: Final = Capture(False) + await MESSAGES_ROUTE.invoke(provider, callbacks=[redacted, plain]) + await drain_logging() + + redacted_kwargs: Final = redacted.kwargs + plain_kwargs: Final = plain.kwargs + assert isinstance(redacted_kwargs, dict) + assert isinstance(plain_kwargs, dict) + redacted_payload: Final = redacted_kwargs["standard_logging_object"] + plain_payload: Final = plain_kwargs["standard_logging_object"] + assert redacted_kwargs is not plain_kwargs + assert redacted_payload is not plain_payload + assert redacted_payload["metadata"] is plain_payload["metadata"] + assert "redacted-by-litellm" in json.dumps(redacted_payload["messages"]) + assert MESSAGES_ROUTE.expected_text in json.dumps(plain_payload["response"]) + assert "redacted-by-litellm" not in json.dumps(plain_payload) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ("python", "rust")) +@pytest.mark.parametrize("accepted", (True, False), ids=("accepted", "rejected")) +async def test_terminal_respects_proxy_deferred_completion_gate(backend: Backend, accepted: bool) -> None: + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + async with isolated_backend(backend): + with recording_service() as provider: + provider.default_response = provider_response(MESSAGES_ROUTE) + recorder: Final = RecordingLogger() + logger: Final = Logging( + model=MESSAGES_ROUTE.provider_model, + messages=MESSAGES, + stream=False, + call_type=MESSAGES_ROUTE.call_type, + start_time=datetime.now(), + litellm_call_id="deferred-completion", + function_id="deferred-completion", + dynamic_async_success_callbacks=[recorder], + dynamic_async_failure_callbacks=[recorder], + ) + logger._defer_async_logging = True + response: Final = await MESSAGES_ROUTE.invoke(provider, litellm_logging_obj=logger) + await asyncio.sleep(0) + assert "async_log_success_event" not in recorder.names + + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging( + logging_obj=logger, + exception_raised=not accepted, + ) + if accepted: + accepted_events: Final = await recorder.wait_for_async("async_log_success_event") + assert len(accepted_events) == 1 + assert "async_log_failure_event" not in recorder.names + else: + error: Final = RuntimeError("post-call guardrail rejected response") + await logger.async_failure_handler(error, str(error)) + rejected_events: Final = await recorder.wait_for_async("async_log_failure_event") + assert len(rejected_events) == 1 + assert "async_log_success_event" not in recorder.names + assert has_rust_response_marker(response) is (backend == "rust") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ("python", "rust")) +async def test_suspended_terminal_callback_owns_payload_until_it_finishes(backend: Backend) -> None: + class Root: + pass + + class Suspended(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.started = asyncio.Event() + self.release = asyncio.Event() + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + assert kwargs["litellm_params"]["metadata"]["retained"] is not None + self.started.set() + await self.release.wait() + + async def invoke() -> weakref.ReferenceType[Root]: + async with isolated_backend(backend): + with recording_service() as provider: + provider.default_response = provider_response(MESSAGES_ROUTE) + callback: Final = Suspended() + root = Root() + reference: Final = weakref.ref(root) + response: Final = await MESSAGES_ROUTE.invoke( + provider, + callbacks=[callback], + metadata={"retained": root}, + ) + del root + await asyncio.wait_for(callback.started.wait(), timeout=2) + gc.collect() + assert reference() is not None + callback.release.set() + await drain_logging() + del response + return reference + + reference: Final = await invoke() + await asyncio.sleep(0) + gc.collect() + assert reference() is None + + +async def observe_retry(backend: Backend) -> tuple[tuple[str, ...], bool]: + async with isolated_backend(backend): + with recording_service() as provider: + provider.expected_requests = 2 + provider.enqueue(ResponseSpec(body={"message": "retry this attempt"}, status=500)) + provider.default_response = ResponseSpec(body=CHAT_RESPONSE) + recorder: Final = RecordingLogger() + response: Final = await litellm.acompletion( + model=CHAT_MODEL, + messages=CHAT_MESSAGES, + api_key="test-key", + api_base=provider.base_url, + callbacks=[recorder], + num_retries=1, + ) + await drain_logging() + return recorder.names, has_rust_response_marker(response) + + +@pytest.mark.asyncio +async def test_retry_attempt_callback_sequence_matches_python() -> None: + python_names, python_rust_dispatch = await observe_retry("python") + rust_names, rust_rust_dispatch = await observe_retry("rust") + assert rust_names == python_names + assert rust_names.count("log_pre_api_call") == 2 + assert rust_names.count("async_log_failure_event") == 1 + assert rust_names.count("async_log_success_event") == 0 + assert python_rust_dispatch is False + assert rust_rust_dispatch is True + + +@pytest.mark.asyncio +async def test_failing_callback_preserves_prior_mutation_and_later_exporters() -> None: + with recording_service() as provider, recording_service() as sink: + provider.default_response = provider_response(OCR_ASYNC) + logger: Final = GenericAPILogger(endpoint=f"{sink.base_url}/logs", batch_size=1, log_format="single") + recorder: Final = RecordingLogger() + response: Final = await OCR_ASYNC.invoke(provider, callbacks=[MutatingFailingLogger(), logger, recorder]) + events: Final = await wait_for_callback(OCR_ASYNC, recorder) + exports: Final = tuple(request for request in sink.requests if request.path == "/logs") + assert response.pages[0].markdown == OCR_ASYNC.expected_text + assert events[0].kwargs["standard_logging_object"]["metadata"]["composition_marker"] == "visible-before-failure" + assert len(exports) == 1 + assert exports[0].body["metadata"]["composition_marker"] == "visible-before-failure" + + +@dataclass(frozen=True, slots=True) +class SerializationSchedule: + callback_order: tuple[str, str] + flush_immediately: bool + gcs_has_tools: bool + + +SERIALIZATION_SCHEDULES: Final = ( + pytest.param( + SerializationSchedule(("gcs", "literalai"), True, True), id="gcs-serializes-before-literalai-mutation" + ), + pytest.param( + SerializationSchedule(("gcs", "literalai"), False, False), id="gcs-serializes-after-literalai-mutation" + ), + pytest.param(SerializationSchedule(("literalai", "gcs"), False, False), id="literalai-mutates-before-gcs-enqueue"), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ("python", "rust")) +@pytest.mark.parametrize("schedule", SERIALIZATION_SCHEDULES) +async def test_gcs_literalai_serialization_schedule( + backend: Backend, schedule: SerializationSchedule, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "premium_user", True) + monkeypatch.setenv("GCS_BATCH_SIZE", "1" if schedule.flush_immediately else "100") + monkeypatch.setenv("GCS_BUCKET_NAME", "composition-bucket") + monkeypatch.setenv("GCS_FLUSH_INTERVAL", "3600") + monkeypatch.setenv("GCS_USE_BATCHED_LOGGING", "true") + monkeypatch.setenv("LITERAL_BATCH_SIZE", "1") + + async with isolated_backend(backend): + with recording_service() as provider: + provider.default_response = provider_response(MESSAGES_ROUTE) + async with gcs_literalai_harness() as harness: + callbacks_by_name: Final = {"gcs": harness.gcs, "literalai": harness.literal} + ordered_callbacks: Final = [callbacks_by_name[name] for name in schedule.callback_order] + callbacks: Final = ( + [harness.gcs, AsyncBoundaryLogger(harness.gcs.flush_queue), harness.literal] + if schedule.flush_immediately + else ordered_callbacks + ) + recorder: Final = RecordingLogger() + response: Final = await MESSAGES_ROUTE.invoke( + provider, tools=MESSAGES_TOOLS, callbacks=[*callbacks, recorder] + ) + await wait_for_callback(MESSAGES_ROUTE, recorder) + if not schedule.flush_immediately: + await harness.gcs.flush_queue() + + gcs_payload_value: Final = harness.storage.posts[0].body + if not isinstance(gcs_payload_value, str): + raise TypeError(f"Expected serialized GCS payload, got {type(gcs_payload_value).__name__}") + gcs_payload: Final = json.loads(gcs_payload_value) + literal_body: Final = harness.literal_sink.posts[0].body + if not isinstance(literal_body, dict): + raise TypeError(f"Expected LiteralAI request body, got {type(literal_body).__name__}") + generation: Final = literal_body["variables"]["generation_0"] + assert ("tools" in gcs_payload["model_parameters"]) is schedule.gcs_has_tools + assert generation["tools"] == MESSAGES_TOOLS + assert len(harness.storage.posts) == 1 + assert len(harness.literal_sink.posts) == 1 + assert has_rust_response_marker(response) is (backend == "rust") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ("python", "rust")) +async def test_interrupted_stream_emits_terminal_only_when_consumer_closes(backend: Backend, otel: OtelHarness) -> None: + async with isolated_backend(backend): + with recording_service() as provider: + provider.default_response = ResponseSpec(body=None, events=MESSAGES_EVENTS) + stream: Final = await MESSAGES_STREAM.open_stream(provider, callbacks=[otel.logger]) + first_chunk: Final = await anext(stream) + await drain_logging() + assert otel.spans() == () + assert has_rust_response_marker(stream) is (backend == "rust") + await stream.aclose() + await drain_logging() + assert first_chunk is not None + assert len(otel.spans()) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ("python", "rust")) +@pytest.mark.parametrize( + ("route", "phase"), + ( + (OCR_ASYNC, "pre"), + (MESSAGES_ROUTE, "pre"), + (MESSAGES_STREAM, "pre"), + (OCR_ASYNC, "success"), + (MESSAGES_ROUTE, "success"), + ), + ids=("ocr-pre", "messages-pre", "messages-stream-pre", "ocr-success", "messages-success"), +) +async def test_deployment_rejection_logs_failure_without_deployment_failure( + backend: Backend, route: Route, phase: str +) -> None: + class Reject(CustomLogger): + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: CallTypes | None + ) -> dict[str, object]: + if phase == "pre": + raise litellm.BadRequestError("deployment rejected", "test-provider", route.provider_model) + return kwargs + + async def async_post_call_success_deployment_hook( + self, request_data: Mapping[str, object], response: object, call_type: CallTypes | None + ) -> object: + raise litellm.BadRequestError("deployment rejected", "test-provider", route.provider_model) + + async with isolated_backend(backend): + with recording_service() as provider: + provider.expected_requests = 0 if phase == "pre" else 1 + provider.default_response = provider_response(route) + recorder: Final = RecordingLogger() + litellm.callbacks.extend((Reject(), recorder)) + + with pytest.raises(litellm.BadRequestError, match="deployment rejected"): + await route.invoke(provider) + await recorder.wait_for_async("async_log_failure_event") + + assert len(provider.requests) == provider.expected_requests + assert "async_log_success_event" not in recorder.names + assert recorder.names.count("async_post_call_failure_deployment_hook") == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ("rust",)) +@pytest.mark.parametrize("ending", ("provider_error", "truncated", "close")) +async def test_established_stream_completion_does_not_notify_deployment_failure(backend: Backend, ending: str) -> None: + async with isolated_backend(backend): + with recording_service() as provider: + provider.default_response = ResponseSpec( + body=None, + events=( + ( + *MESSAGES_EVENTS[:1], + ( + "error", + {"type": "error", "error": {"type": "overloaded_error", "message": "upstream overloaded"}}, + ), + ) + if ending == "provider_error" + else MESSAGES_EVENTS[:-1] + if ending == "truncated" + else MESSAGES_EVENTS + ), + ) + release: Final = threading.Event() + if ending == "close": + provider.enqueue( + ResponseSpec( + body=None, + chunks=tuple( + f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in MESSAGES_EVENTS + ), + release=release, + ) + ) + recorder: Final = RecordingLogger() + litellm.callbacks.append(recorder) + try: + stream: Final = await MESSAGES_STREAM.open_stream(provider) + assert has_rust_response_marker(stream) is (backend == "rust") + assert recorder.names.count("async_pre_call_deployment_hook") == 1 + assert "async_post_call_failure_deployment_hook" not in recorder.names + + if ending == "close": + assert await anext(stream) + await asyncio.wait_for(stream.aclose(), timeout=1) + await stream.aclose() + assert "async_log_success_event" not in recorder.names + assert "async_log_failure_event" not in recorder.names + release.set() + await recorder.wait_for_async("async_log_success_event") + else: + try: + async for _ in stream: + pass + except litellm.APIError: + pass + await drain_logging() + finally: + release.set() + + assert len(provider.requests) == 1 + assert recorder.names.count("async_log_failure_event") == (0 if ending == "close" else 1), recorder.names + assert recorder.names.count("async_log_success_event") == (1 if ending == "close" else 0), recorder.names + assert recorder.names.count("async_post_call_failure_deployment_hook") == 0 diff --git a/tests/test_litellm_rust/integrations/test_exporters.py b/tests/test_litellm_rust/integrations/test_exporters.py new file mode 100644 index 00000000000..08ce78e2a92 --- /dev/null +++ b/tests/test_litellm_rust/integrations/test_exporters.py @@ -0,0 +1,223 @@ +import json +from typing import Final + +import pytest +from opentelemetry.trace import StatusCode +from prometheus_client import CollectorRegistry, Counter + +import litellm +from litellm.integrations.prometheus import PrometheusLogger +from litellm.litellm_core_utils import litellm_logging +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +from tests.test_litellm_rust.support.prometheus import isolated_prometheus_registry +from tests.test_litellm_rust.support.requests import MESSAGES_EVENTS +from tests.test_litellm_rust.integrations import ( + ALL_ROUTES, + ASYNC_ROUTES, + MESSAGES_ROUTE, + MESSAGES_STREAM, + NON_STREAM_ASYNC_ROUTES, + OCR_ASYNC, + OCR_SYNC, + GenericAPIExportHarness, + OtelHarness, + Route, + metric_value, + route_id, + wait_for_callback, +) +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec + +pytestmark = pytest.mark.requires_rust_extension + +FAILURE_RESPONSE: Final = ResponseSpec(body={"message": "provider unavailable"}, status=500) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ASYNC_ROUTES, ids=route_id) +async def test_generic_api_logger_exports_success_over_http( + route: Route, + provider: RecordingServer, + generic_api_export: GenericAPIExportHarness, +) -> None: + recorder: Final = RecordingLogger() + await route.invoke(provider, callbacks=[generic_api_export.logger, recorder]) + await wait_for_callback(route, recorder) + + assert len(generic_api_export.exports) == 1 + payload: Final = generic_api_export.exports[0].body + assert payload["status"] == "success" + assert payload["call_type"] == route.call_type + assert payload["model"] == route.provider_model + assert payload["response_cost"] == pytest.approx(route.expected_cost) + assert route.expected_text in json.dumps(payload["response"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", NON_STREAM_ASYNC_ROUTES, ids=route_id) +async def test_generic_api_logger_exports_provider_failure_over_http( + route: Route, + provider: RecordingServer, + generic_api_export: GenericAPIExportHarness, +) -> None: + provider.enqueue(FAILURE_RESPONSE) + recorder: Final = RecordingLogger() + + try: + with pytest.raises(litellm.InternalServerError): + await route.invoke(provider, callbacks=[generic_api_export.logger, recorder], num_retries=0) + finally: + await drain_logging() + await wait_for_callback(route, recorder, outcome="failure") + + assert len(provider.requests) == 1 + assert len(generic_api_export.exports) == 1 + payload: Final = generic_api_export.exports[0].body + assert payload["status"] == "failure" + assert payload["call_type"] == route.call_type + assert payload["error_information"]["error_class"] == "InternalServerError" + assert payload["error_information"]["error_code"] == "500" + + +def test_prometheus_registry_restores_collectors_after_failure() -> None: + registry: Final = CollectorRegistry() + original: Final = Counter("original", "Original collector", registry=registry) + original.inc(2) + + def failing_test() -> None: + with isolated_prometheus_registry(registry): + assert registry.get_sample_value("original_total") is None + Counter("original", "Temporary replacement", registry=registry).inc(7) + Counter("temporary", "Temporary collector", registry=registry).inc() + raise RuntimeError("test failed") + + with pytest.raises(RuntimeError, match="test failed"): + failing_test() + assert registry.get_sample_value("original_total") == 2 + assert registry.get_sample_value("temporary_total") is None + registry.unregister(original) + assert registry.get_sample_value("original_total") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ALL_ROUTES, ids=route_id) +async def test_otel_emits_one_request_span_on_success( + route: Route, provider: RecordingServer, otel: OtelHarness +) -> None: + await route.invoke(provider, callbacks=[otel.logger]) + spans: Final = await otel.wait_for_spans() + assert len(spans) == 1 + span: Final = spans[0] + assert span.status.status_code is StatusCode.OK + assert span.attributes["llm.request.type"] == route.call_type + assert span.attributes["gen_ai.request.model"] == route.provider_model + assert json.loads(span.attributes["hidden_params"])["response_cost"] == pytest.approx(route.expected_cost) + + +@pytest.mark.asyncio +async def test_otel_stream_span_appears_only_after_exhaustion( + otel: OtelHarness, recording_server: RecordingServer +) -> None: + recording_server.default_response = ResponseSpec(body=None, events=MESSAGES_EVENTS) + stream: Final = await MESSAGES_STREAM.open_stream(recording_server, callbacks=[otel.logger]) + await drain_logging() + assert otel.spans() == () + chunks: Final = [chunk async for chunk in stream] + assert chunks + assert len(await otel.wait_for_spans()) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", (OCR_SYNC, OCR_ASYNC), ids=route_id) +async def test_otel_emits_one_error_span_on_provider_failure( + route: Route, provider: RecordingServer, otel: OtelHarness +) -> None: + provider.enqueue(FAILURE_RESPONSE) + with pytest.raises(litellm.InternalServerError): + await route.invoke(provider, callbacks=[otel.logger]) + spans: Final = await otel.wait_for_spans() + assert len(spans) == 1 + assert spans[0].status.status_code is StatusCode.ERROR + exception_events: Final = [event for event in spans[0].events if event.name == "exception"] + assert len(exception_events) == 1 + assert "InternalServerError" in exception_events[0].attributes["exception.type"] + assert spans[0].attributes["llm.request.type"] == route.call_type + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ALL_ROUTES, ids=route_id) +async def test_otel_and_custom_logger_observe_same_standard_logging_object( + route: Route, provider: RecordingServer, otel: OtelHarness +) -> None: + recorder: Final = RecordingLogger() + await route.invoke(provider, callbacks=[otel.logger, recorder]) + payload: Final = (await wait_for_callback(route, recorder))[0].kwargs["standard_logging_object"] + span: Final = (await otel.wait_for_spans())[0] + assert payload["call_type"] == route.call_type + assert payload["model"] == route.provider_model + assert json.loads(span.attributes["hidden_params"]) == payload["hidden_params"] + assert span.attributes["llm.request.type"] == payload["call_type"] + assert span.attributes["litellm.provider.model"] == payload["model"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ASYNC_ROUTES, ids=route_id) +async def test_prometheus_counts_one_successful_request( + route: Route, provider: RecordingServer, prometheus: PrometheusLogger +) -> None: + before: Final = metric_value("litellm_requests_metric_total", model=route.provider_model) + recorder: Final = RecordingLogger() + await route.invoke(provider, callbacks=[prometheus, recorder]) + await wait_for_callback(route, recorder) + assert metric_value("litellm_requests_metric_total", model=route.provider_model) == before + 1 + assert metric_value("litellm_llm_api_failed_requests_metric_total", model=route.provider_model) == 0 + + +@pytest.mark.asyncio +async def test_prometheus_counts_tokens_from_messages_usage( + recording_server: RecordingServer, prometheus: PrometheusLogger +) -> None: + recording_server.default_response = ResponseSpec(body=MESSAGES_ROUTE.provider_response) + recorder: Final = RecordingLogger() + await MESSAGES_ROUTE.invoke(recording_server, callbacks=[prometheus, recorder]) + await wait_for_callback(MESSAGES_ROUTE, recorder) + assert metric_value("litellm_input_tokens_metric_total", model=MESSAGES_ROUTE.provider_model) == 5 + assert metric_value("litellm_output_tokens_metric_total", model=MESSAGES_ROUTE.provider_model) == 4 + + +@pytest.mark.asyncio +async def test_prometheus_counts_one_failed_request( + otel: OtelHarness, prometheus: PrometheusLogger, recording_server: RecordingServer +) -> None: + recording_server.enqueue(FAILURE_RESPONSE) + with pytest.raises(litellm.InternalServerError): + await OCR_ASYNC.invoke(recording_server, callbacks=[prometheus, otel.logger]) + assert len(await otel.wait_for_spans()) == 1 + assert metric_value("litellm_llm_api_failed_requests_metric_total", model=OCR_ASYNC.provider_model) == 1 + assert metric_value("litellm_requests_metric_total", model=OCR_ASYNC.provider_model) == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", NON_STREAM_ASYNC_ROUTES, ids=route_id) +async def test_prometheus_by_string_name_is_initialized_once(route: Route, provider: RecordingServer) -> None: + provider.expected_requests = 2 + litellm.success_callback = ["prometheus"] # test-quality-ok: public registration; fixture restores globals + recorder: Final = RecordingLogger() + await route.invoke(provider, callbacks=[recorder]) + await route.invoke(provider, callbacks=[recorder]) + await wait_for_callback(route, recorder, count=2) + instances: Final = [cb for cb in litellm_logging._in_memory_loggers if isinstance(cb, PrometheusLogger)] # pyright: ignore[reportPrivateUsage] # string-name cache has no public accessor + assert len(instances) == 1 + assert metric_value("litellm_requests_metric_total", model=route.provider_model) == 2 + assert "prometheus" not in litellm.success_callback + assert instances[0] in litellm._async_success_callback # pyright: ignore[reportPrivateUsage] # callback registry has no public accessor + + +@pytest.mark.asyncio +async def test_sync_ocr_reaches_sync_hooks_only( + recording_server: RecordingServer, otel: OtelHarness, prometheus: PrometheusLogger +) -> None: + recording_server.default_response = ResponseSpec(body=OCR_SYNC.provider_response) + await OCR_SYNC.invoke(recording_server, callbacks=[otel.logger, prometheus]) + assert len(await otel.wait_for_spans()) == 1 + assert metric_value("litellm_requests_metric_total", model=OCR_SYNC.provider_model) == 0 diff --git a/tests/test_litellm_rust/integrations/test_guardrails.py b/tests/test_litellm_rust/integrations/test_guardrails.py new file mode 100644 index 00000000000..03d96efb768 --- /dev/null +++ b/tests/test_litellm_rust/integrations/test_guardrails.py @@ -0,0 +1,697 @@ +import asyncio +import gc +import json +import threading +import weakref +from collections.abc import AsyncIterator +from datetime import datetime +from typing import Final +from unittest.mock import MagicMock + +import pytest +from fastapi import HTTPException + +import litellm +from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException +from litellm.integrations.literal_ai import LiteralAILogger +from litellm.integrations.rubrik import RubrikLogger +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.guardrails.guardrail_hooks.azure.text_moderation import AzureContentSafetyTextModerationGuardrail +from litellm.proxy.guardrails.guardrail_hooks.crowdstrike_aidr.crowdstrike_aidr import CrowdStrikeAIDRHandler +from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail +from litellm.proxy.utils import ProxyLogging +from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks +from litellm.types.utils import CallTypes +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +from tests.test_litellm_rust.support.provenance import has_rust_response_marker +from tests.test_litellm_rust.conftest import Backend, isolated_backend +from tests.test_litellm_rust.support.requests import CHAT_MESSAGES, CHAT_MODEL, CHAT_RESPONSE, MESSAGES +from tests.test_litellm_rust.integrations import ( + MESSAGES_ROUTE, + MESSAGES_STREAM, + NON_STREAM_ASYNC_ROUTES, + GenericAPIExportHarness, + OtelHarness, + ReviewGuardrail, + Route, + provider_response, + purview_harness, + route_id, + wait_for_audits, + wait_for_callback, +) +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, recording_service + +pytestmark = pytest.mark.requires_rust_extension + +AZURE_MODERATION_ALLOW_RESPONSE: Final = { + "blocklistsMatch": [], + "categoriesAnalysis": [ + {"category": "Hate", "severity": 0}, + {"category": "Sexual", "severity": 0}, + {"category": "SelfHarm", "severity": 0}, + {"category": "Violence", "severity": 0}, + ], +} +AZURE_MODERATION_BLOCK_RESPONSE: Final = { + **AZURE_MODERATION_ALLOW_RESPONSE, + "categoriesAnalysis": [{"category": "Violence", "severity": 6}], +} +STREAM_SECRET: Final = "078-05-1120" +STREAM_MASK: Final = "[MASKED]" + + +class BufferedMaskingGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="stream-mask", event_hook=GuardrailEventHooks.post_call, default_on=True) + + async def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + response: AsyncIterator[object], + request_data: dict, + ) -> AsyncIterator[bytes]: + chunks: Final = tuple([chunk async for chunk in response]) + body: Final = b"".join(chunk for chunk in chunks if isinstance(chunk, bytes)) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"masked": True}, + request_data=request_data, + guardrail_status="success", + masked_entity_count={"TEST_IDENTIFIER": 1}, + ) + yield body.replace(STREAM_SECRET.encode(), STREAM_MASK.encode()) + + +class CleanupFailureGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="cleanup-failure", event_hook=GuardrailEventHooks.post_call, default_on=True) + + async def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + response: AsyncIterator[object], + request_data: dict, + ) -> AsyncIterator[object]: + try: + async for chunk in response: + yield chunk + finally: + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response="guardrail cleanup failed", + request_data=request_data, + guardrail_status="guardrail_failed_to_respond", + ) + raise RuntimeError("guardrail cleanup failed") + + +class FailingMaskingGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="mask-failure", event_hook=GuardrailEventHooks.post_call, default_on=True) + + async def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + response: AsyncIterator[object], + request_data: dict, + ) -> AsyncIterator[object]: + tuple([chunk async for chunk in response]) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response="stream masking failed", + request_data=request_data, + guardrail_status="guardrail_failed_to_respond", + ) + raise HTTPException( + status_code=500, + detail={"error": "stream masking failed", "guardrail": self.guardrail_name}, + ) + yield + + +def stream_events(text: str) -> tuple[tuple[str, object], ...]: + return tuple( + (event, {**data, "delta": {**data["delta"], "text": text}}) + if event == "content_block_delta" and isinstance(data, dict) + else (event, data) + for event, data in provider_response(MESSAGES_STREAM).events + ) + + +def streaming_logger(recorder: RecordingLogger, callbacks: list[object] | None = None) -> Logging: + terminal_callbacks: Final = [recorder, *(callbacks or [])] + return Logging( + model=MESSAGES_STREAM.provider_model, + messages=[{"role": "user", "content": "mask the response"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id="stream-mask", + function_id="stream-mask", + dynamic_async_success_callbacks=terminal_callbacks, + dynamic_async_failure_callbacks=terminal_callbacks, + ) + + +def arm_guarded_stream(stream: AsyncIterator[object], request_data: dict, logger: Logging) -> AsyncIterator[object]: + processor: Final = object.__new__(ProxyBaseLLMRequestProcessing) + processor.data = request_data + user: Final = UserAPIKeyAuth(request_route="/v1/messages") + processor._arm_deferred_stream_dispatch(stream, "anthropic_messages", user, logger) + return ProxyLogging(user_api_key_cache=MagicMock()).async_post_call_streaming_iterator_hook( + response=stream, + user_api_key_dict=user, + request_data=request_data, + ) + + +def azure_text_moderation(server: RecordingServer) -> AzureContentSafetyTextModerationGuardrail: + return AzureContentSafetyTextModerationGuardrail( + guardrail_name="azure-text-review", + api_key="test-azure-key", + api_base=server.base_url, + event_hook=GuardrailEventHooks.post_call, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ("python", "rust")) +async def test_stream_masking_reaches_client_loggers_and_exporters( + backend: Backend, + otel: OtelHarness, + generic_api_export: GenericAPIExportHarness, +) -> None: + async with isolated_backend(backend): + with recording_service() as provider: + wire: Final = b"".join( + f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in stream_events(STREAM_SECRET) + ) + split: Final = wire.index(STREAM_SECRET.encode()) + 4 + provider.default_response = ResponseSpec(body=None, chunks=(wire[:split], wire[split:])) + guardrail: Final = BufferedMaskingGuardrail() + recorder: Final = RecordingLogger() + logger: Final = streaming_logger(recorder, [otel.logger, generic_api_export.logger]) + request_data: Final = { + "model": MESSAGES_STREAM.provider_model, + "messages": MESSAGES, + "guardrails": [guardrail.guardrail_name], + "metadata": {}, + "litellm_logging_obj": logger, + } + litellm.callbacks.append(guardrail) + source: Final = await MESSAGES_STREAM.open_stream( + provider, + litellm_logging_obj=logger, + guardrails=[guardrail.guardrail_name], + ) + guarded: Final = arm_guarded_stream(source, request_data, logger) + + assert "async_log_success_event" not in recorder.names + payload: Final = b"".join([chunk async for chunk in guarded if isinstance(chunk, bytes)]) + events: Final = await recorder.wait_for_async("async_log_success_event") + spans: Final = await otel.wait_for_spans() + + assert STREAM_MASK.encode() in payload + assert STREAM_SECRET.encode() not in payload + assert payload.count(b"event: message_stop") == 1 + assert len(events) == 1 + assert events[0].response.choices[0].message.content == STREAM_MASK + assert STREAM_SECRET not in json.dumps(events[0].kwargs, default=str) + assert STREAM_SECRET not in json.dumps(dict(spans[0].attributes), default=str) + assert STREAM_MASK in json.dumps(dict(spans[0].attributes), default=str) + async with asyncio.timeout(10): + while not generic_api_export.exports: + await asyncio.sleep(0.01) + exported: Final = json.dumps(generic_api_export.exports[0].body, default=str) + assert STREAM_SECRET not in exported + assert STREAM_MASK in exported + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ("python", "rust")) +async def test_upstream_stream_failure_is_not_masked_by_guardrail_cleanup( + backend: Backend, +) -> None: + async with isolated_backend(backend): + with recording_service() as provider: + provider.default_response = ResponseSpec( + body=None, + events=( + provider_response(MESSAGES_STREAM).events[0], + ( + "error", + { + "type": "error", + "error": {"type": "overloaded_error", "message": "upstream overloaded"}, + }, + ), + ), + ) + guardrail: Final = CleanupFailureGuardrail() + recorder: Final = RecordingLogger() + logger: Final = streaming_logger(recorder) + request_data: Final = { + "model": MESSAGES_STREAM.provider_model, + "messages": MESSAGES, + "guardrails": [guardrail.guardrail_name], + "metadata": {}, + "litellm_logging_obj": logger, + } + litellm.callbacks.append(guardrail) + source: Final = await MESSAGES_STREAM.open_stream( + provider, + litellm_logging_obj=logger, + guardrails=[guardrail.guardrail_name], + ) + guarded: Final = arm_guarded_stream(source, request_data, logger) + + with pytest.raises(litellm.APIError) as caught: + async for _ in guarded: + pass + guardrail_info: Final = request_data["metadata"]["standard_logging_guardrail_information"] + + assert "guardrail cleanup failed" not in str(caught.value) + assert request_data["metadata"]["stream_guardrail_cleanup_error"] == { + "type": "RuntimeError", + "message": "guardrail cleanup failed", + } + if backend == "rust": + events: Final = await recorder.wait_for_async("async_log_failure_event") + assert len(events) == 1 + assert "guardrail cleanup failed" not in str(events[0].kwargs["exception"]) + else: + assert "async_log_success_event" not in recorder.names + assert guardrail_info[-1]["guardrail_status"] == "guardrail_failed_to_respond" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ("python", "rust")) +async def test_stream_masking_failure_uses_guardrail_error_and_logs_once(backend: Backend) -> None: + async with isolated_backend(backend): + with recording_service() as provider: + provider.default_response = ResponseSpec(body=None, events=stream_events(STREAM_SECRET)) + guardrail: Final = FailingMaskingGuardrail() + recorder: Final = RecordingLogger() + logger: Final = streaming_logger(recorder) + request_data: Final = { + "model": MESSAGES_STREAM.provider_model, + "messages": MESSAGES, + "guardrails": [guardrail.guardrail_name], + "metadata": {}, + "litellm_logging_obj": logger, + } + litellm.callbacks.append(guardrail) + source: Final = await MESSAGES_STREAM.open_stream( + provider, + litellm_logging_obj=logger, + guardrails=[guardrail.guardrail_name], + ) + + with pytest.raises(HTTPException) as caught: + async for _ in arm_guarded_stream(source, request_data, logger): + pass + await logger.async_failure_handler(caught.value, "stream masking failed") + events: Final = await recorder.wait_for_async("async_log_failure_event") + + assert caught.value.status_code == 500 + assert caught.value.detail["error"] == "stream masking failed" + assert caught.value.detail["guardrail"] == guardrail.guardrail_name + assert len(events) == 1 + assert events[0].kwargs["exception"] is caught.value + assert "async_log_success_event" not in recorder.names + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ("python", "rust")) +async def test_guarded_stream_cancel_drains_and_releases_roots(backend: Backend) -> None: + class NamesOnlyLogger(RecordingLogger): + def _record(self, name: str, kwargs: object = None, response: object = None) -> None: + super()._record(name) + + class Root: + pass + + async with isolated_backend(backend): + with recording_service() as provider: + release: Final = threading.Event() + wire: Final = tuple( + f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in stream_events(STREAM_SECRET) + ) + provider.enqueue(ResponseSpec(body=None, chunks=wire, release=release)) + guardrail: Final = BufferedMaskingGuardrail() + recorder: Final = NamesOnlyLogger() + logger: Final = streaming_logger(recorder) + root = Root() + reference: Final = weakref.ref(root) + request_data: Final = { + "model": MESSAGES_STREAM.provider_model, + "messages": MESSAGES, + "guardrails": [guardrail.guardrail_name], + "metadata": {}, + "litellm_logging_obj": logger, + } + litellm.callbacks.append(guardrail) + source: Final = await MESSAGES_STREAM.open_stream( + provider, + litellm_logging_obj=logger, + guardrails=[guardrail.guardrail_name], + metadata={"retained": root}, + ) + first_chunk: Final = asyncio.Event() + + async def observed_source() -> AsyncIterator[bytes]: + async for chunk in source: + first_chunk.set() + yield chunk + + guarded = arm_guarded_stream(observed_source(), request_data, logger) + del root + task: Final = asyncio.create_task(anext(guarded)) + await provider.wait_for_requests(1) + await asyncio.wait_for(first_chunk.wait(), 5) + task.cancel() + + try: + with pytest.raises(asyncio.CancelledError): + await task + await guarded.aclose() + release.set() + events: Final = await recorder.wait_for_async("async_log_success_event") + del guarded + del source + gc.collect() + + assert len(events) == 1 + assert "async_log_failure_event" not in recorder.names + del events + del task + del request_data + del logger + gc.collect() + assert reference() is None + finally: + release.set() + + +@pytest.mark.asyncio +async def test_crowdstrike_redaction_reaches_native_chat_provider_and_exporter( + generic_api_export: GenericAPIExportHarness, +) -> None: + with recording_service() as provider, recording_service() as guard: + provider.default_response = ResponseSpec(body=CHAT_RESPONSE) + guard.default_response = ResponseSpec( + body={ + "result": { + "blocked": False, + "transformed": True, + "guard_output": { + "messages": [ + {"role": "system", "content": "Keep the answer short"}, + {"role": "user", "content": "Employee SSN: "}, + ] + }, + } + } + ) + guardrail: Final = CrowdStrikeAIDRHandler( + guardrail_name="crowdstrike-redaction", + api_key="test-crowdstrike-key", + api_base=guard.base_url, + event_hook=GuardrailEventHooks.pre_call, + ) + recorder: Final = RecordingLogger() + litellm.callbacks.append(guardrail) + + response: Final = await litellm.acompletion( + model=CHAT_MODEL, + messages=CHAT_MESSAGES, + api_key="test-key", + api_base=provider.base_url, + callbacks=[generic_api_export.logger, recorder], + guardrails=[guardrail.guardrail_name], + ) + await recorder.wait_for_async("async_log_success_event") + + provider_messages: Final = provider.requests[0].body["messages"] + guard_messages: Final = guard.requests[0].body["guard_input"]["messages"] + exported_messages: Final = generic_api_export.exports[0].body["messages"] + assert response.choices[0].message.content == "Handled safely" + assert guard_messages == [CHAT_MESSAGES[0], CHAT_MESSAGES[3]] + assert provider_messages == [ + {"role": "user", "content": [{"type": "text", "text": "Earlier safe question"}]}, + {"role": "assistant", "content": [{"type": "text", "text": "Earlier safe answer"}]}, + {"role": "user", "content": [{"type": "text", "text": "Employee SSN: "}]}, + ] + assert provider.requests[0].body["system"] == [{"type": "text", "text": "Keep the answer short"}] + assert exported_messages == CHAT_MESSAGES + assert has_rust_response_marker(response) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ("python", "rust")) +async def test_rubrik_block_preserves_context_for_error_exporters( + backend: Backend, + monkeypatch: pytest.MonkeyPatch, + generic_api_export: GenericAPIExportHarness, +) -> None: + monkeypatch.setenv("RUBRIK_BATCH_SIZE", "1") + async with isolated_backend(backend): + with recording_service() as provider, recording_service() as rubrik_service: + provider.default_response = ResponseSpec(body=CHAT_RESPONSE) + rubrik_service.expected_requests = None + rubrik_service.default_response = ResponseSpec( + body={ + "choices": [ + { + "message": { + "role": "assistant", + "content": "Response blocked by policy", + "tool_calls": [], + } + } + ] + } + ) + rubrik: Final = RubrikLogger( + api_key="test-rubrik-key", + api_base=rubrik_service.base_url, + guardrail_name="rubrik-block", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + recorder: Final = RecordingLogger() + litellm.callbacks.append(rubrik) + + try: + with pytest.raises(ModifyResponseException, match="Response blocked by policy") as raised: + await litellm.acompletion( + model=CHAT_MODEL, + messages=CHAT_MESSAGES, + api_key="test-key", + api_base=provider.base_url, + callbacks=[generic_api_export.logger, recorder], + guardrails=[rubrik.guardrail_name], + ) + assert has_rust_response_marker(raised.value.original_response) is (backend == "rust") + await drain_logging() + await rubrik.flush_queue() + + moderation_requests: Final = tuple( + request for request in rubrik_service.requests if request.path == "/v1/after_completion/openai/v1" + ) + batch_requests: Final = tuple( + request for request in rubrik_service.requests if request.path == "/v1/litellm/batch" + ) + assert len(provider.requests) == 1 + assert len(moderation_requests) == 1 + assert moderation_requests[0].body["request"]["messages"] == CHAT_MESSAGES + assert moderation_requests[0].body["response"]["choices"][0]["message"]["content"] == "Handled safely" + assert len(batch_requests) == 1 + assert "Response blocked by policy" in json.dumps(batch_requests[0].body) + assert len(generic_api_export.exports) == 1 + exported: Final = generic_api_export.exports[0].body + assert exported["status"] == "failure" + assert exported["error_information"]["error_class"] == "ModifyResponseException" + assert "Response blocked by policy" in exported["error_information"]["error_message"] + finally: + await rubrik.aclose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ("python", "rust")) +async def test_purview_audit_retains_payload_after_sdk_response( + backend: Backend, + generic_api_export: GenericAPIExportHarness, +) -> None: + async with isolated_backend(backend): + with recording_service() as provider: + provider.default_response = ResponseSpec(body=CHAT_RESPONSE) + purview: Final = purview_harness("purview-audit") + recorder: Final = RecordingLogger() + litellm.callbacks.append(purview.guardrail) + request_task: Final = asyncio.create_task( + litellm.acompletion( + model=CHAT_MODEL, + messages=CHAT_MESSAGES, + api_key="test-key", + api_base=provider.base_url, + metadata={"user_api_key_user_id": "audit-user"}, + callbacks=[generic_api_export.logger, recorder], + guardrails=[purview.guardrail.guardrail_name], + ) + ) + try: + accepted: Final = await asyncio.to_thread(purview.graph.accepted.wait, 3) + assert accepted + response: Final = await asyncio.wait_for(asyncio.shield(request_task), 2) + assert response.choices[0].message.content == "Handled safely" + assert len(await wait_for_audits(purview.graph, count=1)) == 1 + + purview.graph.release.set() + await recorder.wait_for_async("async_log_success_event") + audits: Final = await wait_for_audits(purview.graph) + activities: Final = tuple( + post.body["contentToProcess"]["activityMetadata"]["activity"] for post in audits + ) + assert activities == ("uploadText", "downloadText") + assert CHAT_MESSAGES[-1]["content"] in json.dumps(audits[0].body) + assert "Handled safely" in json.dumps(audits[1].body) + assert generic_api_export.exports[0].body["status"] == "success" + assert has_rust_response_marker(response) is (backend == "rust") + finally: + purview.graph.release.set() + if not request_task.done(): + request_task.cancel() + await asyncio.gather(request_task, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ("python", "rust")) +async def test_purview_sync_audit_runs_without_caller_event_loop( + backend: Backend, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("LITERAL_BATCH_SIZE", "1") + async with isolated_backend(backend): + with recording_service() as provider, recording_service() as sink: + provider.default_response = ResponseSpec(body=CHAT_RESPONSE) + purview: Final = purview_harness("purview-sync-audit") + exporter: Final = LiteralAILogger(literalai_api_key="test-key", literalai_api_url=sink.base_url) + litellm.callbacks.append(purview.guardrail) + + try: + response: Final = await asyncio.to_thread( + litellm.completion, + model=CHAT_MODEL, + messages=CHAT_MESSAGES, + api_key="test-key", + api_base=provider.base_url, + metadata={"user_api_key_user_id": "audit-user"}, + callbacks=[exporter], + guardrails=[purview.guardrail.guardrail_name], + ) + accepted: Final = await asyncio.to_thread(purview.graph.accepted.wait, 3) + assert accepted + assert response.choices[0].message.content == "Handled safely" + assert len(await wait_for_audits(purview.graph, count=1)) == 1 + + purview.graph.release.set() + audits: Final = await wait_for_audits(purview.graph) + assert CHAT_MESSAGES[-1]["content"] in json.dumps(audits[0].body) + assert "Handled safely" in json.dumps(audits[1].body) + assert sink.requests[0].path == "/api/graphql" + assert "Handled safely" in json.dumps(sink.requests[0].body["variables"]["generation_0"]) + assert has_rust_response_marker(response) is (backend == "rust") + finally: + purview.graph.release.set() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", NON_STREAM_ASYNC_ROUTES, ids=route_id) +async def test_content_filter_post_call_blocks_provider_response(route: Route, provider: RecordingServer) -> None: + guardrail: Final = ContentFilterGuardrail( + guardrail_name="enforced-content-review", + event_hook=GuardrailEventHooks.post_call, + blocked_words=[BlockedWord(keyword=route.expected_text, action=ContentFilterAction.BLOCK)], + ) + litellm.callbacks.append(guardrail) + with pytest.raises(HTTPException, match="Content blocked") as blocked: + await route.invoke(provider, guardrails=["enforced-content-review"]) + assert blocked.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_azure_text_moderation_allows_messages_response_over_http( + recording_server: RecordingServer, otel: OtelHarness +) -> None: + recording_server.expected_requests = None + recording_server.default_response = ResponseSpec(body=AZURE_MODERATION_ALLOW_RESPONSE) + recording_server.enqueue(provider_response(MESSAGES_ROUTE)) + guardrail: Final = azure_text_moderation(recording_server) + recorder: Final = RecordingLogger() + litellm.callbacks.append(guardrail) + + response: Final = await MESSAGES_ROUTE.invoke( + recording_server, callbacks=[otel.logger, recorder], guardrails=[guardrail.guardrail_name] + ) + await wait_for_callback(MESSAGES_ROUTE, recorder) + + assert len(recording_server.requests) == 2 + moderation_request: Final = recording_server.requests[1] + assert moderation_request.path == "/contentsafety/text:analyze?api-version=2024-09-01" + assert moderation_request.headers["ocp-apim-subscription-key"] == "test-azure-key" + assert moderation_request.body == { + "text": MESSAGES_ROUTE.expected_text, + "categories": ["Hate", "Sexual", "SelfHarm", "Violence"], + "blocklistNames": None, + "haltOnBlocklistHit": False, + "outputType": "FourSeverityLevels", + } + assert response["content"][0]["text"] == MESSAGES_ROUTE.expected_text + assert len(await otel.wait_for_spans()) == 1 + + +@pytest.mark.asyncio +async def test_azure_text_moderation_blocks_messages_response_over_http(recording_server: RecordingServer) -> None: + recording_server.expected_requests = None + recording_server.default_response = ResponseSpec(body=AZURE_MODERATION_BLOCK_RESPONSE) + recording_server.enqueue(provider_response(MESSAGES_ROUTE)) + guardrail: Final = azure_text_moderation(recording_server) + litellm.callbacks.append(guardrail) + with pytest.raises(HTTPException, match="Violence crossed severity 2") as blocked: + await MESSAGES_ROUTE.invoke(recording_server, guardrails=[guardrail.guardrail_name]) + assert blocked.value.status_code == 400 + assert recording_server.requests[1].body["text"] == MESSAGES_ROUTE.expected_text + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", NON_STREAM_ASYNC_ROUTES, ids=route_id) +async def test_post_call_guardrail_replacement_is_what_loggers_see( + route: Route, provider: RecordingServer, otel: OtelHarness +) -> None: + async def review(response: object) -> object: + match route.name: + case "ocr-async": + return response.model_copy( + update={"pages": [response.pages[0].model_copy(update={"markdown": "Reviewed OCR"})]} + ) + case _: + return {**response, "content": [{"type": "text", "text": "Reviewed Messages"}]} + + guardrail: Final = ReviewGuardrail(review) + recorder: Final = RecordingLogger() + litellm.callbacks.append(guardrail) + response: Final = await route.invoke(provider, callbacks=[otel.logger, recorder], guardrails=["rust-review"]) + + assert guardrail.call_types == [CallTypes(route.call_type)] + event: Final = (await wait_for_callback(route, recorder))[0] + span: Final = (await otel.wait_for_spans())[0] + match route.name: + case "ocr-async": + assert response.pages[0].markdown == "Reviewed OCR" + assert event.response.pages[0].markdown == "Reviewed OCR" + case _: + assert response["content"][0]["text"] == "Reviewed Messages" + assert event.response.choices[0].message.content == "Reviewed Messages" + assert "Reviewed Messages" in json.dumps(dict(span.attributes)) + assert "guardrails" not in provider.requests[0].body diff --git a/tests/test_litellm_rust/messages/__init__.py b/tests/test_litellm_rust/messages/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/test_litellm_rust/messages/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/test_litellm_rust/messages/test_callback_mutation.py b/tests/test_litellm_rust/messages/test_callback_mutation.py new file mode 100644 index 00000000000..aee034f4b6b --- /dev/null +++ b/tests/test_litellm_rust/messages/test_callback_mutation.py @@ -0,0 +1,682 @@ +import asyncio +import copy +import json +import threading +from collections.abc import Mapping +from typing import Final + +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.types.utils import CallTypes +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +from tests.test_litellm_rust.support.provenance import has_rust_response_marker +from tests.test_litellm_rust.support.requests import ( + MESSAGES, + MESSAGES_EVENTS, + MESSAGES_MODEL, + MESSAGES_RESPONSE, + request_body, + request_headers, +) +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.fixture +def messages_server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) + return recording_server + + +async def call_messages(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object): + return await litellm.anthropic.messages.acreate( + model=MESSAGES_MODEL, + messages=MESSAGES, + max_tokens=64, + api_key="test-key", + api_base=server.base_url, + callbacks=callbacks, + **kwargs, + ) + + +@pytest.mark.asyncio +async def test_messages_pre_call_receives_expected_provider_request(messages_server: RecordingServer) -> None: + observations: Final = [] + + class Observe(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + observations.append((model, messages, copy.deepcopy(kwargs["additional_args"]))) + + await call_messages(messages_server, [Observe()]) + + assert len(observations) == 1 + model, messages, additional_args = observations[0] + assert model == "claude-sonnet-4-5-20250929" + assert messages == MESSAGES + assert additional_args["api_base"] == f"{messages_server.base_url}/v1/messages" + assert additional_args["complete_input_dict"] == { + "model": "claude-sonnet-4-5-20250929", + "messages": MESSAGES, + "max_tokens": 64, + "stream": False, + } + assert additional_args["headers"]["x-api-key"] == "test-key" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("raise_after_edit", [False, True]) +@pytest.mark.parametrize("native", [False, True]) +async def test_messages_pre_call_edits_reach_later_callbacks_and_provider( + messages_server: RecordingServer, raise_after_edit: bool, native: bool, stream: bool +) -> None: + litellm.rust(native) + observed: Final = [] + + class Edit(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + request_body(kwargs)["temperature"] = 0.25 + request_headers(kwargs)["x-audit-tag"] = "reviewed" + if raise_after_edit: + raise RuntimeError("audit exporter unavailable") + + class Observe(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + observed.append((copy.deepcopy(request_body(kwargs)), dict(request_headers(kwargs)))) + + if stream: + messages_server.enqueue(ResponseSpec(body=None, events=MESSAGES_EVENTS)) + response: Final = await call_messages(messages_server, [Edit(), Observe()], stream=stream) + if stream: + async for _ in response: + pass + + assert observed[0][0]["temperature"] == 0.25 + assert observed[0][1]["x-audit-tag"] == "reviewed" + assert "temperature" not in messages_server.requests[0].body + assert messages_server.requests[0].headers["x-audit-tag"] == "reviewed" + + +@pytest.mark.asyncio +async def test_messages_pre_call_rebinding_does_not_replace_inflight_request( + messages_server: RecordingServer, +) -> None: + observed: Final = [] + + class Rebind(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + kwargs["additional_args"]["complete_input_dict"] = {"replacement": True} + + class Observe(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + observed.append(request_body(kwargs)) + + await call_messages(messages_server, [Rebind(), Observe()]) + + assert observed == [{"replacement": True}] + assert messages_server.requests[0].body["messages"] == MESSAGES + + +@pytest.mark.asyncio +async def test_messages_pre_call_state_reaches_terminal_callbacks(messages_server: RecordingServer) -> None: + token: Final = object() + observed: Final = [] + finished: Final = asyncio.Event() + + class Stash(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + kwargs["test-token"] = token + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + observed.append(kwargs["test-token"]) + finished.set() + + await call_messages(messages_server, [Stash()]) + await asyncio.wait_for(finished.wait(), timeout=10) + + assert observed == [token] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native", [False, True]) +@pytest.mark.parametrize("provider", ["anthropic", "azure_ai"]) +@pytest.mark.parametrize("raise_after_edit", [False, True]) +async def test_messages_retained_aliases_preserve_identity_and_snapshot_timing( + messages_server: RecordingServer, native: bool, provider: str, raise_after_edit: bool +) -> None: + litellm.rust(native) + retained: Final = [] + observed: Final = [] + + class Retain(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + body = request_body(kwargs) + message = kwargs["messages"][0] + assert body["messages"][0] is message + assert body["messages"][0]["content"] is message["content"] + assert body["messages"][0]["content"][0] is message["content"][0] + retained.append(message["content"][0]) + + class Edit(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + retained[0]["text"] = "changed through retained reference" + if raise_after_edit: + raise RuntimeError("export failed after mutation") + + class Observe(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + block = request_body(kwargs)["messages"][0]["content"][0] + observed.append((block is retained[0], block["text"])) + + response: Final = await litellm.anthropic.messages.acreate( + model=MESSAGES_MODEL.replace("anthropic/", f"{provider}/"), + messages=[{"role": "user", "content": [{"type": "text", "text": "original"}]}], + max_tokens=64, + api_key="test-key", + api_base=messages_server.base_url, + callbacks=[Retain(), Edit(), Observe()], + ) + assert observed == [(True, "changed through retained reference")] + assert retained[0]["text"] == "changed through retained reference" + assert messages_server.requests[0].body["messages"][0]["content"][0]["text"] == "original" + if native: + assert response["_hidden_params"]["additional_headers"]["x-litellm-rust"] == "true" + + +@pytest.mark.asyncio +async def test_messages_callbacks_run_once(messages_server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + + await call_messages(messages_server, [recorder]) + await recorder.wait_for_async("async_log_success_event") + + assert recorder.names.count("log_pre_api_call") == 1 + assert recorder.names.count("async_logging_hook") == 1 + assert recorder.names.count("async_log_success_event") == 1 + assert "log_failure_event" not in recorder.names + assert "async_log_failure_event" not in recorder.names + + +@pytest.mark.asyncio +async def test_messages_unconsumed_stream_close_drains_and_logs_success_once(messages_server: RecordingServer) -> None: + messages_server.enqueue(ResponseSpec(body=None, events=MESSAGES_EVENTS)) + recorder: Final = RecordingLogger() + stream: Final = await call_messages(messages_server, [recorder], stream=True) + assert "async_log_success_event" not in recorder.names + await stream.aclose() + await stream.aclose() + await recorder.wait_for_async("async_log_success_event") + assert recorder.names.count("async_log_success_event") == 1 + assert "async_log_failure_event" not in recorder.names + + +@pytest.mark.asyncio +async def test_messages_logging_drain_waits_for_suspended_callback(messages_server: RecordingServer) -> None: + started: Final = asyncio.Event() + release: Final = asyncio.Event() + finished: Final = asyncio.Event() + + class SuspendedLogger(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + started.set() + await release.wait() + finished.set() + + await call_messages(messages_server, [SuspendedLogger()]) + draining: Final = asyncio.create_task(drain_logging()) + try: + await asyncio.wait_for(started.wait(), timeout=10) + await asyncio.sleep(0) + assert not draining.done() + assert not finished.is_set() + finally: + release.set() + await asyncio.wait_for(draining, timeout=10) + + assert finished.is_set() + + +@pytest.mark.asyncio +async def test_messages_failure_callbacks_receive_original_provider_error(messages_server: RecordingServer) -> None: + messages_server.default_response = ResponseSpec(body={"error": {"message": "provider unavailable"}}, status=500) + messages_server.expected_requests = None + recorder: Final = RecordingLogger() + + with pytest.raises(litellm.InternalServerError) as caught: + await call_messages(messages_server, [recorder]) + events: Final = await recorder.wait_for_async("async_log_failure_event") + + assert len(events) == 1 + assert events[0].call_type == "anthropic_messages" + assert events[0].kwargs["exception"] is caught.value + assert len(messages_server.requests) == 1 + assert "async_log_success_event" not in recorder.names + + +@pytest.mark.asyncio +async def test_messages_success_callbacks_receive_expected_context(messages_server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + + await call_messages( + messages_server, + [recorder], + litellm_call_id="messages-success", + metadata={"source": "callback-test"}, + ) + async_events: Final = await recorder.wait_for_async("async_log_success_event") + + assert len(async_events) == 1 + event: Final = async_events[0] + assert event.call_type == "anthropic_messages" + assert event.kwargs["litellm_call_id"] == "messages-success" + assert event.kwargs["litellm_params"]["metadata"]["source"] == "callback-test" + assert event.response.choices[0].message.content == "Hello from native Messages" + + +@pytest.mark.asyncio +async def test_messages_pre_call_runs_in_callers_execution_context(messages_server: RecordingServer) -> None: + caller_thread: Final = threading.current_thread() + observations: Final = [] + + class Observe(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + observations.append((asyncio.get_running_loop(), threading.current_thread())) + + caller_loop: Final = asyncio.get_running_loop() + await call_messages(messages_server, [Observe()]) + + assert observations == [(caller_loop, caller_thread)] + + +@pytest.mark.asyncio +async def test_messages_stream_logs_success_after_exhaustion(messages_server: RecordingServer) -> None: + messages_server.enqueue(ResponseSpec(body=None, events=MESSAGES_EVENTS)) + recorder: Final = RecordingLogger() + stream: Final = await call_messages(messages_server, [recorder], stream=True) + + assert "async_log_success_event" not in recorder.names + chunks: Final = [chunk async for chunk in stream] + events: Final = await recorder.wait_for_async("async_log_success_event") + + assert chunks + assert len(events) == 1 + assert "async_log_stream_event" not in recorder.names + assert events[0].kwargs["complete_streaming_response"] is not None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +async def test_messages_compression_hook_replaces_messages_sent_to_provider( + messages_server: RecordingServer, native: bool, stream: bool +) -> None: + litellm.rust(native) + compressed_messages: Final = [{"role": "user", "content": "Compressed context"}] + call_types: Final = [] + recorder: Final = RecordingLogger() + + class CompressMessages(CustomLogger): + async def async_pre_call_deployment_hook(self, kwargs, call_type): + assert not messages_server.requests + call_types.append(call_type) + return {**kwargs, "messages": compressed_messages} + + litellm.callbacks.append(CompressMessages()) + + if stream: + messages_server.enqueue(ResponseSpec(body=None, events=MESSAGES_EVENTS)) + response: Final = await call_messages(messages_server, [recorder], stream=stream) + assert has_rust_response_marker(response) is native + if stream: + assert [chunk async for chunk in response] + await recorder.wait_for_async("async_log_success_event") + + assert call_types == [CallTypes.anthropic_messages] + assert messages_server.requests[0].body["messages"] == compressed_messages + pre_calls: Final = tuple(event for event in recorder.events if event.name == "log_pre_api_call") + assert len(pre_calls) == 1 + assert request_body(pre_calls[0].kwargs)["messages"] == compressed_messages + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +async def test_messages_deployment_rejection_prevents_provider_call( + messages_server: RecordingServer, native: bool, stream: bool +) -> None: + litellm.rust(native) + messages_server.expected_requests = 0 + calls: Final[list[CallTypes | None]] = [] + recorder: Final = RecordingLogger() + + class Reject(CustomLogger): + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: CallTypes | None + ) -> dict[str, object]: + calls.append(call_type) + raise litellm.BadRequestError( + message="deployment policy rejected request", model=MESSAGES_MODEL, llm_provider="anthropic" + ) + + litellm.callbacks.append(Reject()) + + with pytest.raises(litellm.BadRequestError, match="deployment policy rejected request"): + await call_messages(messages_server, [recorder], stream=stream) + + assert calls == [CallTypes.anthropic_messages] + assert not messages_server.requests + assert "log_pre_api_call" not in recorder.names + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +async def test_messages_provider_open_failure_notifies_deployment_once_and_preserves_error( + messages_server: RecordingServer, native: bool, stream: bool +) -> None: + litellm.rust(native) + messages_server.enqueue(ResponseSpec(body={"error": {"message": "provider unavailable"}}, status=500)) + calls: Final[list[tuple[CallTypes | None, Exception, int]]] = [] + + class FailingObserver(CustomLogger): + async def async_post_call_failure_deployment_hook( + self, + request_data: Mapping[str, object], + exception: Exception, + call_type: CallTypes | None, + fallback_depth: int | None = None, + ) -> None: + calls.append((call_type, exception, len(messages_server.requests))) + raise RuntimeError("deployment observer unavailable") + + litellm.callbacks.append(FailingObserver()) + recorder: Final = RecordingLogger() + + with pytest.raises(litellm.InternalServerError) as raised: + await call_messages(messages_server, [recorder], stream=stream) + await recorder.wait_for_async("async_log_failure_event") + + assert len(calls) == 1 + call_type, exception, request_count = calls[0] + assert call_type == CallTypes.anthropic_messages + assert isinstance(exception, litellm.InternalServerError) + assert exception.status_code == raised.value.status_code == 500 + assert request_count == len(messages_server.requests) == 1 + assert recorder.names.count("async_log_failure_event") == 1 + assert "async_log_success_event" not in recorder.names + + +@pytest.mark.asyncio +async def test_messages_post_call_guardrail_replacement_reaches_caller_and_logging( + messages_server: RecordingServer, +) -> None: + recorder: Final = RecordingLogger() + call_types: Final = [] + + class ReviewResponse(CustomLogger): + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + call_types.append(call_type) + response["content"][0]["text"] = "Reviewed response" + return response + + litellm.callbacks.append(ReviewResponse()) + + response: Final = await call_messages(messages_server, [recorder]) + events: Final = await recorder.wait_for_async("async_log_success_event") + + assert call_types == [CallTypes.anthropic_messages] + assert response["content"][0]["text"] == "Reviewed response" + assert events[0].response.choices[0].message.content == "Reviewed response" + + +@pytest.mark.asyncio +async def test_messages_logging_hook_replacement_reaches_later_loggers_only(messages_server: RecordingServer) -> None: + observations: Final = [] + exported: Final = asyncio.Event() + + class RecordGuardrailVerdict(CustomLogger): + async def async_logging_hook(self, kwargs, result, call_type): + observations.append("guardrail") + return {**kwargs, "guardrail-verdict": "allowed"}, result + + class ExportLog(CustomLogger): + async def async_logging_hook(self, kwargs, result, call_type): + return kwargs, result + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + observations.append(("export", kwargs["guardrail-verdict"])) + exported.set() + + response: Final = await call_messages(messages_server, [RecordGuardrailVerdict(), ExportLog()]) + await asyncio.wait_for(exported.wait(), timeout=10) + + assert observations == ["guardrail", ("export", "allowed")] + assert response["content"][0]["text"] == "Hello from native Messages" + + +@pytest.mark.asyncio +async def test_messages_success_callback_failure_does_not_skip_later_loggers( + messages_server: RecordingServer, +) -> None: + recorder: Final = RecordingLogger() + + class UnavailableExporter(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + raise RuntimeError("exporter unavailable") + + response: Final = await call_messages(messages_server, [UnavailableExporter(), recorder]) + events: Final = await recorder.wait_for_async("async_log_success_event") + + assert response["content"][0]["text"] == "Hello from native Messages" + assert len(events) == 1 + assert "async_log_failure_event" not in recorder.names + + +@pytest.mark.asyncio +async def test_messages_stream_success_callback_failure_does_not_skip_later_loggers( + messages_server: RecordingServer, +) -> None: + messages_server.enqueue(ResponseSpec(body=None, events=MESSAGES_EVENTS)) + recorder: Final = RecordingLogger() + + class UnavailableExporter(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + raise RuntimeError("exporter unavailable") + + stream: Final = await call_messages(messages_server, [UnavailableExporter(), recorder], stream=True) + chunks: Final = [chunk async for chunk in stream] + events: Final = await recorder.wait_for_async("async_log_success_event") + + assert chunks + assert len(events) == 1 + assert events[0].response.choices[0].message.content == "Hello from native Messages" + assert "async_log_failure_event" not in recorder.names + + +@pytest.mark.asyncio +async def test_messages_stream_failure_callback_failure_does_not_skip_later_loggers( + messages_server: RecordingServer, +) -> None: + messages_server.enqueue( + ResponseSpec( + body=None, + events=( + MESSAGES_EVENTS[0], + ( + "error", + {"type": "error", "error": {"type": "overloaded_error", "message": "upstream overloaded"}}, + ), + ), + ) + ) + recorder: Final = RecordingLogger() + + class UnavailableExporter(CustomLogger): + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + raise RuntimeError("exporter unavailable") + + stream: Final = await call_messages(messages_server, [UnavailableExporter(), recorder], stream=True) + with pytest.raises(litellm.APIError) as caught: + async for _ in stream: + pass + events: Final = await recorder.wait_for_async("async_log_failure_event") + + assert len(events) == 1 + assert isinstance(events[0].kwargs["exception"], litellm.APIError) + assert type(events[0].kwargs["exception"]) is type(caught.value) + assert "async_log_success_event" not in recorder.names + + +@pytest.mark.asyncio +async def test_messages_concurrent_calls_keep_callback_state_isolated(messages_server: RecordingServer) -> None: + messages_server.expected_requests = 4 + tokens: Final = {f"messages-{index}": object() for index in range(4)} + terminal_state: Final = [] + + class CorrelateCallState(RecordingLogger): + def log_pre_api_call(self, model, messages, kwargs): + kwargs["correlation-token"] = tokens[kwargs["litellm_call_id"]] + super().log_pre_api_call(model, messages, kwargs) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + terminal_state.append((kwargs["litellm_call_id"], kwargs["correlation-token"])) + await super().async_log_success_event(kwargs, response_obj, start_time, end_time) + + correlate: Final = CorrelateCallState() + + await asyncio.gather(*(call_messages(messages_server, [correlate], litellm_call_id=call_id) for call_id in tokens)) + await correlate.wait_for_async("async_log_success_event", count=4) + + assert len(terminal_state) == 4 + assert all(token is tokens[call_id] for call_id, token in terminal_state) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_messages_cancelled_call_runs_no_terminal_callbacks( + messages_server: RecordingServer, stream: bool +) -> None: + messages_server.default_response = ResponseSpec( + body=MESSAGES_RESPONSE, + delay=0.5, + events=MESSAGES_EVENTS if stream else (), + ) + recorder: Final = RecordingLogger() + task: Final = asyncio.create_task(call_messages(messages_server, [recorder], stream=stream)) + + async with asyncio.timeout(10): + while not messages_server.requests: + await asyncio.sleep(0.01) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + + assert recorder.names.count("log_pre_api_call") == 1 + assert "log_success_event" not in recorder.names + assert "async_log_success_event" not in recorder.names + assert "log_failure_event" not in recorder.names + assert "async_log_failure_event" not in recorder.names + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [200, 500]) +async def test_whole_call_with_supplied_logger_still_owns_terminal_callbacks( + messages_server: RecordingServer, status: int +) -> None: + from litellm.rust_bridge import messages as bridge + + events: Final = [] + + class SuppliedLogger: + def pre_call(self, **kwargs): + events.append(("pre", kwargs)) + + async def async_success_handler(self, response, start, end): + events.append(("success", response)) + + def _should_run_sync_callbacks_for_async_calls(self): + return False + + def failure_handler(self, error, trace, start, end): + events.append(("sync_failure", error)) + + async def async_failure_handler(self, error, trace, start, end): + events.append(("async_failure", error)) + + messages_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE, status=status) + arguments: Final = { + "model": "claude-opus-5", + "body": {"model": "claude-opus-5", "messages": MESSAGES, "max_tokens": 64}, + "api_key": "test", + "api_base": messages_server.base_url, + "custom_llm_provider": "anthropic", + "extra_headers": {}, + "timeout": 5.0, + "logging_obj": SuppliedLogger(), + } + if status == 200: + response: Final = await bridge.amessages(**arguments) + await drain_logging() + assert [name for name, value in events] == ["pre", "success"] + assert events[1][1] is response + return + + with pytest.raises(litellm.InternalServerError) as raised: + await bridge.amessages(**arguments) + assert [name for name, value in events] == ["pre", "sync_failure", "async_failure"] + assert events[1][1] is raised.value + assert events[2][1] is raised.value + + +@pytest.mark.asyncio +@pytest.mark.parametrize("detach", (False, True)) +async def test_direct_bridge_stream_owns_callbacks(messages_server: RecordingServer, detach: bool) -> None: + from litellm.rust_bridge import messages as bridge + + recorder: Final = RecordingLogger() + release: Final = threading.Event() + if not detach: + release.set() + messages_server.enqueue( + ResponseSpec( + body=None, + chunks=tuple(f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in MESSAGES_EVENTS), + release=release, + ) + ) + stream: Final = await bridge.amessages( + model=MESSAGES_MODEL.split("/", 1)[-1], + body={"model": MESSAGES_MODEL.split("/", 1)[-1], "messages": MESSAGES, "max_tokens": 64, "stream": True}, + api_key="test-key", + api_base=messages_server.base_url, + custom_llm_provider="anthropic", + extra_headers=None, + timeout=5.0, + request_arguments={"callbacks": [recorder]}, + ) + assert "async_log_success_event" not in recorder.names + try: + if detach: + assert await anext(stream) + await asyncio.wait_for(stream.aclose(), timeout=1) + await stream.aclose() + assert "async_log_success_event" not in recorder.names + assert "async_log_failure_event" not in recorder.names + else: + assert b"message_stop" in b"".join([chunk async for chunk in stream]) + finally: + release.set() + events: Final = await recorder.wait_for_async("async_log_success_event") + assert len(events) == 1 + assert events[0].stream is True + assert events[0].response.usage.total_tokens == 9 + assert events[0].response.choices[0].message.content == "Hello from native Messages" + assert recorder.names.count("log_pre_api_call") == 1 + await stream.aclose() 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..551081a57b1 --- /dev/null +++ b/tests/test_litellm_rust/messages/test_streaming.py @@ -0,0 +1,275 @@ +import json +from collections.abc import AsyncIterator +from typing import Final, cast + +import pytest + +import litellm +from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.fixture +def messages_server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) + return recording_server + + +async def call_messages(server: RecordingServer, **kwargs: object): + return await litellm.anthropic.messages.acreate( + model=MESSAGES_MODEL, + messages=MESSAGES, + max_tokens=64, + api_key="test-key", + api_base=server.base_url, + **kwargs, + ) + + +def assert_native_request(server: RecordingServer) -> None: + assert len(server.requests) == 1 + assert "accept-encoding" not in server.requests[0].headers + + +def assert_native_response(response: object) -> None: + assert isinstance(response, dict) + assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"} + + +@pytest.mark.asyncio +async def test_messages_sends_expected_provider_request(messages_server: RecordingServer) -> None: + response: Final = await call_messages(messages_server) + + assert response["content"] == [{"type": "text", "text": "Hello from native Messages"}] + assert_native_response(response) + assert_native_request(messages_server) + request: Final = messages_server.requests[0] + assert request.path == "/v1/messages" + assert request.headers["x-api-key"] == "test-key" + assert request.headers["anthropic-version"] == "2023-06-01" + assert request.body == { + "model": "claude-sonnet-4-5-20250929", + "messages": MESSAGES, + "max_tokens": 64, + } + + +@pytest.mark.asyncio +async def test_messages_sends_custom_headers(messages_server: RecordingServer) -> None: + response: Final = await call_messages(messages_server, extra_headers={"x-trace-id": "trace-1"}) + + assert_native_response(response) + assert messages_server.requests[0].headers["x-trace-id"] == "trace-1" + + +@pytest.mark.asyncio +async def test_messages_resolves_provider_credentials( + messages_server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("ANTHROPIC_API_KEY", "environment-key") + + response: Final = await litellm.anthropic.messages.acreate( + model=MESSAGES_MODEL, + messages=MESSAGES, + max_tokens=64, + api_base=messages_server.base_url, + ) + + assert_native_response(response) + assert messages_server.requests[0].headers["x-api-key"] == "environment-key" + + +@pytest.mark.asyncio +async def test_messages_explicit_credentials_override_defaults( + messages_server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("ANTHROPIC_API_KEY", "environment-key") + + response: Final = await call_messages(messages_server) + + assert_native_response(response) + assert messages_server.requests[0].headers["x-api-key"] == "test-key" + + +@pytest.mark.asyncio +async def test_azure_messages_uses_foundry_endpoint_and_credentials(messages_server: RecordingServer) -> None: + response: Final = await litellm.anthropic.messages.acreate( + model="azure_ai/claude-opus-4.5", + messages=MESSAGES, + max_tokens=64, + api_key="azure-key", + api_base=f"{messages_server.base_url}/anthropic", + ) + + assert_native_response(response) + assert_native_request(messages_server) + assert messages_server.requests[0].path == "/anthropic/v1/messages" + assert messages_server.requests[0].headers["x-api-key"] == "azure-key" + + +@pytest.mark.asyncio +async def test_messages_stream_yields_anthropic_events(messages_server: RecordingServer) -> None: + messages_server.enqueue(ResponseSpec(body=None, events=MESSAGES_EVENTS)) + stream: Final = cast(AsyncIterator[bytes], await call_messages(messages_server, stream=True)) + payload: Final = b"".join([chunk async for chunk in stream]) + + assert_native_request(messages_server) + assert messages_server.requests[0].body["stream"] is True + assert b"event: message_start" in payload + assert b"event: content_block_delta" in payload + assert b"Hello from native Messages" in payload + assert b"event: message_stop" in payload + message_delta: Final = next( + json.loads(block.split(b"data: ", 1)[1]) + for block in payload.split(b"\n\n") + if block.startswith(b"event: message_delta") + ) + assert message_delta["usage"] == {"input_tokens": 5, "output_tokens": 4} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ["anthropic", "azure_ai"]) +async def test_messages_delivers_before_upstream_finishes(messages_server: RecordingServer, provider: str) -> None: + import asyncio + import threading + + from tests.test_litellm_rust.support.callback_recorder import RecordingLogger + + release: Final = threading.Event() + recorder: Final = RecordingLogger() + chunks: Final = tuple(f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in MESSAGES_EVENTS) + messages_server.enqueue(ResponseSpec(body=None, chunks=chunks, release=release)) + try: + stream: Final = await asyncio.wait_for( + litellm.anthropic.messages.acreate( + model=MESSAGES_MODEL.split("/", 1)[-1], + custom_llm_provider=provider, + messages=MESSAGES, + max_tokens=64, + api_key="test-key", + api_base=messages_server.base_url, + stream=True, + callbacks=[recorder], + ), + timeout=5, + ) + first: Final = await asyncio.wait_for(anext(stream), timeout=2) + assert b"message_start" in first + assert not release.is_set() + assert "async_log_success_event" not in recorder.names + assert messages_server.requests[0].body["stream"] is True + assert stream._hidden_params["additional_headers"]["x-litellm-rust"] == "true" + assert messages_server.requests[0].path == ( + "/v1/messages" if provider == "anthropic" else "/anthropic/v1/messages" + ) + release.set() + remaining: Final = b"".join([chunk async for chunk in stream]) + assert first + remaining == b"".join(chunks) + await stream.aclose() + events: Final = await recorder.wait_for_async("async_log_success_event") + assert len(events) == 1 + assert "async_log_failure_event" not in recorder.names + assert events[0].response.choices[0].message.content == "Hello from native Messages" + assert events[0].response.usage.completion_tokens == 4 + assert len(messages_server.requests) == 1 + finally: + release.set() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ending", ["truncated", "provider_error", "http_error"]) +async def test_messages_stream_failures_do_not_replay(messages_server: RecordingServer, ending: str) -> None: + from tests.test_litellm_rust.support.callback_recorder import RecordingLogger + + recorder: Final = RecordingLogger() + events: Final = ( + MESSAGES_EVENTS[:-1] + if ending == "truncated" + else ( + *MESSAGES_EVENTS[:1], + ("error", {"type": "error", "error": {"type": "overloaded_error", "message": "overloaded"}}), + ) + ) + messages_server.enqueue( + ResponseSpec( + body={"error": "rejected"}, + status=429 if ending == "http_error" else 200, + events=() if ending == "http_error" else events, + ) + ) + + async def consume() -> None: + stream: Final = await call_messages(messages_server, stream=True, callbacks=[recorder]) + async for _ in stream: + pass + + with pytest.raises(litellm.APIError): + await consume() + await recorder.wait_for_async("async_log_failure_event") + assert recorder.names.count("async_log_failure_event") == 1 + assert "async_log_success_event" not in recorder.names + assert len(messages_server.requests) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("abandon", ["close", "cancel", "drop", "exhaust"]) +async def test_messages_stream_abandonment_releases_roots(messages_server: RecordingServer, abandon: str) -> None: + import asyncio + import gc + import threading + import weakref + + from tests.test_litellm_rust.support.callback_recorder import RecordingLogger + + class NamesOnlyLogger(RecordingLogger): + def _record(self, name: str, kwargs: object = None, response: object = None) -> None: + super()._record(name) + + class Root: + pass + + root = Root() + reference: Final = weakref.ref(root) + release: Final = threading.Event() + recorder: Final = NamesOnlyLogger() + chunks: Final = tuple(f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in MESSAGES_EVENTS) + messages_server.enqueue(ResponseSpec(body=None, chunks=chunks, release=release)) + try: + stream = await call_messages(messages_server, stream=True, callbacks=[recorder], metadata={"retained": root}) + del root + assert reference() is not None + assert b"message_start" in await anext(stream) + if abandon == "exhaust": + release.set() + async for _ in stream: + pass + await recorder.wait_for_async("async_log_success_event") + gc.collect() + assert reference() is None + assert recorder.names.count("async_log_success_event") == 1 + return + if abandon == "cancel": + task: Final = asyncio.create_task(anext(stream)) + await asyncio.sleep(0.05) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + elif abandon == "close": + await stream.aclose() + await stream.aclose() + del stream + gc.collect() + assert reference() is not None + assert "async_log_success_event" not in recorder.names + assert "async_log_failure_event" not in recorder.names + release.set() + await recorder.wait_for_async("async_log_success_event") + assert recorder.names.count("async_log_success_event") == 1 + assert "async_log_failure_event" not in recorder.names + assert len(messages_server.requests) == 1 + gc.collect() + assert reference() is None + finally: + release.set() diff --git a/tests/test_litellm_rust/ocr/__init__.py b/tests/test_litellm_rust/ocr/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/test_litellm_rust/ocr/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py new file mode 100644 index 00000000000..84fd01add80 --- /dev/null +++ b/tests/test_litellm_rust/ocr/test_callbacks.py @@ -0,0 +1,502 @@ +import asyncio +import copy +import queue +import threading +from typing import Final + +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.provenance import has_rust_response_marker +from tests.test_litellm_rust.support.requests import ( + OCR_DOCUMENT, + OCR_RESPONSE, + call_native_aocr, + call_native_ocr, + request_body, + request_headers, +) +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_ocr(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object): + return call_native_ocr(server, callbacks=callbacks, **kwargs) + + +async def call_aocr(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object): + return await call_native_aocr(server, callbacks=callbacks, **kwargs) + + +def test_pre_call_receives_expected_provider_request(ocr_server: RecordingServer) -> None: + observations: Final = [] + + class Observe(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + observations.append((model, messages, copy.deepcopy(kwargs["additional_args"]))) + + call_ocr(ocr_server, [Observe()], pages=[0]) + + assert len(observations) == 1 + model, messages, additional_args = observations[0] + assert model == "mistral-ocr-latest" + assert messages == [{"role": "user", "content": "default-message-value"}] + assert additional_args["api_base"] == f"{ocr_server.base_url}/v1/ocr" + assert additional_args["complete_input_dict"] == { + "model": "mistral-ocr-latest", + "document": OCR_DOCUMENT, + "pages": [0], + } + + +@pytest.mark.parametrize("raise_after_edit", [False, True]) +def test_pre_call_body_edits_reach_later_callbacks_and_provider( + ocr_server: RecordingServer, raise_after_edit: bool +) -> None: + observed: Final = [] + + class Edit(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + request_body(kwargs)["include_image_base64"] = True + if raise_after_edit: + raise RuntimeError("audit exporter unavailable") + + class Observe(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + observed.append(copy.deepcopy(request_body(kwargs))) + + call_ocr(ocr_server, [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 + + +def test_pre_call_header_edits_reach_later_callbacks_and_provider(ocr_server: RecordingServer) -> None: + observed: Final = [] + + class Edit(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + request_headers(kwargs)["x-audit-tag"] = "reviewed" + + class Observe(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + observed.append(dict(request_headers(kwargs))) + + call_ocr(ocr_server, [Edit(), Observe()]) + + assert observed[0]["x-audit-tag"] == "reviewed" + assert ocr_server.requests[0].headers["x-audit-tag"] == "reviewed" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("rust_enabled", [False, True]) +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_pre_call_nested_mutation_updates_retained_references( + ocr_server: RecordingServer, rust_enabled: bool, asynchronous: bool +) -> None: + litellm.rust(rust_enabled) + original: Final = dict(OCR_DOCUMENT) + replacement_url: Final = "data:application/pdf;base64,ZGVm" + retained: Final = [] + + class Retain(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + assert request_body(kwargs)["document"] is original + retained.append(request_body(kwargs)["document"]) + + class Edit(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + original["document_url"] = replacement_url + + arguments: Final = { + "model": "mistral/mistral-ocr-latest", + "document": original, + "api_key": "test-key", + "api_base": ocr_server.base_url, + "callbacks": [Retain(), Edit()], + } + response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) + + 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 has_rust_response_marker(response) is rust_enabled + + +def test_pre_call_field_replacement_preserves_original_references(ocr_server: RecordingServer) -> None: + original: Final = dict(OCR_DOCUMENT) + replacement: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,ZGVm"} + retained: Final = [] + + class RetainAndReplace(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + body = request_body(kwargs) + retained.append(body["document"]) + body["document"] = replacement + + call_native_ocr( + ocr_server, + document=original, + callbacks=[RetainAndReplace()], + ) + + assert retained[0] is original + assert original["document_url"] == OCR_DOCUMENT["document_url"] + assert ocr_server.requests[0].body["document"] == replacement + + +def test_pre_call_body_rebinding_does_not_replace_inflight_request(ocr_server: RecordingServer) -> None: + observed: Final = [] + + class Rebind(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + kwargs["additional_args"]["complete_input_dict"] = {"replacement": True} + + class Observe(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + observed.append(request_body(kwargs)) + + call_ocr(ocr_server, [Rebind(), Observe()]) + + assert observed == [{"replacement": True}] + assert ocr_server.requests[0].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT} + + +def test_queued_payload_observes_later_callback_mutations(ocr_server: RecordingServer) -> None: + queued: Final = [] + + class QueuePayload(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + queued.append(request_body(kwargs)) + + class Edit(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + request_body(kwargs)["queued-edit"] = True + + call_ocr(ocr_server, [QueuePayload(), Edit()]) + + assert queued[0]["queued-edit"] is True + + +def test_pre_call_state_reaches_terminal_callbacks(ocr_server: RecordingServer) -> None: + token: Final = object() + terminal_tokens: queue.SimpleQueue[object] = queue.SimpleQueue() + finished: Final = threading.Event() + + class Stash(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + kwargs["test-token"] = token + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + terminal_tokens.put(kwargs["test-token"]) + finished.set() + + call_ocr(ocr_server, [Stash()]) + + assert finished.wait(10) + assert terminal_tokens.get_nowait() is token + + +@pytest.mark.asyncio +async def test_success_callbacks_receive_expected_context_and_response(ocr_server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + + await call_aocr( + ocr_server, + [recorder], + litellm_call_id="ocr-success", + metadata={"source": "callback-test"}, + ) + events: Final = await recorder.wait_for_async("async_log_success_event") + + assert len(events) == 1 + assert events[0].call_type == "aocr" + assert events[0].kwargs["litellm_call_id"] == "ocr-success" + assert events[0].kwargs["litellm_params"]["metadata"]["source"] == "callback-test" + assert events[0].response.pages[0].markdown == "native OCR response" + + +@pytest.mark.asyncio +async def test_failure_callbacks_receive_expected_context_and_error(ocr_server: RecordingServer) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500)) + observations: Final = [] + + class Observe(CustomLogger): + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + observations.append(("sync", kwargs["call_type"], kwargs["exception"], response_obj)) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + observations.append(("async", kwargs["call_type"], kwargs["exception"], response_obj)) + + with pytest.raises(litellm.InternalServerError): + await call_aocr(ocr_server, [Observe()]) + + assert [observation[0] for observation in observations] == ["sync", "async"] + assert all(observation[1] == "aocr" for observation in observations) + assert all(isinstance(observation[2], litellm.InternalServerError) for observation in observations) + assert all(observation[3] is None for observation in observations) + + +@pytest.mark.asyncio +async def test_pre_call_runs_in_callers_execution_context(ocr_server: RecordingServer) -> None: + caller_loop: Final = asyncio.get_running_loop() + caller_thread: Final = threading.current_thread() + recorder: Final = RecordingLogger() + + await call_aocr(ocr_server, [recorder]) + + events: Final = await recorder.wait_for_async("log_pre_api_call") + assert len(events) == 1 + assert events[0].loop is caller_loop + assert events[0].thread is caller_thread + + +@pytest.mark.asyncio +async def test_ocr_failure_callbacks_receive_pre_call_state(ocr_server: RecordingServer) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500)) + token: Final = object() + observed: Final = [] + + class TrackInFlightRequest(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + kwargs["request-token"] = token + + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("sync", kwargs["request-token"])) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("async", kwargs["request-token"])) + + with pytest.raises(litellm.InternalServerError): + await call_aocr(ocr_server, [TrackInFlightRequest()]) + + assert [event for event, _ in observed] == ["sync", "async"] + assert all(observed_token is token for _, observed_token in observed) + + +@pytest.mark.asyncio +async def test_ocr_failure_callback_error_does_not_mask_provider_error_or_later_callbacks( + ocr_server: RecordingServer, +) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500)) + recorder: Final = RecordingLogger() + + class UnavailableExporter(CustomLogger): + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + raise RuntimeError("exporter unavailable") + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + raise RuntimeError("exporter unavailable") + + with pytest.raises(litellm.InternalServerError) as caught: + await call_aocr(ocr_server, [UnavailableExporter(), recorder]) + + sync_events: Final = tuple(event for event in recorder.events if event.name == "log_failure_event") + async_events: Final = tuple(event for event in recorder.events if event.name == "async_log_failure_event") + assert len(sync_events) == 1 + assert len(async_events) == 1 + assert sync_events[0].kwargs["exception"] is caught.value + assert async_events[0].kwargs["exception"] is caught.value + assert "async_log_success_event" not in recorder.names + + +def test_ocr_duplicate_callback_registration_dispatches_once(ocr_server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + + call_ocr( + ocr_server, + [recorder, recorder], + success_callback=[recorder], + failure_callback=[recorder], + ) + recorder.wait_for("log_success_event") + + assert recorder.names.count("log_pre_api_call") == 1 + assert recorder.names.count("logging_hook") == 1 + assert recorder.names.count("log_success_event") == 1 + assert "log_failure_event" not in recorder.names + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_azure_token_callback_precedes_logger_and_preserves_context( + ocr_server: RecordingServer, + isolated_azure_auth: None, + asynchronous: bool, +) -> None: + from contextvars import ContextVar + from tests.test_litellm_rust.support.requests import call_aocr as public_aocr, call_ocr as public_ocr + + context: Final = ContextVar("azure-token-context", default="missing") + context.set("caller") + caller_thread: Final = threading.current_thread() + caller_loop: Final = asyncio.get_running_loop() + observations: Final = [] + + class Provider: + def __call__(self) -> str: + assert context.get() == "caller" + assert threading.current_thread() is caller_thread + assert asyncio.get_running_loop() is caller_loop + observations.append("token") + return "caller-token" + + class Edit(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + assert request_headers(kwargs)["Authorization"] == "Bearer caller-token" + observations.append("pre_call") + request_headers(kwargs)["Authorization"] = "Bearer edited" + + provider: Final = Provider() + arguments: Final = { + "model": "azure_ai/mistral-ocr-latest", + "api_key": None, + "azure_ad_token_provider": provider, + "callbacks": [Edit()], + } + response: Final = ( + await public_aocr(ocr_server, **arguments) if asynchronous else public_ocr(ocr_server, **arguments) + ) + assert has_rust_response_marker(response) + assert observations == ["token", "pre_call"] + assert ocr_server.requests[0].headers["authorization"] == "Bearer edited" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_azure_token_callback_can_reenter_native_sdk( + ocr_server: RecordingServer, + isolated_azure_auth: None, + asynchronous: bool, +) -> None: + from tests.test_litellm_rust.support.requests import call_aocr as public_aocr, call_ocr as public_ocr + + ocr_server.expected_requests = 2 + calls: Final = [] + + def provider() -> str: + calls.append("token") + nested: Final = public_ocr(ocr_server) + assert has_rust_response_marker(nested) + return "outer-token" + + arguments: Final = { + "model": "azure_ai/mistral-ocr-latest", + "api_key": None, + "azure_ad_token_provider": provider, + } + response: Final = ( + await public_aocr(ocr_server, **arguments) if asynchronous else public_ocr(ocr_server, **arguments) + ) + assert has_rust_response_marker(response) + assert calls == ["token"] + assert [request.headers["authorization"] for request in ocr_server.requests] == [ + "Bearer test-key", + "Bearer outer-token", + ] + + +@pytest.mark.asyncio +async def test_concurrent_azure_token_callbacks_keep_results_and_errors_separate( + ocr_server: RecordingServer, + isolated_azure_auth: None, +) -> None: + from tests.test_litellm_rust.support.requests import call_aocr as public_aocr + + ocr_server.expected_requests = 2 + + async def request(token: str, fail: bool) -> object: + def provider() -> str: + if fail: + raise ValueError(token) + return token + + return await public_aocr( + ocr_server, + model="azure_ai/mistral-ocr-latest", + api_key=None, + azure_ad_token_provider=provider, + ) + + responses: Final = await asyncio.gather( + request("first", False), + request("failed", True), + request("second", False), + return_exceptions=True, + ) + assert has_rust_response_marker(responses[0]) + assert isinstance(responses[1], litellm.APIConnectionError) + assert "Failed to get Azure AD token: failed" in str(responses[1]) + assert has_rust_response_marker(responses[2]) + assert sorted(request.headers["authorization"] for request in ocr_server.requests) == [ + "Bearer first", + "Bearer second", + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("outcome", ["success", "failure", "cancellation"]) +async def test_azure_token_provider_is_released_after_request( + ocr_server: RecordingServer, + isolated_azure_auth: None, + outcome: str, +) -> None: + import gc + import weakref + from tests.test_litellm_rust.support.callback_recorder import drain_logging + from tests.test_litellm_rust.support.requests import call_aocr as public_aocr + + class Provider: + def __call__(self) -> str: + if outcome == "failure": + raise ValueError("unavailable") + return "caller-token" + + async def invoke() -> weakref.ReferenceType[Provider]: + provider: Final = Provider() + reference: Final = weakref.ref(provider) + if outcome == "failure": + ocr_server.expected_requests = 0 + with pytest.raises(litellm.APIConnectionError): + await public_aocr( + ocr_server, model="azure_ai/mistral-ocr-latest", api_key=None, azure_ad_token_provider=provider + ) + elif outcome == "cancellation": + ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.1)) + task: Final = asyncio.create_task( + public_aocr( + ocr_server, + model="azure_ai/mistral-ocr-latest", + api_key=None, + azure_ad_token_provider=provider, + ) + ) + await ocr_server.wait_for_requests(1) + assert reference() is provider + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + else: + response: Final = await public_aocr( + ocr_server, + model="azure_ai/mistral-ocr-latest", + api_key=None, + azure_ad_token_provider=provider, + ) + assert has_rust_response_marker(response) + return reference + + reference: Final = await invoke() + await drain_logging() + await asyncio.sleep(0) + gc.collect() + assert reference() is None diff --git a/tests/test_litellm_rust/ocr/test_dispatch.py b/tests/test_litellm_rust/ocr/test_dispatch.py new file mode 100644 index 00000000000..d26a1cf625e --- /dev/null +++ b/tests/test_litellm_rust/ocr/test_dispatch.py @@ -0,0 +1,78 @@ +from typing import Final + +import pytest + +import litellm +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec +from tests.test_litellm_rust.support.provenance import has_rust_response_marker +from tests.test_litellm_rust.support.requests import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE + +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 test_public_ocr_entrypoint_uses_native_transport_when_enabled(ocr_server: RecordingServer) -> None: + response: Final = litellm.ocr( + model=OCR_MODEL, + document=OCR_DOCUMENT, + api_key="test-key", + api_base=ocr_server.base_url, + ) + + assert response.pages[0].markdown == "native OCR response" + assert has_rust_response_marker(response) + + +@pytest.mark.asyncio +async def test_public_aocr_entrypoint_uses_native_transport_when_enabled(ocr_server: RecordingServer) -> None: + response: Final = await litellm.aocr( + model=OCR_MODEL, + document=OCR_DOCUMENT, + api_key="test-key", + api_base=ocr_server.base_url, + ) + + assert response.pages[0].markdown == "native OCR response" + assert has_rust_response_marker(response) + + +def test_public_ocr_falls_back_when_native_transport_declines(ocr_server: RecordingServer) -> None: + response: Final = litellm.ocr( + model=OCR_MODEL, + document={"type": "file", "file": b"%PDF-1.4", "mime_type": "application/pdf"}, + api_key="test-key", + api_base=ocr_server.base_url, + ) + + assert response.pages[0].markdown == "native OCR response" + assert not has_rust_response_marker(response) + + +def test_public_ocr_uses_python_transport_when_disabled(ocr_server: RecordingServer) -> None: + litellm.rust(False) + + response: Final = litellm.ocr( + model=OCR_MODEL, + document=OCR_DOCUMENT, + api_key="test-key", + api_base=ocr_server.base_url, + ) + + assert response.pages[0].markdown == "native OCR response" + assert not has_rust_response_marker(response) + + +def test_native_ocr_requests_an_uncompressed_response(ocr_server: RecordingServer) -> None: + litellm.ocr( + model=OCR_MODEL, + document=OCR_DOCUMENT, + api_key="test-key", + api_base=ocr_server.base_url, + ) + + assert ocr_server.requests[0].headers["accept-encoding"] == "identity" diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py new file mode 100644 index 00000000000..8493e03ce29 --- /dev/null +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -0,0 +1,447 @@ +from typing import Final + +import pytest + +import litellm +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from tests.test_litellm_rust.support.requests import ( + OCR_DOCUMENT, + OCR_MODEL, + OCR_RESPONSE, + call_native_ocr, + call_aocr, + call_ocr as call_public_ocr, +) +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.provenance import has_rust_response_marker +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]) +async def test_azure_ocr_calls_python_token_provider( + ocr_server: RecordingServer, isolated_azure_auth: None, asynchronous: bool +) -> None: + calls: Final = [] + + def token_provider() -> str: + calls.append("token") + return "callback-token" + + arguments: Final = { + "model": "azure_ai/mistral-ocr-latest", + "api_key": None, + "azure_ad_token_provider": token_provider, + } + response: Final = ( + await call_aocr(ocr_server, **arguments) if asynchronous else call_public_ocr(ocr_server, **arguments) + ) + + assert calls == ["token"] + assert has_rust_response_marker(response) + assert response.pages[0].markdown == "native OCR response" + assert_native_request(ocr_server) + assert ocr_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 call_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse: + return call_native_ocr(server, **kwargs) + + +def assert_native_request(server: RecordingServer) -> None: + assert len(server.requests) == 1 + assert not server.requests[0].headers.get("user-agent", "").startswith("python-httpx") + + +def test_ocr_sends_expected_provider_request(ocr_server: RecordingServer) -> None: + response: Final = call_ocr(ocr_server) + + assert response.pages[0].markdown == "native OCR response" + assert_native_request(ocr_server) + assert ocr_server.requests[0].path == "/v1/ocr" + assert ocr_server.requests[0].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT} + + +def test_ocr_rejects_unsupported_file_document_before_callbacks(ocr_server: RecordingServer) -> None: + ocr_server.expected_requests = 0 + recorder: Final = RecordingLogger() + + with pytest.raises(NotImplementedError, match="OCR file document preparation"): + call_native_ocr( + ocr_server, + document={"type": "file", "file": b"%PDF-1.4", "mime_type": "application/pdf"}, + callbacks=[recorder], + ) + + assert ocr_server.requests == [] + assert recorder.events == () + + +def test_ocr_sends_optional_parameters(ocr_server: RecordingServer) -> None: + call_ocr(ocr_server, pages=[0, 2], include_image_base64=True) + + assert ocr_server.requests[0].body["pages"] == [0, 2] + assert ocr_server.requests[0].body["include_image_base64"] is True + + +def test_ocr_sends_custom_headers(ocr_server: RecordingServer) -> None: + call_ocr(ocr_server, extra_headers={"x-trace-id": "trace-1"}) + + assert ocr_server.requests[0].headers["authorization"] == "Bearer test-key" + assert ocr_server.requests[0].headers["x-trace-id"] == "trace-1" + + +def test_ocr_resolves_provider_credentials(ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("MISTRAL_API_KEY", "environment-key") + + call_native_ocr(ocr_server, api_key=None) + + assert ocr_server.requests[0].headers["authorization"] == "Bearer environment-key" + + +def test_ocr_explicit_credentials_override_defaults( + ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("MISTRAL_API_KEY", "environment-key") + + call_ocr(ocr_server) + + assert ocr_server.requests[0].headers["authorization"] == "Bearer test-key" + + +def test_ocr_resolves_provider_endpoint(ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("AZURE_AI_API_KEY", "azure-key") + monkeypatch.setenv("AZURE_AI_API_BASE", ocr_server.base_url) + + call_native_ocr(ocr_server, model="azure_ai/pixtral-12b-2409", api_key=None, api_base=None) + + assert_native_request(ocr_server) + assert ocr_server.requests[0].path == "/providers/mistral/azure/ocr" + assert ocr_server.requests[0].headers["authorization"] == "Bearer azure-key" + + +def test_ocr_resolves_vertex_project_and_location(ocr_server: RecordingServer) -> None: + call_native_ocr( + ocr_server, + model="vertex_ai/mistral-ocr-2505", + api_key="vertex-token", + vertex_project="project-1", + vertex_location="us-central1", + ) + + assert_native_request(ocr_server) + assert ocr_server.requests[0].path == ( + "/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-2505:rawPredict" + ) + + +def test_ocr_returns_normalized_response(ocr_server: RecordingServer) -> None: + response: Final = call_ocr(ocr_server) + + assert isinstance(response, OCRResponse) + assert response.model == "mistral-ocr-latest" + assert response.usage_info.pages_processed == 1 + + +def test_ocr_provider_error_preserves_status_and_context(ocr_server: RecordingServer) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "invalid OCR request"}, status=400)) + + with pytest.raises(litellm.BadRequestError) as caught: + call_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" not in str(caught.value) + + +def test_ocr_honors_request_timeout(ocr_server: RecordingServer) -> None: + ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.2)) + + with pytest.raises(RuntimeError, match="OCR transport failed"): + call_ocr(ocr_server, timeout=0.01) + + assert len(ocr_server.requests) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("backend", ["python", "rust"]) +@pytest.mark.parametrize( + "credentials, expected_token, expected_calls", + [ + ({"api_key": "resource-key"}, "resource-key", 0), + ({"azure_ad_token": "static-token"}, "callback-1", 1), + ({"extra_headers": {"Authorization": "Bearer override"}}, "override", 1), + ], +) +async def test_azure_ocr_token_provider_precedence( + ocr_server: RecordingServer, + isolated_azure_auth: None, + asynchronous: bool, + backend: str, + credentials: dict[str, object], + expected_token: str, + expected_calls: int, +) -> None: + litellm.rust(backend == "rust") + calls: Final = [] + + def token_provider() -> str: + calls.append("token") + return f"callback-{len(calls)}" + + arguments: Final = { + "model": "azure_ai/mistral-ocr-latest", + "api_key": None, + "azure_ad_token_provider": token_provider, + **credentials, + } + response: Final = ( + await call_aocr(ocr_server, **arguments) if asynchronous else call_public_ocr(ocr_server, **arguments) + ) + assert has_rust_response_marker(response) == (backend == "rust") + assert len(calls) == expected_calls + assert len(ocr_server.requests) == 1 + assert ocr_server.requests[0].headers["authorization"] == f"Bearer {expected_token}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_azure_ocr_does_not_cache_caller_tokens( + ocr_server: RecordingServer, + isolated_azure_auth: None, + asynchronous: bool, +) -> None: + calls: Final = [] + ocr_server.expected_requests = 2 + + def token_provider() -> str: + calls.append("token") + return f"callback-{len(calls)}" + + for _ in range(2): + arguments: Final = { + "model": "azure_ai/mistral-ocr-latest", + "api_key": None, + "azure_ad_token_provider": token_provider, + } + response: Final = ( + await call_aocr(ocr_server, **arguments) if asynchronous else call_public_ocr(ocr_server, **arguments) + ) + assert has_rust_response_marker(response) + assert len(calls) == 2 + assert [request.headers["authorization"] for request in ocr_server.requests] == [ + "Bearer callback-1", + "Bearer callback-2", + ] + + +class TokenAbort(BaseException): + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("backend", ["python", "rust"]) +@pytest.mark.parametrize("failure", ["non_string", "type_error", "ordinary", "abort"]) +async def test_azure_ocr_token_failure_stops_execution( + ocr_server: RecordingServer, + isolated_azure_auth: None, + asynchronous: bool, + backend: str, + failure: str, +) -> None: + litellm.rust(backend == "rust") + ocr_server.expected_requests = 0 + calls: Final = [] + recorder: Final = RecordingLogger() + original: Final = { + "type_error": TypeError("token type"), + "ordinary": ValueError("token unavailable"), + "abort": TokenAbort("abort"), + } + + def token_provider() -> object: + calls.append("token") + if failure == "non_string": + return 123 + 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(ocr_server, **arguments) if asynchronous else call_public_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] + elif failure == "abort": + assert caught.value is original[failure] + elif failure == "type_error": + assert caught.value.__context__ is original[failure] + else: + assert isinstance(caught.value.__context__, TypeError) + + +@pytest.mark.parametrize( + "configuration", + [ + {"azure_ad_token": "oidc/assertion", "client_id": "client", "tenant_id": "tenant"}, + {"model": "azure_ai/doc-intelligence/prebuilt-read"}, + {"document": {"type": "file", "file": b"pdf"}}, + ], +) +def test_azure_unsupported_native_auth_does_not_invoke_provider( + ocr_server: RecordingServer, + isolated_azure_auth: None, + configuration: dict[str, object], +) -> None: + ocr_server.expected_requests = 0 + calls: Final = [] + recorder: Final = RecordingLogger() + + def provider() -> str: + calls.append("token") + return "unused" + + arguments: Final = { + "model": "azure_ai/mistral-ocr-latest", + "api_key": None, + "azure_ad_token_provider": provider, + "callbacks": [recorder], + **configuration, + } + with pytest.raises(NotImplementedError): + call_native_ocr(ocr_server, **arguments) + assert calls == [] + assert recorder.events == () + assert ocr_server.requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ["python", "rust"]) +async def test_azure_missing_endpoint_prevents_token_callback( + ocr_server: RecordingServer, + isolated_azure_auth: None, + backend: str, +) -> None: + litellm.rust(backend == "rust") + ocr_server.expected_requests = 0 + calls: Final = [] + + def provider() -> str: + calls.append("token") + return "unused" + + with pytest.raises(litellm.APIConnectionError, match="Missing Azure AI API Base"): + await call_aocr( + ocr_server, + model="azure_ai/mistral-ocr-latest", + api_key=None, + api_base=None, + azure_ad_token_provider=provider, + ) + assert calls == [] + assert ocr_server.requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ["python", "rust"]) +async def test_azure_empty_callback_token_does_not_restore_static_token( + ocr_server: RecordingServer, + isolated_azure_auth: None, + backend: str, +) -> None: + litellm.rust(backend == "rust") + ocr_server.expected_requests = 0 + + def provider() -> str: + return "" + + with pytest.raises(litellm.APIConnectionError, match="Missing Azure AI credentials"): + await call_aocr( + ocr_server, + model="azure_ai/mistral-ocr-latest", + api_key=None, + azure_ad_token="static-token", + azure_ad_token_provider=provider, + ) + assert ocr_server.requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ["python", "rust"]) +async def test_azure_falsey_callable_leaves_static_token_unchanged( + ocr_server: RecordingServer, + isolated_azure_auth: None, + backend: str, +) -> None: + litellm.rust(backend == "rust") + calls: Final = [] + + class Provider: + def __bool__(self) -> bool: + return False + + def __call__(self) -> str: + calls.append("token") + return "unused" + + response: Final = await call_aocr( + ocr_server, + model="azure_ai/mistral-ocr-latest", + api_key=None, + azure_ad_token="static-token", + azure_ad_token_provider=Provider(), + ) + assert has_rust_response_marker(response) == (backend == "rust") + assert calls == [] + assert ocr_server.requests[0].headers["authorization"] == "Bearer static-token" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ["python", "rust"]) +async def test_azure_token_callback_does_not_await_coroutine_result( + ocr_server: RecordingServer, + isolated_azure_auth: None, + backend: str, +) -> None: + litellm.rust(backend == "rust") + ocr_server.expected_requests = 0 + calls: Final = [] + + async def acquire() -> str: + calls.append("awaited") + return "unused" + + coroutine: Final = acquire() + + def provider() -> object: + return coroutine + + try: + with pytest.raises(litellm.APIConnectionError, match="Azure AD token must be a string"): + await call_aocr( + ocr_server, model="azure_ai/mistral-ocr-latest", api_key=None, azure_ad_token_provider=provider + ) + finally: + coroutine.close() + assert calls == [] + assert ocr_server.requests == [] diff --git a/tests/test_litellm_rust/support/__init__.py b/tests/test_litellm_rust/support/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/test_litellm_rust/support/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/test_litellm_rust/support/callback_recorder.py b/tests/test_litellm_rust/support/callback_recorder.py new file mode 100644 index 00000000000..d195cdddbfa --- /dev/null +++ b/tests/test_litellm_rust/support/callback_recorder.py @@ -0,0 +1,178 @@ +import asyncio +import copy +import threading +import time +from dataclasses import dataclass +from typing import Final + +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker + + +async def drain_logging(worker: LoggingWorker = GLOBAL_LOGGING_WORKER) -> None: + # The wrapper's scheduled helper enqueues without awaiting; let it run before joining the queue. + await asyncio.sleep(0) + worker.start() + await asyncio.wait_for(worker.flush(), timeout=10) + + +@dataclass(frozen=True, slots=True) +class HookEvent: + name: str + call_type: str | None + stream: bool | None + thread: threading.Thread + loop: asyncio.AbstractEventLoop | None + has_running_loop: bool + kwargs: object + response: object + + +class RecordingLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self._events: list[HookEvent] = [] + self._condition = threading.Condition() + + @property + def events(self) -> tuple[HookEvent, ...]: + with self._condition: + return tuple(self._events) + + @property + def names(self) -> tuple[str, ...]: + return tuple(event.name for event in self.events) + + def _record(self, name: str, kwargs: object = None, response: object = None) -> None: + details: Final = kwargs if isinstance(kwargs, dict) else {} + try: + snapshot: Final = copy.deepcopy(details) + except Exception: + snapshot = dict(details) + if "exception" in details: + snapshot["exception"] = details["exception"] + try: + loop: Final = asyncio.get_running_loop() + has_running_loop: Final = True + except RuntimeError: + loop = None + has_running_loop = False + event: Final = HookEvent( + name=name, + call_type=details.get("call_type"), + stream=details.get("stream"), + thread=threading.current_thread(), + loop=loop, + has_running_loop=has_running_loop, + kwargs=snapshot, + response=response, + ) + with self._condition: + self._events.append(event) + self._condition.notify_all() + + def wait_for(self, name: str, count: int = 1, timeout: float = 10) -> tuple[HookEvent, ...]: + deadline: Final = time.monotonic() + timeout + with self._condition: + while sum(event.name == name for event in self._events) < count: + remaining: Final = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"Timed out waiting for {count} {name} events; saw {self.names}") + self._condition.wait(remaining) + return tuple(event for event in self._events if event.name == name) + + async def wait_for_async(self, name: str, count: int = 1, timeout: float = 10) -> tuple[HookEvent, ...]: + await asyncio.wait_for(asyncio.to_thread(self.wait_for, name, count, timeout), timeout=timeout + 1) + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=timeout) + return tuple(event for event in self.events if event.name == name) + + def log_pre_api_call(self, model, messages, kwargs): + self._record("log_pre_api_call", kwargs) + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + self._record("log_success_event", kwargs, response_obj) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self._record("async_log_success_event", kwargs, response_obj) + + def log_stream_event(self, kwargs, response_obj, start_time, end_time): + self._record("log_stream_event", kwargs, response_obj) + + async def async_log_stream_event(self, kwargs, response_obj, start_time, end_time): + self._record("async_log_stream_event", kwargs, response_obj) + + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + self._record("log_failure_event", kwargs, response_obj) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + self._record("async_log_failure_event", kwargs, response_obj) + + def logging_hook(self, kwargs, result, call_type): + self._record("logging_hook", kwargs, result) + return kwargs, result + + async def async_logging_hook(self, kwargs, result, call_type): + self._record("async_logging_hook", kwargs, result) + return kwargs, result + + async def async_pre_call_deployment_hook(self, kwargs, call_type): + self._record("async_pre_call_deployment_hook", kwargs) + + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + self._record("async_post_call_success_deployment_hook", request_data, response) + return response + + async def async_post_call_failure_deployment_hook(self, request_data, exception, call_type, fallback_depth=None): + self._record("async_post_call_failure_deployment_hook", request_data, exception) + + +@dataclass(frozen=True, slots=True) +class LiveReferenceEvent: + kwargs: object + response: object + + +class LiveReferenceLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self._events: list[LiveReferenceEvent] = [] + self._condition = threading.Condition() + + @property + def events(self) -> tuple[LiveReferenceEvent, ...]: + with self._condition: + return tuple(self._events) + + def _record(self, kwargs: object, response: object) -> None: + with self._condition: + self._events.append(LiveReferenceEvent(kwargs=kwargs, response=response)) + self._condition.notify_all() + + def wait_for(self, count: int = 1, timeout: float = 10) -> tuple[LiveReferenceEvent, ...]: + deadline: Final = time.monotonic() + timeout + with self._condition: + while len(self._events) < count: + remaining: Final = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"Timed out waiting for {count} live-reference events") + self._condition.wait(remaining) + return tuple(self._events) + + async def wait_for_async(self, count: int = 1, timeout: float = 10) -> tuple[LiveReferenceEvent, ...]: + return await asyncio.wait_for(asyncio.to_thread(self.wait_for, count, timeout), timeout=timeout + 1) + + def release(self) -> None: + with self._condition: + self._events.clear() + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + self._record(kwargs, response_obj) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self._record(kwargs, response_obj) + + +class SecondaryLiveReferenceLogger(LiveReferenceLogger): + pass diff --git a/tests/test_litellm_rust/support/prometheus.py b/tests/test_litellm_rust/support/prometheus.py new file mode 100644 index 00000000000..000e6041a1e --- /dev/null +++ b/tests/test_litellm_rust/support/prometheus.py @@ -0,0 +1,19 @@ +from collections.abc import Iterator +from contextlib import contextmanager +from typing import Final + +from prometheus_client import REGISTRY, CollectorRegistry + + +@contextmanager +def isolated_prometheus_registry(registry: CollectorRegistry = REGISTRY) -> Iterator[None]: + original: Final = tuple(registry._collector_to_names) # pyright: ignore[reportPrivateUsage] # prometheus_client has no public collector enumeration + for collector in original: + registry.unregister(collector) + try: + yield + finally: + for collector in tuple(registry._collector_to_names): # pyright: ignore[reportPrivateUsage] # remove collectors created inside the isolation scope + registry.unregister(collector) + for collector in original: + registry.register(collector) diff --git a/tests/test_litellm_rust/support/provenance.py b/tests/test_litellm_rust/support/provenance.py new file mode 100644 index 00000000000..bffda20e580 --- /dev/null +++ b/tests/test_litellm_rust/support/provenance.py @@ -0,0 +1,4 @@ +def has_rust_response_marker(response: object) -> bool: + from litellm.rust_bridge.provenance import has_rust_response_marker as implementation + + return implementation(response) diff --git a/tests/test_litellm_rust/support/recording_server.py b/tests/test_litellm_rust/support/recording_server.py new file mode 100644 index 00000000000..7266e4b1c90 --- /dev/null +++ b/tests/test_litellm_rust/support/recording_server.py @@ -0,0 +1,145 @@ +import asyncio +import copy +import json +import threading +import time +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass, field +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Final + +import pytest + + +@dataclass +class RecordedRequest: + method: str + path: str + headers: dict[str, str] + raw_body: bytes + body: object | None + + +@dataclass +class ResponseSpec: + body: object + status: int = 200 + headers: dict[str, str] = field(default_factory=dict) + delay: float = 0 + events: tuple[tuple[str, object], ...] = () + chunks: tuple[bytes, ...] = () + accepted: threading.Event | None = None + release_before_response: threading.Event | None = None + release: threading.Event | None = None + + +@dataclass +class RecordingServer: + server: ThreadingHTTPServer + requests: list[RecordedRequest] + responses: list[ResponseSpec] + default_response: ResponseSpec + expected_requests: int | None = 1 + + @property + def base_url(self) -> str: + host, port = self.server.server_address + return f"http://{host}:{port}" + + def enqueue(self, response: ResponseSpec) -> None: + self.responses.append(response) + + async def wait_for_requests(self, count: int) -> None: + async with asyncio.timeout(2): + while len(self.requests) < count: + await asyncio.sleep(0.01) + + +@contextmanager +def recording_service() -> Iterator[RecordingServer]: + requests: list[RecordedRequest] = [] + responses: list[ResponseSpec] = [] + + class Handler(BaseHTTPRequestHandler): + def _handle(self) -> None: + content_length: Final = int(self.headers.get("Content-Length", "0")) + raw_body: Final = self.rfile.read(content_length) if content_length else b"" + body: Final = json.loads(raw_body) if raw_body else None + requests.append( + RecordedRequest( + method=self.command, + path=self.path, + headers={name.lower(): value for name, value in self.headers.items()}, + raw_body=raw_body, + body=body, + ) + ) + response: Final = responses.pop(0) if responses else copy.deepcopy(recording_server.default_response) + if response.accepted is not None: + response.accepted.set() + if response.release_before_response is not None: + response.release_before_response.wait(timeout=10) + if response.delay: + time.sleep(response.delay) + payload: Final = ( + b"".join(response.chunks) + if response.chunks + else ( + b"".join( + f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in response.events + ) + if response.events + else json.dumps(response.body).encode() + ) + ) + self.send_response(response.status) + self.send_header( + "Content-Type", "text/event-stream" if response.events or response.chunks else "application/json" + ) + self.send_header("Content-Length", str(len(payload))) + for name, value in response.headers.items(): + self.send_header(name, value) + self.end_headers() + try: + if response.chunks: + for index, chunk in enumerate(response.chunks): + if index == 1 and response.release is not None: + response.release.wait(timeout=10) + self.wfile.write(chunk) + self.wfile.flush() + else: + self.wfile.write(payload) + except (BrokenPipeError, ConnectionResetError): + pass + + do_GET = _handle + do_POST = _handle + + def log_message(self, format: str, *args: object) -> None: + pass + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True) + thread.start() + try: + recording_server = RecordingServer( + server=server, + requests=requests, + responses=responses, + default_response=ResponseSpec(body={}), + ) + yield recording_server + finally: + server.shutdown() + server.server_close() + thread.join() + if recording_server.expected_requests is not None: + assert len(recording_server.requests) == recording_server.expected_requests + assert recording_server.responses == [] + + +@pytest.fixture +def recording_server() -> Iterator[RecordingServer]: + with recording_service() as server: + yield server diff --git a/tests/test_litellm_rust/support/requests.py b/tests/test_litellm_rust/support/requests.py new file mode 100644 index 00000000000..e34310b854a --- /dev/null +++ b/tests/test_litellm_rust/support/requests.py @@ -0,0 +1,113 @@ +from typing import Final + +import litellm +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.rust_bridge import ocr as native_ocr +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-4-5-20250929" +MESSAGES: Final = [{"role": "user", "content": "Hello"}] +MESSAGES_RESPONSE: Final = { + "id": "msg_native", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [{"type": "text", "text": "Hello from native Messages"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 4}, +} +CHAT_MODEL: Final = "anthropic/claude-opus-5" +CHAT_MESSAGES: Final = [ + {"role": "system", "content": "Keep the answer short"}, + {"role": "user", "content": "Earlier safe question"}, + {"role": "assistant", "content": "Earlier safe answer"}, + {"role": "user", "content": "Employee SSN: 078-05-1120"}, +] +CHAT_RESPONSE: Final = { + "id": "msg_chat_native", + "type": "message", + "role": "assistant", + "model": "claude-opus-5", + "content": [{"type": "text", "text": "Handled safely"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 8, "output_tokens": 2}, +} + + +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 native_ocr.ocr(ocr_arguments(server, **kwargs)) + + +async def call_native_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse: + return await native_ocr.aocr(ocr_arguments(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 + + +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": {"input_tokens": 5, "output_tokens": 4}, + }, + ), + ("message_stop", {"type": "message_stop"}), +) diff --git a/tests/test_litellm_rust/support/routes.py b/tests/test_litellm_rust/support/routes.py new file mode 100644 index 00000000000..509fd3eacf9 --- /dev/null +++ b/tests/test_litellm_rust/support/routes.py @@ -0,0 +1,135 @@ +import asyncio +from collections.abc import AsyncIterator, Mapping +from dataclasses import dataclass +from typing import Final, Literal + +import litellm +from tests.test_litellm_rust.support.callback_recorder import HookEvent, RecordingLogger +from tests.test_litellm_rust.support.requests import ( + MESSAGES, + MESSAGES_EVENTS, + MESSAGES_MODEL, + MESSAGES_RESPONSE, + OCR_RESPONSE, + call_aocr, + call_ocr, +) +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec + +RouteName = Literal["ocr-sync", "ocr-async", "messages", "messages-stream"] + + +@dataclass(frozen=True, slots=True) +class Route: + name: RouteName + call_type: str + provider_model: str + provider_response: Mapping[str, object] + expected_text: str + expected_cost: float + fires_async_hooks: bool + + async def invoke(self, server: RecordingServer, **kwargs: object) -> object: + match self.name: + case "ocr-sync": + return await asyncio.to_thread(call_ocr, server, **kwargs) + case "ocr-async": + return await call_aocr(server, **kwargs) + case "messages": + return await _call_messages(server, **kwargs) + case "messages-stream": + stream: Final = await self.open_stream(server, **kwargs) + return [chunk async for chunk in stream] + + async def open_stream(self, server: RecordingServer, **kwargs: object) -> AsyncIterator[object]: + if self.name != "messages-stream": + raise ValueError(f"{self.name} is not a streaming route") + stream: Final = await _call_messages(server, stream=True, **kwargs) + if not isinstance(stream, AsyncIterator): + raise TypeError(f"Expected async stream, got {type(stream).__name__}") + return stream + + def response_text(self, response: object) -> str: + match self.name: + case "ocr-sync" | "ocr-async": + pages: Final = getattr(response, "pages", None) + if not isinstance(pages, list) or not pages: + raise TypeError(f"Expected OCR response pages, got {type(response).__name__}") + markdown: Final = getattr(pages[0], "markdown", None) + if not isinstance(markdown, str): + raise TypeError(f"Expected OCR markdown, got {type(markdown).__name__}") + return markdown + case "messages": + if not isinstance(response, Mapping): + raise TypeError(f"Expected Messages mapping, got {type(response).__name__}") + content: Final = response.get("content") + if not isinstance(content, list) or not content or not isinstance(content[0], Mapping): + raise TypeError("Expected Messages text content") + text: Final = content[0].get("text") + if not isinstance(text, str): + raise TypeError(f"Expected Messages text, got {type(text).__name__}") + return text + case "messages-stream": + raise ValueError("Streaming response text is assembled by the stream consumer") + + +async def _call_messages(server: RecordingServer, **kwargs: object) -> object: + return await litellm.anthropic.messages.acreate( + model=MESSAGES_MODEL, + messages=MESSAGES, + max_tokens=64, + api_key="test-key", + api_base=server.base_url, + **kwargs, + ) + + +MESSAGES_COST: Final = 5 * 3e-06 + 4 * 1.5e-05 +OCR_COST: Final = 0.004 + +OCR_SYNC: Final = Route("ocr-sync", "ocr", "mistral-ocr-latest", OCR_RESPONSE, "native OCR response", OCR_COST, False) +OCR_ASYNC: Final = Route("ocr-async", "aocr", "mistral-ocr-latest", OCR_RESPONSE, "native OCR response", OCR_COST, True) +MESSAGES_ROUTE: Final = Route( + "messages", + "anthropic_messages", + "claude-sonnet-4-5-20250929", + MESSAGES_RESPONSE, + "Hello from native Messages", + MESSAGES_COST, + True, +) +MESSAGES_STREAM: Final = Route( + "messages-stream", + "anthropic_messages", + "claude-sonnet-4-5-20250929", + MESSAGES_RESPONSE, + "Hello from native Messages", + MESSAGES_COST, + True, +) +ALL_ROUTES: Final = (OCR_SYNC, OCR_ASYNC, MESSAGES_ROUTE, MESSAGES_STREAM) +ASYNC_ROUTES: Final = tuple(route for route in ALL_ROUTES if route.fires_async_hooks) +NON_STREAM_ASYNC_ROUTES: Final = (OCR_ASYNC, MESSAGES_ROUTE) + + +def route_id(route: Route) -> str: + return route.name + + +def provider_response(route: Route) -> ResponseSpec: + return ResponseSpec( + body=route.provider_response, + events=MESSAGES_EVENTS if route.name == "messages-stream" else (), + ) + + +async def wait_for_callback( + route: Route, + recorder: RecordingLogger, + outcome: Literal["success", "failure"] = "success", + count: int = 1, +) -> tuple[HookEvent, ...]: + event: Final = f"{'async_' if route.fires_async_hooks else ''}log_{outcome}_event" + if route.fires_async_hooks: + return await recorder.wait_for_async(event, count=count) + return recorder.wait_for(event, count=count) diff --git a/tests/test_litellm_rust/test_provenance.py b/tests/test_litellm_rust/test_provenance.py new file mode 100644 index 00000000000..aaebacdbb69 --- /dev/null +++ b/tests/test_litellm_rust/test_provenance.py @@ -0,0 +1,20 @@ +from typing import Final + +import pytest + +from litellm.rust_bridge.loader import get_native_bridge + +pytestmark = pytest.mark.requires_rust_extension + + +def test_native_route_catalogue_matches_extension_exports() -> None: + from litellm.rust_bridge.provenance import ( + RUST_NON_RESPONSE_ENTRYPOINTS, + RUST_RESPONSE_ENTRYPOINTS, + ) + + native: Final = get_native_bridge() + assert native is not None + route_exports: Final = frozenset(name for name in dir(native) if not name.startswith("_") and name[0].islower()) + + assert route_exports == RUST_RESPONSE_ENTRYPOINTS | RUST_NON_RESPONSE_ENTRYPOINTS diff --git a/tests/test_litellm_rust/test_support.py b/tests/test_litellm_rust/test_support.py new file mode 100644 index 00000000000..b70ac76a476 --- /dev/null +++ b/tests/test_litellm_rust/test_support.py @@ -0,0 +1,115 @@ +import asyncio +from pathlib import Path +from types import ModuleType +from typing import Final + +import pytest + +import litellm +from litellm.rust_bridge import configuration +from tests.test_litellm_rust import conftest +from tests.test_litellm_rust.conftest import _isolated_list # pyright: ignore[reportPrivateUsage] # exercise fixture restoration without mutating SDK state + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.mark.asyncio +@pytest.mark.parametrize("override", (None, False, True)) +@pytest.mark.parametrize("exception_type", (RuntimeError, asyncio.CancelledError)) +async def test_backend_restores_override_after_nested_failure( + monkeypatch: pytest.MonkeyPatch, + override: bool | None, + exception_type: type[RuntimeError] | type[asyncio.CancelledError], +) -> None: + monkeypatch.setenv("LITELLM_RUST", "1") + configuration.reset_rust_configuration() + if override is not None: + litellm.rust(override) + + async def fail_in_nested_backend() -> None: + async with conftest.isolated_backend("python"): + assert configuration.rust_enabled() is False + with pytest.raises(exception_type, match="scope failed"): + await fail_in_rust_backend() + assert configuration.rust_enabled() is False + raise exception_type("scope failed") + + async def fail_in_rust_backend() -> None: + async with conftest.isolated_backend("rust"): + assert configuration.rust_enabled() is True + raise exception_type("scope failed") + + with pytest.raises(exception_type, match="scope failed"): + await fail_in_nested_backend() + + assert configuration.rust_enabled() is (True if override is None else override) + monkeypatch.setenv("LITELLM_RUST", "0") + assert configuration.rust_enabled() is (False if override is None else override) + + +@pytest.mark.asyncio +async def test_backend_rejects_overlapping_tasks_without_changing_state() -> None: + async def competing_scope() -> None: + async with conftest.isolated_backend("rust"): + pytest.fail("overlapping backend scope was accepted") + + async with conftest.isolated_backend("python"): + callback: Final = object() + litellm.callbacks.append(callback) + with pytest.raises(RuntimeError, match="isolated_backend scopes cannot overlap across tasks"): + await asyncio.create_task(competing_scope()) + assert configuration.rust_enabled() is False + assert litellm.callbacks == [callback] + async with conftest.isolated_backend("rust"): + assert configuration.rust_enabled() is True + assert litellm.callbacks == [] + assert configuration.rust_enabled() is False + assert litellm.callbacks == [callback] + with pytest.raises(RuntimeError, match="isolated_backend scopes cannot overlap across tasks"): + await asyncio.create_task(competing_scope()) + + async def subsequent_scope() -> None: + async with conftest.isolated_backend("rust"): + assert configuration.rust_enabled() is True + + await asyncio.create_task(subsequent_scope()) + + +def test_callback_list_restores_identity_after_rebinding_and_failure() -> None: + container: Final = ModuleType("callback_registry") + callback: Final = object() + original: Final = [callback] + setattr(container, "callbacks", original) + + def failing_test() -> None: + with _isolated_list(container, "callbacks"): + assert getattr(container, "callbacks") is original + assert original == [] + original.append(object()) + setattr(container, "callbacks", [object()]) + raise RuntimeError("test failed") + + with pytest.raises(RuntimeError, match="test failed"): + failing_test() + + assert getattr(container, "callbacks") is original + assert original == [callback] + + +def test_rust_extension_gate_skips_only_rust_suite_items(monkeypatch: pytest.MonkeyPatch) -> None: + class Item: + def __init__(self, path: Path) -> None: + self.path = path + self.markers: list[object] = [] + + def add_marker(self, marker: object) -> None: + self.markers.append(marker) + + monkeypatch.setattr(conftest, "_parse_env_bool", lambda value: False) + rust_item: Final = Item(Path("tests/test_litellm_rust/ocr/test_dispatch.py")) + other_item: Final = Item(Path("tests/test_litellm/test_completion.py")) + + conftest.pytest_collection_modifyitems([rust_item, other_item]) # pyright: ignore[reportArgumentType] # lightweight items isolate hook selection + + assert rust_item.markers + assert not other_item.markers