scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(ocr): support token-priced billing

Browse source
This commit is contained in:
Yujong Lee 2026-09-04 11:53:32 -07:00
parent f74bc9427b
commit 2cc73e2dcc
8 changed files with 208 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
litellm/cost_calculator.py
View file

@ -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:

12
litellm/litellm_core_utils/litellm_logging.py
View file

@ -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()

3
litellm/llms/base_llm/ocr/transformation.py
View file

@ -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"}

1
litellm/model_prices_and_context_window_backup.json
View file

@ -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"

10
litellm/proxy/spend_tracking/spend_tracking_utils.py
View file

@ -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:

1
model_prices_and_context_window.json
View file

@ -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"

86
tests/spend_tracking_tests/test_ocr_spend_tracking.py
View file

@ -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"]

99
tests/test_litellm/llms/vertex_ai/ocr/test_deepseek_transformation.py Normal file
View file

@ -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