diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index bf76f99db69..9f1c91faae7 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -1,10 +1,16 @@ # What is this? ## Common checks for /v1/models and `/model/info` -from typing import Dict, List, Optional, Set +from typing import Any, Dict, List, Optional, Set, Tuple, TypedDict import litellm from litellm._logging import verbose_proxy_logger from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import ( + _is_wildcard_pattern, + _model_matches_any_wildcard_pattern_in_list, + is_model_allowed_by_pattern, +) +from litellm.proxy.utils import PrismaClient from litellm.router import Router from litellm.router_utils.fallback_event_handlers import get_fallback_model_group from litellm.types.router import LiteLLM_Params @@ -379,3 +385,261 @@ def get_all_fallbacks( except Exception as e: verbose_proxy_logger.error(f"Error getting fallbacks for model {model}: {e}") return [] + + +# --- GET /key/{key_id}/models (admin UI): resolve lists and sectioned payload --- + +KEY_RESOLVED_MODELS_DISPLAY_LIMIT = 500 + +_KEY_SECTION_ALL_PROXY = "all_proxy_models" +_KEY_SECTION_ALL_TEAM = "all_team_models" +_KEY_SECTION_ACCESS_GROUP = "access_group" +_KEY_SECTION_UNGROUPED = "ungrouped" + +_KEY_TITLE_ALL_PROXY = "All proxy models" +_KEY_TITLE_ALL_TEAM = "All team models" +_KEY_TITLE_UNGROUPED = "Other models" + + +class KeyResolvedModelDisplaySection(TypedDict): + title: str + section_kind: str + models: List[str] + + +def _filter_key_models_by_search(models: List[str], search: Optional[str]) -> List[str]: + if not search: + return list(models) + needle = search.strip().lower() + if not needle: + return list(models) + return [m for m in models if needle in m.lower()] + + +async def resolve_key_models_for_display( + *, + key_models: List[str], + team_id: Optional[str], + prisma_client: PrismaClient, + llm_router: Optional[Router], +) -> Tuple[List[str], str, bool]: + """ + Returns (resolved_model_names, source, all_team_models_without_team). + """ + all_models: List[str] = [] + if llm_router is not None: + all_models = list(llm_router.get_model_names()) + + source: str = SpecialModelNames.no_default_models.value + resolved: List[str] = list(key_models) + all_team_models_without_team = False + + if SpecialModelNames.all_team_models.value in key_models: + if team_id is not None: + source = SpecialModelNames.all_team_models.value + team_row = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id}, + ) + if team_row is not None and team_row.models is not None: + resolved = list(team_row.models) + else: + source = SpecialModelNames.all_team_models.value + all_team_models_without_team = True + resolved = list(all_models) + + if ( + SpecialModelNames.all_proxy_models.value in key_models + or SpecialModelNames.all_proxy_models.value in resolved + ): + source = SpecialModelNames.all_proxy_models.value + resolved = list(all_models) + + return resolved, source, all_team_models_without_team + + +def _concrete_models_allowed_by_resolved( + resolved: List[str], router_model_names: List[str] +) -> List[str]: + """Concrete router model names the key may call, given resolved patterns (incl. wildcards).""" + resolved_list = list(resolved) + resolved_set = set(resolved) + out: List[str] = [] + seen: set[str] = set() + for m in router_model_names: + if m in seen: + continue + if m in resolved_set: + out.append(m) + seen.add(m) + elif _model_matches_any_wildcard_pattern_in_list(m, resolved_list): + out.append(m) + seen.add(m) + return out + + +def _expand_group_models_for_key_display( + group_models: List[str], + concrete_pool: List[str], +) -> List[str]: + """ + Expand wildcard entries in a group's model list to concrete names from the pool. + Non-wildcard entries are included when present in the pool. + """ + pool_set = set(concrete_pool) + out: List[str] = [] + seen: set[str] = set() + for g in group_models: + if _is_wildcard_pattern(g): + for m in concrete_pool: + if m in seen: + continue + if is_model_allowed_by_pattern(m, g): + out.append(m) + seen.add(m) + else: + if g in seen: + continue + if g in pool_set: + out.append(g) + seen.add(g) + return out + + +def _models_in_any_access_group_for_key_display( + model_access_groups: Dict[str, List[str]], candidate_models: List[str] +) -> set: + """Set of candidate_models that appear in at least one access group list (concrete names).""" + cset = set(candidate_models) + in_group: set = set() + for models_in_group in model_access_groups.values(): + for m in models_in_group: + if m in cset: + in_group.add(m) + return in_group + + +def build_key_resolved_model_display_sections( + *, + display_models: List[str], + source: str, + model_access_groups: Dict[str, List[str]], + compact: bool, +) -> List[KeyResolvedModelDisplaySection]: + """ + Build ordered sections for the admin UI. + + When source is all-proxy or all-team, a scope section lists all display_models first, + then every intersecting **model_access_groups** section is still included (models may + repeat across the sentinel section and each group). Ungrouped is omitted for sentinels + only to avoid repeating the full flat list as "Other models". + """ + sections: List[KeyResolvedModelDisplaySection] = [] + empty_models: List[str] = [] if compact else [] + + display_set = set(display_models) + + if source == SpecialModelNames.all_proxy_models.value: + sections.append( + KeyResolvedModelDisplaySection( + title=_KEY_TITLE_ALL_PROXY, + section_kind=_KEY_SECTION_ALL_PROXY, + models=empty_models if compact else list(display_models), + ) + ) + elif source == SpecialModelNames.all_team_models.value: + sections.append( + KeyResolvedModelDisplaySection( + title=_KEY_TITLE_ALL_TEAM, + section_kind=_KEY_SECTION_ALL_TEAM, + models=empty_models if compact else list(display_models), + ) + ) + + for group_name, group_models in model_access_groups.items(): + intersected = [m for m in group_models if m in display_set] + if not intersected: + continue + sections.append( + KeyResolvedModelDisplaySection( + title=group_name, + section_kind=_KEY_SECTION_ACCESS_GROUP, + models=empty_models if compact else intersected, + ) + ) + + skip_ungrouped = source in ( + SpecialModelNames.all_proxy_models.value, + SpecialModelNames.all_team_models.value, + ) + if not skip_ungrouped: + in_any = _models_in_any_access_group_for_key_display( + model_access_groups, display_models + ) + ungrouped = [m for m in display_models if m not in in_any] + if ungrouped: + sections.append( + KeyResolvedModelDisplaySection( + title=_KEY_TITLE_UNGROUPED, + section_kind=_KEY_SECTION_UNGROUPED, + models=empty_models if compact else ungrouped, + ) + ) + + return sections + + +def prepare_key_models_response_payload( + *, + resolved: List[str], + source: str, + all_team_models_without_team: bool, + model_access_groups: Dict[str, List[str]], + search: Optional[str], + compact: bool, + all_router_model_names: List[str], +) -> Dict[str, Any]: + """ + Apply search, truncation, and build the JSON-serializable response dict for GET /key/{id}/models. + + Access-group sections list concrete model names: wildcards in router group metadata + are expanded against models allowed for this key (resolved ∩ router names). + """ + base_concrete = _concrete_models_allowed_by_resolved(resolved, all_router_model_names) + + expanded_groups: Dict[str, List[str]] = {} + for group_name, group_models in model_access_groups.items(): + expanded = _expand_group_models_for_key_display(group_models, base_concrete) + if expanded: + expanded_groups[group_name] = expanded + + filtered_concrete = _filter_key_models_by_search(base_concrete, search) + matched_count = len(filtered_concrete) + models_truncated = matched_count > KEY_RESOLVED_MODELS_DISPLAY_LIMIT + display_models = ( + filtered_concrete[:KEY_RESOLVED_MODELS_DISPLAY_LIMIT] + if models_truncated + else filtered_concrete + ) + + display_set = set(display_models) + intersected_groups: Dict[str, List[str]] = {} + for group_name, models in expanded_groups.items(): + in_slice = [m for m in models if m in display_set] + if in_slice: + intersected_groups[group_name] = in_slice + + model_display_sections = build_key_resolved_model_display_sections( + display_models=display_models, + source=source, + model_access_groups=intersected_groups, + compact=compact, + ) + + return { + "model_display_sections": model_display_sections, + "source": source, + "resolved_total_count": len(resolved), + "matched_count": matched_count, + "models_truncated": models_truncated, + "all_team_models_without_team": all_team_models_without_team, + } diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 5c2d19b357e..e1e2e834115 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -56,7 +56,7 @@ from litellm.proxy.management_endpoints.common_utils import ( _is_user_team_admin, _set_object_metadata_field, ) -from litellm.proxy.management_endpoints.key_resolved_models_helpers import ( +from litellm.proxy.auth.model_checks import ( prepare_key_models_response_payload, resolve_key_models_for_display, ) diff --git a/litellm/proxy/management_endpoints/key_resolved_models_helpers.py b/litellm/proxy/management_endpoints/key_resolved_models_helpers.py deleted file mode 100644 index 8f4031cbd38..00000000000 --- a/litellm/proxy/management_endpoints/key_resolved_models_helpers.py +++ /dev/null @@ -1,271 +0,0 @@ -""" -Helpers for GET /key/{key_id}/models: resolve key model lists and build sectioned UI payloads. -""" - -from __future__ import annotations - -from typing import Any, Dict, List, Optional, Tuple, TypedDict - -from litellm.proxy._types import SpecialModelNames -from litellm.proxy.auth.auth_checks import ( - _is_wildcard_pattern, - _model_matches_any_wildcard_pattern_in_list, - is_model_allowed_by_pattern, -) -from litellm.proxy.utils import PrismaClient -from litellm.router import Router - -KEY_RESOLVED_MODELS_DISPLAY_LIMIT = 500 - -SECTION_ALL_PROXY = "all_proxy_models" -SECTION_ALL_TEAM = "all_team_models" -SECTION_ACCESS_GROUP = "access_group" -SECTION_UNGROUPED = "ungrouped" - -TITLE_ALL_PROXY = "All proxy models" -TITLE_ALL_TEAM = "All team models" -TITLE_UNGROUPED = "Other models" - - -class ModelDisplaySection(TypedDict): - title: str - section_kind: str - models: List[str] - - -def _filter_models_by_search(models: List[str], search: Optional[str]) -> List[str]: - if not search: - return list(models) - needle = search.strip().lower() - if not needle: - return list(models) - return [m for m in models if needle in m.lower()] - - -async def resolve_key_models_for_display( - *, - key_models: List[str], - team_id: Optional[str], - prisma_client: PrismaClient, - llm_router: Optional[Router], -) -> Tuple[List[str], str, bool]: - """ - Returns (resolved_model_names, source, all_team_models_without_team). - """ - all_models: List[str] = [] - if llm_router is not None: - all_models = list(llm_router.get_model_names()) - - source: str = SpecialModelNames.no_default_models.value - resolved: List[str] = list(key_models) - all_team_models_without_team = False - - if SpecialModelNames.all_team_models.value in key_models: - if team_id is not None: - source = SpecialModelNames.all_team_models.value - team_row = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id}, - ) - if team_row is not None and team_row.models is not None: - resolved = list(team_row.models) - else: - source = SpecialModelNames.all_team_models.value - all_team_models_without_team = True - resolved = list(all_models) - - if ( - SpecialModelNames.all_proxy_models.value in key_models - or SpecialModelNames.all_proxy_models.value in resolved - ): - source = SpecialModelNames.all_proxy_models.value - resolved = list(all_models) - - return resolved, source, all_team_models_without_team - - -def _concrete_models_allowed_by_resolved( - resolved: List[str], router_model_names: List[str] -) -> List[str]: - """Concrete router model names the key may call, given resolved patterns (incl. wildcards).""" - resolved_list = list(resolved) - resolved_set = set(resolved) - out: List[str] = [] - seen: set[str] = set() - for m in router_model_names: - if m in seen: - continue - if m in resolved_set: - out.append(m) - seen.add(m) - elif _model_matches_any_wildcard_pattern_in_list(m, resolved_list): - out.append(m) - seen.add(m) - return out - - -def _expand_group_models_for_display( - group_models: List[str], - concrete_pool: List[str], -) -> List[str]: - """ - Expand wildcard entries in a group's model list to concrete names from the pool. - Non-wildcard entries are included when present in the pool. - """ - pool_set = set(concrete_pool) - out: List[str] = [] - seen: set[str] = set() - for g in group_models: - if _is_wildcard_pattern(g): - for m in concrete_pool: - if m in seen: - continue - if is_model_allowed_by_pattern(m, g): - out.append(m) - seen.add(m) - else: - if g in seen: - continue - if g in pool_set: - out.append(g) - seen.add(g) - return out - - -def _models_in_any_access_group( - model_access_groups: Dict[str, List[str]], candidate_models: List[str] -) -> set: - """Set of candidate_models that appear in at least one access group list (concrete names).""" - cset = set(candidate_models) - in_group: set = set() - for models_in_group in model_access_groups.values(): - for m in models_in_group: - if m in cset: - in_group.add(m) - return in_group - - -def build_model_display_sections( - *, - display_models: List[str], - source: str, - model_access_groups: Dict[str, List[str]], - compact: bool, -) -> List[ModelDisplaySection]: - """ - Build ordered sections for the admin UI. - - When source is all-proxy or all-team, a scope section lists all display_models first, - then every intersecting **model_access_groups** section is still included (models may - repeat across the sentinel section and each group). Ungrouped is omitted for sentinels - only to avoid repeating the full flat list as "Other models". - """ - sections: List[ModelDisplaySection] = [] - empty_models: List[str] = [] if compact else [] - - display_set = set(display_models) - - if source == SpecialModelNames.all_proxy_models.value: - sections.append( - ModelDisplaySection( - title=TITLE_ALL_PROXY, - section_kind=SECTION_ALL_PROXY, - models=empty_models if compact else list(display_models), - ) - ) - elif source == SpecialModelNames.all_team_models.value: - sections.append( - ModelDisplaySection( - title=TITLE_ALL_TEAM, - section_kind=SECTION_ALL_TEAM, - models=empty_models if compact else list(display_models), - ) - ) - - for group_name, group_models in model_access_groups.items(): - # group_models are models in this group that appear in resolved; further restrict to display slice - intersected = [m for m in group_models if m in display_set] - if not intersected: - continue - sections.append( - ModelDisplaySection( - title=group_name, - section_kind=SECTION_ACCESS_GROUP, - models=empty_models if compact else intersected, - ) - ) - - # Sentinel scope sections already list the full display set; skip ungrouped to avoid duplicating it. - skip_ungrouped = source in ( - SpecialModelNames.all_proxy_models.value, - SpecialModelNames.all_team_models.value, - ) - if not skip_ungrouped: - in_any = _models_in_any_access_group(model_access_groups, display_models) - ungrouped = [m for m in display_models if m not in in_any] - if ungrouped: - sections.append( - ModelDisplaySection( - title=TITLE_UNGROUPED, - section_kind=SECTION_UNGROUPED, - models=empty_models if compact else ungrouped, - ) - ) - - return sections - - -def prepare_key_models_response_payload( - *, - resolved: List[str], - source: str, - all_team_models_without_team: bool, - model_access_groups: Dict[str, List[str]], - search: Optional[str], - compact: bool, - all_router_model_names: List[str], -) -> Dict[str, Any]: - """ - Apply search, truncation, and build the JSON-serializable response dict. - - Access-group sections list concrete model names: wildcards in router group metadata - are expanded against models allowed for this key (resolved ∩ router names). - """ - base_concrete = _concrete_models_allowed_by_resolved(resolved, all_router_model_names) - - expanded_groups: Dict[str, List[str]] = {} - for group_name, group_models in model_access_groups.items(): - expanded = _expand_group_models_for_display(group_models, base_concrete) - if expanded: - expanded_groups[group_name] = expanded - - filtered_concrete = _filter_models_by_search(base_concrete, search) - matched_count = len(filtered_concrete) - models_truncated = matched_count > KEY_RESOLVED_MODELS_DISPLAY_LIMIT - display_models = ( - filtered_concrete[:KEY_RESOLVED_MODELS_DISPLAY_LIMIT] - if models_truncated - else filtered_concrete - ) - - display_set = set(display_models) - intersected_groups: Dict[str, List[str]] = {} - for group_name, models in expanded_groups.items(): - in_slice = [m for m in models if m in display_set] - if in_slice: - intersected_groups[group_name] = in_slice - - model_display_sections = build_model_display_sections( - display_models=display_models, - source=source, - model_access_groups=intersected_groups, - compact=compact, - ) - - return { - "model_display_sections": model_display_sections, - "source": source, - "resolved_total_count": len(resolved), - "matched_count": matched_count, - "models_truncated": models_truncated, - "all_team_models_without_team": all_team_models_without_team, - } diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_resolved_models_helpers.py b/tests/test_litellm/proxy/management_endpoints/test_key_resolved_models_helpers.py index 90e84c9c672..8822ec8a0b3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_resolved_models_helpers.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_resolved_models_helpers.py @@ -1,5 +1,5 @@ """ -Unit tests for key_resolved_models_helpers. +Unit tests for GET /key/{id}/models payload helpers (litellm.proxy.auth.model_checks). """ from unittest.mock import AsyncMock, MagicMock @@ -7,20 +7,20 @@ from unittest.mock import AsyncMock, MagicMock import pytest from litellm.proxy._types import SpecialModelNames -from litellm.proxy.management_endpoints.key_resolved_models_helpers import ( +from litellm.proxy.auth.model_checks import ( KEY_RESOLVED_MODELS_DISPLAY_LIMIT, prepare_key_models_response_payload, resolve_key_models_for_display, - _filter_models_by_search, + _filter_key_models_by_search, ) def test_filter_models_by_search(): models = ["GPT-4", "claude-3", "embed-small"] - assert _filter_models_by_search(models, None) == models - assert _filter_models_by_search(models, " ") == models - assert _filter_models_by_search(models, "gpt") == ["GPT-4"] - assert _filter_models_by_search(models, "CLAUDE") == ["claude-3"] + assert _filter_key_models_by_search(models, None) == models + assert _filter_key_models_by_search(models, " ") == models + assert _filter_key_models_by_search(models, "gpt") == ["GPT-4"] + assert _filter_key_models_by_search(models, "CLAUDE") == ["claude-3"] @pytest.mark.asyncio