From f89ca64481d6064770eeb49ee8795e7dbcc10432 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 23:57:41 -0700 Subject: [PATCH] fix(batches): honor deployment OCR page pricing in batch cost and answer 400 for unsupported Mistral file purposes --- litellm/cost_calculator.py | 54 ++++++++++++------- litellm/litellm_core_utils/litellm_logging.py | 8 ++- litellm/llms/mistral/files/transformation.py | 6 ++- .../provider_endpoints_support_backup.json | 2 +- provider_endpoints_support.json | 2 +- .../test_litellm/batches/test_batch_utils.py | 29 ++++++++++ .../test_litellm_logging.py | 32 +++++++++++ .../test_mistral_files_transformation.py | 10 ++-- 8 files changed, 116 insertions(+), 27 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 30f2cb8489c..04c9675cc83 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2115,12 +2115,8 @@ def ocr_cost( return ocr_pages_cost + annotation_pages_cost, 0.0 -_OCR_PRICING_KEYS: Final = ( - "ocr_cost_per_page", - "ocr_cost_per_page_batches", - "annotation_cost_per_page", - "annotation_cost_per_page_batches", -) +_OCR_BATCH_PAGE_RATE_KEYS: Final = ("ocr_cost_per_page_batches", "ocr_cost_per_page") +_OCR_BATCH_ANNOTATION_RATE_KEYS: Final = ("annotation_cost_per_page_batches", "annotation_cost_per_page") def ocr_batch_cost( @@ -2133,17 +2129,27 @@ def ocr_batch_cost( Batch OCR is billed per page at the ``*_batches`` rate, falling back to the synchronous per-page rate when a model has no batch price recorded, the same - fallback ``batch_cost_calculator`` applies to per-token batch pricing. Returns - ``(prompt_cost, completion_cost)`` with the whole cost in the first slot, like - ``ocr_cost``. + fallback ``batch_cost_calculator`` applies to per-token batch pricing. Each + per-page family (OCR pages, annotation pages) belongs to the deployment's + ``model_info`` when it prices that family at either rate and to the published + cost map otherwise, so a deployment overriding one family keeps the model's + published rate for the other, and the cost map is only consulted for a family + the deployment leaves out. Returns ``(prompt_cost, completion_cost)`` with the + whole cost in the first slot, like ``ocr_cost``. """ - has_ocr_pricing: Final = model_info is not None and any(model_info.get(k) is not None for k in _OCR_PRICING_KEYS) - resolved_info: Final = ( - model_info - if has_ocr_pricing - else _lookup_model_info_or_none(model=model, custom_llm_provider=custom_llm_provider) + pages_processed: Final = usage_info.pages_processed or 0 + annotation_pages: Final = usage_info.pages_processed_annotation or 0 + deployment_page_rate: Final = _first_price(model_info, *_OCR_BATCH_PAGE_RATE_KEYS) + deployment_annotation_rate: Final = _first_price(model_info, *_OCR_BATCH_ANNOTATION_RATE_KEYS) + needs_published_pricing: Final = (pages_processed > 0 and deployment_page_rate is None) or ( + annotation_pages > 0 and deployment_annotation_rate is None ) - if resolved_info is None: + published: Final = ( + _lookup_model_info_or_none(model=model, custom_llm_provider=custom_llm_provider) + if needs_published_pricing + else None + ) + if needs_published_pricing and published is None: verbose_logger.warning( "OCR batch cost: model=%s custom_llm_provider=%s has no pricing entry; returning 0.0 cost.", _single_log_line(model), @@ -2151,10 +2157,16 @@ def ocr_batch_cost( ) return 0.0, 0.0 - page_rate: Final = _first_price(resolved_info, "ocr_cost_per_page_batches", "ocr_cost_per_page") - annotation_rate: Final = _first_price(resolved_info, "annotation_cost_per_page_batches", "annotation_cost_per_page") - pages_processed: Final = usage_info.pages_processed or 0 - annotation_pages: Final = usage_info.pages_processed_annotation or 0 + page_rate: Final = ( + deployment_page_rate + if deployment_page_rate is not None + else _first_price(published, *_OCR_BATCH_PAGE_RATE_KEYS) + ) + annotation_rate: Final = ( + deployment_annotation_rate + if deployment_annotation_rate is not None + else _first_price(published, *_OCR_BATCH_ANNOTATION_RATE_KEYS) + ) if page_rate is None and pages_processed > 0: verbose_logger.warning( "OCR batch cost: model=%s custom_llm_provider=%s reported pages_processed=%s but no " @@ -2178,7 +2190,9 @@ def _lookup_model_info_or_none(model: str, custom_llm_provider: str | None) -> M return None -def _first_price(model_info: ModelInfo, *keys: str) -> float | None: +def _first_price(model_info: ModelInfo | None, *keys: str) -> float | None: + if model_info is None: + return None return next((price for price in (model_info.get(k) for k in keys) if isinstance(price, (int, float))), None) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index eeebcec50b1..f7679b31f69 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -371,6 +371,10 @@ _DEPLOYMENT_PRICING_KEYS: Final = ( "output_cost_per_token", "input_cost_per_token_batches", "output_cost_per_token_batches", + "ocr_cost_per_page", + "ocr_cost_per_page_batches", + "annotation_cost_per_page", + "annotation_cost_per_page_batches", ) @@ -386,7 +390,9 @@ def deployment_pricing_model_info(model_id: str | None, deployment_model: str | the model's published rates instead of billing as zero. Ownership is per token direction: declaring either rate for a direction takes that whole direction, so a published batch rate can never displace a standard rate - the deployment configured itself. + the deployment configured itself. OCR per-page rates count as declared + pricing too; they pass through as registered and ``ocr_batch_cost`` layers + the published rate under each per-page family the deployment leaves out. """ if model_id is None: return None diff --git a/litellm/llms/mistral/files/transformation.py b/litellm/llms/mistral/files/transformation.py index 6e64961485a..88867ea4802 100644 --- a/litellm/llms/mistral/files/transformation.py +++ b/litellm/llms/mistral/files/transformation.py @@ -100,7 +100,11 @@ def _to_mistral_purpose(purpose: str) -> MistralFilePurpose: only run when the caller says ``purpose=batch``.""" mistral_purpose: Final = _MISTRAL_PURPOSE_BY_OPENAI.get(purpose) if mistral_purpose is None: - raise ValueError(f"Mistral does not support purpose={purpose!r}. Use one of: {_SUPPORTED_PURPOSES}") + raise mistral_error( + f"Mistral does not support purpose={purpose!r}. Use one of: {_SUPPORTED_PURPOSES}", + status_code=400, + headers=httpx.Headers(), + ) return mistral_purpose diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index dbeaccdda2d..30c1e0b894e 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -1462,7 +1462,7 @@ "audio_transcriptions": false, "audio_speech": false, "moderations": false, - "batches": false, + "batches": true, "rerank": false, "ocr": true, "a2a": true, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index c71f4a82a4a..af9b194bbee 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1577,7 +1577,7 @@ "audio_transcriptions": false, "audio_speech": false, "moderations": false, - "batches": false, + "batches": true, "rerank": false, "ocr": true, "a2a": true, diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 4cad45ed809..1e7e0200754 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1940,6 +1940,35 @@ def test_ocr_rows_use_deployment_model_info_pricing_over_cost_map(monkeypatch): assert result.cost == pytest.approx(0.01) +def test_ocr_rows_keep_the_published_page_rate_when_the_deployment_prices_only_annotations(monkeypatch): + monkeypatch.setattr( + litellm, + "get_model_info", + lambda model, custom_llm_provider=None: { + "ocr_cost_per_page_batches": 0.002, + "annotation_cost_per_page_batches": 0.0025, + }, + ) + result = bu._aggregate_batch_cost_usage_models( + entries=[_ocr_row(4, annotation_pages=4)], + custom_llm_provider="mistral", + model_info={"annotation_cost_per_page_batches": 0.01}, + ) + assert result.cost == pytest.approx(4 * 0.002 + 4 * 0.01) + + +def test_ocr_rows_bill_the_deployment_sync_page_rate_over_the_published_batch_rate(monkeypatch): + monkeypatch.setattr( + litellm, "get_model_info", lambda model, custom_llm_provider=None: pytest.fail("cost map must not be consulted") + ) + result = bu._aggregate_batch_cost_usage_models( + entries=[_ocr_row(3)], + custom_llm_provider="mistral", + model_info={"ocr_cost_per_page": 0.0912}, + ) + assert result.cost == pytest.approx(3 * 0.0912) + + def test_ocr_rows_without_pricing_bill_zero_but_count_as_successful(monkeypatch): monkeypatch.setattr(litellm, "get_model_info", lambda model, custom_llm_provider=None: {"mode": "ocr"}) result = bu._aggregate_batch_cost_usage_models(entries=[_ocr_row(3)], custom_llm_provider="mistral") diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 473d16a43f2..999adbdd935 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -17,12 +17,14 @@ from openai._legacy_response import HttpxBinaryResponseContent import litellm from litellm._logging import session_id_var, trace_id_var from litellm.constants import SENTRY_DENYLIST, SENTRY_PII_DENYLIST +from litellm.cost_calculator import ocr_batch_cost from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.litellm_core_utils.litellm_logging import ( _get_status_fields, set_callbacks, ) +from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse from litellm.types.utils import ( CallTypes, @@ -529,6 +531,36 @@ class TestGetRouterDeploymentModelInfo: finally: litellm.model_cost.pop(deployment_id, None) + def test_ocr_only_deployment_pricing_reaches_batch_ocr_cost(self, logging_obj) -> None: + """Regression: a deployment priced only per page was treated as unpriced, so a retrieved OCR batch + billed at the published rate while the same deployment's synchronous OCR calls billed at its own.""" + deployment_id = "deploy-ocr-only-pricing-1" + litellm.model_cost[deployment_id] = { + "id": deployment_id, + "litellm_provider": "mistral", + "mode": "ocr", + "ocr_cost_per_page": 0.0456, + "ocr_cost_per_page_batches": 0.0123, + } + logging_obj.litellm_params = { + "litellm_metadata": {"model_info": {"id": deployment_id}}, + "model": "mistral/mistral-ocr-latest", + } + logging_obj.model_call_details["model"] = "mistral/mistral-ocr-latest" + published_annotation_rate = litellm.model_cost["mistral/mistral-ocr-latest"]["annotation_cost_per_page_batches"] + try: + info = logging_obj.get_router_deployment_model_info() + assert info is not None + assert info["ocr_cost_per_page_batches"] == 0.0123 + pages_only = OCRUsageInfo(pages_processed=3) + assert ocr_batch_cost("mistral-ocr-latest", "mistral", pages_only, info)[0] == pytest.approx(3 * 0.0123) + with_annotations = OCRUsageInfo(pages_processed=3, pages_processed_annotation=2) + assert ocr_batch_cost("mistral-ocr-latest", "mistral", with_annotations, info)[0] == pytest.approx( + 3 * 0.0123 + 2 * published_annotation_rate + ) + finally: + litellm.model_cost.pop(deployment_id, None) + class TestRetrieveBatchCostPassesModelIdentity: """Regression: retrieving a batch priced it with no model identity at all. diff --git a/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py b/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py index 1dcd6d92a4b..5043cc583ef 100644 --- a/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py +++ b/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py @@ -13,6 +13,7 @@ import httpx import pytest from openai.types.file_deleted import FileDeleted +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.mistral.files.transformation import MistralFilesConfig from litellm.types.llms.openai import CreateFileRequest, FileContentRequest, OpenAIFileObject from litellm.types.utils import LlmProviders @@ -115,14 +116,16 @@ def test_upload_request_maps_user_data_onto_ocr(config): @pytest.mark.parametrize("purpose", ["assistants", "vision", "evals"]) def test_upload_request_rejects_purposes_mistral_lacks(config, purpose): """Regression: these used to be silently rewritten to ``batch``, so an upload that skipped the - proxy's batch-only validation and guardrails still landed on Mistral as a batch input file.""" - with pytest.raises(ValueError, match=f"purpose={purpose!r}"): + proxy's batch-only validation and guardrails still landed on Mistral as a batch input file. The + rejection is a 400 provider error, so the proxy answers invalid_request_error instead of a 500.""" + with pytest.raises(BaseLLMException, match=f"purpose={purpose!r}") as exc_info: config.transform_create_file_request( model="", create_file_data=CreateFileRequest(file=("f.bin", b"x"), purpose=purpose), optional_params={}, litellm_params={}, ) + assert exc_info.value.status_code == 400 def test_upload_request_requires_file(config): @@ -212,8 +215,9 @@ def test_list_request_accepts_the_purpose_an_ocr_file_reads_back_as(config): def test_list_request_rejects_purposes_mistral_lacks(config): - with pytest.raises(ValueError, match="purpose='assistants'"): + with pytest.raises(BaseLLMException, match="purpose='assistants'") as exc_info: config.transform_list_files_request(purpose="assistants", optional_params={}, litellm_params={}) + assert exc_info.value.status_code == 400 def test_list_response(config):