diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 68f52e4a10e..6d35a9e89aa 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, cast from httpx import Response from pydantic import BaseModel +from typing_extensions import ReadOnly, TypedDict import litellm import litellm._logging @@ -306,6 +307,15 @@ def _transcription_usage_has_token_details( return (prompt_tokens_val > 0) or (completion_tokens_val > 0) +OCRPricingField = Literal["ocr_cost_per_page", "ocr_cost_per_credit", "annotation_cost_per_page"] + + +class OCRPricing(TypedDict, total=False): + ocr_cost_per_page: ReadOnly[float | None] + ocr_cost_per_credit: ReadOnly[float | None] + annotation_cost_per_page: ReadOnly[float | None] + + def cost_per_token( model: str = "", prompt_tokens: int = 0, @@ -341,7 +351,7 @@ def cost_per_token( ### REQUEST MODEL ### request_model: str | None = None, # original request model for router detection ### DEPLOYMENT-SPECIFIC PRICING ### - custom_model_info: ModelInfo | None = None, + custom_model_info: OCRPricing | None = None, ) -> tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. @@ -1652,7 +1662,7 @@ def completion_cost( vertex_location=vertex_location, response=completion_response, request_model=request_model_for_cost, - custom_model_info=_deployment_model_info(litellm_logging_obj, custom_pricing, router_model_id), + custom_model_info=_ocr_model_info(litellm_logging_obj, custom_pricing, router_model_id), ) # Get additional costs from provider (e.g., routing fees, infrastructure costs) @@ -1911,6 +1921,35 @@ def _deployment_model_info( ) +def _ocr_model_info( + litellm_logging_obj: LitellmLoggingObject | None, + custom_pricing: bool | None, + router_model_id: str | None, +) -> OCRPricing | None: + deployment_info: Final = _deployment_model_info(litellm_logging_obj, custom_pricing, router_model_id) + litellm_params: Final = getattr(litellm_logging_obj, "litellm_params", None) if custom_pricing else None + if litellm_params is None: + return deployment_info + return OCRPricing( + ocr_cost_per_page=_request_or_deployment_price("ocr_cost_per_page", litellm_params, deployment_info), + ocr_cost_per_credit=_request_or_deployment_price("ocr_cost_per_credit", litellm_params, deployment_info), + annotation_cost_per_page=_request_or_deployment_price( + "annotation_cost_per_page", litellm_params, deployment_info + ), + ) + + +def _request_or_deployment_price( + field: OCRPricingField, + litellm_params: Mapping[str, object], + deployment_info: ModelInfo | None, +) -> float | None: + request_price: Final = litellm_params.get(field) + if isinstance(request_price, int | float): + return request_price + return deployment_info.get(field) if deployment_info is not None else None + + def _cost_map_model_info(model: str, custom_llm_provider: str | None) -> ModelInfo | None: try: return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) @@ -1922,14 +1961,14 @@ def ocr_cost( model: str, custom_llm_provider: str | None, response: object | None = None, - model_info: ModelInfo | None = None, + model_info: OCRPricing | None = None, ) -> tuple[float, float]: """ Args: model: str - model name custom_llm_provider: Optional[str] - custom LLM provider response: Optional[Any] - response object - model_info: Optional[ModelInfo] - deployment-specific model info; its OCR pricing + model_info: Optional[OCRPricing] - deployment-specific model info; its OCR pricing takes precedence over the model cost map Returns: diff --git a/tests/test_litellm/ocr/test_main.py b/tests/test_litellm/ocr/test_main.py index 0007d98dcb0..de4b28dafdd 100644 --- a/tests/test_litellm/ocr/test_main.py +++ b/tests/test_litellm/ocr/test_main.py @@ -1,9 +1,13 @@ from typing import Final +import pytest + +import litellm 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.ocr.main import _prepare_ocr_request -OCR_MODEL: Final = "mistral/mistral-ocr-4-1" +OCR_MODEL: Final = "mistral/some-unmapped-ocr-model-for-testing" DOCUMENT: Final = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} @@ -47,3 +51,15 @@ def test_prepare_ocr_request_without_custom_pricing_leaves_logging_params_unpric 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 OCR_MODEL not in litellm.model_cost + logging_obj: Final = _prepare({"ocr_cost_per_page": 0.05}) + response: Final = OCRResponse( + pages=[OCRPage(index=index, markdown=f"page {index}") for index in range(3)], + model=OCR_MODEL, + usage_info=OCRUsageInfo(pages_processed=3), + ) + + assert logging_obj._response_cost_calculator(result=response) == pytest.approx(0.05 * 3) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index cb9217b2f41..6db91d7a775 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4661,8 +4661,8 @@ def _ocr_response(model: str, pages_processed: int, credits: float | None = None ) -def _ocr_logging_obj(litellm_params: dict[str, dict[str, ModelInfo]]) -> Logging: - logging_obj = Logging( +def _ocr_logging_obj(litellm_params: dict[str, object]) -> Logging: + logging_obj: Final = Logging( model=UNMAPPED_OCR_MODEL, messages=[], stream=False, @@ -4671,7 +4671,7 @@ def _ocr_logging_obj(litellm_params: dict[str, dict[str, ModelInfo]]) -> Logging litellm_call_id="test-ocr-custom-pricing", function_id="1234", ) - logging_obj.litellm_params = litellm_params + logging_obj.update_environment_variables(litellm_params=litellm_params, optional_params={}) return logging_obj @@ -4796,6 +4796,34 @@ def test_completion_cost_ocr_prefers_pricing_registered_under_router_model_id(mo assert cost == pytest.approx(0.05 * 3) +def test_completion_cost_ocr_bills_request_level_pricing_for_direct_sdk_call(): + logging_obj = _ocr_logging_obj({"ocr_cost_per_page": 0.05}) + + cost = completion_cost( + completion_response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=3), + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + call_type="ocr", + custom_pricing=True, + litellm_logging_obj=logging_obj, + ) + assert cost == pytest.approx(0.05 * 3) + + +def test_completion_cost_ocr_request_level_pricing_fills_in_deployment_model_info_without_ocr_pricing(): + logging_obj = _ocr_logging_obj({"ocr_cost_per_page": 0.05, "metadata": {"model_info": {"mode": "ocr"}}}) + + cost = completion_cost( + completion_response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=3), + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + call_type="ocr", + custom_pricing=True, + litellm_logging_obj=logging_obj, + ) + assert cost == pytest.approx(0.05 * 3) + + def test_completion_cost_ocr_ignores_deployment_pricing_without_custom_pricing_flag(): logging_obj = _ocr_logging_obj({"metadata": {"model_info": {"ocr_cost_per_page": 0.004}}})