From c33454fdec73751610b3d823bf2285360d3af43a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 13:22:45 -0700 Subject: [PATCH] fix(cost_calculator): keep custom-priced router ids when resolving slash aliases --- litellm/cost_calculator.py | 57 +++++++--------------- tests/test_litellm/test_cost_calculator.py | 48 ++++++++++++++++++ 2 files changed, 65 insertions(+), 40 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 98a872a1603..1e13cb9307e 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -799,50 +799,27 @@ 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_model = _strip_unregistered_leading_segments(return_model, region_name) 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 +def _strip_unregistered_leading_segments(model: str, region_name: str | None) -> str: + """Resolve a provider-prefixed slash alias like "vertex_ai/vertex/claude-opus-5" to the + registered cost key ("vertex_ai/claude-opus-5"), keeping the model unchanged when it already + resolves downstream (custom-priced router ids) or no stripped candidate is registered (#38069).""" + segments: Final = model.split("/") + if "/".join(segments[1:]) in litellm.model_cost: + return model + head_len: Final = 2 if region_name is not None and len(segments) > 2 and segments[1] == region_name else 1 + head: Final = "/".join(segments[:head_len]) + tail: Final = segments[head_len:] + strippable: Final = next( + (index for index, segment in enumerate(tail) if segment in LlmProvidersSet or segment == region_name), + len(tail), + ) + candidates: Final = tuple(f"{head}/{'/'.join(tail[start:])}" for start in range(min(strippable, len(tail) - 1) + 1)) + return next((candidate for candidate in candidates if candidate in litellm.model_cost), model) @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 14c587ccd14..9aa627b2b23 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -3866,3 +3866,51 @@ def test_select_model_name_unresolvable_alias_unchanged(_local_model_cost_map): ) assert selected == "vertex_ai/team/nonsense-model" + + +def test_completion_cost_keeps_custom_priced_slash_router_id(_local_model_cost_map): + """A custom-priced router id containing "/" keeps its custom pricing instead of being + rewritten to the built-in key its suffix happens to match.""" + + from litellm.cost_calculator import _select_model_name_for_cost_calc + + litellm.register_model( + model_cost={ + "vertex/claude-opus-5": { + "input_cost_per_token": 7e-6, + "output_cost_per_token": 8e-6, + "litellm_provider": "vertex_ai", + } + } + ) + + selected = _select_model_name_for_cost_calc( + model="vertex_ai/claude-opus-5", + completion_response=None, + custom_pricing=True, + custom_llm_provider="vertex_ai", + router_model_id="vertex/claude-opus-5", + ) + assert selected == "vertex_ai/vertex/claude-opus-5" + + 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", + custom_pricing=True, + router_model_id="vertex/claude-opus-5", + ) + assert cost == pytest.approx(100 * 7e-6 + 50 * 8e-6, rel=1e-9)