fix(cost): price Vertex AI DeepSeek OCR by token usage

DeepSeek OCR on Vertex AI reports prompt and completion tokens and never a
page count, so the OCR cost calculator either raised (cost-map name, no cost
header at all) or returned 0.0 (short name, no cost-map entry). OCR cost now
falls back to the model's token rates when no page pricing applies, the
DeepSeek transform reports the canonical deepseek-ai/ model on the response so
the short deployment name resolves the cost-map entry, and the unsourced
ocr_cost_per_page is dropped from that entry since Google bills it per token.
The Rust OCR port mirrors the canonical response model.
This commit is contained in:
mateo-berri 2026-09-02 14:50:42 -07:00
parent ba2e5d2d8e
commit 24299c3c5e
8 changed files with 207 additions and 15 deletions

View file

@ -298,7 +298,8 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
})?;
let usage = response.get("usage").cloned();
let content = first_choice_content(&response_json)?;
let mut ocr_data = ocr_data_from_content(content.clone(), usage.clone(), model);
let provider_model = deepseek_model_name(model);
let mut ocr_data = ocr_data_from_content(content.clone(), usage.clone(), &provider_model);
if !ocr_data.get("pages").is_some_and(Value::is_array) {
ocr_data = json!({
@ -309,7 +310,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
other => other.to_string(),
}
}],
"model": ocr_data.get("model").and_then(Value::as_str).unwrap_or(model),
"model": ocr_data.get("model").and_then(Value::as_str).unwrap_or(&provider_model),
"usage_info": ocr_data.get("usage_info").cloned().or(usage).unwrap_or_else(|| json!({})),
});
}
@ -332,7 +333,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
model: object
.get("model")
.and_then(Value::as_str)
.unwrap_or(model)
.unwrap_or(&provider_model)
.to_string(),
document_annotation: object.get("document_annotation").cloned(),
usage_info,
@ -429,7 +430,7 @@ mod tests {
response.pages,
vec![json!({"index": 0, "markdown": "# OCR text"})]
);
assert_eq!(response.model, "deepseek-ocr-maas");
assert_eq!(response.model, "deepseek-ai/deepseek-ocr-maas");
assert_eq!(response.usage_info, Some(json!({"prompt_tokens": 1})));
}
}

View file

@ -133,6 +133,7 @@ if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import (
Logging as LitellmLoggingObject,
)
from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo
else:
LitellmLoggingObject = Any
@ -1871,6 +1872,23 @@ def response_cost_calculator(
raise e
def _ocr_token_cost(usage_info: "OCRUsageInfo", model_info: ModelInfo | None) -> tuple[float, float] | None:
if model_info is None:
return None
input_cost_per_token: Final = model_info.get("input_cost_per_token") or 0.0
output_cost_per_token: Final = model_info.get("output_cost_per_token") or 0.0
if input_cost_per_token == 0.0 and output_cost_per_token == 0.0:
return None
token_counts: Final = usage_info.model_extra
if token_counts is None:
return None
prompt_tokens: Final = token_counts.get("prompt_tokens")
completion_tokens: Final = token_counts.get("completion_tokens")
if not isinstance(prompt_tokens, int) or not isinstance(completion_tokens, int):
return None
return prompt_tokens * input_cost_per_token, completion_tokens * output_cost_per_token
def ocr_cost(
model: str,
custom_llm_provider: str | None,
@ -1883,9 +1901,7 @@ def ocr_cost(
response: Optional[Any] - response object
Returns:
Tuple[float, float]: cost of OCR processing
(Parent function requires a tuple, so we return a tuple. Cost is only in the first element.)
Tuple[float, float]: (prompt cost, completion cost) when priced per token, otherwise the page cost in the first element
"""
from litellm.llms.base_llm.ocr.transformation import OCRResponse
@ -1917,6 +1933,12 @@ def ocr_cost(
pages_processed: Final = response.usage_info.pages_processed
annotation_pages: Final = response.usage_info.pages_processed_annotation or 0
has_billable_annotation_pages: Final = annotation_rate is not None and annotation_pages > 0
has_page_pricing: Final = (
pages_processed is not None and ocr_cost_per_page is not None
) or has_billable_annotation_pages
token_cost: Final = _ocr_token_cost(response.usage_info, model_info)
if not has_page_pricing and token_cost is not None:
return token_cost
if pages_processed is None and not has_billable_annotation_pages:
if cost_per_credit is not None or ocr_cost_per_page is None:

View file

@ -26,6 +26,10 @@ else:
LiteLLMLoggingObj = Any
def _provider_model_name(model: str) -> str:
return model if model.startswith("deepseek-ai/") else f"deepseek-ai/{model}"
class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
"""
Vertex AI DeepSeek OCR transformation configuration.
@ -177,7 +181,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
content_item = {"type": "image_url", "image_url": document_url}
# Build DeepSeek OCR request
provider_model: Final = model if model.startswith("deepseek-ai/") else f"deepseek-ai/{model}"
provider_model: Final = _provider_model_name(model)
data: Final = {
"model": provider_model,
"messages": [{"role": "user", "content": [content_item]}],
@ -261,6 +265,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
"""
verbose_logger.debug("Vertex AI DeepSeek OCR transform_ocr_response called")
verbose_logger.debug("Raw response: %s", raw_response.text)
provider_model: Final = _provider_model_name(model)
try:
response_json: Final = raw_response.json()
@ -288,14 +293,14 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
# If content is markdown text, create a single page with the markdown
ocr_data = {
"pages": [{"index": 0, "markdown": content}],
"model": model,
"model": provider_model,
"usage_info": response_json.get("usage", {}),
}
except json.JSONDecodeError:
# If JSON parsing fails, treat content as markdown
ocr_data = {
"pages": [{"index": 0, "markdown": content}],
"model": model,
"model": provider_model,
"usage_info": response_json.get("usage", {}),
}
@ -309,7 +314,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
"markdown": (content if isinstance(content, str) else json.dumps(content)),
}
],
"model": ocr_data.get("model", model),
"model": ocr_data.get("model", provider_model),
"usage_info": ocr_data.get("usage_info", response_json.get("usage", {})),
}
@ -339,7 +344,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
return OCRResponse(
pages=pages,
model=ocr_data.get("model", model),
model=ocr_data.get("model", provider_model),
document_annotation=ocr_data.get("document_annotation"),
usage_info=usage_info,
object="ocr",

View file

@ -45806,7 +45806,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"

View file

@ -45806,7 +45806,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"

View file

@ -0,0 +1,67 @@
from unittest.mock import MagicMock
import httpx
import pytest
from litellm.cost_calculator import completion_cost
from litellm.llms.vertex_ai.ocr.deepseek_transformation import VertexAIDeepSeekOCRConfig
PROMPT_TOKENS = 901
COMPLETION_TOKENS = 212
INPUT_COST_PER_TOKEN = 3e-07
OUTPUT_COST_PER_TOKEN = 1.2e-06
def _deepseek_chat_response() -> httpx.Response:
return httpx.Response(
status_code=200,
json={
"choices": [{"message": {"role": "assistant", "content": "# OCR text"}}],
"usage": {
"prompt_tokens": PROMPT_TOKENS,
"completion_tokens": COMPLETION_TOKENS,
"total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS,
},
},
request=httpx.Request("POST", "https://us-central1-aiplatform.googleapis.com"),
)
@pytest.mark.parametrize("model", ["deepseek-ocr-maas", "deepseek-ai/deepseek-ocr-maas"])
def test_response_is_priced_from_token_usage_for_either_model_name(local_model_cost_map: None, model: str) -> None:
response = VertexAIDeepSeekOCRConfig().transform_ocr_response(
model=model,
raw_response=_deepseek_chat_response(),
logging_obj=MagicMock(),
)
cost = completion_cost(
completion_response=response,
model=f"vertex_ai/{model}",
custom_llm_provider="vertex_ai",
call_type="ocr",
)
assert response.model == "deepseek-ai/deepseek-ocr-maas"
assert cost == pytest.approx(PROMPT_TOKENS * INPUT_COST_PER_TOKEN + COMPLETION_TOKENS * OUTPUT_COST_PER_TOKEN)
assert cost > 0
def test_json_content_without_pages_reports_the_canonical_model(local_model_cost_map: None) -> None:
raw_response = httpx.Response(
200,
json={
"choices": [{"message": {"role": "assistant", "content": '{"text": "# OCR text"}'}}],
"usage": {"prompt_tokens": PROMPT_TOKENS, "completion_tokens": COMPLETION_TOKENS},
},
request=httpx.Request("POST", "https://example.invalid"),
)
response = VertexAIDeepSeekOCRConfig().transform_ocr_response(
model="deepseek-ocr-maas",
raw_response=raw_response,
logging_obj=MagicMock(),
)
assert response.model == "deepseek-ai/deepseek-ocr-maas"
assert response.pages[0].markdown == '{"text": "# OCR text"}'

View file

@ -10,12 +10,14 @@ from pydantic import BaseModel
import litellm
from litellm.cost_calculator import (
BaseTokenUsageProcessor,
RealtimeAPITokenUsageProcessor,
completion_cost,
cost_per_token,
handle_realtime_stream_cost_calculation,
ocr_cost,
RealtimeAPITokenUsageProcessor,
response_cost_calculator,
)
from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo
from litellm.types.llms.openai import OpenAIRealtimeStreamList
from litellm.types.utils import (
CacheCreationTokenDetails,
@ -4473,3 +4475,100 @@ def test_explicit_pricing_precedes_private_provider_response_model(
)
assert selected == expected
def test_ocr_cost_prices_token_usage_when_pages_are_not_reported(_local_model_cost_map):
response = OCRResponse(
pages=[OCRPage(index=0, markdown="# OCR text")],
model="vertex_ai/deepseek-ai/deepseek-ocr-maas",
usage_info=OCRUsageInfo.model_validate({"prompt_tokens": 901, "completion_tokens": 278, "total_tokens": 1179}),
)
cost = completion_cost(
completion_response=response,
model="vertex_ai/deepseek-ai/deepseek-ocr-maas",
custom_llm_provider="vertex_ai",
call_type="ocr",
)
assert cost == pytest.approx(901 * 3e-07 + 278 * 1.2e-06)
def test_ocr_cost_does_not_price_tokens_for_page_priced_models(_local_model_cost_map):
response = OCRResponse(
pages=[OCRPage(index=0, markdown="# OCR text")],
model="mistral/mistral-ocr-latest",
usage_info=OCRUsageInfo.model_validate({"prompt_tokens": 901, "completion_tokens": 278}),
)
with pytest.raises(ValueError, match="pages_processed is None"):
completion_cost(
completion_response=response,
model="mistral/mistral-ocr-latest",
custom_llm_provider="mistral",
call_type="ocr",
)
def test_ocr_cost_prefers_page_pricing_when_pages_are_reported(_local_model_cost_map, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setitem(
litellm.model_cost,
"vertex_ai/page-and-token-ocr",
{
"mode": "ocr",
"litellm_provider": "vertex_ai",
"input_cost_per_token": 3e-07,
"output_cost_per_token": 1.2e-06,
"ocr_cost_per_page": 0.001,
},
)
response = OCRResponse(
pages=[OCRPage(index=0, markdown="# OCR text")],
model="vertex_ai/page-and-token-ocr",
usage_info=OCRUsageInfo.model_validate(
{"pages_processed": 2, "prompt_tokens": 901, "completion_tokens": 278}
),
)
cost = completion_cost(
completion_response=response,
model="vertex_ai/page-and-token-ocr",
custom_llm_provider="vertex_ai",
call_type="ocr",
)
assert cost == pytest.approx(2 * 0.001)
def test_ocr_cost_splits_token_cost_into_prompt_and_completion(_local_model_cost_map):
response = OCRResponse(
pages=[OCRPage(index=0, markdown="# OCR text")],
model="vertex_ai/deepseek-ai/deepseek-ocr-maas",
usage_info=OCRUsageInfo.model_validate({"prompt_tokens": 901, "completion_tokens": 278}),
)
prompt_cost, completion_cost_value = ocr_cost(
model="vertex_ai/deepseek-ai/deepseek-ocr-maas",
custom_llm_provider="vertex_ai",
response=response,
)
assert prompt_cost == pytest.approx(901 * 3e-07)
assert completion_cost_value == pytest.approx(278 * 1.2e-06)
def test_ocr_cost_stays_zero_when_token_priced_response_lacks_token_counts(_local_model_cost_map):
response = OCRResponse(
pages=[OCRPage(index=0, markdown="# OCR text")],
model="vertex_ai/deepseek-ai/deepseek-ocr-maas",
usage_info=OCRUsageInfo(),
)
cost = completion_cost(
completion_response=response,
model="vertex_ai/deepseek-ai/deepseek-ocr-maas",
custom_llm_provider="vertex_ai",
call_type="ocr",
)
assert cost == 0.0