From 9867cc6cf98422bb29aa326ec75a1c5068facce1 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 2 Dec 2025 15:58:27 -0800 Subject: [PATCH] Change edit team models to match create team models --- .../src/components/team/team_info.test.tsx | 115 +++++++++++++++++- .../src/components/team/team_info.tsx | 101 +++++++++++++-- 2 files changed, 202 insertions(+), 14 deletions(-) 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 54362193350..3a87b42d253 100644 --- a/ui/litellm-dashboard/src/components/team/team_info.test.tsx +++ b/ui/litellm-dashboard/src/components/team/team_info.test.tsx @@ -13,6 +13,7 @@ vi.mock("@/components/networking", () => ({ getGuardrailsList: vi.fn(), fetchMCPAccessGroups: vi.fn(), getTeamPermissionsCall: vi.fn(), + organizationInfoCall: vi.fn(), })); describe("TeamInfoView", () => { @@ -161,6 +162,116 @@ describe("TeamInfoView", () => { const allProxyModelsOption = screen.queryByText("All Proxy Models"); expect(allProxyModelsOption).not.toBeInTheDocument(); - }, // This is a workaround to fix the flaky test issue. TODO: Remove this once we have a better solution. - 10000); + }, 10000); // This is a workaround to fix the flaky test issue. TODO: Remove this once we have a better solution. + + it("should only show organization models in dropdown when team is in organization with limited models", async () => { + const organizationId = "org-123"; + const organizationModels = ["gpt-4", "claude-3-opus"]; + const userModels = ["gpt-4", "gpt-3.5-turbo", "claude-3-opus", "claude-2"]; + + vi.mocked(networking.teamInfoCall).mockResolvedValue({ + team_id: "123", + team_info: { + team_alias: "Test Team", + team_id: "123", + organization_id: organizationId, + admins: ["admin@test.com"], + members: ["user1@test.com"], + members_with_roles: [ + { + user_id: "user1@test.com", + user_email: "user1@test.com", + role: "member", + spend: 0, + budget_id: "budget1", + }, + ], + metadata: {}, + tpm_limit: null, + rpm_limit: null, + max_budget: null, + budget_duration: null, + models: ["gpt-4"], + blocked: false, + spend: 0, + max_parallel_requests: null, + budget_reset_at: null, + model_id: null, + litellm_model_table: null, + created_at: "2024-01-01T00:00:00Z", + team_member_budget_table: null, + }, + keys: [], + 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.getGuardrailsList).mockResolvedValue({ guardrails: [] }); + vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]); + + render( + {}} + onClose={() => {}} + accessToken="123" + is_team_admin={true} + is_proxy_admin={true} + userModels={userModels} + editTeam={false} + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getAllByText("Test Team")).not.toBeNull(); + }); + + const settingsTab = screen.getByRole("tab", { name: "Settings" }); + act(() => { + fireEvent.click(settingsTab); + }); + + await waitFor(() => { + expect(screen.getByText("Team Settings")).toBeInTheDocument(); + }); + + const editButton = screen.getByRole("button", { name: "Edit Settings" }); + act(() => { + fireEvent.click(editButton); + }); + + await waitFor(() => { + expect(screen.getByLabelText("Models")).toBeInTheDocument(); + }); + + const modelsSelect = screen.getByLabelText("Models"); + act(() => { + fireEvent.mouseDown(modelsSelect); + }); + + await waitFor(() => { + const dropdownOptions = screen.getAllByRole("option"); + const optionTexts = dropdownOptions.map((option) => option.textContent); + + 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 a53d35e57bb..581722e8fa7 100644 --- a/ui/litellm-dashboard/src/components/team/team_info.tsx +++ b/ui/litellm-dashboard/src/components/team/team_info.tsx @@ -2,6 +2,8 @@ import UserSearchModal from "@/components/common_components/user_search_modal"; import { getGuardrailsList, Member, + Organization, + organizationInfoCall, teamInfoCall, teamMemberAddCall, teamMemberDeleteCall, @@ -28,11 +30,11 @@ import { } from "@tremor/react"; import { Button, Form, Input, message, Select, Switch, Tooltip } from "antd"; import { CheckIcon, CopyIcon } from "lucide-react"; -import React, { useEffect, useState } from "react"; +import React, { useEffect, useMemo, useState } from "react"; import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils"; import DeleteResourceModal from "../common_components/DeleteResourceModal"; import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSelector"; -import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; +import { getModelDisplayName, 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"; @@ -117,6 +119,29 @@ export interface TeamInfoProps { premiumUser?: boolean; } +const getOrganizationModels = (organization: Organization | null, userModels: string[]) => { + let tempModelsToPick = []; + + if (organization) { + // Check if organization has "all-proxy-models" in its models array + if (organization.models.includes("all-proxy-models")) { + // Treat as all-proxy-models (use userModels) + tempModelsToPick = userModels; + } else if (organization.models.length > 0) { + // Organization has specific models + tempModelsToPick = organization.models; + } else { + // Empty array [] is treated as all-proxy-models + tempModelsToPick = userModels; + } + } else { + // No organization, show all available models + tempModelsToPick = userModels; + } + + return unfurlWildcardModelsInList(tempModelsToPick, userModels); +}; + const TeamInfoView: React.FC = ({ teamId, onClose, @@ -143,6 +168,7 @@ const TeamInfoView: React.FC = ({ const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); const [isDeleting, setIsDeleting] = useState(false); const [isTeamSaving, setIsTeamSaving] = useState(false); + const [organization, setOrganization] = useState(null); console.log("userModels in team info", userModels); @@ -166,6 +192,31 @@ const TeamInfoView: React.FC = ({ fetchTeamInfo(); }, [teamId, accessToken]); + // Fetch organization data when team has organization_id + useEffect(() => { + const fetchOrganization = async () => { + if (!accessToken || !teamData?.team_info?.organization_id) { + setOrganization(null); + return; + } + + try { + const orgData = await organizationInfoCall(accessToken, teamData.team_info.organization_id); + setOrganization(orgData); + } catch (error) { + console.error("Error fetching organization info:", error); + setOrganization(null); + } + }; + + fetchOrganization(); + }, [accessToken, teamData?.team_info?.organization_id]); + + // Compute modelsToPick based on organization and userModels + const modelsToPick = useMemo(() => { + return getOrganizationModels(organization, userModels); + }, [organization, userModels]); + const fetchMcpAccessGroups = async () => { if (!accessToken) return; if (mcpAccessGroupsLoaded) return; @@ -596,15 +647,41 @@ const TeamInfoView: React.FC = ({ rules={[{ required: true, message: "Please select at least one model" }]} > +