tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-21 12:20:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Feat] Add Cost Tracking for /ocr endpoints (#15678)

Browse Source 
* mistral/mistral-ocr-latest

* fix: add _hidden_params to OCRResponses

* test: hidden params exists

* feat: add mistral/mistral-ocr-2505-completion

* fix test

* add ModelInfoBase fields

* fix get OCR cost

* check response cost from OCR

* add handling for OCR costs

* add mistral-document-ai-2505

* docs OCR

* ruff check fix
This commit is contained in:
Ishaan Jaff
2025-10-17 15:54:10 -07:00
committed by GitHub
parent 6c26971cd4
commit a8be3ae412
8 changed files with 152 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/ocr.md
+6
View File
@@ -1,5 +1,11 @@
# /ocr
| Feature | Supported |
|---------|-----------|
| Cost Tracking | ✅ |
| Logging | ✅ (Basic Logging not supported) |
| Load Balancing | ✅ |
:::tip
LiteLLM follows the [Mistral API request/response for the OCR API](https://docs.mistral.ai/capabilities/vision/#optical-character-recognition-ocr)
litellm/cost_calculator.py
+54
View File
@@ -153,6 +153,7 @@ def cost_per_token( # noqa: PLR0915
audio_transcription_file_duration: float = 0.0, # for audio transcription calls - the file time in seconds
### SERVICE TIER ###
service_tier: Optional[str] = None, # for OpenAI service tier pricing
response: Optional[Any] = None,
) -> Tuple[float, float]: # type: ignore
"""
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
@@ -293,6 +294,12 @@ def cost_per_token( # noqa: PLR0915
custom_llm_provider=custom_llm_provider,
billed_units=rerank_billed_units,
)
elif call_type == "ocr" or call_type == "aocr":
return ocr_cost(
model=model,
custom_llm_provider=custom_llm_provider,
response=response,
)
elif (
call_type == "aretrieve_batch"
or call_type == "retrieve_batch"
@@ -1001,6 +1008,7 @@ def completion_cost( # noqa: PLR0915
audio_transcription_file_duration=audio_transcription_file_duration,
rerank_billed_units=rerank_billed_units,
service_tier=service_tier,
response=completion_response,
)
_final_cost = (
prompt_tokens_cost_usd_dollar + completion_tokens_cost_usd_dollar
@@ -1158,6 +1166,52 @@ def response_cost_calculator(
raise e
def ocr_cost(
model: str,
custom_llm_provider: Optional[str],
response: Optional[Any] = None,
) -> Tuple[float, float]:
"""
Args:
model: str - model name
custom_llm_provider: Optional[str] - custom LLM provider
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.)
"""
from litellm.llms.base_llm.ocr.transformation import OCRResponse
#########################################################
# validate it's an OCR response
#########################################################
if response is None or not isinstance(response, OCRResponse):
raise ValueError(f"response must be of type OCRResponse got type={type(response)}")
if response.usage_info is None:
raise ValueError("OCR response usage_info is None")
pages_processed = response.usage_info.pages_processed
if pages_processed is None:
raise ValueError("OCR response pages_processed is None")
try:
model_info: Optional[ModelInfo] = litellm.get_model_info(
model=model, custom_llm_provider=custom_llm_provider
)
except Exception:
model_info = None
ocr_cost_per_page: float = 0.0
if model_info is not None:
ocr_cost_per_page = model_info.get("ocr_cost_per_page") or 0.0
total_ocr_processing_cost: float = ocr_cost_per_page * pages_processed
return total_ocr_processing_cost, 0.0
def rerank_cost(
model: str,
custom_llm_provider: Optional[str],
litellm/llms/base_llm/ocr/transformation.py
+11 -7
View File
@@ -4,9 +4,10 @@ Base OCR transformation configuration.
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
import httpx
from pydantic import BaseModel
from pydantic import PrivateAttr
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.llms.base import LiteLLMPydanticObjectBase
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@@ -18,14 +19,14 @@ else:
DocumentType = Dict[str, str]
class OCRPageDimensions(BaseModel):
class OCRPageDimensions(LiteLLMPydanticObjectBase):
"""Page dimensions from OCR response."""
dpi: Optional[int] = None
height: Optional[int] = None
width: Optional[int] = None
class OCRPageImage(BaseModel):
class OCRPageImage(LiteLLMPydanticObjectBase):
"""Image extracted from OCR page."""
image_base64: Optional[str] = None
bbox: Optional[Dict[str, Any]] = None
@@ -33,7 +34,7 @@ class OCRPageImage(BaseModel):
model_config = {"extra": "allow"}
class OCRPage(BaseModel):
class OCRPage(LiteLLMPydanticObjectBase):
"""Single page from OCR response."""
index: int
markdown: str
@@ -43,7 +44,7 @@ class OCRPage(BaseModel):
model_config = {"extra": "allow"}
class OCRUsageInfo(BaseModel):
class OCRUsageInfo(LiteLLMPydanticObjectBase):
"""Usage information from OCR response."""
pages_processed: Optional[int] = None
doc_size_bytes: Optional[int] = None
@@ -51,7 +52,7 @@ class OCRUsageInfo(BaseModel):
model_config = {"extra": "allow"}
class OCRResponse(BaseModel):
class OCRResponse(LiteLLMPydanticObjectBase):
"""
Standard OCR response format.
Standardized to Mistral OCR format - other providers should transform to this format.
@@ -64,8 +65,11 @@ class OCRResponse(BaseModel):
model_config = {"extra": "allow"}
# Define private attributes using PrivateAttr
_hidden_params: dict = PrivateAttr(default_factory=dict)
class OCRRequestData(BaseModel):
class OCRRequestData(LiteLLMPydanticObjectBase):
"""OCR request data structure."""
data: Optional[Union[Dict, bytes]] = None
files: Optional[Dict[str, Any]] = None
litellm/model_prices_and_context_window_backup.json
+30 -1
View File
@@ -3275,6 +3275,15 @@
"supports_tool_choice": true,
"supports_reasoning": true
},
"azure_ai/mistral-document-ai-2505": {
"litellm_provider": "azure_ai",
"ocr_cost_per_page": 3e-3,
"mode": "ocr",
"supported_endpoints": [
"/v1/ocr"
],
"source": "https://devblogs.microsoft.com/foundry/whats-new-in-azure-ai-foundry-august-2025/#mistral-document-ai-(ocr)-%E2%80%94-serverless-in-foundry"
},
"azure_ai/MAI-DS-R1": {
"input_cost_per_token": 1.35e-06,
"litellm_provider": "azure_ai",
@@ -9649,7 +9658,7 @@
"supports_tool_choice": true,
"supports_url_context": true,
"supports_vision": true,
"supports_web_search": true,
"supports_web_search": false,
"tpm": 8000000
},
"gemini-2.5-flash-image-preview": {
@@ -15570,6 +15579,26 @@
"supports_response_schema": true,
"supports_tool_choice": true
},
"mistral/mistral-ocr-latest": {
"litellm_provider": "mistral",
"ocr_cost_per_page": 1e-3,
"annotation_cost_per_page": 3e-3,
"mode": "ocr",
"supported_endpoints": [
"/v1/ocr"
],
"source": "https://mistral.ai/pricing#api-pricing"
},
"mistral/mistral-ocr-2505-completion": {
"litellm_provider": "mistral",
"ocr_cost_per_page": 1e-3,
"annotation_cost_per_page": 3e-3,
"mode": "ocr",
"supported_endpoints": [
"/v1/ocr"
],
"source": "https://mistral.ai/pricing#api-pricing"
},
"mistral/magistral-medium-latest": {
"input_cost_per_token": 2e-06,
"litellm_provider": "mistral",
litellm/types/utils.py
+4
View File
@@ -173,6 +173,8 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
output_cost_per_video_per_second: Optional[float] # only for vertex ai models
output_cost_per_audio_per_second: Optional[float] # only for vertex ai models
output_cost_per_second: Optional[float] # for OpenAI Speech models
ocr_cost_per_page: Optional[float] # for OCR models
annotation_cost_per_page: Optional[float] # for OCR models
search_context_cost_per_query: Optional[
SearchContextCostPerQuery
] # Cost for using web search tool
@@ -330,6 +332,8 @@ CallTypesLiteral = Literal[
"agenerate_content",
"generate_content_stream",
"agenerate_content_stream",
"ocr",
"aocr",
]
litellm/utils.py
+2
View File
@@ -4994,6 +4994,8 @@ def _get_model_info_helper( # noqa: PLR0915
),
tpm=_model_info.get("tpm", None),
rpm=_model_info.get("rpm", None),
ocr_cost_per_page=_model_info.get("ocr_cost_per_page", None),
annotation_cost_per_page=_model_info.get("annotation_cost_per_page", None),
)
except Exception as e:
verbose_logger.debug(f"Error getting model info: {e}")
model_prices_and_context_window.json
+29
View File
@@ -3275,6 +3275,15 @@
"supports_tool_choice": true,
"supports_reasoning": true
},
"azure_ai/mistral-document-ai-2505": {
"litellm_provider": "azure_ai",
"ocr_cost_per_page": 3e-3,
"mode": "ocr",
"supported_endpoints": [
"/v1/ocr"
],
"source": "https://devblogs.microsoft.com/foundry/whats-new-in-azure-ai-foundry-august-2025/#mistral-document-ai-(ocr)-%E2%80%94-serverless-in-foundry"
},
"azure_ai/MAI-DS-R1": {
"input_cost_per_token": 1.35e-06,
"litellm_provider": "azure_ai",
@@ -15570,6 +15579,26 @@
"supports_response_schema": true,
"supports_tool_choice": true
},
"mistral/mistral-ocr-latest": {
"litellm_provider": "mistral",
"ocr_cost_per_page": 1e-3,
"annotation_cost_per_page": 3e-3,
"mode": "ocr",
"supported_endpoints": [
"/v1/ocr"
],
"source": "https://mistral.ai/pricing#api-pricing"
},
"mistral/mistral-ocr-2505-completion": {
"litellm_provider": "mistral",
"ocr_cost_per_page": 1e-3,
"annotation_cost_per_page": 3e-3,
"mode": "ocr",
"supported_endpoints": [
"/v1/ocr"
],
"source": "https://mistral.ai/pricing#api-pricing"
},
"mistral/magistral-medium-latest": {
"input_cost_per_token": 2e-06,
"litellm_provider": "mistral",
tests/ocr_tests/base_ocr_unit_tests.py
+16
View File
@@ -5,6 +5,7 @@ This follows the same pattern as BaseLLMChatTest in tests/llm_translation/base_l
"""
import pytest
import litellm
import os
from abc import ABC, abstractmethod
@@ -45,6 +46,8 @@ class BaseOCRTest(ABC):
litellm._turn_on_debug()
base_ocr_call_args = self.get_base_ocr_call_args()
print("BASE OCR Call args=", base_ocr_call_args)
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
try:
if sync_mode:
@@ -96,6 +99,19 @@ class BaseOCRTest(ABC):
assert len(total_text) > 0, "Should extract some text from the document"
#########################################################
# validate we get a response cost in hidden parameters
#########################################################
hidden_params = response._hidden_params
assert isinstance(hidden_params, dict), "Hidden parameters should be a dictionary"
print("response usage_info:", response.usage_info)
response_cost = hidden_params.get("response_cost")
assert response_cost is not None, "Response cost should be in hidden parameters"
assert response_cost > 0, "Response cost should be greater than 0"
print("response_cost=", response_cost)
except Exception as e:
pytest.fail(f"OCR call failed: {str(e)}")