From 2cc73e2dcc36b6b4e0c6ba3740733fcf7e6d02bd Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 4 Sep 2026 11:53:32 -0700 Subject: [PATCH] fix(ocr): support token-priced billing --- litellm/cost_calculator.py | 18 ++++ litellm/litellm_core_utils/litellm_logging.py | 12 ++- litellm/llms/base_llm/ocr/transformation.py | 3 + ...odel_prices_and_context_window_backup.json | 1 - .../spend_tracking/spend_tracking_utils.py | 10 +- model_prices_and_context_window.json | 1 - .../test_ocr_spend_tracking.py | 86 +++++++++++++--- .../ocr/test_deepseek_transformation.py | 99 +++++++++++++++++++ 8 files changed, 208 insertions(+), 22 deletions(-) create mode 100644 tests/test_litellm/llms/vertex_ai/ocr/test_deepseek_transformation.py diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index cb7da32f857..7ccc1f29c7c 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1913,6 +1913,24 @@ def ocr_cost( except Exception: model_info = None + prompt_tokens: Final = response.usage_info.prompt_tokens + completion_tokens: Final = response.usage_info.completion_tokens + has_token_usage: Final = prompt_tokens is not None or completion_tokens is not None + has_token_pricing: Final = model_info is not None and ( + model_info.get("input_cost_per_token") is not None or model_info.get("output_cost_per_token") is not None + ) + if has_token_usage and has_token_pricing and model_info is not None: + return generic_cost_per_token( + model=model, + usage=Usage( + prompt_tokens=prompt_tokens or 0, + completion_tokens=completion_tokens or 0, + total_tokens=response.usage_info.total_tokens, + ), + custom_llm_provider=custom_llm_provider or str(model_info.get("litellm_provider") or ""), + model_info=model_info, + ) + credits: Final = getattr(response.usage_info, "credits", None) cost_per_credit = None if model_info is not None: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 17a19f05fa3..f756cfe5893 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -13,6 +13,7 @@ import traceback from collections.abc import Callable, Iterator, Mapping, Sequence from datetime import datetime as dt_object from functools import lru_cache +from itertools import chain from types import MappingProxyType, TracebackType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast @@ -5409,7 +5410,9 @@ class StandardLoggingPayloadSetup: return combined_usage_object.model_dump() if not response_obj: return _empty - _raw: Final = response_obj.get("usage", None) + usage: Final = response_obj.get("usage") + uses_usage_info: Final = usage is None + _raw: Final = response_obj.get("usage_info") if uses_usage_info else usage if _raw is None: return _empty if isinstance(_raw, ResponseAPIUsage): @@ -5419,6 +5422,13 @@ class StandardLoggingPayloadSetup: return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(_raw).model_dump() if InteractionsUsageObjectTransformation.is_interactions_usage_object(_raw): return InteractionsUsageObjectTransformation.transform_interactions_usage_object(_raw).model_dump() + if uses_usage_info: + token_fields: Final = frozenset({"prompt_tokens", "completion_tokens", "total_tokens"}) + return { # mutable-ok: the existing logging API returns a mutable usage dictionary + key: value + for key, value in chain(_empty.items(), _raw.items()) + if key not in token_fields or value is not None + } return _raw if isinstance(_raw, Usage): return _raw.model_dump() diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index 75306cd572a..d439ac39421 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -79,6 +79,9 @@ class OCRUsageInfo(LiteLLMPydanticObjectBase): pages_processed_annotation: int | None = None credits: float | None = None doc_size_bytes: int | None = None + prompt_tokens: int | None = None + completion_tokens: int | None = None + total_tokens: int | None = None model_config = {"extra": "allow"} diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 44c4f10ec38..6e51b7ad12d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -45975,7 +45975,6 @@ "mode": "ocr", "input_cost_per_token": 3e-07, "output_cost_per_token": 1.2e-06, - "ocr_cost_per_page": 0.0003, "source": "https://cloud.google.com/vertex-ai/pricing", "supported_regions": [ "us-central1" diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index a37c3ba4405..698e2db1f86 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -280,10 +280,14 @@ def _extract_usage_for_ocr_call(response_obj: object, response_obj_dict: dict) - "total_tokens": 0, } # Add all fields from usage_info, including pages_processed + token_fields: Final = frozenset({"prompt_tokens", "completion_tokens", "total_tokens"}) for key, value in usage_info.items(): - result[key] = value - # Ensure pages_processed exists - if "pages_processed" not in result: + if key not in token_fields or value is not None: + result[key] = value + has_token_usage: Final = any( + usage_info.get(key) is not None for key in ("prompt_tokens", "completion_tokens", "total_tokens") + ) + if "pages_processed" not in result and not has_token_usage: result["pages_processed"] = 0 return result else: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 44c4f10ec38..6e51b7ad12d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -45975,7 +45975,6 @@ "mode": "ocr", "input_cost_per_token": 3e-07, "output_cost_per_token": 1.2e-06, - "ocr_cost_per_page": 0.0003, "source": "https://cloud.google.com/vertex-ai/pricing", "supported_regions": [ "us-central1" diff --git a/tests/spend_tracking_tests/test_ocr_spend_tracking.py b/tests/spend_tracking_tests/test_ocr_spend_tracking.py index 3ce77c56361..b0439b526cc 100644 --- a/tests/spend_tracking_tests/test_ocr_spend_tracking.py +++ b/tests/spend_tracking_tests/test_ocr_spend_tracking.py @@ -2,18 +2,20 @@ Unit tests for OCR spend tracking in get_logging_payload. This test file verifies that OCR/AOCR calls correctly extract usage_info -and populate the spend logs payload with pages_processed instead of token counts. +and populate spend logs with page or token billing usage. """ -import pytest +import json from datetime import datetime, timezone -from unittest.mock import Mock -from pydantic import BaseModel from typing import Optional +from unittest.mock import Mock + +import pytest +from pydantic import BaseModel from litellm.proxy.spend_tracking.spend_tracking_utils import ( - get_logging_payload, _extract_usage_for_ocr_call, + get_logging_payload, ) @@ -50,9 +52,7 @@ class TestExtractUsageForOCRCall: def test_extract_usage_from_pydantic_model(self): """Test extracting usage from Pydantic model response""" usage_info = MockUsageInfo(pages_processed=10, doc_size_bytes=1024) - response_obj = MockOCRResponse( - id="ocr-123", object="ocr", model="test-ocr-model", usage_info=usage_info - ) + response_obj = MockOCRResponse(id="ocr-123", object="ocr", model="test-ocr-model", usage_info=usage_info) response_obj_dict = response_obj.model_dump() usage = _extract_usage_for_ocr_call(response_obj, response_obj_dict) @@ -102,6 +102,42 @@ class TestExtractUsageForOCRCall: assert usage.get("total_tokens") == 0 assert usage.get("pages_processed") == 0 + def test_extract_usage_preserves_token_billed_ocr(self): + response_obj_dict = { + "usage_info": { + "prompt_tokens": 281, + "completion_tokens": 6, + "total_tokens": 287, + } + } + + usage = _extract_usage_for_ocr_call(response_obj_dict, response_obj_dict) + + assert usage == { + "prompt_tokens": 281, + "completion_tokens": 6, + "total_tokens": 287, + } + + def test_extract_usage_keeps_zero_token_defaults_for_page_billed_ocr(self): + response_obj_dict = { + "usage_info": { + "pages_processed": 5, + "prompt_tokens": None, + "completion_tokens": None, + "total_tokens": None, + } + } + + usage = _extract_usage_for_ocr_call(response_obj_dict, response_obj_dict) + + assert usage == { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + "pages_processed": 5, + } + class TestGetLoggingPayloadOCR: """Test get_logging_payload with OCR call types""" @@ -144,8 +180,6 @@ class TestGetLoggingPayloadOCR: assert payload["spend"] == 0.05 # Verify pages_processed is in additional_usage_values - import json - metadata = json.loads(payload["metadata"]) assert "additional_usage_values" in metadata assert metadata["additional_usage_values"]["pages_processed"] == 7 @@ -175,8 +209,6 @@ class TestGetLoggingPayloadOCR: assert payload["total_tokens"] == 0 # Verify pages_processed is in additional_usage_values - import json - metadata = json.loads(payload["metadata"]) assert "additional_usage_values" in metadata assert metadata["additional_usage_values"]["pages_processed"] == 12 @@ -223,8 +255,6 @@ class TestGetLoggingPayloadOCR: assert payload["total_tokens"] == 0 # Verify pages_processed is 0 - import json - metadata = json.loads(payload["metadata"]) assert metadata["additional_usage_values"]["pages_processed"] == 0 @@ -289,8 +319,32 @@ class TestGetLoggingPayloadOCR: assert payload["completion_tokens"] == 0 # Verify pages_processed and doc_size_bytes are both in additional_usage_values - import json - metadata = json.loads(payload["metadata"]) assert metadata["additional_usage_values"]["pages_processed"] == 5 assert metadata["additional_usage_values"]["doc_size_bytes"] == 1024 + + def test_token_billed_ocr_maps_usage_without_pages(self, mock_datetime, base_kwargs): + response_obj = { + "id": "ocr-token-billed", + "object": "ocr", + "model": "deepseek-ocr-maas", + "usage_info": { + "prompt_tokens": 281, + "completion_tokens": 6, + "total_tokens": 287, + }, + } + + payload = get_logging_payload( + kwargs=base_kwargs, + response_obj=response_obj, + start_time=mock_datetime, + end_time=mock_datetime, + ) + + assert payload["prompt_tokens"] == 281 + assert payload["completion_tokens"] == 6 + assert payload["total_tokens"] == 287 + + metadata = json.loads(payload["metadata"]) + assert "pages_processed" not in metadata["additional_usage_values"] diff --git a/tests/test_litellm/llms/vertex_ai/ocr/test_deepseek_transformation.py b/tests/test_litellm/llms/vertex_ai/ocr/test_deepseek_transformation.py new file mode 100644 index 00000000000..a6aa4d1e595 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/ocr/test_deepseek_transformation.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +import math +from collections.abc import Callable +from typing import Final, cast + +import httpx + +import litellm +import litellm.cost_calculator as cost_calculator +from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.llms.vertex_ai.ocr.deepseek_transformation import VertexAIDeepSeekOCRConfig + +MODEL: Final = "vertex_ai/deepseek-ai/deepseek-ocr-maas" +PROMPT_TOKENS: Final = 281 +COMPLETION_TOKENS: Final = 6 + + +def _response() -> OCRResponse: + transform_response: Final = cast( # cast-ok: the legacy provider method has untyped variadic parameters + Callable[..., OCRResponse], VertexAIDeepSeekOCRConfig().transform_ocr_response + ) + return transform_response( + model="deepseek-ai/deepseek-ocr-maas", + raw_response=httpx.Response( + status_code=200, + json={ + "choices": [{"message": {"content": "invoice 123"}}], + "usage": { + "prompt_tokens": PROMPT_TOKENS, + "completion_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, + }, + }, + ), + logging_obj=None, + ) + + +def test_deepseek_ocr_preserves_billable_token_usage() -> None: + response: Final = _response() + + assert response.usage_info is not None + assert response.usage_info.prompt_tokens == PROMPT_TOKENS + assert response.usage_info.completion_tokens == COMPLETION_TOKENS + assert response.usage_info.total_tokens == PROMPT_TOKENS + COMPLETION_TOKENS + assert response.usage_info.pages_processed is None + + +def test_ocr_usage_maps_to_standard_logging() -> None: + normalize_usage: Final = cast( # cast-ok: the legacy logging helper has an untyped dictionary signature + Callable[..., dict[str, object]], StandardLoggingPayloadSetup.get_usage_as_dict + ) + token_usage: Final = normalize_usage(response_obj=_response().model_dump()) + page_usage: Final = normalize_usage( + response_obj={ + "usage_info": { + "pages_processed": 5, + "prompt_tokens": None, + "completion_tokens": None, + "total_tokens": None, + } + } + ) + + assert token_usage["prompt_tokens"] == PROMPT_TOKENS + assert token_usage["completion_tokens"] == COMPLETION_TOKENS + assert token_usage["total_tokens"] == PROMPT_TOKENS + COMPLETION_TOKENS + assert page_usage == { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + "pages_processed": 5, + } + + +def test_deepseek_ocr_cost_uses_provider_token_rates(local_model_cost_map: object) -> None: + model_info: Final = litellm.get_model_info(model=MODEL, custom_llm_provider="vertex_ai") + response: Final = _response() + calculate_cost: Final = cast( # cast-ok: completion_cost has legacy untyped optional parameters + Callable[..., float], cost_calculator.completion_cost + ) + + cost: Final = calculate_cost( + completion_response=response, + model=MODEL, + custom_llm_provider="vertex_ai", + call_type="ocr", + ) + + input_cost_per_token: Final = model_info["input_cost_per_token"] + output_cost_per_token: Final = model_info["output_cost_per_token"] + assert input_cost_per_token is not None + assert output_cost_per_token is not None + expected: Final = PROMPT_TOKENS * input_cost_per_token + COMPLETION_TOKENS * output_cost_per_token + assert model_info.get("ocr_cost_per_page") is None + assert math.isclose(cost, expected, rel_tol=1e-12) + assert cost > 0