From d71f4aeff9587ae189446598de23b35e6d8d8142 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 9 Sep 2026 19:47:00 -0700
Subject: [PATCH] fix(cost): layer deployment OCR rates over the cost map field
 by field

---
 litellm/cost_calculator.py                 | 45 +++++++++-------------
 tests/test_litellm/test_cost_calculator.py | 18 +++++++++
 2 files changed, 37 insertions(+), 26 deletions(-)

diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py
index 09700ad6b72..93d9609df9e 100644
--- a/litellm/cost_calculator.py
+++ b/litellm/cost_calculator.py
@@ -1939,24 +1939,22 @@ def _ocr_model_info(
     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
-        ),
+    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 _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 _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:
@@ -1977,8 +1975,8 @@ def ocr_cost(
         model: str - model name
         custom_llm_provider: Optional[str] - custom LLM provider
         response: Optional[Any] - response object
-        model_info: Optional[OCRPricing] - deployment-specific model info; its OCR pricing
-            takes precedence over the model cost map
+        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
@@ -1997,19 +1995,14 @@ def ocr_cost(
         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("annotation_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)
+    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 pricing is not None else None
+    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 = 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
+    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/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py
index a05fadce42d..d07ff231429 100644
--- a/tests/test_litellm/test_cost_calculator.py
+++ b/tests/test_litellm/test_cost_calculator.py
@@ -4813,6 +4813,24 @@ def test_ocr_cost_uses_deployment_annotation_only_pricing_for_unmapped_model():
     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