diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/providers/useProviderFields.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/providers/useProviderFields.ts new file mode 100644 index 00000000000..5d219f3183e --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/providers/useProviderFields.ts @@ -0,0 +1,14 @@ +import { getProviderCreateMetadata, ProviderCreateInfo } from "@/components/networking"; +import { useQuery } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; + +const providerFieldsKeys = createQueryKeys("providerFields"); + +export const useProviderFields = () => { + return useQuery({ + queryKey: providerFieldsKeys.list({}), + queryFn: async () => await getProviderCreateMetadata(), + staleTime: 24 * 60 * 60 * 1000, // 24 hours - data rarely changes + gcTime: 24 * 60 * 60 * 1000, // 24 hours - keep in cache for 24 hours + }); +}; diff --git a/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx index 87c5b4c0aa0..201bbb3c180 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx @@ -1,13 +1,13 @@ -import { render, renderHook, waitFor } from "@testing-library/react"; -import { describe, it, vi, expect } from "vitest"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render, renderHook, screen } from "@testing-library/react"; import { Form } from "antd"; -import AddModelTab from "./add_model_tab"; -import { Providers } from "../provider_info_helpers"; +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 type { UploadProps } from "antd/es/upload"; +import { Providers } from "../provider_info_helpers"; +import AddModelTab from "./add_model_tab"; -// Mock the networking module vi.mock("../networking", async () => { const actual = await vi.importActual("../networking"); return { @@ -19,6 +19,12 @@ vi.mock("../networking", async () => { 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", @@ -31,91 +37,198 @@ vi.mock("../networking", async () => { }; }); -describe("Add Model Tab", () => { - it( - "should render", - async () => { - // Create a form instance using renderHook - const { result } = renderHook(() => Form.useForm()); - const [form] = result.current; - - // Mock functions - const handleOk = vi.fn(); - const setSelectedProvider = vi.fn(); - const setProviderModelsFn = vi.fn(); - const getPlaceholder = vi.fn((provider: Providers) => `Enter ${provider} model name`); - const setShowAdvancedSettings = vi.fn(); - - // Mock data - const selectedProvider = Providers.OpenAI; - const providerModels = ["gpt-4", "gpt-3.5-turbo"]; - const showAdvancedSettings = false; - - const teams: 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 credentials: CredentialItem[] = [ - { - credential_name: "test-credential", - credential_values: {}, - credential_info: { - custom_llm_provider: "openai", - description: "Test credential", - }, - }, - ]; - - const uploadProps: UploadProps = { - beforeUpload: () => false, - showUploadList: false, - }; - - const accessToken = "test-access-token"; - const userRole = "Admin"; - const premiumUser = true; - - const { container, findByText } = render( - , - ); - - // Wait for the tabs to render which indicates the component loaded - await waitFor( - () => { - const tabs = container.querySelectorAll('[role="tab"]'); - expect(tabs.length).toBeGreaterThan(0); - }, - { timeout: 10000 }, - ); +const createQueryClient = () => + new QueryClient({ + defaultOptions: { + queries: { + retry: false, + gcTime: 0, + }, }, - 15000, - ); + }); + +const createTestProps = () => { + const { result } = renderHook(() => Form.useForm()); + const [form] = result.current; + + const handleOk = vi.fn(); + const setSelectedProvider = vi.fn(); + const setProviderModelsFn = vi.fn(); + const getPlaceholder = vi.fn((provider: Providers) => `Enter ${provider} model name`); + const setShowAdvancedSettings = vi.fn(); + + const selectedProvider = Providers.OpenAI; + const providerModels = ["gpt-4", "gpt-3.5-turbo"]; + const showAdvancedSettings = false; + + const teams: 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 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, + setSelectedProvider, + setProviderModelsFn, + getPlaceholder, + setShowAdvancedSettings, + selectedProvider, + providerModels, + showAdvancedSettings, + teams, + credentials, + uploadProps, + accessToken: "test-access-token", + userRole: "Admin", + premiumUser: true, + }; +}; + +describe("Add Model Tab", () => { + it("should render", async () => { + const props = createTestProps(); + const queryClient = createQueryClient(); + + render( + + + , + ); + + expect(await screen.findByRole("tab", { name: "Add Model" })).toBeInTheDocument(); + }); + + it("should display both Add Model and Add Auto Router tabs", async () => { + const props = createTestProps(); + const queryClient = createQueryClient(); + + render( + + + , + ); + + expect(await screen.findByRole("tab", { name: "Add Model" })).toBeInTheDocument(); + expect(await screen.findByRole("tab", { name: "Add Auto Router" })).toBeInTheDocument(); + }); + + it("should display provider selection field", async () => { + const props = createTestProps(); + const queryClient = createQueryClient(); + + render( + + + , + ); + + expect(await screen.findByText("Provider")).toBeInTheDocument(); + }); + + it("should display Test Connect and Add Model buttons", async () => { + const props = createTestProps(); + const queryClient = createQueryClient(); + + render( + + + , + ); + + const testConnectButtons = await screen.findAllByRole("button", { name: "Test Connect" }); + expect(testConnectButtons.length).toBeGreaterThan(0); + expect(await screen.findByRole("button", { name: "Add Model" })).toBeInTheDocument(); + }, 10000); // 10 seconds timeout for complex logic }); diff --git a/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx index 4efcd7be907..b2e1dec2827 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx @@ -1,31 +1,29 @@ -import React, { useEffect, useMemo, useState } from "react"; -import { Card, Form, Button, Tooltip, Typography, Select as AntdSelect, Modal } from "antd"; +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 type { FormInstance } from "antd"; +import { Select as AntdSelect, Button, Card, Col, Form, Modal, Row, Tooltip, Typography } from "antd"; import type { UploadProps } from "antd/es/upload"; -import { TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react"; -import LiteLLMModelNameField from "./litellm_model_name"; -import ConditionalPublicModelName from "./conditional_public_model_name"; -import ProviderSpecificFields from "./provider_specific_fields"; -import AdvancedSettings from "./advanced_settings"; -import { Providers, providerLogoMap } from "../provider_info_helpers"; +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, getGuardrailsList, - getProviderCreateMetadata, modelAvailableCall, tagListCall, } from "../networking"; -import ConnectionErrorDisplay from "./model_connection_test"; -import { TEST_MODES } from "./add_model_modes"; -import { Row, Col } from "antd"; -import { Text, Switch } from "@tremor/react"; -import TeamDropdown from "../common_components/team_dropdown"; -import { all_admin_roles } from "@/utils/roles"; -import AddAutoRouterTab from "./add_auto_router_tab"; -import { handleAddAutoRouterSubmit } from "./handle_add_auto_router_submit"; +import { Providers, providerLogoMap } from "../provider_info_helpers"; import { Tag } from "../tag_management/types"; +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 { 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 @@ -76,9 +74,11 @@ const AddModelTab: React.FC = ({ const [connectionTestId, setConnectionTestId] = useState(""); // Provider metadata for driving the provider select from backend config - const [providerMetadata, setProviderMetadata] = useState(null); - const [isProviderMetadataLoading, setIsProviderMetadataLoading] = useState(false); - const [providerMetadataError, setProviderMetadataError] = useState(null); + const { + data: providerMetadata, + isLoading: isProviderMetadataLoading, + error: providerMetadataError, + } = useProviderFields(); useEffect(() => { const fetchGuardrails = async () => { @@ -107,37 +107,6 @@ const AddModelTab: React.FC = ({ fetchTags(); }, [accessToken]); - useEffect(() => { - let isMounted = true; - - const fetchProviderMetadata = async () => { - setIsProviderMetadataLoading(true); - setProviderMetadataError(null); - try { - const metadata = await getProviderCreateMetadata(); - if (!isMounted) { - return; - } - setProviderMetadata(metadata); - } catch (error) { - console.error("Failed to fetch provider metadata:", error); - if (isMounted) { - setProviderMetadataError("Failed to load providers"); - } - } finally { - if (isMounted) { - setIsProviderMetadataLoading(false); - } - } - }; - - fetchProviderMetadata(); - - return () => { - isMounted = false; - }; - }, []); - // Test connection when button is clicked const handleTestConnection = async () => { setIsTestingConnection(true); @@ -168,6 +137,12 @@ const AddModelTab: React.FC = ({ 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 = () => { @@ -232,9 +207,9 @@ const AddModelTab: React.FC = ({ }); }} > - {providerMetadataError && sortedProviderMetadata.length === 0 && ( + {providerMetadataErrorText && sortedProviderMetadata.length === 0 && ( - {providerMetadataError} + {providerMetadataErrorText} )} {sortedProviderMetadata.map((providerInfo) => { diff --git a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.test.tsx b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.test.tsx index ab9b3e92e98..4590121acf2 100644 --- a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.test.tsx @@ -1,6 +1,7 @@ -import { render, waitFor } from "@testing-library/react"; -import { describe, it, expect, beforeAll, vi } from "vitest"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render, screen, waitFor } from "@testing-library/react"; import { Form } from "antd"; +import { beforeAll, describe, expect, it, vi } from "vitest"; import { Providers } from "../provider_info_helpers"; import ProviderSpecificFields from "./provider_specific_fields"; @@ -97,7 +98,6 @@ vi.mock("../networking", async () => { }; }); -// Mock window.matchMedia for Ant Design components beforeAll(() => { Object.defineProperty(window, "matchMedia", { writable: true, @@ -105,8 +105,8 @@ beforeAll(() => { matches: false, media: query, onchange: null, - addListener: () => {}, // deprecated - removeListener: () => {}, // deprecated + addListener: () => {}, + removeListener: () => {}, addEventListener: () => {}, removeEventListener: () => {}, dispatchEvent: () => false, @@ -114,80 +114,105 @@ beforeAll(() => { }); }); +const createQueryClient = () => + new QueryClient({ + defaultOptions: { + queries: { + retry: false, + gcTime: 0, + }, + }, + }); + describe("ProviderSpecificFields", () => { - it("should render the provider specific fields for OpenAI", async () => { - const { getByLabelText, getByPlaceholderText } = render( -
- - , + it("should render", async () => { + const queryClient = createQueryClient(); + render( + +
+ + +
, ); await waitFor(() => { - // Check for the API Base text input - const apiBaseInput = getByPlaceholderText("https://api.openai.com/v1"); + expect(screen.getByLabelText("OpenAI API Key")).toBeInTheDocument(); + }); + }); + + it("should render the provider specific fields for OpenAI", async () => { + const queryClient = createQueryClient(); + render( + +
+ + +
, + ); + + await waitFor(() => { + const apiKeyLabel = screen.getByLabelText("OpenAI API Key"); + expect(apiKeyLabel).toBeInTheDocument(); + + const apiBaseInput = screen.getByPlaceholderText("https://api.openai.com/v1"); expect(apiBaseInput).toBeInTheDocument(); expect(apiBaseInput).toHaveAttribute("type", "text"); - // Check for Organization field - const orgInput = getByPlaceholderText("[OPTIONAL] my-unique-org"); + const orgInput = screen.getByPlaceholderText("[OPTIONAL] my-unique-org"); expect(orgInput).toBeInTheDocument(); - - // Check for API Key field - const apiKeyLabel = getByLabelText("OpenAI API Key"); - expect(apiKeyLabel).toBeInTheDocument(); }); }); it("should render the provider specific fields for vLLM", async () => { - const { getByLabelText, getByPlaceholderText } = render( -
- - , + const queryClient = createQueryClient(); + render( + +
+ + +
, ); await waitFor(() => { - const apiBaseInput = getByPlaceholderText("https://..."); + const apiKeyLabel = screen.getByLabelText("vLLM API Key"); + expect(apiKeyLabel).toBeInTheDocument(); + + const apiBaseInput = screen.getByPlaceholderText("https://..."); expect(apiBaseInput).toBeInTheDocument(); expect(apiBaseInput).toHaveAttribute("type", "text"); - - // Check for API Key field - const apiKeyLabel = getByLabelText("vLLM API Key"); - expect(apiKeyLabel).toBeInTheDocument(); }); }); it("should render the provider specific fields for Azure", async () => { - const { getByLabelText, getByPlaceholderText } = render( -
- - , + const queryClient = createQueryClient(); + render( + +
+ + +
, ); await waitFor(() => { - // Check for API Base field - const apiBaseInput = getByPlaceholderText("https://..."); - expect(apiBaseInput).toBeInTheDocument(); - expect(apiBaseInput).toHaveAttribute("type", "text"); - - // Check for API Version field - const apiVersionInput = getByPlaceholderText("2023-07-01-preview"); - expect(apiVersionInput).toBeInTheDocument(); - - // Check for Base Model field - const baseModelInput = getByPlaceholderText("azure/gpt-3.5-turbo"); - expect(baseModelInput).toBeInTheDocument(); - - // Check for API Key field - const apiKeyInput = getByLabelText("Azure API Key"); + const apiKeyInput = screen.getByLabelText("Azure API Key"); expect(apiKeyInput).toBeInTheDocument(); expect(apiKeyInput).toHaveAttribute("type", "password"); expect(apiKeyInput).toHaveAttribute("placeholder", "Enter your Azure API Key"); - // Check for Azure AD Token field - const azureAdTokenInput = getByLabelText("Azure AD Token"); + const azureAdTokenInput = screen.getByLabelText("Azure AD Token"); expect(azureAdTokenInput).toBeInTheDocument(); expect(azureAdTokenInput).toHaveAttribute("type", "password"); expect(azureAdTokenInput).toHaveAttribute("placeholder", "Enter your Azure AD Token"); + + const apiBaseInput = screen.getByPlaceholderText("https://..."); + expect(apiBaseInput).toBeInTheDocument(); + expect(apiBaseInput).toHaveAttribute("type", "text"); + + const apiVersionInput = screen.getByPlaceholderText("2023-07-01-preview"); + expect(apiVersionInput).toBeInTheDocument(); + + const baseModelInput = screen.getByPlaceholderText("azure/gpt-3.5-turbo"); + expect(baseModelInput).toBeInTheDocument(); }); }); }); diff --git a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx index c1b17ff4419..892e9e0197e 100644 --- a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx +++ b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx @@ -1,15 +1,10 @@ -import React from "react"; -import { Form, Select } from "antd"; -import { TextInput, Text } from "@tremor/react"; -import { Row, Col, Typography, Button as Button2, Upload, UploadProps } from "antd"; +import { useProviderFields } from "@/app/(dashboard)/hooks/providers/useProviderFields"; import { UploadOutlined } from "@ant-design/icons"; +import { Text, TextInput } from "@tremor/react"; +import { Button as Button2, Col, Form, Row, Select, Typography, Upload, UploadProps } from "antd"; +import React from "react"; +import { CredentialItem, ProviderCredentialFieldMetadata } from "../networking"; import { provider_map, Providers } from "../provider_info_helpers"; -import { - CredentialItem, - ProviderCreateInfo, - ProviderCredentialFieldMetadata, - getProviderCreateMetadata, -} from "../networking"; const { Link } = Typography; interface ProviderSpecificFieldsProps { @@ -99,65 +94,42 @@ const ProviderSpecificFields: React.FC = ({ selecte const selectedProviderEnum = Providers[selectedProvider as keyof typeof Providers] as Providers; const form = Form.useFormInstance(); // Get form instance from context - const [providerMetadata, setProviderMetadata] = React.useState(null); - const [isLoading, setIsLoading] = React.useState(false); - const [loadError, setLoadError] = React.useState(null); + const { data: providerMetadata, isLoading, error: loadError } = useProviderFields(); + // Memoize the expensive cache computation + const cacheEntries = React.useMemo(() => { + if (!providerMetadata) { + return null; + } + + // Compute cache entries keyed by provider display name and identifiers + const entries: Record = {}; + providerMetadata.forEach((providerInfo) => { + const displayName = providerInfo.provider_display_name; + const mappedFields = providerInfo.credential_fields.map(mapFieldMetadataToUiField); + + // Primary key: human-readable display name + entries[displayName] = mappedFields; + + // Also cache by backend identifiers so lookups by provider slug work + if (providerInfo.provider) { + entries[providerInfo.provider] = mappedFields; + } + if (providerInfo.litellm_provider) { + entries[providerInfo.litellm_provider] = mappedFields; + } + }); + return entries; + }, [providerMetadata]); + + // Sync memoized cache entries to module-level cache React.useEffect(() => { - const hasCachedFields = Object.keys(providerFieldsByDisplayName).length > 0; - if (hasCachedFields) { - // We already have fields cached globally; no need to refetch. - // This is important so we can reuse credential field definitions - // across mounts and in non-React helpers. + if (!cacheEntries) { return; } - let isMounted = true; - - const fetchProviderFields = async () => { - setIsLoading(true); - setLoadError(null); - try { - const metadata = await getProviderCreateMetadata(); - if (!isMounted) { - return; - } - setProviderMetadata(metadata); - - // Populate cache keyed by provider display name and identifiers - metadata.forEach((providerInfo) => { - const displayName = providerInfo.provider_display_name; - const mappedFields = providerInfo.credential_fields.map(mapFieldMetadataToUiField); - - // Primary key: human-readable display name - providerFieldsByDisplayName[displayName] = mappedFields; - - // Also cache by backend identifiers so lookups by provider slug work - if (providerInfo.provider) { - providerFieldsByDisplayName[providerInfo.provider] = mappedFields; - } - if (providerInfo.litellm_provider) { - providerFieldsByDisplayName[providerInfo.litellm_provider] = mappedFields; - } - }); - } catch (error) { - console.error("Failed to load provider credential fields:", error); - if (isMounted) { - setLoadError("Failed to load provider credential fields"); - } - } finally { - if (isMounted) { - setIsLoading(false); - } - } - }; - - fetchProviderFields(); - - return () => { - isMounted = false; - }; - }, []); + Object.assign(providerFieldsByDisplayName, cacheEntries); + }, [cacheEntries]); const allFields = React.useMemo(() => { // First try to resolve from the in-memory cache. We support both the @@ -234,7 +206,9 @@ const ProviderSpecificFields: React.FC = ({ selecte {loadError && allFields.length === 0 && ( - {loadError} + + {loadError instanceof Error ? loadError.message : "Failed to load provider credential fields"} + )}