mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
4ecc55ec70
The Rust bridge imported httpx to construct the provider error response, which fails in the isolated wheel check where httpx is absent. Rust now raises RustUpstreamError with a headers attribute and the Python lifecycle wraps it in a typed UpstreamFailure carrying the httpx.Response before legacy mapping. Test helpers gained call_native so pytest.raises blocks hold a single call Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
808 lines
28 KiB
Python
808 lines
28 KiB
Python
import json
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
import httpx
|
|
import pytest
|
|
from pydantic import JsonValue
|
|
|
|
import litellm
|
|
from litellm.llms.base_llm.ocr.transformation import OCRResponse
|
|
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
|
|
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
|
|
from tests.test_litellm_rust.support.requests import (
|
|
OCR_DOCUMENT,
|
|
OCR_RESPONSE,
|
|
call_native,
|
|
call_native_aocr,
|
|
call_native_ocr,
|
|
)
|
|
|
|
pytestmark = pytest.mark.requires_rust_extension
|
|
|
|
|
|
@pytest.fixture(params=[False, True], ids=["python", "rust"])
|
|
def ocr_backend(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> bool:
|
|
enabled: Final = bool(request.param)
|
|
monkeypatch.setenv("LITELLM_RUST", "1" if enabled else "0")
|
|
return enabled
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
|
async def test_ocr_contract_upstream_status(
|
|
ocr_server: RecordingServer,
|
|
ocr_backend: bool,
|
|
asynchronous: bool,
|
|
) -> None:
|
|
upstream: Final = ResponseSpec(body={"detail": "invalid provider option"}, status=422)
|
|
ocr_server.enqueue(upstream)
|
|
arguments: Final = {
|
|
"model": "vertex_ai/mistral-ocr-latest",
|
|
"vertex_project": "test-project",
|
|
"vertex_location": "us-central1",
|
|
"num_retries": 0,
|
|
}
|
|
with pytest.raises(litellm.BadRequestError) as caught:
|
|
await call_native(ocr_server, asynchronous, **arguments)
|
|
assert caught.value.status_code == upstream.status
|
|
assert caught.value.response.status_code == upstream.status
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
|
@pytest.mark.parametrize("preserved", ["body", "headers"])
|
|
async def test_ocr_contract_provider_error_details(
|
|
ocr_server: RecordingServer,
|
|
ocr_backend: bool,
|
|
asynchronous: bool,
|
|
preserved: str,
|
|
) -> None:
|
|
payload: Final = {"message": "rate limited"}
|
|
headers: Final = {"Retry-After": "17", "X-Request-ID": "ocr-request-123", "X-Future-Header": "retained"}
|
|
ocr_server.enqueue(ResponseSpec(body=payload, status=429, headers=headers))
|
|
with pytest.raises(litellm.RateLimitError) as caught:
|
|
await call_native(ocr_server, asynchronous, num_retries=0)
|
|
response: Final = caught.value.response
|
|
assert isinstance(response, httpx.Response)
|
|
if preserved == "body":
|
|
assert response.content == json.dumps(payload).encode()
|
|
else:
|
|
for name, value in headers.items():
|
|
assert response.headers.get(name.lower()) == value
|
|
assert response.headers.get(name.upper()) == value
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
|
async def test_ocr_contract_invalid_response_format(
|
|
ocr_server: RecordingServer,
|
|
ocr_backend: bool,
|
|
asynchronous: bool,
|
|
) -> None:
|
|
ocr_server.expected_requests = 0
|
|
with pytest.raises(litellm.UnsupportedParamsError) as caught:
|
|
await call_native(ocr_server, asynchronous, req_format="bogus", num_retries=0)
|
|
assert caught.value.status_code == 400
|
|
for value in ("req_format", "bogus", "native", "litellm"):
|
|
assert value in str(caught.value)
|
|
assert ocr_server.requests == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
|
@pytest.mark.parametrize(
|
|
"document,field",
|
|
[
|
|
([], "document"),
|
|
({"document_url": "https://example.com/a.pdf"}, "type"),
|
|
({"type": "text"}, "type"),
|
|
],
|
|
)
|
|
async def test_ocr_contract_malformed_document_is_actionable(
|
|
ocr_server: RecordingServer,
|
|
ocr_backend: bool,
|
|
asynchronous: bool,
|
|
document: JsonValue,
|
|
field: str,
|
|
) -> None:
|
|
ocr_server.expected_requests = None
|
|
with pytest.raises(litellm.BadRequestError) as caught:
|
|
await call_native(ocr_server, asynchronous, document=document, num_retries=0)
|
|
assert caught.value.status_code == 400
|
|
assert field.lower() in str(caught.value).lower()
|
|
assert "NoneType: None" not in str(caught.value)
|
|
assert "indices must be" not in str(caught.value)
|
|
assert ocr_server.requests == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
|
@pytest.mark.parametrize("option,value,field", [("pages", [-1], "pages"), ("features", [1], "features")])
|
|
async def test_ocr_contract_azure_invalid_options_are_bad_requests(
|
|
ocr_server: RecordingServer,
|
|
ocr_backend: bool,
|
|
asynchronous: bool,
|
|
option: str,
|
|
value: JsonValue,
|
|
field: str,
|
|
) -> None:
|
|
ocr_server.expected_requests = 0
|
|
arguments: Final = {"model": "azure_ai/doc-intelligence/prebuilt-read", option: value, "num_retries": 0}
|
|
with pytest.raises(litellm.BadRequestError) as caught:
|
|
await call_native(ocr_server, asynchronous, **arguments)
|
|
assert caught.value.status_code == 400
|
|
assert field in str(caught.value)
|
|
assert ocr_server.requests == []
|
|
|
|
|
|
@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(
|
|
ocr_server: RecordingServer,
|
|
ocr_backend: bool,
|
|
asynchronous: bool,
|
|
model: str,
|
|
) -> None:
|
|
ocr_server.expected_requests = None
|
|
payload: Final = (
|
|
{"result": {"chunks": [{"content": "native OCR response"}]}, "usage": {"num_pages": 1}}
|
|
if model.startswith("reducto/")
|
|
else OCR_RESPONSE
|
|
)
|
|
ocr_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_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
|
|
)
|
|
assert response.pages[0].markdown == "native OCR response"
|
|
assert response.get_provider_native_response() == payload
|
|
assert len(ocr_server.requests) == 1
|
|
if ocr_backend:
|
|
assert_native_request(ocr_server)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
|
async def test_ocr_contract_unknown_reducto_model_reaches_provider(
|
|
ocr_server: RecordingServer,
|
|
ocr_backend: bool,
|
|
asynchronous: bool,
|
|
) -> None:
|
|
ocr_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_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
|
|
)
|
|
assert response.model == "future-parse-model"
|
|
assert response.pages[0].markdown == "future model response"
|
|
assert len(ocr_server.requests) == 1
|
|
assert ocr_server.requests[0].path == "/parse"
|
|
assert ocr_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(
|
|
ocr_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_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
|
|
)
|
|
|
|
assert calls == ["token"]
|
|
assert response.pages[0].markdown == "native OCR response"
|
|
assert_native_request(ocr_server)
|
|
assert ocr_server.requests[0].headers["authorization"] == "Bearer callback-token"
|
|
|
|
|
|
@pytest.fixture
|
|
def ocr_server(recording_server: RecordingServer) -> RecordingServer:
|
|
recording_server.default_response = ResponseSpec(body=OCR_RESPONSE)
|
|
return recording_server
|
|
|
|
|
|
def assert_native_request(server: RecordingServer) -> None:
|
|
assert len(server.requests) == 1
|
|
assert not server.requests[0].headers.get("user-agent", "").startswith("python-httpx")
|
|
|
|
|
|
def test_native_ocr_sends_model_and_document_to_mistral_ocr_path(ocr_server: RecordingServer) -> None:
|
|
response: Final = call_native_ocr(ocr_server)
|
|
|
|
assert response.pages[0].markdown == "native OCR response"
|
|
assert_native_request(ocr_server)
|
|
assert ocr_server.requests[0].path == "/v1/ocr"
|
|
assert ocr_server.requests[0].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT}
|
|
|
|
|
|
def test_native_ocr_prepares_file_document_like_python(ocr_server: RecordingServer) -> None:
|
|
response: Final = call_native_ocr(
|
|
ocr_server,
|
|
document={"type": "file", "file": b"%PDF-1.4", "mime_type": "application/pdf"},
|
|
)
|
|
|
|
assert response.pages[0].markdown == "native OCR response"
|
|
assert_native_request(ocr_server)
|
|
assert ocr_server.requests[0].body == {
|
|
"model": "mistral-ocr-latest",
|
|
"document": {
|
|
"type": "document_url",
|
|
"document_url": "data:application/pdf;base64,JVBERi0xLjQ=",
|
|
},
|
|
}
|
|
|
|
|
|
def test_native_ocr_reads_sdk_path_input(ocr_server: RecordingServer, tmp_path: Path) -> None:
|
|
document_path: Final = tmp_path / "document.pdf"
|
|
document_path.write_bytes(b"%PDF-1.4")
|
|
|
|
response: Final = call_native_ocr(
|
|
ocr_server,
|
|
document={"type": "file", "file": document_path},
|
|
)
|
|
|
|
assert response.pages[0].markdown == "native OCR response"
|
|
assert ocr_server.requests[0].body["document"] == {
|
|
"type": "document_url",
|
|
"document_url": "data:application/pdf;base64,JVBERi0xLjQ=",
|
|
}
|
|
|
|
|
|
def test_native_ocr_sends_pages_and_image_options(ocr_server: RecordingServer) -> None:
|
|
call_native_ocr(ocr_server, pages=[0, 2], include_image_base64=True)
|
|
|
|
assert ocr_server.requests[0].body["pages"] == [0, 2]
|
|
assert ocr_server.requests[0].body["include_image_base64"] is True
|
|
|
|
|
|
def test_native_ocr_merges_custom_headers_with_authorization(ocr_server: RecordingServer) -> None:
|
|
call_native_ocr(ocr_server, extra_headers={"x-trace-id": "trace-1"})
|
|
|
|
assert ocr_server.requests[0].headers["authorization"] == "Bearer test-key"
|
|
assert ocr_server.requests[0].headers["x-trace-id"] == "trace-1"
|
|
|
|
|
|
def test_native_mistral_ocr_uses_environment_api_key_when_argument_is_missing(
|
|
ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.setenv("MISTRAL_API_KEY", "environment-key")
|
|
|
|
call_native_ocr(ocr_server, api_key=None)
|
|
|
|
assert ocr_server.requests[0].headers["authorization"] == "Bearer environment-key"
|
|
|
|
|
|
def test_native_mistral_ocr_prefers_explicit_api_key_over_environment(
|
|
ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.setenv("MISTRAL_API_KEY", "environment-key")
|
|
|
|
call_native_ocr(ocr_server)
|
|
|
|
assert ocr_server.requests[0].headers["authorization"] == "Bearer test-key"
|
|
|
|
|
|
def test_native_azure_ocr_uses_environment_endpoint_and_api_key(
|
|
ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.setenv("AZURE_AI_API_KEY", "azure-key")
|
|
monkeypatch.setenv("AZURE_AI_API_BASE", ocr_server.base_url)
|
|
|
|
call_native_ocr(ocr_server, model="azure_ai/pixtral-12b-2409", api_key=None, api_base=None)
|
|
|
|
assert_native_request(ocr_server)
|
|
assert ocr_server.requests[0].path == "/providers/mistral/azure/ocr"
|
|
assert ocr_server.requests[0].headers["authorization"] == "Bearer azure-key"
|
|
|
|
|
|
def test_native_vertex_ocr_builds_path_from_project_and_location(ocr_server: RecordingServer) -> None:
|
|
call_native_ocr(
|
|
ocr_server,
|
|
model="vertex_ai/mistral-ocr-2505",
|
|
api_key="vertex-token",
|
|
vertex_project="project-1",
|
|
vertex_location="us-central1",
|
|
)
|
|
|
|
assert_native_request(ocr_server)
|
|
assert ocr_server.requests[0].path == (
|
|
"/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-2505:rawPredict"
|
|
)
|
|
|
|
|
|
def test_native_ocr_normalizes_provider_response_model_and_usage(ocr_server: RecordingServer) -> None:
|
|
response: Final = call_native_ocr(ocr_server)
|
|
|
|
assert isinstance(response, OCRResponse)
|
|
assert response.model == "mistral-ocr-latest"
|
|
assert response.usage_info.pages_processed == 1
|
|
|
|
|
|
def test_native_ocr_maps_provider_400_with_public_provider_details(ocr_server: RecordingServer) -> None:
|
|
ocr_server.enqueue(ResponseSpec(body={"message": "invalid OCR request"}, status=400))
|
|
|
|
with pytest.raises(litellm.BadRequestError) as caught:
|
|
call_native_ocr(ocr_server)
|
|
|
|
assert caught.value.status_code == 400
|
|
assert caught.value.model == "mistral-ocr-latest"
|
|
assert caught.value.llm_provider == "mistral"
|
|
assert "invalid OCR request" in str(caught.value)
|
|
|
|
|
|
def test_native_ocr_rejects_unknown_response_format_before_provider_request(ocr_server: RecordingServer) -> None:
|
|
ocr_server.expected_requests = 0
|
|
|
|
with pytest.raises(litellm.BadRequestError, match="Invalid `req_format`"):
|
|
call_native_ocr(ocr_server, req_format="raw")
|
|
|
|
assert ocr_server.requests == []
|
|
|
|
|
|
def test_ocr_raises_public_timeout_when_request_exceeds_timeout(ocr_server: RecordingServer) -> None:
|
|
litellm.rust(True)
|
|
ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.2))
|
|
|
|
with pytest.raises(litellm.Timeout):
|
|
call_native_ocr(ocr_server, timeout=0.01)
|
|
|
|
assert len(ocr_server.requests) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
|
@pytest.mark.parametrize(
|
|
"credentials, expected_token, expected_calls",
|
|
[
|
|
({"api_key": "resource-key"}, "resource-key", 0),
|
|
({"azure_ad_token": "static-token"}, "callback-1", 1),
|
|
({"extra_headers": {"Authorization": "Bearer override"}}, "override", 1),
|
|
],
|
|
ids=["api-key-skips-provider", "provider-overrides-static-token", "header-overrides-provider"],
|
|
)
|
|
async def test_native_azure_ocr_applies_python_credential_precedence(
|
|
ocr_server: RecordingServer,
|
|
isolated_azure_auth: None,
|
|
asynchronous: bool,
|
|
credentials: dict[str, object],
|
|
expected_token: str,
|
|
expected_calls: int,
|
|
) -> None:
|
|
calls: Final = []
|
|
|
|
def token_provider() -> str:
|
|
calls.append("token")
|
|
return f"callback-{len(calls)}"
|
|
|
|
arguments: Final = {
|
|
"model": "azure_ai/mistral-ocr-latest",
|
|
"api_key": None,
|
|
"azure_ad_token_provider": token_provider,
|
|
**credentials,
|
|
}
|
|
response: Final = (
|
|
await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
|
|
)
|
|
assert response.pages[0].markdown == "native OCR response"
|
|
assert len(calls) == expected_calls
|
|
assert len(ocr_server.requests) == 1
|
|
assert ocr_server.requests[0].headers["authorization"] == f"Bearer {expected_token}"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
|
async def test_native_azure_ocr_calls_token_provider_for_each_request(
|
|
ocr_server: RecordingServer,
|
|
isolated_azure_auth: None,
|
|
asynchronous: bool,
|
|
) -> None:
|
|
calls: Final = []
|
|
ocr_server.expected_requests = 2
|
|
|
|
def token_provider() -> str:
|
|
calls.append("token")
|
|
return f"callback-{len(calls)}"
|
|
|
|
for _ in range(2):
|
|
arguments: Final = {
|
|
"model": "azure_ai/mistral-ocr-latest",
|
|
"api_key": None,
|
|
"azure_ad_token_provider": token_provider,
|
|
}
|
|
response: Final = (
|
|
await call_native_aocr(ocr_server, **arguments)
|
|
if asynchronous
|
|
else call_native_ocr(ocr_server, **arguments)
|
|
)
|
|
assert response.pages[0].markdown == "native OCR response"
|
|
assert len(calls) == 2
|
|
assert [request.headers["authorization"] for request in ocr_server.requests] == [
|
|
"Bearer callback-1",
|
|
"Bearer callback-2",
|
|
]
|
|
|
|
|
|
class TokenAbort(BaseException):
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
|
@pytest.mark.parametrize(
|
|
"failure",
|
|
["non_string", "type_error", "ordinary", "abort"],
|
|
ids=["non-string-result", "type-error", "value-error", "base-exception"],
|
|
)
|
|
async def test_native_azure_ocr_token_provider_failure_prevents_pre_call_callback_and_request(
|
|
ocr_server: RecordingServer,
|
|
isolated_azure_auth: None,
|
|
asynchronous: bool,
|
|
failure: str,
|
|
) -> None:
|
|
ocr_server.expected_requests = 0
|
|
calls: Final = []
|
|
recorder: Final = RecordingLogger()
|
|
original: Final = {
|
|
"type_error": TypeError("token type"),
|
|
"ordinary": ValueError("token unavailable"),
|
|
"abort": TokenAbort("abort"),
|
|
}
|
|
|
|
def token_provider() -> object:
|
|
calls.append("token")
|
|
if failure == "non_string":
|
|
return 123
|
|
raise original[failure]
|
|
|
|
arguments: Final = {
|
|
"model": "azure_ai/mistral-ocr-latest",
|
|
"api_key": None,
|
|
"azure_ad_token_provider": token_provider,
|
|
"callbacks": [recorder],
|
|
}
|
|
expected: Final = TokenAbort if failure == "abort" else litellm.APIConnectionError
|
|
with pytest.raises(expected) as caught:
|
|
await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
|
|
assert calls == ["token"]
|
|
assert ocr_server.requests == []
|
|
assert "log_pre_api_call" not in recorder.names
|
|
if failure == "ordinary":
|
|
assert "Failed to get Azure AD token: token unavailable" in str(caught.value)
|
|
assert isinstance(caught.value.__context__, RuntimeError)
|
|
assert caught.value.__context__.__cause__ is original[failure]
|
|
elif failure == "abort":
|
|
assert caught.value is original[failure]
|
|
elif failure == "type_error":
|
|
assert caught.value.__context__ is original[failure]
|
|
else:
|
|
assert isinstance(caught.value.__context__, TypeError)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"configuration",
|
|
[{"azure_ad_token": "oidc/assertion", "client_id": "client", "tenant_id": "tenant"}],
|
|
ids=["invalid-oidc-assertion"],
|
|
)
|
|
def test_public_azure_ocr_maps_invalid_oidc_configuration_before_token_or_request(
|
|
ocr_server: RecordingServer,
|
|
isolated_azure_auth: None,
|
|
configuration: dict[str, object],
|
|
) -> None:
|
|
ocr_server.expected_requests = 0
|
|
calls: Final = []
|
|
recorder: Final = RecordingLogger()
|
|
|
|
def provider() -> str:
|
|
calls.append("token")
|
|
return "unused"
|
|
|
|
arguments: Final = {
|
|
"model": "azure_ai/mistral-ocr-latest",
|
|
"api_key": None,
|
|
"azure_ad_token_provider": provider,
|
|
"callbacks": [recorder],
|
|
**configuration,
|
|
}
|
|
with pytest.raises(litellm.APIConnectionError):
|
|
call_native_ocr(ocr_server, **arguments)
|
|
assert calls == []
|
|
assert "log_pre_api_call" not in recorder.names
|
|
assert ocr_server.requests == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_native_azure_ocr_validates_endpoint_before_calling_token_provider(
|
|
ocr_server: RecordingServer,
|
|
isolated_azure_auth: None,
|
|
) -> None:
|
|
ocr_server.expected_requests = 0
|
|
calls: Final = []
|
|
|
|
def provider() -> str:
|
|
calls.append("token")
|
|
return "unused"
|
|
|
|
with pytest.raises(litellm.APIConnectionError, match="Missing Azure AI API Base"):
|
|
await call_native_aocr(
|
|
ocr_server,
|
|
model="azure_ai/mistral-ocr-latest",
|
|
api_key=None,
|
|
api_base=None,
|
|
azure_ad_token_provider=provider,
|
|
)
|
|
assert calls == []
|
|
assert ocr_server.requests == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_native_azure_ocr_does_not_fall_back_to_static_token_after_empty_provider_result(
|
|
ocr_server: RecordingServer,
|
|
isolated_azure_auth: None,
|
|
) -> None:
|
|
ocr_server.expected_requests = 0
|
|
|
|
def provider() -> str:
|
|
return ""
|
|
|
|
with pytest.raises(litellm.APIConnectionError, match="Missing Azure AI credentials"):
|
|
await call_native_aocr(
|
|
ocr_server,
|
|
model="azure_ai/mistral-ocr-latest",
|
|
api_key=None,
|
|
azure_ad_token="static-token",
|
|
azure_ad_token_provider=provider,
|
|
)
|
|
assert ocr_server.requests == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_native_azure_ocr_ignores_falsey_token_provider_and_uses_static_token(
|
|
ocr_server: RecordingServer,
|
|
isolated_azure_auth: None,
|
|
) -> None:
|
|
calls: Final = []
|
|
|
|
class Provider:
|
|
def __bool__(self) -> bool:
|
|
return False
|
|
|
|
def __call__(self) -> str:
|
|
calls.append("token")
|
|
return "unused"
|
|
|
|
response: Final = await call_native_aocr(
|
|
ocr_server,
|
|
model="azure_ai/mistral-ocr-latest",
|
|
api_key=None,
|
|
azure_ad_token="static-token",
|
|
azure_ad_token_provider=Provider(),
|
|
)
|
|
assert response.pages[0].markdown == "native OCR response"
|
|
assert calls == []
|
|
assert ocr_server.requests[0].headers["authorization"] == "Bearer static-token"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_native_azure_ocr_rejects_coroutine_returned_by_sync_token_provider(
|
|
ocr_server: RecordingServer,
|
|
isolated_azure_auth: None,
|
|
) -> None:
|
|
ocr_server.expected_requests = 0
|
|
calls: Final = []
|
|
|
|
async def acquire() -> str:
|
|
calls.append("awaited")
|
|
return "unused"
|
|
|
|
coroutine: Final = acquire()
|
|
|
|
def provider() -> object:
|
|
return coroutine
|
|
|
|
try:
|
|
with pytest.raises(litellm.APIConnectionError, match="Azure AD token must be a string"):
|
|
await call_native_aocr(
|
|
ocr_server, model="azure_ai/mistral-ocr-latest", api_key=None, azure_ad_token_provider=provider
|
|
)
|
|
finally:
|
|
coroutine.close()
|
|
assert calls == []
|
|
assert ocr_server.requests == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True])
|
|
@pytest.mark.parametrize(
|
|
"override, expected_key",
|
|
[
|
|
({}, "credential-key"),
|
|
({"api_key": "explicit-key"}, "explicit-key"),
|
|
({"api_key": None}, "environment-key"),
|
|
],
|
|
ids=["inherit", "explicit", "explicit-none"],
|
|
)
|
|
async def test_native_ocr_inherits_named_credentials_without_overwriting_arguments(
|
|
ocr_server: RecordingServer,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
asynchronous: bool,
|
|
override: dict[str, object],
|
|
expected_key: str,
|
|
) -> None:
|
|
from litellm.models.credentials import CredentialItem
|
|
|
|
pages: Final = [0]
|
|
opaque: Final = object()
|
|
monkeypatch.setenv("MISTRAL_API_KEY", "environment-key")
|
|
monkeypatch.setattr(
|
|
litellm,
|
|
"credential_list",
|
|
[
|
|
CredentialItem(credential_name="other", credential_info={}, credential_values={"api_key": "wrong-key"}),
|
|
CredentialItem(
|
|
credential_name="ocr-test",
|
|
credential_info={},
|
|
credential_values={
|
|
"api_key": "credential-key",
|
|
"api_base": ocr_server.base_url,
|
|
"pages": pages,
|
|
"opaque": opaque,
|
|
},
|
|
),
|
|
CredentialItem(credential_name="ocr-test", credential_info={}, credential_values={"api_key": "later-key"}),
|
|
],
|
|
)
|
|
|
|
class Observer(RecordingLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
super().log_pre_api_call(model, messages, kwargs)
|
|
pages.append(2)
|
|
|
|
arguments: Final = {
|
|
"model": "mistral/mistral-ocr-latest",
|
|
"document": OCR_DOCUMENT,
|
|
"litellm_credential_name": "ocr-test",
|
|
"callbacks": [Observer()],
|
|
**override,
|
|
}
|
|
response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments)
|
|
assert response.pages[0].markdown == "native OCR response"
|
|
assert ocr_server.requests[0].headers["authorization"] == f"Bearer {expected_key}"
|
|
assert ocr_server.requests[0].body["pages"] == [0, 2]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"filename,field,mime",
|
|
[("scan.PNG", "image_url", "image/png"), ("document.pdf", "document_url", "application/pdf")],
|
|
)
|
|
def test_native_ocr_infers_mime_type_from_reader_name(
|
|
ocr_server: RecordingServer, filename: str, field: str, mime: str
|
|
) -> None:
|
|
from io import BytesIO
|
|
|
|
file: Final = BytesIO(b"abc")
|
|
file.name = filename
|
|
call_native_ocr(ocr_server, document={"type": "file", "file": file})
|
|
assert ocr_server.requests[0].body["document"] == {"type": field, field: f"data:{mime};base64,YWJj"}
|
|
|
|
|
|
def test_native_ocr_encodes_str_reader_results_as_utf8(ocr_server: RecordingServer) -> None:
|
|
from io import StringIO
|
|
|
|
call_native_ocr(ocr_server, document={"type": "file", "file": StringIO("abc"), "mime_type": "text/plain"})
|
|
assert ocr_server.requests[0].body["document"] == {
|
|
"type": "document_url",
|
|
"document_url": "data:text/plain;base64,YWJj",
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize("attribute", ["read", "name"])
|
|
def test_native_file_preparation_preserves_property_errors(ocr_server: RecordingServer, attribute: str) -> None:
|
|
ocr_server.expected_requests = 0
|
|
failure: Final = LookupError("file property failed")
|
|
|
|
class File:
|
|
def __getattribute__(self, name: str):
|
|
if name == attribute:
|
|
raise failure
|
|
return super().__getattribute__(name)
|
|
|
|
def read(self):
|
|
return b"abc"
|
|
|
|
with pytest.raises(litellm.APIConnectionError, match="file property failed") as caught:
|
|
call_native_ocr(ocr_server, document={"type": "file", "file": File()})
|
|
assert caught.value.__context__ is failure
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True])
|
|
async def test_native_file_preparation_preserves_reader_exception(
|
|
ocr_server: RecordingServer, asynchronous: bool
|
|
) -> None:
|
|
ocr_server.expected_requests = 0
|
|
failure: Final = RuntimeError("reader failed")
|
|
|
|
class Reader:
|
|
def read(self) -> bytes:
|
|
raise failure
|
|
|
|
document: Final = {"type": "file", "file": Reader()}
|
|
with pytest.raises(litellm.APIConnectionError, match="reader failed") as caught:
|
|
await call_native_aocr(ocr_server, document=document) if asynchronous else call_native_ocr(
|
|
ocr_server, document=document
|
|
)
|
|
assert caught.value.__context__ is failure
|
|
|
|
|
|
def test_native_file_preparation_rejects_unsupported_reader_results(ocr_server: RecordingServer) -> None:
|
|
ocr_server.expected_requests = 0
|
|
|
|
class Reader:
|
|
def read(self) -> int:
|
|
return 1
|
|
|
|
with pytest.raises(litellm.APIConnectionError, match="bytes or str") as caught:
|
|
call_native_ocr(ocr_server, document={"type": "file", "file": Reader()})
|
|
assert isinstance(caught.value.__context__, TypeError)
|
|
|
|
|
|
@pytest.mark.parametrize("kind", ["bytes", "path", "reader"])
|
|
def test_native_file_preparation_rejects_oversized_input(
|
|
ocr_server: RecordingServer, kind: str, tmp_path: Path
|
|
) -> None:
|
|
ocr_server.expected_requests = 0
|
|
limit: Final = 50 * 1024 * 1024
|
|
path: Final = tmp_path / "large.pdf"
|
|
with path.open("wb") as stream:
|
|
stream.truncate(limit + 1)
|
|
|
|
class Reader:
|
|
def read(self) -> bytes:
|
|
return b"a" * (limit + 1)
|
|
|
|
document: Final = {
|
|
"type": "file",
|
|
"file": path if kind == "path" else Reader() if kind == "reader" else b"a" * (limit + 1),
|
|
}
|
|
with pytest.raises(litellm.BadRequestError, match="exceeds the size limit"):
|
|
call_native_ocr(ocr_server, document=document)
|
|
|
|
|
|
def test_native_file_preparation_reports_missing_paths(ocr_server: RecordingServer, tmp_path: Path) -> None:
|
|
ocr_server.expected_requests = 0
|
|
missing: Final = tmp_path / "missing.pdf"
|
|
with pytest.raises(litellm.APIConnectionError, match=f"File not found: {missing}") as caught:
|
|
call_native_ocr(ocr_server, document={"type": "file", "file": missing})
|
|
assert isinstance(caught.value.__context__, FileNotFoundError)
|
|
|
|
|
|
def test_native_file_preparation_rejects_empty_readers(ocr_server: RecordingServer) -> None:
|
|
from io import BytesIO
|
|
|
|
ocr_server.expected_requests = 0
|
|
with pytest.raises(litellm.BadRequestError, match="File is empty"):
|
|
call_native_ocr(ocr_server, document={"type": "file", "file": BytesIO(b"")})
|