mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
refactor(cost): share the deployment model_info lookup between video and OCR cost paths
Extract one immutable helper for reading the deployment's model_info off the logging object, stop rebinding the model_info parameter inside ocr_cost, drop the explanatory comment blocks, and move the OCR custom pricing regression tests into tests/test_litellm/test_cost_calculator.py
This commit is contained in:
mateo-berri
parent
f5637c85bf
commit
39239ab974
3 changed files with 166 additions and 201 deletions
|
|
@ -341,7 +341,7 @@ def cost_per_token(
|
|||
### REQUEST MODEL ###
|
||||
request_model: str | None = None, # original request model for router detection
|
||||
### DEPLOYMENT-SPECIFIC PRICING ###
|
||||
custom_model_info: ModelInfo | None = None, # deployment model_info, for non-token custom pricing
|
||||
custom_model_info: ModelInfo | None = None,
|
||||
) -> tuple[float, float]:
|
||||
"""
|
||||
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
|
||||
|
|
@ -1429,20 +1429,7 @@ def completion_cost(
|
|||
)
|
||||
elif call_type in _VIDEO_CALL_TYPES:
|
||||
### VIDEO GENERATION COST CALCULATION ###
|
||||
# Extract custom model_info for deployment-specific pricing
|
||||
_video_model_info: ModelInfo | None = None
|
||||
if custom_pricing and litellm_logging_obj is not None:
|
||||
_litellm_params = getattr(litellm_logging_obj, "litellm_params", None)
|
||||
if _litellm_params is not None:
|
||||
_video_model_info = next(
|
||||
(
|
||||
model_info
|
||||
for _metadata_key in ("metadata", "litellm_metadata")
|
||||
if (model_info := (_litellm_params.get(_metadata_key) or {}).get("model_info"))
|
||||
is not None
|
||||
),
|
||||
None,
|
||||
)
|
||||
_video_model_info: ModelInfo | None = _deployment_model_info(litellm_logging_obj, custom_pricing)
|
||||
|
||||
usage_obj = getattr(completion_response, "usage", None)
|
||||
duration_seconds: float | None = None
|
||||
|
|
@ -1638,24 +1625,6 @@ def completion_cost(
|
|||
if litellm_logging_obj is not None:
|
||||
request_model_for_cost = litellm_logging_obj.model
|
||||
|
||||
# Deployment-specific model_info, for modalities whose pricing is
|
||||
# not token-based and so cannot travel via custom_cost_per_token
|
||||
# (e.g. OCR per-page pricing). Same extraction as the video path
|
||||
# above, minus its `or {}` default: truthiness on the value adds
|
||||
# no mutable-collection construction (LIT002) and reads the same.
|
||||
# Checked under both keys: router calls that go through
|
||||
# `_ageneric_api_call_with_fallbacks` (OCR included) store the
|
||||
# deployment's model_info under `litellm_metadata`, not `metadata`.
|
||||
_custom_model_info: ModelInfo | None = None
|
||||
if custom_pricing and litellm_logging_obj is not None:
|
||||
_cm_litellm_params = getattr(litellm_logging_obj, "litellm_params", None)
|
||||
if _cm_litellm_params is not None:
|
||||
for _cm_metadata_key in ("metadata", "litellm_metadata"):
|
||||
_cm_metadata = _cm_litellm_params.get(_cm_metadata_key)
|
||||
if _cm_metadata and _cm_metadata.get("model_info") is not None:
|
||||
_custom_model_info = _cm_metadata.get("model_info")
|
||||
break
|
||||
|
||||
(
|
||||
prompt_tokens_cost_usd_dollar,
|
||||
completion_tokens_cost_usd_dollar,
|
||||
|
|
@ -1681,7 +1650,7 @@ def completion_cost(
|
|||
vertex_location=vertex_location,
|
||||
response=completion_response,
|
||||
request_model=request_model_for_cost,
|
||||
custom_model_info=_custom_model_info,
|
||||
custom_model_info=_deployment_model_info(litellm_logging_obj, custom_pricing),
|
||||
)
|
||||
|
||||
# Get additional costs from provider (e.g., routing fees, infrastructure costs)
|
||||
|
|
@ -1911,6 +1880,32 @@ def response_cost_calculator(
|
|||
raise e
|
||||
|
||||
|
||||
def _deployment_model_info(
|
||||
litellm_logging_obj: LitellmLoggingObject | None,
|
||||
custom_pricing: bool | None,
|
||||
) -> ModelInfo | None:
|
||||
if not custom_pricing or litellm_logging_obj is None:
|
||||
return None
|
||||
litellm_params: Final = getattr(litellm_logging_obj, "litellm_params", None)
|
||||
if litellm_params is None:
|
||||
return None
|
||||
return next(
|
||||
(
|
||||
model_info
|
||||
for metadata_key in ("metadata", "litellm_metadata")
|
||||
if (metadata := litellm_params.get(metadata_key)) and (model_info := metadata.get("model_info")) is not None
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def _cost_map_model_info(model: str, custom_llm_provider: str | None) -> ModelInfo | None:
|
||||
try:
|
||||
return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def ocr_cost(
|
||||
model: str,
|
||||
custom_llm_provider: str | None,
|
||||
|
|
@ -1922,9 +1917,8 @@ def ocr_cost(
|
|||
model: str - model name
|
||||
custom_llm_provider: Optional[str] - custom LLM provider
|
||||
response: Optional[Any] - response object
|
||||
model_info: Optional[ModelInfo] - deployment-specific model info, used for
|
||||
custom pricing. Takes precedence over the model cost map, mirroring
|
||||
the video generation cost path.
|
||||
model_info: Optional[ModelInfo] - deployment-specific model info; its OCR pricing
|
||||
takes precedence over the model cost map
|
||||
|
||||
Returns:
|
||||
Tuple[float, float]: cost of OCR processing
|
||||
|
|
@ -1942,34 +1936,18 @@ def ocr_cost(
|
|||
if response.usage_info is None:
|
||||
raise ValueError("OCR response usage_info is None")
|
||||
|
||||
#########################################################
|
||||
# Deployment-specific pricing wins over the cost map.
|
||||
#
|
||||
# Custom pricing set on a deployment is registered under the router's
|
||||
# deployment id, while the shared "{provider}/{model}" key has its pricing
|
||||
# fields stripped (see _register_custom_pricing_for_request). A cost map
|
||||
# lookup therefore cannot see it, so an OCR model that is not in the map
|
||||
# bills $0 no matter how it is priced in config. Prefer the caller-supplied
|
||||
# model_info when it carries OCR pricing.
|
||||
#########################################################
|
||||
has_custom_ocr_pricing: Final[bool] = model_info is not None and (
|
||||
has_custom_ocr_pricing: Final = model_info is not None and (
|
||||
model_info.get("ocr_cost_per_page") is not None or model_info.get("ocr_cost_per_credit") is not None
|
||||
)
|
||||
if not has_custom_ocr_pricing:
|
||||
try:
|
||||
model_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
|
||||
except Exception:
|
||||
model_info = None
|
||||
pricing: Final = model_info if has_custom_ocr_pricing else _cost_map_model_info(model, custom_llm_provider)
|
||||
|
||||
credits: Final = getattr(response.usage_info, "credits", None)
|
||||
cost_per_credit = None
|
||||
if model_info is not None:
|
||||
cost_per_credit = model_info.get("ocr_cost_per_credit")
|
||||
cost_per_credit: Final = pricing.get("ocr_cost_per_credit") if pricing is not None else None
|
||||
if credits is not None and cost_per_credit is not None:
|
||||
return cost_per_credit * credits, 0.0
|
||||
|
||||
ocr_cost_per_page: Final = model_info.get("ocr_cost_per_page") if model_info is not None else None
|
||||
annotation_cost_per_page: Final = model_info.get("annotation_cost_per_page") if model_info is not None else None
|
||||
ocr_cost_per_page: Final = pricing.get("ocr_cost_per_page") if pricing is not None else None
|
||||
annotation_cost_per_page: Final = pricing.get("annotation_cost_per_page") if pricing is not None else None
|
||||
annotation_rate: Final = annotation_cost_per_page if annotation_cost_per_page is not None else ocr_cost_per_page
|
||||
|
||||
pages_processed: Final = response.usage_info.pages_processed
|
||||
|
|
|
|||
|
|
@ -4645,3 +4645,133 @@ def test_collect_and_combine_realtime_usage_stores_partitioned_text_tokens() ->
|
|||
assert combined.completion_tokens_details.reasoning_tokens == 95
|
||||
assert combined.completion_tokens_details.text_tokens == 38
|
||||
assert combined.completion_tokens_details.audio_tokens == 0
|
||||
|
||||
|
||||
UNMAPPED_OCR_MODEL: Final = "azure_ai/some-unmapped-ocr-model-for-testing"
|
||||
MAPPED_OCR_MODEL: Final = "mistral/mistral-ocr-4-0"
|
||||
|
||||
|
||||
def _ocr_response(model: str, pages_processed: int, credits: float | None = None):
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo
|
||||
|
||||
return OCRResponse(
|
||||
pages=[OCRPage(index=index, markdown=f"page {index}") for index in range(pages_processed)],
|
||||
model=model,
|
||||
usage_info=OCRUsageInfo(pages_processed=pages_processed, credits=credits),
|
||||
)
|
||||
|
||||
|
||||
def _ocr_logging_obj(litellm_params: dict):
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
||||
logging_obj = Logging(
|
||||
model=UNMAPPED_OCR_MODEL,
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type="ocr",
|
||||
start_time=None,
|
||||
litellm_call_id="test-ocr-custom-pricing",
|
||||
function_id="1234",
|
||||
)
|
||||
logging_obj.litellm_params = litellm_params
|
||||
return logging_obj
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pages_processed", [1, 3, 10])
|
||||
def test_ocr_cost_uses_deployment_per_page_pricing_for_unmapped_model(pages_processed: int):
|
||||
from litellm.cost_calculator import ocr_cost
|
||||
|
||||
assert UNMAPPED_OCR_MODEL not in litellm.model_cost
|
||||
cost, _ = ocr_cost(
|
||||
model=UNMAPPED_OCR_MODEL,
|
||||
custom_llm_provider="azure_ai",
|
||||
response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=pages_processed),
|
||||
model_info={"ocr_cost_per_page": 0.004},
|
||||
)
|
||||
assert cost == pytest.approx(0.004 * pages_processed)
|
||||
|
||||
|
||||
def test_ocr_cost_uses_deployment_per_credit_pricing_for_unmapped_model():
|
||||
from litellm.cost_calculator import ocr_cost
|
||||
|
||||
cost, _ = ocr_cost(
|
||||
model=UNMAPPED_OCR_MODEL,
|
||||
custom_llm_provider="azure_ai",
|
||||
response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=2, credits=4),
|
||||
model_info={"ocr_cost_per_credit": 0.25},
|
||||
)
|
||||
assert cost == pytest.approx(0.25 * 4)
|
||||
|
||||
|
||||
def test_ocr_cost_unmapped_model_without_deployment_pricing_bills_zero():
|
||||
from litellm.cost_calculator import ocr_cost
|
||||
|
||||
cost, _ = ocr_cost(
|
||||
model=UNMAPPED_OCR_MODEL,
|
||||
custom_llm_provider="azure_ai",
|
||||
response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=5),
|
||||
model_info={"id": "some-deployment-id"},
|
||||
)
|
||||
assert cost == 0.0
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_local_model_cost_map")
|
||||
def test_ocr_cost_deployment_pricing_overrides_cost_map_for_mapped_model():
|
||||
from litellm.cost_calculator import ocr_cost
|
||||
|
||||
map_price: Final = litellm.get_model_info(MAPPED_OCR_MODEL)["ocr_cost_per_page"]
|
||||
assert map_price is not None
|
||||
override_price: Final = map_price * 10
|
||||
|
||||
cost, _ = ocr_cost(
|
||||
model=MAPPED_OCR_MODEL,
|
||||
custom_llm_provider="mistral",
|
||||
response=_ocr_response(MAPPED_OCR_MODEL, pages_processed=2),
|
||||
model_info={"ocr_cost_per_page": override_price},
|
||||
)
|
||||
assert cost == pytest.approx(override_price * 2)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_local_model_cost_map")
|
||||
def test_ocr_cost_falls_through_to_cost_map_when_deployment_has_no_ocr_pricing():
|
||||
from litellm.cost_calculator import ocr_cost
|
||||
|
||||
map_price: Final = litellm.get_model_info(MAPPED_OCR_MODEL)["ocr_cost_per_page"]
|
||||
assert map_price is not None
|
||||
|
||||
cost, _ = ocr_cost(
|
||||
model=MAPPED_OCR_MODEL,
|
||||
custom_llm_provider="mistral",
|
||||
response=_ocr_response(MAPPED_OCR_MODEL, pages_processed=2),
|
||||
model_info={"id": "some-deployment-id"},
|
||||
)
|
||||
assert cost == pytest.approx(map_price * 2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
|
||||
def test_completion_cost_ocr_reads_deployment_pricing_from_logging_metadata(metadata_key: str):
|
||||
logging_obj = _ocr_logging_obj({metadata_key: {"model_info": {"ocr_cost_per_page": 0.004}}})
|
||||
|
||||
cost = completion_cost(
|
||||
completion_response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=3),
|
||||
model=UNMAPPED_OCR_MODEL,
|
||||
custom_llm_provider="azure_ai",
|
||||
call_type="ocr",
|
||||
custom_pricing=True,
|
||||
litellm_logging_obj=logging_obj,
|
||||
)
|
||||
assert cost == pytest.approx(0.004 * 3)
|
||||
|
||||
|
||||
def test_completion_cost_ocr_ignores_deployment_pricing_without_custom_pricing_flag():
|
||||
logging_obj = _ocr_logging_obj({"metadata": {"model_info": {"ocr_cost_per_page": 0.004}}})
|
||||
|
||||
cost = completion_cost(
|
||||
completion_response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=3),
|
||||
model=UNMAPPED_OCR_MODEL,
|
||||
custom_llm_provider="azure_ai",
|
||||
call_type="ocr",
|
||||
custom_pricing=False,
|
||||
litellm_logging_obj=logging_obj,
|
||||
)
|
||||
assert cost == 0.0
|
||||
|
|
|
|||
|
|
@ -1,143 +0,0 @@
|
|||
"""
|
||||
Regression tests: OCR cost must honour deployment-specific custom pricing.
|
||||
|
||||
Before the fix, `ocr_cost()` resolved pricing exclusively through
|
||||
`litellm.get_model_info(model=..., custom_llm_provider=...)`, i.e. a cost map
|
||||
lookup keyed by model name. Custom pricing set on a deployment is registered
|
||||
under the router's deployment id, and the shared "{provider}/{model}" key has
|
||||
its pricing fields stripped, so the lookup could never see it. An OCR model
|
||||
absent from the cost map therefore billed $0 no matter how it was priced in
|
||||
config, even though `ocr_cost_per_page` / `ocr_cost_per_credit` are declared
|
||||
fields of `CustomPricingLiteLLMParams`.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.cost_calculator import completion_cost, ocr_cost
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo
|
||||
|
||||
# A model deliberately absent from the cost map.
|
||||
UNMAPPED_MODEL = "azure_ai/some-unmapped-ocr-model-for-testing"
|
||||
CUSTOM_COST_PER_PAGE = 0.004
|
||||
CUSTOM_COST_PER_CREDIT = 0.25
|
||||
|
||||
|
||||
def _ocr_response(model: str, pages_processed: int = 1, credits: int | None = None) -> OCRResponse:
|
||||
# NOTE: model_construct() is used rather than OCRResponse(...) because the
|
||||
# OCRResponse field `object: str = "ocr"` shadows the builtin `object` used
|
||||
# in the `tables` / `keyValuePairs` annotations above it, so pydantic tries
|
||||
# to resolve "ocr" as a forward-referenced type and schema building fails.
|
||||
# That is an unrelated defect; validation is not what these tests exercise.
|
||||
usage_info = OCRUsageInfo(pages_processed=pages_processed)
|
||||
if credits is not None:
|
||||
usage_info.credits = credits
|
||||
return OCRResponse.model_construct(
|
||||
pages=[OCRPage(index=i, markdown=f"page {i}") for i in range(pages_processed)],
|
||||
model=model,
|
||||
usage_info=usage_info,
|
||||
)
|
||||
|
||||
|
||||
def test_unmapped_ocr_model_has_no_map_pricing() -> None:
|
||||
"""Guard the premise: the model really is absent from the cost map."""
|
||||
assert UNMAPPED_MODEL not in litellm.model_cost
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pages_processed", [1, 3, 10])
|
||||
def test_ocr_cost_uses_custom_per_page_pricing(pages_processed: int) -> None:
|
||||
cost, _ = ocr_cost(
|
||||
model=UNMAPPED_MODEL,
|
||||
custom_llm_provider="azure_ai",
|
||||
response=_ocr_response(UNMAPPED_MODEL, pages_processed=pages_processed),
|
||||
model_info={"ocr_cost_per_page": CUSTOM_COST_PER_PAGE},
|
||||
)
|
||||
assert cost == pytest.approx(CUSTOM_COST_PER_PAGE * pages_processed)
|
||||
|
||||
|
||||
def test_ocr_cost_uses_custom_per_credit_pricing() -> None:
|
||||
cost, _ = ocr_cost(
|
||||
model=UNMAPPED_MODEL,
|
||||
custom_llm_provider="azure_ai",
|
||||
response=_ocr_response(UNMAPPED_MODEL, pages_processed=2, credits=4),
|
||||
model_info={"ocr_cost_per_credit": CUSTOM_COST_PER_CREDIT},
|
||||
)
|
||||
assert cost == pytest.approx(CUSTOM_COST_PER_CREDIT * 4)
|
||||
|
||||
|
||||
def test_unmapped_ocr_model_without_custom_pricing_still_bills_zero() -> None:
|
||||
"""Unchanged behaviour when nothing is configured — no map entry, no override."""
|
||||
cost, _ = ocr_cost(
|
||||
model=UNMAPPED_MODEL,
|
||||
custom_llm_provider="azure_ai",
|
||||
response=_ocr_response(UNMAPPED_MODEL, pages_processed=5),
|
||||
)
|
||||
assert cost == 0.0
|
||||
|
||||
|
||||
def test_custom_pricing_does_not_override_a_mapped_model_when_absent() -> None:
|
||||
"""model_info without OCR pricing must fall through to the cost map."""
|
||||
mapped_model = "mistral/mistral-ocr-4-0"
|
||||
cost, _ = ocr_cost(
|
||||
model=mapped_model,
|
||||
custom_llm_provider="mistral",
|
||||
response=_ocr_response(mapped_model, pages_processed=2),
|
||||
model_info={"id": "some-deployment-id"},
|
||||
)
|
||||
assert cost == pytest.approx(0.004 * 2)
|
||||
|
||||
|
||||
def test_ocr_custom_pricing_end_to_end_through_completion_cost() -> None:
|
||||
"""The whole path: litellm_params.metadata.model_info -> ocr_cost."""
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
|
||||
logging_obj = LiteLLMLogging(
|
||||
model=UNMAPPED_MODEL,
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type="ocr",
|
||||
start_time=None,
|
||||
litellm_call_id="test-ocr-custom-pricing",
|
||||
function_id="1234",
|
||||
)
|
||||
logging_obj.litellm_params = {"metadata": {"model_info": {"ocr_cost_per_page": CUSTOM_COST_PER_PAGE}}}
|
||||
|
||||
cost = completion_cost(
|
||||
completion_response=_ocr_response(UNMAPPED_MODEL, pages_processed=3),
|
||||
model=UNMAPPED_MODEL,
|
||||
custom_llm_provider="azure_ai",
|
||||
call_type="ocr",
|
||||
custom_pricing=True,
|
||||
litellm_logging_obj=logging_obj,
|
||||
)
|
||||
assert cost == pytest.approx(CUSTOM_COST_PER_PAGE * 3)
|
||||
|
||||
|
||||
def test_ocr_custom_pricing_end_to_end_via_litellm_metadata() -> None:
|
||||
"""Router OCR calls go through `_ageneric_api_call_with_fallbacks`, which
|
||||
stores the deployment's model_info under `litellm_metadata` rather than
|
||||
`metadata`. The extraction must read that key too."""
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
|
||||
logging_obj = LiteLLMLogging(
|
||||
model=UNMAPPED_MODEL,
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type="ocr",
|
||||
start_time=None,
|
||||
litellm_call_id="test-ocr-custom-pricing-litellm-metadata",
|
||||
function_id="1234",
|
||||
)
|
||||
logging_obj.litellm_params = {
|
||||
"litellm_metadata": {"model_info": {"ocr_cost_per_page": CUSTOM_COST_PER_PAGE}},
|
||||
}
|
||||
|
||||
cost = completion_cost(
|
||||
completion_response=_ocr_response(UNMAPPED_MODEL, pages_processed=3),
|
||||
model=UNMAPPED_MODEL,
|
||||
custom_llm_provider="azure_ai",
|
||||
call_type="ocr",
|
||||
custom_pricing=True,
|
||||
litellm_logging_obj=logging_obj,
|
||||
)
|
||||
assert cost == pytest.approx(CUSTOM_COST_PER_PAGE * 3)
|
||||
Loading…
Add table
Reference in a new issue