scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

test(ocr): isolate native Rust bridge contracts (#40410)

Browse source 
* test(rust): add retained callback suite as expected failures

* fix(tests): narrow retained callback xfails

* test(rust): clarify retained callback contracts

* test(ocr): clarify retained Rust contracts

* test(ocr): restore guardrail contracts

* test(ocr): require Rust file input parity

* test(ocr): isolate native bridge contracts

* fix(ci): repair Rust dispatch and OSV checks

* test(ocr): assert explicit backend dispatch
This commit is contained in:
yujonglee 2026-09-10 16:51:37 -07:00 • committed by GitHub
parent b294a51834
commit 61b0def867
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 1533 additions and 38 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
.github/workflows/test-rust.yml vendored
View file

@ -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"

13
tests/test_litellm_rust/README.md Normal file
View file

@ -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

134
tests/test_litellm_rust/conftest.py
View file

@ -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)

1
tests/test_litellm_rust/ocr/__init__.py Normal file
View file

@ -0,0 +1 @@

513
tests/test_litellm_rust/ocr/test_callbacks.py Normal file
View file

@ -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

44
tests/test_litellm_rust/ocr/test_dispatch.py Normal file
View file

@ -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

74
tests/test_litellm_rust/ocr/test_guardrails.py Normal file
View file

@ -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

434
tests/test_litellm_rust/ocr/test_requests.py Normal file
View file

@ -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 == []

1
tests/test_litellm_rust/support/__init__.py Normal file
View file

@ -0,0 +1 @@

101
tests/test_litellm_rust/support/callback_recorder.py Normal file
View file

@ -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

108
tests/test_litellm_rust/support/recording_server.py Normal file
View file

@ -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 == []

59
tests/test_litellm_rust/support/requests.py Normal file
View file

@ -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

61
tests/test_litellm_rust/test_ocr.py
View file

@ -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"},
}