From 93c852f75af82e89483534eecd148d661d9277a0 Mon Sep 17 00:00:00 2001 From: MarkXian Date: Fri, 28 Aug 2026 11:52:34 +0800 Subject: [PATCH] fix(ocr): include annotation page costs --- litellm/cost_calculator.py | 22 ++++++++++++---- .../llms/mistral/ocr/test_mistral_ocr_cost.py | 26 ++++++++++++++++++- 2 files changed, 42 insertions(+), 6 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index f8f9de7fbec..a419e6d802e 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1899,12 +1899,17 @@ def ocr_cost( return cost_per_credit * credits, 0.0 ocr_cost_per_page: float | None = None + annotation_cost_per_page: float | None = None if model_info is not None: ocr_cost_per_page = model_info.get("ocr_cost_per_page") + annotation_cost_per_page = model_info.get("annotation_cost_per_page") - pages_processed: Final = response.usage_info.pages_processed - if pages_processed is None: - if cost_per_credit is not None or ocr_cost_per_page is None: + raw_pages_processed: Final = response.usage_info.pages_processed + pages_processed_annotation: Final = getattr(response.usage_info, "pages_processed_annotation", None) or 0 + if raw_pages_processed is None: + if pages_processed_annotation > 0 and (annotation_cost_per_page is not None or ocr_cost_per_page is not None): + pages_processed = 0 + elif cost_per_credit is not None or ocr_cost_per_page is None: # Surface missing usage data instead of silently under-reporting # cost. The previous behavior raised ValueError; we now return 0.0 # for credit-priced or unpriced models, so log a warning to keep @@ -1917,7 +1922,10 @@ def ocr_cost( credits, ) return 0.0, 0.0 - raise ValueError("OCR response pages_processed is None") + else: + raise ValueError("OCR response pages_processed is None") + else: + pages_processed = raw_pages_processed if ocr_cost_per_page is None: # No per-page pricing configured. Either the model is on credit-based @@ -1935,7 +1943,11 @@ def ocr_cost( ) return 0.0, 0.0 - total_ocr_processing_cost: Final[float] = ocr_cost_per_page * pages_processed + total_ocr_processing_cost: float = ocr_cost_per_page * pages_processed + if pages_processed_annotation > 0: + annotation_rate = annotation_cost_per_page if annotation_cost_per_page is not None else ocr_cost_per_page + total_ocr_processing_cost += annotation_rate * pages_processed_annotation + return total_ocr_processing_cost, 0.0 diff --git a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py index 890df597933..16cd2b59799 100644 --- a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py +++ b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py @@ -33,6 +33,17 @@ def _ocr_response(model: str, pages_processed: int) -> OCRResponse: ) +def _annotated_ocr_response(model: str, pages_processed: int, pages_processed_annotation: int) -> OCRResponse: + return OCRResponse( + pages=[OCRPage(index=i, markdown=f"page {i}") for i in range(pages_processed)], + model=model, + usage_info=OCRUsageInfo( + pages_processed=pages_processed, + pages_processed_annotation=pages_processed_annotation, + ), + ) + + @pytest.mark.parametrize("model", ["mistral-ocr-4-0", "mistral-ocr-latest"]) def test_model_info_ocr4_price(model: str) -> None: info = litellm.get_model_info(model=f"mistral/{model}", custom_llm_provider="mistral") @@ -51,7 +62,6 @@ def test_ocr4_cost_scales_with_pages(model: str, pages_processed: int) -> None: assert cost == pytest.approx(OCR4_COST_PER_PAGE * pages_processed) - @pytest.mark.parametrize("cost_map_path", [MAIN_COST_MAP, BACKUP_COST_MAP]) def test_ocr3_pricing_entry(cost_map_path: Path) -> None: with open(cost_map_path) as f: @@ -68,6 +78,7 @@ def test_ocr3_pricing_entry(cost_map_path: Path) -> None: def test_ocr3_model_info_price(local_model_cost_map) -> None: info = litellm.get_model_info(model=OCR3_MODEL, custom_llm_provider="mistral") assert info["ocr_cost_per_page"] == OCR3_COST_PER_PAGE + assert info["annotation_cost_per_page"] == OCR3_ANNOTATION_COST_PER_PAGE @pytest.mark.parametrize("pages_processed", [1, 3, 10]) @@ -79,3 +90,16 @@ def test_ocr3_cost_scales_with_pages(local_model_cost_map, pages_processed: int) call_type="ocr", ) assert cost == pytest.approx(OCR3_COST_PER_PAGE * pages_processed) + + +def test_ocr3_cost_includes_annotation_pages(local_model_cost_map) -> None: + cost = completion_cost( + completion_response=_annotated_ocr_response( + "mistral-ocr-2512", pages_processed=2, pages_processed_annotation=3 + ), + model=OCR3_MODEL, + custom_llm_provider="mistral", + call_type="ocr", + ) + + assert cost == pytest.approx(OCR3_COST_PER_PAGE * 2 + OCR3_ANNOTATION_COST_PER_PAGE * 3)