diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs index 1cbe8a179e3..c7e5f123c19 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs @@ -58,6 +58,13 @@ impl PythonLogger { params.set_item(name, value)?; } } + for name in custom_pricing_fields(py)? { + if let Some(value) = kwargs.bind(py).get_item(&name)? + && !value.is_none() + { + params.set_item(name, value)?; + } + } update.set_item("litellm_params", params)?; update.set_item("custom_llm_provider", &pre_call.custom_llm_provider)?; self.object(py) @@ -120,6 +127,17 @@ impl PythonLogger { } } +fn custom_pricing_fields(py: Python<'_>) -> PyResult> { + py.import("litellm.types.utils")? + .getattr("CustomPricingLiteLLMParams")? + .getattr("model_fields")? + .cast_into::()? + .keys() + .iter() + .map(|name| name.extract::()) + .collect() +} + fn redact( py: Python<'_>, params: &Bound<'_, PyDict>, diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 814eaaf76f7..22bdb016dc1 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 @@ -310,6 +311,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, @@ -344,6 +354,7 @@ def cost_per_token( response: Any | None = None, ### REQUEST MODEL ### request_model: str | None = None, # original request model for router detection + custom_model_info: OCRPricing | None = None, ) -> tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. @@ -558,6 +569,7 @@ def cost_per_token( model=model, custom_llm_provider=custom_llm_provider, response=response, + model_info=custom_model_info, ) elif ( call_type == "aretrieve_batch" @@ -1432,20 +1444,9 @@ def completion_cost( ) elif call_type in _VIDEO_CALL_TYPES: ### VIDEO GENERATION COST CALCULATION ### - # Extract custom model_info for deployment-specific pricing - _video_model_info: ModelInfo | None = None - if custom_pricing and litellm_logging_obj is not None: - _litellm_params = getattr(litellm_logging_obj, "litellm_params", None) - if _litellm_params is not None: - _video_model_info = next( - ( - model_info - for _metadata_key in ("metadata", "litellm_metadata") - if (model_info := (_litellm_params.get(_metadata_key) or {}).get("model_info")) - is not None - ), - None, - ) + _video_model_info: ModelInfo | None = _deployment_model_info( + litellm_logging_obj, custom_pricing, router_model_id + ) usage_obj = getattr(completion_response, "usage", None) duration_seconds: float | None = None @@ -1665,6 +1666,7 @@ def completion_cost( data_residency=data_residency, vertex_location=vertex_location, response=completion_response, + 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) @@ -1898,16 +1900,82 @@ def response_cost_calculator( raise e +def _deployment_model_info( + litellm_logging_obj: LitellmLoggingObject | None, + custom_pricing: bool | None, + router_model_id: str | None, +) -> ModelInfo | None: + if not custom_pricing: + return None + registered_deployment_info: Final = ( + _cost_map_model_info(router_model_id, None) + if router_model_id is not None and router_model_id in litellm.model_cost + else None + ) + if registered_deployment_info is not None: + return registered_deployment_info + if litellm_logging_obj is None: + return None + litellm_params: Final = getattr(litellm_logging_obj, "litellm_params", None) + if litellm_params is None: + return None + return next( + ( + model_info + for metadata_key in ("metadata", "litellm_metadata") + if (metadata := litellm_params.get(metadata_key)) and (model_info := metadata.get("model_info")) is not None + ), + None, + ) + + +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 _layered_ocr_pricing(litellm_params, deployment_info) + + +def _first_ocr_price(field: OCRPricingField, *sources: Mapping[str, object] | None) -> float | None: + return next( + (price for source in sources if source is not None and isinstance(price := source.get(field), int | float)), + None, + ) + + +def _layered_ocr_pricing(*sources: Mapping[str, object] | None) -> OCRPricing: + return OCRPricing( + ocr_cost_per_page=_first_ocr_price("ocr_cost_per_page", *sources), + ocr_cost_per_credit=_first_ocr_price("ocr_cost_per_credit", *sources), + annotation_cost_per_page=_first_ocr_price("annotation_cost_per_page", *sources), + ) + + +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) + except Exception: + return None + + def ocr_cost( model: str, custom_llm_provider: str | None, response: object | 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[OCRPricing] - deployment-specific OCR pricing; each rate it sets + overrides the model cost map's, the rest fall back to the map Returns: Tuple[float, float]: cost of OCR processing @@ -1925,20 +1993,15 @@ def ocr_cost( if response.usage_info is None: raise ValueError("OCR response usage_info is None") - try: - model_info: ModelInfo | None = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) - except Exception: - model_info = None - credits: Final = getattr(response.usage_info, "credits", None) - cost_per_credit = None - if model_info is not None: - cost_per_credit = model_info.get("ocr_cost_per_credit") + pricing: Final = _layered_ocr_pricing(model_info, _cost_map_model_info(model, custom_llm_provider)) + + cost_per_credit: Final = pricing.get("ocr_cost_per_credit") if credits is not None and cost_per_credit is not None: return cost_per_credit * credits, 0.0 - ocr_cost_per_page: Final = model_info.get("ocr_cost_per_page") if model_info is not None else None - annotation_cost_per_page: Final = model_info.get("annotation_cost_per_page") if model_info is not None else None + ocr_cost_per_page: Final = pricing.get("ocr_cost_per_page") + annotation_cost_per_page: Final = pricing.get("annotation_cost_per_page") annotation_rate: Final = annotation_cost_per_page if annotation_cost_per_page is not None else ocr_cost_per_page pages_processed: Final = response.usage_info.pages_processed diff --git a/litellm/ocr/legacy.py b/litellm/ocr/legacy.py index ddf6016dce3..a742be274b3 100644 --- a/litellm/ocr/legacy.py +++ b/litellm/ocr/legacy.py @@ -28,6 +28,7 @@ from litellm.llms.base_llm.ocr.transformation import ( from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.ocr.input import FileReader from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import CustomPricingLiteLLMParams from litellm.utils import ProviderConfigManager, client base_llm_http_handler: Final = BaseLLMHTTPHandler() @@ -149,6 +150,7 @@ def _prepare_ocr_request( litellm_params={ "litellm_call_id": litellm_call_id, "api_base": resolved_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_legacy.py b/tests/test_litellm/ocr/test_legacy.py index a30976f89af..4b0b78f5a0f 100644 --- a/tests/test_litellm/ocr/test_legacy.py +++ b/tests/test_litellm/ocr/test_legacy.py @@ -11,7 +11,8 @@ import pytest import litellm from litellm.integrations.custom_logger import CustomLogger -from litellm.llms.base_llm.ocr.transformation import OCRResponse +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.legacy import _prepare_ocr_request @@ -198,3 +199,61 @@ def test_generic_azure_connection_still_applies_to_foundry_ocr(monkeypatch: pyte 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) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index f610821e06a..3ef768790f8 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -17,6 +17,8 @@ from litellm.cost_calculator import ( handle_realtime_stream_cost_calculation, response_cost_calculator, ) +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo from litellm.types.llms.openai import OpenAIRealtimeStreamList from litellm.types.rerank import RerankResponse from litellm.types.utils import ( @@ -4768,3 +4770,228 @@ def test_collect_and_combine_realtime_usage_stores_partitioned_text_tokens() -> assert combined.completion_tokens_details.reasoning_tokens == 95 assert combined.completion_tokens_details.text_tokens == 38 assert combined.completion_tokens_details.audio_tokens == 0 + + +UNMAPPED_OCR_MODEL: Final = "azure_ai/some-unmapped-ocr-model-for-testing" +MAPPED_OCR_MODEL: Final = "mistral/mistral-ocr-4-0" + + +def _ocr_response(model: str, pages_processed: int, credits: float | None = None) -> OCRResponse: + return OCRResponse( + pages=[OCRPage(index=index, markdown=f"page {index}") for index in range(pages_processed)], + model=model, + usage_info=OCRUsageInfo(pages_processed=pages_processed, credits=credits), + ) + + +def _ocr_logging_obj(litellm_params: dict[str, object]) -> Logging: + logging_obj: Final = Logging( + model=UNMAPPED_OCR_MODEL, + messages=[], + stream=False, + call_type="ocr", + start_time=None, + litellm_call_id="test-ocr-custom-pricing", + function_id="1234", + ) + logging_obj.update_environment_variables(litellm_params=litellm_params, optional_params={}) + return logging_obj + + +@pytest.mark.parametrize("pages_processed", [1, 3, 10]) +def test_ocr_cost_uses_deployment_per_page_pricing_for_unmapped_model(pages_processed: int): + from litellm.cost_calculator import ocr_cost + + assert UNMAPPED_OCR_MODEL not in litellm.model_cost + cost, _ = ocr_cost( + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=pages_processed), + model_info={"ocr_cost_per_page": 0.004}, + ) + assert cost == pytest.approx(0.004 * pages_processed) + + +def test_ocr_cost_uses_deployment_annotation_only_pricing_for_unmapped_model(): + from litellm.cost_calculator import ocr_cost + + assert UNMAPPED_OCR_MODEL not in litellm.model_cost + response: Final = OCRResponse( + pages=[OCRPage(index=index, markdown=f"page {index}") for index in range(3)], + model=UNMAPPED_OCR_MODEL, + usage_info=OCRUsageInfo(pages_processed=3, pages_processed_annotation=2), + ) + cost, _ = ocr_cost( + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + response=response, + model_info={"annotation_cost_per_page": 0.01}, + ) + assert cost == pytest.approx(0.01 * 2) + + +def test_ocr_cost_annotation_only_override_keeps_mapped_per_page_rate(): + from litellm.cost_calculator import ocr_cost + + map_price: Final = litellm.model_cost[MAPPED_OCR_MODEL]["ocr_cost_per_page"] + response: Final = OCRResponse( + pages=[OCRPage(index=index, markdown=f"page {index}") for index in range(3)], + model=MAPPED_OCR_MODEL, + usage_info=OCRUsageInfo(pages_processed=3, pages_processed_annotation=2), + ) + cost, _ = ocr_cost( + model=MAPPED_OCR_MODEL, + custom_llm_provider="mistral", + response=response, + model_info={"annotation_cost_per_page": 0.01}, + ) + assert cost == pytest.approx(map_price * 3 + 0.01 * 2) + + +def test_ocr_cost_uses_deployment_per_credit_pricing_for_unmapped_model(): + from litellm.cost_calculator import ocr_cost + + cost, _ = ocr_cost( + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=2, credits=4), + model_info={"ocr_cost_per_credit": 0.25}, + ) + assert cost == pytest.approx(0.25 * 4) + + +def test_ocr_cost_unmapped_model_without_deployment_pricing_bills_zero(): + from litellm.cost_calculator import ocr_cost + + cost, _ = ocr_cost( + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=5), + model_info={"id": "some-deployment-id"}, + ) + assert cost == 0.0 + + +@pytest.mark.usefixtures("_local_model_cost_map") +def test_ocr_cost_deployment_pricing_overrides_cost_map_for_mapped_model(): + 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 + override_price: Final = map_price * 10 + + 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_page": override_price}, + ) + assert cost == pytest.approx(override_price * 2) + + +@pytest.mark.usefixtures("_local_model_cost_map") +def test_ocr_cost_falls_through_to_cost_map_when_deployment_has_no_ocr_pricing(): + 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={"id": "some-deployment-id"}, + ) + 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}}}) + + 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.004 * 3) + + +def test_completion_cost_ocr_prefers_pricing_registered_under_router_model_id(monkeypatch: pytest.MonkeyPatch): + deployment_id: Final = "ocr-deployment-priced-through-litellm-params" + monkeypatch.setitem( + litellm.model_cost, deployment_id, {"mode": "ocr", "litellm_provider": "azure_ai", "ocr_cost_per_page": 0.05} + ) + logging_obj = _ocr_logging_obj({"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, + router_model_id=deployment_id, + litellm_logging_obj=logging_obj, + ) + 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}}}) + + 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=False, + litellm_logging_obj=logging_obj, + ) + assert cost == 0.0 diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index 1acad5527d8..e7ebc5b3018 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -79,6 +79,20 @@ async def test_proxy_metadata_remains_python_owned(ocr_server: RecordingServer) assert "metadata" not in ocr_server.requests[0].body +@pytest.mark.asyncio +async def test_request_level_custom_pricing_reaches_logging_params_and_bills_the_call( + ocr_server: RecordingServer, +) -> None: + recorder: Final = RecordingLogger() + response: Final = await call_aocr(ocr_server, callbacks=[recorder], ocr_cost_per_page=0.05) + events: Final = await recorder.wait_for_async("async_log_success_event") + + assert response.usage_info is not None and response.usage_info.pages_processed == 1 + assert events[0].kwargs["litellm_params"]["ocr_cost_per_page"] == 0.05 + assert response._hidden_params["response_cost"] == pytest.approx(0.05) + assert "ocr_cost_per_page" not in ocr_server.requests[0].body + + @pytest.mark.asyncio async def test_response_replacement_finalized_before_dispatch_in_caller_task(ocr_server: RecordingServer) -> None: caller: Final = asyncio.current_task()