mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
7d79ccb7cc
* feat(rust-bridge): scope catalog rules with a contextvar Co-Authored-By: Yujong Lee <yujong@berri.ai> * test(rust-bridge): select native integration routes with scoped rules Co-Authored-By: Yujong Lee <yujong@berri.ai> * test(rust-bridge): run the Rust extension suite without LITELLM_RUST Co-Authored-By: Yujong Lee <yujong@berri.ai> * test(rust-bridge): use scoped catalog rules in dispatch tests Co-Authored-By: Yujong Lee <yujong@berri.ai> * refactor(rust-bridge): keep catalog rules a static global and override them once per test Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
104 lines
3.9 KiB
Python
104 lines
3.9 KiB
Python
from io import BytesIO
|
|
from typing import Final
|
|
|
|
import pytest
|
|
|
|
from tests.test_litellm_rust.ocr.data import (
|
|
OCR_DOCUMENT,
|
|
OCR_RESPONSE,
|
|
call_aocr,
|
|
call_ocr,
|
|
)
|
|
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], ids=["sync", "async"])
|
|
@pytest.mark.parametrize("model", ["mistral/mistral-ocr-latest", "azure_ai/mistral-ocr-latest", "reducto/parse-v3"])
|
|
async def test_ocr_contract_native_format_supported(
|
|
server: RecordingServer,
|
|
asynchronous: bool,
|
|
model: str,
|
|
) -> None:
|
|
server.expected_requests = None
|
|
payload: Final = (
|
|
{"result": {"chunks": [{"content": "native OCR response"}]}, "usage": {"num_pages": 1}}
|
|
if model.startswith("reducto/")
|
|
else OCR_RESPONSE
|
|
)
|
|
server.default_response = ResponseSpec(body=payload)
|
|
arguments: Final = {
|
|
"model": model,
|
|
"req_format": "native",
|
|
"num_retries": 0,
|
|
"document": {"type": "document_url", "document_url": "reducto://ready.pdf"}
|
|
if model.startswith("reducto/")
|
|
else OCR_DOCUMENT,
|
|
}
|
|
response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments)
|
|
assert response.pages[0].markdown == "native OCR response"
|
|
assert response.get_provider_native_response() == payload
|
|
assert len(server.requests) == 1
|
|
assert_served_natively(server, response)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
|
async def test_ocr_contract_unknown_reducto_model_reaches_provider(
|
|
server: RecordingServer,
|
|
asynchronous: bool,
|
|
) -> None:
|
|
server.default_response = ResponseSpec(body={"result": {"chunks": [{"content": "future model response"}]}})
|
|
arguments: Final = {
|
|
"model": "reducto/future-parse-model",
|
|
"document": {"type": "document_url", "document_url": "reducto://ready.pdf"},
|
|
"num_retries": 0,
|
|
}
|
|
response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments)
|
|
assert response.model == "future-parse-model"
|
|
assert response.pages[0].markdown == "future model response"
|
|
assert len(server.requests) == 1
|
|
assert server.requests[0].path == "/parse"
|
|
assert server.requests[0].body == {"input": "reducto://ready.pdf"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
|
async def test_native_azure_ocr_uses_token_provider_result_as_bearer_token(
|
|
server: RecordingServer, isolated_azure_auth: None, asynchronous: bool
|
|
) -> None:
|
|
calls: Final = []
|
|
|
|
def token_provider() -> str:
|
|
calls.append("token")
|
|
return "callback-token"
|
|
|
|
arguments: Final = {
|
|
"model": "azure_ai/mistral-ocr-latest",
|
|
"api_key": None,
|
|
"azure_ad_token_provider": token_provider,
|
|
}
|
|
response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments)
|
|
|
|
assert calls == ["token"]
|
|
assert response.pages[0].markdown == "native OCR response"
|
|
assert_served_natively(server, response)
|
|
assert server.requests[0].headers["authorization"] == "Bearer callback-token"
|
|
|
|
|
|
def test_native_ocr_encodes_python_file_input_and_drops_unknown_arguments(server: RecordingServer) -> None:
|
|
response: Final = call_ocr(
|
|
server,
|
|
document={"type": "file", "file": BytesIO(b"abc"), "mime_type": "image/png"},
|
|
opaque_extension=object(),
|
|
)
|
|
|
|
assert response.pages[0].markdown == "native OCR response"
|
|
assert_served_natively(server, response)
|
|
assert server.requests[0].body == {
|
|
"model": "mistral-ocr-latest",
|
|
"document": {"type": "image_url", "image_url": "data:image/png;base64,YWJj"},
|
|
}
|