tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-13 02:22:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf(teams): single-pass access group resolution + asyncio.gather in list endpoint

Browse Source 
- Fetch each access group object once and extract all 3 resource fields
  in a single pass instead of 3 separate calls (3N → N lookups)
- Use asyncio.gather to resolve access groups across teams concurrently
  in list_team_v2 instead of sequential awaits
- Add 5 unit tests for _resolve_access_group_resources
This commit is contained in:
Ryan Crabbe
2026-04-02 14:52:32 -07:00
parent f62658795a
commit bbe708b093
2 changed files with 212 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/management_endpoints/team_endpoints.py
+49 -17
View File
@@ -62,11 +62,9 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.auth.auth_checks import (
_get_agent_ids_from_access_groups,
_get_mcp_server_ids_from_access_groups,
_get_models_from_access_groups,
allowed_route_check_inside_route,
can_org_access_model,
get_access_object,
get_org_object,
get_team_object,
get_user_object,
@@ -3349,6 +3347,9 @@ async def _resolve_access_group_resources(
"""
Resolve resources inherited from access groups.
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.
"""
@@ -3360,16 +3361,38 @@ async def _resolve_access_group_resources(
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
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(getattr(ag, "access_model_names", []))
mcp_ids.extend(getattr(ag, "access_mcp_server_ids", []))
agent_ids.extend(getattr(ag, "access_agent_ids", []))
except Exception:
verbose_proxy_logger.debug(
"Could not fetch access group %s for resource resolution",
ag_id,
)
return {
"access_group_models": await _get_models_from_access_groups(
access_group_ids=access_group_ids,
),
"access_group_mcp_server_ids": await _get_mcp_server_ids_from_access_groups(
access_group_ids=access_group_ids,
),
"access_group_agent_ids": await _get_agent_ids_from_access_groups(
access_group_ids=access_group_ids,
),
"access_group_models": list(set(models)),
"access_group_mcp_server_ids": list(set(mcp_ids)),
"access_group_agent_ids": list(set(agent_ids)),
}
@@ -3601,11 +3624,20 @@ async def list_team_v2(
# Resolve resources inherited from access groups for each team
if not use_deleted_table:
for team_item in team_list:
if isinstance(team_item, TeamListItem):
resolved = await _resolve_access_group_resources(
access_group_ids=team_item.access_group_ids,
)
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"
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
+163
View File
@@ -6491,3 +6491,166 @@ 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 _resolve_access_group_resources
# ---------------------------------------------------------------------------
class TestResolveAccessGroupResources:
"""Tests for the single-pass 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."""
from litellm.proxy.management_endpoints.team_endpoints import (
_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": [],
}
@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,
)
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"],
)
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(),
):
result = await _resolve_access_group_resources(
access_group_ids=["ag-1"],
)
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"]
@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
from litellm.proxy.management_endpoints.team_endpoints import (
_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"],
)
async def fake_get_access_object(access_group_id, **kwargs):
return {"ag-1": ag1, "ag-2": ag2}[access_group_id]
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(),
):
result = await _resolve_access_group_resources(
access_group_ids=["ag-1", "ag-2"],
)
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"]
@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
from litellm.proxy.management_endpoints.team_endpoints import (
_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=[],
)
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")
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(),
):
result = await _resolve_access_group_resources(
access_group_ids=["ag-1", "ag-missing"],
)
assert result["access_group_models"] == ["gpt-4"]
assert result["access_group_mcp_server_ids"] == []
assert result["access_group_agent_ids"] == []
@pytest.mark.asyncio
async def test_returns_empty_when_cache_unavailable(self):
"""If user_api_key_cache is None, should return empty results."""
from litellm.proxy.management_endpoints.team_endpoints import (
_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"],
)
assert result == {
"access_group_models": [],
"access_group_mcp_server_ids": [],
"access_group_agent_ids": [],
}