diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d2dfa177515..31dab435b65 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3187,6 +3187,8 @@ class PrismaClient: t.organization_id as org_id, p.project_alias AS project_alias, tm.spend AS team_member_spend, + b_tm.tpm_limit AS team_member_tpm_limit, + b_tm.rpm_limit AS team_member_rpm_limit, m.aliases AS team_model_aliases, -- Added comma to separate b.* columns b.max_budget AS litellm_budget_table_max_budget, @@ -3202,6 +3204,7 @@ class PrismaClient: FROM "LiteLLM_VerificationToken" AS v LEFT JOIN "LiteLLM_TeamTable" AS t ON v.team_id = t.team_id LEFT JOIN "LiteLLM_TeamMembership" AS tm ON v.team_id = tm.team_id AND tm.user_id = v.user_id + LEFT JOIN "LiteLLM_BudgetTable" AS b_tm ON tm.budget_id = b_tm.budget_id LEFT JOIN "LiteLLM_ModelTable" m ON t.model_id = m.id LEFT JOIN "LiteLLM_BudgetTable" AS b ON v.budget_id = b.budget_id LEFT JOIN "LiteLLM_ProjectTable" AS p ON v.project_id = p.project_id 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 d4fb5b72719..370477c3605 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 @@ -1029,6 +1029,83 @@ async def test_team_member_rate_limits_v3(): ), "Team member TPM limit should be set" +@pytest.mark.asyncio +async def test_team_member_rate_limits_v3_raises_429_when_over_limit(): + """ + When should_rate_limit reports OVER_LIMIT for the team_member descriptor, the + pre-call hook raises HTTP 429 with rate_limit headers — same contract as + test_rpm_api_key_rate_limits_v3 / test_tpm_api_key_rate_limits_v3. + """ + _api_key = hash_token("sk-12345") + _team_id = "team_123" + _user_id = "user_456" + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + team_id=_team_id, + user_id=_user_id, + team_member_rpm_limit=10, + team_member_tpm_limit=1000, + ) + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + captured_descriptors = None + + async def mock_should_rate_limit(descriptors, **kwargs): + nonlocal captured_descriptors + captured_descriptors = descriptors + return { + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "current_limit": 10, + "limit_remaining": -1, + "rate_limit_type": "requests", + "descriptor_key": "team_member", + }, + { + "code": "OK", + "current_limit": 1000, + "limit_remaining": 500, + "rate_limit_type": "tokens", + "descriptor_key": "team_member", + }, + ], + } + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + error = None + try: + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, + call_type="", + ) + except HTTPException as e: + error = e + assert e.status_code == 429 + assert "rate_limit_type" in e.headers + assert e.headers.get("rate_limit_type") == "requests" + assert "retry-after" in e.headers + + assert error is not None, "An Exception must be thrown" + assert captured_descriptors is not None, "Rate limit descriptors should be captured" + team_member_descriptor = None + for descriptor in captured_descriptors: + if descriptor["key"] == "team_member": + team_member_descriptor = descriptor + break + assert team_member_descriptor is not None + assert team_member_descriptor["value"] == f"{_team_id}:{_user_id}" + + @pytest.mark.asyncio async def test_dynamic_rate_limiting_v3(): """