Merge pull request #35880 from BerriAI/devin_ai_fix_cost_estimate_onprem_provider_35210

fix(proxy): forward resolved provider and deployment pricing in /cost/estimate
This commit is contained in:
Mateo Wang 2026-08-12 00:08:40 -07:00 committed by GitHub
commit 9bfe593241
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 280 additions and 59 deletions

View file

@ -10,6 +10,8 @@ PATCH /config/cost_margin_config - Update cost margin configuration
POST /cost/estimate - Estimate cost for a given model and token counts
"""
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Final
from fastapi import APIRouter, Depends, HTTPException
@ -24,29 +26,65 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.types.utils import LlmProvidersSet
from litellm.types.utils import CostPerToken, LlmProvidersSet, ModelInfo
router: Final = APIRouter()
def _resolve_model_for_cost_lookup(model: str) -> tuple[str, str | None]:
@dataclass(frozen=True, slots=True)
class ResolvedCostModel:
model: str
provider: str | None
custom_cost_per_token: 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.
"""
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
return CostPerToken(
input_cost_per_token=input_price or 0.0,
output_cost_per_token=output_price or 0.0,
)
def _lookup_model_info(model: str) -> ModelInfo | None:
try:
return litellm.get_model_info(model=model)
except Exception:
return None
def _resolve_model_for_cost_lookup(model: str) -> ResolvedCostModel:
"""
Resolve a model name (which may be a router alias/model_group) to the
underlying litellm model name for cost lookup.
underlying litellm model name, provider, and any deployment-configured
pricing used for cost lookup.
Args:
model: The model name from the request (could be a router alias like 'e-model-router'
or an actual model name like 'azure_ai/gpt-4')
Returns:
Tuple of (resolved_model_name, custom_llm_provider)
- resolved_model_name: The actual model name to use for cost lookup
- custom_llm_provider: The provider if resolved from router, None otherwise
"""
from litellm.proxy.proxy_server import llm_router
custom_llm_provider: str | None = None
# Try to resolve from router if available
if llm_router is not None:
try:
@ -57,31 +95,25 @@ def _resolve_model_for_cost_lookup(model: str) -> tuple[str, str | None]:
first_deployment: Final = deployments[0]
litellm_params: Final = first_deployment.get("litellm_params", {})
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, 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")
if base_model:
verbose_proxy_logger.debug("Resolved model '%s' to base_model '%s' from router", model, base_model)
custom_llm_provider = litellm_params.get("custom_llm_provider")
return (
str(base_model),
(str(custom_llm_provider) if custom_llm_provider is not None else None),
)
return ResolvedCostModel(str(base_model), provider, custom_cost_per_token)
resolved_model: Final = litellm_params.get("model")
if resolved_model:
verbose_proxy_logger.debug("Resolved model '%s' to '%s' from router", model, resolved_model)
custom_llm_provider = litellm_params.get("custom_llm_provider")
return (
str(resolved_model),
(str(custom_llm_provider) if custom_llm_provider is not None else None),
)
return ResolvedCostModel(str(resolved_model), provider, custom_cost_per_token)
except Exception as e:
verbose_proxy_logger.debug("Could not resolve model '%s' from router: %s", model, e)
# Return original model if not resolved
return model, custom_llm_provider
return ResolvedCostModel(model, None, None)
def _calculate_period_costs(num_requests, cost_per_request, input_cost, output_cost, margin_cost):
@ -450,7 +482,9 @@ async def estimate_cost(
from litellm.types.utils import ModelResponse, Usage
# Resolve model name (handles router aliases like 'e-model-router' -> 'azure_ai/gpt-4')
resolved_model, resolved_provider = _resolve_model_for_cost_lookup(request.model)
resolved: Final = _resolve_model_for_cost_lookup(request.model)
resolved_model: Final = resolved.model
resolved_provider: Final = resolved.provider
verbose_proxy_logger.debug("Cost estimate: request.model='%s' resolved to '%s'", request.model, resolved_model)
@ -480,6 +514,8 @@ async def estimate_cost(
cost_per_request: Final = completion_cost(
completion_response=mock_response,
model=resolved_model,
custom_llm_provider=resolved_provider,
custom_cost_per_token=resolved.custom_cost_per_token,
litellm_logging_obj=litellm_logging_obj,
)
except Exception as e:
@ -497,20 +533,22 @@ async def estimate_cost(
output_cost: Final = cost_breakdown.get("output_cost", 0.0) if cost_breakdown else 0.0
margin_cost: Final = cost_breakdown.get("margin_total_amount", 0.0) if cost_breakdown else 0.0
# Get model info for per-token pricing display
try:
model_info: Final = litellm.get_model_info(model=resolved_model)
input_cost_per_token = model_info.get("input_cost_per_token")
output_cost_per_token = model_info.get("output_cost_per_token")
custom_llm_provider = model_info.get("litellm_provider")
except Exception:
input_cost_per_token = None
output_cost_per_token = None
custom_llm_provider = None
model_info: Final = _lookup_model_info(resolved_model)
mapped_input_price: Final = model_info.get("input_cost_per_token") if model_info is not None else None
mapped_output_price: Final = model_info.get("output_cost_per_token") if model_info is not None else None
mapped_provider: Final = model_info.get("litellm_provider") if model_info is not None else None
# Use provider from router resolution if not found in model_info
if custom_llm_provider is None and resolved_provider is not None:
custom_llm_provider = resolved_provider
input_cost_per_token: Final = (
resolved.custom_cost_per_token["input_cost_per_token"]
if resolved.custom_cost_per_token is not None
else mapped_input_price
)
output_cost_per_token: Final = (
resolved.custom_cost_per_token["output_cost_per_token"]
if resolved.custom_cost_per_token is not None
else mapped_output_price
)
custom_llm_provider: Final = mapped_provider if mapped_provider is not None else resolved_provider
# Calculate daily and monthly costs
(

View file

@ -322,9 +322,9 @@ class TestResolveModelForCostLookup:
"litellm.proxy.proxy_server.llm_router",
mock_router,
):
resolved_model, provider = _resolve_model_for_cost_lookup("gpt-5.3-codex")
resolved = _resolve_model_for_cost_lookup("gpt-5.3-codex")
assert resolved_model == "azure/gpt-4o"
assert resolved.model == "azure/gpt-4o"
mock_router.get_model_list.assert_called_once_with(model_name="gpt-5.3-codex")
def test_falls_back_to_litellm_params_model_when_no_base_model(self):
@ -352,9 +352,9 @@ class TestResolveModelForCostLookup:
"litellm.proxy.proxy_server.llm_router",
mock_router,
):
resolved_model, provider = _resolve_model_for_cost_lookup("gpt-4")
resolved = _resolve_model_for_cost_lookup("gpt-4")
assert resolved_model == "openai/gpt-4"
assert resolved.model == "openai/gpt-4"
def test_resolves_base_model_from_litellm_params(self):
"""
@ -383,9 +383,9 @@ class TestResolveModelForCostLookup:
"litellm.proxy.proxy_server.llm_router",
mock_router,
):
resolved_model, provider = _resolve_model_for_cost_lookup("my-azure-model")
resolved = _resolve_model_for_cost_lookup("my-azure-model")
assert resolved_model == "azure/gpt-4o-mini"
assert resolved.model == "azure/gpt-4o-mini"
def test_returns_original_model_when_no_router(self):
"""
@ -399,12 +399,10 @@ class TestResolveModelForCostLookup:
"litellm.proxy.proxy_server.llm_router",
None,
):
resolved_model, provider = _resolve_model_for_cost_lookup(
"azure/openai/gpt-5.3-codex"
)
resolved = _resolve_model_for_cost_lookup("azure/openai/gpt-5.3-codex")
assert resolved_model == "azure/openai/gpt-5.3-codex"
assert provider is None
assert resolved.model == "azure/openai/gpt-5.3-codex"
assert resolved.provider is None
def test_returns_custom_llm_provider_on_base_model_path(self):
"""base_model path: the custom_llm_provider from litellm_params is
@ -427,10 +425,10 @@ class TestResolveModelForCostLookup:
]
with patch("litellm.proxy.proxy_server.llm_router", mock_router):
resolved_model, provider = _resolve_model_for_cost_lookup("my-azure-model")
resolved = _resolve_model_for_cost_lookup("my-azure-model")
assert resolved_model == "azure/gpt-4o"
assert provider == "azure"
assert resolved.model == "azure/gpt-4o"
assert resolved.provider == "azure"
def test_returns_custom_llm_provider_on_resolved_model_path(self):
"""resolved-model path (no base_model): the custom_llm_provider from
@ -452,10 +450,10 @@ class TestResolveModelForCostLookup:
]
with patch("litellm.proxy.proxy_server.llm_router", mock_router):
resolved_model, provider = _resolve_model_for_cost_lookup("gpt-4")
resolved = _resolve_model_for_cost_lookup("gpt-4")
assert resolved_model == "openai/gpt-4"
assert provider == "openai"
assert resolved.model == "openai/gpt-4"
assert resolved.provider == "openai"
def test_resolves_base_model_when_deployment_has_no_litellm_params(self):
"""A deployment can omit litellm_params entirely; base_model from
@ -474,10 +472,10 @@ class TestResolveModelForCostLookup:
]
with patch("litellm.proxy.proxy_server.llm_router", mock_router):
resolved_model, provider = _resolve_model_for_cost_lookup("my-azure-model")
resolved = _resolve_model_for_cost_lookup("my-azure-model")
assert resolved_model == "azure/gpt-4o"
assert provider is None
assert resolved.model == "azure/gpt-4o"
assert resolved.provider is None
def test_resolves_model_when_deployment_has_no_model_info(self):
"""A deployment can omit model_info entirely; litellm_params.model must
@ -496,7 +494,192 @@ class TestResolveModelForCostLookup:
]
with patch("litellm.proxy.proxy_server.llm_router", mock_router):
resolved_model, provider = _resolve_model_for_cost_lookup("gpt-4")
resolved = _resolve_model_for_cost_lookup("gpt-4")
assert resolved_model == "openai/gpt-4"
assert provider is None
assert resolved.model == "openai/gpt-4"
assert resolved.provider is None
class TestEstimateCostOnPremProvider:
"""Regression tests for LIT-5210: /cost/estimate on on-prem deployment aliases."""
@pytest.mark.asyncio
async def test_estimate_cost_onprem_model_without_pricing(self):
"""
On-prem deployments (custom_llm_provider set, model absent from the cost map)
must not 500 with "LLM Provider NOT provided". The resolved provider has to be
forwarded to completion_cost so provider inference doesn't run on the bare model.
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": {},
}
]
saved_model_cost = dict(litellm.model_cost)
litellm.register_model(
{
"openai/zai-org/GLM-5.2": {
"input_cost_per_token": 0.0,
"output_cost_per_token": 0.0,
"litellm_provider": "openai",
"mode": "chat",
}
}
)
try:
with patch("litellm.proxy.proxy_server.llm_router", mock_router):
response = await estimate_cost(request=request, user_api_key_dict=MagicMock())
finally:
litellm.model_cost = saved_model_cost
assert response.model == "nvidia/zai-org/glm-5.2"
assert response.provider == "openai"
assert response.cost_per_request == 0.0
@pytest.mark.asyncio
async def test_estimate_cost_onprem_model_with_configured_pricing(self):
"""
On-prem deployments with input/output_cost_per_token configured must estimate a
real cost using that pricing, not fall back to 0.0.
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,
num_requests_per_day=100,
)
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": {},
}
]
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.002)
assert response.input_cost_per_request == pytest.approx(0.001)
assert response.output_cost_per_request == pytest.approx(0.001)
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)