fix(ui): test complexity auto router connections

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<string, string>) => void }) => (
+    <button
+      type="button"
+      onClick={() =>
+        onChange({
+          SIMPLE: "gpt-3.5-turbo",
+          MEDIUM: "gpt-4o-mini",
+          COMPLEX: "",
+          REASONING: "",
+        })
+      }
+    >
+      Select complexity tiers
+    </button>
+  ),
+}));
+
+const renderAddAutoRouterTab = () => {
+  const Component = () => {
+    const [form] = Form.useForm();
+    return <AddAutoRouterTab form={form} handleOk={vi.fn()} accessToken="test-token" userRole="proxy_admin" />;
+  };
+
+  return renderWithProviders(<Component />);
+};
+
+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,
+      ),
+    );
+  });
+});

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<string, unknown>;
 
 interface ComplexityTiers {
   SIMPLE: string;
@@ -35,6 +36,14 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({ form, handleOk, acc
   const [isResultModalVisible, setIsResultModalVisible] = useState<boolean>(false);
   const [isTestingConnection, setIsTestingConnection] = useState<boolean>(false);
   const [connectionTestId, setConnectionTestId] = useState<string>("");
+  const [preparedConnectionRequest, setPreparedConnectionRequest] = useState<
+    | {
+        litellmParamsObj: ConnectionParams;
+        modelInfoObj: ConnectionParams;
+        mode?: string;
+      }
+    | undefined
+  >();
 
   const [modelAccessGroups, setModelAccessGroups] = useState<string[]>([]);
   const [modelInfo, setModelInfo] = useState<ModelGroup[]>([]);
@@ -80,6 +89,27 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({ 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<AddAutoRouterTabProps> = ({ form, handleOk, acc
             accessToken={accessToken}
             testMode="chat"
             modelName={form.getFieldValue("auto_router_name")}
+            preparedConnectionRequest={preparedConnectionRequest}
             onClose={() => {
               setIsResultModalVisible(false);
               setIsTestingConnection(false);

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<string, unknown> = {
+    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);
 

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<string, unknown>;
+
 interface ModelConnectionTestProps {
   formValues: Record<string, any>;
   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<ModelConnectionTestProps> = ({
   accessToken,
   testMode,
   modelName = "this model",
+  preparedConnectionRequest,
   onClose,
   onTestComplete,
 }) => {
@@ -43,9 +51,24 @@ const ModelConnectionTest: React.FC<ModelConnectionTestProps> = ({
 
     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<ModelConnectionTestProps> = ({
         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);