mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
338 lines
14 KiB
Python
338 lines
14 KiB
Python
from collections.abc import AsyncGenerator
|
|
from datetime import datetime
|
|
from io import BytesIO
|
|
from typing import Final
|
|
from unittest.mock import Mock
|
|
|
|
import httpx
|
|
import orjson
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
from litellm.litellm_core_utils.litellm_logging import Logging, use_custom_pricing_for_model
|
|
from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo
|
|
from litellm.llms.custom_httpx import llm_http_handler
|
|
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
|
from litellm.ocr.main import _prepare_ocr_request
|
|
from litellm.rust_bridge import bindings, configuration, runtime
|
|
from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR
|
|
from litellm.utils import ProviderConfigManager
|
|
|
|
|
|
@pytest.fixture
|
|
async def provider(monkeypatch: pytest.MonkeyPatch) -> AsyncGenerator[Mock]:
|
|
configuration.reset_rust_configuration()
|
|
monkeypatch.setenv("LITELLM_RUST", "0")
|
|
monkeypatch.setattr(bindings, "get_native_bridge", Mock(side_effect=AssertionError("Rust must not load")))
|
|
handler: Final = Mock(
|
|
return_value=httpx.Response(
|
|
200,
|
|
json={
|
|
"pages": [{"index": 0, "markdown": "parsed document"}],
|
|
"model": "mistral-ocr-latest",
|
|
"usage_info": {"pages_processed": 1},
|
|
},
|
|
)
|
|
)
|
|
transport: Final = httpx.MockTransport(handler)
|
|
with httpx.Client(transport=transport) as sync_client:
|
|
async with httpx.AsyncClient(transport=transport) as async_client:
|
|
sync_handler: Final = HTTPHandler(client=sync_client)
|
|
async_handler: Final = AsyncHTTPHandler()
|
|
await async_handler.client.aclose()
|
|
async_handler.client = async_client
|
|
monkeypatch.setattr(llm_http_handler, "_get_httpx_client", lambda: sync_handler)
|
|
monkeypatch.setattr(llm_http_handler, "get_async_httpx_client", lambda llm_provider: async_handler)
|
|
yield handler
|
|
NATIVE_OCR.reset()
|
|
NATIVE_AOCR.reset()
|
|
configuration.reset_rust_configuration()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("mode", ["sync", "async", "sync_async"])
|
|
@pytest.mark.parametrize("dispatch", ["disabled", "declined", "unavailable"])
|
|
async def test_python_request_response_and_callbacks(
|
|
provider: Mock, monkeypatch: pytest.MonkeyPatch, mode: str, dispatch: str
|
|
) -> None:
|
|
class Declined(Exception):
|
|
pass
|
|
|
|
if dispatch != "disabled":
|
|
monkeypatch.setenv("LITELLM_RUST", "1")
|
|
binding: Final = NATIVE_AOCR if mode == "async" else NATIVE_OCR
|
|
binding.override(Mock(side_effect=Declined()) if dispatch == "declined" else None)
|
|
monkeypatch.setattr(runtime, "native_exception_types", lambda: (Declined, RuntimeError))
|
|
logger: Final = Mock(spec=CustomLogger)
|
|
monkeypatch.setattr(litellm, "input_callback", [logger])
|
|
arguments: Final = {
|
|
"model": "mistral/mistral-ocr-latest",
|
|
"document": {"type": "file", "file": BytesIO(b"pdf"), "mime_type": "application/pdf"},
|
|
"api_key": "test-key",
|
|
"api_base": "https://ocr.test/v1",
|
|
"timeout": 7.0,
|
|
"pages": [0, 2],
|
|
"include_image_base64": True,
|
|
"extra_headers": {"x-test-header": "preserved"},
|
|
}
|
|
|
|
async def call() -> OCRResponse:
|
|
if mode == "async":
|
|
return await litellm.aocr(**arguments)
|
|
if mode == "sync_async":
|
|
from litellm.litellm_core_utils.litellm_logging import Logging
|
|
|
|
logging_obj: Final = Logging(
|
|
model=arguments["model"],
|
|
messages=[],
|
|
stream=False,
|
|
call_type="aocr",
|
|
start_time=datetime.now(),
|
|
litellm_call_id="test-call",
|
|
function_id="test-function",
|
|
)
|
|
return await litellm.ocr(**arguments, aocr=True, litellm_logging_obj=logging_obj)
|
|
return litellm.ocr(**arguments)
|
|
|
|
response: Final = await call()
|
|
assert response.pages[0].markdown == "parsed document"
|
|
assert response.usage_info.pages_processed == 1
|
|
assert provider.call_count == 1
|
|
request: Final = provider.call_args.args[0]
|
|
assert str(request.url) == "https://ocr.test/v1/ocr"
|
|
assert request.headers["authorization"] == "Bearer test-key"
|
|
assert request.headers["x-test-header"] == "preserved"
|
|
assert request.extensions["timeout"] == {"connect": 7.0, "read": 7.0, "write": 7.0, "pool": 7.0}
|
|
assert orjson.loads(request.content) == {
|
|
"model": "mistral-ocr-latest",
|
|
"document": {"type": "document_url", "document_url": "data:application/pdf;base64,cGRm"},
|
|
"pages": [0, 2],
|
|
"include_image_base64": True,
|
|
}
|
|
assert logger.log_pre_api_call.call_count == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True])
|
|
async def test_python_provider_errors_keep_public_exception(provider: Mock, asynchronous: bool) -> None:
|
|
provider.return_value = httpx.Response(429, json={"error": "rate limited"})
|
|
arguments: Final = {
|
|
"model": "mistral/mistral-ocr-latest",
|
|
"document": {"type": "document_url", "document_url": "https://example.com/file.pdf"},
|
|
"api_key": "test-key",
|
|
"api_base": "https://ocr.test/v1",
|
|
"num_retries": 0,
|
|
}
|
|
|
|
async def call() -> object:
|
|
if asynchronous:
|
|
return await litellm.aocr(**arguments)
|
|
return litellm.ocr(**arguments)
|
|
|
|
with pytest.raises(litellm.RateLimitError) as error:
|
|
await call()
|
|
assert error.value.status_code == 429
|
|
assert error.value.model == "mistral-ocr-latest"
|
|
assert error.value.llm_provider == "mistral"
|
|
assert provider.call_count == 1
|
|
|
|
|
|
def test_document_intelligence_environment_key_is_not_replaced_by_generic_azure_key(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("AZURE_AI_API_KEY", "generic-key")
|
|
monkeypatch.setenv("AZURE_DOCUMENT_INTELLIGENCE_API_KEY", "document-key")
|
|
monkeypatch.setenv("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT", "https://document.example.com")
|
|
prepared: Final = _prepare_ocr_request(
|
|
model="azure_ai/doc-intelligence/prebuilt-layout",
|
|
document={"type": "document_url", "document_url": "https://example.com/file.pdf"},
|
|
api_key=None,
|
|
api_base=None,
|
|
timeout=None,
|
|
custom_llm_provider=None,
|
|
extra_headers=None,
|
|
kwargs={"litellm_logging_obj": Mock()},
|
|
)
|
|
|
|
assert prepared.api_key is None
|
|
headers: Final = prepared.provider_config.validate_environment(
|
|
headers={},
|
|
model=prepared.model,
|
|
api_key=prepared.api_key,
|
|
api_base=prepared.api_base,
|
|
litellm_params=prepared.litellm_params,
|
|
)
|
|
assert headers["Ocp-Apim-Subscription-Key"] == "document-key"
|
|
|
|
|
|
def test_document_intelligence_explicit_connection_is_preserved(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("AZURE_AI_API_KEY", "generic-key")
|
|
monkeypatch.setenv("AZURE_AI_API_BASE", "https://generic.example.com")
|
|
prepared: Final = _prepare_ocr_request(
|
|
model="azure_ai/doc-intelligence/prebuilt-layout",
|
|
document={"type": "document_url", "document_url": "https://example.com/file.pdf"},
|
|
api_key="explicit-key",
|
|
api_base="https://document.example.com",
|
|
timeout=None,
|
|
custom_llm_provider=None,
|
|
extra_headers=None,
|
|
kwargs={"litellm_logging_obj": Mock()},
|
|
)
|
|
|
|
assert prepared.api_key == "explicit-key"
|
|
assert prepared.api_base == "https://document.example.com"
|
|
|
|
|
|
def test_generic_azure_connection_still_applies_to_foundry_ocr(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("AZURE_AI_API_KEY", "generic-key")
|
|
monkeypatch.setenv("AZURE_AI_API_BASE", "https://generic.example.com")
|
|
prepared: Final = _prepare_ocr_request(
|
|
model="azure_ai/mistral-document-ai-2505",
|
|
document={"type": "document_url", "document_url": "https://example.com/file.pdf"},
|
|
api_key=None,
|
|
api_base=None,
|
|
timeout=None,
|
|
custom_llm_provider=None,
|
|
extra_headers=None,
|
|
kwargs={"litellm_logging_obj": Mock()},
|
|
)
|
|
|
|
assert prepared.api_key == "generic-key"
|
|
assert prepared.api_base == "https://generic.example.com"
|
|
|
|
|
|
PRICING_OCR_MODEL: Final = "mistral/some-unmapped-ocr-model-for-testing"
|
|
PRICING_DOCUMENT: Final = {"type": "document_url", "document_url": "https://example.com/doc.pdf"}
|
|
|
|
|
|
def _pricing_logging_obj() -> Logging:
|
|
return Logging(
|
|
model=PRICING_OCR_MODEL,
|
|
messages=[],
|
|
stream=False,
|
|
call_type="ocr",
|
|
start_time=None,
|
|
litellm_call_id="test-ocr-request-pricing",
|
|
function_id="1234",
|
|
)
|
|
|
|
|
|
def _prepare_with_pricing(kwargs: dict[str, object]) -> Logging:
|
|
logging_obj: Final = _pricing_logging_obj()
|
|
_prepare_ocr_request(
|
|
model=PRICING_OCR_MODEL,
|
|
document=dict(PRICING_DOCUMENT),
|
|
api_key="test-key",
|
|
api_base=None,
|
|
timeout=None,
|
|
custom_llm_provider=None,
|
|
extra_headers=None,
|
|
kwargs={"litellm_logging_obj": logging_obj, **kwargs},
|
|
)
|
|
return logging_obj
|
|
|
|
|
|
def test_prepare_ocr_request_forwards_custom_pricing_to_logging_params() -> None:
|
|
logging_obj: Final = _prepare_with_pricing({"ocr_cost_per_page": 0.05, "ocr_cost_per_credit": 0.5})
|
|
|
|
assert logging_obj.litellm_params["ocr_cost_per_page"] == 0.05
|
|
assert logging_obj.litellm_params["ocr_cost_per_credit"] == 0.5
|
|
assert use_custom_pricing_for_model(logging_obj.litellm_params) is True
|
|
|
|
|
|
def test_prepare_ocr_request_without_custom_pricing_leaves_logging_params_unpriced() -> None:
|
|
logging_obj: Final = _prepare_with_pricing({})
|
|
|
|
assert "ocr_cost_per_page" not in logging_obj.litellm_params
|
|
assert use_custom_pricing_for_model(logging_obj.litellm_params) is False
|
|
|
|
|
|
def test_direct_ocr_call_bills_request_level_per_page_pricing() -> None:
|
|
assert PRICING_OCR_MODEL not in litellm.model_cost
|
|
logging_obj: Final = _prepare_with_pricing({"ocr_cost_per_page": 0.05})
|
|
response: Final = OCRResponse(
|
|
pages=[OCRPage(index=index, markdown=f"page {index}") for index in range(3)],
|
|
model=PRICING_OCR_MODEL,
|
|
usage_info=OCRUsageInfo(pages_processed=3),
|
|
)
|
|
|
|
assert logging_obj._response_cost_calculator(result=response) == pytest.approx(0.05 * 3)
|
|
|
|
|
|
def _prepare(model: str, document: object, **kwargs: object) -> object:
|
|
return _prepare_ocr_request(
|
|
model=model,
|
|
document=document, # pyright: ignore[reportArgumentType] # exercises the runtime guard for untyped callers
|
|
api_key="test-key",
|
|
api_base=None,
|
|
timeout=None,
|
|
custom_llm_provider=None,
|
|
extra_headers=None,
|
|
kwargs={"litellm_logging_obj": Mock(), **kwargs},
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("document", "match"),
|
|
(
|
|
("https://example.com/file.pdf", "document must be a dict"),
|
|
({"type": "video_url", "video_url": "https://example.com/clip.mp4"}, "Invalid document type: video_url"),
|
|
({"type": "document_url", "document_url": ""}, "Document URL is required"),
|
|
),
|
|
)
|
|
def test_prepare_ocr_request_rejects_malformed_documents(document: object, match: str) -> None:
|
|
with pytest.raises(litellm.BadRequestError, match=match):
|
|
_prepare("mistral/mistral-ocr-latest", document)
|
|
|
|
|
|
def test_prepare_ocr_request_maps_param_mapping_errors_to_bad_request(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
config: Final = Mock()
|
|
config.resolve_connection_params.return_value = ("test-key", None)
|
|
config.get_supported_ocr_params.return_value = ["pages"]
|
|
config.map_ocr_params.side_effect = ValueError("pages must be a list")
|
|
monkeypatch.setattr(ProviderConfigManager, "get_provider_ocr_config", Mock(return_value=config))
|
|
|
|
with pytest.raises(litellm.BadRequestError, match="pages must be a list") as error:
|
|
_prepare("mistral/mistral-ocr-latest", dict(PRICING_DOCUMENT), pages="1")
|
|
|
|
assert error.value.llm_provider == "mistral"
|
|
assert isinstance(error.value.__cause__, ValueError)
|
|
|
|
|
|
def test_prepare_ocr_request_rejects_provider_without_ocr_support() -> None:
|
|
with pytest.raises(ValueError, match="OCR is not supported for provider: openai"):
|
|
_prepare("openai/gpt-4o", dict(PRICING_DOCUMENT))
|
|
|
|
|
|
def test_prepare_ocr_request_rejects_invalid_request_format() -> None:
|
|
with pytest.raises(litellm.UnsupportedParamsError, match="Invalid `req_format`"):
|
|
_prepare("mistral/mistral-ocr-latest", dict(PRICING_DOCUMENT), req_format="markdown")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_python_none_provider_response_raises_public_error(
|
|
provider: Mock, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
from litellm.ocr import main
|
|
|
|
monkeypatch.setattr(main.base_llm_http_handler, "ocr", Mock(return_value=None))
|
|
|
|
with pytest.raises(litellm.APIConnectionError, match="unexpected None response") as error:
|
|
await litellm.aocr(model="mistral/mistral-ocr-latest", document=dict(PRICING_DOCUMENT), api_key="test-key")
|
|
assert error.value.llm_provider == "mistral"
|
|
assert provider.call_count == 0
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("model", "expected_provider"),
|
|
(("mistral-ocr-latest", "mistral"), ("azure_ai/doc-intelligence/prebuilt-layout", "azure_ai")),
|
|
)
|
|
def test_preparation_errors_map_to_public_exception_for_inferred_provider(
|
|
provider: Mock, model: str, expected_provider: str
|
|
) -> None:
|
|
with pytest.raises(litellm.BadRequestError) as error:
|
|
litellm.ocr(model=model, document="not-a-document") # pyright: ignore[reportArgumentType] # exercises the runtime guard
|
|
assert error.value.llm_provider == expected_provider
|
|
assert "document must be a dict" in str(error.value)
|
|
assert provider.call_count == 0
|