diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index afafc12623a..68f52e4a10e 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1948,12 +1948,13 @@ def ocr_cost( if response.usage_info is None: raise ValueError("OCR response usage_info is None") + credits: Final = getattr(response.usage_info, "credits", None) has_custom_ocr_pricing: Final = model_info is not None and ( - model_info.get("ocr_cost_per_page") is not None or model_info.get("ocr_cost_per_credit") is not None + model_info.get("ocr_cost_per_page") is not None + or (credits is not None and model_info.get("ocr_cost_per_credit") is not None) ) pricing: Final = model_info if has_custom_ocr_pricing else _cost_map_model_info(model, custom_llm_provider) - credits: Final = getattr(response.usage_info, "credits", None) cost_per_credit: Final = pricing.get("ocr_cost_per_credit") if pricing is not None else None if credits is not None and cost_per_credit is not None: return cost_per_credit * credits, 0.0 diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index df3f9d2096b..8b0e6950801 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -32,6 +32,7 @@ from litellm.rust_bridge import ocr as rust_ocr_bridge from litellm.rust_bridge.bindings import native_exception_types from litellm.rust_bridge.configuration import rust_enabled from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import CustomPricingLiteLLMParams from litellm.utils import ProviderConfigManager, client ####### ENVIRONMENT VARIABLES ################### @@ -171,6 +172,7 @@ def _prepare_ocr_request( litellm_params={ "litellm_call_id": litellm_call_id, "api_base": api_base, + **litellm_params.model_dump(include=frozenset(CustomPricingLiteLLMParams.model_fields), exclude_none=True), }, custom_llm_provider=custom_llm_provider, ) diff --git a/tests/test_litellm/ocr/test_main.py b/tests/test_litellm/ocr/test_main.py new file mode 100644 index 00000000000..0007d98dcb0 --- /dev/null +++ b/tests/test_litellm/ocr/test_main.py @@ -0,0 +1,49 @@ +from typing import Final + +from litellm.litellm_core_utils.litellm_logging import Logging, use_custom_pricing_for_model +from litellm.ocr.main import _prepare_ocr_request + +OCR_MODEL: Final = "mistral/mistral-ocr-4-1" +DOCUMENT: Final = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} + + +def _logging_obj() -> Logging: + return Logging( + model=OCR_MODEL, + messages=[], + stream=False, + call_type="ocr", + start_time=None, + litellm_call_id="test-ocr-request-pricing", + function_id="1234", + ) + + +def _prepare(kwargs: dict[str, object]) -> Logging: + logging_obj: Final = _logging_obj() + _prepare_ocr_request( + model=OCR_MODEL, + document=dict(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({"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({}) + + assert "ocr_cost_per_page" not in logging_obj.litellm_params + assert use_custom_pricing_for_model(logging_obj.litellm_params) is False diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 04a547c222c..cb9217b2f41 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4746,6 +4746,22 @@ def test_ocr_cost_falls_through_to_cost_map_when_deployment_has_no_ocr_pricing() assert cost == pytest.approx(map_price * 2) +@pytest.mark.usefixtures("_local_model_cost_map") +def test_ocr_cost_ignores_deployment_credit_pricing_when_response_reports_no_credits(): + from litellm.cost_calculator import ocr_cost + + map_price: Final = litellm.get_model_info(MAPPED_OCR_MODEL)["ocr_cost_per_page"] + assert map_price is not None + + cost, _ = ocr_cost( + model=MAPPED_OCR_MODEL, + custom_llm_provider="mistral", + response=_ocr_response(MAPPED_OCR_MODEL, pages_processed=2), + model_info={"ocr_cost_per_credit": 0.5}, + ) + assert cost == pytest.approx(map_price * 2) + + @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) def test_completion_cost_ocr_reads_deployment_pricing_from_logging_metadata(metadata_key: str): logging_obj = _ocr_logging_obj({metadata_key: {"model_info": {"ocr_cost_per_page": 0.004}}})