From 24ec0bd74fed1b92337e7bbdc2b287eb6f884b8a Mon Sep 17 00:00:00 2001 From: YankhangSIT <112969742+YankhangSIT@users.noreply.github.com> Date: Wed, 17 Jun 2026 23:04:09 +0800 Subject: [PATCH] fix(ui): test complexity auto router connections --- .../add_model/add_auto_router_tab.test.tsx | 88 +++++++++++++++++ .../add_model/add_auto_router_tab.tsx | 33 ++++++- .../handle_add_auto_router_submit.tsx | 98 ++++++++++--------- .../add_model/model_connection_test.tsx | 34 +++++-- 4 files changed, 200 insertions(+), 53 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx new file mode 100644 index 00000000000..86d045525f6 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -0,0 +1,88 @@ +import { Form } from "antd"; +import userEvent from "@testing-library/user-event"; +import { renderWithProviders, screen, waitFor } from "../../../tests/test-utils"; +import { describe, expect, it, vi } from "vitest"; +import AddAutoRouterTab from "./add_auto_router_tab"; +import { modelAvailableCall, testConnectionRequest } from "../networking"; +import { fetchAvailableModels } from "@/components/llm_calls/fetch_models"; + +vi.mock("../networking", () => ({ + modelAvailableCall: vi.fn(), + testConnectionRequest: vi.fn(), +})); + +vi.mock("@/components/llm_calls/fetch_models", () => ({ + fetchAvailableModels: vi.fn(), +})); + +vi.mock("../molecules/notifications_manager", () => ({ + default: { + fromBackend: vi.fn(), + success: vi.fn(), + }, +})); + +vi.mock("./ComplexityRouterConfig", () => ({ + default: ({ onChange }: { onChange: (tiers: Record) => void }) => ( + + ), +})); + +const renderAddAutoRouterTab = () => { + const Component = () => { + const [form] = Form.useForm(); + return ; + }; + + return renderWithProviders(); +}; + +describe("AddAutoRouterTab", () => { + it("tests complexity router connections with auto-router params", async () => { + vi.mocked(modelAvailableCall).mockResolvedValue({ data: [] }); + vi.mocked(fetchAvailableModels).mockResolvedValue([ + { model_group: "gpt-3.5-turbo" }, + { model_group: "gpt-4o-mini" }, + ]); + vi.mocked(testConnectionRequest).mockResolvedValue({ status: "success" }); + + const user = userEvent.setup(); + renderAddAutoRouterTab(); + + await user.type(screen.getByPlaceholderText("e.g., smart_router, auto_router_1"), "smart_router"); + await user.click(screen.getByRole("button", { name: "Select complexity tiers" })); + await user.click(screen.getByRole("button", { name: "Test Connection" })); + + await waitFor(() => + expect(testConnectionRequest).toHaveBeenCalledWith( + "test-token", + { + model: "auto_router/complexity_router", + complexity_router_config: { + tiers: { + SIMPLE: "gpt-3.5-turbo", + MEDIUM: "gpt-4o-mini", + COMPLEX: "", + REASONING: "", + }, + }, + complexity_router_default_model: "gpt-4o-mini", + }, + {}, + undefined, + ), + ); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index 9e9a097e3fe..b14f658347a 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -5,7 +5,7 @@ import { Text, TextInput } from "@tremor/react"; import { modelAvailableCall } from "../networking"; import ConnectionErrorDisplay from "./model_connection_test"; import { all_admin_roles } from "@/utils/roles"; -import { handleAddAutoRouterSubmit } from "./handle_add_auto_router_submit"; +import { buildAutoRouterModelConfig, handleAddAutoRouterSubmit } from "./handle_add_auto_router_submit"; import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; import RouterConfigBuilder from "./RouterConfigBuilder"; import ComplexityRouterConfig from "./ComplexityRouterConfig"; @@ -20,6 +20,7 @@ interface AddAutoRouterTabProps { } type RouterType = "complexity" | "semantic"; +type ConnectionParams = Record; interface ComplexityTiers { SIMPLE: string; @@ -35,6 +36,14 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc const [isResultModalVisible, setIsResultModalVisible] = useState(false); const [isTestingConnection, setIsTestingConnection] = useState(false); const [connectionTestId, setConnectionTestId] = useState(""); + const [preparedConnectionRequest, setPreparedConnectionRequest] = useState< + | { + litellmParamsObj: ConnectionParams; + modelInfoObj: ConnectionParams; + mode?: string; + } + | undefined + >(); const [modelAccessGroups, setModelAccessGroups] = useState([]); const [modelInfo, setModelInfo] = useState([]); @@ -80,6 +89,27 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc // Test connection when button is clicked const handleTestConnection = async () => { + const currentFormValues = form.getFieldsValue(); + setPreparedConnectionRequest(undefined); + + if (routerType === "complexity") { + const defaultModel = + complexityTiers.MEDIUM || complexityTiers.SIMPLE || complexityTiers.COMPLEX || complexityTiers.REASONING; + const autoRouterConfig = buildAutoRouterModelConfig({ + ...currentFormValues, + auto_router_default_model: defaultModel, + model_type: "complexity_router", + complexity_router_config: { + tiers: complexityTiers, + }, + }); + + setPreparedConnectionRequest({ + litellmParamsObj: autoRouterConfig.litellm_params as ConnectionParams, + modelInfoObj: (autoRouterConfig.model_info || {}) as ConnectionParams, + }); + } + setIsTestingConnection(true); setConnectionTestId(`test-${Date.now()}`); setIsResultModalVisible(true); @@ -444,6 +474,7 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc accessToken={accessToken} testMode="chat" modelName={form.getFieldValue("auto_router_name")} + preparedConnectionRequest={preparedConnectionRequest} onClose={() => { setIsResultModalVisible(false); setIsTestingConnection(false); diff --git a/ui/litellm-dashboard/src/components/add_model/handle_add_auto_router_submit.tsx b/ui/litellm-dashboard/src/components/add_model/handle_add_auto_router_submit.tsx index e0ce97b5f70..271e0e63423 100644 --- a/ui/litellm-dashboard/src/components/add_model/handle_add_auto_router_submit.tsx +++ b/ui/litellm-dashboard/src/components/add_model/handle_add_auto_router_submit.tsx @@ -1,65 +1,71 @@ import { modelCreateCall, Model } from "../networking"; import NotificationManager from "../molecules/notifications_manager"; +interface AutoRouterFormValues { + auto_router_name: string; + auto_router_config?: unknown; + auto_router_default_model?: string; + auto_router_embedding_model?: string; + complexity_router_config?: unknown; + custom_embedding_model?: string; + model_access_group?: string[]; + model_type?: string; + team_id?: string; +} + +export const buildAutoRouterModelConfig = (values: AutoRouterFormValues): Model => { + if (values.model_type === "complexity_router") { + return { + model_name: values.auto_router_name, + litellm_params: { + model: "auto_router/complexity_router", + complexity_router_config: values.complexity_router_config, + complexity_router_default_model: values.auto_router_default_model, + }, + model_info: { + ...(values.team_id ? { team_id: values.team_id } : {}), + ...(values.model_access_group?.length ? { access_groups: values.model_access_group } : {}), + }, + }; + } + + const litellmParams: Record = { + model: `auto_router/${values.auto_router_name}`, + auto_router_config: JSON.stringify(values.auto_router_config), + auto_router_default_model: values.auto_router_default_model, + }; + + if (values.auto_router_embedding_model && values.auto_router_embedding_model !== "custom") { + litellmParams.auto_router_embedding_model = values.auto_router_embedding_model; + } else if (values.custom_embedding_model) { + litellmParams.auto_router_embedding_model = values.custom_embedding_model; + } + + return { + model_name: values.auto_router_name, + litellm_params: litellmParams, + model_info: { + ...(values.team_id ? { team_id: values.team_id } : {}), + ...(values.model_access_group?.length ? { access_groups: values.model_access_group } : {}), + }, + }; +}; + export const handleAddAutoRouterSubmit = async (values: any, accessToken: string, form: any, callback?: () => void) => { try { console.log("=== AUTO ROUTER SUBMIT HANDLER CALLED ==="); console.log("handling auto router submit for formValues:", values); console.log("Model type:", values.model_type); - let autoRouterConfig: any; - if (values.model_type === "complexity_router") { - // Complexity Router configuration console.log("Creating complexity router configuration"); - - autoRouterConfig = { - model_name: values.auto_router_name, - litellm_params: { - // Use special prefix for complexity router - model: `auto_router/complexity_router`, - // Pass the complexity router config as a JSON object (not stringified) - complexity_router_config: values.complexity_router_config, - // Default model for fallback (use MEDIUM or first available tier) - complexity_router_default_model: values.auto_router_default_model, - }, - model_info: {}, - }; - console.log("Complexity router config:", values.complexity_router_config); } else { - // Semantic Router configuration (existing behavior) console.log("Creating semantic router configuration"); - - autoRouterConfig = { - model_name: values.auto_router_name, - litellm_params: { - model: `auto_router/${values.auto_router_name}`, - auto_router_config: JSON.stringify(values.auto_router_config), // Convert JSON object to string as expected by backend - auto_router_default_model: values.auto_router_default_model, - }, - model_info: {}, - }; - - // Add optional embedding model if provided - if (values.auto_router_embedding_model && values.auto_router_embedding_model !== "custom") { - autoRouterConfig.litellm_params.auto_router_embedding_model = values.auto_router_embedding_model; - } else if (values.custom_embedding_model) { - autoRouterConfig.litellm_params.auto_router_embedding_model = values.custom_embedding_model; - } - - console.log("Semantic router config (stringified):", autoRouterConfig.litellm_params.auto_router_config); + console.log("Semantic router config (stringified):", JSON.stringify(values.auto_router_config)); } - // Add team information if provided - if (values.team_id) { - autoRouterConfig.model_info.team_id = values.team_id; - } - - // Add model access groups if provided - if (values.model_access_group && values.model_access_group.length > 0) { - autoRouterConfig.model_info.access_groups = values.model_access_group; - } + const autoRouterConfig = buildAutoRouterModelConfig(values); console.log("Auto router configuration to be created:", autoRouterConfig); diff --git a/ui/litellm-dashboard/src/components/add_model/model_connection_test.tsx b/ui/litellm-dashboard/src/components/add_model/model_connection_test.tsx index b71874e625c..75e8cefec92 100644 --- a/ui/litellm-dashboard/src/components/add_model/model_connection_test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/model_connection_test.tsx @@ -6,11 +6,18 @@ import { prepareModelAddRequest } from "./handle_add_model_submit"; import NotificationsManager from "../molecules/notifications_manager"; const { Text } = Typography; +type ConnectionParams = Record; + interface ModelConnectionTestProps { formValues: Record; accessToken: string; testMode: string; modelName?: string; + preparedConnectionRequest?: { + litellmParamsObj: ConnectionParams; + modelInfoObj: ConnectionParams; + mode?: string; + }; onClose?: () => void; onTestComplete?: () => void; } @@ -20,6 +27,7 @@ const ModelConnectionTest: React.FC = ({ accessToken, testMode, modelName = "this model", + preparedConnectionRequest, onClose, onTestComplete, }) => { @@ -43,9 +51,24 @@ const ModelConnectionTest: React.FC = ({ try { console.log("Testing connection with form values:", formValues); - const result = await prepareModelAddRequest(formValues, accessToken, null); + let connectionRequest = preparedConnectionRequest; - if (!result) { + if (!connectionRequest) { + const result = await prepareModelAddRequest(formValues, accessToken, null); + if (!result || result.length === 0) { + console.log("No result from prepareModelAddRequest"); + setError("Failed to prepare model data. Please check your form inputs."); + setIsSuccess(false); + setIsLoading(false); + return; + } + + console.log("Result from prepareModelAddRequest:", result); + const { litellmParamsObj, modelInfoObj } = result[0]; + connectionRequest = { litellmParamsObj, modelInfoObj }; + } + + if (!connectionRequest) { console.log("No result from prepareModelAddRequest"); setError("Failed to prepare model data. Please check your form inputs."); setIsSuccess(false); @@ -53,11 +76,10 @@ const ModelConnectionTest: React.FC = ({ return; } - console.log("Result from prepareModelAddRequest:", result); + const { litellmParamsObj, modelInfoObj, mode } = connectionRequest; + const requestMode = mode ?? (typeof modelInfoObj.mode === "string" ? modelInfoObj.mode : undefined); - const { litellmParamsObj, modelInfoObj, modelName: returnedModelName } = result[0]; - - const response = await testConnectionRequest(accessToken, litellmParamsObj, modelInfoObj, modelInfoObj?.mode); + const response = await testConnectionRequest(accessToken, litellmParamsObj, modelInfoObj, requestMode as string); if (response.status === "success") { NotificationsManager.success("Connection test successful!"); setError(null);