litellm/tests/test_litellm/rust_bridge/test_legacy_callbacks.py
Yujong Lee 63d994ade4 refactor(rust): run OCR through a route-neutral callback contract and a legacy Logging adapter
Extracted from #41733 without the router loop, the cache machine layer, streaming, or the
error, timeout and route-pruning work that moved to #41745

litellm-callbacks holds the contract a native call and its host share: Machine, HostOp,
CallEvent, the in-process run loop, and Passthrough, which is built only by comparing the
caller's inputs with the body the route sends, so a route can never mark a key it rewrote.
litellm-host-python (formerly python-interop) owns the CPython driver and the Execution
handle, and litellm-callbacks-legacy is the @client wrapper as the native call sees it:
function_setup, the deployment hooks, pre_call and post_call, the success and failure fan-out
and the deferred proxy release. OCR is the one route on it, and the old core and bridge
lifecycles are gone

The passthrough rule is the structural fix for the bug #41719 patched in core and #41716
reworks: an inlined remote document no longer counts as the caller's value, so the legacy
adapter never hands the caller's URL back into the body. core/tests/ocr/passthrough.rs pins
it for every route and document source, including that unchanged values stay passthrough,
and callbacks-legacy/tests/payload.rs pins the adapter side with a real pre_call callback

Python OCR integration tests that only exercised core behavior now live as Rust tests, so
tests/test_litellm_rust keeps the cases that need the full Python stack
2026-09-17 21:13:16 -07:00

79 lines
2.5 KiB
Python

import datetime
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
import pytest
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.rust_bridge.legacy_callbacks import check_limits, setup
_OCR_KWARGS: Final = MappingProxyType(
{
"model": "mistral/mistral-ocr-latest",
"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
}
)
@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
@pytest.mark.parametrize(
"cap, request_retry_count, refused",
[(5, 5, True), (5, 4, False), (0, 0, False), (0, 1, True)],
ids=[
"cap-above-four-reached",
"cap-above-four-not-reached",
"first-attempt-passes-cap-of-zero",
"cap-of-zero-refuses-first-retry",
],
)
def test_check_limits_reads_request_retry_count(
monkeypatch: pytest.MonkeyPatch, metadata_key: str, cap: int, request_retry_count: int, refused: bool
) -> None:
monkeypatch.setattr(litellm, "num_retries_per_request", cap)
monkeypatch.setattr(litellm, "max_budget", None)
kwargs: Final = {
"model": "mistral/mistral-ocr-latest",
metadata_key: {"request_retry_count": request_retry_count},
}
if refused:
with pytest.raises(RuntimeError, match="Max retries per request hit!"):
check_limits(kwargs)
else:
check_limits(kwargs)
def _supplied_logger() -> Logging:
return Logging(
model="mistral/mistral-ocr-latest",
messages=[],
stream=False,
call_type="aocr",
start_time=datetime.datetime.now(),
litellm_call_id="supplied",
function_id="supplied",
)
def test_setup_adopts_a_supplied_logger_as_caller_owned() -> None:
supplied: Final = _supplied_logger()
result: Final = setup(
"aocr", (), {**_OCR_KWARGS, "litellm_logging_obj": supplied}, datetime.datetime.now(), asynchronous=True
)
assert result.logger is supplied
assert result.bridge_owned is False
@pytest.mark.parametrize(
"call_type, kwargs",
[
("aocr", _OCR_KWARGS),
("aembedding", MappingProxyType({"model": "text-embedding-3-large", "input": ["hi"]})),
],
ids=["ocr", "embedding"],
)
def test_setup_owns_every_logger_it_builds(call_type: str, kwargs: Mapping[str, object]) -> None:
result: Final = setup(call_type, (), kwargs, datetime.datetime.now(), asynchronous=True)
assert result.bridge_owned is True
assert result.logger.litellm_call_id == result.kwargs["litellm_call_id"]