Merge pull request #25027 from BerriAI/litellm_add-access-group-to-model

feat(teams): resolve access group resources in team endpoints

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):

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,

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):

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

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<boolean>(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 (
+        <Badge key={index} size={"xs"} color="red">
+          <Text>All Proxy Models</Text>
+        </Badge>
+      );
+    }
+    const displayName = getModelDisplayName(entry.name);
+    const truncated = displayName.length > 30 ? `${displayName.slice(0, 30)}...` : displayName;
+    return (
+      <Badge
+        key={index}
+        size={"xs"}
+        color={entry.source === "access_group" ? "green" : "blue"}
+        title={entry.source === "access_group" ? "From access group" : "Direct assignment"}
+      >
+        <Text>{truncated}</Text>
+      </Badge>
+    );
+  };
+
   return (
     <TableCell
       style={{
@@ -18,78 +59,46 @@ const ModelsCell = ({ team }: ModelsCellProps) => {
         whiteSpace: "pre-wrap",
         overflow: "hidden",
       }}
-      className={team.models.length > 3 ? "px-0" : ""}
+      className={modelEntries.length > 3 ? "px-0" : ""}
     >
       <div className="flex flex-col">
-        {Array.isArray(team.models) ? (
+        {modelEntries.length === 0 ? (
+          <Badge size={"xs"} className="mb-1" color="red">
+            <Text>All Proxy Models</Text>
+          </Badge>
+        ) : (
           <div className="flex flex-col">
-            {team.models.length === 0 ? (
-              <Badge size={"xs"} className="mb-1" color="red">
-                <Text>All Proxy Models</Text>
-              </Badge>
-            ) : (
-              <>
-                <div className="flex items-start">
-                  {team.models.length > 3 && (
-                    <div>
-                      <Icon
-                        icon={expandedAccordion ? ChevronDownIcon : ChevronRightIcon}
-                        className="cursor-pointer"
-                        size="xs"
-                        onClick={() => {
-                          setExpandedAccordion((prev) => !prev);
-                        }}
-                      />
-                    </div>
-                  )}
-                  <div className="flex flex-wrap gap-1">
-                    {team.models.slice(0, 3).map((model: string, index: number) =>
-                      model === "all-proxy-models" ? (
-                        <Badge key={index} size={"xs"} color="red">
-                          <Text>All Proxy Models</Text>
-                        </Badge>
-                      ) : (
-                        <Badge key={index} size={"xs"} color="blue">
-                          <Text>
-                            {model.length > 30
-                              ? `${getModelDisplayName(model).slice(0, 30)}...`
-                              : getModelDisplayName(model)}
-                          </Text>
-                        </Badge>
-                      ),
-                    )}
-                    {team.models.length > 3 && !expandedAccordion && (
-                      <Badge size={"xs"} color="gray" className="cursor-pointer">
-                        <Text>
-                          +{team.models.length - 3} {team.models.length - 3 === 1 ? "more model" : "more models"}
-                        </Text>
-                      </Badge>
-                    )}
-                    {expandedAccordion && (
-                      <div className="flex flex-wrap gap-1">
-                        {team.models.slice(3).map((model: string, index: number) =>
-                          model === "all-proxy-models" ? (
-                            <Badge key={index + 3} size={"xs"} color="red">
-                              <Text>All Proxy Models</Text>
-                            </Badge>
-                          ) : (
-                            <Badge key={index + 3} size={"xs"} color="blue">
-                              <Text>
-                                {model.length > 30
-                                  ? `${getModelDisplayName(model).slice(0, 30)}...`
-                                  : getModelDisplayName(model)}
-                              </Text>
-                            </Badge>
-                          ),
-                        )}
-                      </div>
-                    )}
-                  </div>
+            <div className="flex items-start">
+              {modelEntries.length > 3 && (
+                <div>
+                  <Icon
+                    icon={expandedAccordion ? ChevronDownIcon : ChevronRightIcon}
+                    className="cursor-pointer"
+                    size="xs"
+                    onClick={() => {
+                      setExpandedAccordion((prev) => !prev);
+                    }}
+                  />
                 </div>
-              </>
-            )}
+              )}
+              <div className="flex flex-wrap gap-1">
+                {modelEntries.slice(0, 3).map((entry, index) => renderBadge(entry, index))}
+                {modelEntries.length > 3 && !expandedAccordion && (
+                  <Badge size={"xs"} color="gray" className="cursor-pointer">
+                    <Text>
+                      +{modelEntries.length - 3} {modelEntries.length - 3 === 1 ? "more model" : "more models"}
+                    </Text>
+                  </Badge>
+                )}
+                {expandedAccordion && (
+                  <div className="flex flex-wrap gap-1">
+                    {modelEntries.slice(3).map((entry, index) => renderBadge(entry, index + 3))}
+                  </div>
+                )}
+              </div>
+            </div>
           </div>
-        ) : null}
+        )}
       </div>
     </TableCell>
   );

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 {

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<TeamInfoProps> = ({
                 <Card>
                   <Text>Models</Text>
                   <div className="mt-2 flex flex-wrap gap-2">
-                    {info.models.length === 0 ? (
+                    {info.models.length === 0 || info.models.includes("all-proxy-models") ? (
                       <Badge color="red">All proxy models</Badge>
                     ) : (
-                      info.models.map((model, index) => (
-                        <Badge key={index} color="red">
-                          {model}
-                        </Badge>
-                      ))
+                      <>
+                        {info.models.map((model: string, index: number) => (
+                          <Badge key={`direct-${index}`} color="blue">
+                            {model}
+                          </Badge>
+                        ))}
+                        {(info.access_group_models || []).map((model: string, index: number) => (
+                          <Badge key={`ag-${index}`} color="green" title="From access group">
+                            {model}
+                          </Badge>
+                        ))}
+                      </>
                     )}
                   </div>
                 </Card>