diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 32f76dd1fe8..bf85245abd4 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -56,6 +56,10 @@ 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 ( + prepare_key_models_response_payload, + resolve_key_models_for_display, +) from litellm.proxy.management_endpoints.model_management_endpoints import ( _add_model_to_db, ) @@ -2807,14 +2811,25 @@ async def info_key_fn( async def key_resolved_models_fn( key_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + search: Optional[str] = Query( + default=None, + description="Case-insensitive substring filter on resolved model names. Omit or blank for no filter.", + ), + compact: bool = Query( + default=False, + description="When true, returns section structure and counts but omits per-section model name lists.", + ), ): """ - + Return resolved models for a virtual key as `model_display_sections` for the admin UI. + + - **source**: How models were resolved (`all-team-models`, `all-proxy-models`, `no-default-models`). + - **model_display_sections**: Ordered sections (`all_proxy_models`, `all_team_models`, `access_group`, `ungrouped`); a model may appear under multiple access-group sections. + - **search**: Filters the resolved list before truncation and sectioning. + - **compact**: Metadata-only payload (empty `models` arrays in each section) for fast initial load. + - **all_team_models_without_team**: True when the key uses `all-team-models` but has no `team_id` (assign a team in key settings). """ - from litellm.proxy.proxy_server import ( - llm_router, - prisma_client, - ) + from litellm.proxy.proxy_server import llm_router, prisma_client try: if prisma_client is None: @@ -2850,33 +2865,32 @@ async def key_resolved_models_fn( ) key_models = list[str](key_info.models or []) - all_models: List[str] = [] - + model_access_groups: Dict[str, List[str]] = {} if llm_router is not None: - all_models = llm_router.get_model_names() + model_access_groups = llm_router.get_model_access_groups() - source: str = SpecialModelNames.no_default_models.value - resolved: List[str] = key_models + resolved, source, all_team_models_without_team = ( + await resolve_key_models_for_display( + key_models=key_models, + team_id=key_info.team_id, + prisma_client=prisma_client, + llm_router=llm_router, + ) + ) - #Team Models - if (SpecialModelNames.all_team_models.value in key_models): - if key_info.team_id is not None: - source = SpecialModelNames.all_team_models.value - team_row = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": key_info.team_id}, - ) - if team_row is not None and team_row.models is not None: - resolved = list[str](team_row.models) - else: - source = SpecialModelNames.all_team_models.value - resolved = all_models - - #Proxy 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 = all_models + all_router_model_names: List[str] = [] + if llm_router is not None: + all_router_model_names = list(llm_router.get_model_names()) - return {"models": resolved, "source": source} + return prepare_key_models_response_payload( + resolved=resolved, + source=source, + all_team_models_without_team=all_team_models_without_team, + model_access_groups=model_access_groups, + search=search, + compact=compact, + all_router_model_names=all_router_model_names, + ) except Exception as e: raise handle_exception_on_proxy(e) diff --git a/litellm/proxy/management_endpoints/key_resolved_models_helpers.py b/litellm/proxy/management_endpoints/key_resolved_models_helpers.py new file mode 100644 index 00000000000..8f4031cbd38 --- /dev/null +++ b/litellm/proxy/management_endpoints/key_resolved_models_helpers.py @@ -0,0 +1,271 @@ +""" +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 new file mode 100644 index 00000000000..90e84c9c672 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_key_resolved_models_helpers.py @@ -0,0 +1,206 @@ +""" +Unit tests for key_resolved_models_helpers. +""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy._types import SpecialModelNames +from litellm.proxy.management_endpoints.key_resolved_models_helpers import ( + KEY_RESOLVED_MODELS_DISPLAY_LIMIT, + prepare_key_models_response_payload, + resolve_key_models_for_display, + _filter_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"] + + +@pytest.mark.asyncio +async def test_resolve_all_team_models_without_team(): + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + mock_router = MagicMock() + mock_router.get_model_names.return_value = ["a", "b"] + + resolved, source, no_team = await resolve_key_models_for_display( + key_models=["all-team-models"], + team_id=None, + prisma_client=mock_prisma, + llm_router=mock_router, + ) + assert source == SpecialModelNames.all_team_models.value + assert resolved == ["a", "b"] + assert no_team is True + + +@pytest.mark.asyncio +async def test_resolve_all_team_models_with_team(): + mock_prisma = MagicMock() + team_row = MagicMock() + team_row.models = ["team-m1"] + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_router = MagicMock() + mock_router.get_model_names.return_value = ["proxy-wide"] + + resolved, source, no_team = await resolve_key_models_for_display( + key_models=["all-team-models"], + team_id="team-1", + prisma_client=mock_prisma, + llm_router=mock_router, + ) + assert source == SpecialModelNames.all_team_models.value + assert resolved == ["team-m1"] + assert no_team is False + + +@pytest.mark.asyncio +async def test_resolve_all_proxy_models(): + mock_prisma = MagicMock() + mock_router = MagicMock() + mock_router.get_model_names.return_value = ["m1", "m2"] + + resolved, source, no_team = await resolve_key_models_for_display( + key_models=["all-proxy-models"], + team_id=None, + prisma_client=mock_prisma, + llm_router=mock_router, + ) + assert source == SpecialModelNames.all_proxy_models.value + assert resolved == ["m1", "m2"] + assert no_team is False + + +def test_prepare_payload_search_and_truncation(): + big = [f"m{i}" for i in range(KEY_RESOLVED_MODELS_DISPLAY_LIMIT + 50)] + out = prepare_key_models_response_payload( + resolved=big, + source=SpecialModelNames.no_default_models.value, + all_team_models_without_team=False, + model_access_groups={}, + search=None, + compact=False, + all_router_model_names=big, + ) + assert out["resolved_total_count"] == len(big) + assert out["matched_count"] == len(big) + assert out["models_truncated"] is True + ung = out["model_display_sections"][0] + assert ung["section_kind"] == "ungrouped" + assert len(ung["models"]) == KEY_RESOLVED_MODELS_DISPLAY_LIMIT + + +def test_prepare_payload_compact(): + out = prepare_key_models_response_payload( + resolved=["x", "y"], + source=SpecialModelNames.no_default_models.value, + all_team_models_without_team=False, + model_access_groups={}, + search=None, + compact=True, + all_router_model_names=["x", "y"], + ) + assert out["matched_count"] == 2 + for sec in out["model_display_sections"]: + assert sec["models"] == [] + + +def test_prepare_payload_explicit_models_with_access_groups(): + out = prepare_key_models_response_payload( + resolved=["a", "b", "z"], + source=SpecialModelNames.no_default_models.value, + all_team_models_without_team=False, + model_access_groups={"G": ["a", "b"], "H": ["b"]}, + search=None, + compact=False, + all_router_model_names=["a", "b", "z"], + ) + kinds = [s["section_kind"] for s in out["model_display_sections"]] + assert "access_group" in kinds + assert "ungrouped" in kinds + ung = next(s for s in out["model_display_sections"] if s["section_kind"] == "ungrouped") + assert ung["models"] == ["z"] + + +def test_all_proxy_models_includes_access_group_sections(): + """Sentinel 'all proxy' section is first; access groups are not skipped.""" + out = prepare_key_models_response_payload( + resolved=["m1", "m2", "m3"], + source=SpecialModelNames.all_proxy_models.value, + all_team_models_without_team=False, + model_access_groups={ + "grp-a": ["m1", "not-in-resolved"], + "grp-b": ["m2", "m1"], + }, + search=None, + compact=False, + all_router_model_names=["m1", "m2", "m3"], + ) + sections = out["model_display_sections"] + kinds = [s["section_kind"] for s in sections] + assert kinds[0] == "all_proxy_models" + assert sections[0]["models"] == ["m1", "m2", "m3"] + assert kinds.count("access_group") == 2 + by_title = {s["title"]: s["models"] for s in sections} + assert by_title["grp-a"] == ["m1"] + assert by_title["grp-b"] == ["m2", "m1"] + assert "ungrouped" not in kinds + + +def test_all_team_models_includes_access_group_sections(): + out = prepare_key_models_response_payload( + resolved=["x", "y"], + source=SpecialModelNames.all_team_models.value, + all_team_models_without_team=False, + model_access_groups={"T1": ["x"], "T2": ["x", "y"]}, + search=None, + compact=False, + all_router_model_names=["x", "y"], + ) + sections = out["model_display_sections"] + kinds = [s["section_kind"] for s in sections] + assert kinds[0] == "all_team_models" + assert sections[0]["models"] == ["x", "y"] + assert kinds.count("access_group") == 2 + assert "ungrouped" not in kinds + + +def test_same_model_listed_under_multiple_access_groups(): + """Overlap across groups is allowed; model 'b' appears in G and H.""" + out = prepare_key_models_response_payload( + resolved=["a", "b", "c"], + source=SpecialModelNames.no_default_models.value, + all_team_models_without_team=False, + model_access_groups={"G": ["a", "b"], "H": ["b", "c"]}, + search=None, + compact=False, + all_router_model_names=["a", "b", "c"], + ) + g = next(s for s in out["model_display_sections"] if s["title"] == "G") + h = next(s for s in out["model_display_sections"] if s["title"] == "H") + assert "b" in g["models"] and "b" in h["models"] + + +def test_access_group_wildcard_expands_to_concrete_models(): + router_models = ["openai/gpt-4o", "openai/gpt-3.5-turbo", "anthropic/claude-3"] + out = prepare_key_models_response_payload( + resolved=["openai/*"], + source=SpecialModelNames.no_default_models.value, + all_team_models_without_team=False, + model_access_groups={"openai_group": ["openai/*"]}, + search=None, + compact=False, + all_router_model_names=router_models, + ) + sec = next(s for s in out["model_display_sections"] if s["title"] == "openai_group") + assert "openai/*" not in sec["models"] + assert "openai/gpt-4o" in sec["models"] + assert "openai/gpt-3.5-turbo" in sec["models"] + assert "anthropic/claude-3" not in sec["models"] diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/KeyModelList.test.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/KeyModelList.test.tsx new file mode 100644 index 00000000000..908ff04c800 --- /dev/null +++ b/ui/litellm-dashboard/src/components/key_team_helpers/KeyModelList.test.tsx @@ -0,0 +1,208 @@ +/* @vitest-environment jsdom */ +import React from "react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { render, screen } from "@testing-library/react"; +import KeyModelList from "./KeyModelList"; +import * as useGetKeyModelsHook from "@/hooks/keys/useGetKeyModels"; + +vi.mock("@/hooks/keys/useGetKeyModels", () => ({ + useGetKeyModels: vi.fn(), +})); + +const mockUseGetKeyModels = vi.mocked(useGetKeyModelsHook.useGetKeyModels); + +describe("KeyModelList", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should render grouped models by default without requiring search", () => { + mockUseGetKeyModels.mockReturnValue({ + searchInput: "", + setSearchInput: vi.fn(), + defaultModelsQuery: { + data: { + model_display_sections: [ + { title: "All proxy models", section_kind: "all_proxy_models", models: ["m1", "m2"] }, + { title: "grp-a", section_kind: "access_group", models: ["m1"] }, + ], + source: "all-proxy-models", + resolved_total_count: 2, + matched_count: 2, + models_truncated: false, + all_team_models_without_team: false, + }, + isLoading: false, + isError: false, + isSuccess: true, + } as any, + searchQuery: { data: undefined, isLoading: false, isFetching: false, isFetched: false } as any, + hasActiveSearch: false, + isInitialLoading: false, + searchInputLoading: false, + }); + + render(); + expect(screen.getByText("Models")).toBeInTheDocument(); + expect(screen.getByText("All proxy models")).toBeInTheDocument(); + expect(screen.getByText("access_group")).toBeInTheDocument(); + expect(screen.getByText("grp-a")).toBeInTheDocument(); + expect(screen.getAllByText("m1").length).toBeGreaterThanOrEqual(1); + const innerCards = document.querySelectorAll(".ant-card-type-inner"); + expect(innerCards.length).toBeGreaterThanOrEqual(2); + }); + + it("should show explanation icon next to All team models when no team is assigned", () => { + mockUseGetKeyModels.mockReturnValue({ + searchInput: "", + setSearchInput: vi.fn(), + defaultModelsQuery: { + data: { + model_display_sections: [ + { title: "All team models", section_kind: "all_team_models", models: ["a"] }, + ], + source: "all-team-models", + resolved_total_count: 1, + matched_count: 1, + models_truncated: false, + all_team_models_without_team: true, + }, + isLoading: false, + isError: false, + isSuccess: true, + } as any, + searchQuery: { data: undefined, isLoading: false, isFetching: false, isFetched: false } as any, + hasActiveSearch: false, + isInitialLoading: false, + searchInputLoading: false, + }); + + render(); + expect(screen.getByLabelText("Why this matters")).toBeInTheDocument(); + expect(screen.getByText("All team models")).toBeInTheDocument(); + }); + + it("should use a scroll-capped list region with overflow-y-auto", () => { + mockUseGetKeyModels.mockReturnValue({ + searchInput: "", + setSearchInput: vi.fn(), + defaultModelsQuery: { + data: { + model_display_sections: [ + { title: "Other models", section_kind: "ungrouped", models: ["x"] }, + ], + source: "no-default-models", + resolved_total_count: 1, + matched_count: 1, + models_truncated: false, + all_team_models_without_team: false, + }, + isLoading: false, + isError: false, + isSuccess: true, + } as any, + searchQuery: { data: undefined, isLoading: false, isFetching: false, isFetched: false } as any, + hasActiveSearch: false, + isInitialLoading: false, + searchInputLoading: false, + }); + + render(); + const scroll = screen.getByTestId("key-model-list-scroll"); + expect(scroll.className).toMatch(/overflow-y-auto/); + expect(scroll.className).toMatch(/min-h-0/); + }); + + it("should show search results when hasActiveSearch", () => { + mockUseGetKeyModels.mockReturnValue({ + searchInput: "gpt", + setSearchInput: vi.fn(), + defaultModelsQuery: { data: undefined, isLoading: false, isError: false, isSuccess: true } as any, + searchQuery: { + data: { + model_display_sections: [ + { title: "Other models", section_kind: "ungrouped", models: ["gpt-4"] }, + ], + source: "no-default-models", + resolved_total_count: 1, + matched_count: 1, + models_truncated: false, + all_team_models_without_team: false, + }, + isLoading: false, + isFetching: false, + isFetched: true, + isSuccess: true, + } as any, + hasActiveSearch: true, + isInitialLoading: false, + searchInputLoading: false, + }); + + render(); + expect(screen.getByText("gpt-4")).toBeInTheDocument(); + }); + + it("should show nothing found when search succeeds with zero matches", () => { + mockUseGetKeyModels.mockReturnValue({ + searchInput: "zzz", + setSearchInput: vi.fn(), + defaultModelsQuery: { data: undefined, isLoading: false, isError: false, isSuccess: true } as any, + searchQuery: { + data: { + model_display_sections: [], + source: "no-default-models", + resolved_total_count: 5, + matched_count: 0, + models_truncated: false, + all_team_models_without_team: false, + }, + isLoading: false, + isFetching: false, + isFetched: true, + isSuccess: true, + } as any, + hasActiveSearch: true, + isInitialLoading: false, + searchInputLoading: false, + }); + + render(); + expect(screen.getByText(/Nothing found for "zzz"/i)).toBeInTheDocument(); + }); + + it("should render all proxy scope and access group sections together", () => { + mockUseGetKeyModels.mockReturnValue({ + searchInput: "m", + setSearchInput: vi.fn(), + defaultModelsQuery: { data: undefined, isLoading: false, isError: false, isSuccess: true } as any, + searchQuery: { + data: { + model_display_sections: [ + { title: "All proxy models", section_kind: "all_proxy_models", models: ["m1", "m2"] }, + { title: "grp-a", section_kind: "access_group", models: ["m1"] }, + { title: "grp-b", section_kind: "access_group", models: ["m2", "m1"] }, + ], + source: "all-proxy-models", + resolved_total_count: 2, + matched_count: 2, + models_truncated: false, + all_team_models_without_team: false, + }, + isLoading: false, + isFetching: false, + isFetched: true, + isSuccess: true, + } as any, + hasActiveSearch: true, + isInitialLoading: false, + searchInputLoading: false, + }); + + render(); + expect(screen.getByText("All proxy models")).toBeInTheDocument(); + expect(screen.getAllByText("access_group").length).toBeGreaterThanOrEqual(2); + const m1Tags = screen.getAllByText("m1"); + expect(m1Tags.length).toBeGreaterThanOrEqual(2); + }); +}); diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/KeyModelList.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/KeyModelList.tsx index ed1e34e9c36..3a8e95b6847 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/KeyModelList.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/KeyModelList.tsx @@ -1,36 +1,192 @@ -import React from "react"; -import {UseGetKeyModels} from "@/hooks/keys/useGetKeyModels"; -import { Card, Tag } from 'antd'; +import React from 'react'; +import { QuestionCircleOutlined } from '@ant-design/icons'; +import { Alert, Card, Empty, Input, Tag, Tooltip, Typography } from 'antd'; +import { useGetKeyModels } from '@/hooks/keys/useGetKeyModels'; +import type { KeyModelDisplaySection } from '@/components/networking'; + +const { Text } = Typography; + +/** Card body (search + list) max height — list area scrolls inside flex layout. */ +const CARD_BODY_MAX_HEIGHT = 'min(72vh, 580px)'; interface KeyModelListProps { - key_id: string; + key_id: string; } -const extractDefaultTags = (source: string) => { - if (source === 'all-proxy-models') { - return All proxy models - } else if (source === 'all-team-models') { - return All team models - } - return '' -} +const INNER_CARD_CLASS = 'mb-3 last:mb-0'; +const SectionBlock: React.FC<{ + section: KeyModelDisplaySection; + warnNoTeam: boolean; +}> = ({ section, warnNoTeam }) => { + const tagRow = ( +
+ {section.models.map((item, index) => ( + {item} + ))} +
+ ); + if (section.section_kind === 'access_group') { + return ( + + access_group + + {section.title} + + + } + > + {tagRow} + + ); + } + + if (section.section_kind === 'all_proxy_models' || section.section_kind === 'all_team_models') { + const showNoTeamHint = section.section_kind === 'all_team_models' && warnNoTeam; + const title = ( + + {section.title} + {showNoTeamHint ? ( + + + + ) : null} + + ); + return ( + + {tagRow} + + ); + } + + return ( + {section.title}} + > + {tagRow} + + ); +}; const KeyModelList: React.FC = ({ key_id }) => { - const { data: keyModels, isLoading} = UseGetKeyModels(key_id) + const { + searchInput, + setSearchInput, + defaultModelsQuery, + searchQuery, + hasActiveSearch, + isInitialLoading, + searchInputLoading, + } = useGetKeyModels(key_id); - + const defaultData = defaultModelsQuery.data; + const searchData = searchQuery.data; + const warnNoTeam = + (searchData?.all_team_models_without_team ?? defaultData?.all_team_models_without_team) === true; - const title = keyModels ? <>Model {extractDefaultTags(keyModels.source)} : 'Model' - return ( - -
- {keyModels && keyModels.models.map((item:string)=> {return {item}})} -
- - ); + const sections = hasActiveSearch ? searchData?.model_display_sections : defaultData?.model_display_sections; + const truncated = hasActiveSearch ? searchData?.models_truncated : defaultData?.models_truncated; + const nothingFound = + hasActiveSearch && + searchQuery.isFetched && + !searchQuery.isFetching && + searchData !== undefined && + searchData.matched_count === 0; + + const renderBody = () => { + if (defaultModelsQuery.isError) { + return ; + } + + if (nothingFound) { + return ; + } + + if (searchQuery.isError && hasActiveSearch) { + return ; + } + + if (sections && sections.length > 0) { + return ( + <> + {truncated ? ( + + ) : null} + {sections.map((section) => ( + + ))} + + ); + } + + if (!isInitialLoading && defaultModelsQuery.isSuccess && (!sections || sections.length === 0)) { + return ; + } + + return null; + }; + + return ( + +
+ setSearchInput(e.target.value)} + onSearch={(v) => setSearchInput(v)} + /> +
+
+ {renderBody()} +
+ + ); }; export default KeyModelList; - diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts index cb0673484ce..6c2166d9367 100644 --- a/ui/litellm-dashboard/src/components/networking.test.ts +++ b/ui/litellm-dashboard/src/components/networking.test.ts @@ -406,7 +406,7 @@ describe("individualModelHealthCheckCall", () => { }); -describe("fetching models from key information", () => { +describe("fetchKeyModelCall", () => { const originalFetch = global.fetch; beforeEach(() => { @@ -417,23 +417,32 @@ describe("fetching models from key information", () => { global.fetch = originalFetch; }); - it("should call /health with model_id query param so health checks run by deployment id", async () => { + it("should request /key/{id}/models with optional compact and search query params", async () => { const mockFetch = vi.fn().mockResolvedValue({ ok: true, json: vi.fn().mockResolvedValue({ - healthy_count: 1, - unhealthy_count: 0, - healthy_endpoints: [], - unhealthy_endpoints: [], + model_display_sections: [], + source: "no-default-models", + resolved_total_count: 0, + matched_count: 0, + models_truncated: false, + all_team_models_without_team: false, }), } as any); global.fetch = mockFetch as any; - await Networking.fetchKeyModelCall("token-123", "key-abc-456"); - - expect(mockFetch).toHaveBeenCalledOnce(); - const [url] = mockFetch.mock.calls[0]; - const urlStr = typeof url === "string" ? url : (url as Request).url; + await Networking.fetchKeyModelCall("token-123", "key-abc-456", { compact: true }); + let [url] = mockFetch.mock.calls[0]; + let urlStr = typeof url === "string" ? url : (url as Request).url; expect(urlStr).toContain("key-abc-456"); + expect(urlStr).toContain("compact=true"); + + vi.clearAllMocks(); + global.fetch = mockFetch as any; + await Networking.fetchKeyModelCall("token-123", "key-abc-456", { search: "gpt-4" }); + [url] = mockFetch.mock.calls[0]; + urlStr = typeof url === "string" ? url : (url as Request).url; + const parsed = new URL(urlStr, "http://example.com"); + expect(parsed.searchParams.get("search")).toBe("gpt-4"); }); }); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index a10967ec3bb..f33a87559dc 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -3308,18 +3308,48 @@ export const keyAliasesCall = async ( } }; -export interface keyModelResponse { - source: string; +export type KeyModelDisplaySectionKind = + | "all_proxy_models" + | "all_team_models" + | "access_group" + | "ungrouped"; + +export interface KeyModelDisplaySection { + title: string; + section_kind: KeyModelDisplaySectionKind; models: string[]; } +export interface KeyModelResponse { + model_display_sections: KeyModelDisplaySection[]; + source: string; + resolved_total_count: number; + matched_count: number; + models_truncated: boolean; + all_team_models_without_team: boolean; +} + +export type FetchKeyModelCallOptions = { + search?: string; + compact?: boolean; +}; + export const fetchKeyModelCall = async ( accessToken: string, - key_id: string -): Promise => { - + key_id: string, + options?: FetchKeyModelCallOptions +): Promise => { try { - let url = proxyBaseUrl ? `${proxyBaseUrl}/key/${key_id}/models` : `/key/${key_id}/models`; + const params = new URLSearchParams(); + if (options?.search !== undefined && options.search.trim() !== "") { + params.set("search", options.search.trim()); + } + if (options?.compact === true) { + params.set("compact", "true"); + } + const qs = params.toString(); + const base = proxyBaseUrl ? `${proxyBaseUrl}/key/${key_id}/models` : `/key/${key_id}/models`; + let url = qs ? `${base}?${qs}` : base; const response = await fetch(url, { method: "GET", diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 1c9f4ac3187..6fe3cea901a 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -369,7 +369,7 @@ export default function KeyInfoView({ }; return ( -
+
({ @@ -16,8 +16,20 @@ vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ })), })); +const emptyKeyModelResponse = { + model_display_sections: [], + source: "no-default-models", + resolved_total_count: 0, + matched_count: 0, + models_truncated: false, + all_team_models_without_team: false, +}; + const createQueryClient = () => new QueryClient({ + defaultOptions: { + queries: { retry: false }, + }, }); const wrapper = ({ children }: { children: React.ReactNode }) => { @@ -25,27 +37,31 @@ const wrapper = ({ children }: { children: React.ReactNode }) => { return React.createElement(QueryClientProvider, { client: queryClient }, children); }; -const mockAccessToken = 'test-key-id'; -const mockAccessGroups = ["group-1", "group-2", "group-3"]; - describe("useGetKeyModels", () => { beforeEach(async () => { vi.clearAllMocks(); const useAuthorizedModule = await import("@/app/(dashboard)/hooks/useAuthorized"); vi.mocked(useAuthorizedModule.default).mockReturnValue({ - accessToken: mockAccessToken, + accessToken: "test-token-456", } as any); }); - it("should return hook result without errors", () => { - vi.mocked(networking.fetchKeyModelCall).mockResolvedValue({source: '', models: []}); + it("should load default full model list without compact", async () => { + vi.mocked(networking.fetchKeyModelCall).mockResolvedValue({ + ...emptyKeyModelResponse, + resolved_total_count: 3, + model_display_sections: [ + { title: "Other models", section_kind: "ungrouped", models: ["a", "b", "c"] }, + ], + }); - const { result } = renderHook(() => UseGetKeyModels('test-key-id'), { wrapper }); + const { result } = renderHook(() => useGetKeyModels("test-key-id"), { wrapper }); - expect(result.current).toBeDefined(); - expect(result.current).toHaveProperty("data"); - expect(result.current).toHaveProperty("isSuccess"); - expect(result.current).toHaveProperty("isError"); - expect(result.current).toHaveProperty("status"); + expect(result.current).toHaveProperty("defaultModelsQuery"); + expect(result.current).toHaveProperty("searchQuery"); + + await waitFor(() => expect(result.current.defaultModelsQuery.isSuccess).toBe(true)); + expect(networking.fetchKeyModelCall).toHaveBeenCalledWith("test-token-456", "test-key-id"); + expect(result.current.defaultModelsQuery.data?.resolved_total_count).toBe(3); }); -}); \ No newline at end of file +}); diff --git a/ui/litellm-dashboard/src/hooks/keys/useGetKeyModels.ts b/ui/litellm-dashboard/src/hooks/keys/useGetKeyModels.ts index 684861a86ac..2a2c2754e27 100644 --- a/ui/litellm-dashboard/src/hooks/keys/useGetKeyModels.ts +++ b/ui/litellm-dashboard/src/hooks/keys/useGetKeyModels.ts @@ -1,15 +1,61 @@ import { useQuery } from '@tanstack/react-query'; +import { useDebouncedState } from '@tanstack/react-pacer/debouncer'; +import { useCallback, useState } from 'react'; import { fetchKeyModelCall } from '@/components/networking'; import useAuthorized from '@/app/(dashboard)/hooks/useAuthorized'; -export const UseGetKeyModels = (key_id: string) => { +/** Wait after last keystroke before calling the search API (avoids spzzy refetches). */ +const SEARCH_DEBOUNCE_MS = 450; + +export const useGetKeyModels = (key_id: string) => { const { accessToken } = useAuthorized(); - - return useQuery({ - queryKey: ['keyModels', key_id], - queryFn: () => { - if (!accessToken) throw new Error("Access Token required"); - return fetchKeyModelCall(accessToken, key_id); - }, + const [searchInput, setSearchInputState] = useState(''); + const [debouncedSearch, setDebouncedSearch] = useDebouncedState('', { + wait: SEARCH_DEBOUNCE_MS, }); + + const setSearchInput = useCallback( + (value: string) => { + setSearchInputState(value); + setDebouncedSearch(value); + }, + [setDebouncedSearch], + ); + + const trimmedDebounced = debouncedSearch.trim(); + const hasActiveSearch = trimmedDebounced.length > 0; + /** True while the user is still typing and the debounced value has not caught up yet. */ + const isSearchDebouncing = searchInput.trim() !== trimmedDebounced; + + const defaultModelsQuery = useQuery({ + queryKey: ['keyModelsDefault', key_id], + queryFn: () => { + if (!accessToken) throw new Error('Access Token required'); + return fetchKeyModelCall(accessToken, key_id); + }, + enabled: Boolean(accessToken && key_id), + }); + + const searchQuery = useQuery({ + queryKey: ['keyModelsSearch', key_id, trimmedDebounced], + queryFn: () => { + if (!accessToken) throw new Error('Access Token required'); + return fetchKeyModelCall(accessToken, key_id, { search: trimmedDebounced }); + }, + enabled: Boolean(accessToken && key_id && hasActiveSearch), + }); + + const isSearchFetching = hasActiveSearch && searchQuery.isFetching; + const searchInputLoading = isSearchDebouncing || isSearchFetching; + + return { + searchInput, + setSearchInput, + debouncedSearch: trimmedDebounced, + defaultModelsQuery, + searchQuery, + hasActiveSearch, + isInitialLoading: defaultModelsQuery.isLoading, + searchInputLoading, + }; };