ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(teams): expose keys_count on /v2/team/list and wire UI Resources badge (#28502)

Browse Source
This commit is contained in:
michelligabriele 2026-05-29 14:09:07 +02:00 committed by GitHub
parent f27df8d516
commit 68852ef165
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 305 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

51
litellm/proxy/management_endpoints/team_endpoints.py
View File

@ -3978,11 +3978,13 @@ async def _batch_resolve_access_group_resources(
def _convert_teams_to_response_models(
teams: list,
use_deleted_table: bool,
keys_count_by_team: Optional[Dict[str, int]] = None,
) -> List[Union[TeamListItem, LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]]:
"""Convert raw Prisma team rows to response models."""
team_list: List[
Union[TeamListItem, LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]
] = []
counts = keys_count_by_team or {}
for team in teams:
try:
team_dict = team.model_dump()
@ -3997,10 +3999,45 @@ def _convert_teams_to_response_models(
members_with_roles = []
team_dict["members_with_roles"] = members_with_roles
members_count = len(members_with_roles)
team_list.append(TeamListItem(**team_dict, members_count=members_count))
keys_count = counts.get(team_dict.get("team_id") or "", 0)
team_list.append(
TeamListItem(
**team_dict,
members_count=members_count,
keys_count=keys_count,
)
)
return team_list
async def _get_keys_count_by_team(
prisma_client: Any,
teams: list,
) -> Dict[str, int]:
"""Aggregate virtual-key counts per team for the given page of teams.
Runs a single GROUP BY against LiteLLM_VerificationToken. The IN clause is
bounded by page_size and uses the existing @@index([team_id]), so this is
one DB round-trip per page. Returns an empty map when the page has no teams.
"""
page_team_ids = [
getattr(t, "team_id", None) for t in teams if getattr(t, "team_id", None)
]
if not page_team_ids:
return {}
grouped = await prisma_client.db.litellm_verificationtoken.group_by(
by=["team_id"],
where={"team_id": {"in": page_team_ids}},
count={"team_id": True},
)
return {
row["team_id"]: row.get("_count", {}).get("team_id", 0)
for row in grouped
if row.get("team_id")
}
async def _enforce_list_team_v2_access(
user_api_key_dict: UserAPIKeyAuth,
user_id: Optional[str],
@ -4228,8 +4265,16 @@ async def list_team_v2(
# Calculate total pages
total_pages = -(-total_count // page_size) # Ceiling division
# Convert Prisma models to response models with members_count
team_list = _convert_teams_to_response_models(teams, use_deleted_table)
# Aggregate virtual-key counts per team for the current page. The deleted
# table does not carry keys_count, so it is skipped.
keys_count_by_team: Dict[str, int] = {}
if not use_deleted_table:
keys_count_by_team = await _get_keys_count_by_team(prisma_client, teams)
# Convert Prisma models to response models with members_count and keys_count
team_list = _convert_teams_to_response_models(
teams, use_deleted_table, keys_count_by_team=keys_count_by_team
)
# Resolve resources inherited from access groups (single batch query)
if not use_deleted_table:

1
litellm/types/proxy/management_endpoints/team_endpoints.py
View File

@ -69,6 +69,7 @@ class TeamListItem(LiteLLM_TeamTable):
"""A team item in the paginated list response, enriched with computed fields."""
members_count: int = 0
keys_count: int = 0
# Resources inherited from access groups (separate from direct assignments)
access_group_models: Optional[List[str]] = None
access_group_mcp_server_ids: Optional[List[str]] = None

161
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
View File

@ -2832,6 +2832,7 @@ async def test_list_team_v2_security_check_non_admin_user_own_teams():
]
mock_db.litellm_teamtable.find_many = AsyncMock(return_value=mock_teams)
mock_db.litellm_teamtable.count = AsyncMock(return_value=2)
mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[])
with patch(
"litellm.proxy.management_endpoints.team_endpoints.get_user_object",
@ -2888,6 +2889,7 @@ async def test_list_team_v2_security_check_admin_user():
]
mock_db.litellm_teamtable.find_many = AsyncMock(return_value=mock_teams)
mock_db.litellm_teamtable.count = AsyncMock(return_value=2)
mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[])
# Should NOT raise an exception
result = await list_team_v2(
@ -3036,6 +3038,7 @@ async def test_list_team_v2_org_admin_sees_org_teams():
}
mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team])
mock_db.litellm_teamtable.count = AsyncMock(return_value=1)
mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[])
result = await list_team_v2(
http_request=mock_request,
@ -3211,6 +3214,7 @@ async def test_list_team_v2_org_admin_with_user_id_returns_user_teams():
}
mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team])
mock_db.litellm_teamtable.count = AsyncMock(return_value=1)
mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[])
result = await list_team_v2(
http_request=mock_request,
@ -3390,6 +3394,163 @@ async def test_list_team_v2_search_composes_with_user_id_filter():
assert where["team_id"] == {"in": ["team_a", "team_b"]}
@pytest.mark.asyncio
async def test_list_team_v2_populates_keys_count():
"""
Test that list_team_v2 returns a keys_count per team derived from a single
batched group_by against LiteLLM_VerificationToken.
"""
from unittest.mock import AsyncMock, Mock, patch
from fastapi import Request
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.team_endpoints import list_team_v2
mock_request = Mock(spec=Request)
mock_user_api_key_dict_admin = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="admin_user_123",
)
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client:
mock_db = Mock()
mock_prisma_client.db = mock_db
team_a = Mock()
team_a.team_id = "team_a"
team_a.model_dump = lambda: {
"team_id": "team_a",
"team_alias": "Team A",
"members_with_roles": [{"user_id": "u1", "role": "user"}],
}
team_b = Mock()
team_b.team_id = "team_b"
team_b.model_dump = lambda: {
"team_id": "team_b",
"team_alias": "Team B",
"members_with_roles": [],
}
mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[team_a, team_b])
mock_db.litellm_teamtable.count = AsyncMock(return_value=2)
mock_db.litellm_verificationtoken.group_by = AsyncMock(
return_value=[
{"team_id": "team_a", "_count": {"team_id": 3}},
# team_b intentionally absent → expect 0
]
)
result = await list_team_v2(
http_request=mock_request,
user_id=None,
user_api_key_dict=mock_user_api_key_dict_admin,
page=1,
page_size=10,
status=None,
)
assert result["total"] == 2
by_id = {t.team_id: t for t in result["teams"]}
assert by_id["team_a"].keys_count == 3
assert by_id["team_b"].keys_count == 0
# The aggregate is one batched query, filtered by the page's team IDs.
group_by_kwargs = mock_db.litellm_verificationtoken.group_by.call_args.kwargs
assert group_by_kwargs["by"] == ["team_id"]
assert group_by_kwargs["where"] == {"team_id": {"in": ["team_a", "team_b"]}}
assert group_by_kwargs["count"] == {"team_id": True}
@pytest.mark.asyncio
async def test_list_team_v2_keys_count_skipped_for_empty_page():
"""
When the page has no teams, the keys-count group_by must not be issued.
"""
from unittest.mock import AsyncMock, Mock, patch
from fastapi import Request
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.team_endpoints import list_team_v2
mock_request = Mock(spec=Request)
mock_user_api_key_dict_admin = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="admin_user_123",
)
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client:
mock_db = Mock()
mock_prisma_client.db = mock_db
mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[])
mock_db.litellm_teamtable.count = AsyncMock(return_value=0)
mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[])
result = await list_team_v2(
http_request=mock_request,
user_id=None,
user_api_key_dict=mock_user_api_key_dict_admin,
page=1,
page_size=10,
status=None,
)
assert result["total"] == 0
assert result["teams"] == []
mock_db.litellm_verificationtoken.group_by.assert_not_called()
@pytest.mark.asyncio
async def test_list_team_v2_keys_count_skipped_for_deleted_status():
"""
The deleted-table branch returns LiteLLM_DeletedTeamTable items, which do
not carry keys_count — group_by must not be issued.
"""
from unittest.mock import AsyncMock, Mock, patch
from fastapi import Request
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.team_endpoints import list_team_v2
mock_request = Mock(spec=Request)
mock_user_api_key_dict_admin = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="admin_user_123",
)
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client:
mock_db = Mock()
mock_prisma_client.db = mock_db
mock_deleted = Mock()
mock_deleted.team_id = "team_d"
mock_deleted.model_dump = lambda: {
"team_id": "team_d",
"team_alias": "Deleted Team",
}
mock_db.litellm_deletedteamtable.find_many = AsyncMock(
return_value=[mock_deleted]
)
mock_db.litellm_deletedteamtable.count = AsyncMock(return_value=1)
mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[])
result = await list_team_v2(
http_request=mock_request,
user_id=None,
user_api_key_dict=mock_user_api_key_dict_admin,
page=1,
page_size=10,
status="deleted",
)
assert result["total"] == 1
mock_db.litellm_verificationtoken.group_by.assert_not_called()
@pytest.mark.asyncio
async def test_team_member_delete_cleans_membership(mock_db_client, mock_admin_auth):
"""

80
ui/litellm-dashboard/src/components/OldTeams.test.tsx
View File

@ -1018,3 +1018,83 @@ describe("OldTeams - organization alias display", () => {
});
});
});
describe("OldTeams - Resources column keys badge", () => {
beforeEach(() => {
vi.clearAllMocks();
mockUseOrganizations.mockReturnValue({ data: [] });
});
it("renders keys_count from the v2 payload in the Resources badge", async () => {
const { container } = renderWithQueryClient(
<OldTeams
teams={[
{
team_id: "1",
team_alias: "Team With Keys",
organization_id: "org-123",
models: ["gpt-4"],
max_budget: 100,
budget_duration: "1d",
tpm_limit: 1000,
rpm_limit: 1000,
created_at: new Date().toISOString(),
keys: [],
keys_count: 3,
members_with_roles: [],
spend: 0,
} as any,
]}
searchParams={{}}
accessToken="test-token"
setTeams={vi.fn()}
userID="user-123"
userRole="Admin"
organizations={[]}
/>,
);
await waitFor(() => {
expect(screen.getByText("Team With Keys")).toBeInTheDocument();
});
const cyanTag = container.querySelector(".ant-tag-cyan");
expect(cyanTag).not.toBeNull();
expect(cyanTag?.textContent).toContain("3");
});
it("falls back to keys.length when keys_count is absent", async () => {
const { container } = renderWithQueryClient(
<OldTeams
teams={[
{
team_id: "2",
team_alias: "Legacy Team",
organization_id: "org-123",
models: ["gpt-4"],
max_budget: 100,
budget_duration: "1d",
tpm_limit: 1000,
rpm_limit: 1000,
created_at: new Date().toISOString(),
keys: [{ token: "t1" } as any, { token: "t2" } as any],
members_with_roles: [],
spend: 0,
} as any,
]}
searchParams={{}}
accessToken="test-token"
setTeams={vi.fn()}
userID="user-123"
userRole="Admin"
organizations={[]}
/>,
);
await waitFor(() => {
expect(screen.getByText("Legacy Team")).toBeInTheDocument();
});
const cyanTag = container.querySelector(".ant-tag-cyan");
expect(cyanTag).not.toBeNull();
expect(cyanTag?.textContent).toContain("2");
});
});

20
ui/litellm-dashboard/src/components/OldTeams.tsx
View File

@ -105,6 +105,7 @@ interface TeamInfo {
interface PerTeamInfo {
keys: KeyResponse[];
keys_count: number;
team_info: TeamInfo;
}
@ -364,6 +365,7 @@ const Teams: React.FC<TeamProps> = ({
(acc, team) => {
acc[team.team_id] = {
keys: team.keys || [],
keys_count: team.keys_count ?? team.keys?.length ?? 0,
team_info: {
members_with_roles: team.members_with_roles || [],
},
@ -745,7 +747,7 @@ const Teams: React.FC<TeamProps> = ({
render: (_: unknown, record: Team) => {
const memberCount = perTeamInfo?.[record.team_id]?.team_info?.members_with_roles?.length ?? 0;
const modelCount = record.models?.length ?? 0;
const keyCount = perTeamInfo?.[record.team_id]?.keys?.length ?? 0;
const keyCount = perTeamInfo?.[record.team_id]?.keys_count ?? 0;
return (
<Flex gap={12} align="center">
<Tooltip title={`${memberCount} Members`}>
@ -977,17 +979,23 @@ const Teams: React.FC<TeamProps> = ({
<DeleteResourceModal
isOpen={isDeleteModalOpen}
title="Delete Team?"
alertMessage={
teamToDelete?.keys?.length === 0
alertMessage={(() => {
const deleteKeyCount =
teamToDelete?.keys_count ?? teamToDelete?.keys?.length ?? 0;
return deleteKeyCount === 0
? undefined
: `Warning: This team has ${teamToDelete?.keys?.length} keys associated with it. Deleting the team will also delete all associated keys. This action is irreversible.`
}
: `Warning: This team has ${deleteKeyCount} keys associated with it. Deleting the team will also delete all associated keys. This action is irreversible.`;
})()}
message="Are you sure you want to delete this team and all its keys? This action cannot be undone."
resourceInformationTitle="Team Information"
resourceInformation={[
{ label: "Team ID", value: teamToDelete?.team_id, code: true },
{ label: "Team Name", value: teamToDelete?.team_alias },
{ label: "Keys", value: teamToDelete?.keys?.length },
{
label: "Keys",
value:
teamToDelete?.keys_count ?? teamToDelete?.keys?.length ?? 0,
},
{ label: "Members", value: teamToDelete?.members_with_roles?.length },
]}
requiredConfirmation={teamToDelete?.team_alias}

1
ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx
View File

@ -13,6 +13,7 @@ export interface Team {
organization_id: string;
created_at: string;
keys: KeyResponse[];
keys_count?: number;
members_with_roles: Member[];
spend: number;
access_group_ids?: string[];