diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7953102b0..ae50326ca 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2033,6 +2033,7 @@ class TeamInfoResponseObject(TypedDict): class TeamListResponseObject(LiteLLM_TeamTable): team_memberships: List[LiteLLM_TeamMembership] + keys: List # list of keys that belong to the team class CurrentItemRateLimit(TypedDict): diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 965f07da7..74289c90a 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1275,10 +1275,17 @@ async def list_team( for tm in returned_tm: if tm.team_id == team.team_id: _team_memberships.append(tm) + + # add all keys that belong to the team + keys = await prisma_client.db.litellm_verificationtoken.find_many( + where={"team_id": team.team_id} + ) + returned_responses.append( TeamListResponseObject( **team.model_dump(), team_memberships=_team_memberships, + keys=keys, ) ) diff --git a/litellm/router.py b/litellm/router.py index d129ee556..82a37a9f4 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -5204,7 +5204,6 @@ class Router: parent_otel_span=parent_otel_span, ) raise exception - verbose_router_logger.info( f"get_available_deployment for model: {model}, Selected deployment: {self.print_deployment(deployment)} for model: {model}" ) diff --git a/litellm/router_utils/handle_error.py b/litellm/router_utils/handle_error.py index 2a15fcce0..321ba5dc5 100644 --- a/litellm/router_utils/handle_error.py +++ b/litellm/router_utils/handle_error.py @@ -64,7 +64,6 @@ async def send_llm_exception_alert( ) - async def async_raise_no_deployment_exception( litellm_router_instance: LitellmRouter, model: str, parent_otel_span: Optional[Span] ): @@ -74,7 +73,6 @@ async def async_raise_no_deployment_exception( verbose_router_logger.info( f"get_available_deployment for model: {model}, No deployment available" ) - model_ids = litellm_router_instance.get_model_ids(model_name=model) _cooldown_time = litellm_router_instance.cooldown_cache.get_min_cooldown( model_ids=model_ids, parent_otel_span=parent_otel_span diff --git a/tests/llm_translation/test_optional_params.py b/tests/llm_translation/test_optional_params.py index 163e7d919..fdda7b171 100644 --- a/tests/llm_translation/test_optional_params.py +++ b/tests/llm_translation/test_optional_params.py @@ -786,7 +786,6 @@ def test_unmapped_vertex_anthropic_model(): assert "max_retries" not in optional_params - @pytest.mark.parametrize("provider", ["anthropic", "vertex_ai"]) def test_anthropic_parallel_tool_calls(provider): optional_params = get_optional_params( diff --git a/tests/proxy_admin_ui_tests/test_key_management.py b/tests/proxy_admin_ui_tests/test_key_management.py index ecf81372f..bc7371843 100644 --- a/tests/proxy_admin_ui_tests/test_key_management.py +++ b/tests/proxy_admin_ui_tests/test_key_management.py @@ -25,6 +25,8 @@ import logging import pytest import litellm from litellm._logging import verbose_proxy_logger +from litellm.proxy.management_endpoints.team_endpoints import list_team +from litellm.proxy._types import * from litellm.proxy.management_endpoints.internal_user_endpoints import ( new_user, user_info, @@ -421,3 +423,120 @@ async def test_get_users_key_count(prisma_client): assert ( updated_key_count == initial_key_count + 1 ), f"Expected key count to increase by 1, but got {updated_key_count} (was {initial_key_count})" + + +async def cleanup_existing_teams(prisma_client): + all_teams = await prisma_client.db.litellm_teamtable.find_many() + for team in all_teams: + await prisma_client.delete_data(team_id_list=[team.team_id], table_name="team") + + +@pytest.mark.asyncio +async def test_list_teams(prisma_client): + """ + Tests /team/list endpoint to verify it returns both keys and members_with_roles + """ + litellm.set_verbose = True + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + await litellm.proxy.proxy_server.prisma_client.connect() + + # Delete all existing teams first + await cleanup_existing_teams(prisma_client) + + # Create a test team with members + team_id = f"test_team_{uuid.uuid4()}" + team_alias = f"test_team_alias_{uuid.uuid4()}" + test_team = await new_team( + data=NewTeamRequest( + team_id=team_id, + team_alias=team_alias, + members_with_roles=[ + Member(role="admin", user_id="test_user_1"), + Member(role="user", user_id="test_user_2"), + ], + models=["gpt-4"], + tpm_limit=1000, + rpm_limit=1000, + budget_duration="30d", + max_budget=1000, + ), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="admin" + ), + ) + + # Create a key for the team + test_key = await generate_key_fn( + data=GenerateKeyRequest( + team_id=team_id, + key_alias=f"test_key_{uuid.uuid4()}", + ), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="admin" + ), + ) + + # Get team list + teams = await list_team( + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="admin" + ), + user_id=None, + ) + + print("teams", teams) + + # Find our test team in the response + test_team_response = None + for team in teams: + if team.team_id == team_id: + test_team_response = team + break + + assert ( + test_team_response is not None + ), f"Could not find test team {team_id} in response" + + # Verify members_with_roles + assert ( + len(test_team_response.members_with_roles) == 3 + ), "Expected 3 members in team" # 2 members + 1 team admin + member_roles = {m.role for m in test_team_response.members_with_roles} + assert "admin" in member_roles, "Expected admin role in members" + assert "user" in member_roles, "Expected user role in members" + + # Verify all required fields in TeamListResponseObject + assert ( + test_team_response.team_id == team_id + ), f"team_id should be expected value {team_id}" + assert ( + test_team_response.team_alias == team_alias + ), f"team_alias should be expected value {team_alias}" + assert test_team_response.spend is not None, "spend should not be None" + assert ( + test_team_response.max_budget == 1000 + ), f"max_budget should be expected value 1000" + assert test_team_response.models == [ + "gpt-4" + ], f"models should be expected value ['gpt-4']" + assert ( + test_team_response.tpm_limit == 1000 + ), f"tpm_limit should be expected value 1000" + assert ( + test_team_response.rpm_limit == 1000 + ), f"rpm_limit should be expected value 1000" + assert ( + test_team_response.budget_reset_at is not None + ), "budget_reset_at should not be None since budget_duration is 30d" + + # Verify keys are returned + assert len(test_team_response.keys) > 0, "Expected at least one key for team" + assert any( + k.team_id == team_id for k in test_team_response.keys + ), "Expected to find team key in response" + + # Clean up + await prisma_client.delete_data(team_id_list=[team_id], table_name="team") diff --git a/ui/litellm-dashboard/src/components/teams.tsx b/ui/litellm-dashboard/src/components/teams.tsx index 2b2c57fe7..90a29de32 100644 --- a/ui/litellm-dashboard/src/components/teams.tsx +++ b/ui/litellm-dashboard/src/components/teams.tsx @@ -551,11 +551,8 @@ const Team: React.FC = ({ {perTeamInfo && team.team_id && perTeamInfo[team.team_id] && - perTeamInfo[team.team_id].team_info && - perTeamInfo[team.team_id].team_info - .members_with_roles && - perTeamInfo[team.team_id].team_info - .members_with_roles.length}{" "} + perTeamInfo[team.team_id].members_with_roles && + perTeamInfo[team.team_id].members_with_roles.length}{" "} Members