diff --git a/litellm/utils.py b/litellm/utils.py index d93c88e05a0..75de4f333ac 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4894,8 +4894,31 @@ def get_max_tokens(model: str) -> int | None: ) -def _strip_stable_vertex_version(model_name) -> str: - return re.sub(r"-\d+$", "", model_name) +def _strip_stable_vertex_version(model_name: str) -> str: + """Strip Vertex AI *publish* version suffixes, not product generation numbers. + + Vertex stable/versioned ids look like ``gemini-1.5-flash-001`` or + ``gemini-2.0-flash-001`` (3+ digit build ids). Anthropic-on-Vertex product + names use short generation numbers (``claude-sonnet-5``, ``claude-opus-4-1``) + and optional region resource names (``claude-sonnet-5@default``). + + The old ``-\\d+$`` pattern treated ``-5`` as a version and rewrote + ``claude-sonnet-5`` → ``claude-sonnet``, while leaving + ``claude-sonnet-5@default`` untouched. That made bare and ``@default`` + cost-map lookups diverge (issue #35758): bare ids missed priced entries + and logged $0, while ``@default`` still resolved. + + Only strip: + - 3+ digit trailing version codes (``-001``, ``-002``) + - 8-digit date suffixes (``-20250929``) + Leave region suffixes after ``@`` alone. + """ + if not model_name: + return model_name + if "@" in model_name: + # Region / resource name (e.g. model@default) is not a publish version. + return model_name + return re.sub(r"-(?:\d{8}|\d{3,})$", "", model_name) def _get_base_bedrock_model(model_name) -> str: @@ -5159,6 +5182,23 @@ def _get_model_info_from_generalization( return None +def _cost_map_provider_prefix(custom_llm_provider: str | None) -> str | None: + """Map sub-provider tags to the cost-map key prefix. + + Cost-map keys use ``vertex_ai/...`` even when ``litellm_provider`` is the + more specific ``vertex_ai-anthropic_models`` (or other ``vertex_ai-*`` + tags). Using the sub-provider as a key prefix produces miss keys like + ``vertex_ai-anthropic_models/claude-sonnet-5`` and falls through to $0. + """ + if custom_llm_provider is None: + return None + if custom_llm_provider.startswith("vertex_ai"): + return "vertex_ai" + if custom_llm_provider.startswith("fireworks_ai"): + return "fireworks_ai" + return custom_llm_provider + + def _get_potential_model_names(model: str, custom_llm_provider: str | None) -> PotentialModelNamesAndCustomLLMProvider: if custom_llm_provider is None: # Get custom_llm_provider @@ -5170,18 +5210,21 @@ def _get_potential_model_names(model: str, custom_llm_provider: str | None) -> P combined_model_name = model stripped_model_name = _strip_model_name(model=model, custom_llm_provider=custom_llm_provider) combined_stripped_model_name = stripped_model_name - elif custom_llm_provider and model.startswith( - custom_llm_provider + "/" - ): # handle case where custom_llm_provider is provided and model starts with custom_llm_provider - split_model = model.split("/", 1)[1] - combined_model_name = model - stripped_model_name = _strip_model_name(model=split_model, custom_llm_provider=custom_llm_provider) - combined_stripped_model_name = f"{custom_llm_provider}/{stripped_model_name}" else: - split_model = model - combined_model_name = f"{custom_llm_provider}/{model}" - stripped_model_name = _strip_model_name(model=model, custom_llm_provider=custom_llm_provider) - combined_stripped_model_name = f"{custom_llm_provider}/{stripped_model_name}" + cost_map_prefix: Final = _cost_map_provider_prefix(custom_llm_provider) or custom_llm_provider + # Prefer the cost-map prefix (e.g. vertex_ai/) over sub-provider tags when + # stripping a leading provider segment from the model string. + if model.startswith(custom_llm_provider + "/"): + split_model = model.split("/", 1)[1] + combined_model_name = f"{cost_map_prefix}/{split_model}" + elif cost_map_prefix != custom_llm_provider and model.startswith(cost_map_prefix + "/"): + split_model = model.split("/", 1)[1] + combined_model_name = model if model.startswith(cost_map_prefix + "/") else f"{cost_map_prefix}/{split_model}" + else: + split_model = model + combined_model_name = f"{cost_map_prefix}/{model}" + stripped_model_name = _strip_model_name(model=split_model, custom_llm_provider=cost_map_prefix) + combined_stripped_model_name = f"{cost_map_prefix}/{stripped_model_name}" if custom_llm_provider in ("bedrock", "bedrock_converse"): from litellm.llms.bedrock.common_utils import strip_bedrock_routing_prefix diff --git a/tests/test_litellm/test_vertex_claude_sonnet5_cost_parity.py b/tests/test_litellm/test_vertex_claude_sonnet5_cost_parity.py new file mode 100644 index 00000000000..6a3eb901745 --- /dev/null +++ b/tests/test_litellm/test_vertex_claude_sonnet5_cost_parity.py @@ -0,0 +1,149 @@ +""" +Regression for https://github.com/BerriAI/litellm/issues/35758 + +``vertex_ai/claude-sonnet-5`` must cost the same as ``vertex_ai/claude-sonnet-5@default``. + +Root cause: ``_strip_stable_vertex_version`` treated the product generation +suffix ``-5`` as a Vertex publish version and rewrote bare ids to +``claude-sonnet``, while ``@default`` resource names were left intact. +""" + +from __future__ import annotations + +import os + +import pytest + +import litellm +from litellm.cost_calculator import completion_cost, cost_per_token +from litellm.types.utils import Choices, Message, ModelResponse, Usage +from litellm.utils import ( + _get_potential_model_names, + _strip_model_name, + _strip_stable_vertex_version, + get_model_info, +) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Use the bundled cost map so tests do not depend on network-fetched main.""" + original = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original + litellm.get_model_info.cache_clear() + + +@pytest.mark.parametrize( + "model_name,expected", + [ + # Product generation numbers must be preserved + ("claude-sonnet-5", "claude-sonnet-5"), + ("claude-sonnet-5@default", "claude-sonnet-5@default"), + ("claude-sonnet-4-5", "claude-sonnet-4-5"), + ("claude-opus-4-1", "claude-opus-4-1"), + # Real Vertex publish / date version suffixes still strip + ("gemini-1.5-flash-001", "gemini-1.5-flash"), + ("gemini-2.0-flash-001", "gemini-2.0-flash"), + ("some-model-20250929", "some-model"), + ], +) +def test_strip_stable_vertex_version_preserves_generation_numbers(model_name, expected): + assert _strip_stable_vertex_version(model_name) == expected + assert _strip_model_name(model_name, "vertex_ai") == expected + + +def test_potential_names_for_vertex_anthropic_subprovider(local_model_cost_map): + """``vertex_ai-anthropic_models`` must still form ``vertex_ai/...`` cost keys.""" + names = _get_potential_model_names( + model="claude-sonnet-5", + custom_llm_provider="vertex_ai-anthropic_models", + ) + assert names["combined_model_name"] == "vertex_ai/claude-sonnet-5" + assert names["stripped_model_name"] == "claude-sonnet-5" + assert names["combined_stripped_model_name"] == "vertex_ai/claude-sonnet-5" + + +def test_get_model_info_parity_bare_vs_default(local_model_cost_map): + bare = get_model_info(model="vertex_ai/claude-sonnet-5") + default = get_model_info(model="vertex_ai/claude-sonnet-5@default") + assert bare["input_cost_per_token"] == default["input_cost_per_token"] + assert bare["output_cost_per_token"] == default["output_cost_per_token"] + assert bare["input_cost_per_token"] and bare["input_cost_per_token"] > 0 + # Provider-only bare id must resolve via vertex_ai cost-map prefix + via_sub = get_model_info( + model="claude-sonnet-5", + custom_llm_provider="vertex_ai-anthropic_models", + ) + assert via_sub["input_cost_per_token"] == bare["input_cost_per_token"] + assert via_sub["key"] == "vertex_ai/claude-sonnet-5" + + +def _usage() -> Usage: + return Usage( + prompt_tokens=1000, + completion_tokens=500, + total_tokens=1500, + cache_creation_input_tokens=100, + cache_read_input_tokens=200, + ) + + +@pytest.mark.parametrize( + "model", + [ + "vertex_ai/claude-sonnet-5", + "vertex_ai/claude-sonnet-5@default", + ], +) +def test_completion_cost_nonzero_for_vertex_sonnet5(model, local_model_cost_map): + resp = ModelResponse( + model=model, + choices=[Choices(message=Message(role="assistant", content="ok"), finish_reason="stop")], + usage=_usage(), + ) + cost = completion_cost( + completion_response=resp, + model=model, + custom_llm_provider="vertex_ai", + ) + assert cost > 0 + + +def test_completion_cost_parity_bare_vs_default(local_model_cost_map): + usage = _usage() + costs = [] + for model in ("vertex_ai/claude-sonnet-5", "vertex_ai/claude-sonnet-5@default"): + resp = ModelResponse( + model=model.split("/", 1)[1], + choices=[Choices(message=Message(role="assistant", content="ok"), finish_reason="stop")], + usage=usage, + ) + costs.append( + completion_cost( + completion_response=resp, + model=model, + custom_llm_provider="vertex_ai", + ) + ) + assert costs[0] == costs[1] + assert costs[0] > 0 + + +def test_cost_per_token_with_vertex_anthropic_provider_tag(local_model_cost_map): + """Provider tag from the cost map must not zero out bare Claude-on-Vertex ids.""" + prompt, completion = cost_per_token( + model="claude-sonnet-5", + prompt_tokens=1000, + completion_tokens=500, + cache_creation_input_tokens=100, + cache_read_input_tokens=200, + custom_llm_provider="vertex_ai-anthropic_models", + usage_object=_usage(), + ) + assert prompt + completion > 0