diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index ac606a71ba6..ae89f09bb44 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -110,16 +110,6 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( router = APIRouter() -def _get_access_object(*args, **kwargs): - """ - Lazily import and delegate to `get_access_object` from - `litellm.proxy.auth.auth_checks` to avoid module-level cyclic imports. - """ - from litellm.proxy.auth.auth_checks import get_access_object as _inner_get_access_object - - return _inner_get_access_object(*args, **kwargs) - - class TeamMemberBudgetHandler: """Helper class to handle team member budget, RPM, and TPM limit operations""" @@ -3053,12 +3043,17 @@ async def team_info( ) # Resolve resources inherited from access groups - resolved = await _resolve_access_group_resources( - access_group_ids=_team_info.access_group_ids, - ) - _team_info.access_group_models = resolved["access_group_models"] - _team_info.access_group_mcp_server_ids = resolved["access_group_mcp_server_ids"] - _team_info.access_group_agent_ids = resolved["access_group_agent_ids"] + 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, @@ -3350,62 +3345,34 @@ async def _build_team_list_where_conditions( return where_conditions -async def _resolve_access_group_resources( - access_group_ids: Optional[List[str]], -) -> Dict[str, List[str]]: +async def _batch_resolve_access_group_resources( + all_access_group_ids: List[str], +) -> Dict[str, Dict[str, List[str]]]: """ - Resolve resources inherited from access groups. + Batch-fetch access groups in a single DB query and return a per-group + resource map. - Fetches each access group object once and extracts all three resource - fields in a single pass (models, MCP servers, agents). - - Returns only the access-group-sourced resources (not direct assignments). - Keeps them separate so callers can distinguish where each resource comes from. + Returns {ag_id: {"models": [...], "mcp_server_ids": [...], "agent_ids": [...]}}. + Missing/invalid groups are silently omitted. """ - empty: Dict[str, List[str]] = { - "access_group_models": [], - "access_group_mcp_server_ids": [], - "access_group_agent_ids": [], - } - if not access_group_ids: - return empty - from litellm.proxy.proxy_server import prisma_client as _prisma_client - from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj - from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache - if _user_api_key_cache is None: - return empty + if not all_access_group_ids or _prisma_client is None: + return {} - if _prisma_client is None: - return empty + unique_ids = list(set(all_access_group_ids)) + rows = await _prisma_client.db.litellm_accessgrouptable.find_many( + where={"access_group_id": {"in": unique_ids}}, + ) - models: List[str] = [] - mcp_ids: List[str] = [] - agent_ids: List[str] = [] - - for ag_id in access_group_ids: - try: - ag = await _get_access_object( - access_group_id=ag_id, - prisma_client=_prisma_client, - user_api_key_cache=_user_api_key_cache, - proxy_logging_obj=_proxy_logging_obj, - ) - models.extend(ag.access_model_names or []) - mcp_ids.extend(ag.access_mcp_server_ids or []) - agent_ids.extend(ag.access_agent_ids or []) - except Exception: - verbose_proxy_logger.debug( - "Could not fetch access group %s for resource resolution", - ag_id, - ) - - return { - "access_group_models": list(set(models)), - "access_group_mcp_server_ids": list(set(mcp_ids)), - "access_group_agent_ids": list(set(agent_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( @@ -3634,27 +3601,29 @@ 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 for each team + # 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: - results = await asyncio.gather( - *[ - _resolve_access_group_resources( - access_group_ids=t.access_group_ids, - ) - for t in team_items_with_ag - ] - ) - for team_item, resolved in zip(team_items_with_ag, results): - team_item.access_group_models = resolved["access_group_models"] - team_item.access_group_mcp_server_ids = resolved[ - "access_group_mcp_server_ids" - ] - team_item.access_group_agent_ids = resolved["access_group_agent_ids"] + 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, 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 a2c8c4d560c..8f0a045cf86 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -6494,175 +6494,128 @@ async def test_create_team_member_budget_table_with_duration(): # --------------------------------------------------------------------------- -# Tests for _resolve_access_group_resources +# Tests for _batch_resolve_access_group_resources # --------------------------------------------------------------------------- -class TestResolveAccessGroupResources: - """Tests for the single-pass access group resource resolution helper.""" +class TestBatchResolveAccessGroupResources: + """Tests for the batch access group resource resolution helper.""" @pytest.mark.asyncio - async def test_returns_empty_when_no_access_group_ids(self): - """None or empty list should return empty lists for all resource types.""" + async def test_returns_empty_when_no_ids(self): + """Empty list should return empty dict.""" from litellm.proxy.management_endpoints.team_endpoints import ( - _resolve_access_group_resources, + _batch_resolve_access_group_resources, ) - result_none = await _resolve_access_group_resources(access_group_ids=None) - assert result_none == { - "access_group_models": [], - "access_group_mcp_server_ids": [], - "access_group_agent_ids": [], - } - - result_empty = await _resolve_access_group_resources(access_group_ids=[]) - assert result_empty == { - "access_group_models": [], - "access_group_mcp_server_ids": [], - "access_group_agent_ids": [], - } + 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._types import LiteLLM_AccessGroupTable from litellm.proxy.management_endpoints.team_endpoints import ( - _resolve_access_group_resources, + _batch_resolve_access_group_resources, ) - fake_ag = LiteLLM_AccessGroupTable( - access_group_id="ag-1", - access_group_name="test-group", - access_model_names=["gpt-4", "claude-3"], - access_mcp_server_ids=["mcp-1"], - access_agent_ids=["agent-1", "agent-2"], - ) + 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"] - with patch( - "litellm.proxy.management_endpoints.team_endpoints._get_access_object", - new_callable=AsyncMock, - return_value=fake_ag, - ): - with patch( - "litellm.proxy.proxy_server.user_api_key_cache", - MagicMock(), - ): - with patch( - "litellm.proxy.proxy_server.prisma_client", - MagicMock(), - ): - result = await _resolve_access_group_resources( - access_group_ids=["ag-1"], - ) + fake_prisma = MagicMock() + fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock(return_value=[fake_row]) - assert sorted(result["access_group_models"]) == ["claude-3", "gpt-4"] - assert result["access_group_mcp_server_ids"] == ["mcp-1"] - assert sorted(result["access_group_agent_ids"]) == ["agent-1", "agent-2"] + 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_deduplicates(self): - """Multiple access groups with overlapping resources should deduplicate.""" - from litellm.proxy._types import LiteLLM_AccessGroupTable + async def test_multiple_access_groups(self): + """Multiple access groups returned in a single query.""" from litellm.proxy.management_endpoints.team_endpoints import ( - _resolve_access_group_resources, + _batch_resolve_access_group_resources, ) - ag1 = LiteLLM_AccessGroupTable( - access_group_id="ag-1", - access_group_name="group-1", - access_model_names=["gpt-4", "claude-3"], - access_mcp_server_ids=["mcp-1"], - access_agent_ids=["agent-1"], - ) - ag2 = LiteLLM_AccessGroupTable( - access_group_id="ag-2", - access_group_name="group-2", - access_model_names=["gpt-4", "gemini"], - access_mcp_server_ids=["mcp-1", "mcp-2"], - access_agent_ids=["agent-2"], - ) + 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"] - async def fake_get_access_object(access_group_id, **kwargs): - return {"ag-1": ag1, "ag-2": ag2}[access_group_id] + 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"] - with patch( - "litellm.proxy.management_endpoints.team_endpoints._get_access_object", - side_effect=fake_get_access_object, - ): - with patch( - "litellm.proxy.proxy_server.user_api_key_cache", - MagicMock(), - ): - with patch( - "litellm.proxy.proxy_server.prisma_client", - MagicMock(), - ): - result = await _resolve_access_group_resources( - access_group_ids=["ag-1", "ag-2"], - ) + fake_prisma = MagicMock() + fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock(return_value=[row1, row2]) - assert sorted(result["access_group_models"]) == ["claude-3", "gemini", "gpt-4"] - assert sorted(result["access_group_mcp_server_ids"]) == ["mcp-1", "mcp-2"] - assert sorted(result["access_group_agent_ids"]) == ["agent-1", "agent-2"] + 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_skipped(self): - """If an access group doesn't exist, it should be skipped gracefully.""" - from litellm.proxy._types import LiteLLM_AccessGroupTable + 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 ( - _resolve_access_group_resources, + _batch_resolve_access_group_resources, ) - ag1 = LiteLLM_AccessGroupTable( - access_group_id="ag-1", - access_group_name="group-1", - access_model_names=["gpt-4"], - access_mcp_server_ids=[], - access_agent_ids=[], - ) + row1 = MagicMock() + row1.access_group_id = "ag-1" + row1.access_model_names = ["gpt-4"] + row1.access_mcp_server_ids = [] + row1.access_agent_ids = [] - async def fake_get_access_object(access_group_id, **kwargs): - if access_group_id == "ag-1": - return ag1 - raise HTTPException(status_code=404, detail="Not found") + fake_prisma = MagicMock() + fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock(return_value=[row1]) - with patch( - "litellm.proxy.management_endpoints.team_endpoints._get_access_object", - side_effect=fake_get_access_object, - ): - with patch( - "litellm.proxy.proxy_server.user_api_key_cache", - MagicMock(), - ): - with patch( - "litellm.proxy.proxy_server.prisma_client", - MagicMock(), - ): - result = await _resolve_access_group_resources( - access_group_ids=["ag-1", "ag-missing"], - ) + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): + result = await _batch_resolve_access_group_resources(["ag-1", "ag-missing"]) - assert result["access_group_models"] == ["gpt-4"] - assert result["access_group_mcp_server_ids"] == [] - assert result["access_group_agent_ids"] == [] + assert "ag-1" in result + assert "ag-missing" not in result @pytest.mark.asyncio - async def test_returns_empty_when_cache_unavailable(self): - """If user_api_key_cache is None, should return empty results.""" + 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 ( - _resolve_access_group_resources, + _batch_resolve_access_group_resources, ) - with patch( - "litellm.proxy.proxy_server.user_api_key_cache", - None, - ): - result = await _resolve_access_group_resources( - access_group_ids=["ag-1"], - ) + with patch("litellm.proxy.proxy_server.prisma_client", None): + result = await _batch_resolve_access_group_resources(["ag-1"]) - assert result == { - "access_group_models": [], - "access_group_mcp_server_ids": [], - "access_group_agent_ids": [], - } + 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