diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index a592b128f2e..0efdc227335 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -19,8 +19,8 @@ class RustOcr(Protocol): extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, - callbacks: list[object] | None, - guardrails: list[object] | None, + callbacks: list[object] | None = None, + guardrails: list[object] | None = None, ) -> dict[str, object]: raise NotImplementedError @@ -36,8 +36,8 @@ class RustAocr(Protocol): extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, - callbacks: list[object] | None, - guardrails: list[object] | None, + callbacks: list[object] | None = None, + guardrails: list[object] | None = None, ) -> Awaitable[dict[str, object]]: raise NotImplementedError @@ -127,18 +127,21 @@ def ocr( rust_ocr = load_rust_ocr() if rust_ocr is None: return None - return rust_ocr( - model=model, - document=cast(dict[str, object], document), - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=cast(dict[str, object] | None, extra_headers), - optional_params=optional_params, - timeout_seconds=_timeout_to_seconds(timeout), - callbacks=callbacks, - guardrails=guardrails, - ) + kwargs: dict[str, object] = { + "model": model, + "document": cast(dict[str, object], document), + "api_key": api_key, + "api_base": api_base, + "custom_llm_provider": custom_llm_provider, + "extra_headers": cast(dict[str, object] | None, extra_headers), + "optional_params": optional_params, + "timeout_seconds": _timeout_to_seconds(timeout), + } + if callbacks: + kwargs["callbacks"] = callbacks + if guardrails: + kwargs["guardrails"] = guardrails + return rust_ocr(**kwargs) async def aocr( @@ -157,15 +160,18 @@ async def aocr( rust_aocr = load_rust_aocr() if rust_aocr is None: return None - return await rust_aocr( - model=model, - document=cast(dict[str, object], document), - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=cast(dict[str, object] | None, extra_headers), - optional_params=optional_params, - timeout_seconds=_timeout_to_seconds(timeout), - callbacks=callbacks, - guardrails=guardrails, - ) + kwargs: dict[str, object] = { + "model": model, + "document": cast(dict[str, object], document), + "api_key": api_key, + "api_base": api_base, + "custom_llm_provider": custom_llm_provider, + "extra_headers": cast(dict[str, object] | None, extra_headers), + "optional_params": optional_params, + "timeout_seconds": _timeout_to_seconds(timeout), + } + if callbacks: + kwargs["callbacks"] = callbacks + if guardrails: + kwargs["guardrails"] = guardrails + return await rust_aocr(**kwargs) diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 7fdd85323ef..acad249a2bb 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -53,8 +53,6 @@ class RecordingBridge: extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, - callbacks: list[object] | None = None, - guardrails: list[object] | None = None, ) -> dict[str, object]: self.calls.append( { @@ -66,8 +64,6 @@ class RecordingBridge: "extra_headers": extra_headers, "optional_params": optional_params, "timeout_seconds": timeout_seconds, - "callbacks": callbacks or [], - "guardrails": guardrails or [], } ) return dict(FAKE_OCR_RESPONSE) @@ -89,8 +85,6 @@ class RecordingAsyncBridge: extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, - callbacks: list[object] | None = None, - guardrails: list[object] | None = None, ) -> dict[str, object]: self.calls.append( { @@ -102,8 +96,6 @@ class RecordingAsyncBridge: "extra_headers": extra_headers, "optional_params": optional_params, "timeout_seconds": timeout_seconds, - "callbacks": callbacks or [], - "guardrails": guardrails or [], } ) return dict(FAKE_OCR_RESPONSE) @@ -120,8 +112,6 @@ class RaisingBridge: extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, - callbacks: list[object] | None = None, - guardrails: list[object] | None = None, ) -> dict[str, object]: raise RuntimeError("bridge failed") @@ -137,8 +127,6 @@ class RaisingAsyncBridge: extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, - callbacks: list[object] | None = None, - guardrails: list[object] | None = None, ) -> dict[str, object]: raise RuntimeError("bridge failed") @@ -226,11 +214,9 @@ def build_prepared_request( @pytest.fixture(autouse=True) def _reset_rust_flag(): """Keep the global toggle isolated between tests.""" - litellm.logging_callback_manager._reset_all_callbacks() rust_bridge.use_litellm_rust(False, ocr=None, aocr=None) rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL yield - litellm.logging_callback_manager._reset_all_callbacks() rust_bridge.use_litellm_rust(False, ocr=None, aocr=None) rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL @@ -414,8 +400,6 @@ def test_bridge_wrapper_forwards_prepared_args_and_wraps_response(): }, "optional_params": {"include_image_base64": True, "pages": [0]}, "timeout_seconds": 12.5, - "callbacks": [], - "guardrails": [], } @@ -445,8 +429,6 @@ async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response(): "extra_headers": None, "optional_params": {"vertex_project": "project-1"}, "timeout_seconds": 42.0, - "callbacks": [], - "guardrails": [], } @@ -480,8 +462,6 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): }, "optional_params": {"include_image_base64": True}, "timeout_seconds": 12.5, - "callbacks": [], - "guardrails": [], }