mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
perf(teams): batch-fetch access groups in single DB query
Replace per-ID _resolve_access_group_resources loop with a single find_many call that deduplicates IDs across all teams. Removes the N+1 query pattern on cold cache for the team list endpoint.
This commit is contained in:
Ryan Crabbe
parent
bb03a11d7c
commit
93369bf60d
2 changed files with 135 additions and 213 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue