scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

Keep Rust OCR callback proof in OCR tests

Browse source
This commit is contained in:
Ishaan Jaff 2026-06-25 19:59:45 -07:00
parent 36809efcf8
commit 07dfb0ef35
No known key found for this signature in database
2 changed files with 34 additions and 48 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

62
litellm/rust_bridge/ocr.py
View file

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

20
tests/test_litellm/ocr/test_rust_bridge.py
View file

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