fix(cost): read OCR pricing registered under the router deployment id

This commit is contained in:
mateo-berri 2026-09-08 19:43:31 -07:00
parent 9a78bf638a
commit a29103cbbf
2 changed files with 34 additions and 3 deletions

View file

@ -1429,7 +1429,9 @@ def completion_cost(
)
elif call_type in _VIDEO_CALL_TYPES:
### VIDEO GENERATION COST CALCULATION ###
_video_model_info: ModelInfo | None = _deployment_model_info(litellm_logging_obj, custom_pricing)
_video_model_info: ModelInfo | None = _deployment_model_info(
litellm_logging_obj, custom_pricing, router_model_id
)
usage_obj = getattr(completion_response, "usage", None)
duration_seconds: float | None = None
@ -1650,7 +1652,7 @@ def completion_cost(
vertex_location=vertex_location,
response=completion_response,
request_model=request_model_for_cost,
custom_model_info=_deployment_model_info(litellm_logging_obj, custom_pricing),
custom_model_info=_deployment_model_info(litellm_logging_obj, custom_pricing, router_model_id),
)
# Get additional costs from provider (e.g., routing fees, infrastructure costs)
@ -1883,8 +1885,18 @@ def response_cost_calculator(
def _deployment_model_info(
litellm_logging_obj: LitellmLoggingObject | None,
custom_pricing: bool | None,
router_model_id: str | None,
) -> ModelInfo | None:
if not custom_pricing or litellm_logging_obj is None:
if not custom_pricing:
return None
registered_deployment_info: Final = (
_cost_map_model_info(router_model_id, None)
if router_model_id is not None and router_model_id in litellm.model_cost
else None
)
if registered_deployment_info is not None:
return registered_deployment_info
if litellm_logging_obj is None:
return None
litellm_params: Final = getattr(litellm_logging_obj, "litellm_params", None)
if litellm_params is None:

View file

@ -4761,6 +4761,25 @@ def test_completion_cost_ocr_reads_deployment_pricing_from_logging_metadata(meta
assert cost == pytest.approx(0.004 * 3)
def test_completion_cost_ocr_prefers_pricing_registered_under_router_model_id(monkeypatch: pytest.MonkeyPatch):
deployment_id: Final = "ocr-deployment-priced-through-litellm-params"
monkeypatch.setitem(
litellm.model_cost, deployment_id, {"mode": "ocr", "litellm_provider": "azure_ai", "ocr_cost_per_page": 0.05}
)
logging_obj = _ocr_logging_obj({"metadata": {"model_info": {"mode": "ocr"}}})
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,
router_model_id=deployment_id,
litellm_logging_obj=logging_obj,
)
assert cost == pytest.approx(0.05 * 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}}})