diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index c8765eb0bd0..32b8eb3d4d0 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45110 + "limit": 45098 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 39838 + "limit": 39826 }, "reportUnknownParameterType": { "limit": 20237 }, "reportUnknownVariableType": { - "limit": 31383 + "limit": 31371 }, "reportUnnecessaryCast": { "limit": 122 diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c6cb527b2e6..b84e9a91c0f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3894,12 +3894,19 @@ class OrganizationMemberUpdateResponse(MemberUpdateResponse): ########################################## +class TeamAccessGroupModelGrant(LiteLLMPydanticObjectBase): + access_group_id: str + access_group_name: str + models: tuple[str, ...] + + class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable): team_member_budget_table: LiteLLM_BudgetTableFull | None = None # Resources inherited from access groups (separate from direct assignments) access_group_models: list[str] | None = None access_group_mcp_server_ids: list[str] | None = None access_group_agent_ids: list[str] | None = None + access_group_details: tuple[TeamAccessGroupModelGrant, ...] | None = None class TeamInfoResponseObject(TypedDict): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index acebcf40ecc..b012d25ee10 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -13,6 +13,7 @@ import asyncio import math import re import time +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast from fastapi import HTTPException, Request, status @@ -2834,7 +2835,7 @@ async def get_org_object( async def _get_resources_from_access_groups( - access_group_ids: list[str], + access_group_ids: Sequence[str], resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"], prisma_client: PrismaClient | None = None, user_api_key_cache: UserApiKeyCache | None = None, @@ -2893,7 +2894,7 @@ async def _get_resources_from_access_groups( async def _get_models_from_access_groups( - access_group_ids: list[str], + access_group_ids: Sequence[str], prisma_client: PrismaClient | None = None, user_api_key_cache: UserApiKeyCache | None = None, proxy_logging_obj: ProxyLogging | None = None, diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index a6e5eb2a0a0..ff9211742f3 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -1,6 +1,7 @@ # What is this? ## Common checks for /v1/models and `/model/info` import copy +from collections.abc import Sequence from typing import Any, Final import litellm @@ -178,8 +179,8 @@ def get_team_models( def get_complete_model_list( - key_models: list[str], - team_models: list[str], + key_models: Sequence[str], + team_models: Sequence[str], proxy_model_list: list[str], user_model: str | None, infer_model_from_keys: bool | None, @@ -203,7 +204,7 @@ def get_complete_model_list( def append_unique(models): for model in models: - if model not in unique_models: + if model not in unique_models and model != SpecialModelNames.no_default_models.value: unique_models.append(model) if key_models: diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index fe5a0e06d2e..0b99879f9fe 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -57,6 +57,7 @@ from litellm.proxy._types import ( SpecialManagementEndpointEnums, SpecialModelNames, SpecialProxyStrings, + TeamAccessGroupModelGrant, TeamAddMemberResponse, TeamInfoResponseObject, TeamInfoResponseObjectTeamTable, @@ -3829,7 +3830,7 @@ async def _add_team_member_budget_table( ) -> TeamInfoResponseObjectTeamTable: try: team_budget: Final = await _budget_db(prisma_client).find_unique(where={"budget_id": team_member_budget_id}) - team_info_response_object.team_member_budget_table = team_budget + return team_info_response_object.model_copy(update={"team_member_budget_table": team_budget}) except Exception: verbose_proxy_logger.info( "Team member budget table not found, passed team_member_budget_id=%s", team_member_budget_id @@ -3838,21 +3839,34 @@ async def _add_team_member_budget_table( return team_info_response_object -async def _resolve_team_access_group_resources(_team_info: TeamInfoResponseObjectTeamTable) -> None: - """Populate access_group_models / mcp_server_ids / agent_ids on the team - info response by resolving inherited resources from its access groups.""" +async def _resolve_team_access_group_resources( + _team_info: TeamInfoResponseObjectTeamTable, +) -> TeamInfoResponseObjectTeamTable: + """Return a copy of the team info with access_group_models / mcp_server_ids / + agent_ids / details resolved from its access groups.""" if not _team_info.access_group_ids: - return + return _team_info ag_lookup: Final = 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) + resolved_groups: Final = tuple( + ag_lookup[ag_id] for ag_id in dict.fromkeys(_team_info.access_group_ids) if ag_id in ag_lookup + ) + return _team_info.model_copy( + update={ + "access_group_models": list({m for group in resolved_groups for m in (group.access_model_names or [])}), + "access_group_mcp_server_ids": list( + {s for group in resolved_groups for s in (group.access_mcp_server_ids or [])} + ), + "access_group_agent_ids": list({a for group in resolved_groups for a in (group.access_agent_ids or [])}), + "access_group_details": tuple( + TeamAccessGroupModelGrant( + access_group_id=group.access_group_id, + access_group_name=group.access_group_name, + models=tuple(group.access_model_names or ()), + ) + for group in resolved_groups + ), + } + ) @router.get("/team/info", tags=["team management"], dependencies=[Depends(user_api_key_auth)]) @@ -3958,11 +3972,11 @@ async def team_info( ) # Resolve resources inherited from access groups - await _resolve_team_access_group_resources(_team_info) + resolved_team_info: Final = await _resolve_team_access_group_resources(_team_info) response_object: Final = TeamInfoResponseObject( team_id=team_id, - team_info=_team_info, + team_info=resolved_team_info, keys=keys, team_memberships=returned_tm, ) @@ -4391,32 +4405,21 @@ async def _build_team_list_where_conditions( async def _batch_resolve_access_group_resources( all_access_group_ids: list[str], -) -> dict[str, dict[str, list[str]]]: +) -> dict[str, LiteLLM_AccessGroupTable]: """ - 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. + Batch-fetch access groups in a single DB query and return them keyed by + access_group_id. 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: Final = list(set(all_access_group_ids)) + unique_ids: Final = tuple(frozenset(all_access_group_ids)) rows: Final = await _access_group_db(_prisma_client).find_many( where={"access_group_id": {"in": unique_ids}}, ) - - result: Final[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 + return {row.access_group_id: row for row in rows} def _convert_teams_to_response_models( @@ -4710,15 +4713,18 @@ async def list_team_v2( all_ag_ids: Final = [ag_id for t in team_items_with_ag for ag_id in (t.access_group_ids or [])] ag_lookup: Final = 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) + team_groups = tuple( + ag_lookup[ag_id] for ag_id in (team_item.access_group_ids or []) if ag_id in ag_lookup + ) + team_item.access_group_models = list( + {m for group in team_groups for m in (group.access_model_names or [])} + ) + team_item.access_group_mcp_server_ids = list( + {s for group in team_groups for s in (group.access_mcp_server_ids or [])} + ) + team_item.access_group_agent_ids = list( + {a for group in team_groups for a in (group.access_agent_ids or [])} + ) return { "teams": team_list, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 605455a2f73..e59c6adaf22 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -165,6 +165,7 @@ if TYPE_CHECKING: from prisma.client import TransactionManager from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction @@ -6419,6 +6420,74 @@ def construct_database_url_from_env_vars() -> str | None: return None +async def _get_validated_team_object( + user_api_key_dict: "UserAPIKeyAuth", + team_id: str, + prisma_client: "PrismaClient", + user_api_key_cache: "UserApiKeyCache", + proxy_logging_obj: "ProxyLogging", +) -> "LiteLLM_TeamTableCachedObj": + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.management_endpoints.team_endpoints import validate_membership + + team_object: Final = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_object) + return team_object + + +async def _get_team_object_for_access_groups( + team_id: str | None, + prisma_client: Optional["PrismaClient"], + user_api_key_cache: Optional["UserApiKeyCache"], + proxy_logging_obj: Optional["ProxyLogging"], +) -> Optional["LiteLLM_TeamTableCachedObj"]: + from litellm.proxy.auth.auth_checks import get_team_object + + if team_id is None or prisma_client is None or user_api_key_cache is None or proxy_logging_obj is None: + return None + try: + return await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except HTTPException: + verbose_proxy_logger.debug("Could not fetch team %s while listing models", team_id) + return None + + +async def _get_access_group_models( + user_api_key_dict: "UserAPIKeyAuth", + team_object: Optional["LiteLLM_TeamTableCachedObj"], + prisma_client: Optional["PrismaClient"], + user_api_key_cache: Optional["UserApiKeyCache"], + proxy_logging_obj: Optional["ProxyLogging"], +) -> tuple[str, ...]: + from litellm.proxy.auth.auth_checks import ( + _get_models_from_access_groups, + get_authorized_resources_from_key_access_groups, + ) + + team_group_models: Final = await _get_models_from_access_groups( + access_group_ids=(team_object.access_group_ids or ()) if team_object is not None else (), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + key_group_models: Final = await get_authorized_resources_from_key_access_groups( + valid_token=user_api_key_dict, + team_object=team_object, + resource_field="access_model_names", + ) + return tuple(dict.fromkeys((*team_group_models, *key_group_models))) + + async def get_available_models_for_user( user_api_key_dict: "UserAPIKeyAuth", llm_router: Optional["Router"], @@ -6450,13 +6519,11 @@ async def get_available_models_for_user( Returns: List of model names available to the user """ - from litellm.proxy.auth.auth_checks import get_team_object from litellm.proxy.auth.model_checks import ( get_complete_model_list, get_key_models, get_team_models, ) - from litellm.proxy.management_endpoints.team_endpoints import validate_membership # Get proxy model list and access groups if llm_router is None: @@ -6466,31 +6533,33 @@ async def get_available_models_for_user( proxy_model_list = llm_router.get_model_names() model_access_groups = llm_router.get_model_access_groups() - # Get key models - key_models = get_key_models( - user_api_key_dict=user_api_key_dict, - proxy_model_list=proxy_model_list, - model_access_groups=model_access_groups, - include_model_access_groups=include_model_access_groups, - ) - - # Get team models - team_models: list[str] = user_api_key_dict.team_models - - # If specific team_id is provided, validate and get team models - if team_id and prisma_client and proxy_logging_obj and user_api_key_cache: - key_models = [] - team_object: Final = await get_team_object( + requested_team_object: Final = ( + await _get_validated_team_object( + user_api_key_dict=user_api_key_dict, team_id=team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) - await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_object) - team_models = team_object.models + if team_id and prisma_client and proxy_logging_obj and user_api_key_cache + else None + ) - team_models = get_team_models( - team_models=team_models, + key_models: Final[Sequence[str]] = ( + () + if requested_team_object is not None + else get_key_models( + user_api_key_dict=user_api_key_dict, + proxy_model_list=proxy_model_list, + model_access_groups=model_access_groups, + include_model_access_groups=include_model_access_groups, + ) + ) + + team_models: Final = get_team_models( + team_models=( + requested_team_object.models if requested_team_object is not None else user_api_key_dict.team_models + ), proxy_model_list=proxy_model_list, model_access_groups=model_access_groups, include_model_access_groups=include_model_access_groups, @@ -6498,10 +6567,31 @@ async def get_available_models_for_user( effective_team_id: Final = team_id or user_api_key_dict.team_id + access_group_models: Final = ( + await _get_access_group_models( + user_api_key_dict=user_api_key_dict, + team_object=requested_team_object + or await _get_team_object_for_access_groups( + team_id=effective_team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if key_models or team_models + else () + ) + + granted_key_models: Final = (*key_models, *access_group_models) if key_models else key_models + granted_team_models: Final = (*team_models, *access_group_models) if team_models else team_models + # Get complete model list all_models: Final = get_complete_model_list( - key_models=key_models, - team_models=team_models, + key_models=granted_key_models, + team_models=granted_team_models, proxy_model_list=proxy_model_list, user_model=user_model, infer_model_from_keys=general_settings.get("infer_model_from_keys", False), diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index f56b8e113a9..e6c0eaee3c4 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -735,3 +735,28 @@ def test_add_known_models_refreshes_models_by_provider_for_wildcard_expansion(): litellm.vertex_language_models.discard(fake_model) litellm.add_known_models(model_cost_map={}) assert fake_model not in litellm.models_by_provider["vertex_ai"] + +def test_get_complete_model_list_drops_no_default_models_sentinel(): + from litellm.proxy.auth.model_checks import get_complete_model_list + + result = get_complete_model_list( + key_models=["no-default-models", "model-a"], + team_models=[], + proxy_model_list=["model-a", "model-b"], + user_model=None, + infer_model_from_keys=False, + ) + assert result == ["model-a"] + + +def test_get_complete_model_list_sentinel_only_grants_nothing(): + from litellm.proxy.auth.model_checks import get_complete_model_list + + result = get_complete_model_list( + key_models=["no-default-models"], + team_models=["no-default-models"], + proxy_model_list=["model-a", "model-b"], + user_model=None, + infer_model_from_keys=False, + ) + assert result == [] 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 eedffa1ea5f..a1cbc77b7a5 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -8637,9 +8637,9 @@ class TestBatchResolveAccessGroupResources: 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"] + assert sorted(result["ag-1"].access_model_names) == ["claude-3", "gpt-4"] + assert result["ag-1"].access_mcp_server_ids == ["mcp-1"] + assert sorted(result["ag-1"].access_agent_ids) == ["agent-1", "agent-2"] @pytest.mark.asyncio async def test_multiple_access_groups(self): @@ -8668,8 +8668,8 @@ class TestBatchResolveAccessGroupResources: 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"] + assert result["ag-1"].access_model_names == ["gpt-4"] + assert result["ag-2"].access_model_names == ["gemini"] @pytest.mark.asyncio async def test_missing_access_group_omitted(self): @@ -8735,6 +8735,75 @@ class TestBatchResolveAccessGroupResources: assert "ag-1" in result +class TestResolveTeamAccessGroupResources: + """Tests for the per-team access group resolution on /team/info.""" + + @pytest.mark.asyncio + async def test_populates_flat_lists_and_per_group_details(self): + """access_group_details must attribute each model to the group granting it, + so the UI can show provenance on hover; flat lists stay for back-compat. + Duplicated ids must collapse to one entry (response amplification), and the + input object must stay untouched (resolution returns a copy).""" + from litellm.proxy._types import TeamInfoResponseObjectTeamTable + from litellm.proxy.management_endpoints.team_endpoints import ( + _resolve_team_access_group_resources, + ) + + row1 = MagicMock() + row1.access_group_id = "ag-1" + row1.access_group_name = "shared-models" + row1.access_model_names = ["gpt-4", "claude-3"] + row1.access_mcp_server_ids = ["mcp-1"] + row1.access_agent_ids = [] + + row2 = MagicMock() + row2.access_group_id = "ag-2" + row2.access_group_name = "extra-models" + row2.access_model_names = ["claude-3", "gemini"] + row2.access_mcp_server_ids = [] + row2.access_agent_ids = ["agent-1"] + + fake_prisma = MagicMock() + fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock( + return_value=[row1, row2] + ) + + team_info = TeamInfoResponseObjectTeamTable( + team_id="team-1", access_group_ids=["ag-1", "ag-2", "ag-1", "ag-missing"] + ) + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): + resolved = await _resolve_team_access_group_resources(team_info) + + assert team_info.access_group_details is None + assert sorted(resolved.access_group_models or []) == [ + "claude-3", + "gemini", + "gpt-4", + ] + assert resolved.access_group_mcp_server_ids == ["mcp-1"] + assert resolved.access_group_agent_ids == ["agent-1"] + assert [ + (d.access_group_id, d.access_group_name, d.models) + for d in (resolved.access_group_details or []) + ] == [ + ("ag-1", "shared-models", ("gpt-4", "claude-3")), + ("ag-2", "extra-models", ("claude-3", "gemini")), + ] + + @pytest.mark.asyncio + async def test_no_access_groups_leaves_details_unset(self): + from litellm.proxy._types import TeamInfoResponseObjectTeamTable + from litellm.proxy.management_endpoints.team_endpoints import ( + _resolve_team_access_group_resources, + ) + + team_info = TeamInfoResponseObjectTeamTable(team_id="team-1", access_group_ids=[]) + resolved = await _resolve_team_access_group_resources(team_info) + + assert resolved.access_group_details is None + assert resolved.access_group_models is None + + @pytest.mark.asyncio async def test_verify_team_access_denies_unauthorized_user(): """ diff --git a/tests/test_litellm/proxy/utils/helpers/test_model_access.py b/tests/test_litellm/proxy/utils/helpers/test_model_access.py index 59268e1427b..5fb4392eec6 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_model_access.py +++ b/tests/test_litellm/proxy/utils/helpers/test_model_access.py @@ -9,6 +9,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ( create_model_info_response, get_available_models_for_user, + hash_token, is_known_model, is_known_vector_store_index, model_dump_with_preserved_fields, @@ -404,3 +405,119 @@ async def test_get_available_models_for_user_error_path_complete_list_raises( general_settings={}, user_model=None, ) + + +@pytest.mark.asyncio +async def test_get_available_models_for_user_resolves_team_access_group_models( + monkeypatch, +): + from litellm.models.access_group import LiteLLM_AccessGroupTable + from litellm.models.team import LiteLLM_TeamTableCachedObj + + team = LiteLLM_TeamTableCachedObj( + team_id="team-1", + models=["no-default-models"], + access_group_ids=["ag-1"], + ) + access_group = LiteLLM_AccessGroupTable( + access_group_id="ag-1", + access_group_name="repro-group", + access_model_names=["model-a", "model-b"], + assigned_team_ids=["team-1"], + ) + + async def _get_team_object(**_kwargs): + return team + + async def _get_access_object(**_kwargs): + return access_group + + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object) + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_access_object", _get_access_object) + + result = await get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-test-key", + user_id="user-1", + team_id="team-1", + models=["all-team-models"], + team_models=["no-default-models"], + ), + llm_router=_router_with_models(["model-a", "model-b", "model-c"]), + general_settings={}, + user_model=None, + prisma_client=MagicMock(), + proxy_logging_obj=MagicMock(), + user_api_key_cache=MagicMock(), + ) + assert sorted(result) == ["model-a", "model-b"] + + +@pytest.mark.asyncio +async def test_get_available_models_for_user_without_access_groups_grants_nothing( + monkeypatch, +): + from litellm.models.team import LiteLLM_TeamTableCachedObj + + async def _get_team_object(**_kwargs): + return LiteLLM_TeamTableCachedObj(team_id="team-1", models=["no-default-models"]) + + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object) + + result = await get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-test-key", + user_id="user-1", + team_id="team-1", + models=["all-team-models"], + team_models=["no-default-models"], + ), + llm_router=_router_with_models(["model-a", "model-b"]), + general_settings={}, + user_model=None, + prisma_client=MagicMock(), + proxy_logging_obj=MagicMock(), + user_api_key_cache=MagicMock(), + ) + assert result == [] + +@pytest.mark.asyncio +async def test_get_available_models_for_user_resolves_key_access_group_models( + monkeypatch, +): + from litellm.models.access_group import LiteLLM_AccessGroupTable + from litellm.models.team import LiteLLM_TeamTableCachedObj + + async def _get_team_object(**_kwargs): + return LiteLLM_TeamTableCachedObj(team_id="team-1", models=["no-default-models"]) + + async def _get_access_object(**_kwargs): + return LiteLLM_AccessGroupTable( + access_group_id="ag-1", + access_group_name="key-group", + access_model_names=["model-b"], + assigned_key_ids=[hash_token("sk-test-key")], + ) + + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object) + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_access_object", _get_access_object) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + + result = await get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-test-key", + user_id="user-1", + team_id="team-1", + models=["no-default-models"], + team_models=["no-default-models"], + access_group_ids=["ag-1"], + ), + llm_router=_router_with_models(["model-a", "model-b"]), + general_settings={}, + user_model=None, + prisma_client=MagicMock(), + proxy_logging_obj=MagicMock(), + user_api_key_cache=MagicMock(), + ) + assert result == ["model-b"] diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 991f8eaa934..0a0cfe9a617 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23245 + "limit": 23235 }, "LIT002": { - "limit": 27179 + "limit": 27176 }, "LIT003": { "limit": 269 @@ -30,6 +30,6 @@ "limit": 16769 }, "LIT011": { - "limit": 5602 + "limit": 5598 } } diff --git a/ui/litellm-dashboard/src/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx index 2bda0f72cec..e8331294972 100644 --- a/ui/litellm-dashboard/src/components/Teams.test.tsx +++ b/ui/litellm-dashboard/src/components/Teams.test.tsx @@ -613,6 +613,34 @@ describe("Teams - access_group_ids in team create", () => { ); }); }); + + it("creates a team with no models selected, sending the no-default-models sentinel instead of an empty list", async () => { + renderWithQueryClient(); + + const createButton = screen.getAllByRole("button", { name: /create team/i })[0]; + act(() => { + fireEvent.click(createButton); + }); + + await waitFor(() => { + expect(screen.getByLabelText(/team name/i)).toBeInTheDocument(); + }); + + fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "Group Only Team" } }); + + const createTeamSubmitButtons = screen.getAllByRole("button", { name: /create team/i }); + fireEvent.click(createTeamSubmitButtons[createTeamSubmitButtons.length - 1]); + + await waitFor(() => { + expect(teamCreateCall).toHaveBeenCalledWith( + "test-token", + expect.objectContaining({ + team_alias: "Group Only Team", + models: ["no-default-models"], + }), + ); + }); + }); }); describe("Teams - metadata key-value pairs in team create", () => { diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index 95642b93019..42ff8bcc4c2 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -42,6 +42,7 @@ interface TeamProps { import DeleteResourceModal from "./common_components/DeleteResourceModal"; import { teamCreateCall } from "./networking"; +import { normalizeTeamModelSelection } from "./team/teamModelAccess"; import { ModelSelect } from "./ModelSelect/ModelSelect"; const canCreateOrManageTeams = ( @@ -351,7 +352,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser } } - await teamCreateCall(accessToken, formValues); + await teamCreateCall(accessToken, { ...formValues, models: normalizeTeamModelSelection(formValues.models) }); NotificationsManager.success("Team created"); await refreshTeams(); form.resetFields(); @@ -618,17 +619,11 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser label={ Models{" "} - + } - rules={[ - { - required: true, - message: "Please select at least one model", - }, - ]} name="models" > = new Set([ "disable_global_guardrails", ]); +const TEAM_MODEL_BADGE_COLORS: Record = { + "all-proxy": "red", + "no-default": "gray", + direct: "blue", + "access-group": "green", +}; + export interface TeamMembership { user_id: string; team_id: string; @@ -132,6 +145,7 @@ export interface TeamData { access_group_models?: string[]; access_group_mcp_server_ids?: string[]; access_group_agent_ids?: string[]; + access_group_details?: TeamAccessGroupModelGrant[]; router_settings?: Record; guardrails?: string[]; policies?: string[]; @@ -483,7 +497,7 @@ const TeamInfoView: React.FC = ({ const updateData: any = { team_id: teamId, team_alias: values.team_alias, - models: values.models, + models: normalizeTeamModelSelection(values.models), tpm_limit: sanitizeNumeric(values.tpm_limit), rpm_limit: sanitizeNumeric(values.rpm_limit), model_tpm_limit: modelTpmLimit, @@ -764,21 +778,14 @@ const TeamInfoView: React.FC = ({ Models
- {info.models.length === 0 || info.models.includes("all-proxy-models") ? ( - All proxy models - ) : ( - <> - {info.models.map((model: string, index: number) => ( - - {model} - - ))} - {(info.access_group_models || []).map((model: string, index: number) => ( - - {model} - - ))} - + {computeTeamModelBadges(info.models, info.access_group_models || [], info.access_group_details).map( + (badge, index) => ( + + + {badge.label} + + + ), )}
@@ -982,7 +989,7 @@ const TeamInfoView: React.FC = ({ { + it("substitutes the no-default-models sentinel for an empty selection", () => { + expect(normalizeTeamModelSelection([])).toEqual(["no-default-models"]); + expect(normalizeTeamModelSelection(undefined)).toEqual(["no-default-models"]); + }); + + it("passes a non-empty selection through untouched", () => { + expect(normalizeTeamModelSelection(["gpt-4o-mini"])).toEqual(["gpt-4o-mini"]); + expect(normalizeTeamModelSelection(["all-proxy-models"])).toEqual(["all-proxy-models"]); + }); +}); + +describe("computeTeamModelBadges", () => { + it("attributes group-only models to the groups granting them", () => { + const badges = computeTeamModelBadges(["sonnet-direct"], [], GRANTS); + expect(badges).toEqual([ + { + label: "sonnet-direct", + kind: "direct", + tooltip: "Granted directly in the team's model list", + }, + { label: "haiku", kind: "access-group", tooltip: "Granted via access groups shared, extra" }, + { label: "gpt-4o-mini", kind: "access-group", tooltip: "Granted via access group shared" }, + { label: "sonnet", kind: "access-group", tooltip: "Granted via access group extra" }, + ]); + }); + + it("marks a model both direct and group-granted on the direct badge, without a duplicate badge", () => { + const badges = computeTeamModelBadges(["haiku"], [], GRANTS); + expect(badges).toEqual([ + { + label: "haiku", + kind: "direct", + tooltip: "Granted directly in the team's model list, and also via access groups shared, extra", + }, + { label: "gpt-4o-mini", kind: "access-group", tooltip: "Granted via access group shared" }, + { label: "sonnet", kind: "access-group", tooltip: "Granted via access group extra" }, + ]); + }); + + it("shows the no-default-models sentinel as its own badge and keeps group badges visible", () => { + const badges = computeTeamModelBadges(["no-default-models"], [], [GRANTS[0]]); + expect(badges.map((b) => [b.label, b.kind])).toEqual([ + ["No default models", "no-default"], + ["haiku", "access-group"], + ["gpt-4o-mini", "access-group"], + ]); + }); + + it("still shows group badges when the empty model list grants everything", () => { + const badges = computeTeamModelBadges([], [], [GRANTS[0]]); + expect(badges[0]).toEqual({ + label: "All proxy models", + kind: "all-proxy", + tooltip: "The team's model list is empty, so it can access every model on the proxy", + }); + expect(badges.slice(1).map((b) => b.label)).toEqual(["haiku", "gpt-4o-mini"]); + }); + + it("distinguishes the all-proxy-models sentinel from an empty list in the tooltip", () => { + const badges = computeTeamModelBadges(["all-proxy-models"], [], []); + expect(badges).toEqual([ + { + label: "All proxy models", + kind: "all-proxy", + tooltip: "Granted by the All Proxy Models entry in the team's model list", + }, + ]); + }); + + it("falls back to the flat access_group_models list when per-group details are absent", () => { + const badges = computeTeamModelBadges(["direct-model"], ["haiku"], undefined); + expect(badges).toEqual([ + { label: "direct-model", kind: "direct", tooltip: "Granted directly in the team's model list" }, + { label: "haiku", kind: "access-group", tooltip: "Granted via an access group" }, + ]); + }); +}); diff --git a/ui/litellm-dashboard/src/components/team/teamModelAccess.ts b/ui/litellm-dashboard/src/components/team/teamModelAccess.ts new file mode 100644 index 00000000000..91ddf0f4045 --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/teamModelAccess.ts @@ -0,0 +1,82 @@ +export const ALL_PROXY_MODELS = "all-proxy-models"; +export const NO_DEFAULT_MODELS = "no-default-models"; + +export interface TeamAccessGroupModelGrant { + access_group_id: string; + access_group_name: string; + models: string[]; +} + +export type TeamModelBadgeKind = "all-proxy" | "no-default" | "direct" | "access-group"; + +export interface TeamModelBadge { + label: string; + kind: TeamModelBadgeKind; + tooltip: string; +} + +export function normalizeTeamModelSelection(models: string[] | undefined): string[] { + return models && models.length > 0 ? models : [NO_DEFAULT_MODELS]; +} + +const describeGroups = (names: string[]): string => + names.length > 1 ? `access groups ${names.join(", ")}` : `access group ${names[0]}`; + +export function computeTeamModelBadges( + models: string[], + accessGroupModels: string[], + accessGroupDetails: TeamAccessGroupModelGrant[] | undefined, +): TeamModelBadge[] { + const grants = accessGroupDetails ?? []; + const groupNamesFor = (model: string): string[] => + grants.filter((g) => g.models.includes(model)).map((g) => g.access_group_name); + const viaGroups = (model: string): string => { + const names = groupNamesFor(model); + return names.length > 0 ? describeGroups(names) : "an access group"; + }; + + const allProxy = models.length === 0 || models.includes(ALL_PROXY_MODELS); + const directModels = allProxy ? [] : models.filter((m) => m !== NO_DEFAULT_MODELS); + const groupModels = [...new Set(grants.length > 0 ? grants.flatMap((g) => g.models) : accessGroupModels)].filter( + (m) => !directModels.includes(m), + ); + + const allProxyBadge: TeamModelBadge = { + label: "All proxy models", + kind: "all-proxy", + tooltip: models.includes(ALL_PROXY_MODELS) + ? "Granted by the All Proxy Models entry in the team's model list" + : "The team's model list is empty, so it can access every model on the proxy", + }; + const noDefaultBadge: TeamModelBadge = { + label: "No default models", + kind: "no-default", + tooltip: "No models are granted directly. Access comes only from access groups", + }; + const headBadge = (): TeamModelBadge[] => { + if (allProxy) return [allProxyBadge]; + if (models.includes(NO_DEFAULT_MODELS)) return [noDefaultBadge]; + return []; + }; + + return [ + ...headBadge(), + ...directModels.map( + (m): TeamModelBadge => ({ + label: m, + kind: "direct", + tooltip: + groupNamesFor(m).length > 0 + ? `Granted directly in the team's model list, and also via ${viaGroups(m)}` + : "Granted directly in the team's model list", + }), + ), + ...groupModels.map( + (m): TeamModelBadge => ({ + label: m, + kind: "access-group", + tooltip: `Granted via ${viaGroups(m)}`, + }), + ), + ]; +}