diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 254126c5847..9f6f7d7e95e 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -287,19 +287,30 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): def _get_matching_model_group_budget_configs( self, model: str, internal_model_max_budget: GenericBudgetConfigType ) -> tuple[tuple[str, BudgetConfig], ...]: - model_without_provider = self._get_model_without_custom_llm_provider(model) + request_variants = self._model_name_variants(model) return tuple( (_group_name, _config) for _group_name, _config in internal_model_max_budget.items() if _config.models and any( - _member == model - or _member == model_without_provider - or self._get_model_without_custom_llm_provider(_member) == model - for _member in _config.models + _member in request_variants or model in self._model_name_variants(_member) for _member in _config.models ) ) + def _model_name_variants(self, model: str) -> frozenset[str]: + """ + The name itself plus its provider-stripped forms: without the first + segment (`huggingface/meta-llama/Llama-3.1-8B` -> `meta-llama/Llama-3.1-8B`) + and without everything before the last slash (`openai/gpt-4` -> `gpt-4`). + + Group matching compares one side's raw name against the other side's + variants, never stripped-vs-stripped, so a provider-prefixed member + still pins its group to that provider's route. + """ + if "/" not in model: + return frozenset({model}) + return frozenset({model, model.split("/", 1)[1], self._get_model_without_custom_llm_provider(model)}) + def _get_model_without_custom_llm_provider(self, model: str) -> str: if "/" in model: return model.split("/")[-1] diff --git a/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py b/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py index a7aadf1192f..d599852fb9b 100644 --- a/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py @@ -253,3 +253,40 @@ async def test_end_user_group_budget_within_budget_passes(): assert await limiter.is_end_user_within_model_budget(END_USER_ID, OPUS_GROUP_BUDGET, "anthropic-opus-4-7") is True assert await limiter.is_end_user_within_model_budget(END_USER_ID, OPUS_GROUP_BUDGET, "anthropic-opus-4-8") is True + + +@pytest.mark.asyncio +async def test_group_budget_matches_namespaced_model_with_provider_prefix(): + namespaced_group_budget = { + "llama-family": { + "models": ["meta-llama/Llama-3.1-8B", "meta-llama/Llama-3.1-70B"], + "budget_limit": 10.0, + "time_period": "30d", + } + } + limiter = _make_limiter() + key = _make_key(namespaced_group_budget) + + await _log_spend(limiter, "huggingface/meta-llama/Llama-3.1-8B", 11.0, namespaced_group_budget) + + with pytest.raises(litellm.BudgetExceededError, match="model group=llama-family"): + await limiter.is_key_within_model_budget(key, "huggingface/meta-llama/Llama-3.1-70B") + with pytest.raises(litellm.BudgetExceededError, match="model group=llama-family"): + await limiter.is_key_within_model_budget(key, "meta-llama/Llama-3.1-8B") + + +@pytest.mark.asyncio +async def test_group_member_with_provider_prefix_does_not_match_other_provider(): + pinned_group_budget = { + "openai-gpt4": { + "models": ["openai/gpt-4"], + "budget_limit": 10.0, + "time_period": "30d", + } + } + limiter = _make_limiter() + key = _make_key(pinned_group_budget) + + await _log_spend(limiter, "gpt-4", 11.0, pinned_group_budget) + + assert await limiter.is_key_within_model_budget(key, "azure/gpt-4") is True