From f81aefd48a73057896d17e2cfdae678e5e70717c Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 15 Jan 2026 21:41:09 -0800 Subject: [PATCH] Team Settings Model Select --- .../ModelSelect/ModelSelect.test.tsx | 243 ++++++++++++++++-- .../components/ModelSelect/ModelSelect.tsx | 212 +++++++++------ .../components/ModelSelect/modelUtils.test.ts | 34 +-- .../src/components/ModelSelect/modelUtils.ts | 14 +- .../src/components/organizations.tsx | 6 +- .../src/components/team/team_info.test.tsx | 135 ++++++++-- .../src/components/team/team_info.tsx | 64 ++--- .../src/components/view_users/types.ts | 1 + 8 files changed, 505 insertions(+), 204 deletions(-) diff --git a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.test.tsx b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.test.tsx index 1c4bae557a..9a56d219ba 100644 --- a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.test.tsx +++ b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.test.tsx @@ -18,6 +18,10 @@ vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ useOrganization: vi.fn(), })); +vi.mock("@/app/(dashboard)/hooks/users/useCurrentUser", () => ({ + useCurrentUser: vi.fn(), +})); + vi.mock("antd", async (importOriginal) => { const actual = await importOriginal(); return { @@ -71,10 +75,12 @@ vi.mock("antd", async (importOriginal) => { import { useAllProxyModels } from "@/app/(dashboard)/hooks/models/useModels"; import { useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import { useTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser"; const mockUseAllProxyModels = vi.mocked(useAllProxyModels); const mockUseTeam = vi.mocked(useTeam); const mockUseOrganization = vi.mocked(useOrganization); +const mockUseCurrentUser = vi.mocked(useCurrentUser); describe("ModelSelect", () => { const mockProxyModels: ProxyModel[] = [ @@ -100,10 +106,16 @@ describe("ModelSelect", () => { data: undefined, isLoading: false, } as any); + mockUseCurrentUser.mockReturnValue({ + data: { models: [] }, + isLoading: false, + } as any); }); it("should render", async () => { - renderWithProviders(); + renderWithProviders( + , + ); await waitFor(() => { expect(screen.getByTestId("model-select")).toBeInTheDocument(); @@ -116,7 +128,7 @@ describe("ModelSelect", () => { isLoading: true, } as any); - renderWithProviders(); + renderWithProviders(); expect(screen.getByTestId("skeleton-input")).toBeInTheDocument(); expect(screen.queryByTestId("model-select")).not.toBeInTheDocument(); @@ -128,7 +140,7 @@ describe("ModelSelect", () => { isLoading: true, } as any); - renderWithProviders(); + renderWithProviders(); expect(screen.getByTestId("skeleton-input")).toBeInTheDocument(); }); @@ -139,13 +151,49 @@ describe("ModelSelect", () => { isLoading: true, } as any); - renderWithProviders(); + renderWithProviders(); + + expect(screen.getByTestId("skeleton-input")).toBeInTheDocument(); + }); + + it("should show skeleton loader when current user is loading", () => { + mockUseCurrentUser.mockReturnValue({ + data: undefined, + isLoading: true, + } as any); + + renderWithProviders(); expect(screen.getByTestId("skeleton-input")).toBeInTheDocument(); }); it("should render special options group", async () => { - renderWithProviders(); + const mockOrganization: Organization = { + organization_id: "org-1", + organization_alias: "Test Org", + budget_id: "budget-1", + metadata: {}, + models: ["all-proxy-models"], + spend: 0, + model_spend: {}, + created_at: "2024-01-01", + created_by: "user-1", + updated_at: "2024-01-01", + updated_by: "user-1", + litellm_budget_table: null, + teams: null, + users: null, + members: null, + }; + + mockUseOrganization.mockReturnValue({ + data: mockOrganization, + isLoading: false, + } as any); + + renderWithProviders( + , + ); await waitFor(() => { const select = screen.getByTestId("model-select"); @@ -156,7 +204,9 @@ describe("ModelSelect", () => { }); it("should render wildcard options group", async () => { - renderWithProviders(); + renderWithProviders( + , + ); await waitFor(() => { expect(screen.getByText("All Openai models")).toBeInTheDocument(); @@ -165,7 +215,9 @@ describe("ModelSelect", () => { }); it("should render regular models group", async () => { - renderWithProviders(); + renderWithProviders( + , + ); await waitFor(() => { expect(screen.getByText("gpt-4")).toBeInTheDocument(); @@ -175,7 +227,9 @@ describe("ModelSelect", () => { it("should call onChange when selecting a regular model", async () => { const user = userEvent.setup(); - renderWithProviders(); + renderWithProviders( + , + ); await waitFor(() => { expect(screen.getByTestId("model-select")).toBeInTheDocument(); @@ -189,7 +243,37 @@ describe("ModelSelect", () => { it("should call onChange with only last special option when multiple special options are selected", async () => { const user = userEvent.setup(); - renderWithProviders(); + const mockOrganization: Organization = { + organization_id: "org-1", + organization_alias: "Test Org", + budget_id: "budget-1", + metadata: {}, + models: ["all-proxy-models"], + spend: 0, + model_spend: {}, + created_at: "2024-01-01", + created_by: "user-1", + updated_at: "2024-01-01", + updated_by: "user-1", + litellm_budget_table: null, + teams: null, + users: null, + members: null, + }; + + mockUseOrganization.mockReturnValue({ + data: mockOrganization, + isLoading: false, + } as any); + + renderWithProviders( + , + ); await waitFor(() => { expect(screen.getByTestId("model-select")).toBeInTheDocument(); @@ -203,7 +287,12 @@ describe("ModelSelect", () => { it("should disable regular models when special option is selected", async () => { renderWithProviders( - , + , ); await waitFor(() => { @@ -214,7 +303,12 @@ describe("ModelSelect", () => { it("should disable wildcard models when special option is selected", async () => { renderWithProviders( - , + , ); await waitFor(() => { @@ -224,8 +318,37 @@ describe("ModelSelect", () => { }); it("should disable other special options when one special option is selected", async () => { + const mockOrganization: Organization = { + organization_id: "org-1", + organization_alias: "Test Org", + budget_id: "budget-1", + metadata: {}, + models: ["all-proxy-models"], + spend: 0, + model_spend: {}, + created_at: "2024-01-01", + created_by: "user-1", + updated_at: "2024-01-01", + updated_by: "user-1", + litellm_budget_table: null, + teams: null, + users: null, + members: null, + }; + + mockUseOrganization.mockReturnValue({ + data: mockOrganization, + isLoading: false, + } as any); + renderWithProviders( - , + , ); await waitFor(() => { @@ -235,7 +358,9 @@ describe("ModelSelect", () => { }); it("should filter models when showAllProxyModelsOverride is true", async () => { - renderWithProviders(); + renderWithProviders( + , + ); await waitFor(() => { expect(screen.getByText("gpt-4")).toBeInTheDocument(); @@ -267,7 +392,9 @@ describe("ModelSelect", () => { isLoading: false, } as any); - renderWithProviders(); + renderWithProviders( + , + ); await waitFor(() => { expect(screen.getByText("gpt-4")).toBeInTheDocument(); @@ -275,7 +402,7 @@ describe("ModelSelect", () => { }); }); - it("should return empty models array when organization does not have all-proxy-models", async () => { + it("should filter models when organization has specific models", async () => { const mockOrganization: Organization = { organization_id: "org-1", organization_alias: "Test Org", @@ -299,17 +426,24 @@ describe("ModelSelect", () => { isLoading: false, } as any); - renderWithProviders(); + renderWithProviders( + , + ); await waitFor(() => { - expect(screen.queryByText("gpt-4")).not.toBeInTheDocument(); + expect(screen.getByText("gpt-4")).toBeInTheDocument(); expect(screen.queryByText("claude-3")).not.toBeInTheDocument(); }); }); it("should use custom dataTestId when provided", async () => { renderWithProviders( - , + , ); await waitFor(() => { @@ -319,7 +453,9 @@ describe("ModelSelect", () => { it("should handle multiple model selections", async () => { const user = userEvent.setup(); - renderWithProviders(); + renderWithProviders( + , + ); await waitFor(() => { expect(screen.getByTestId("model-select")).toBeInTheDocument(); @@ -337,7 +473,9 @@ describe("ModelSelect", () => { }); it("should capitalize provider name in wildcard options", async () => { - renderWithProviders(); + renderWithProviders( + , + ); await waitFor(() => { expect(screen.getByText("All Openai models")).toBeInTheDocument(); @@ -356,11 +494,74 @@ describe("ModelSelect", () => { isLoading: false, } as any); - renderWithProviders(); + renderWithProviders( + , + ); await waitFor(() => { const gpt4Options = screen.getAllByText("gpt-4"); expect(gpt4Options.length).toBeGreaterThan(0); }); }); + + it("should filter models based on user context with includeUserModels option", async () => { + mockUseCurrentUser.mockReturnValue({ + data: { models: ["gpt-4"] }, + isLoading: false, + } as any); + + renderWithProviders( + , + ); + + await waitFor(() => { + expect(screen.getByText("gpt-4")).toBeInTheDocument(); + expect(screen.queryByText("claude-3")).not.toBeInTheDocument(); + }); + }); + + it("should filter models based on team context", async () => { + const mockTeam = { + team_id: "team-1", + team_alias: "Test Team", + models: ["gpt-4"], + }; + + const mockOrganization: Organization = { + organization_id: "org-1", + organization_alias: "Test Org", + budget_id: "budget-1", + metadata: {}, + models: ["gpt-4"], + spend: 0, + model_spend: {}, + created_at: "2024-01-01", + created_by: "user-1", + updated_at: "2024-01-01", + updated_by: "user-1", + litellm_budget_table: null, + teams: null, + users: null, + members: null, + }; + + mockUseTeam.mockReturnValue({ + data: mockTeam, + isLoading: false, + } as any); + + mockUseOrganization.mockReturnValue({ + data: mockOrganization, + isLoading: false, + } as any); + + renderWithProviders( + , + ); + + await waitFor(() => { + expect(screen.getByText("gpt-4")).toBeInTheDocument(); + expect(screen.queryByText("claude-3")).not.toBeInTheDocument(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx index 5aa1ba6a30..a3ddeff322 100644 --- a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx +++ b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx @@ -1,77 +1,110 @@ import { ProxyModel, useAllProxyModels } from "@/app/(dashboard)/hooks/models/useModels"; +import { useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import { useTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser"; import { Select, Skeleton, Tooltip, type SelectProps } from "antd"; import { Organization, Team } from "../networking"; -import { useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import { splitWildcardModels } from "./modelUtils"; -const MODEL_SELECT_SPECIAL_VALUES = { - ALL_PROXY_MODELS: { - label: "All Proxy Models", - value: "all-proxy-models", - }, - NO_DEFAULT_MODELS: { - label: "No Default Models", - value: "no-default-models", - }, -}; +const MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE = { + label: "All Proxy Models", + value: "all-proxy-models", +} as const; -const MODEL_SELECT_SPECIAL_VALUES_ARRAY = Object.values(MODEL_SELECT_SPECIAL_VALUES); +const MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE = { + label: "No Default Models", + value: "no-default-models", +} as const; -export interface ModelSelectContext { +const MODEL_SELECT_SPECIAL_VALUES_ARRAY = [ + MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE, + MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE, +] as const; + +export interface ModelSelectProps { teamID?: string; organizationID?: string; - includeUserModels?: boolean; - showAllTeamModelsOption?: boolean; - showAllProxyModelsOverride?: boolean; - includeSpecialOptions?: boolean; + options?: { + includeUserModels?: boolean; + showAllTeamModelsOption?: boolean; + showAllProxyModelsOverride?: boolean; + includeSpecialOptions?: boolean; + }; + context: "team" | "organization" | "user"; dataTestId?: string; value?: string[]; onChange: (values: string[]) => void; } -const filterModels = ( - allProxyModels: ProxyModel[], - ctx: ModelSelectContext, - { - selectedTeam, - selectedOrganization, - userModels, - }: { selectedTeam?: Team; selectedOrganization?: Organization; userModels?: ProxyModel[] }, -): ProxyModel[] => { - const deduplicatedProxyModels = Array.from(new Map(allProxyModels.map((model) => [model.id, model])).values()); - if (ctx.showAllProxyModelsOverride) { - return deduplicatedProxyModels; - } - - if (selectedOrganization) { - if (selectedOrganization.models.includes(MODEL_SELECT_SPECIAL_VALUES.ALL_PROXY_MODELS.value)) { - return deduplicatedProxyModels; - } - } - - return []; +type FilterContextArgs = { + allProxyModels: string[]; + selectedTeam?: Team; + selectedOrganization?: Organization; + userModels?: string[]; + options?: ModelSelectProps["options"]; }; -export const ModelSelect = (ctx: ModelSelectContext) => { - const { - teamID, - organizationID, - includeUserModels, - showAllTeamModelsOption, - showAllProxyModelsOverride, - includeSpecialOptions, - dataTestId, - value = [], - onChange, - } = ctx; +const contextFilters: Record string[]> = { + user: ({ allProxyModels, userModels, options }) => { + if (!userModels) return []; + if (options?.includeUserModels) return userModels; + return []; + }, + + team: ({ allProxyModels, selectedOrganization, userModels }) => { + if (selectedOrganization) { + if (selectedOrganization.models.includes(MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value)) { + return allProxyModels; + } + // Return organization's models (filtered from allProxyModels) + return allProxyModels.filter((model) => selectedOrganization.models.includes(model)); + } + + return userModels ?? []; + }, + + organization: ({ allProxyModels, selectedOrganization, options }) => { + if (!selectedOrganization) return []; + + if (selectedOrganization.models.includes(MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value)) { + return allProxyModels; + } + + return allProxyModels.filter((model) => selectedOrganization.models.includes(model)); + }, +}; + +const filterModels = ( + allProxyModels: ProxyModel[], + ctx: ModelSelectProps, + extra: { selectedTeam?: Team; selectedOrganization?: Organization; userModels?: string[] }, +): string[] => { + const deduplicatedProxyModels = Array.from(new Map(allProxyModels.map((m) => [m.id, m])).values()).map( + (model) => model.id, + ); + if (ctx.options?.showAllProxyModelsOverride) return deduplicatedProxyModels; + + const filterFn = contextFilters[ctx.context]; + if (!filterFn) return []; + + return filterFn({ allProxyModels: deduplicatedProxyModels, ...extra, options: ctx.options }); +}; + +export const ModelSelect = (props: ModelSelectProps) => { + const { teamID, organizationID, options, context, dataTestId, value = [], onChange } = props; + const { includeUserModels, showAllTeamModelsOption, showAllProxyModelsOverride, includeSpecialOptions } = + options || {}; const { data: allProxyModels, isLoading: isLoadingAllProxyModels } = useAllProxyModels(); const { data: team, isLoading: isLoadingTeam } = useTeam(teamID); const { data: organization, isLoading: isLoadingOrganization } = useOrganization(organizationID); + const { data: currentUser, isLoading: isCurrentUserLoading } = useCurrentUser(); const isSpecialOption = (value: string) => MODEL_SELECT_SPECIAL_VALUES_ARRAY.some((sv) => sv.value === value); const hasSpecialOptionSelected = value.some(isSpecialOption); - const isLoading = isLoadingAllProxyModels || isLoadingTeam || isLoadingOrganization; + const isLoading = isLoadingAllProxyModels || isLoadingTeam || isLoadingOrganization || isCurrentUserLoading; + const shouldShowAllProxyModels = + showAllProxyModelsOverride || + (organization?.models.includes(MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value) && includeSpecialOptions); if (isLoading) { return ; @@ -95,9 +128,10 @@ export const ModelSelect = (ctx: ModelSelectContext) => { onChange(finalValues); }; - const filteredModels = filterModels(allProxyModels?.data ?? [], ctx, { + const filteredModels = filterModels(allProxyModels?.data ?? [], props, { selectedTeam: team, selectedOrganization: organization, + userModels: currentUser?.models, }); const { wildcard, regular } = splitWildcardModels(filteredModels); @@ -107,36 +141,60 @@ export const ModelSelect = (ctx: ModelSelectContext) => { value={value} onChange={handleChange} options={[ - { - label: Special Options, - title: "Special Options", - options: MODEL_SELECT_SPECIAL_VALUES_ARRAY.map((specialValue) => ({ - label: {specialValue.label}, - value: specialValue.value, - disabled: value.length > 0 && value.some((v) => isSpecialOption(v) && v !== specialValue.value), - key: specialValue.value, - })), - }, - { - label: Wildcard Options, - title: "Wildcard Options", - options: wildcard.map((model) => { - const provider = model.id.replace("/*", ""); - const capitalizedProvider = provider.charAt(0).toUpperCase() + provider.slice(1); + includeSpecialOptions + ? { + label: Special Options, + title: "Special Options", + options: [ + ...(shouldShowAllProxyModels + ? [ + { + label: All Proxy Models, + value: MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value, + disabled: + value.length > 0 && + value.some( + (v) => isSpecialOption(v) && v !== MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value, + ), + key: MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value, + }, + ] + : []), + { + label: No Default Models, + value: MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE.value, + disabled: + value.length > 0 && + value.some((v) => isSpecialOption(v) && v !== MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE.value), + key: MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE.value, + }, + ], + } + : [], + ...(wildcard.length > 0 + ? [ + { + label: Wildcard Options, + title: "Wildcard Options", + options: wildcard.map((model) => { + const provider = model.replace("/*", ""); + const capitalizedProvider = provider.charAt(0).toUpperCase() + provider.slice(1); - return { - label: {`All ${capitalizedProvider} models`}, - value: model.id, - disabled: hasSpecialOptionSelected, - }; - }), - }, + return { + label: {`All ${capitalizedProvider} models`}, + value: model, + disabled: hasSpecialOptionSelected, + }; + }), + }, + ] + : []), { label: Models, title: "Models", options: regular.map((model) => ({ - label: {model.id}, - value: model.id, + label: {model}, + value: model, disabled: hasSpecialOptionSelected, })), }, diff --git a/ui/litellm-dashboard/src/components/ModelSelect/modelUtils.test.ts b/ui/litellm-dashboard/src/components/ModelSelect/modelUtils.test.ts index 3da6477f43..0e5f269cc6 100644 --- a/ui/litellm-dashboard/src/components/ModelSelect/modelUtils.test.ts +++ b/ui/litellm-dashboard/src/components/ModelSelect/modelUtils.test.ts @@ -1,4 +1,3 @@ -import type { ProxyModel } from "@/app/(dashboard)/hooks/models/useModels"; import { describe, expect, it } from "vitest"; import { splitWildcardModels } from "./modelUtils"; @@ -9,29 +8,21 @@ describe("splitWildcardModels", () => { }); it("should split models into wildcard and regular groups", () => { - const models: ProxyModel[] = [ - { id: "gpt-4", object: "model", created: 1234567890, owned_by: "openai" }, - { id: "openai/*", object: "model", created: 1234567890, owned_by: "openai" }, - { id: "claude-3", object: "model", created: 1234567890, owned_by: "anthropic" }, - { id: "anthropic/*", object: "model", created: 1234567890, owned_by: "anthropic" }, - ]; + const models: string[] = ["gpt-4", "openai/*", "claude-3", "anthropic/*"]; const result = splitWildcardModels(models); expect(result.wildcard).toHaveLength(2); - expect(result.wildcard[0].id).toBe("openai/*"); - expect(result.wildcard[1].id).toBe("anthropic/*"); + expect(result.wildcard[0]).toBe("openai/*"); + expect(result.wildcard[1]).toBe("anthropic/*"); expect(result.regular).toHaveLength(2); - expect(result.regular[0].id).toBe("gpt-4"); - expect(result.regular[1].id).toBe("claude-3"); + expect(result.regular[0]).toBe("gpt-4"); + expect(result.regular[1]).toBe("claude-3"); }); it("should return only wildcard models when all models are wildcard", () => { - const models: ProxyModel[] = [ - { id: "openai/*", object: "model", created: 1234567890, owned_by: "openai" }, - { id: "anthropic/*", object: "model", created: 1234567890, owned_by: "anthropic" }, - ]; + const models: string[] = ["openai/*", "anthropic/*"]; const result = splitWildcardModels(models); @@ -40,10 +31,7 @@ describe("splitWildcardModels", () => { }); it("should return only regular models when no models are wildcard", () => { - const models: ProxyModel[] = [ - { id: "gpt-4", object: "model", created: 1234567890, owned_by: "openai" }, - { id: "claude-3", object: "model", created: 1234567890, owned_by: "anthropic" }, - ]; + const models: string[] = ["gpt-4", "claude-3"]; const result = splitWildcardModels(models); @@ -52,16 +40,12 @@ describe("splitWildcardModels", () => { }); it("should correctly identify wildcard models ending with /*", () => { - const models: ProxyModel[] = [ - { id: "provider/*", object: "model", created: 1234567890, owned_by: "provider" }, - { id: "not-wildcard", object: "model", created: 1234567890, owned_by: "provider" }, - { id: "also-not/*/wildcard", object: "model", created: 1234567890, owned_by: "provider" }, - ]; + const models: string[] = ["provider/*", "not-wildcard", "also-not/*/wildcard"]; const result = splitWildcardModels(models); expect(result.wildcard).toHaveLength(1); - expect(result.wildcard[0].id).toBe("provider/*"); + expect(result.wildcard[0]).toBe("provider/*"); expect(result.regular).toHaveLength(2); }); }); diff --git a/ui/litellm-dashboard/src/components/ModelSelect/modelUtils.ts b/ui/litellm-dashboard/src/components/ModelSelect/modelUtils.ts index 54d176bd44..afaac14d6b 100644 --- a/ui/litellm-dashboard/src/components/ModelSelect/modelUtils.ts +++ b/ui/litellm-dashboard/src/components/ModelSelect/modelUtils.ts @@ -1,16 +1,14 @@ -import { ProxyModel } from "@/app/(dashboard)/hooks/models/useModels"; - export interface GroupedModels { - wildcard: ProxyModel[]; - regular: ProxyModel[]; + wildcard: string[]; + regular: string[]; } -export const splitWildcardModels = (models: ProxyModel[]): GroupedModels => { - const wildcard: ProxyModel[] = []; - const regular: ProxyModel[] = []; +export const splitWildcardModels = (models: string[]): GroupedModels => { + const wildcard: string[] = []; + const regular: string[] = []; for (const model of models) { - if (model.id.endsWith("/*")) { + if (model.endsWith("/*")) { wildcard.push(model); } else { regular.push(model); diff --git a/ui/litellm-dashboard/src/components/organizations.tsx b/ui/litellm-dashboard/src/components/organizations.tsx index 783a728258..fc825bdf92 100644 --- a/ui/litellm-dashboard/src/components/organizations.tsx +++ b/ui/litellm-dashboard/src/components/organizations.tsx @@ -29,12 +29,12 @@ import DeleteResourceModal from "./common_components/DeleteResourceModal"; import TableIconActionButton from "./common_components/IconActionButton/TableIconActionButtons/TableIconActionButton"; import { getModelDisplayName } from "./key_team_helpers/fetch_available_models_team_key"; import MCPServerSelector from "./mcp_server_management/MCPServerSelector"; +import { ModelSelect } from "./ModelSelect/ModelSelect"; import NotificationsManager from "./molecules/notifications_manager"; import { Organization, organizationCreateCall, organizationDeleteCall, organizationListCall } from "./networking"; import OrganizationInfoView from "./organization/organization_view"; import NumericalInput from "./shared/numerical_input"; import VectorStoreSelector from "./vector_store_management/VectorStoreSelector"; -import { ModelSelect } from "./ModelSelect/ModelSelect"; interface OrganizationsTableProps { organizations: Organization[]; @@ -475,10 +475,10 @@ const OrganizationsTable: React.FC = ({ form.setFieldValue("models", values)} + context="organization" /> diff --git a/ui/litellm-dashboard/src/components/team/team_info.test.tsx b/ui/litellm-dashboard/src/components/team/team_info.test.tsx index 4f7d70ba6a..6d5e55f734 100644 --- a/ui/litellm-dashboard/src/components/team/team_info.test.tsx +++ b/ui/litellm-dashboard/src/components/team/team_info.test.tsx @@ -1,7 +1,7 @@ import * as networking from "@/components/networking"; import { act, fireEvent, screen, waitFor } from "@testing-library/react"; import { renderWithProviders } from "../../../tests/test-utils"; -import { afterEach, describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import TeamInfoView from "./team_info"; // Mock the networking module @@ -17,7 +17,54 @@ vi.mock("@/components/networking", () => ({ organizationInfoCall: vi.fn(), })); +// Mock hooks used by ModelSelect +vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ + useAllProxyModels: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ + useTeam: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ + useOrganization: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/users/useCurrentUser", () => ({ + useCurrentUser: vi.fn(), +})); + +import { useAllProxyModels } from "@/app/(dashboard)/hooks/models/useModels"; +import { useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; +import { useTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser"; + +const mockUseAllProxyModels = vi.mocked(useAllProxyModels); +const mockUseTeam = vi.mocked(useTeam); +const mockUseOrganization = vi.mocked(useOrganization); +const mockUseCurrentUser = vi.mocked(useCurrentUser); + describe("TeamInfoView", () => { + beforeEach(() => { + // Set up default mock implementations + mockUseAllProxyModels.mockReturnValue({ + data: { data: [] }, + isLoading: false, + } as any); + mockUseTeam.mockReturnValue({ + data: undefined, + isLoading: false, + } as any); + mockUseOrganization.mockReturnValue({ + data: undefined, + isLoading: false, + } as any); + mockUseCurrentUser.mockReturnValue({ + data: { models: [] }, + isLoading: false, + } as any); + }); + afterEach(() => { vi.clearAllMocks(); }); @@ -158,7 +205,7 @@ describe("TeamInfoView", () => { }); await waitFor(() => { - expect(screen.getByLabelText("Models")).toBeInTheDocument(); + expect(screen.getByTestId("models-select")).toBeInTheDocument(); }); const allProxyModelsOption = screen.queryByText("All Proxy Models"); @@ -170,6 +217,40 @@ describe("TeamInfoView", () => { const organizationModels = ["gpt-4", "claude-3-opus"]; const userModels = ["gpt-4", "gpt-3.5-turbo", "claude-3-opus", "claude-2"]; + // Mock all proxy models - should include all user models + const allProxyModels = userModels.map((id) => ({ + id, + object: "model", + created: 1234567890, + owned_by: "openai", + })); + + mockUseAllProxyModels.mockReturnValue({ + data: { data: allProxyModels }, + isLoading: false, + } as any); + + mockUseCurrentUser.mockReturnValue({ + data: { models: userModels }, + isLoading: false, + } as any); + + const organizationData = { + organization_id: organizationId, + organization_name: "Test Organization", + spend: 0, + max_budget: null, + models: organizationModels, + tpm_limit: null, + rpm_limit: null, + members: null, + }; + + mockUseOrganization.mockReturnValue({ + data: organizationData, + isLoading: false, + } as any); + vi.mocked(networking.teamInfoCall).mockResolvedValue({ team_id: "123", team_info: { @@ -206,16 +287,7 @@ describe("TeamInfoView", () => { team_memberships: [], }); - vi.mocked(networking.organizationInfoCall).mockResolvedValue({ - organization_id: organizationId, - organization_name: "Test Organization", - spend: 0, - max_budget: null, - models: organizationModels, - tpm_limit: null, - rpm_limit: null, - members: null, - }); + vi.mocked(networking.organizationInfoCall).mockResolvedValue(organizationData); vi.mocked(networking.getGuardrailsList).mockResolvedValue({ guardrails: [] }); vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]); @@ -253,26 +325,39 @@ describe("TeamInfoView", () => { }); await waitFor(() => { - expect(screen.getByLabelText("Models")).toBeInTheDocument(); + expect(screen.getByTestId("models-select")).toBeInTheDocument(); }); - const modelsSelect = screen.getByLabelText("Models"); + // Find the Ant Design Select selector element to open the dropdown + // The data-testid is on the Select component, we need to find the selector inside it + const modelsSelectElement = screen.getByTestId("models-select"); + const selectSelector = modelsSelectElement.querySelector(".ant-select-selector"); + expect(selectSelector).toBeTruthy(); + + // Open the dropdown by clicking on the selector act(() => { - fireEvent.mouseDown(modelsSelect); + fireEvent.mouseDown(selectSelector!); }); - await waitFor(() => { - const dropdownOptions = screen.getAllByRole("option"); - const optionTexts = dropdownOptions.map((option) => option.textContent); + // Wait for dropdown to open - Ant Design renders options in a portal + await waitFor( + () => { + const dropdownOptions = document.querySelectorAll(".ant-select-item-option"); + expect(dropdownOptions.length).toBeGreaterThan(0); + }, + { timeout: 5000 }, + ); - organizationModels.forEach((model) => { - expect(optionTexts).toContain(model); - }); + const dropdownOptions = document.querySelectorAll(".ant-select-item-option"); + const optionTexts = Array.from(dropdownOptions).map((option) => option.textContent?.trim() || ""); - const modelsNotInOrganization = userModels.filter((m) => !organizationModels.includes(m)); - modelsNotInOrganization.forEach((model) => { - expect(optionTexts).not.toContain(model); - }); + organizationModels.forEach((model) => { + expect(optionTexts).toContain(model); + }); + + const modelsNotInOrganization = userModels.filter((m) => !organizationModels.includes(m)); + modelsNotInOrganization.forEach((model) => { + expect(optionTexts).not.toContain(model); }); }, 10000); diff --git a/ui/litellm-dashboard/src/components/team/team_info.tsx b/ui/litellm-dashboard/src/components/team/team_info.tsx index fd7994aa1d..819c9f0f9b 100644 --- a/ui/litellm-dashboard/src/components/team/team_info.tsx +++ b/ui/litellm-dashboard/src/components/team/team_info.tsx @@ -1,3 +1,4 @@ +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import UserSearchModal from "@/components/common_components/user_search_modal"; import { getGuardrailsList, @@ -12,6 +13,7 @@ import { } from "@/components/networking"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import { mapEmptyStringToNull } from "@/utils/keyUpdateUtils"; +import { isProxyAdminRole } from "@/utils/roles"; import { InfoCircleOutlined } from "@ant-design/icons"; import { ArrowLeftIcon } from "@heroicons/react/outline"; import { @@ -34,11 +36,13 @@ import React, { useEffect, useMemo, useState } from "react"; import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils"; import AgentSelector from "../agent_management/AgentSelector"; import DeleteResourceModal from "../common_components/DeleteResourceModal"; +import DurationSelect from "../common_components/DurationSelect"; import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSelector"; -import { getModelDisplayName, unfurlWildcardModelsInList } from "../key_team_helpers/fetch_available_models_team_key"; +import { unfurlWildcardModelsInList } from "../key_team_helpers/fetch_available_models_team_key"; import LoggingSettingsView from "../logging_settings_view"; import MCPServerSelector from "../mcp_server_management/MCPServerSelector"; import MCPToolPermissions from "../mcp_server_management/MCPToolPermissions"; +import { ModelSelect } from "../ModelSelect/ModelSelect"; import NotificationsManager from "../molecules/notifications_manager"; import { fetchMCPAccessGroups } from "../networking"; import ObjectPermissionsView from "../object_permissions_view"; @@ -48,7 +52,6 @@ import EditLoggingSettings from "./EditLoggingSettings"; import MemberModal from "./EditMembership"; import MemberPermissions from "./member_permissions"; import TeamMembersComponent from "./team_member_view"; -import DurationSelect from "../common_components/DurationSelect"; export interface TeamMembership { user_id: string; @@ -173,8 +176,7 @@ const TeamInfoView: React.FC = ({ const [isDeleting, setIsDeleting] = useState(false); const [isTeamSaving, setIsTeamSaving] = useState(false); const [organization, setOrganization] = useState(null); - - console.log("userModels in team info", userModels); + const { userRole } = useAuthorized(); const canEditTeam = is_team_admin || is_proxy_admin; @@ -696,47 +698,19 @@ const TeamInfoView: React.FC = ({ name="models" rules={[{ required: true, message: "Please select at least one model" }]} > - + form.setFieldValue("models", values)} + teamID={teamId} + organizationID={teamData?.team_info?.organization_id || undefined} + options={{ + includeSpecialOptions: true, + includeUserModels: !teamData?.team_info?.organization_id, + showAllProxyModelsOverride: isProxyAdminRole(userRole) && !teamData?.team_info?.organization_id, + }} + context="team" + dataTestId="models-select" + /> diff --git a/ui/litellm-dashboard/src/components/view_users/types.ts b/ui/litellm-dashboard/src/components/view_users/types.ts index d674db5c7d..744aa00a88 100644 --- a/ui/litellm-dashboard/src/components/view_users/types.ts +++ b/ui/litellm-dashboard/src/components/view_users/types.ts @@ -5,6 +5,7 @@ export interface UserInfo { user_role: string; spend: number; max_budget: number | null; + models: string[]; key_count: number; created_at: string; updated_at: string;