fix(cost_calculator): resolve real cost key when model_name alias contains '/'

When a router-facing model_name alias contains a '/' whose leading segment
is not a registered provider (e.g. 'vertex/claude-opus-5' for deployment
'vertex_ai/claude-opus-5'), _select_model_name_for_cost_calc re-prefixed
it into a non-existent key ('vertex_ai/vertex/claude-opus-5'), so cost
lookup silently priced every streamed request at $0 - token counts were
recorded, no error raised, budgets never tripped.

After prefixing, walk the alias tail and return the first assembly that
exists in litellm.model_cost ('vertex_ai/claude-opus-5'). Provider/region
segments in the head are preserved, and an alias that resolves to no
known key keeps the previous behavior (no crash, legacy double-prefix).

Fixes #38069
This commit is contained in:
ksk2023 2026-08-26 22:06:06 +08:00
parent da91d4b6c9
commit a76e224a5f
2 changed files with 128 additions and 0 deletions

View file

@ -799,9 +799,52 @@ def _select_model_name_for_cost_calc(
else:
return_model = f"{custom_llm_provider}/{return_model}"
# A router-facing model_name alias may itself contain a "/" (e.g. "vertex/claude-opus-5")
# whose leading segment is NOT a registered provider, so the alias was re-prefixed above
# into a non-existent key like "vertex_ai/vertex/claude-opus-5". Strip unregistered
# leading segments until the assembled name is a known model_cost key (or nothing
# strippable remains), so cost lookup resolves to the real model instead of pricing
# the request at $0. See #38069.
stripped = _strip_unregistered_leading_segments(return_model, region_name)
if stripped is not None:
return_model = stripped
return return_model
def _strip_unregistered_leading_segments(model: str, region_name: str | None) -> str | None:
"""Find the real cost-map key hidden inside a re-prefixed alias.
After the provider (and optional region) was prepended to a router-facing alias that
itself contains "/" segments (e.g. "vertex_ai/vertex/claude-opus-5"), walk the
remaining segments: keep the provider/region head, then try successively shorter
tails "vertex_ai/vertex/claude-opus-5", "vertex_ai/claude-opus-5" returning the
first that exists in `litellm.model_cost`. Segments that are themselves providers or
regions are never dropped from the head, so "vertex_ai/us-east-1/model" keeps its region.
Returns None if no candidate resolves; the caller keeps its original assembly.
"""
segments = model.split("/")
# head = leading provider (+ optional region) that must be preserved
head_len = 1
if region_name is not None and len(segments) > 1 and segments[1] == region_name:
head_len = 2
head = segments[:head_len]
tail = segments[head_len:]
while tail:
candidate = "/".join(head + tail)
if candidate in litellm.model_cost:
return candidate
# only strip a leading tail segment that is NOT itself meaningful:
# a provider/region segment inside the tail (e.g. "azure_ai/o3") stays
if tail[0] in LlmProvidersSet or tail[0] == region_name:
return None
tail = tail[1:]
return None
@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
def _model_contains_known_llm_provider(model: str) -> bool:
"""

View file

@ -3781,3 +3781,88 @@ def test_completion_cost_prices_anthropic_shaped_cache_read_tokens(_local_model_
)
assert cost == pytest.approx(3 * 4e-6 + 4014 * 4e-7 + 5 * 2e-5, rel=1e-9)
def test_select_model_name_strips_unregistered_alias_prefix(_local_model_cost_map):
"""A router-facing model_name alias containing "/" whose leading segment is NOT a
registered provider must not be double-prefixed into a non-existent cost key.
Regression test for #38069: alias "vertex/claude-opus-5" (real deployment
"vertex_ai/claude-opus-5") was re-prefixed into "vertex_ai/vertex/claude-opus-5",
silently pricing every streamed request at $0.
"""
from litellm.cost_calculator import _select_model_name_for_cost_calc
response = litellm.ModelResponse(
id="x",
choices=[
{
"index": 0,
"message": {"role": "assistant", "content": "hi"},
"finish_reason": "stop",
}
],
model="vertex/claude-opus-5",
)
response._hidden_params = {}
selected = _select_model_name_for_cost_calc(
model=None,
completion_response=response,
custom_llm_provider="vertex_ai",
)
assert selected == "vertex_ai/claude-opus-5"
def test_completion_cost_nonzero_for_slash_alias_model_name(_local_model_cost_map):
"""End-to-end cost through a "/"-containing alias must price above zero (#38069)."""
response = litellm.ModelResponse(
id="x",
choices=[
{
"index": 0,
"message": {"role": "assistant", "content": "hi"},
"finish_reason": "stop",
}
],
model="vertex/claude-opus-5",
)
response._hidden_params = {"custom_llm_provider": "vertex_ai"}
response.usage = litellm.Usage(prompt_tokens=100, completion_tokens=50)
cost = litellm.completion_cost(
completion_response=response,
custom_llm_provider="vertex_ai",
)
assert cost == pytest.approx(100 * 5e-6 + 50 * 2.5e-5, rel=1e-9)
def test_select_model_name_unresolvable_alias_unchanged(_local_model_cost_map):
"""An alias that resolves to no known cost key keeps the legacy double-prefixed name."""
from litellm.cost_calculator import _select_model_name_for_cost_calc
response = litellm.ModelResponse(
id="x",
choices=[
{
"index": 0,
"message": {"role": "assistant", "content": "hi"},
"finish_reason": "stop",
}
],
model="team/nonsense-model",
)
response._hidden_params = {}
selected = _select_model_name_for_cost_calc(
model=None,
completion_response=response,
custom_llm_provider="vertex_ai",
)
assert selected == "vertex_ai/team/nonsense-model"