From 0fdbe03c508106bb6ab48ace26ae0b57329a23eb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 11 Aug 2026 23:25:01 -0700 Subject: [PATCH] fix(proxy): honor model_info custom pricing in /cost/estimate --- .../cost_tracking_settings.py | 22 +++-- .../test_cost_tracking_settings.py | 83 +++++++++++++++++++ 2 files changed, 98 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/management_endpoints/cost_tracking_settings.py b/litellm/proxy/management_endpoints/cost_tracking_settings.py index 5d0feecfdf3..56439172b63 100644 --- a/litellm/proxy/management_endpoints/cost_tracking_settings.py +++ b/litellm/proxy/management_endpoints/cost_tracking_settings.py @@ -38,16 +38,24 @@ class ResolvedCostModel: custom_cost_per_token: CostPerToken | None -def _extract_custom_pricing(litellm_params: Mapping[str, object]) -> CostPerToken | None: +def _configured_price(key: str, sources: tuple[Mapping[str, object], ...]) -> float | None: + values: Final = (source.get(key) for source in sources) + numeric: Final = (float(value) for value in values if isinstance(value, (int, float))) + return next(numeric, None) + + +def _extract_custom_pricing( + litellm_params: Mapping[str, object], model_info: Mapping[str, object] +) -> CostPerToken | None: """ Pull per-token pricing configured on a deployment so on-prem / self-hosted models (absent from the public cost map) still estimate a real cost. + Pricing may live on ``litellm_params`` or ``model_info``; ``litellm_params`` + wins, matching the router's cost-map registration precedence. """ - input_cost: Final = litellm_params.get("input_cost_per_token") - output_cost: Final = litellm_params.get("output_cost_per_token") - - input_price: Final = float(input_cost) if isinstance(input_cost, (int, float)) else None - output_price: Final = float(output_cost) if isinstance(output_cost, (int, float)) else None + sources: Final = (litellm_params, model_info) + input_price: Final = _configured_price("input_cost_per_token", sources) + output_price: Final = _configured_price("output_cost_per_token", sources) if input_price is None and output_price is None: return None @@ -89,7 +97,7 @@ def _resolve_model_for_cost_lookup(model: str) -> ResolvedCostModel: model_info: Final = first_deployment.get("model_info", {}) custom_llm_provider: Final = litellm_params.get("custom_llm_provider") provider: Final = str(custom_llm_provider) if custom_llm_provider is not None else None - custom_cost_per_token: Final = _extract_custom_pricing(litellm_params) + custom_cost_per_token: Final = _extract_custom_pricing(litellm_params, model_info) # Check base_model first (needed for Azure custom deployment names) base_model: Final = model_info.get("base_model") or litellm_params.get("base_model") diff --git a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py index c62f0370d46..7e83180bfcd 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py @@ -600,3 +600,86 @@ class TestEstimateCostOnPremProvider: assert response.daily_cost == pytest.approx(0.2) assert response.input_cost_per_token == pytest.approx(0.000001) assert response.output_cost_per_token == pytest.approx(0.000002) + + @pytest.mark.asyncio + async def test_estimate_cost_onprem_model_with_model_info_pricing(self): + """ + Custom pricing configured under model_info (how DB / Admin UI added + deployments store it) must be honored, not just litellm_params pricing. + + completion_cost is intentionally NOT mocked. + """ + from litellm.proxy._types import CostEstimateRequest + from litellm.proxy.management_endpoints.cost_tracking_settings import ( + estimate_cost, + ) + + request = CostEstimateRequest( + model="nvidia/zai-org/glm-5.2", + input_tokens=1000, + output_tokens=500, + ) + + mock_router = MagicMock() + mock_router.get_model_list.return_value = [ + { + "model_name": "nvidia/zai-org/glm-5.2", + "litellm_params": { + "model": "zai-org/GLM-5.2", + "custom_llm_provider": "openai", + }, + "model_info": { + "input_cost_per_token": 0.000003, + "output_cost_per_token": 0.000004, + }, + } + ] + + with patch("litellm.proxy.proxy_server.llm_router", mock_router): + response = await estimate_cost(request=request, user_api_key_dict=MagicMock()) + + assert response.provider == "openai" + assert response.cost_per_request == pytest.approx(0.005) + assert response.input_cost_per_token == pytest.approx(0.000003) + assert response.output_cost_per_token == pytest.approx(0.000004) + + @pytest.mark.asyncio + async def test_estimate_cost_litellm_params_pricing_overrides_model_info(self): + """ + When pricing is set in both places, litellm_params wins, matching the + router's cost-map registration precedence. + """ + from litellm.proxy._types import CostEstimateRequest + from litellm.proxy.management_endpoints.cost_tracking_settings import ( + estimate_cost, + ) + + request = CostEstimateRequest( + model="nvidia/zai-org/glm-5.2", + input_tokens=1000, + output_tokens=500, + ) + + mock_router = MagicMock() + mock_router.get_model_list.return_value = [ + { + "model_name": "nvidia/zai-org/glm-5.2", + "litellm_params": { + "model": "zai-org/GLM-5.2", + "custom_llm_provider": "openai", + "input_cost_per_token": 0.000001, + "output_cost_per_token": 0.000002, + }, + "model_info": { + "input_cost_per_token": 0.000003, + "output_cost_per_token": 0.000004, + }, + } + ] + + with patch("litellm.proxy.proxy_server.llm_router", mock_router): + response = await estimate_cost(request=request, user_api_key_dict=MagicMock()) + + assert response.cost_per_request == pytest.approx(0.002) + assert response.input_cost_per_token == pytest.approx(0.000001) + assert response.output_cost_per_token == pytest.approx(0.000002)