From 39239ab974374b52d37dafb944852c6ce62369f8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 18:36:05 -0700 Subject: [PATCH] refactor(cost): share the deployment model_info lookup between video and OCR cost paths Extract one immutable helper for reading the deployment's model_info off the logging object, stop rebinding the model_info parameter inside ocr_cost, drop the explanatory comment blocks, and move the OCR custom pricing regression tests into tests/test_litellm/test_cost_calculator.py --- litellm/cost_calculator.py | 94 +++++------- tests/test_litellm/test_cost_calculator.py | 130 ++++++++++++++++ tests/test_litellm/test_ocr_custom_pricing.py | 143 ------------------ 3 files changed, 166 insertions(+), 201 deletions(-) delete mode 100644 tests/test_litellm/test_ocr_custom_pricing.py diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 496f060d5ca..9aed82f37e3 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -341,7 +341,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, # deployment model_info, for non-token custom pricing + custom_model_info: ModelInfo | None = None, ) -> tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. @@ -1429,20 +1429,7 @@ 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) usage_obj = getattr(completion_response, "usage", None) duration_seconds: float | None = None @@ -1638,24 +1625,6 @@ def completion_cost( if litellm_logging_obj is not None: request_model_for_cost = litellm_logging_obj.model - # Deployment-specific model_info, for modalities whose pricing is - # not token-based and so cannot travel via custom_cost_per_token - # (e.g. OCR per-page pricing). Same extraction as the video path - # above, minus its `or {}` default: truthiness on the value adds - # no mutable-collection construction (LIT002) and reads the same. - # Checked under both keys: router calls that go through - # `_ageneric_api_call_with_fallbacks` (OCR included) store the - # deployment's model_info under `litellm_metadata`, not `metadata`. - _custom_model_info: ModelInfo | None = None - if custom_pricing and litellm_logging_obj is not None: - _cm_litellm_params = getattr(litellm_logging_obj, "litellm_params", None) - if _cm_litellm_params is not None: - for _cm_metadata_key in ("metadata", "litellm_metadata"): - _cm_metadata = _cm_litellm_params.get(_cm_metadata_key) - if _cm_metadata and _cm_metadata.get("model_info") is not None: - _custom_model_info = _cm_metadata.get("model_info") - break - ( prompt_tokens_cost_usd_dollar, completion_tokens_cost_usd_dollar, @@ -1681,7 +1650,7 @@ def completion_cost( vertex_location=vertex_location, response=completion_response, request_model=request_model_for_cost, - custom_model_info=_custom_model_info, + custom_model_info=_deployment_model_info(litellm_logging_obj, custom_pricing), ) # Get additional costs from provider (e.g., routing fees, infrastructure costs) @@ -1911,6 +1880,32 @@ def response_cost_calculator( raise e +def _deployment_model_info( + litellm_logging_obj: LitellmLoggingObject | None, + custom_pricing: bool | None, +) -> ModelInfo | None: + if not custom_pricing or 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 _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, @@ -1922,9 +1917,8 @@ def ocr_cost( 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, used for - custom pricing. Takes precedence over the model cost map, mirroring - the video generation cost path. + model_info: Optional[ModelInfo] - deployment-specific model info; its OCR pricing + takes precedence over the model cost map Returns: Tuple[float, float]: cost of OCR processing @@ -1942,34 +1936,18 @@ def ocr_cost( if response.usage_info is None: raise ValueError("OCR response usage_info is None") - ######################################################### - # Deployment-specific pricing wins over the cost map. - # - # Custom pricing set on a deployment is registered under the router's - # deployment id, while the shared "{provider}/{model}" key has its pricing - # fields stripped (see _register_custom_pricing_for_request). A cost map - # lookup therefore cannot see it, so an OCR model that is not in the map - # bills $0 no matter how it is priced in config. Prefer the caller-supplied - # model_info when it carries OCR pricing. - ######################################################### - has_custom_ocr_pricing: Final[bool] = model_info is not None and ( + 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 ) - if not has_custom_ocr_pricing: - try: - model_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) - except Exception: - model_info = 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 = None - if model_info is not None: - cost_per_credit = model_info.get("ocr_cost_per_credit") + 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 - 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") if pricing is not None else None + annotation_cost_per_page: Final = pricing.get("annotation_cost_per_page") if pricing is not None else None 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/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index f8fa2231597..824108ab0ec 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4645,3 +4645,133 @@ 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): + from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo + + 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): + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = 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.litellm_params = litellm_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_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.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_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/test_ocr_custom_pricing.py b/tests/test_litellm/test_ocr_custom_pricing.py deleted file mode 100644 index 4dffc43c23f..00000000000 --- a/tests/test_litellm/test_ocr_custom_pricing.py +++ /dev/null @@ -1,143 +0,0 @@ -""" -Regression tests: OCR cost must honour deployment-specific custom pricing. - -Before the fix, `ocr_cost()` resolved pricing exclusively through -`litellm.get_model_info(model=..., custom_llm_provider=...)`, i.e. a cost map -lookup keyed by model name. Custom pricing set on a deployment is registered -under the router's deployment id, and the shared "{provider}/{model}" key has -its pricing fields stripped, so the lookup could never see it. An OCR model -absent from the cost map therefore billed $0 no matter how it was priced in -config, even though `ocr_cost_per_page` / `ocr_cost_per_credit` are declared -fields of `CustomPricingLiteLLMParams`. -""" - -import pytest - -import litellm -from litellm.cost_calculator import completion_cost, ocr_cost -from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo - -# A model deliberately absent from the cost map. -UNMAPPED_MODEL = "azure_ai/some-unmapped-ocr-model-for-testing" -CUSTOM_COST_PER_PAGE = 0.004 -CUSTOM_COST_PER_CREDIT = 0.25 - - -def _ocr_response(model: str, pages_processed: int = 1, credits: int | None = None) -> OCRResponse: - # NOTE: model_construct() is used rather than OCRResponse(...) because the - # OCRResponse field `object: str = "ocr"` shadows the builtin `object` used - # in the `tables` / `keyValuePairs` annotations above it, so pydantic tries - # to resolve "ocr" as a forward-referenced type and schema building fails. - # That is an unrelated defect; validation is not what these tests exercise. - usage_info = OCRUsageInfo(pages_processed=pages_processed) - if credits is not None: - usage_info.credits = credits - return OCRResponse.model_construct( - pages=[OCRPage(index=i, markdown=f"page {i}") for i in range(pages_processed)], - model=model, - usage_info=usage_info, - ) - - -def test_unmapped_ocr_model_has_no_map_pricing() -> None: - """Guard the premise: the model really is absent from the cost map.""" - assert UNMAPPED_MODEL not in litellm.model_cost - - -@pytest.mark.parametrize("pages_processed", [1, 3, 10]) -def test_ocr_cost_uses_custom_per_page_pricing(pages_processed: int) -> None: - cost, _ = ocr_cost( - model=UNMAPPED_MODEL, - custom_llm_provider="azure_ai", - response=_ocr_response(UNMAPPED_MODEL, pages_processed=pages_processed), - model_info={"ocr_cost_per_page": CUSTOM_COST_PER_PAGE}, - ) - assert cost == pytest.approx(CUSTOM_COST_PER_PAGE * pages_processed) - - -def test_ocr_cost_uses_custom_per_credit_pricing() -> None: - cost, _ = ocr_cost( - model=UNMAPPED_MODEL, - custom_llm_provider="azure_ai", - response=_ocr_response(UNMAPPED_MODEL, pages_processed=2, credits=4), - model_info={"ocr_cost_per_credit": CUSTOM_COST_PER_CREDIT}, - ) - assert cost == pytest.approx(CUSTOM_COST_PER_CREDIT * 4) - - -def test_unmapped_ocr_model_without_custom_pricing_still_bills_zero() -> None: - """Unchanged behaviour when nothing is configured — no map entry, no override.""" - cost, _ = ocr_cost( - model=UNMAPPED_MODEL, - custom_llm_provider="azure_ai", - response=_ocr_response(UNMAPPED_MODEL, pages_processed=5), - ) - assert cost == 0.0 - - -def test_custom_pricing_does_not_override_a_mapped_model_when_absent() -> None: - """model_info without OCR pricing must fall through to the cost map.""" - mapped_model = "mistral/mistral-ocr-4-0" - cost, _ = ocr_cost( - model=mapped_model, - custom_llm_provider="mistral", - response=_ocr_response(mapped_model, pages_processed=2), - model_info={"id": "some-deployment-id"}, - ) - assert cost == pytest.approx(0.004 * 2) - - -def test_ocr_custom_pricing_end_to_end_through_completion_cost() -> None: - """The whole path: litellm_params.metadata.model_info -> ocr_cost.""" - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging - - logging_obj = LiteLLMLogging( - model=UNMAPPED_MODEL, - messages=[], - stream=False, - call_type="ocr", - start_time=None, - litellm_call_id="test-ocr-custom-pricing", - function_id="1234", - ) - logging_obj.litellm_params = {"metadata": {"model_info": {"ocr_cost_per_page": CUSTOM_COST_PER_PAGE}}} - - cost = completion_cost( - completion_response=_ocr_response(UNMAPPED_MODEL, pages_processed=3), - model=UNMAPPED_MODEL, - custom_llm_provider="azure_ai", - call_type="ocr", - custom_pricing=True, - litellm_logging_obj=logging_obj, - ) - assert cost == pytest.approx(CUSTOM_COST_PER_PAGE * 3) - - -def test_ocr_custom_pricing_end_to_end_via_litellm_metadata() -> None: - """Router OCR calls go through `_ageneric_api_call_with_fallbacks`, which - stores the deployment's model_info under `litellm_metadata` rather than - `metadata`. The extraction must read that key too.""" - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging - - logging_obj = LiteLLMLogging( - model=UNMAPPED_MODEL, - messages=[], - stream=False, - call_type="ocr", - start_time=None, - litellm_call_id="test-ocr-custom-pricing-litellm-metadata", - function_id="1234", - ) - logging_obj.litellm_params = { - "litellm_metadata": {"model_info": {"ocr_cost_per_page": CUSTOM_COST_PER_PAGE}}, - } - - cost = completion_cost( - completion_response=_ocr_response(UNMAPPED_MODEL, pages_processed=3), - model=UNMAPPED_MODEL, - custom_llm_provider="azure_ai", - call_type="ocr", - custom_pricing=True, - litellm_logging_obj=logging_obj, - ) - assert cost == pytest.approx(CUSTOM_COST_PER_PAGE * 3)