diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index c6901411167..17b6481a2bf 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -4,8 +4,22 @@ on: push: paths: - "litellm-rust/**" + - "litellm/rust_bridge/**" + - "tests/test_litellm_rust/**" + - "litellm/integrations/custom_logger.py" + - "litellm/litellm_core_utils/litellm_logging.py" + - "litellm/litellm_core_utils/logging_worker.py" + - "litellm/proxy/guardrails/**" + - "litellm/utils.py" + - "litellm/ocr/**" + - "litellm/llms/base_llm/ocr/**" + - "litellm/llms/custom_httpx/llm_http_handler.py" + - "tests/test_litellm/ocr/**" + - "tests/test_litellm/conftest.py" + - "Makefile" - ".cargo/**" - "pyproject.toml" + - "uv.lock" - "rust-toolchain.toml" - ".github/actions/setup-uv-with-retries/**" - ".github/scripts/smoke_test_native_wheel.py" @@ -20,8 +34,22 @@ on: - "litellm_**" paths: - "litellm-rust/**" + - "litellm/rust_bridge/**" + - "tests/test_litellm_rust/**" + - "litellm/integrations/custom_logger.py" + - "litellm/litellm_core_utils/litellm_logging.py" + - "litellm/litellm_core_utils/logging_worker.py" + - "litellm/proxy/guardrails/**" + - "litellm/utils.py" + - "litellm/ocr/**" + - "litellm/llms/base_llm/ocr/**" + - "litellm/llms/custom_httpx/llm_http_handler.py" + - "tests/test_litellm/ocr/**" + - "tests/test_litellm/conftest.py" + - "Makefile" - ".cargo/**" - "pyproject.toml" + - "uv.lock" - "rust-toolchain.toml" - ".github/actions/setup-uv-with-retries/**" - ".github/scripts/smoke_test_native_wheel.py" diff --git a/tests/test_litellm_rust/README.md b/tests/test_litellm_rust/README.md new file mode 100644 index 00000000000..4c117fb846b --- /dev/null +++ b/tests/test_litellm_rust/README.md @@ -0,0 +1,13 @@ +# Rust OCR bridge tests + +This suite covers OCR requests through LiteLLM's compiled Rust extension. OCR behavior tests live under `ocr/`; reusable OCR request, callback, and recording-server fixtures live under `support/` + +A test name identifies the OCR entrypoint or callback under test and its expected observable result. Parameter IDs state the execution mode or credential case. Keep multiple assertions together only when they prove one request, mutation, failure, or callback lifecycle behavior. Record callback observations and assert them after the callback returns because production logging can swallow callback exceptions + +`ocr/test_requests.py` covers provider payloads, file preparation, endpoint and credential resolution, normalized responses, errors, timeouts, and Azure token-provider behavior. `ocr/test_callbacks.py` covers OCR callback inputs, mutations, ordering, context, failure handling, concurrency, and cleanup. `ocr/test_guardrails.py` covers OCR post-call blocking and response replacement. These contract modules call the Rust bridge directly. `ocr/test_dispatch.py` has the single public API dispatch test, covering enabled native dispatch and disabled Python dispatch. `test_ocr.py` is a strict smoke test of the compiled Rust OCR transport + +Run `make test-rust-extension` as the acceptance command. It builds a fresh wheel, installs that wheel into a temporary environment, requires `LITELLM_RUST=1`, and runs this suite with isolated Python imports + +Collection fails when `LITELLM_RUST=1` is set but the compiled `_native` module cannot be imported. The autouse fixture isolates callback and configuration state but does not select a backend. Native contract tests call `litellm.rust_bridge.ocr` directly, while the strict dispatch test explicitly enables and disables Rust and records which OCR entrypoint runs + +The OCR contract modules are non-strict expected failures until the retained callback implementation from #40070 lands. The public dispatch test remains strict. Passing contract cases appear as XPASS so staging coverage stays visible diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py index 02d274bb405..b0c75d9d2f5 100644 --- a/tests/test_litellm_rust/conftest.py +++ b/tests/test_litellm_rust/conftest.py @@ -1,24 +1,132 @@ +import asyncio import os +from collections.abc import AsyncIterator, Generator, Iterator +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack, contextmanager +from types import ModuleType +from typing import Final, cast import pytest +import pytest_asyncio + +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.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", +) +EXPECTED_FAILURE_REASONS: Final = { + "ocr/test_callbacks.py": "requires the OCR callback lifecycle implementation from #40070", + "ocr/test_guardrails.py": "requires the OCR guardrail lifecycle implementation from #40070", + "ocr/test_requests.py": "requires the OCR request and Azure authentication implementation from #40070", +} -def pytest_collection_modifyitems(items): - rust_enabled = os.environ.get("LITELLM_RUST", "").strip().lower() in { - "1", - "true", - "yes", - "on", - } - if not rust_enabled: - skip = pytest.mark.skip(reason="requires LITELLM_RUST=1 and a compiled Rust extension") +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) -> Iterator[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) -> Iterator[None]: + original: Final[object] = getattr(container, attribute) + setattr(container, attribute, value) + try: + yield + finally: + setattr(container, attribute, original) + + +@pytest_asyncio.fixture(autouse=True, loop_scope="function") +async def isolate_ocr_test_state() -> AsyncIterator[None]: + 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] # no public callback-cache accessor + stack.enter_context(_rebound(utils, "callback_list", [])) # rebind-ok: isolate legacy callback registry + stack.enter_context(_rebound(litellm, "cache", None)) # test-quality-ok: isolate process-global cache + stack.enter_context(_rebound(_CONFIGURATION, "override", None)) + executor: Final = ThreadPoolExecutor(thread_name_prefix="rust-ocr-test-logging") + stack.enter_context(_rebound(utils, "executor", executor)) + try: + yield + finally: + try: + await drain_logging() + finally: + await asyncio.to_thread(executor.shutdown, wait=True) + await GLOBAL_LOGGING_WORKER.stop() + + +@pytest.fixture +def recording_server() -> Generator[RecordingServer]: + with recording_service() as server: + yield server + + +def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: + for item in items: + if "test_litellm_rust" not in item.path.parts: + continue + relative_path: Final = "/".join(item.path.parts[item.path.parts.index("test_litellm_rust") + 1 :]) + reason: Final = EXPECTED_FAILURE_REASONS.get(relative_path) + if reason is not None: + item.add_marker(pytest.mark.xfail(reason=reason, strict=False)) + + if not _parse_env_bool(os.environ.get("LITELLM_RUST")): + skip: Final = pytest.mark.skip(reason="requires LITELLM_RUST=1 and a compiled Rust extension") 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/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..b08446412c0 --- /dev/null +++ b/tests/test_litellm_rust/ocr/test_callbacks.py @@ -0,0 +1,513 @@ +import asyncio +import copy +import queue +import threading +from typing import Final + +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +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_native_ocr_with_callbacks(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object): + return call_native_ocr(server, callbacks=callbacks, **kwargs) + + +async def call_native_aocr_with_callbacks(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object): + return await call_native_aocr(server, callbacks=callbacks, **kwargs) + + +def test_native_ocr_pre_call_callback_receives_transformed_provider_request(ocr_server: RecordingServer) -> None: + observations: Final = [] + + class Observe(CustomLogger): + def log_pre_api_call(self, model, _messages, kwargs): + observations.append((model, copy.deepcopy(kwargs["additional_args"]))) + + call_native_ocr_with_callbacks(ocr_server, [Observe()], pages=[0]) + + assert len(observations) == 1 + model, additional_args = observations[0] + assert model == "mistral-ocr-latest" + assert additional_args["api_base"] == f"{ocr_server.base_url}/v1/ocr" + assert additional_args["complete_input_dict"] == { + "model": "mistral-ocr-latest", + "document": OCR_DOCUMENT, + "pages": [0], + } + + +@pytest.mark.parametrize("raise_after_edit", [False, True], ids=["callback-returns", "callback-raises"]) +def test_native_ocr_pre_call_body_edit_reaches_next_callback_and_provider( + ocr_server: RecordingServer, raise_after_edit: bool +) -> 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("pre-call callback failed") + + class Observe(CustomLogger): + def log_pre_api_call(self, model, _messages, kwargs): + observed.append(copy.deepcopy(request_body(kwargs))) + + call_native_ocr_with_callbacks(ocr_server, [Edit(), Observe()], include_image_base64=False) + + assert observed[0]["include_image_base64"] is True + assert ocr_server.requests[0].body["include_image_base64"] is True + + +def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(ocr_server: RecordingServer) -> None: + 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_native_ocr_with_callbacks(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("asynchronous", [False, True], ids=["sync", "async"]) +async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_and_provider_references( + ocr_server: RecordingServer, asynchronous: bool +) -> None: + original: Final = dict(OCR_DOCUMENT) + replacement_url: Final = "data:application/pdf;base64,ZGVm" + retained: Final = [] + aliases: Final = [] + + class Retain(CustomLogger): + def log_pre_api_call(self, model, _messages, kwargs): + aliases.append(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 call_native_aocr(ocr_server, **arguments) + if asynchronous + else call_native_ocr(ocr_server, **arguments) + ) + + assert aliases == [True] + assert retained[0]["document_url"] == replacement_url + assert original["document_url"] == replacement_url + assert ocr_server.requests[0].body["document"]["document_url"] == replacement_url + assert response.pages[0].markdown == "native OCR response" + + +def test_native_ocr_pre_call_document_replacement_does_not_mutate_original_document( + 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_native_ocr_pre_call_body_rebinding_is_visible_to_callbacks_but_not_provider( + 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_native_ocr_with_callbacks(ocr_server, [Rebind(), Observe()]) + + assert observed == [{"replacement": True}] + assert ocr_server.requests[0].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT} + + +def test_native_ocr_callback_retained_body_observes_later_callback_mutation(ocr_server: RecordingServer) -> None: + 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_native_ocr_with_callbacks(ocr_server, [QueuePayload(), Edit()]) + + assert queued[0]["queued-edit"] is True + + +def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(ocr_server: RecordingServer) -> None: + 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_native_ocr_with_callbacks(ocr_server, [Stash()]) + + assert finished.wait(10) + assert terminal_tokens.get_nowait() is token + + +@pytest.mark.asyncio +async def test_native_aocr_success_callback_receives_call_id_metadata_and_response( + ocr_server: RecordingServer, +) -> None: + recorder: Final = RecordingLogger() + + await call_native_aocr_with_callbacks( + 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_native_aocr_failure_callbacks_receive_call_type_error_and_no_response( + 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_native_aocr_with_callbacks(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_native_aocr_pre_call_callback_runs_on_caller_loop_and_thread(ocr_server: RecordingServer) -> None: + caller_loop: Final = asyncio.get_running_loop() + caller_thread: Final = threading.current_thread() + recorder: Final = RecordingLogger() + + await call_native_aocr_with_callbacks(ocr_server, [recorder]) + + 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_native_aocr_failure_callbacks_receive_state_added_by_pre_call_callback( + 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_native_aocr_with_callbacks(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_native_aocr_callback_error_does_not_mask_provider_error_or_skip_later_failure_callbacks( + ocr_server: RecordingServer, +) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500)) + recorder: Final = RecordingLogger() + + class FailingCallback(CustomLogger): + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + raise RuntimeError("failure callback failed") + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + raise RuntimeError("failure callback failed") + + with pytest.raises(litellm.InternalServerError) as caught: + await call_native_aocr_with_callbacks(ocr_server, [FailingCallback(), 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_native_ocr_dispatches_each_callback_phase_once_when_logger_is_registered_multiple_times( + ocr_server: RecordingServer, +) -> None: + recorder: Final = RecordingLogger() + + call_native_ocr_with_callbacks( + 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], ids=["sync", "async"]) +async def test_native_azure_ocr_resolves_token_before_pre_call_on_caller_context( + ocr_server: RecordingServer, + isolated_azure_auth: None, + asynchronous: bool, +) -> None: + from contextvars import ContextVar + 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 call_native_aocr(ocr_server, **arguments) + if asynchronous + else call_native_ocr(ocr_server, **arguments) + ) + assert response.pages[0].markdown == "native OCR response" + assert observations == ["token", "pre_call"] + assert ocr_server.requests[0].headers["authorization"] == "Bearer edited" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_native_azure_ocr_token_provider_can_make_nested_native_ocr_call( + ocr_server: RecordingServer, + isolated_azure_auth: None, + asynchronous: bool, +) -> None: + ocr_server.expected_requests = 2 + calls: Final = [] + + def provider() -> str: + calls.append("token") + nested: Final = call_native_ocr(ocr_server) + assert nested.pages[0].markdown == "native OCR response" + return "outer-token" + + arguments: Final = { + "model": "azure_ai/mistral-ocr-latest", + "api_key": None, + "azure_ad_token_provider": provider, + } + response: Final = ( + await call_native_aocr(ocr_server, **arguments) + if asynchronous + else call_native_ocr(ocr_server, **arguments) + ) + assert response.pages[0].markdown == "native OCR 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_native_azure_ocr_calls_isolate_token_results_and_error( + ocr_server: RecordingServer, + isolated_azure_auth: None, +) -> None: + 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 call_native_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 isinstance(responses[0], OCRResponse) + assert isinstance(responses[1], litellm.APIConnectionError) + assert "Failed to get Azure AD token: failed" in str(responses[1]) + assert isinstance(responses[2], OCRResponse) + assert sorted(request.headers["authorization"] for request in ocr_server.requests) == [ + "Bearer first", + "Bearer second", + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("outcome", ["success", "failure", "cancellation"]) +async def test_native_azure_ocr_releases_token_provider_after_terminal_outcome( + 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 + 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 call_native_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( + call_native_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 call_native_aocr( + ocr_server, + model="azure_ai/mistral-ocr-latest", + api_key=None, + azure_ad_token_provider=provider, + ) + assert response.pages[0].markdown == "native OCR 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..a6c76bc5d0e --- /dev/null +++ b/tests/test_litellm_rust/ocr/test_dispatch.py @@ -0,0 +1,44 @@ +from typing import Final +from unittest.mock import Mock + +import pytest + +import litellm +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.ocr import main as ocr_main +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec +from tests.test_litellm_rust.support.requests import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.fixture +def ocr_server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=OCR_RESPONSE) + return recording_server + + +@pytest.mark.parametrize("rust_enabled", [True, False], ids=["enabled", "disabled"]) +def test_public_ocr_dispatches_according_to_rust_setting( + ocr_server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + rust_enabled: bool, +) -> None: + rust_call: Final = Mock(wraps=ocr_main.rust_ocr_bridge.ocr) + python_call: Final = Mock(wraps=ocr_main.base_llm_http_handler.ocr) + monkeypatch.setattr(ocr_main.rust_ocr_bridge, "ocr", rust_call) + monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", python_call) + litellm.rust(rust_enabled) + + response: Final = litellm.ocr( + model=OCR_MODEL, + document=OCR_DOCUMENT, + api_key="test-key", + api_base=ocr_server.base_url, + ) + + assert isinstance(response, OCRResponse) + assert response.pages[0].markdown == "native OCR response" + assert rust_call.call_count == int(rust_enabled) + assert python_call.call_count == int(not rust_enabled) + assert len(ocr_server.requests) == 1 diff --git a/tests/test_litellm_rust/ocr/test_guardrails.py b/tests/test_litellm_rust/ocr/test_guardrails.py new file mode 100644 index 00000000000..f6fc1c7cb8d --- /dev/null +++ b/tests/test_litellm_rust/ocr/test_guardrails.py @@ -0,0 +1,74 @@ +from typing import Final + +import pytest +from fastapi import HTTPException + +import litellm +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail +from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks +from litellm.types.utils import CallTypes +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec +from tests.test_litellm_rust.support.requests import OCR_RESPONSE, call_native_aocr + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.fixture +def ocr_server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=OCR_RESPONSE) + return recording_server + + +class ReplaceOCRMarkdown(CustomGuardrail): + def __init__(self) -> None: + super().__init__( + guardrail_name="replace-ocr-markdown", event_hook=GuardrailEventHooks.post_call, default_on=True + ) + self.call_types: list[CallTypes] = [] + + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + self.call_types.append(call_type) + reviewed_page: Final = response.pages[0].model_copy(update={"markdown": "Reviewed OCR"}) + return response.model_copy(update={"pages": [reviewed_page]}) + + +@pytest.mark.asyncio +async def test_native_aocr_post_call_content_filter_blocks_matching_markdown( + ocr_server: RecordingServer, +) -> None: + guardrail: Final = ContentFilterGuardrail( + guardrail_name="block-native-ocr-markdown", + event_hook=GuardrailEventHooks.post_call, + blocked_words=[BlockedWord(keyword="native OCR response", action=ContentFilterAction.BLOCK)], + ) + litellm.callbacks.append(guardrail) + + with pytest.raises(HTTPException, match="Content blocked") as blocked: + await call_native_aocr(ocr_server, guardrails=[guardrail.guardrail_name]) + + assert blocked.value.status_code == 400 + assert len(ocr_server.requests) == 1 + + +@pytest.mark.asyncio +async def test_native_aocr_post_call_replacement_reaches_caller_and_success_callback( + ocr_server: RecordingServer, +) -> None: + guardrail: Final = ReplaceOCRMarkdown() + recorder: Final = RecordingLogger() + litellm.callbacks.append(guardrail) + + response: Final = await call_native_aocr( + ocr_server, + callbacks=[recorder], + guardrails=[guardrail.guardrail_name], + ) + success_events: Final = await recorder.wait_for_async("async_log_success_event") + + assert guardrail.call_types == [CallTypes.aocr] + assert response.pages[0].markdown == "Reviewed OCR" + assert len(success_events) == 1 + assert success_events[0].response.pages[0].markdown == "Reviewed OCR" + assert "guardrails" not in ocr_server.requests[0].body 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..d241fe08fc8 --- /dev/null +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -0,0 +1,434 @@ +from typing import Final + +import pytest + +import litellm +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.requests import ( + OCR_DOCUMENT, + OCR_RESPONSE, + call_native_aocr, + call_native_ocr, +) +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_native_azure_ocr_uses_token_provider_result_as_bearer_token( + 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_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments) + ) + + assert calls == ["token"] + 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 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_native_ocr_sends_model_and_document_to_mistral_ocr_path(ocr_server: RecordingServer) -> None: + response: Final = call_native_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_native_ocr_prepares_file_document_like_python(ocr_server: RecordingServer) -> None: + response: Final = call_native_ocr( + ocr_server, + document={"type": "file", "file": b"%PDF-1.4", "mime_type": "application/pdf"}, + ) + + assert response.pages[0].markdown == "native OCR response" + assert_native_request(ocr_server) + assert ocr_server.requests[0].body == { + "model": "mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "data:application/pdf;base64,JVBERi0xLjQ=", + }, + } + + +def test_native_ocr_sends_pages_and_image_options(ocr_server: RecordingServer) -> None: + call_native_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_native_ocr_merges_custom_headers_with_authorization(ocr_server: RecordingServer) -> None: + call_native_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_native_mistral_ocr_uses_environment_api_key_when_argument_is_missing( + 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_native_mistral_ocr_prefers_explicit_api_key_over_environment( + ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("MISTRAL_API_KEY", "environment-key") + + call_native_ocr(ocr_server) + + assert ocr_server.requests[0].headers["authorization"] == "Bearer test-key" + + +def test_native_azure_ocr_uses_environment_endpoint_and_api_key( + 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_native_vertex_ocr_builds_path_from_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_native_ocr_normalizes_provider_response_model_and_usage(ocr_server: RecordingServer) -> None: + response: Final = call_native_ocr(ocr_server) + + assert isinstance(response, OCRResponse) + assert response.model == "mistral-ocr-latest" + assert response.usage_info.pages_processed == 1 + + +def test_native_ocr_maps_provider_400_without_exposing_response_body(ocr_server: RecordingServer) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "invalid OCR request"}, status=400)) + + with pytest.raises(litellm.BadRequestError) as caught: + call_native_ocr(ocr_server) + + assert caught.value.status_code == 400 + assert caught.value.model == "mistral-ocr-latest" + assert caught.value.llm_provider == "mistral" + assert "invalid OCR request" not in str(caught.value) + + +def test_native_ocr_raises_transport_error_when_request_exceeds_timeout(ocr_server: RecordingServer) -> None: + ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.2)) + + with pytest.raises(RuntimeError, match="OCR transport failed"): + call_native_ocr(ocr_server, timeout=0.01) + + assert len(ocr_server.requests) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@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), + ], + ids=["api-key-skips-provider", "provider-overrides-static-token", "header-overrides-provider"], +) +async def test_native_azure_ocr_applies_python_credential_precedence( + ocr_server: RecordingServer, + isolated_azure_auth: None, + asynchronous: bool, + credentials: dict[str, object], + expected_token: str, + expected_calls: int, +) -> None: + 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_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments) + ) + assert response.pages[0].markdown == "native OCR response" + 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], ids=["sync", "async"]) +async def test_native_azure_ocr_calls_token_provider_for_each_request( + 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_native_aocr(ocr_server, **arguments) + if asynchronous + else call_native_ocr(ocr_server, **arguments) + ) + assert response.pages[0].markdown == "native OCR 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], ids=["sync", "async"]) +@pytest.mark.parametrize( + "failure", + ["non_string", "type_error", "ordinary", "abort"], + ids=["non-string-result", "type-error", "value-error", "base-exception"], +) +async def test_native_azure_ocr_token_provider_failure_prevents_pre_call_callback_and_request( + ocr_server: RecordingServer, + isolated_azure_auth: None, + asynchronous: bool, + failure: str, +) -> None: + ocr_server.expected_requests = 0 + calls: Final = [] + recorder: Final = RecordingLogger() + original: Final = { + "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_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments) + assert calls == ["token"] + assert ocr_server.requests == [] + assert "log_pre_api_call" not in recorder.names + if failure == "ordinary": + assert "Failed to get Azure AD token: token unavailable" in str(caught.value) + assert isinstance(caught.value.__context__, RuntimeError) + assert caught.value.__context__.__cause__ is original[failure] + 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"}, + ], + ids=["oidc-assertion", "document-intelligence-model"], +) +def test_native_azure_ocr_rejects_unsupported_configuration_before_token_or_callbacks( + 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 +async def test_native_azure_ocr_validates_endpoint_before_calling_token_provider( + ocr_server: RecordingServer, + isolated_azure_auth: None, +) -> None: + 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_native_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 +async def test_native_azure_ocr_does_not_fall_back_to_static_token_after_empty_provider_result( + ocr_server: RecordingServer, + isolated_azure_auth: None, +) -> None: + ocr_server.expected_requests = 0 + + def provider() -> str: + return "" + + with pytest.raises(litellm.APIConnectionError, match="Missing Azure AI credentials"): + await call_native_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 +async def test_native_azure_ocr_ignores_falsey_token_provider_and_uses_static_token( + ocr_server: RecordingServer, + isolated_azure_auth: None, +) -> None: + calls: Final = [] + + class Provider: + def __bool__(self) -> bool: + return False + + def __call__(self) -> str: + calls.append("token") + return "unused" + + response: Final = await call_native_aocr( + ocr_server, + model="azure_ai/mistral-ocr-latest", + api_key=None, + azure_ad_token="static-token", + azure_ad_token_provider=Provider(), + ) + assert response.pages[0].markdown == "native OCR response" + assert calls == [] + assert ocr_server.requests[0].headers["authorization"] == "Bearer static-token" + + +@pytest.mark.asyncio +async def test_native_azure_ocr_rejects_coroutine_returned_by_sync_token_provider( + ocr_server: RecordingServer, + isolated_azure_auth: None, +) -> None: + 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_native_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..6de011b1414 --- /dev/null +++ b/tests/test_litellm_rust/support/callback_recorder.py @@ -0,0 +1,101 @@ +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: + 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 + thread: threading.Thread + loop: asyncio.AbstractEventLoop | None + 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() + except RuntimeError: + loop = None + event: Final = HookEvent( + name=name, + call_type=details.get("call_type"), + thread=threading.current_thread(), + loop=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_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 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..5a9b9497c6e --- /dev/null +++ b/tests/test_litellm_rust/support/recording_server.py @@ -0,0 +1,108 @@ +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 + + +@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 + + +@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.delay: + time.sleep(response.delay) + payload: Final = json.dumps(response.body).encode() + self.send_response(response.status) + self.send_header("Content-Type", "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: + self.wfile.write(payload) + except (BrokenPipeError, ConnectionResetError): + pass + + 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 == [] diff --git a/tests/test_litellm_rust/support/requests.py b/tests/test_litellm_rust/support/requests.py new file mode 100644 index 00000000000..d681d752ffd --- /dev/null +++ b/tests/test_litellm_rust/support/requests.py @@ -0,0 +1,59 @@ +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}, +} + + +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 diff --git a/tests/test_litellm_rust/test_ocr.py b/tests/test_litellm_rust/test_ocr.py index d5b1fce1139..1293de9ee0e 100644 --- a/tests/test_litellm_rust/test_ocr.py +++ b/tests/test_litellm_rust/test_ocr.py @@ -1,31 +1,34 @@ import json import threading +from collections.abc import Generator +from dataclasses import dataclass from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Final import pytest -import litellm +from litellm.rust_bridge import ocr as rust_ocr_bridge pytestmark = pytest.mark.requires_rust_extension +@dataclass(frozen=True, slots=True) +class RecordedOCRRequest: + body: object + + @pytest.fixture -def ocr_server(): - requests = [] +def ocr_server() -> Generator[tuple[ThreadingHTTPServer, list[RecordedOCRRequest]]]: + requests: Final[list[RecordedOCRRequest]] = [] class Handler(BaseHTTPRequestHandler): - def do_POST(self): + def do_POST(self) -> None: requests.append( - { - "headers": {name.lower(): value for name, value in self.headers.items()}, - "body": json.loads(self.rfile.read(int(self.headers["Content-Length"]))), - } + RecordedOCRRequest( + body=json.loads(self.rfile.read(int(self.headers["Content-Length"]))), + ) ) - if self.headers.get("User-Agent", "").startswith("python-httpx"): - self.send_response(418) - self.end_headers() - return - response = json.dumps( + response: Final = json.dumps( { "pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}], "model": "mistral-ocr-latest", @@ -38,11 +41,11 @@ def ocr_server(): self.end_headers() self.wfile.write(response) - def log_message(self, format, *args): + def log_message(self, format: str, *args: object) -> None: pass - server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) - thread = threading.Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True) + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread: Final = threading.Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True) thread.start() try: yield server, requests @@ -52,21 +55,29 @@ def ocr_server(): thread.join() -def test_ocr_with_rust_extension(ocr_server): +def test_native_ocr_with_compiled_rust_extension( + ocr_server: tuple[ThreadingHTTPServer, list[RecordedOCRRequest]], +) -> None: server, requests = ocr_server - host, port = server.server_address + address: Final = server.server_address + host: Final = str(address[0]) + port: Final = int(address[1]) - response = litellm.ocr( - model="mistral/mistral-ocr-latest", + response: Final = rust_ocr_bridge.ocr( + model="mistral-ocr-latest", document={"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, api_key="test-key", api_base=f"http://{host}:{port}", + custom_llm_provider="mistral", + extra_headers=None, + optional_params={}, + timeout=None, ) - assert response.pages[0].markdown == "native OCR response" + assert response is not None + assert response["pages"][0]["markdown"] == "native OCR response" assert len(requests) == 1 - assert not requests[0]["headers"].get("user-agent", "").startswith("python-httpx") - assert requests[0]["body"] == { - "model": "mistral-ocr-latest", - "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + assert requests[0].body == { + "model": "mistral-ocr-latest", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, }