scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(proxy): honor model_info custom pricing in /cost/estimate

Browse source
This commit is contained in:
mateo-berri 2026-08-11 23:25:01 -07:00
parent 464a4cf207
commit 0fdbe03c50
2 changed files with 98 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
litellm/proxy/management_endpoints/cost_tracking_settings.py
View file

@ -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")

83
tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py
View file

@ -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)