Merge pull request #18532 from BerriAI/litellm_ui_add_model_team_admin

[Feature] UI - Add Model Clearer Flow For Team Admins

ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.ts

@@ -0,0 +1,18 @@
+import { useQuery, UseQueryResult } from "@tanstack/react-query";
+import { createQueryKeys } from "../common/queryKeysFactory";
+import { getGuardrailsList } from "@/components/networking";
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
+
+const guardrailKeys = createQueryKeys("guardrails");
+
+export const useGuardrails = (): UseQueryResult<string[]> => {
+  const { accessToken, userId, userRole } = useAuthorized();
+  return useQuery<string[]>({
+    queryKey: guardrailKeys.list({}),
+    queryFn: async () => {
+      const response = await getGuardrailsList(accessToken!);
+      return response.guardrails.map((g: { guardrail_name: string }) => g.guardrail_name);
+    },
+    enabled: Boolean(accessToken && userId && userRole),
+  });
+};

ui/litellm-dashboard/src/app/(dashboard)/hooks/tags/useTags.ts

@@ -0,0 +1,16 @@
+import { useQuery, UseQueryResult } from "@tanstack/react-query";
+import { createQueryKeys } from "../common/queryKeysFactory";
+import { tagListCall } from "@/components/networking";
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
+import { TagListResponse } from "@/components/tag_management/types";
+
+const tagKeys = createQueryKeys("tags");
+
+export const useTags = (): UseQueryResult<TagListResponse> => {
+  const { accessToken, userId, userRole } = useAuthorized();
+  return useQuery<TagListResponse>({
+    queryKey: tagKeys.list({}),
+    queryFn: async () => await tagListCall(accessToken!),
+    enabled: Boolean(accessToken && userId && userRole),
+  });
+};

ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx

@@ -536,21 +536,19 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({
     );
   };
 
-  const handleOk = () => {
-    addModelForm
-      .validateFields()
-      .then((values: any) => {
-        handleAddModelSubmit(values, accessToken, addModelForm, handleRefreshClick);
-      })
-      .catch((error: any) => {
-        const errorMessages =
-          error.errorFields
-            ?.map((field: any) => {
-              return `${field.name.join(".")}: ${field.errors.join(", ")}`;
-            })
-            .join(" | ") || "Unknown validation error";
-        NotificationsManager.fromBackend(`Please fill in the following required fields: ${errorMessages}`);
-      });
+  const handleOk = async () => {
+    try {
+      const values = await addModelForm.validateFields();
+      await handleAddModelSubmit(values, accessToken, addModelForm, handleRefreshClick);
+    } catch (error: any) {
+      const errorMessages =
+        error.errorFields
+          ?.map((field: any) => {
+            return `${field.name.join(".")}: ${field.errors.join(", ")}`;
+          })
+          .join(" | ") || "Unknown validation error";
+      NotificationsManager.fromBackend(`Please fill in the following required fields: ${errorMessages}`);
+    }
   };
 
   Object.keys(Providers).find((key) => (Providers as { [index: string]: any })[key] === selectedProvider);
@@ -679,7 +677,6 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({
                       credentials={credentialsList}
                       accessToken={accessToken}
                       userRole={userRole}
-                      premiumUser={premiumUser}
                     />
                   </TabPanel>
                 )}

ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx

@@ -0,0 +1,276 @@
+import { renderHook, screen, waitFor, renderWithProviders } from "../../../tests/test-utils";
+import userEvent from "@testing-library/user-event";
+import { Form } from "antd";
+import type { UploadProps } from "antd/es/upload";
+import { describe, expect, it, vi } from "vitest";
+import type { Team } from "../key_team_helpers/key_list";
+import type { CredentialItem } from "../networking";
+import { Providers } from "../provider_info_helpers";
+import AddModelForm from "./AddModelForm";
+
+vi.mock("../molecules/models/ProviderLogo", () => ({
+  ProviderLogo: ({ provider, className }: { provider: string; className?: string }) => (
+    <div className={className} data-testid={`provider-logo-${provider}`}>
+      {provider}
+    </div>
+  ),
+}));
+
+vi.mock("../networking", async () => {
+  const actual = await vi.importActual("../networking");
+  return {
+    ...actual,
+    getGuardrailsList: vi.fn().mockResolvedValue({
+      guardrails: [{ guardrail_name: "test-guardrail-1" }, { guardrail_name: "test-guardrail-2" }],
+    }),
+    tagListCall: vi.fn().mockResolvedValue({}),
+    modelAvailableCall: vi.fn().mockResolvedValue({
+      data: [{ id: "model-group-1" }, { id: "model-group-2" }],
+    }),
+    modelHubCall: vi.fn().mockResolvedValue({
+      data: [
+        { model_group: "gpt-4", mode: "chat" },
+        { model_group: "gpt-3.5-turbo", mode: "chat" },
+      ],
+    }),
+    getProviderCreateMetadata: vi.fn().mockResolvedValue([
+      {
+        provider: "OpenAI",
+        provider_display_name: "OpenAI",
+        litellm_provider: "openai",
+        default_model_placeholder: "gpt-3.5-turbo",
+        credential_fields: [],
+      },
+    ]),
+  };
+});
+
+vi.mock("@/app/(dashboard)/hooks/providers/useProviderFields", () => ({
+  useProviderFields: vi.fn().mockReturnValue({
+    data: [
+      {
+        provider: "OpenAI",
+        provider_display_name: "OpenAI",
+        litellm_provider: "openai",
+        default_model_placeholder: "gpt-3.5-turbo",
+        credential_fields: [],
+      },
+    ],
+    isLoading: false,
+    error: null,
+  }),
+}));
+
+vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
+  default: vi.fn(),
+}));
+
+vi.mock("@/app/(dashboard)/hooks/guardrails/useGuardrails", () => ({
+  useGuardrails: vi.fn().mockReturnValue({
+    data: [{ guardrail_name: "test-guardrail" }],
+    isLoading: false,
+    error: null,
+  }),
+}));
+
+vi.mock("@/app/(dashboard)/hooks/tags/useTags", () => ({
+  useTags: vi.fn().mockReturnValue({
+    data: { tag1: ["model1", "model2"] },
+    isLoading: false,
+    error: null,
+  }),
+}));
+
+const mockAuthorizedUser = (userRole: string, userId: string, premiumUser: boolean) => ({
+  token: "test-token",
+  accessToken: "test-access-token",
+  userId,
+  userEmail: "test@example.com",
+  userRole,
+  premiumUser,
+  disabledPersonalKeyCreation: false,
+  showSSOBanner: false,
+});
+
+const testTeam: Team = {
+  team_id: "team-1",
+  team_alias: "Test Team",
+  models: ["gpt-4"],
+  max_budget: 100,
+  budget_duration: "monthly",
+  tpm_limit: null,
+  rpm_limit: null,
+  organization_id: "org-1",
+  created_at: "2024-01-01T00:00:00Z",
+  keys: [],
+  members_with_roles: [],
+};
+
+const createTestProps = (userRole = "proxy_admin", userId = "user-1", isTeamAdmin = false) => {
+  const { result } = renderHook(() => Form.useForm());
+  const [form] = result.current;
+
+  const teams = [
+    {
+      ...testTeam,
+      members_with_roles: isTeamAdmin ? [{ user_id: userId, role: "admin" }] : [],
+    },
+  ];
+
+  const credentials: CredentialItem[] = [
+    {
+      credential_name: "test-credential",
+      credential_values: {},
+      credential_info: {
+        custom_llm_provider: "openai",
+        description: "Test credential",
+      },
+    },
+  ];
+
+  const uploadProps: UploadProps = {
+    beforeUpload: () => false,
+    showUploadList: false,
+  };
+
+  return {
+    form,
+    handleOk: vi.fn(),
+    setSelectedProvider: vi.fn(),
+    setProviderModelsFn: vi.fn(),
+    getPlaceholder: vi.fn((provider: Providers) => `Enter ${provider} model name`),
+    setShowAdvancedSettings: vi.fn(),
+    selectedProvider: Providers.OpenAI,
+    providerModels: ["gpt-4", "gpt-3.5-turbo"],
+    showAdvancedSettings: false,
+    teams,
+    credentials,
+    uploadProps,
+    userRole,
+    userId,
+  };
+};
+
+describe("AddModelForm", () => {
+  it("should render", async () => {
+    const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized"));
+    mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("proxy_admin", "user-1", true));
+
+    const props = createTestProps();
+
+    renderWithProviders(<AddModelForm {...props} />);
+
+    expect(await screen.findByRole("heading", { name: "Add Model" })).toBeInTheDocument();
+  });
+
+  it("should show proxy admin only (not team admin) - should not see Select Team dropdown unless switch is toggled", async () => {
+    const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized"));
+    mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("proxy_admin", "user-1", true));
+
+    const props = createTestProps("proxy_admin", "user-1", false);
+
+    renderWithProviders(<AddModelForm {...props} />);
+
+    await screen.findByText("Provider");
+
+    expect(screen.queryByText("Team Selection Required")).not.toBeInTheDocument();
+    expect(screen.queryByText("Select Team")).not.toBeInTheDocument();
+
+    const teamSwitch = screen.getByRole("switch");
+    expect(teamSwitch).toBeInTheDocument();
+
+    expect(screen.queryByText("Select Team")).not.toBeInTheDocument();
+
+    await userEvent.click(teamSwitch);
+
+    expect(await screen.findByText("Select Team")).toBeInTheDocument();
+  });
+
+  it("should show proxy admin who is also team admin - should not see Select Team dropdown unless switch is toggled", async () => {
+    const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized"));
+    mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("proxy_admin", "user-1", true));
+
+    const props = createTestProps("proxy_admin", "user-1", true);
+
+    renderWithProviders(<AddModelForm {...props} />);
+
+    await screen.findByText("Provider");
+
+    expect(screen.queryByText("Team Selection Required")).not.toBeInTheDocument();
+    expect(screen.queryByText("Select Team")).not.toBeInTheDocument();
+
+    const teamSwitch = screen.getByRole("switch");
+    expect(teamSwitch).toBeInTheDocument();
+
+    expect(screen.queryByText("Select Team")).not.toBeInTheDocument();
+
+    await userEvent.click(teamSwitch);
+
+    expect(await screen.findByText("Select Team")).toBeInTheDocument();
+  });
+
+  it("should show team admin (not proxy admin) - should see alert and team select, must select team before seeing remaining fields", async () => {
+    const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized"));
+    mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("team_member", "user-1", true));
+
+    const props = createTestProps("team_member", "user-1", true);
+
+    renderWithProviders(<AddModelForm {...props} />);
+
+    await screen.findByRole("heading", { name: "Add Model" });
+
+    expect(screen.getByText("Team Selection Required")).toBeInTheDocument();
+
+    expect(screen.getByText("Select Team")).toBeInTheDocument();
+
+    expect(screen.queryByText("Provider")).not.toBeInTheDocument();
+
+    const teamSelect = screen.getByRole("combobox");
+    await userEvent.click(teamSelect);
+    await userEvent.click(screen.getByText("Test Team"));
+
+    await waitFor(() => {
+      expect(screen.getByText("Provider")).toBeInTheDocument();
+    });
+  });
+
+  it("should show team admin (not proxy admin) - should not see team-BYOK switch", async () => {
+    const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized"));
+    mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("team_member", "user-1", true));
+
+    const props = createTestProps("team_member", "user-1", true);
+
+    renderWithProviders(<AddModelForm {...props} />);
+
+    await screen.findByText("Select Team");
+
+    const teamSelect = screen.getByRole("combobox");
+    await userEvent.click(teamSelect);
+    await userEvent.click(screen.getByText("Test Team"));
+
+    await waitFor(() => {
+      expect(screen.getByText("Provider")).toBeInTheDocument();
+    });
+
+    expect(screen.queryByRole("switch")).not.toBeInTheDocument();
+  });
+
+  it("should handle non-admin, non-team-admin users - should not see team selection or switch", async () => {
+    const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized"));
+    mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("user", "user-1", false));
+
+    const props = createTestProps("user", "user-1", false);
+
+    renderWithProviders(<AddModelForm {...props} />);
+
+    await screen.findByRole("heading", { name: "Add Model" });
+
+    expect(screen.queryByText("Team Selection Required")).not.toBeInTheDocument();
+
+    expect(screen.queryByText("Select Team")).not.toBeInTheDocument();
+
+    expect(screen.queryByText("Provider")).not.toBeInTheDocument();
+
+    expect(screen.queryByRole("switch")).not.toBeInTheDocument();
+  });
+});

ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx

@@ -0,0 +1,421 @@
+import { useProviderFields } from "@/app/(dashboard)/hooks/providers/useProviderFields";
+import { useGuardrails } from "@/app/(dashboard)/hooks/guardrails/useGuardrails";
+import { useTags } from "@/app/(dashboard)/hooks/tags/useTags";
+import { all_admin_roles, isUserTeamAdminForAnyTeam } from "@/utils/roles";
+import { Switch, Text } from "@tremor/react";
+import type { FormInstance } from "antd";
+import { Select as AntdSelect, Button, Card, Col, Form, Modal, Row, Tooltip, Typography, Alert } from "antd";
+import type { UploadProps } from "antd/es/upload";
+import React, { useEffect, useMemo, useState } from "react";
+import TeamDropdown from "../common_components/team_dropdown";
+import type { Team } from "../key_team_helpers/key_list";
+import { type CredentialItem, type ProviderCreateInfo, modelAvailableCall } from "../networking";
+import { Providers, providerLogoMap } from "../provider_info_helpers";
+import { ProviderLogo } from "../molecules/models/ProviderLogo";
+import AdvancedSettings from "./advanced_settings";
+import ConditionalPublicModelName from "./conditional_public_model_name";
+import LiteLLMModelNameField from "./litellm_model_name";
+import ConnectionErrorDisplay from "./model_connection_test";
+import ProviderSpecificFields from "./provider_specific_fields";
+import { TEST_MODES } from "./add_model_modes";
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
+
+interface AddModelFormProps {
+  form: FormInstance; // For the Add Model tab
+  handleOk: () => Promise<void>;
+  selectedProvider: Providers;
+  setSelectedProvider: (provider: Providers) => void;
+  providerModels: string[];
+  setProviderModelsFn: (provider: Providers) => void;
+  getPlaceholder: (provider: Providers) => string;
+  uploadProps: UploadProps;
+  showAdvancedSettings: boolean;
+  setShowAdvancedSettings: (show: boolean) => void;
+  teams: Team[] | null;
+  credentials: CredentialItem[];
+}
+
+const { Title, Link } = Typography;
+
+const AddModelForm: React.FC<AddModelFormProps> = ({
+  form,
+  handleOk,
+  selectedProvider,
+  setSelectedProvider,
+  providerModels,
+  setProviderModelsFn,
+  getPlaceholder,
+  uploadProps,
+  showAdvancedSettings,
+  setShowAdvancedSettings,
+  teams,
+  credentials,
+}) => {
+  const [testMode, setTestMode] = useState<string>("chat");
+  const [isResultModalVisible, setIsResultModalVisible] = useState<boolean>(false);
+  const [isTestingConnection, setIsTestingConnection] = useState<boolean>(false);
+  // Using a unique ID to force the ConnectionErrorDisplay to remount and run a fresh test
+  const [connectionTestId, setConnectionTestId] = useState<string>("");
+
+  const { accessToken, userRole, premiumUser, userId } = useAuthorized();
+  const {
+    data: providerMetadata,
+    isLoading: isProviderMetadataLoading,
+    error: providerMetadataError,
+  } = useProviderFields();
+  const { data: guardrailsList, isLoading: isGuardrailsLoading, error: guardrailsError } = useGuardrails();
+  const { data: tagsList, isLoading: isTagsLoading, error: tagsError } = useTags();
+
+  const handleTestConnection = async () => {
+    setIsTestingConnection(true);
+    setConnectionTestId(`test-${Date.now()}`);
+    setIsResultModalVisible(true);
+  };
+
+  const [isTeamOnly, setIsTeamOnly] = useState<boolean>(false);
+  const [modelAccessGroups, setModelAccessGroups] = useState<string[]>([]);
+  // Team admin specific state
+  const [teamAdminSelectedTeam, setTeamAdminSelectedTeam] = useState<string | null>(null);
+
+  useEffect(() => {
+    const fetchModelAccessGroups = async () => {
+      const response = await modelAvailableCall(accessToken, "", "", false, null, true, true);
+      setModelAccessGroups(response["data"].map((model: any) => model["id"]));
+    };
+    fetchModelAccessGroups();
+  }, [accessToken]);
+
+  const sortedProviderMetadata: ProviderCreateInfo[] = useMemo(() => {
+    if (!providerMetadata) {
+      return [];
+    }
+    return [...providerMetadata].sort((a, b) => a.provider_display_name.localeCompare(b.provider_display_name));
+  }, [providerMetadata]);
+
+  const providerMetadataErrorText = providerMetadataError
+    ? providerMetadataError instanceof Error
+      ? providerMetadataError.message
+      : "Failed to load providers"
+    : null;
+
+  const isAdmin = all_admin_roles.includes(userRole);
+  const isTeamAdmin = isUserTeamAdminForAnyTeam(teams, userId);
+
+  return (
+    <>
+      <Title level={2}>Add Model</Title>
+
+      <Card>
+        <Form
+          form={form}
+          onFinish={async (values) => {
+            console.log("🔥 Form onFinish triggered with values:", values);
+            await handleOk().then(() => {
+              setTeamAdminSelectedTeam(null);
+            });
+          }}
+          onFinishFailed={(errorInfo) => {
+            console.log("💥 Form onFinishFailed triggered:", errorInfo);
+          }}
+          labelCol={{ span: 10 }}
+          wrapperCol={{ span: 16 }}
+          labelAlign="left"
+        >
+          <>
+            {isTeamAdmin && !isAdmin && (
+              <>
+                <Form.Item
+                  label="Select Team"
+                  name="team_id"
+                  rules={[{ required: true, message: "Please select a team to continue" }]}
+                  tooltip="Select the team for which you want to add this model"
+                >
+                  <TeamDropdown
+                    teams={teams}
+                    onChange={(value) => {
+                      setTeamAdminSelectedTeam(value);
+                    }}
+                  />
+                </Form.Item>
+                {!teamAdminSelectedTeam && (
+                  <Alert
+                    message="Team Selection Required"
+                    description="As a team admin, you need to select your team first before adding models."
+                    type="info"
+                    showIcon
+                    className="mb-4"
+                  />
+                )}
+              </>
+            )}
+            {(isAdmin || (isTeamAdmin && teamAdminSelectedTeam)) && (
+              <>
+                <Form.Item
+                  rules={[{ required: true, message: "Required" }]}
+                  label="Provider:"
+                  name="custom_llm_provider"
+                  tooltip="E.g. OpenAI, Azure OpenAI, Anthropic, Bedrock, etc."
+                  labelCol={{ span: 10 }}
+                  labelAlign="left"
+                >
+                  <AntdSelect
+                    virtual={false}
+                    showSearch
+                    loading={isProviderMetadataLoading}
+                    placeholder={isProviderMetadataLoading ? "Loading providers..." : "Select a provider"}
+                    optionFilterProp="data-label"
+                    onChange={(value) => {
+                      setSelectedProvider(value as Providers);
+                      setProviderModelsFn(value as Providers);
+                      form.setFieldsValue({
+                        custom_llm_provider: value,
+                      });
+                      form.setFieldsValue({
+                        model: [],
+                        model_name: undefined,
+                      });
+                    }}
+                  >
+                    {providerMetadataErrorText && sortedProviderMetadata.length === 0 && (
+                      <AntdSelect.Option key="__error" value="">
+                        {providerMetadataErrorText}
+                      </AntdSelect.Option>
+                    )}
+                    {sortedProviderMetadata.map((providerInfo) => {
+                      const displayName = providerInfo.provider_display_name;
+                      const providerKey = providerInfo.provider;
+                      const logoSrc = providerLogoMap[displayName] ?? "";
+
+                      return (
+                        <AntdSelect.Option key={providerKey} value={providerKey} data-label={displayName}>
+                          <div className="flex items-center space-x-2">
+                            <ProviderLogo provider={providerKey} className="w-5 h-5" />
+                            <span>{displayName}</span>
+                          </div>
+                        </AntdSelect.Option>
+                      );
+                    })}
+                  </AntdSelect>
+                </Form.Item>
+                <LiteLLMModelNameField
+                  selectedProvider={selectedProvider}
+                  providerModels={providerModels}
+                  getPlaceholder={getPlaceholder}
+                />
+
+                {/* Conditionally Render "Public Model Name" */}
+                <ConditionalPublicModelName />
+
+                {/* Select Mode */}
+                <Form.Item label="Mode" name="mode" className="mb-1">
+                  <AntdSelect
+                    style={{ width: "100%" }}
+                    value={testMode}
+                    onChange={(value) => setTestMode(value)}
+                    options={TEST_MODES}
+                  />
+                </Form.Item>
+                <Row>
+                  <Col span={10}></Col>
+                  <Col span={10}>
+                    <Text className="mb-5 mt-1">
+                      <strong>Optional</strong> - LiteLLM endpoint to use when health checking this model{" "}
+                      <Link href="https://docs.litellm.ai/docs/proxy/health#health" target="_blank">
+                        Learn more
+                      </Link>
+                    </Text>
+                  </Col>
+                </Row>
+
+                {/* Credentials */}
+                <div className="mb-4">
+                  <Typography.Text className="text-sm text-gray-500 mb-2">
+                    Either select existing credentials OR enter new provider credentials below
+                  </Typography.Text>
+                </div>
+
+                <Form.Item label="Existing Credentials" name="litellm_credential_name" initialValue={null}>
+                  <AntdSelect
+                    showSearch
+                    placeholder="Select or search for existing credentials"
+                    optionFilterProp="children"
+                    filterOption={(input, option) => (option?.label ?? "").toLowerCase().includes(input.toLowerCase())}
+                    options={[
+                      { value: null, label: "None" },
+                      ...credentials.map((credential) => ({
+                        value: credential.credential_name,
+                        label: credential.credential_name,
+                      })),
+                    ]}
+                    allowClear
+                  />
+                </Form.Item>
+
+                <Form.Item
+                  noStyle
+                  shouldUpdate={(prevValues, currentValues) =>
+                    prevValues.litellm_credential_name !== currentValues.litellm_credential_name ||
+                    prevValues.provider !== currentValues.provider
+                  }
+                >
+                  {({ getFieldValue }) => {
+                    const credentialName = getFieldValue("litellm_credential_name");
+                    console.log("🔑 Credential Name Changed:", credentialName);
+                    // Only show provider specific fields if no credentials selected
+                    if (!credentialName) {
+                      return (
+                        <>
+                          <div className="flex items-center my-4">
+                            <div className="flex-grow border-t border-gray-200"></div>
+                            <span className="px-4 text-gray-500 text-sm">OR</span>
+                            <div className="flex-grow border-t border-gray-200"></div>
+                          </div>
+                          <ProviderSpecificFields selectedProvider={selectedProvider} uploadProps={uploadProps} />
+                        </>
+                      );
+                    }
+                    return null;
+                  }}
+                </Form.Item>
+                <div className="flex items-center my-4">
+                  <div className="flex-grow border-t border-gray-200"></div>
+                  <span className="px-4 text-gray-500 text-sm">Additional Model Info Settings</span>
+                  <div className="flex-grow border-t border-gray-200"></div>
+                </div>
+                {/* Team-only Model Switch - Only show for proxy admins, not team admins */}
+                {(isAdmin || !isTeamAdmin) && (
+                  <Form.Item
+                    label="Team-BYOK Model"
+                    tooltip="Only use this model + credential combination for this team. Useful when teams want to onboard their own OpenAI keys."
+                    className="mb-4"
+                  >
+                    <Tooltip
+                      title={
+                        !premiumUser
+                          ? "This is an enterprise-only feature. Upgrade to premium to restrict model+credential combinations to a specific team."
+                          : ""
+                      }
+                      placement="top"
+                    >
+                      <Switch
+                        checked={isTeamOnly}
+                        onChange={(checked) => {
+                          setIsTeamOnly(checked);
+                          if (!checked) {
+                            form.setFieldValue("team_id", undefined);
+                          }
+                        }}
+                        disabled={!premiumUser}
+                      />
+                    </Tooltip>
+                  </Form.Item>
+                )}
+
+                {/* Conditional Team Selection */}
+                {isTeamOnly && (isAdmin || !isTeamAdmin) && (
+                  <Form.Item
+                    label="Select Team"
+                    name="team_id"
+                    className="mb-4"
+                    tooltip="Only keys for this team will be able to call this model."
+                    rules={[
+                      {
+                        required: isTeamOnly && !isAdmin,
+                        message: "Please select a team.",
+                      },
+                    ]}
+                  >
+                    <TeamDropdown teams={teams} disabled={!premiumUser} />
+                  </Form.Item>
+                )}
+                {isAdmin && (
+                  <>
+                    <Form.Item
+                      label="Model Access Group"
+                      name="model_access_group"
+                      className="mb-4"
+                      tooltip="Use model access groups to give users access to select models, and add new ones to the group over time."
+                    >
+                      <AntdSelect
+                        mode="tags"
+                        showSearch
+                        placeholder="Select existing groups or type to create new ones"
+                        optionFilterProp="children"
+                        tokenSeparators={[","]}
+                        options={modelAccessGroups.map((group) => ({
+                          value: group,
+                          label: group,
+                        }))}
+                        maxTagCount="responsive"
+                        allowClear
+                      />
+                    </Form.Item>
+                  </>
+                )}
+                <AdvancedSettings
+                  showAdvancedSettings={showAdvancedSettings}
+                  setShowAdvancedSettings={setShowAdvancedSettings}
+                  teams={teams}
+                  guardrailsList={guardrailsList || []}
+                  tagsList={tagsList || {}}
+                />
+              </>
+            )}
+            <div className="flex justify-between items-center mb-4">
+              <Tooltip title="Get help on our github">
+                <Typography.Link href="https://github.com/BerriAI/litellm/issues">Need Help?</Typography.Link>
+              </Tooltip>
+              <div className="space-x-2">
+                <Button onClick={handleTestConnection} loading={isTestingConnection}>
+                  Test Connect
+                </Button>
+                <Button htmlType="submit">Add Model</Button>
+              </div>
+            </div>
+          </>
+        </Form>
+      </Card>
+
+      {/* Test Connection Results Modal */}
+      <Modal
+        title="Connection Test Results"
+        open={isResultModalVisible}
+        onCancel={() => {
+          setIsResultModalVisible(false);
+          setIsTestingConnection(false);
+        }}
+        footer={[
+          <Button
+            key="close"
+            onClick={() => {
+              setIsResultModalVisible(false);
+              setIsTestingConnection(false);
+            }}
+          >
+            Close
+          </Button>,
+        ]}
+        width={700}
+      >
+        {/* Only render the ConnectionErrorDisplay when modal is visible and we have a test ID */}
+        {isResultModalVisible && (
+          <ConnectionErrorDisplay
+            // The key prop tells React to create a fresh component instance when it changes
+            key={connectionTestId}
+            formValues={form.getFieldsValue()}
+            accessToken={accessToken}
+            testMode={testMode}
+            modelName={form.getFieldValue("model_name") || form.getFieldValue("model")}
+            onClose={() => {
+              setIsResultModalVisible(false);
+              setIsTestingConnection(false);
+            }}
+            onTestComplete={() => setIsTestingConnection(false)}
+          />
+        )}
+      </Modal>
+    </>
+  );
+};
+
+export default AddModelForm;

ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx

@@ -62,6 +62,14 @@ vi.mock("@/app/(dashboard)/hooks/providers/useProviderFields", () => ({
   }),
 }));
 
+vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
+  default: vi.fn().mockReturnValue({
+    accessToken: "test-access-token",
+    userRole: "Admin",
+    premiumUser: true,
+  }),
+}));
+
 const createQueryClient = () =>
   new QueryClient({
     defaultOptions: {
@@ -137,7 +145,6 @@ const createTestProps = () => {
     uploadProps,
     accessToken: "test-access-token",
     userRole: "Admin",
-    premiumUser: true,
   };
 };
 
@@ -163,7 +170,6 @@ describe("Add Model Tab", () => {
           credentials={props.credentials}
           accessToken={props.accessToken}
           userRole={props.userRole}
-          premiumUser={props.premiumUser}
         />
       </QueryClientProvider>,
     );
@@ -192,7 +198,6 @@ describe("Add Model Tab", () => {
           credentials={props.credentials}
           accessToken={props.accessToken}
           userRole={props.userRole}
-          premiumUser={props.premiumUser}
         />
       </QueryClientProvider>,
     );
@@ -222,7 +227,6 @@ describe("Add Model Tab", () => {
           credentials={props.credentials}
           accessToken={props.accessToken}
           userRole={props.userRole}
-          premiumUser={props.premiumUser}
         />
       </QueryClientProvider>,
     );
@@ -251,7 +255,6 @@ describe("Add Model Tab", () => {
           credentials={props.credentials}
           accessToken={props.accessToken}
           userRole={props.userRole}
-          premiumUser={props.premiumUser}
         />
       </QueryClientProvider>,
     );
@@ -289,7 +292,6 @@ describe("Add Model Tab", () => {
           credentials={props.credentials}
           accessToken={props.accessToken}
           userRole={props.userRole}
-          premiumUser={props.premiumUser}
         />
       </QueryClientProvider>,
     );

ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx

@@ -1,34 +1,18 @@
-import { useProviderFields } from "@/app/(dashboard)/hooks/providers/useProviderFields";
-import { all_admin_roles } from "@/utils/roles";
-import { Switch, Tab, TabGroup, TabList, TabPanel, TabPanels, Text } from "@tremor/react";
+import { Tab, TabGroup, TabList, TabPanel, TabPanels } from "@tremor/react";
 import type { FormInstance } from "antd";
-import { Select as AntdSelect, Button, Card, Col, Form, Modal, Row, Tooltip, Typography } from "antd";
+import { Form } from "antd";
 import type { UploadProps } from "antd/es/upload";
-import React, { useEffect, useMemo, useState } from "react";
-import TeamDropdown from "../common_components/team_dropdown";
+import React from "react";
 import type { Team } from "../key_team_helpers/key_list";
-import {
-  type CredentialItem,
-  type ProviderCreateInfo,
-  getGuardrailsList,
-  modelAvailableCall,
-  tagListCall,
-} from "../networking";
-import { Providers, providerLogoMap } from "../provider_info_helpers";
-import { ProviderLogo } from "../molecules/models/ProviderLogo";
-import { Tag } from "../tag_management/types";
+import { type CredentialItem } from "../networking";
+import { Providers } from "../provider_info_helpers";
 import AddAutoRouterTab from "./add_auto_router_tab";
-import { TEST_MODES } from "./add_model_modes";
-import AdvancedSettings from "./advanced_settings";
-import ConditionalPublicModelName from "./conditional_public_model_name";
+import AddModelForm from "./AddModelForm";
 import { handleAddAutoRouterSubmit } from "./handle_add_auto_router_submit";
-import LiteLLMModelNameField from "./litellm_model_name";
-import ConnectionErrorDisplay from "./model_connection_test";
-import ProviderSpecificFields from "./provider_specific_fields";
 
 interface AddModelTabProps {
   form: FormInstance; // For the Add Model tab
-  handleOk: () => void;
+  handleOk: (values?: any) => Promise<void>;
   selectedProvider: Providers;
   setSelectedProvider: (provider: Providers) => void;
   providerModels: string[];
@@ -41,11 +25,8 @@ interface AddModelTabProps {
   credentials: CredentialItem[];
   accessToken: string;
   userRole: string;
-  premiumUser: boolean;
 }
 
-const { Title, Link } = Typography;
-
 const AddModelTab: React.FC<AddModelTabProps> = ({
   form,
   handleOk,
@@ -61,90 +42,9 @@ const AddModelTab: React.FC<AddModelTabProps> = ({
   credentials,
   accessToken,
   userRole,
-  premiumUser,
 }) => {
   // Create separate form instance for auto router
   const [autoRouterForm] = Form.useForm();
-  // State for test mode and connection testing
-  const [testMode, setTestMode] = useState<string>("chat");
-  const [isResultModalVisible, setIsResultModalVisible] = useState<boolean>(false);
-  const [isTestingConnection, setIsTestingConnection] = useState<boolean>(false);
-  const [guardrailsList, setGuardrailsList] = useState<string[]>([]);
-  const [tagsList, setTagsList] = useState<Record<string, Tag>>({});
-  // Using a unique ID to force the ConnectionErrorDisplay to remount and run a fresh test
-  const [connectionTestId, setConnectionTestId] = useState<string>("");
-
-  // Provider metadata for driving the provider select from backend config
-  const {
-    data: providerMetadata,
-    isLoading: isProviderMetadataLoading,
-    error: providerMetadataError,
-  } = useProviderFields();
-
-  useEffect(() => {
-    const fetchGuardrails = async () => {
-      try {
-        const response = await getGuardrailsList(accessToken);
-        const guardrailNames = response.guardrails.map((g: { guardrail_name: string }) => g.guardrail_name);
-        setGuardrailsList(guardrailNames);
-      } catch (error) {
-        console.error("Failed to fetch guardrails:", error);
-      }
-    };
-
-    fetchGuardrails();
-  }, [accessToken]);
-
-  useEffect(() => {
-    const fetchTags = async () => {
-      try {
-        const response = await tagListCall(accessToken);
-        setTagsList(response);
-      } catch (error) {
-        console.error("Failed to fetch tags:", error);
-      }
-    };
-
-    fetchTags();
-  }, [accessToken]);
-
-  // Test connection when button is clicked
-  const handleTestConnection = async () => {
-    setIsTestingConnection(true);
-    // Generate a new test ID (using timestamp for uniqueness)
-    // This forces React to create a new instance of ConnectionErrorDisplay
-    setConnectionTestId(`test-${Date.now()}`);
-    // Show the modal with the fresh test
-    setIsResultModalVisible(true);
-  };
-
-  // State for team-only switch
-  const [isTeamOnly, setIsTeamOnly] = useState<boolean>(false);
-
-  const [modelAccessGroups, setModelAccessGroups] = useState<string[]>([]);
-
-  useEffect(() => {
-    const fetchModelAccessGroups = async () => {
-      const response = await modelAvailableCall(accessToken, "", "", false, null, true, true);
-      setModelAccessGroups(response["data"].map((model: any) => model["id"]));
-    };
-    fetchModelAccessGroups();
-  }, [accessToken]);
-
-  const sortedProviderMetadata: ProviderCreateInfo[] = useMemo(() => {
-    if (!providerMetadata) {
-      return [];
-    }
-    return [...providerMetadata].sort((a, b) => a.provider_display_name.localeCompare(b.provider_display_name));
-  }, [providerMetadata]);
-
-  const providerMetadataErrorText = providerMetadataError
-    ? providerMetadataError instanceof Error
-      ? providerMetadataError.message
-      : "Failed to load providers"
-    : null;
-
-  const isAdmin = all_admin_roles.includes(userRole);
 
   const handleAutoRouterOk = () => {
     autoRouterForm
@@ -166,247 +66,20 @@ const AddModelTab: React.FC<AddModelTabProps> = ({
         </TabList>
         <TabPanels>
           <TabPanel>
-            <Title level={2}>Add Model</Title>
-            <Card>
-              <Form
-                form={form}
-                onFinish={(values) => {
-                  console.log("🔥 Form onFinish triggered with values:", values);
-                  handleOk();
-                }}
-                onFinishFailed={(errorInfo) => {
-                  console.log("💥 Form onFinishFailed triggered:", errorInfo);
-                }}
-                labelCol={{ span: 10 }}
-                wrapperCol={{ span: 16 }}
-                labelAlign="left"
-              >
-                <>
-                  {/* Provider Selection */}
-                  <Form.Item
-                    rules={[{ required: true, message: "Required" }]}
-                    label="Provider:"
-                    name="custom_llm_provider"
-                    tooltip="E.g. OpenAI, Azure OpenAI, Anthropic, Bedrock, etc."
-                    labelCol={{ span: 10 }}
-                    labelAlign="left"
-                  >
-                    <AntdSelect
-                      virtual={false}
-                      showSearch
-                      loading={isProviderMetadataLoading}
-                      placeholder={isProviderMetadataLoading ? "Loading providers..." : "Select a provider"}
-                      optionFilterProp="data-label"
-                      onChange={(value) => {
-                        setSelectedProvider(value as Providers);
-                        setProviderModelsFn(value as Providers);
-                        form.setFieldsValue({
-                          custom_llm_provider: value,
-                        });
-                        form.setFieldsValue({
-                          model: [],
-                          model_name: undefined,
-                        });
-                      }}
-                    >
-                      {providerMetadataErrorText && sortedProviderMetadata.length === 0 && (
-                        <AntdSelect.Option key="__error" value="">
-                          {providerMetadataErrorText}
-                        </AntdSelect.Option>
-                      )}
-                      {sortedProviderMetadata.map((providerInfo) => {
-                        const displayName = providerInfo.provider_display_name;
-                        const providerKey = providerInfo.provider;
-                        const logoSrc = providerLogoMap[displayName] ?? "";
-
-                        return (
-                          <AntdSelect.Option key={providerKey} value={providerKey} data-label={displayName}>
-                            <div className="flex items-center space-x-2">
-                              <ProviderLogo provider={providerKey} className="w-5 h-5" />
-                              <span>{displayName}</span>
-                            </div>
-                          </AntdSelect.Option>
-                        );
-                      })}
-                    </AntdSelect>
-                  </Form.Item>
-                  <LiteLLMModelNameField
-                    selectedProvider={selectedProvider}
-                    providerModels={providerModels}
-                    getPlaceholder={getPlaceholder}
-                  />
-
-                  {/* Conditionally Render "Public Model Name" */}
-                  <ConditionalPublicModelName />
-
-                  {/* Select Mode */}
-                  <Form.Item label="Mode" name="mode" className="mb-1">
-                    <AntdSelect
-                      style={{ width: "100%" }}
-                      value={testMode}
-                      onChange={(value) => setTestMode(value)}
-                      options={TEST_MODES}
-                    />
-                  </Form.Item>
-                  <Row>
-                    <Col span={10}></Col>
-                    <Col span={10}>
-                      <Text className="mb-5 mt-1">
-                        <strong>Optional</strong> - LiteLLM endpoint to use when health checking this model{" "}
-                        <Link href="https://docs.litellm.ai/docs/proxy/health#health" target="_blank">
-                          Learn more
-                        </Link>
-                      </Text>
-                    </Col>
-                  </Row>
-
-                  {/* Credentials */}
-                  <div className="mb-4">
-                    <Typography.Text className="text-sm text-gray-500 mb-2">
-                      Either select existing credentials OR enter new provider credentials below
-                    </Typography.Text>
-                  </div>
-
-                  <Form.Item label="Existing Credentials" name="litellm_credential_name" initialValue={null}>
-                    <AntdSelect
-                      showSearch
-                      placeholder="Select or search for existing credentials"
-                      optionFilterProp="children"
-                      filterOption={(input, option) =>
-                        (option?.label ?? "").toLowerCase().includes(input.toLowerCase())
-                      }
-                      options={[
-                        { value: null, label: "None" },
-                        ...credentials.map((credential) => ({
-                          value: credential.credential_name,
-                          label: credential.credential_name,
-                        })),
-                      ]}
-                      allowClear
-                    />
-                  </Form.Item>
-
-                  <Form.Item
-                    noStyle
-                    shouldUpdate={(prevValues, currentValues) =>
-                      prevValues.litellm_credential_name !== currentValues.litellm_credential_name ||
-                      prevValues.provider !== currentValues.provider
-                    }
-                  >
-                    {({ getFieldValue }) => {
-                      const credentialName = getFieldValue("litellm_credential_name");
-                      console.log("🔑 Credential Name Changed:", credentialName);
-                      // Only show provider specific fields if no credentials selected
-                      if (!credentialName) {
-                        return (
-                          <>
-                            <div className="flex items-center my-4">
-                              <div className="flex-grow border-t border-gray-200"></div>
-                              <span className="px-4 text-gray-500 text-sm">OR</span>
-                              <div className="flex-grow border-t border-gray-200"></div>
-                            </div>
-                            <ProviderSpecificFields selectedProvider={selectedProvider} uploadProps={uploadProps} />
-                          </>
-                        );
-                      }
-                      return null;
-                    }}
-                  </Form.Item>
-                  <div className="flex items-center my-4">
-                    <div className="flex-grow border-t border-gray-200"></div>
-                    <span className="px-4 text-gray-500 text-sm">Additional Model Info Settings</span>
-                    <div className="flex-grow border-t border-gray-200"></div>
-                  </div>
-                  {/* Team-only Model Switch */}
-                  <Form.Item
-                    label="Team-BYOK Model"
-                    tooltip="Only use this model + credential combination for this team. Useful when teams want to onboard their own OpenAI keys."
-                    className="mb-4"
-                  >
-                    <Tooltip
-                      title={
-                        !premiumUser
-                          ? "This is an enterprise-only feature. Upgrade to premium to restrict model+credential combinations to a specific team."
-                          : ""
-                      }
-                      placement="top"
-                    >
-                      <Switch
-                        checked={isTeamOnly}
-                        onChange={(checked) => {
-                          setIsTeamOnly(checked);
-                          if (!checked) {
-                            form.setFieldValue("team_id", undefined);
-                          }
-                        }}
-                        disabled={!premiumUser}
-                      />
-                    </Tooltip>
-                  </Form.Item>
-
-                  {/* Conditional Team Selection */}
-                  {isTeamOnly && (
-                    <Form.Item
-                      label="Select Team"
-                      name="team_id"
-                      className="mb-4"
-                      tooltip="Only keys for this team will be able to call this model."
-                      rules={[
-                        {
-                          required: isTeamOnly && !isAdmin,
-                          message: "Please select a team.",
-                        },
-                      ]}
-                    >
-                      <TeamDropdown teams={teams} disabled={!premiumUser} />
-                    </Form.Item>
-                  )}
-                  {isAdmin && (
-                    <>
-                      <Form.Item
-                        label="Model Access Group"
-                        name="model_access_group"
-                        className="mb-4"
-                        tooltip="Use model access groups to give users access to select models, and add new ones to the group over time."
-                      >
-                        <AntdSelect
-                          mode="tags"
-                          showSearch
-                          placeholder="Select existing groups or type to create new ones"
-                          optionFilterProp="children"
-                          tokenSeparators={[","]}
-                          options={modelAccessGroups.map((group) => ({
-                            value: group,
-                            label: group,
-                          }))}
-                          maxTagCount="responsive"
-                          allowClear
-                        />
-                      </Form.Item>
-                    </>
-                  )}
-                  <AdvancedSettings
-                    showAdvancedSettings={showAdvancedSettings}
-                    setShowAdvancedSettings={setShowAdvancedSettings}
-                    teams={teams}
-                    guardrailsList={guardrailsList}
-                    tagsList={tagsList}
-                  />
-
-                  <div className="flex justify-between items-center mb-4">
-                    <Tooltip title="Get help on our github">
-                      <Typography.Link href="https://github.com/BerriAI/litellm/issues">Need Help?</Typography.Link>
-                    </Tooltip>
-                    <div className="space-x-2">
-                      <Button onClick={handleTestConnection} loading={isTestingConnection}>
-                        Test Connect
-                      </Button>
-                      <Button htmlType="submit">Add Model</Button>
-                    </div>
-                  </div>
-                </>
-              </Form>
-            </Card>
+            <AddModelForm
+              form={form}
+              handleOk={handleOk}
+              selectedProvider={selectedProvider}
+              setSelectedProvider={setSelectedProvider}
+              providerModels={providerModels}
+              setProviderModelsFn={setProviderModelsFn}
+              getPlaceholder={getPlaceholder}
+              uploadProps={uploadProps}
+              showAdvancedSettings={showAdvancedSettings}
+              setShowAdvancedSettings={setShowAdvancedSettings}
+              teams={teams}
+              credentials={credentials}
+            />
           </TabPanel>
           <TabPanel>
             <AddAutoRouterTab
@@ -418,45 +91,6 @@ const AddModelTab: React.FC<AddModelTabProps> = ({
           </TabPanel>
         </TabPanels>
       </TabGroup>
-
-      {/* Test Connection Results Modal */}
-      <Modal
-        title="Connection Test Results"
-        open={isResultModalVisible}
-        onCancel={() => {
-          setIsResultModalVisible(false);
-          setIsTestingConnection(false);
-        }}
-        footer={[
-          <Button
-            key="close"
-            onClick={() => {
-              setIsResultModalVisible(false);
-              setIsTestingConnection(false);
-            }}
-          >
-            Close
-          </Button>,
-        ]}
-        width={700}
-      >
-        {/* Only render the ConnectionErrorDisplay when modal is visible and we have a test ID */}
-        {isResultModalVisible && (
-          <ConnectionErrorDisplay
-            // The key prop tells React to create a fresh component instance when it changes
-            key={connectionTestId}
-            formValues={form.getFieldsValue()}
-            accessToken={accessToken}
-            testMode={testMode}
-            modelName={form.getFieldValue("model_name") || form.getFieldValue("model")}
-            onClose={() => {
-              setIsResultModalVisible(false);
-              setIsTestingConnection(false);
-            }}
-            onTestComplete={() => setIsTestingConnection(false)}
-          />
-        )}
-      </Modal>
     </>
   );
 };

ui/litellm-dashboard/src/components/templates/model_dashboard.tsx

@@ -950,7 +950,7 @@ const OldModelDashboard: React.FC<ModelDashboardProps> = ({
     );
   };
 
-  const handleOk = () => {
+  const handleOk = async () => {
     console.log("🚀 handleOk called from model dashboard!");
     console.log("Current form values:", addModelForm.getFieldsValue());
 
@@ -1354,7 +1354,6 @@ const OldModelDashboard: React.FC<ModelDashboardProps> = ({
                     credentials={credentialsList}
                     accessToken={accessToken}
                     userRole={userRole}
-                    premiumUser={premiumUser}
                   />
                 </TabPanel>
                 <TabPanel>