mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
336 lines
12 KiB
Python
336 lines
12 KiB
Python
from collections.abc import Mapping
|
|
from typing import Final, NoReturn
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm.llms.base_llm.ocr.transformation import OCRResponse
|
|
from litellm.ocr.dispatch import (
|
|
_ADISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
|
|
_DISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
|
|
)
|
|
from litellm.rust_bridge import runtime
|
|
from litellm.rust_bridge.bindings import NativeBinding
|
|
from litellm.rust_bridge.catalog import Route, RouteRule, Rules
|
|
from litellm.rust_bridge.configuration import Rollout
|
|
from litellm.rust_bridge.ocr.entrypoints import NativeAocr, NativeOcr
|
|
from litellm.rust_bridge.public_call import NativeCall
|
|
from tests.test_litellm_rust.ocr.data import OCR_DOCUMENT, OCR_MODEL
|
|
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
|
|
from tests.test_litellm_rust.support.routing import assert_served_natively
|
|
|
|
pytestmark = pytest.mark.requires_rust_extension
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True])
|
|
@pytest.mark.parametrize("caching", [None, False, True])
|
|
async def test_ocr_does_not_depend_on_chat_cache(
|
|
server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool, caching: bool | None
|
|
) -> None:
|
|
from litellm.caching.caching import Cache
|
|
|
|
monkeypatch.setattr(litellm, "cache", Cache(type="local", supported_call_types=["completion", "acompletion"]))
|
|
arguments: Final = {
|
|
"model": OCR_MODEL,
|
|
"document": OCR_DOCUMENT,
|
|
"api_key": "test-key",
|
|
"api_base": server.base_url,
|
|
"caching": caching,
|
|
}
|
|
response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments)
|
|
assert response.pages[0].markdown == "native OCR response"
|
|
assert len(server.requests) == 1
|
|
|
|
|
|
@pytest.mark.parametrize("failure", ["invalid", "live", "schema"])
|
|
def test_native_projection_errors_never_select_python(
|
|
server: RecordingServer, monkeypatch: pytest.MonkeyPatch, failure: str
|
|
) -> None:
|
|
import dataclasses
|
|
import ssl
|
|
|
|
from litellm.rust_bridge import runtime
|
|
from litellm.rust_bridge.catalog import Route, RouteContext
|
|
from litellm.rust_bridge.host import settings
|
|
from litellm.rust_bridge.ocr.entrypoints import NATIVE_OCR
|
|
from litellm.rust_bridge.public_call import NativeCall
|
|
|
|
server.expected_requests = 0
|
|
snapshot: Final = dataclasses.replace(settings.http_settings(), user_agent=1)
|
|
if failure == "schema":
|
|
monkeypatch.setattr(settings, "http_settings", lambda: snapshot)
|
|
else:
|
|
monkeypatch.setattr(
|
|
litellm, "ssl_verify", ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) if failure == "live" else object()
|
|
)
|
|
request: Final = NativeCall(
|
|
args=(),
|
|
kwargs={},
|
|
base={
|
|
"model": "mistral/mistral-ocr-latest",
|
|
"document": OCR_DOCUMENT,
|
|
"api_key": "test-key",
|
|
"api_base": server.base_url,
|
|
"timeout": None,
|
|
"custom_llm_provider": "mistral",
|
|
"extra_headers": None,
|
|
},
|
|
)
|
|
|
|
def python_fallback() -> NoReturn:
|
|
pytest.fail("projection failures must not select Python")
|
|
|
|
with pytest.raises(RuntimeError if failure == "schema" else ValueError, match="http_settings"):
|
|
runtime.run(
|
|
RouteContext(Route.OCR, provider="mistral"),
|
|
binding=NATIVE_OCR,
|
|
native=lambda native: native(request),
|
|
python=python_fallback,
|
|
)
|
|
assert server.requests == []
|
|
|
|
|
|
MODELS: Final = ("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0")
|
|
|
|
|
|
IMAGE: Final = {"type": "image_url", "image_url": "data:image/png;base64,YWJj"}
|
|
|
|
|
|
BOX: Final = {"top_left_x": 0, "top_left_y": 0, "bottom_right_x": 32, "bottom_right_y": 32}
|
|
|
|
|
|
PAYLOAD: Final = {
|
|
"pages": [
|
|
{
|
|
"index": 4,
|
|
"markdown": {"content": "receipt", "images": [{"id": "image", "bounding_box": BOX, "description": "scan"}]},
|
|
},
|
|
{"markdown": {"content": "page two"}},
|
|
],
|
|
"meta": {"billed_units": {"pages": 3}},
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("model", MODELS)
|
|
async def test_ocr_health_check_uses_the_native_route(recording_server: RecordingServer, model: str) -> None:
|
|
recording_server.enqueue(ResponseSpec(body=PAYLOAD))
|
|
response: Final = await litellm.ahealth_check(
|
|
model_params={"model": model, "api_key": "test-key", "api_base": recording_server.base_url}, mode="ocr"
|
|
)
|
|
assert "error" not in response
|
|
assert_served_natively(recording_server)
|
|
assert recording_server.requests[0].body["document"]["image_url"].startswith("data:image/png;base64,")
|
|
|
|
|
|
RUST_RULES: Final[Rules] = (RouteRule(Route.OCR, Rollout.RUST_REQUIRED),)
|
|
|
|
|
|
def ocr_binding(native: NativeOcr | None) -> NativeBinding[NativeOcr]:
|
|
binding: Final[NativeBinding[NativeOcr]] = NativeBinding("ocr", validate=lambda _: None)
|
|
binding.override(native)
|
|
return binding
|
|
|
|
|
|
def aocr_binding(native: NativeAocr | None) -> NativeBinding[NativeAocr]:
|
|
binding: Final[NativeBinding[NativeAocr]] = NativeBinding("aocr", validate=lambda _: None)
|
|
binding.override(native)
|
|
return binding
|
|
|
|
|
|
def response(model: str = "mistral/mistral-ocr-latest") -> OCRResponse:
|
|
return OCRResponse(pages=[], model=model)
|
|
|
|
|
|
def test_native_receives_normalized_positional_request_and_original_call_shape() -> None:
|
|
document: Final[Mapping[str, object]] = {"type": "file", "file": b"pdf"}
|
|
timeout: Final = httpx.Timeout(30)
|
|
extra_headers: Final[dict[str, object]] = {"x-test": "1"}
|
|
pages: Final = [0, 2]
|
|
args: Final[tuple[object, ...]] = ("mistral/mistral-ocr-latest", document)
|
|
kwargs: Final[Mapping[str, object]] = {
|
|
"api_key": "test-key",
|
|
"api_base": "https://example.invalid",
|
|
"timeout": timeout,
|
|
"custom_llm_provider": "mistral",
|
|
"extra_headers": extra_headers,
|
|
"pages": pages,
|
|
}
|
|
captured: Final[list[tuple[NativeCall, tuple[object, ...], Mapping[str, object]]]] = []
|
|
expected: Final = response()
|
|
|
|
def native(
|
|
request: NativeCall,
|
|
) -> OCRResponse:
|
|
args: Final = request.args
|
|
kwargs: Final = request.kwargs
|
|
captured.append((request, args, kwargs))
|
|
return expected
|
|
|
|
result: Final = _DISPATCH.run(
|
|
args,
|
|
kwargs,
|
|
python=runtime.NO_PYTHON,
|
|
binding=ocr_binding(native),
|
|
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
|
rules=RUST_RULES,
|
|
)
|
|
|
|
request, call_args, call_kwargs = captured[0]
|
|
assert result is expected
|
|
assert request.resolved["model"] == "mistral/mistral-ocr-latest"
|
|
assert request.resolved["document"] is document
|
|
assert request.resolved["api_key"] == "test-key"
|
|
assert request.resolved["api_base"] == "https://example.invalid"
|
|
assert request.resolved["timeout"] is timeout
|
|
assert request.resolved["custom_llm_provider"] == "mistral"
|
|
assert request.resolved["extra_headers"] is extra_headers
|
|
assert request.kwargs == kwargs
|
|
assert request.kwargs["pages"] is pages
|
|
assert call_args is args
|
|
assert call_kwargs is kwargs
|
|
|
|
|
|
def test_native_preserves_keyword_model_and_document_in_original_call_shape() -> None:
|
|
document: Final[Mapping[str, object]] = {
|
|
"type": "document_url",
|
|
"document_url": "https://example.invalid/document.pdf",
|
|
}
|
|
pages: Final = [1]
|
|
args: Final[tuple[object, ...]] = ()
|
|
kwargs: Final[Mapping[str, object]] = {
|
|
"model": "mistral/mistral-ocr-latest",
|
|
"document": document,
|
|
"pages": pages,
|
|
}
|
|
captured: Final[list[tuple[NativeCall, tuple[object, ...], Mapping[str, object]]]] = []
|
|
expected: Final = response()
|
|
|
|
def native(
|
|
request: NativeCall,
|
|
) -> OCRResponse:
|
|
args: Final = request.args
|
|
kwargs: Final = request.kwargs
|
|
captured.append((request, args, kwargs))
|
|
return expected
|
|
|
|
result: Final = _DISPATCH.run(
|
|
args,
|
|
kwargs,
|
|
python=runtime.NO_PYTHON,
|
|
binding=ocr_binding(native),
|
|
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
|
rules=RUST_RULES,
|
|
)
|
|
|
|
request, call_args, call_kwargs = captured[0]
|
|
assert result is expected
|
|
assert request.resolved["model"] == "mistral/mistral-ocr-latest"
|
|
assert request.resolved["document"] is document
|
|
assert request.kwargs == kwargs
|
|
assert call_args is args
|
|
assert call_kwargs is kwargs
|
|
assert call_kwargs["model"] == "mistral/mistral-ocr-latest"
|
|
assert call_kwargs["document"] is document
|
|
|
|
|
|
def test_missing_native_binding_is_a_required_rust_error() -> None:
|
|
args: Final[tuple[object, ...]] = ("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"})
|
|
|
|
with pytest.raises(RuntimeError, match="Rust ocr bridge"):
|
|
_DISPATCH.run(
|
|
args,
|
|
{},
|
|
python=runtime.NO_PYTHON,
|
|
binding=ocr_binding(None),
|
|
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
|
rules=RUST_RULES,
|
|
)
|
|
|
|
|
|
def test_non_required_rule_cannot_be_served_without_python() -> None:
|
|
args: Final[tuple[object, ...]] = ("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"})
|
|
|
|
def native(
|
|
request: NativeCall,
|
|
) -> OCRResponse:
|
|
return response()
|
|
|
|
with pytest.raises(runtime.NoPythonImplementationError, match="RUST_REQUIRED"):
|
|
_DISPATCH.run(
|
|
args,
|
|
{},
|
|
python=runtime.NO_PYTHON,
|
|
binding=ocr_binding(native),
|
|
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
|
rules=(RouteRule(Route.OCR, Rollout.PYTHON_ONLY),),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("args", "kwargs", "message"),
|
|
(
|
|
(
|
|
("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"}),
|
|
{"model": "duplicate"},
|
|
r"ocr\(\) got multiple values for argument 'model'",
|
|
),
|
|
(
|
|
("mistral/mistral-ocr-latest",),
|
|
{},
|
|
r"ocr\(\) missing 1 required positional argument: 'document'",
|
|
),
|
|
),
|
|
)
|
|
def test_ocr_parser_errors_before_native(args: tuple[object, ...], kwargs: Mapping[str, object], message: str) -> None:
|
|
def native(
|
|
request: NativeCall,
|
|
) -> OCRResponse:
|
|
pytest.fail("OCR parser failures must not call native")
|
|
|
|
with pytest.raises(TypeError, match=message):
|
|
_DISPATCH.run(
|
|
args,
|
|
kwargs,
|
|
python=runtime.NO_PYTHON,
|
|
binding=ocr_binding(native),
|
|
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
|
rules=RUST_RULES,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("args", "kwargs", "message"),
|
|
(
|
|
(
|
|
("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"}),
|
|
{"model": "duplicate"},
|
|
r"aocr\(\) got multiple values for argument 'model'",
|
|
),
|
|
(
|
|
("mistral/mistral-ocr-latest",),
|
|
{},
|
|
r"aocr\(\) missing 1 required positional argument: 'document'",
|
|
),
|
|
),
|
|
)
|
|
async def test_aocr_parser_errors_before_native(
|
|
args: tuple[object, ...], kwargs: Mapping[str, object], message: str
|
|
) -> None:
|
|
async def native(
|
|
request: NativeCall,
|
|
) -> OCRResponse:
|
|
pytest.fail("OCR parser failures must not call native")
|
|
|
|
with pytest.raises(TypeError, match=message):
|
|
await _ADISPATCH.arun(
|
|
args,
|
|
kwargs,
|
|
python=runtime.NO_PYTHON,
|
|
binding=aocr_binding(native),
|
|
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
|
rules=RUST_RULES,
|
|
)
|