diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e4702364a3e..a2dbf3c169d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3805,6 +3805,10 @@ class OrganizationMemberUpdateResponse(MemberUpdateResponse): class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable): team_member_budget_table: Optional[LiteLLM_BudgetTable] = None + # 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 + access_group_agent_ids: Optional[List[str]] = None class TeamInfoResponseObject(TypedDict): diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 3643373be65..ae89f09bb44 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -3042,6 +3042,19 @@ async def team_info( team_info_response_object=_team_info, ) + # Resolve resources inherited from access groups + if _team_info.access_group_ids: + ag_lookup = await _batch_resolve_access_group_resources(_team_info.access_group_ids) + models, mcp_ids, agent_ids = set(), set(), set() + for ag_id in _team_info.access_group_ids: + if ag_id in ag_lookup: + models.update(ag_lookup[ag_id]["models"]) + mcp_ids.update(ag_lookup[ag_id]["mcp_server_ids"]) + agent_ids.update(ag_lookup[ag_id]["agent_ids"]) + _team_info.access_group_models = list(models) + _team_info.access_group_mcp_server_ids = list(mcp_ids) + _team_info.access_group_agent_ids = list(agent_ids) + response_object = TeamInfoResponseObject( team_id=team_id, team_info=_team_info, @@ -3332,6 +3345,36 @@ async def _build_team_list_where_conditions( return where_conditions +async def _batch_resolve_access_group_resources( + all_access_group_ids: List[str], +) -> Dict[str, Dict[str, List[str]]]: + """ + Batch-fetch access groups in a single DB query and return a per-group + resource map. + + Returns {ag_id: {"models": [...], "mcp_server_ids": [...], "agent_ids": [...]}}. + Missing/invalid groups are silently omitted. + """ + from litellm.proxy.proxy_server import prisma_client as _prisma_client + + if not all_access_group_ids or _prisma_client is None: + return {} + + unique_ids = list(set(all_access_group_ids)) + rows = await _prisma_client.db.litellm_accessgrouptable.find_many( + where={"access_group_id": {"in": unique_ids}}, + ) + + result: Dict[str, Dict[str, List[str]]] = {} + for row in rows: + result[row.access_group_id] = { + "models": list(row.access_model_names or []), + "mcp_server_ids": list(row.access_mcp_server_ids or []), + "agent_ids": list(row.access_agent_ids or []), + } + return result + + def _convert_teams_to_response_models( teams: list, use_deleted_table: bool, @@ -3558,6 +3601,30 @@ async def list_team_v2( # Convert Prisma models to response models with members_count team_list = _convert_teams_to_response_models(teams, use_deleted_table) + # Resolve resources inherited from access groups (single batch query) + if not use_deleted_table: + team_items_with_ag = [ + t for t in team_list + if isinstance(t, TeamListItem) and t.access_group_ids + ] + if team_items_with_ag: + all_ag_ids = [ + ag_id + for t in team_items_with_ag + for ag_id in (t.access_group_ids or []) + ] + ag_lookup = await _batch_resolve_access_group_resources(all_ag_ids) + for team_item in team_items_with_ag: + models, mcp_ids, agent_ids = set(), set(), set() + for ag_id in (team_item.access_group_ids or []): + if ag_id in ag_lookup: + models.update(ag_lookup[ag_id]["models"]) + mcp_ids.update(ag_lookup[ag_id]["mcp_server_ids"]) + agent_ids.update(ag_lookup[ag_id]["agent_ids"]) + team_item.access_group_models = list(models) + team_item.access_group_mcp_server_ids = list(mcp_ids) + team_item.access_group_agent_ids = list(agent_ids) + return { "teams": team_list, "total": total_count, diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index 5055a65783f..2455eb495d1 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -47,6 +47,10 @@ class TeamListItem(LiteLLM_TeamTable): """A team item in the paginated list response, enriched with computed fields.""" members_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 + access_group_agent_ids: Optional[List[str]] = None class TeamListResponse(BaseModel): diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index e11dbbd8915..232c698603a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -6497,3 +6497,131 @@ async def test_create_team_member_budget_table_with_duration(): assert budget_request.budget_duration == "30d" assert budget_request.max_budget == 20.0 assert result["metadata"]["team_member_budget_id"] == "budget-abc" + + +# --------------------------------------------------------------------------- +# Tests for _batch_resolve_access_group_resources +# --------------------------------------------------------------------------- + + +class TestBatchResolveAccessGroupResources: + """Tests for the batch access group resource resolution helper.""" + + @pytest.mark.asyncio + async def test_returns_empty_when_no_ids(self): + """Empty list should return empty dict.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _batch_resolve_access_group_resources, + ) + + assert await _batch_resolve_access_group_resources([]) == {} + + @pytest.mark.asyncio + async def test_single_access_group(self): + """Single access group should return its resources.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _batch_resolve_access_group_resources, + ) + + fake_row = MagicMock() + fake_row.access_group_id = "ag-1" + fake_row.access_model_names = ["gpt-4", "claude-3"] + fake_row.access_mcp_server_ids = ["mcp-1"] + fake_row.access_agent_ids = ["agent-1", "agent-2"] + + fake_prisma = MagicMock() + fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock(return_value=[fake_row]) + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): + result = await _batch_resolve_access_group_resources(["ag-1"]) + + assert sorted(result["ag-1"]["models"]) == ["claude-3", "gpt-4"] + assert result["ag-1"]["mcp_server_ids"] == ["mcp-1"] + assert sorted(result["ag-1"]["agent_ids"]) == ["agent-1", "agent-2"] + + @pytest.mark.asyncio + async def test_multiple_access_groups(self): + """Multiple access groups returned in a single query.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _batch_resolve_access_group_resources, + ) + + row1 = MagicMock() + row1.access_group_id = "ag-1" + row1.access_model_names = ["gpt-4"] + row1.access_mcp_server_ids = ["mcp-1"] + row1.access_agent_ids = ["agent-1"] + + row2 = MagicMock() + row2.access_group_id = "ag-2" + row2.access_model_names = ["gemini"] + row2.access_mcp_server_ids = ["mcp-2"] + row2.access_agent_ids = ["agent-2"] + + fake_prisma = MagicMock() + fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock(return_value=[row1, row2]) + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): + result = await _batch_resolve_access_group_resources(["ag-1", "ag-2"]) + + assert result["ag-1"]["models"] == ["gpt-4"] + assert result["ag-2"]["models"] == ["gemini"] + + @pytest.mark.asyncio + async def test_missing_access_group_omitted(self): + """If an access group doesn't exist in DB, it's simply not in the result.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _batch_resolve_access_group_resources, + ) + + row1 = MagicMock() + row1.access_group_id = "ag-1" + row1.access_model_names = ["gpt-4"] + row1.access_mcp_server_ids = [] + row1.access_agent_ids = [] + + fake_prisma = MagicMock() + fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock(return_value=[row1]) + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): + result = await _batch_resolve_access_group_resources(["ag-1", "ag-missing"]) + + assert "ag-1" in result + assert "ag-missing" not in result + + @pytest.mark.asyncio + async def test_returns_empty_when_prisma_unavailable(self): + """If prisma_client is None, should return empty dict.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _batch_resolve_access_group_resources, + ) + + with patch("litellm.proxy.proxy_server.prisma_client", None): + result = await _batch_resolve_access_group_resources(["ag-1"]) + + assert result == {} + + @pytest.mark.asyncio + async def test_deduplicates_input_ids(self): + """Duplicate IDs in input should result in a single DB lookup.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _batch_resolve_access_group_resources, + ) + + row1 = MagicMock() + row1.access_group_id = "ag-1" + row1.access_model_names = ["gpt-4"] + row1.access_mcp_server_ids = [] + row1.access_agent_ids = [] + + fake_find_many = AsyncMock(return_value=[row1]) + fake_prisma = MagicMock() + fake_prisma.db.litellm_accessgrouptable.find_many = fake_find_many + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): + result = await _batch_resolve_access_group_resources(["ag-1", "ag-1", "ag-1"]) + + # Should have been called with deduplicated list + call_args = fake_find_many.call_args + assert len(call_args.kwargs["where"]["access_group_id"]["in"]) == 1 + assert "ag-1" in result diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/ModelsCell.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/ModelsCell.tsx index 5cabe4c4a8f..62a7fdb783f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/ModelsCell.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/ModelsCell.tsx @@ -1,16 +1,57 @@ import { Badge, Icon, TableCell, Text } from "@tremor/react"; import { ChevronDownIcon, ChevronRightIcon } from "@heroicons/react/outline"; import { getModelDisplayName } from "@/components/key_team_helpers/fetch_available_models_team_key"; -import React, { useState } from "react"; +import React, { useMemo, useState } from "react"; import { Team } from "@/components/key_team_helpers/key_list"; interface ModelsCellProps { team: Team; } +interface ModelEntry { + name: string; + source: "direct" | "access_group"; +} + const ModelsCell = ({ team }: ModelsCellProps) => { const [expandedAccordion, setExpandedAccordion] = useState(false); + const isAllModels = !team.models || team.models.length === 0 || team.models.includes("all-proxy-models"); + + const modelEntries: ModelEntry[] = useMemo(() => { + if (isAllModels) return []; + const entries: ModelEntry[] = team.models.map((m) => ({ + name: m, + source: "direct" as const, + })); + for (const m of team.access_group_models || []) { + entries.push({ name: m, source: "access_group" }); + } + return entries; + }, [team.models, team.access_group_models, isAllModels]); + + const renderBadge = (entry: ModelEntry, index: number) => { + if (entry.name === "all-proxy-models") { + return ( + + All Proxy Models + + ); + } + const displayName = getModelDisplayName(entry.name); + const truncated = displayName.length > 30 ? `${displayName.slice(0, 30)}...` : displayName; + return ( + + {truncated} + + ); + }; + return ( { whiteSpace: "pre-wrap", overflow: "hidden", }} - className={team.models.length > 3 ? "px-0" : ""} + className={modelEntries.length > 3 ? "px-0" : ""} >
- {Array.isArray(team.models) ? ( + {modelEntries.length === 0 ? ( + + All Proxy Models + + ) : (
- {team.models.length === 0 ? ( - - All Proxy Models - - ) : ( - <> -
- {team.models.length > 3 && ( -
- { - setExpandedAccordion((prev) => !prev); - }} - /> -
- )} -
- {team.models.slice(0, 3).map((model: string, index: number) => - model === "all-proxy-models" ? ( - - All Proxy Models - - ) : ( - - - {model.length > 30 - ? `${getModelDisplayName(model).slice(0, 30)}...` - : getModelDisplayName(model)} - - - ), - )} - {team.models.length > 3 && !expandedAccordion && ( - - - +{team.models.length - 3} {team.models.length - 3 === 1 ? "more model" : "more models"} - - - )} - {expandedAccordion && ( -
- {team.models.slice(3).map((model: string, index: number) => - model === "all-proxy-models" ? ( - - All Proxy Models - - ) : ( - - - {model.length > 30 - ? `${getModelDisplayName(model).slice(0, 30)}...` - : getModelDisplayName(model)} - - - ), - )} -
- )} -
+
+ {modelEntries.length > 3 && ( +
+ { + setExpandedAccordion((prev) => !prev); + }} + />
- - )} + )} +
+ {modelEntries.slice(0, 3).map((entry, index) => renderBadge(entry, index))} + {modelEntries.length > 3 && !expandedAccordion && ( + + + +{modelEntries.length - 3} {modelEntries.length - 3 === 1 ? "more model" : "more models"} + + + )} + {expandedAccordion && ( +
+ {modelEntries.slice(3).map((entry, index) => renderBadge(entry, index + 3))} +
+ )} +
+
- ) : null} + )}
); diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index a681e438cd1..04b9a5c9962 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -15,6 +15,10 @@ export interface Team { keys: KeyResponse[]; members_with_roles: Member[]; spend: number; + access_group_ids?: string[]; + access_group_models?: string[]; + access_group_mcp_server_ids?: string[]; + access_group_agent_ids?: string[]; } export interface KeyResponse { diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 2d51bc31f05..a6fbe331b9c 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -95,6 +95,9 @@ export interface TeamData { } | null; created_at: string; access_group_ids?: string[]; + access_group_models?: string[]; + access_group_mcp_server_ids?: string[]; + access_group_agent_ids?: string[]; guardrails?: string[]; policies?: string[]; object_permission?: { @@ -657,14 +660,21 @@ const TeamInfoView: React.FC = ({ Models
- {info.models.length === 0 ? ( + {info.models.length === 0 || info.models.includes("all-proxy-models") ? ( All proxy models ) : ( - info.models.map((model, index) => ( - - {model} - - )) + <> + {info.models.map((model: string, index: number) => ( + + {model} + + ))} + {(info.access_group_models || []).map((model: string, index: number) => ( + + {model} + + ))} + )}