diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index b8ac75aec3a..d6f5bb1c5bf 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -55,6 +55,8 @@ _GLOB_WILDCARD: Final = "*" ModelAccessGroups: TypeAlias = Callable[[str], Collection[str]] +AccessGroupMembers: TypeAlias = Callable[[str], Collection[str]] + _BUDGET_START_TIME_KEY_PREFIXES: Final = MappingProxyType( { Litellm_EntityType.KEY: "virtual_key_budget_start_time", @@ -153,6 +155,16 @@ def proxy_model_access_groups(model: str) -> frozenset[str]: return frozenset(llm_router.get_model_access_groups(model_name=model)) +def proxy_access_group_members(access_group: str) -> frozenset[str]: + try: + from litellm.proxy.proxy_server import llm_router + except ImportError: + return frozenset() + if llm_router is None: + return frozenset() + return frozenset(llm_router.get_model_access_groups(model_access_group=access_group).get(access_group, ())) + + def resolve_model_budget( model: str, model_max_budget: Mapping[str, object], @@ -223,6 +235,47 @@ def _glob_precedence(pattern: str) -> tuple[tuple[int, int], str]: return PatternUtils.calculate_pattern_specificity(pattern), pattern +def entries_outranked_by( + entry: str, + model_max_budget: Collection[str], + access_group_members: AccessGroupMembers = proxy_access_group_members, +) -> Iterator[str]: + """Entries that `entry` can take precedence over for some request model. + + A pattern can match a provider-prefixed or wildcard-routed name whose + candidates or access groups include any other entry, so patterns are taken + to overlap: a pattern outranks every less specific pattern, and a model or + group entry outranks every pattern. + """ + if _GLOB_WILDCARD in entry: + entry_precedence: Final = _glob_precedence(entry) + yield from ( + other + for other in model_max_budget + if _GLOB_WILDCARD in other and _glob_precedence(other) < entry_precedence + ) + return + yield from (other for other in model_max_budget if _GLOB_WILDCARD in other) + entry_members: Final = frozenset(access_group_members(entry)) + yield from ( + other + for other in model_max_budget + if _GLOB_WILDCARD not in other + and _ranks_ahead_of_group( + entry=entry, entry_members=entry_members, group=other, group_members=access_group_members(other) + ) + ) + + +def _ranks_ahead_of_group(entry: str, entry_members: frozenset[str], group: str, group_members: Iterable[str]) -> bool: + return any( + _GLOB_WILDCARD in member + or entry in _budget_model_candidates(member) + or (member in entry_members and entry < group) + for member in group_members + ) + + def _usable_entry(entry: str, raw_budget_config: object) -> ResolvedModelBudget | None: if (budget_config := _usable_budget_config(raw_budget_config)) is None: # An entry that will not validate cannot be keyed, so it cannot be diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index d9763f72741..c1b81a345d0 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -137,7 +137,10 @@ from litellm.proxy.common_utils.user_api_key_cache import ( from litellm.proxy.db.db_span import db_span from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.hooks.model_max_budget_limiter import ( + AccessGroupMembers, build_model_max_budget_usage, + entries_outranked_by, + proxy_access_group_members, resolve_model_budget, ) from litellm.proxy.management.teams.authz import TEAM_OR_ORG_ADMIN, TeamRole, is_team_admin, team_access_denied @@ -1427,6 +1430,7 @@ def _check_team_model_budget_update_authority( data: UpdateTeamRequest, user_api_key_dict: UserAPIKeyAuth, existing_model_max_budget: Mapping[str, object] | None, + access_group_members: AccessGroupMembers = proxy_access_group_members, ) -> None: """Like `_check_team_budget_update_authority`: only a proxy admin may raise, re-window or drop a per-model cap.""" if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: @@ -1454,11 +1458,7 @@ def _check_team_model_budget_update_authority( cap = governing.budget_config if cap.max_budget is None: continue - if ( - proposed.max_budget is None - or proposed.max_budget > cap.max_budget - or proposed.budget_duration != cap.budget_duration - ): + if _looser_model_cap(proposed=proposed, cap=cap): raise HTTPException( status_code=403, detail={ @@ -1470,6 +1470,73 @@ def _check_team_model_budget_update_authority( ) }, ) + for model_name, proposed in requested.items(): + _check_outranked_model_caps( + entry=model_name, + proposed=proposed, + existing_model_max_budget=existing_model_max_budget, + access_group_members=access_group_members, + ) + + +def _check_outranked_model_caps( + entry: str, + proposed: BudgetConfig, + existing_model_max_budget: Mapping[str, object], + access_group_members: AccessGroupMembers, +) -> None: + if _keeps_or_tightens(proposed=proposed, existing=_existing_model_cap(existing_model_max_budget.get(entry))): + return + outranked: Final = next( + _looser_than_outranked_caps(entry, proposed, existing_model_max_budget, access_group_members), None + ) + if outranked is None: + return + outranked_name, cap = outranked + raise HTTPException( + status_code=403, + detail={ + "error": ( + f"Only a proxy admin can set a team's model_max_budget entry {entry!r} looser than " + f"{outranked_name!r}, which it can take precedence over. Current max_budget={cap.max_budget} per " + f"{cap.budget_duration} (entry {outranked_name!r}), requested={proposed.max_budget} per " + f"{proposed.budget_duration}." + ) + }, + ) + + +def _looser_than_outranked_caps( + entry: str, + proposed: BudgetConfig, + existing_model_max_budget: Mapping[str, object], + access_group_members: AccessGroupMembers, +) -> Iterator[tuple[str, BudgetConfig]]: + outranked_caps: Final = ( + (outranked, _existing_model_cap(existing_model_max_budget[outranked])) + for outranked in entries_outranked_by( + entry=entry, model_max_budget=existing_model_max_budget, access_group_members=access_group_members + ) + ) + yield from ( + (outranked, cap) + for outranked, cap in outranked_caps + if cap is not None and cap.max_budget is not None and _looser_model_cap(proposed=proposed, cap=cap) + ) + + +def _keeps_or_tightens(proposed: BudgetConfig, existing: BudgetConfig | None) -> bool: + return existing is not None and ( + existing.max_budget is None or not _looser_model_cap(proposed=proposed, cap=existing) + ) + + +def _looser_model_cap(proposed: BudgetConfig, cap: BudgetConfig) -> bool: + return ( + proposed.max_budget is None + or (cap.max_budget is not None and proposed.max_budget > cap.max_budget) + or proposed.budget_duration != cap.budget_duration + ) def _should_auto_add_team_creator( diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 37028558af3..2d728b506e9 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -17119,6 +17119,104 @@ def test_team_admin_can_tighten_or_keep_team_model_caps(requested) -> None: ) +_TEAM_OPUS_FAMILY_CAP: Final = MappingProxyType({"claude-opus*": {"max_budget": 10.0, "budget_duration": "1d"}}) +_TEAM_OPUS_GROUP_CAP: Final = MappingProxyType({"z-opus-group": {"max_budget": 10.0, "budget_duration": "1d"}}) +_TEAM_ACCESS_GROUP_MEMBERS: Final = MappingProxyType( + { + "z-opus-group": ("claude-opus-4-8",), + "a-opus-group": ("claude-opus-4-8", "claude-opus-5-5"), + "zz-opus-group": ("claude-opus-4-8",), + "a-gemini-group": ("gemini-3-pro",), + "z-anthropic-group": ("anthropic/*",), + } +) + + +def _team_access_group_members(access_group: str) -> tuple[str, ...]: + return _TEAM_ACCESS_GROUP_MEMBERS.get(access_group, ()) + + +def _check_team_admin_model_caps(existing: Mapping[str, object], requested: Mapping[str, object]) -> None: + from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority + + _check_team_model_budget_update_authority( + data=UpdateTeamRequest(team_id="t1", model_max_budget=dict(requested)), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin"), + existing_model_max_budget=existing, + access_group_members=_team_access_group_members, + ) + + +@pytest.mark.parametrize( + ("existing", "added"), + [ + (_TEAM_OPUS_FAMILY_CAP, {"opus-group": {"max_budget": 1000.0, "budget_duration": "1d"}}), + (_TEAM_OPUS_FAMILY_CAP, {"opus-group": {"max_budget": 10.0, "budget_duration": "7d"}}), + (_TEAM_OPUS_FAMILY_CAP, {"*laude-opus-4-8": {"max_budget": 1000.0, "budget_duration": "1d"}}), + (_TEAM_OPUS_GROUP_CAP, {"a-opus-group": {"max_budget": 1000.0, "budget_duration": "1d"}}), + (_TEAM_OPUS_GROUP_CAP, {"claude-opus-4-8": {"max_budget": 1000.0, "budget_duration": "1d"}}), + ( + {"z-anthropic-group": {"max_budget": 10.0, "budget_duration": "1d"}}, + {"claude-opus-4-8": {"max_budget": 1000.0, "budget_duration": "1d"}}, + ), + ], + ids=[ + "group_over_pattern", + "group_rewindows_pattern", + "more_specific_pattern", + "earlier_group_sharing_a_model", + "model_inside_group", + "model_under_wildcard_group", + ], +) +def test_team_admin_cannot_add_an_entry_looser_than_a_cap_it_outranks( + existing: Mapping[str, object], added: Mapping[str, object] +) -> None: + with pytest.raises(HTTPException) as exc: + _check_team_admin_model_caps(existing=existing, requested={**existing, **added}) + assert exc.value.status_code == 403 + assert "proxy admin" in exc.value.detail["error"] + assert repr(next(iter(existing))) in exc.value.detail["error"] + + +@pytest.mark.parametrize( + ("existing", "requested"), + [ + (_TEAM_OPUS_FAMILY_CAP, {**_TEAM_OPUS_FAMILY_CAP, "opus-group": {"max_budget": 5.0, "budget_duration": "1d"}}), + (_TEAM_OPUS_FAMILY_CAP, {**_TEAM_OPUS_FAMILY_CAP, "claude*": {"max_budget": 1000.0, "budget_duration": "1d"}}), + (_TEAM_OPUS_GROUP_CAP, {**_TEAM_OPUS_GROUP_CAP, "a-gemini-group": {"max_budget": 1000.0, "budget_duration": "1d"}}), + (_TEAM_OPUS_GROUP_CAP, {**_TEAM_OPUS_GROUP_CAP, "zz-opus-group": {"max_budget": 1000.0, "budget_duration": "1d"}}), + (_TEAM_OPUS_GROUP_CAP, {**_TEAM_OPUS_GROUP_CAP, "gpt-4o": {"max_budget": 1000.0, "budget_duration": "1d"}}), + ( + {**_TEAM_OPUS_FAMILY_CAP, "opus-group": {"max_budget": 100.0, "budget_duration": "1d"}}, + {**_TEAM_OPUS_FAMILY_CAP, "opus-group": {"max_budget": 100.0, "budget_duration": "1d"}}, + ), + ( + {**_TEAM_OPUS_FAMILY_CAP, "opus-group": {"max_budget": 100.0, "budget_duration": "1d"}}, + {**_TEAM_OPUS_FAMILY_CAP, "opus-group": {"max_budget": 50.0, "budget_duration": "1d"}}, + ), + ( + {**_TEAM_OPUS_FAMILY_CAP, "opus-group": {"rpm_limit": 10}}, + {**_TEAM_OPUS_FAMILY_CAP, "opus-group": {"rpm_limit": 20}}, + ), + ], + ids=[ + "tighter_group", + "less_specific_pattern", + "group_without_shared_models", + "later_group_sharing_a_model", + "model_outside_group", + "unchanged_looser_entry", + "tightened_looser_entry", + "uncapped_entry_edited", + ], +) +def test_team_admin_can_add_entries_that_never_lift_an_outranked_cap( + existing: Mapping[str, object], requested: Mapping[str, object] +) -> None: + assert _check_team_admin_model_caps(existing=existing, requested=requested) is None + + def test_team_model_cap_authority_skips_omitted_field_malformed_rows_and_proxy_admins() -> None: from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority