From b4d5b418bc9a127e8b2c30f04bc41214dc3eccc3 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Thu, 10 Sep 2026 13:50:26 -0700 Subject: [PATCH] fix(ocr): retain callback mutations for native calls --- litellm/ocr/main.py | 45 +++++++++---- litellm/rust_bridge/ocr.py | 66 +++++++++++++------ tests/test_litellm_rust/README.md | 2 +- tests/test_litellm_rust/conftest.py | 22 +++++-- tests/test_litellm_rust/ocr/test_callbacks.py | 33 +++++----- .../support/callback_recorder.py | 2 +- 6 files changed, 116 insertions(+), 54 deletions(-) diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index df3f9d2096b..bed9c34c213 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -58,8 +58,8 @@ class _PreparedOCRRequest: class _PreparedRustOCRCall: api_key: str | None api_base: str | None + body: dict[str, object] headers: dict[str, object] - optional_params: dict[str, object] _RUST_OCR_PROVIDERS: Final = { @@ -261,15 +261,16 @@ def _prepare_rust_ocr_call( ) rust_api_base: Final = _rust_bridge_api_base(prepared_request, resolve_api_key) rust_optional_params: Final = _rust_bridge_optional_params(prepared_request, resolve_api_key) + body: Final[dict[str, object]] = { + "model": prepared_request.model, + "document": prepared_request.document, + **rust_optional_params, + } prepared_request.litellm_logging_obj.pre_call( input="OCR document processing", api_key=resolved_api_key, additional_args={ - "complete_input_dict": { - "model": prepared_request.model, - "document": prepared_request.document, - **rust_optional_params, - }, + "complete_input_dict": body, "api_base": resolved_complete_url, "headers": resolved_headers, }, @@ -277,11 +278,29 @@ def _prepare_rust_ocr_call( return _PreparedRustOCRCall( api_key=resolved_api_key, api_base=rust_api_base, + body=body, headers=cast(dict[str, object], resolved_headers), - optional_params=rust_optional_params, ) +def _rust_ocr_model(body: Mapping[str, object]) -> str: + model: Final = body.get("model") + if not isinstance(model, str): + raise TypeError("OCR callback produced a non-string model") + return model + + +def _rust_ocr_document(body: Mapping[str, object]) -> dict[str, object]: + document: Final = body.get("document") + if not isinstance(document, dict): + raise TypeError("OCR callback produced a non-dict document") + return cast(dict[str, object], document) # cast-ok: the native bridge validates the retained document fields + + +def _rust_ocr_optional_params(body: Mapping[str, object]) -> dict[str, object]: + return {name: value for name, value in body.items() if name not in {"model", "document"}} + + def _map_rust_ocr_error( error: Exception, prepared_request: _PreparedOCRRequest, @@ -321,13 +340,13 @@ def _run_rust_ocr( ) try: rust_response: Final = rust_ocr_bridge.ocr( - model=prepared_request.model, - document=prepared_request.document, + model=_rust_ocr_model(prepared.body), + document=_rust_ocr_document(prepared.body), api_key=prepared.api_key, api_base=prepared.api_base, custom_llm_provider=prepared_request.custom_llm_provider, extra_headers=prepared.headers, - optional_params=prepared.optional_params, + optional_params=_rust_ocr_optional_params(prepared.body), timeout=prepared_request.effective_timeout, ) except Exception as error: @@ -349,13 +368,13 @@ async def _run_rust_aocr( ) try: rust_response: Final = await rust_ocr_bridge.aocr( - model=prepared_request.model, - document=prepared_request.document, + model=_rust_ocr_model(prepared.body), + document=_rust_ocr_document(prepared.body), api_key=prepared.api_key, api_base=prepared.api_base, custom_llm_provider=prepared_request.custom_llm_provider, extra_headers=prepared.headers, - optional_params=prepared.optional_params, + optional_params=_rust_ocr_optional_params(prepared.body), timeout=prepared_request.effective_timeout, ) except Exception as error: diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index b7fdb5a98ef..26ca4ab4628 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -2,7 +2,7 @@ from __future__ import annotations -from collections.abc import Awaitable +from collections.abc import Awaitable, Callable from typing import Final, Protocol, cast # noqa: TID251 # native extension exposes dynamically typed callables import httpx @@ -61,17 +61,40 @@ def load_rust_aocr() -> RustAocr | None: return _AOCR.load() +def _public_ocr(arguments: dict[str, object]) -> object: + import litellm + + public_ocr: Final = cast( # cast-ok: the decorated public OCR callable accepts the retained argument bag + Callable[..., object], litellm.ocr + ) + return public_ocr(**arguments) + + +async def _public_aocr(arguments: dict[str, object]) -> object: + import litellm + + public_aocr: Final = cast( # cast-ok: the decorated public OCR callable accepts the retained argument bag + Callable[..., Awaitable[object]], litellm.aocr + ) + return await public_aocr(**arguments) + + def ocr( + arguments: dict[str, object] | None = None, *, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout: float | httpx.Timeout | None, -) -> dict[str, object] | None: + model: str | None = None, + document: dict[str, object] | None = None, + api_key: str | None = None, + api_base: str | None = None, + custom_llm_provider: str | None = None, + extra_headers: dict[str, object] | None = None, + optional_params: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, +) -> object: + if arguments is not None: + return _public_ocr(arguments) + if model is None or document is None or optional_params is None: + raise TypeError("Native OCR requires model, document, and optional_params") rust_ocr: Final = load_rust_ocr() if rust_ocr is None: return None @@ -88,16 +111,21 @@ def ocr( async def aocr( + arguments: dict[str, object] | None = None, *, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout: float | httpx.Timeout | None, -) -> dict[str, object] | None: + model: str | None = None, + document: dict[str, object] | None = None, + api_key: str | None = None, + api_base: str | None = None, + custom_llm_provider: str | None = None, + extra_headers: dict[str, object] | None = None, + optional_params: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, +) -> object: + if arguments is not None: + return await _public_aocr(arguments) + if model is None or document is None or optional_params is None: + raise TypeError("Native OCR requires model, document, and optional_params") rust_aocr: Final = load_rust_aocr() if rust_aocr is None: return None diff --git a/tests/test_litellm_rust/README.md b/tests/test_litellm_rust/README.md index 4c117fb846b..9ea24a95126 100644 --- a/tests/test_litellm_rust/README.md +++ b/tests/test_litellm_rust/README.md @@ -10,4 +10,4 @@ Run `make test-rust-extension` as the acceptance command. It builds a fresh whee 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 +The OCR callback contracts and supported request behavior are strict. Five request and guardrail cases remain non-strict expected failures for post-call route attribution, sanitized provider errors, timeout-specific mapping, and native Azure validation diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py index b0c75d9d2f5..0291ec68da0 100644 --- a/tests/test_litellm_rust/conftest.py +++ b/tests/test_litellm_rust/conftest.py @@ -30,9 +30,22 @@ CALLBACK_ATTRIBUTES: Final = ( "_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", + ( + "ocr/test_guardrails.py", + "test_native_aocr_post_call_content_filter_blocks_matching_markdown", + ): "requires OCR post-call guardrail route attribution", + ( + "ocr/test_requests.py", + "test_native_ocr_maps_provider_400_without_exposing_response_body", + ): "requires sanitized native OCR provider errors", + ( + "ocr/test_requests.py", + "test_native_ocr_raises_transport_error_when_request_exceeds_timeout", + ): "requires timeout-specific native OCR exception mapping", + ( + "ocr/test_requests.py", + "test_native_azure_ocr_rejects_unsupported_configuration_before_token_or_callbacks", + ): "requires native Azure validation before host authentication", } @@ -98,7 +111,8 @@ def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: 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) + test_name: Final = item.name.partition("[")[0] + reason: Final = EXPECTED_FAILURE_REASONS.get((relative_path, test_name)) if reason is not None: item.add_marker(pytest.mark.xfail(reason=reason, strict=False)) diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py index b08446412c0..eaeb40060e7 100644 --- a/tests/test_litellm_rust/ocr/test_callbacks.py +++ b/tests/test_litellm_rust/ocr/test_callbacks.py @@ -10,6 +10,7 @@ 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.recording_server import RecordingServer, ResponseSpec from tests.test_litellm_rust.support.requests import ( OCR_DOCUMENT, OCR_RESPONSE, @@ -18,7 +19,6 @@ from tests.test_litellm_rust.support.requests import ( request_body, request_headers, ) -from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec pytestmark = pytest.mark.requires_rust_extension @@ -41,7 +41,7 @@ def test_native_ocr_pre_call_callback_receives_transformed_provider_request(ocr_ observations: Final = [] class Observe(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + 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]) @@ -64,13 +64,13 @@ def test_native_ocr_pre_call_body_edit_reaches_next_callback_and_provider( observed: Final = [] class Edit(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + 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): + 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) @@ -83,11 +83,11 @@ def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(ocr_ observed: Final = [] class Edit(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + 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): + def log_pre_api_call(self, model, messages, kwargs): observed.append(dict(request_headers(kwargs))) call_native_ocr_with_callbacks(ocr_server, [Edit(), Observe()]) @@ -107,12 +107,12 @@ async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_ aliases: Final = [] class Retain(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + 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): + def log_pre_api_call(self, model, messages, kwargs): original["document_url"] = replacement_url arguments: Final = { @@ -143,7 +143,7 @@ def test_native_ocr_pre_call_document_replacement_does_not_mutate_original_docum retained: Final = [] class RetainAndReplace(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): body = request_body(kwargs) retained.append(body["document"]) body["document"] = replacement @@ -165,11 +165,11 @@ def test_native_ocr_pre_call_body_rebinding_is_visible_to_callbacks_but_not_prov observed: Final = [] class Rebind(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + 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): + def log_pre_api_call(self, model, messages, kwargs): observed.append(request_body(kwargs)) call_native_ocr_with_callbacks(ocr_server, [Rebind(), Observe()]) @@ -182,11 +182,11 @@ def test_native_ocr_callback_retained_body_observes_later_callback_mutation(ocr_ queued: Final = [] class QueuePayload(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + 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): + def log_pre_api_call(self, model, messages, kwargs): request_body(kwargs)["queued-edit"] = True call_native_ocr_with_callbacks(ocr_server, [QueuePayload(), Edit()]) @@ -200,7 +200,7 @@ def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(o finished: Final = threading.Event() class Stash(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + 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): @@ -280,7 +280,7 @@ async def test_native_aocr_failure_callbacks_receive_state_added_by_pre_call_cal observed: Final = [] class TrackInFlightRequest(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + 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): @@ -364,7 +364,7 @@ async def test_native_azure_ocr_resolves_token_before_pre_call_on_caller_context return "caller-token" class Edit(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + 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" @@ -465,6 +465,7 @@ async def test_native_azure_ocr_releases_token_provider_after_terminal_outcome( ) -> None: import gc import weakref + from tests.test_litellm_rust.support.callback_recorder import drain_logging class Provider: def __call__(self) -> str: diff --git a/tests/test_litellm_rust/support/callback_recorder.py b/tests/test_litellm_rust/support/callback_recorder.py index 6de011b1414..d3749ccc095 100644 --- a/tests/test_litellm_rust/support/callback_recorder.py +++ b/tests/test_litellm_rust/support/callback_recorder.py @@ -81,7 +81,7 @@ class RecordingLogger(CustomLogger): 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): + 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):