mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
test(ocr): isolate native bridge contracts
This commit is contained in:
Yujong Lee
parent
bc0ba47e33
commit
795d0ddc40
7 changed files with 91 additions and 194 deletions
|
|
@ -4,10 +4,10 @@ This suite covers OCR requests through LiteLLM's compiled Rust extension. OCR be
|
|||
|
||||
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. `ocr/test_dispatch.py` covers public sync and async native dispatch and explicit Python dispatch. `test_ocr.py` is the strict wire-level smoke test
|
||||
`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. `test_ocr.py` has the single public API dispatch test, covering enabled native dispatch and disabled Python dispatch
|
||||
|
||||
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 selects Rust for every test unless a fallback or parity case explicitly selects Python. Backend selection alone does not prove native execution because a public OCR request can fall back. Tests that claim native dispatch assert either the Rust response marker or a wire property that distinguishes the native client. `test_public_ocr_executes_the_compiled_extension_without_python_fallback` is strict and its server rejects Python HTTPX requests
|
||||
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 verifies the selected transport at the wire boundary
|
||||
|
||||
The OCR contract modules are non-strict expected failures until the retained callback implementation from #40070 lands. The compiled-extension OCR smoke test remains strict. Passing contract cases appear as XPASS so staging coverage stays visible
|
||||
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
|
||||
|
|
|
|||
|
|
@ -31,7 +31,6 @@ CALLBACK_ATTRIBUTES: Final = (
|
|||
)
|
||||
EXPECTED_FAILURE_REASONS: Final = {
|
||||
"ocr/test_callbacks.py": "requires the OCR callback lifecycle implementation from #40070",
|
||||
"ocr/test_dispatch.py": "requires the OCR native dispatch 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",
|
||||
}
|
||||
|
|
@ -67,12 +66,6 @@ def _rebound(container: object, attribute: str, value: object) -> Iterator[None]
|
|||
setattr(container, attribute, original)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _rust_mode(enabled: bool) -> Iterator[None]:
|
||||
with _rebound(_CONFIGURATION, "override", enabled):
|
||||
yield
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(autouse=True, loop_scope="function")
|
||||
async def isolate_ocr_test_state() -> AsyncIterator[None]:
|
||||
with ExitStack() as stack:
|
||||
|
|
@ -81,7 +74,7 @@ async def isolate_ocr_test_state() -> AsyncIterator[None]:
|
|||
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(_rust_mode(True))
|
||||
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:
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ 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,
|
||||
|
|
@ -17,7 +18,6 @@ from tests.test_litellm_rust.support.requests import (
|
|||
request_body,
|
||||
request_headers,
|
||||
)
|
||||
from tests.test_litellm_rust.support.response_marker import has_rust_response_marker
|
||||
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
|
||||
|
||||
pytestmark = pytest.mark.requires_rust_extension
|
||||
|
|
@ -29,11 +29,11 @@ def ocr_server(recording_server: RecordingServer) -> RecordingServer:
|
|||
return recording_server
|
||||
|
||||
|
||||
def call_ocr(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object):
|
||||
def call_native_ocr_with_callbacks(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object):
|
||||
return call_native_ocr(server, callbacks=callbacks, **kwargs)
|
||||
|
||||
|
||||
async def call_aocr(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object):
|
||||
async def call_native_aocr_with_callbacks(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object):
|
||||
return await call_native_aocr(server, callbacks=callbacks, **kwargs)
|
||||
|
||||
|
||||
|
|
@ -44,7 +44,7 @@ def test_native_ocr_pre_call_callback_receives_transformed_provider_request(ocr_
|
|||
def log_pre_api_call(self, model, _messages, kwargs):
|
||||
observations.append((model, copy.deepcopy(kwargs["additional_args"])))
|
||||
|
||||
call_ocr(ocr_server, [Observe()], pages=[0])
|
||||
call_native_ocr_with_callbacks(ocr_server, [Observe()], pages=[0])
|
||||
|
||||
assert len(observations) == 1
|
||||
model, additional_args = observations[0]
|
||||
|
|
@ -73,7 +73,7 @@ def test_native_ocr_pre_call_body_edit_reaches_next_callback_and_provider(
|
|||
def log_pre_api_call(self, model, _messages, kwargs):
|
||||
observed.append(copy.deepcopy(request_body(kwargs)))
|
||||
|
||||
call_ocr(ocr_server, [Edit(), Observe()], include_image_base64=False)
|
||||
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
|
||||
|
|
@ -90,19 +90,17 @@ def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(ocr_
|
|||
def log_pre_api_call(self, model, _messages, kwargs):
|
||||
observed.append(dict(request_headers(kwargs)))
|
||||
|
||||
call_ocr(ocr_server, [Edit(), Observe()])
|
||||
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("rust_enabled", [False, True], ids=["python-backend", "rust-backend"])
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
async def test_ocr_pre_call_nested_document_edit_updates_caller_callback_and_provider_references(
|
||||
ocr_server: RecordingServer, rust_enabled: bool, asynchronous: bool
|
||||
async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_and_provider_references(
|
||||
ocr_server: RecordingServer, asynchronous: bool
|
||||
) -> None:
|
||||
litellm.rust(rust_enabled)
|
||||
original: Final = dict(OCR_DOCUMENT)
|
||||
replacement_url: Final = "data:application/pdf;base64,ZGVm"
|
||||
retained: Final = []
|
||||
|
|
@ -124,13 +122,17 @@ async def test_ocr_pre_call_nested_document_edit_updates_caller_callback_and_pro
|
|||
"api_base": ocr_server.base_url,
|
||||
"callbacks": [Retain(), Edit()],
|
||||
}
|
||||
response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments)
|
||||
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 has_rust_response_marker(response) is rust_enabled
|
||||
assert response.pages[0].markdown == "native OCR response"
|
||||
|
||||
|
||||
def test_native_ocr_pre_call_document_replacement_does_not_mutate_original_document(
|
||||
|
|
@ -170,7 +172,7 @@ def test_native_ocr_pre_call_body_rebinding_is_visible_to_callbacks_but_not_prov
|
|||
def log_pre_api_call(self, model, _messages, kwargs):
|
||||
observed.append(request_body(kwargs))
|
||||
|
||||
call_ocr(ocr_server, [Rebind(), Observe()])
|
||||
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}
|
||||
|
|
@ -187,7 +189,7 @@ def test_native_ocr_callback_retained_body_observes_later_callback_mutation(ocr_
|
|||
def log_pre_api_call(self, model, _messages, kwargs):
|
||||
request_body(kwargs)["queued-edit"] = True
|
||||
|
||||
call_ocr(ocr_server, [QueuePayload(), Edit()])
|
||||
call_native_ocr_with_callbacks(ocr_server, [QueuePayload(), Edit()])
|
||||
|
||||
assert queued[0]["queued-edit"] is True
|
||||
|
||||
|
|
@ -205,7 +207,7 @@ def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(o
|
|||
terminal_tokens.put(kwargs["test-token"])
|
||||
finished.set()
|
||||
|
||||
call_ocr(ocr_server, [Stash()])
|
||||
call_native_ocr_with_callbacks(ocr_server, [Stash()])
|
||||
|
||||
assert finished.wait(10)
|
||||
assert terminal_tokens.get_nowait() is token
|
||||
|
|
@ -217,7 +219,7 @@ async def test_native_aocr_success_callback_receives_call_id_metadata_and_respon
|
|||
) -> None:
|
||||
recorder: Final = RecordingLogger()
|
||||
|
||||
await call_aocr(
|
||||
await call_native_aocr_with_callbacks(
|
||||
ocr_server,
|
||||
[recorder],
|
||||
litellm_call_id="ocr-success",
|
||||
|
|
@ -247,7 +249,7 @@ async def test_native_aocr_failure_callbacks_receive_call_type_error_and_no_resp
|
|||
observations.append(("async", kwargs["call_type"], kwargs["exception"], response_obj))
|
||||
|
||||
with pytest.raises(litellm.InternalServerError):
|
||||
await call_aocr(ocr_server, [Observe()])
|
||||
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)
|
||||
|
|
@ -261,7 +263,7 @@ async def test_native_aocr_pre_call_callback_runs_on_caller_loop_and_thread(ocr_
|
|||
caller_thread: Final = threading.current_thread()
|
||||
recorder: Final = RecordingLogger()
|
||||
|
||||
await call_aocr(ocr_server, [recorder])
|
||||
await call_native_aocr_with_callbacks(ocr_server, [recorder])
|
||||
|
||||
events: Final = await recorder.wait_for_async("log_pre_api_call")
|
||||
assert len(events) == 1
|
||||
|
|
@ -288,7 +290,7 @@ async def test_native_aocr_failure_callbacks_receive_state_added_by_pre_call_cal
|
|||
observed.append(("async", kwargs["request-token"]))
|
||||
|
||||
with pytest.raises(litellm.InternalServerError):
|
||||
await call_aocr(ocr_server, [TrackInFlightRequest()])
|
||||
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)
|
||||
|
|
@ -309,7 +311,7 @@ async def test_native_aocr_callback_error_does_not_mask_provider_error_or_skip_l
|
|||
raise RuntimeError("failure callback failed")
|
||||
|
||||
with pytest.raises(litellm.InternalServerError) as caught:
|
||||
await call_aocr(ocr_server, [FailingCallback(), recorder])
|
||||
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")
|
||||
|
|
@ -325,7 +327,7 @@ def test_native_ocr_dispatches_each_callback_phase_once_when_logger_is_registere
|
|||
) -> None:
|
||||
recorder: Final = RecordingLogger()
|
||||
|
||||
call_ocr(
|
||||
call_native_ocr_with_callbacks(
|
||||
ocr_server,
|
||||
[recorder, recorder],
|
||||
success_callback=[recorder],
|
||||
|
|
@ -341,14 +343,12 @@ def test_native_ocr_dispatches_each_callback_phase_once_when_logger_is_registere
|
|||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
async def test_public_azure_ocr_resolves_token_before_pre_call_on_caller_context(
|
||||
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
|
||||
from tests.test_litellm_rust.support.requests import call_aocr as public_aocr, call_ocr as public_ocr
|
||||
|
||||
context: Final = ContextVar("azure-token-context", default="missing")
|
||||
context.set("caller")
|
||||
caller_thread: Final = threading.current_thread()
|
||||
|
|
@ -377,29 +377,29 @@ async def test_public_azure_ocr_resolves_token_before_pre_call_on_caller_context
|
|||
"callbacks": [Edit()],
|
||||
}
|
||||
response: Final = (
|
||||
await public_aocr(ocr_server, **arguments) if asynchronous else public_ocr(ocr_server, **arguments)
|
||||
await call_native_aocr(ocr_server, **arguments)
|
||||
if asynchronous
|
||||
else call_native_ocr(ocr_server, **arguments)
|
||||
)
|
||||
assert has_rust_response_marker(response)
|
||||
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_public_azure_ocr_token_provider_can_make_nested_native_ocr_call(
|
||||
async def test_native_azure_ocr_token_provider_can_make_nested_native_ocr_call(
|
||||
ocr_server: RecordingServer,
|
||||
isolated_azure_auth: None,
|
||||
asynchronous: bool,
|
||||
) -> None:
|
||||
from tests.test_litellm_rust.support.requests import call_aocr as public_aocr, call_ocr as public_ocr
|
||||
|
||||
ocr_server.expected_requests = 2
|
||||
calls: Final = []
|
||||
|
||||
def provider() -> str:
|
||||
calls.append("token")
|
||||
nested: Final = public_ocr(ocr_server)
|
||||
assert has_rust_response_marker(nested)
|
||||
nested: Final = call_native_ocr(ocr_server)
|
||||
assert nested.pages[0].markdown == "native OCR response"
|
||||
return "outer-token"
|
||||
|
||||
arguments: Final = {
|
||||
|
|
@ -408,9 +408,11 @@ async def test_public_azure_ocr_token_provider_can_make_nested_native_ocr_call(
|
|||
"azure_ad_token_provider": provider,
|
||||
}
|
||||
response: Final = (
|
||||
await public_aocr(ocr_server, **arguments) if asynchronous else public_ocr(ocr_server, **arguments)
|
||||
await call_native_aocr(ocr_server, **arguments)
|
||||
if asynchronous
|
||||
else call_native_ocr(ocr_server, **arguments)
|
||||
)
|
||||
assert has_rust_response_marker(response)
|
||||
assert response.pages[0].markdown == "native OCR response"
|
||||
assert calls == ["token"]
|
||||
assert [request.headers["authorization"] for request in ocr_server.requests] == [
|
||||
"Bearer test-key",
|
||||
|
|
@ -419,12 +421,10 @@ async def test_public_azure_ocr_token_provider_can_make_nested_native_ocr_call(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_public_azure_ocr_calls_isolate_token_results_and_error(
|
||||
async def test_concurrent_native_azure_ocr_calls_isolate_token_results_and_error(
|
||||
ocr_server: RecordingServer,
|
||||
isolated_azure_auth: None,
|
||||
) -> None:
|
||||
from tests.test_litellm_rust.support.requests import call_aocr as public_aocr
|
||||
|
||||
ocr_server.expected_requests = 2
|
||||
|
||||
async def request(token: str, fail: bool) -> object:
|
||||
|
|
@ -433,7 +433,7 @@ async def test_concurrent_public_azure_ocr_calls_isolate_token_results_and_error
|
|||
raise ValueError(token)
|
||||
return token
|
||||
|
||||
return await public_aocr(
|
||||
return await call_native_aocr(
|
||||
ocr_server,
|
||||
model="azure_ai/mistral-ocr-latest",
|
||||
api_key=None,
|
||||
|
|
@ -446,10 +446,10 @@ async def test_concurrent_public_azure_ocr_calls_isolate_token_results_and_error
|
|||
request("second", False),
|
||||
return_exceptions=True,
|
||||
)
|
||||
assert has_rust_response_marker(responses[0])
|
||||
assert isinstance(responses[0], OCRResponse)
|
||||
assert isinstance(responses[1], litellm.APIConnectionError)
|
||||
assert "Failed to get Azure AD token: failed" in str(responses[1])
|
||||
assert has_rust_response_marker(responses[2])
|
||||
assert isinstance(responses[2], OCRResponse)
|
||||
assert sorted(request.headers["authorization"] for request in ocr_server.requests) == [
|
||||
"Bearer first",
|
||||
"Bearer second",
|
||||
|
|
@ -458,7 +458,7 @@ async def test_concurrent_public_azure_ocr_calls_isolate_token_results_and_error
|
|||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("outcome", ["success", "failure", "cancellation"])
|
||||
async def test_public_azure_ocr_releases_token_provider_after_terminal_outcome(
|
||||
async def test_native_azure_ocr_releases_token_provider_after_terminal_outcome(
|
||||
ocr_server: RecordingServer,
|
||||
isolated_azure_auth: None,
|
||||
outcome: str,
|
||||
|
|
@ -466,8 +466,6 @@ async def test_public_azure_ocr_releases_token_provider_after_terminal_outcome(
|
|||
import gc
|
||||
import weakref
|
||||
from tests.test_litellm_rust.support.callback_recorder import drain_logging
|
||||
from tests.test_litellm_rust.support.requests import call_aocr as public_aocr
|
||||
|
||||
class Provider:
|
||||
def __call__(self) -> str:
|
||||
if outcome == "failure":
|
||||
|
|
@ -480,13 +478,13 @@ async def test_public_azure_ocr_releases_token_provider_after_terminal_outcome(
|
|||
if outcome == "failure":
|
||||
ocr_server.expected_requests = 0
|
||||
with pytest.raises(litellm.APIConnectionError):
|
||||
await public_aocr(
|
||||
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(
|
||||
public_aocr(
|
||||
call_native_aocr(
|
||||
ocr_server,
|
||||
model="azure_ai/mistral-ocr-latest",
|
||||
api_key=None,
|
||||
|
|
@ -499,13 +497,13 @@ async def test_public_azure_ocr_releases_token_provider_after_terminal_outcome(
|
|||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
else:
|
||||
response: Final = await public_aocr(
|
||||
response: Final = await call_native_aocr(
|
||||
ocr_server,
|
||||
model="azure_ai/mistral-ocr-latest",
|
||||
api_key=None,
|
||||
azure_ad_token_provider=provider,
|
||||
)
|
||||
assert has_rust_response_marker(response)
|
||||
assert response.pages[0].markdown == "native OCR response"
|
||||
return reference
|
||||
|
||||
reference: Final = await invoke()
|
||||
|
|
|
|||
|
|
@ -1,71 +0,0 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
|
||||
from tests.test_litellm_rust.support.response_marker import has_rust_response_marker
|
||||
from tests.test_litellm_rust.support.requests import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE
|
||||
|
||||
pytestmark = pytest.mark.requires_rust_extension
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ocr_server(recording_server: RecordingServer) -> RecordingServer:
|
||||
recording_server.default_response = ResponseSpec(body=OCR_RESPONSE)
|
||||
return recording_server
|
||||
|
||||
|
||||
def test_sync_ocr_response_is_marked_as_rust_when_native_dispatch_is_enabled(ocr_server: RecordingServer) -> None:
|
||||
response: Final = litellm.ocr(
|
||||
model=OCR_MODEL,
|
||||
document=OCR_DOCUMENT,
|
||||
api_key="test-key",
|
||||
api_base=ocr_server.base_url,
|
||||
)
|
||||
|
||||
assert response.pages[0].markdown == "native OCR response"
|
||||
assert has_rust_response_marker(response)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_ocr_response_is_marked_as_rust_when_native_dispatch_is_enabled(
|
||||
ocr_server: RecordingServer,
|
||||
) -> None:
|
||||
response: Final = await litellm.aocr(
|
||||
model=OCR_MODEL,
|
||||
document=OCR_DOCUMENT,
|
||||
api_key="test-key",
|
||||
api_base=ocr_server.base_url,
|
||||
)
|
||||
|
||||
assert response.pages[0].markdown == "native OCR response"
|
||||
assert has_rust_response_marker(response)
|
||||
|
||||
|
||||
def test_public_ocr_uses_rust_for_file_document(ocr_server: RecordingServer) -> None:
|
||||
response: Final = litellm.ocr(
|
||||
model=OCR_MODEL,
|
||||
document={"type": "file", "file": b"%PDF-1.4", "mime_type": "application/pdf"},
|
||||
api_key="test-key",
|
||||
api_base=ocr_server.base_url,
|
||||
)
|
||||
|
||||
assert response.pages[0].markdown == "native OCR response"
|
||||
assert has_rust_response_marker(response)
|
||||
|
||||
|
||||
def test_public_ocr_response_has_no_rust_marker_when_native_dispatch_is_disabled(
|
||||
ocr_server: RecordingServer,
|
||||
) -> None:
|
||||
litellm.rust(False)
|
||||
|
||||
response: Final = litellm.ocr(
|
||||
model=OCR_MODEL,
|
||||
document=OCR_DOCUMENT,
|
||||
api_key="test-key",
|
||||
api_base=ocr_server.base_url,
|
||||
)
|
||||
|
||||
assert response.pages[0].markdown == "native OCR response"
|
||||
assert not has_rust_response_marker(response)
|
||||
|
|
@ -10,7 +10,7 @@ from litellm.types.guardrails import BlockedWord, ContentFilterAction, Guardrail
|
|||
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_aocr
|
||||
from tests.test_litellm_rust.support.requests import OCR_RESPONSE, call_native_aocr
|
||||
|
||||
pytestmark = pytest.mark.requires_rust_extension
|
||||
|
||||
|
|
@ -35,7 +35,7 @@ class ReplaceOCRMarkdown(CustomGuardrail):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_public_aocr_post_call_content_filter_blocks_matching_markdown(
|
||||
async def test_native_aocr_post_call_content_filter_blocks_matching_markdown(
|
||||
ocr_server: RecordingServer,
|
||||
) -> None:
|
||||
guardrail: Final = ContentFilterGuardrail(
|
||||
|
|
@ -46,21 +46,21 @@ async def test_public_aocr_post_call_content_filter_blocks_matching_markdown(
|
|||
litellm.callbacks.append(guardrail)
|
||||
|
||||
with pytest.raises(HTTPException, match="Content blocked") as blocked:
|
||||
await call_aocr(ocr_server, guardrails=[guardrail.guardrail_name])
|
||||
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_public_aocr_post_call_replacement_reaches_caller_and_success_callback(
|
||||
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_aocr(
|
||||
response: Final = await call_native_aocr(
|
||||
ocr_server,
|
||||
callbacks=[recorder],
|
||||
guardrails=[guardrail.guardrail_name],
|
||||
|
|
|
|||
|
|
@ -8,19 +8,17 @@ from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
|
|||
from tests.test_litellm_rust.support.requests import (
|
||||
OCR_DOCUMENT,
|
||||
OCR_RESPONSE,
|
||||
call_aocr,
|
||||
call_native_aocr,
|
||||
call_native_ocr,
|
||||
call_ocr as call_public_ocr,
|
||||
)
|
||||
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
|
||||
from tests.test_litellm_rust.support.response_marker import has_rust_response_marker
|
||||
|
||||
pytestmark = pytest.mark.requires_rust_extension
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
async def test_public_azure_ocr_uses_token_provider_result_as_bearer_token(
|
||||
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 = []
|
||||
|
|
@ -35,11 +33,10 @@ async def test_public_azure_ocr_uses_token_provider_result_as_bearer_token(
|
|||
"azure_ad_token_provider": token_provider,
|
||||
}
|
||||
response: Final = (
|
||||
await call_aocr(ocr_server, **arguments) if asynchronous else call_public_ocr(ocr_server, **arguments)
|
||||
await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
|
||||
)
|
||||
|
||||
assert calls == ["token"]
|
||||
assert has_rust_response_marker(response)
|
||||
assert response.pages[0].markdown == "native OCR response"
|
||||
assert_native_request(ocr_server)
|
||||
assert ocr_server.requests[0].headers["authorization"] == "Bearer callback-token"
|
||||
|
|
@ -51,17 +48,13 @@ def ocr_server(recording_server: RecordingServer) -> RecordingServer:
|
|||
return recording_server
|
||||
|
||||
|
||||
def call_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
|
||||
return call_native_ocr(server, **kwargs)
|
||||
|
||||
|
||||
def assert_native_request(server: RecordingServer) -> None:
|
||||
assert len(server.requests) == 1
|
||||
assert not server.requests[0].headers.get("user-agent", "").startswith("python-httpx")
|
||||
|
||||
|
||||
def test_native_ocr_sends_model_and_document_to_mistral_ocr_path(ocr_server: RecordingServer) -> None:
|
||||
response: Final = call_ocr(ocr_server)
|
||||
response: Final = call_native_ocr(ocr_server)
|
||||
|
||||
assert response.pages[0].markdown == "native OCR response"
|
||||
assert_native_request(ocr_server)
|
||||
|
|
@ -87,14 +80,14 @@ def test_native_ocr_prepares_file_document_like_python(ocr_server: RecordingServ
|
|||
|
||||
|
||||
def test_native_ocr_sends_pages_and_image_options(ocr_server: RecordingServer) -> None:
|
||||
call_ocr(ocr_server, pages=[0, 2], include_image_base64=True)
|
||||
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_ocr(ocr_server, extra_headers={"x-trace-id": "trace-1"})
|
||||
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"
|
||||
|
|
@ -115,7 +108,7 @@ def test_native_mistral_ocr_prefers_explicit_api_key_over_environment(
|
|||
) -> None:
|
||||
monkeypatch.setenv("MISTRAL_API_KEY", "environment-key")
|
||||
|
||||
call_ocr(ocr_server)
|
||||
call_native_ocr(ocr_server)
|
||||
|
||||
assert ocr_server.requests[0].headers["authorization"] == "Bearer test-key"
|
||||
|
||||
|
|
@ -149,7 +142,7 @@ def test_native_vertex_ocr_builds_path_from_project_and_location(ocr_server: Rec
|
|||
|
||||
|
||||
def test_native_ocr_normalizes_provider_response_model_and_usage(ocr_server: RecordingServer) -> None:
|
||||
response: Final = call_ocr(ocr_server)
|
||||
response: Final = call_native_ocr(ocr_server)
|
||||
|
||||
assert isinstance(response, OCRResponse)
|
||||
assert response.model == "mistral-ocr-latest"
|
||||
|
|
@ -160,7 +153,7 @@ def test_native_ocr_maps_provider_400_without_exposing_response_body(ocr_server:
|
|||
ocr_server.enqueue(ResponseSpec(body={"message": "invalid OCR request"}, status=400))
|
||||
|
||||
with pytest.raises(litellm.BadRequestError) as caught:
|
||||
call_ocr(ocr_server)
|
||||
call_native_ocr(ocr_server)
|
||||
|
||||
assert caught.value.status_code == 400
|
||||
assert caught.value.model == "mistral-ocr-latest"
|
||||
|
|
@ -172,14 +165,13 @@ def test_native_ocr_raises_transport_error_when_request_exceeds_timeout(ocr_serv
|
|||
ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.2))
|
||||
|
||||
with pytest.raises(RuntimeError, match="OCR transport failed"):
|
||||
call_ocr(ocr_server, timeout=0.01)
|
||||
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("backend", ["python", "rust"], ids=["python-backend", "rust-backend"])
|
||||
@pytest.mark.parametrize(
|
||||
"credentials, expected_token, expected_calls",
|
||||
[
|
||||
|
|
@ -189,16 +181,14 @@ def test_native_ocr_raises_transport_error_when_request_exceeds_timeout(ocr_serv
|
|||
],
|
||||
ids=["api-key-skips-provider", "provider-overrides-static-token", "header-overrides-provider"],
|
||||
)
|
||||
async def test_public_azure_ocr_applies_same_credential_precedence_on_python_and_rust(
|
||||
async def test_native_azure_ocr_applies_python_credential_precedence(
|
||||
ocr_server: RecordingServer,
|
||||
isolated_azure_auth: None,
|
||||
asynchronous: bool,
|
||||
backend: str,
|
||||
credentials: dict[str, object],
|
||||
expected_token: str,
|
||||
expected_calls: int,
|
||||
) -> None:
|
||||
litellm.rust(backend == "rust")
|
||||
calls: Final = []
|
||||
|
||||
def token_provider() -> str:
|
||||
|
|
@ -212,9 +202,9 @@ async def test_public_azure_ocr_applies_same_credential_precedence_on_python_and
|
|||
**credentials,
|
||||
}
|
||||
response: Final = (
|
||||
await call_aocr(ocr_server, **arguments) if asynchronous else call_public_ocr(ocr_server, **arguments)
|
||||
await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
|
||||
)
|
||||
assert has_rust_response_marker(response) == (backend == "rust")
|
||||
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}"
|
||||
|
|
@ -222,7 +212,7 @@ async def test_public_azure_ocr_applies_same_credential_precedence_on_python_and
|
|||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
async def test_public_azure_ocr_calls_token_provider_for_each_request(
|
||||
async def test_native_azure_ocr_calls_token_provider_for_each_request(
|
||||
ocr_server: RecordingServer,
|
||||
isolated_azure_auth: None,
|
||||
asynchronous: bool,
|
||||
|
|
@ -241,9 +231,11 @@ async def test_public_azure_ocr_calls_token_provider_for_each_request(
|
|||
"azure_ad_token_provider": token_provider,
|
||||
}
|
||||
response: Final = (
|
||||
await call_aocr(ocr_server, **arguments) if asynchronous else call_public_ocr(ocr_server, **arguments)
|
||||
await call_native_aocr(ocr_server, **arguments)
|
||||
if asynchronous
|
||||
else call_native_ocr(ocr_server, **arguments)
|
||||
)
|
||||
assert has_rust_response_marker(response)
|
||||
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",
|
||||
|
|
@ -257,20 +249,17 @@ class TokenAbort(BaseException):
|
|||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
@pytest.mark.parametrize("backend", ["python", "rust"], ids=["python-backend", "rust-backend"])
|
||||
@pytest.mark.parametrize(
|
||||
"failure",
|
||||
["non_string", "type_error", "ordinary", "abort"],
|
||||
ids=["non-string-result", "type-error", "value-error", "base-exception"],
|
||||
)
|
||||
async def test_public_azure_ocr_token_provider_failure_prevents_pre_call_callback_and_request(
|
||||
async def test_native_azure_ocr_token_provider_failure_prevents_pre_call_callback_and_request(
|
||||
ocr_server: RecordingServer,
|
||||
isolated_azure_auth: None,
|
||||
asynchronous: bool,
|
||||
backend: str,
|
||||
failure: str,
|
||||
) -> None:
|
||||
litellm.rust(backend == "rust")
|
||||
ocr_server.expected_requests = 0
|
||||
calls: Final = []
|
||||
recorder: Final = RecordingLogger()
|
||||
|
|
@ -294,7 +283,7 @@ async def test_public_azure_ocr_token_provider_failure_prevents_pre_call_callbac
|
|||
}
|
||||
expected: Final = TokenAbort if failure == "abort" else litellm.APIConnectionError
|
||||
with pytest.raises(expected) as caught:
|
||||
await call_aocr(ocr_server, **arguments) if asynchronous else call_public_ocr(ocr_server, **arguments)
|
||||
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
|
||||
|
|
@ -346,13 +335,10 @@ def test_native_azure_ocr_rejects_unsupported_configuration_before_token_or_call
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("backend", ["python", "rust"], ids=["python-backend", "rust-backend"])
|
||||
async def test_public_azure_ocr_validates_endpoint_before_calling_token_provider(
|
||||
async def test_native_azure_ocr_validates_endpoint_before_calling_token_provider(
|
||||
ocr_server: RecordingServer,
|
||||
isolated_azure_auth: None,
|
||||
backend: str,
|
||||
) -> None:
|
||||
litellm.rust(backend == "rust")
|
||||
ocr_server.expected_requests = 0
|
||||
calls: Final = []
|
||||
|
||||
|
|
@ -361,7 +347,7 @@ async def test_public_azure_ocr_validates_endpoint_before_calling_token_provider
|
|||
return "unused"
|
||||
|
||||
with pytest.raises(litellm.APIConnectionError, match="Missing Azure AI API Base"):
|
||||
await call_aocr(
|
||||
await call_native_aocr(
|
||||
ocr_server,
|
||||
model="azure_ai/mistral-ocr-latest",
|
||||
api_key=None,
|
||||
|
|
@ -373,20 +359,17 @@ async def test_public_azure_ocr_validates_endpoint_before_calling_token_provider
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("backend", ["python", "rust"], ids=["python-backend", "rust-backend"])
|
||||
async def test_public_azure_ocr_does_not_fall_back_to_static_token_after_empty_provider_result(
|
||||
async def test_native_azure_ocr_does_not_fall_back_to_static_token_after_empty_provider_result(
|
||||
ocr_server: RecordingServer,
|
||||
isolated_azure_auth: None,
|
||||
backend: str,
|
||||
) -> None:
|
||||
litellm.rust(backend == "rust")
|
||||
ocr_server.expected_requests = 0
|
||||
|
||||
def provider() -> str:
|
||||
return ""
|
||||
|
||||
with pytest.raises(litellm.APIConnectionError, match="Missing Azure AI credentials"):
|
||||
await call_aocr(
|
||||
await call_native_aocr(
|
||||
ocr_server,
|
||||
model="azure_ai/mistral-ocr-latest",
|
||||
api_key=None,
|
||||
|
|
@ -397,13 +380,10 @@ async def test_public_azure_ocr_does_not_fall_back_to_static_token_after_empty_p
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("backend", ["python", "rust"], ids=["python-backend", "rust-backend"])
|
||||
async def test_public_azure_ocr_ignores_falsey_token_provider_and_uses_static_token(
|
||||
async def test_native_azure_ocr_ignores_falsey_token_provider_and_uses_static_token(
|
||||
ocr_server: RecordingServer,
|
||||
isolated_azure_auth: None,
|
||||
backend: str,
|
||||
) -> None:
|
||||
litellm.rust(backend == "rust")
|
||||
calls: Final = []
|
||||
|
||||
class Provider:
|
||||
|
|
@ -414,26 +394,23 @@ async def test_public_azure_ocr_ignores_falsey_token_provider_and_uses_static_to
|
|||
calls.append("token")
|
||||
return "unused"
|
||||
|
||||
response: Final = await call_aocr(
|
||||
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 has_rust_response_marker(response) == (backend == "rust")
|
||||
assert response.pages[0].markdown == "native OCR response"
|
||||
assert calls == []
|
||||
assert ocr_server.requests[0].headers["authorization"] == "Bearer static-token"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("backend", ["python", "rust"], ids=["python-backend", "rust-backend"])
|
||||
async def test_public_azure_ocr_rejects_coroutine_returned_by_sync_token_provider(
|
||||
async def test_native_azure_ocr_rejects_coroutine_returned_by_sync_token_provider(
|
||||
ocr_server: RecordingServer,
|
||||
isolated_azure_auth: None,
|
||||
backend: str,
|
||||
) -> None:
|
||||
litellm.rust(backend == "rust")
|
||||
ocr_server.expected_requests = 0
|
||||
calls: Final = []
|
||||
|
||||
|
|
@ -448,7 +425,7 @@ async def test_public_azure_ocr_rejects_coroutine_returned_by_sync_token_provide
|
|||
|
||||
try:
|
||||
with pytest.raises(litellm.APIConnectionError, match="Azure AD token must be a string"):
|
||||
await call_aocr(
|
||||
await call_native_aocr(
|
||||
ocr_server, model="azure_ai/mistral-ocr-latest", api_key=None, azure_ad_token_provider=provider
|
||||
)
|
||||
finally:
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from typing import Final
|
|||
import pytest
|
||||
|
||||
import litellm
|
||||
from tests.test_litellm_rust.support.response_marker import has_rust_response_marker
|
||||
|
||||
pytestmark = pytest.mark.requires_rust_extension
|
||||
|
||||
|
|
@ -19,7 +20,7 @@ class RecordedOCRRequest:
|
|||
|
||||
|
||||
@pytest.fixture
|
||||
def native_only_ocr_server() -> Generator[tuple[ThreadingHTTPServer, list[RecordedOCRRequest]]]:
|
||||
def ocr_server() -> Generator[tuple[ThreadingHTTPServer, list[RecordedOCRRequest]]]:
|
||||
requests: Final[list[RecordedOCRRequest]] = []
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
|
|
@ -30,10 +31,6 @@ def native_only_ocr_server() -> Generator[tuple[ThreadingHTTPServer, list[Record
|
|||
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: Final = json.dumps(
|
||||
{
|
||||
"pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}],
|
||||
|
|
@ -61,10 +58,12 @@ def native_only_ocr_server() -> Generator[tuple[ThreadingHTTPServer, list[Record
|
|||
thread.join()
|
||||
|
||||
|
||||
def test_public_ocr_executes_the_compiled_extension_without_python_fallback(
|
||||
native_only_ocr_server: tuple[ThreadingHTTPServer, list[RecordedOCRRequest]],
|
||||
@pytest.mark.parametrize("rust_enabled", [True, False], ids=["enabled", "disabled"])
|
||||
def test_public_ocr_dispatches_according_to_rust_setting(
|
||||
ocr_server: tuple[ThreadingHTTPServer, list[RecordedOCRRequest]], rust_enabled: bool
|
||||
) -> None:
|
||||
server, requests = native_only_ocr_server
|
||||
litellm.rust(rust_enabled)
|
||||
server, requests = ocr_server
|
||||
address: Final = server.server_address
|
||||
host: Final = str(address[0])
|
||||
port: Final = int(address[1])
|
||||
|
|
@ -77,8 +76,9 @@ def test_public_ocr_executes_the_compiled_extension_without_python_fallback(
|
|||
)
|
||||
|
||||
assert response.pages[0].markdown == "native OCR response"
|
||||
assert has_rust_response_marker(response) is rust_enabled
|
||||
assert len(requests) == 1
|
||||
assert not requests[0].headers.get("user-agent", "").startswith("python-httpx")
|
||||
assert requests[0].headers.get("user-agent", "").startswith("python-httpx") == (not rust_enabled)
|
||||
assert requests[0].body == {
|
||||
"model": "mistral-ocr-latest",
|
||||
"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue