diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 644253ceac7..ca0fb91118e 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -901,6 +901,40 @@ def _get_deployment_default_tpm_limit(model_name: str) -> Optional[int]: return _get_deployment_default_limit(model_name, "default_api_key_tpm_limit") +def get_key_own_model_rate_limits( + user_api_key_dict: UserAPIKeyAuth, + rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"], +) -> dict[str, int] | None: + """ + Get the per-model limits configured on the key itself, ignoring anything it + only inherits (team metadata, deployment defaults). + + Priority order (returns first found): + 1. Key metadata (model_rpm_limit / model_tpm_limit) + 2. Key model_max_budget (rpm_limit / tpm_limit per model) + + Callers that must know whether the key *overrides* an inherited limit use + this; callers that want the effective limit use get_key_model_rpm_limit / + get_key_model_tpm_limit, which continue the chain past the key. + """ + if user_api_key_dict.metadata: + metadata_limits = user_api_key_dict.metadata.get(rate_limit_key) + if metadata_limits: + return metadata_limits + + if user_api_key_dict.model_max_budget: + budget_field = "rpm_limit" if rate_limit_key == "model_rpm_limit" else "tpm_limit" + budget_limits = { + model: budget[budget_field] + for model, budget in user_api_key_dict.model_max_budget.items() + if isinstance(budget, dict) and budget.get(budget_field) is not None + } + if budget_limits: + return budget_limits + + return None + + def get_key_model_rpm_limit( user_api_key_dict: UserAPIKeyAuth, model_name: Optional[str] = None, @@ -914,20 +948,9 @@ def get_key_model_rpm_limit( 3. Team metadata (model_rpm_limit) 4. Deployment default_api_key_rpm_limit (when model_name is provided) """ - # 1. Check key metadata first (takes priority) - if user_api_key_dict.metadata: - result = user_api_key_dict.metadata.get("model_rpm_limit") - if result: - return result - - # 2. Check model_max_budget - if user_api_key_dict.model_max_budget: - model_rpm_limit: Dict[str, Any] = {} - for model, budget in user_api_key_dict.model_max_budget.items(): - if isinstance(budget, dict) and budget.get("rpm_limit") is not None: - model_rpm_limit[model] = budget["rpm_limit"] - if model_rpm_limit: - return model_rpm_limit + key_own_limits = get_key_own_model_rate_limits(user_api_key_dict, "model_rpm_limit") + if key_own_limits is not None: + return key_own_limits # 3. Fallback to team metadata if user_api_key_dict.team_metadata: @@ -957,20 +980,9 @@ def get_key_model_tpm_limit( 3. Team metadata (model_tpm_limit) 4. Deployment default_api_key_tpm_limit (when model_name is provided) """ - # 1. Check key metadata first (takes priority) - if user_api_key_dict.metadata: - result = user_api_key_dict.metadata.get("model_tpm_limit") - if result: - return result - - # 2. Check model_max_budget (iterate per-model like RPM does) - if user_api_key_dict.model_max_budget: - model_tpm_limit: Dict[str, Any] = {} - for model, budget in user_api_key_dict.model_max_budget.items(): - if isinstance(budget, dict) and budget.get("tpm_limit") is not None: - model_tpm_limit[model] = budget["tpm_limit"] - if model_tpm_limit: - return model_tpm_limit + key_own_limits = get_key_own_model_rate_limits(user_api_key_dict, "model_tpm_limit") + if key_own_limits is not None: + return key_own_limits # 3. Fallback to team metadata if user_api_key_dict.team_metadata: diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 3477af36285..f5c1609bcf7 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -209,13 +209,25 @@ class _PROXY_BatchRateLimiter(CustomLogger): user_api_key_dict: UserAPIKeyAuth, data: Dict, ) -> List["RateLimitDescriptor"]: - return self.parallel_request_limiter._create_rate_limit_descriptors( + """Batch descriptors, mirroring what ``async_pre_call_hook`` enforces for a chat request. + + Team per-model limits live in their own helper (the same one the chat + path calls), so they have to be added explicitly here too. + """ + descriptors = self.parallel_request_limiter._create_rate_limit_descriptors( user_api_key_dict=user_api_key_dict, data=data, rpm_limit_type=None, tpm_limit_type=None, model_has_failures=False, ) + requested_model = data.get("model") + self.parallel_request_limiter._add_team_model_rate_limit_descriptor_from_metadata( + user_api_key_dict=user_api_key_dict, + requested_model=requested_model if isinstance(requested_model, str) else None, + descriptors=descriptors, + ) + return descriptors def _should_skip_batch_input_file_processing( self, diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 3e7edd8cd06..8be2ed10e9c 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -35,6 +35,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( + get_key_own_model_rate_limits, get_key_tag_rpm_limit, get_model_rate_limit_from_metadata, ) @@ -376,6 +377,7 @@ class RequestRateLimiterStash: reserved_tokens: int = 0 reserved_model: Optional[str] = None reserved_scopes: FrozenSet[Tuple[str, str]] = field(default_factory=frozenset) + suppressed_tpm_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset) reservation_released: bool = False @@ -2031,11 +2033,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Returns list of descriptors for API key, user, team, team member, end user, model-specific, agent, and agent-session limits. """ - from litellm.proxy.auth.auth_utils import ( - get_team_model_rpm_limit, - get_team_model_tpm_limit, - ) - descriptors = [] # API Key rate limits @@ -2161,37 +2158,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): descriptors=descriptors, ) - if ( - get_team_model_rpm_limit(user_api_key_dict) is not None - or get_team_model_tpm_limit(user_api_key_dict) is not None - ): - _tpm_limit_for_team_model = get_team_model_tpm_limit(user_api_key_dict) or {} - _rpm_limit_for_team_model = get_team_model_rpm_limit(user_api_key_dict) or {} - should_check_rate_limit = False - if requested_model in _tpm_limit_for_team_model: - should_check_rate_limit = True - elif requested_model in _rpm_limit_for_team_model: - should_check_rate_limit = True - - if should_check_rate_limit: - model_specific_tpm_limit = None - model_specific_rpm_limit = None - if requested_model in _tpm_limit_for_team_model: - model_specific_tpm_limit = _tpm_limit_for_team_model[requested_model] - if requested_model in _rpm_limit_for_team_model: - model_specific_rpm_limit = _rpm_limit_for_team_model[requested_model] - descriptors.append( - RateLimitDescriptor( - key="model_per_team", - value=f"{user_api_key_dict.team_id}:{requested_model}", - rate_limit={ - "requests_per_unit": model_specific_rpm_limit, - "tokens_per_unit": model_specific_tpm_limit, - "window_size": self.window_size, - }, - ) - ) - # Agent-level and session-level rate limits resolved_agent_id = self._get_resolved_agent_id(user_api_key_dict, data) @@ -2266,6 +2232,58 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return batch_limiter return None + def _team_model_limit_for_metric( + self, + user_api_key_dict: UserAPIKeyAuth, + requested_model: str, + rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"], + ) -> int | None: + """ + Team per-model limit that still applies to this key for one metric. + + A key that sets its own limit for this model (key metadata or + model_max_budget) already gets it enforced through the model_per_key + descriptor, and documented precedence is key > team. Enforcing the team + counter too would clamp the key back to the stricter of the two, so the + team limit is dropped for that metric only - a key overriding RPM alone + still inherits the team's TPM ceiling. + """ + team_limits = get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", rate_limit_key) or {} + team_limit = team_limits.get(requested_model) + if team_limit is None: + return None + + key_own_limits = get_key_own_model_rate_limits(user_api_key_dict, rate_limit_key) or {} + if key_own_limits.get(requested_model) is not None: + return None + + return team_limit + + def _suppressed_team_model_tpm_scopes( + self, + user_api_key_dict: UserAPIKeyAuth, + requested_model: str | None, + ) -> frozenset[tuple[str, str]]: + """ + Scopes whose team per-model TPM ceiling this key's own override replaced. + + Post-call reconciliation charges every candidate scope regardless of what + was enforced pre-call, so without this an overriding key keeps filling the + team counter it is no longer gated by and starves the keys that do inherit + the team limit. RPM needs no equivalent: its counters are incremented only + for descriptors that were actually emitted. + """ + if requested_model is None or not user_api_key_dict.team_id: + return frozenset() + + team_limits = get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_tpm_limit") or {} + if team_limits.get(requested_model) is None: + return frozenset() + if self._team_model_limit_for_metric(user_api_key_dict, requested_model, "model_tpm_limit") is not None: + return frozenset() + + return frozenset({("model_per_team", f"{user_api_key_dict.team_id}:{requested_model}")}) + def _add_team_model_rate_limit_descriptor_from_metadata( self, user_api_key_dict: UserAPIKeyAuth, @@ -2273,34 +2291,29 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): descriptors: List[RateLimitDescriptor], ) -> None: """Add team model rate limit descriptor from team_metadata if applicable.""" - if ( - get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_rpm_limit") is not None - or get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_tpm_limit") is not None - ): - _tpm_limit_for_team_model = ( - get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_tpm_limit") or {} - ) - _rpm_limit_for_team_model = ( - get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_rpm_limit") or {} - ) - should_check_rate_limit = ( - requested_model in _tpm_limit_for_team_model or requested_model in _rpm_limit_for_team_model - ) + if requested_model is None: + return - if should_check_rate_limit and requested_model is not None: - model_specific_tpm_limit = _tpm_limit_for_team_model.get(requested_model) - model_specific_rpm_limit = _rpm_limit_for_team_model.get(requested_model) - descriptors.append( - RateLimitDescriptor( - key="model_per_team", - value=f"{user_api_key_dict.team_id}:{requested_model}", - rate_limit={ - "requests_per_unit": model_specific_rpm_limit, - "tokens_per_unit": model_specific_tpm_limit, - "window_size": self.window_size, - }, - ) - ) + model_specific_rpm_limit = self._team_model_limit_for_metric( + user_api_key_dict, requested_model, "model_rpm_limit" + ) + model_specific_tpm_limit = self._team_model_limit_for_metric( + user_api_key_dict, requested_model, "model_tpm_limit" + ) + if model_specific_rpm_limit is None and model_specific_tpm_limit is None: + return + + descriptors.append( + RateLimitDescriptor( + key="model_per_team", + value=f"{user_api_key_dict.team_id}:{requested_model}", + rate_limit={ + "requests_per_unit": model_specific_rpm_limit, + "tokens_per_unit": model_specific_tpm_limit, + "window_size": self.window_size, + }, + ) + ) def _add_project_model_rate_limit_descriptor_from_metadata( self, @@ -2447,6 +2460,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): requested_model=requested_model, descriptors=descriptors, ) + stash.suppressed_tpm_scopes = self._suppressed_team_model_tpm_scopes( + user_api_key_dict=user_api_key_dict, + requested_model=requested_model, + ) # Project Level Rate Limits self._add_project_model_rate_limit_descriptor_from_metadata( @@ -2941,6 +2958,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): reserved_tokens = stash.reserved_tokens if stash is not None else 0 reserved_model = stash.reserved_model if stash is not None else None reserved_scopes: FrozenSet[Tuple[str, str]] = stash.reserved_scopes if stash is not None else frozenset() + suppressed_scopes: frozenset[tuple[str, str]] = ( + stash.suppressed_tpm_scopes if stash is not None else frozenset() + ) # Reconciliation must target the same model-scoped counter that the # pre-call reservation incremented. If a reservation was made, # ``reserved_model`` is authoritative; otherwise fall back to the @@ -2960,11 +2980,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # is empty, so every scope falls through the unreserved branch and # gets the full actual charge — matching pre-PR behavior. # ---------------------------------------------------------------- - targets = self._collect_tpm_scope_targets( - standard_logging_metadata=standard_logging_metadata, - kwargs=kwargs, - model_group=reconcile_model, - ) + targets = [ + target + for target in self._collect_tpm_scope_targets( + standard_logging_metadata=standard_logging_metadata, + kwargs=kwargs, + model_group=reconcile_model, + ) + if target not in suppressed_scopes + ] if reserved_tokens > 0 and total_tokens < reserved_tokens: verbose_proxy_logger.debug( f"Releasing unused TPM budget on success: " diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 1610d76efb7..183d1b7acda 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -20,6 +20,7 @@ from litellm.proxy.auth.auth_utils import ( get_key_mcp_rpm_limit, get_key_model_rpm_limit, get_key_model_tpm_limit, + get_key_own_model_rate_limits, get_key_tag_rpm_limit, get_model_from_request, get_project_model_rpm_limit, @@ -169,6 +170,55 @@ class TestGetKeyModelRpmLimit: assert result == {} +class TestGetKeyOwnModelRateLimits: + """Only limits set on the key itself, so callers can tell an override from an inherited limit.""" + + def test_metadata_limits_are_the_keys_own(self): + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-123", + metadata={"model_rpm_limit": {"gpt-4": 100}}, + team_metadata={"model_rpm_limit": {"gpt-4": 5}}, + ) + assert get_key_own_model_rate_limits(user_api_key_dict, "model_rpm_limit") == { + "gpt-4": 100 + } + + def test_model_max_budget_limits_are_the_keys_own(self): + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-123", + model_max_budget={ + "gpt-4": {"rpm_limit": 100, "tpm_limit": 2000}, + "gpt-3.5-turbo": {"tpm_limit": 3000}, + }, + ) + assert get_key_own_model_rate_limits(user_api_key_dict, "model_rpm_limit") == { + "gpt-4": 100 + } + assert get_key_own_model_rate_limits(user_api_key_dict, "model_tpm_limit") == { + "gpt-4": 2000, + "gpt-3.5-turbo": 3000, + } + + def test_team_metadata_is_not_the_keys_own(self): + team_only = UserAPIKeyAuth( + api_key="sk-123", + team_metadata={ + "model_rpm_limit": {"gpt-4": 5}, + "model_tpm_limit": {"gpt-4": 500}, + }, + ) + assert get_key_own_model_rate_limits(team_only, "model_rpm_limit") is None + assert get_key_own_model_rate_limits(team_only, "model_tpm_limit") is None + + def test_returns_none_when_nothing_configured_on_the_key(self): + assert ( + get_key_own_model_rate_limits( + UserAPIKeyAuth(api_key="sk-123"), "model_rpm_limit" + ) + is None + ) + + class TestGetKeyMcpRpmLimit: def test_empty_dict_limits_are_returned(self): key_override = UserAPIKeyAuth( diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 1c29c287c3a..9c79ddf350a 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -5100,3 +5100,281 @@ async def test_reserve_tpm_tokens_never_evaluates_the_requests_dimension(): f"reservation pass, got: {response}" ) assert [s["rate_limit_type"] for s in response["statuses"]] == ["tokens"] + + +def _make_team_model_limit_handler(): + """Limiter backed only by in-memory cache, plus the descriptors it enforced.""" + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + captured: List[Dict[str, Any]] = [] + + original_should_rate_limit = handler.should_rate_limit + + async def capturing_should_rate_limit(descriptors, **kwargs): + captured.extend(descriptors) + return await original_should_rate_limit(descriptors, **kwargs) + + handler.should_rate_limit = capturing_should_rate_limit + return handler, captured + + +async def _run_pre_call_hook(handler, user_api_key_dict, model="gpt-4"): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={"model": model, "messages": [{"role": "user", "content": "hi"}]}, + call_type="", + ) + + +def _descriptors_with_key(captured, key): + return [d for d in captured if d["key"] == key] + + +@pytest.mark.asyncio +async def test_key_model_rpm_override_wins_over_team_model_rpm(): + """Key metadata model_rpm_limit beats the team's for the same model, so the + team counter must not also cap the key at the (lower) team value.""" + handler, captured = _make_team_model_limit_handler() + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-key-rpm-override"), + team_id="team-precedence", + metadata={"model_rpm_limit": {"gpt-4": 100}}, + team_metadata={"model_rpm_limit": {"gpt-4": 3}}, + ) + + await _run_pre_call_hook(handler, user_api_key_dict) + + key_descriptors = _descriptors_with_key(captured, "model_per_key") + assert len(key_descriptors) == 1 + assert key_descriptors[0]["rate_limit"]["requests_per_unit"] == 100 + + team_descriptors = _descriptors_with_key(captured, "model_per_team") + assert ( + team_descriptors == [] + ), f"the team's only limit is suppressed by the key override, so no team descriptor: {team_descriptors}" + + +@pytest.mark.asyncio +async def test_key_rpm_override_still_inherits_team_model_tpm(): + """Overriding RPM on the key must not release the team's TPM ceiling.""" + handler, captured = _make_team_model_limit_handler() + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-key-rpm-only"), + team_id="team-precedence", + metadata={"model_rpm_limit": {"gpt-4": 100}}, + team_metadata={ + "model_rpm_limit": {"gpt-4": 3}, + "model_tpm_limit": {"gpt-4": 500}, + }, + ) + + await _run_pre_call_hook(handler, user_api_key_dict) + + team_descriptors = _descriptors_with_key(captured, "model_per_team") + assert len(team_descriptors) == 1, f"expected one team descriptor, got: {team_descriptors}" + assert team_descriptors[0]["rate_limit"]["requests_per_unit"] is None + assert team_descriptors[0]["rate_limit"]["tokens_per_unit"] == 500 + + +@pytest.mark.asyncio +async def test_key_model_max_budget_rpm_override_wins_over_team_model_rpm(): + """The override may be expressed as model_max_budget[model].rpm_limit.""" + handler, captured = _make_team_model_limit_handler() + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-key-budget-override"), + team_id="team-precedence", + model_max_budget={"gpt-4": {"rpm_limit": 100}}, + team_metadata={"model_rpm_limit": {"gpt-4": 3}}, + ) + + await _run_pre_call_hook(handler, user_api_key_dict) + + key_descriptors = _descriptors_with_key(captured, "model_per_key") + assert len(key_descriptors) == 1 + assert key_descriptors[0]["rate_limit"]["requests_per_unit"] == 100 + + team_descriptors = _descriptors_with_key(captured, "model_per_team") + assert ( + team_descriptors == [] + ), f"the team's only limit is suppressed by the key override, so no team descriptor: {team_descriptors}" + + +@pytest.mark.asyncio +async def test_key_override_on_one_model_leaves_team_limit_on_other_model(): + """Suppression is per model: an override on gpt-4 must not free claude-3.""" + handler, captured = _make_team_model_limit_handler() + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-key-other-model"), + team_id="team-precedence", + metadata={"model_rpm_limit": {"gpt-4": 100}}, + team_metadata={"model_rpm_limit": {"gpt-4": 3, "claude-3": 7}}, + ) + + await _run_pre_call_hook(handler, user_api_key_dict, model="claude-3") + + team_descriptors = _descriptors_with_key(captured, "model_per_team") + assert len(team_descriptors) == 1, f"expected one team descriptor, got: {team_descriptors}" + assert team_descriptors[0]["value"] == "team-precedence:claude-3" + assert team_descriptors[0]["rate_limit"]["requests_per_unit"] == 7 + + +@pytest.mark.asyncio +async def test_team_model_descriptor_is_emitted_exactly_once(): + """Two code paths used to emit the same model_per_team descriptor, so every + request burned two slots of the team's per-model budget.""" + handler, captured = _make_team_model_limit_handler() + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-team-only"), + team_id="team-precedence", + team_metadata={"model_rpm_limit": {"gpt-4": 3}}, + ) + + await _run_pre_call_hook(handler, user_api_key_dict) + + team_descriptors = _descriptors_with_key(captured, "model_per_team") + assert len(team_descriptors) == 1, f"expected one team descriptor, got: {team_descriptors}" + assert team_descriptors[0]["rate_limit"]["requests_per_unit"] == 3 + + +@pytest.mark.asyncio +async def test_team_model_rpm_limit_admits_exactly_the_configured_requests(): + """A team model_rpm_limit of N admits N requests, not N/2.""" + handler, _ = _make_team_model_limit_handler() + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-team-budget-burn"), + team_id="team-burn", + team_metadata={"model_rpm_limit": {"gpt-4": 4}}, + ) + + for _ in range(4): + await _run_pre_call_hook(handler, user_api_key_dict) + + with pytest.raises(HTTPException) as exc_info: + await _run_pre_call_hook(handler, user_api_key_dict) + assert exc_info.value.status_code == 429 + + +@pytest.mark.asyncio +async def test_key_model_rpm_override_admits_its_own_higher_limit(): + """A key raising its per-model RPM above the team's gets the key's budget.""" + handler, _ = _make_team_model_limit_handler() + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-key-budget-burn"), + team_id="team-burn", + metadata={"model_rpm_limit": {"gpt-4": 5}}, + team_metadata={"model_rpm_limit": {"gpt-4": 2}}, + ) + + for _ in range(5): + await _run_pre_call_hook(handler, user_api_key_dict) + + with pytest.raises(HTTPException) as exc_info: + await _run_pre_call_hook(handler, user_api_key_dict) + assert exc_info.value.status_code == 429 + + +async def _increments_for_completed_call(handler, user_api_key_dict, model="gpt-4"): + """Drive one full request (pre-call gate, then success logging) and return the + counter increments the post-call reconciliation emitted.""" + await _run_pre_call_hook(handler, user_api_key_dict, model=model) + + increments: List[Dict[str, Any]] = [] + + async def mock_increment(increment_list, **kwargs): + increments.extend( + {"key": op["key"], "increment": op["increment_value"]} for op in increment_list + ) + + handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = mock_increment + + await handler.async_log_success_event( + kwargs={ + "litellm_params": {"metadata": {"model_group": model}}, + "standard_logging_object": { + "metadata": { + "user_api_key_hash": user_api_key_dict.api_key, + "user_api_key_team_id": user_api_key_dict.team_id, + } + }, + }, + response_obj=ModelResponse( + usage=Usage(prompt_tokens=40, completion_tokens=60, total_tokens=100) + ), + start_time=datetime.now(), + end_time=datetime.now(), + ) + return increments + + +@pytest.mark.asyncio +async def test_key_tpm_override_does_not_charge_the_team_token_counter(): + """A key exempt from the team's per-model TPM ceiling must not fill it either, + or its traffic starves the keys that do inherit the limit.""" + handler, _ = _make_team_model_limit_handler() + team_id = "team-tpm-accounting" + overriding_key = UserAPIKeyAuth( + api_key=hash_token("sk-tpm-override"), + team_id=team_id, + metadata={"model_tpm_limit": {"gpt-4": 100000}}, + team_metadata={"model_tpm_limit": {"gpt-4": 500}}, + ) + + increments = await _increments_for_completed_call(handler, overriding_key) + + team_tokens_key = handler.create_rate_limit_keys("model_per_team", f"{team_id}:gpt-4", "tokens") + assert not [ + i for i in increments if i["key"] == team_tokens_key + ], f"the team token counter must not be charged by an exempt key, got: {increments}" + + key_tokens_key = handler.create_rate_limit_keys( + "model_per_key", f"{overriding_key.api_key}:gpt-4", "tokens" + ) + assert [ + i for i in increments if i["key"] == key_tokens_key + ], f"the key's own token counter must still be charged, got: {increments}" + + +@pytest.mark.asyncio +async def test_inheriting_key_still_charges_the_team_token_counter(): + """The suppression is scoped to keys that override; an inheriting key still + spends the team's per-model token budget.""" + handler, _ = _make_team_model_limit_handler() + team_id = "team-tpm-accounting" + inheriting_key = UserAPIKeyAuth( + api_key=hash_token("sk-tpm-inherits"), + team_id=team_id, + team_metadata={"model_tpm_limit": {"gpt-4": 500}}, + ) + + increments = await _increments_for_completed_call(handler, inheriting_key) + + team_tokens_key = handler.create_rate_limit_keys("model_per_team", f"{team_id}:gpt-4", "tokens") + assert [ + i for i in increments if i["key"] == team_tokens_key + ], f"an inheriting key must charge the team token counter, got: {increments}" + + +@pytest.mark.asyncio +async def test_key_rpm_override_still_charges_the_team_token_counter(): + """Suppression is per metric: overriding RPM leaves the team's TPM ceiling in + force, so its counter must keep being charged.""" + handler, _ = _make_team_model_limit_handler() + team_id = "team-tpm-accounting" + rpm_only_override = UserAPIKeyAuth( + api_key=hash_token("sk-rpm-override-tpm-inherits"), + team_id=team_id, + metadata={"model_rpm_limit": {"gpt-4": 100}}, + team_metadata={ + "model_rpm_limit": {"gpt-4": 3}, + "model_tpm_limit": {"gpt-4": 500}, + }, + ) + + increments = await _increments_for_completed_call(handler, rpm_only_override) + + team_tokens_key = handler.create_rate_limit_keys("model_per_team", f"{team_id}:gpt-4", "tokens") + assert [ + i for i in increments if i["key"] == team_tokens_key + ], f"an RPM-only override must still charge the team token counter, got: {increments}" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 109d638fb9c..90552b6eae1 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -23578,7 +23578,7 @@ export interface components { * @description Default role assigned to new users created * @default internal_user_viewer */ - user_role: ("internal_user" | "internal_user_viewer" | "proxy_admin" | "proxy_admin_viewer") | null; + user_role: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null; }; /** * DefaultTeamSSOParams @@ -26172,7 +26172,7 @@ export interface components { [key: string]: unknown; } | null; /** Stream Timeout */ - stream_timeout?: number | string | null; + stream_timeout?: string | number | null; /** Tag Regex */ tag_regex?: string[] | null; /** Tags */ @@ -34304,7 +34304,7 @@ export interface components { [key: string]: unknown; } | null; /** Stream Timeout */ - stream_timeout?: number | string | null; + stream_timeout?: string | number | null; /** Tag Regex */ tag_regex?: string[] | null; /** Tags */