mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
feat(ui): working Test Connection for the complexity auto router
The consolidated auto-router tab dropped the Test Connection button because the shared prepareModelAddRequest helper returns an empty array for an auto router (it has no model_mappings), so the caller crashed destructuring result[0].litellmParamsObj. That is the crash in #31590 and the open PR #31794. #31794 only silenced the crash by pointing the test at auto_router/complexity_router, which is not a provider model, so the /health/test_connection health check (a real litellm.ahealth_check completion) would still error. Bring the button back and make it meaningful: an auto router dispatches to saved model groups, so Test Connection now probes those directly. It builds a deduped target list from the configured tiers (tiers sharing a model group collapse to one probe) plus the embedding model when semantic keyword matching is on, then runs a live /health/test_connection against each and shows per-target pass/fail. This never touches prepareModelAddRequest, so the original destructure crash cannot recur. Scope is the recommended complexity router only; the to-be-deprecated semantic router is untouched. No backend changes. Supersedes #31794. Resolves #31590.
This commit is contained in:
Abhimanyu Kapur
parent
92dfbdbb21
commit
2b2e8cf2bf
5 changed files with 387 additions and 5 deletions
|
|
@ -1,5 +1,5 @@
|
|||
import React, { useEffect, useState } from "react";
|
||||
import { Card, Form, Button, Tooltip, Typography, Select as AntdSelect, Radio, Badge, Space } from "antd";
|
||||
import { Card, Form, Button, Tooltip, Typography, Select as AntdSelect, Radio, Badge, Space, Modal } from "antd";
|
||||
import type { FormInstance } from "antd";
|
||||
import { ThunderboltOutlined, BranchesOutlined } from "@ant-design/icons";
|
||||
import { Text, TextInput } from "@tremor/react";
|
||||
|
|
@ -12,6 +12,8 @@ import ComplexityRouterConfig, { ComplexityRouterConfigValue } from "./Complexit
|
|||
import { KeywordTierRule } from "./KeywordTierRules";
|
||||
import { DEFAULT_MATCH_THRESHOLD } from "./SemanticKeywordMatching";
|
||||
import { buildComplexityRouterConfig, getSemanticConfigError } from "./build_complexity_router_config";
|
||||
import { buildAutoRouterTestTargets, AutoRouterTestTarget } from "./build_auto_router_test_targets";
|
||||
import AutoRouterConnectionTest from "./auto_router_connection_test";
|
||||
import NotificationManager from "../molecules/notifications_manager";
|
||||
|
||||
interface AddAutoRouterTabProps {
|
||||
|
|
@ -45,6 +47,11 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({ form, handleOk, acc
|
|||
// Semantic router config (existing)
|
||||
const [routerConfig, setRouterConfig] = useState<any>(null);
|
||||
|
||||
const [isTestModalVisible, setIsTestModalVisible] = useState<boolean>(false);
|
||||
const [isTestingConnection, setIsTestingConnection] = useState<boolean>(false);
|
||||
const [connectionTestId, setConnectionTestId] = useState<number>(0);
|
||||
const [testTargets, setTestTargets] = useState<AutoRouterTestTarget[]>([]);
|
||||
|
||||
useEffect(() => {
|
||||
const fetchModelAccessGroups = async () => {
|
||||
const response = await modelAvailableCall(accessToken, "", "", false, null, true, true);
|
||||
|
|
@ -194,6 +201,24 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({ form, handleOk, acc
|
|||
}
|
||||
};
|
||||
|
||||
const handleTestConnection = () => {
|
||||
const targets = buildAutoRouterTestTargets({
|
||||
tiers: complexityRouterConfig.tiers,
|
||||
semanticMatchingEnabled,
|
||||
embeddingModel,
|
||||
});
|
||||
|
||||
if (targets.length === 0) {
|
||||
NotificationManager.fromBackend("Please select at least one model for a complexity tier");
|
||||
return;
|
||||
}
|
||||
|
||||
setTestTargets(targets);
|
||||
setConnectionTestId((id) => id + 1);
|
||||
setIsTestingConnection(true);
|
||||
setIsTestModalVisible(true);
|
||||
};
|
||||
|
||||
return (
|
||||
<>
|
||||
<Title level={2}>Add Auto Router</Title>
|
||||
|
|
@ -355,10 +380,15 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({ form, handleOk, acc
|
|||
<Typography.Link href="https://github.com/BerriAI/litellm/issues">Need Help?</Typography.Link>
|
||||
</Tooltip>
|
||||
<div className="space-x-2">
|
||||
{/* TODO: add back a Test Connection or JSON preview action here. Test Connection was removed
|
||||
because prepareModelAddRequest can't build a valid pre-save payload for an auto router
|
||||
(tiers are model-group references, not litellm_params); a JSON preview of the
|
||||
complexity_router_config would be a good alternative. */}
|
||||
{routerType === "recommended" && (
|
||||
<Button
|
||||
data-testid="auto-router-test-connect-btn"
|
||||
onClick={handleTestConnection}
|
||||
loading={isTestingConnection}
|
||||
>
|
||||
Test Connection
|
||||
</Button>
|
||||
)}
|
||||
<Button
|
||||
type="primary"
|
||||
onClick={() => {
|
||||
|
|
@ -371,6 +401,36 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({ form, handleOk, acc
|
|||
</div>
|
||||
</Form>
|
||||
</Card>
|
||||
|
||||
<Modal
|
||||
title="Connection Test Results"
|
||||
open={isTestModalVisible}
|
||||
onCancel={() => {
|
||||
setIsTestModalVisible(false);
|
||||
setIsTestingConnection(false);
|
||||
}}
|
||||
footer={[
|
||||
<Button
|
||||
key="close"
|
||||
onClick={() => {
|
||||
setIsTestModalVisible(false);
|
||||
setIsTestingConnection(false);
|
||||
}}
|
||||
>
|
||||
Close
|
||||
</Button>,
|
||||
]}
|
||||
width={700}
|
||||
>
|
||||
{isTestModalVisible && (
|
||||
<AutoRouterConnectionTest
|
||||
key={connectionTestId}
|
||||
accessToken={accessToken}
|
||||
targets={testTargets}
|
||||
onTestComplete={() => setIsTestingConnection(false)}
|
||||
/>
|
||||
)}
|
||||
</Modal>
|
||||
</>
|
||||
);
|
||||
};
|
||||
|
|
|
|||
|
|
@ -0,0 +1,80 @@
|
|||
import { renderWithProviders, screen, waitFor } from "../../../tests/test-utils";
|
||||
import { vi } from "vitest";
|
||||
import AutoRouterConnectionTest from "./auto_router_connection_test";
|
||||
import { AutoRouterTestTarget } from "./build_auto_router_test_targets";
|
||||
|
||||
vi.mock("../networking", async () => {
|
||||
const actual = await vi.importActual("../networking");
|
||||
return {
|
||||
...actual,
|
||||
testConnectionRequest: vi.fn(),
|
||||
};
|
||||
});
|
||||
|
||||
const getMock = async () => vi.mocked((await import("../networking")).testConnectionRequest);
|
||||
|
||||
const targets: AutoRouterTestTarget[] = [
|
||||
{ labels: ["SIMPLE"], modelGroup: "gpt-4o-mini", mode: "chat" },
|
||||
{ labels: ["MEDIUM", "COMPLEX"], modelGroup: "claude-sonnet-4", mode: "chat" },
|
||||
{ labels: ["Embedding"], modelGroup: "voyage-3-5", mode: "embedding" },
|
||||
];
|
||||
|
||||
describe("AutoRouterConnectionTest", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it("probes each target once with the right model and mode (chat for tiers, embedding for the embedding model)", async () => {
|
||||
const mock = await getMock();
|
||||
mock.mockResolvedValue({ status: "success" });
|
||||
|
||||
renderWithProviders(<AutoRouterConnectionTest accessToken="sk-test" targets={targets} />);
|
||||
|
||||
await waitFor(() => expect(mock).toHaveBeenCalledTimes(3));
|
||||
|
||||
expect(mock).toHaveBeenCalledWith("sk-test", { model: "gpt-4o-mini" }, {}, "chat");
|
||||
expect(mock).toHaveBeenCalledWith("sk-test", { model: "claude-sonnet-4" }, {}, "chat");
|
||||
expect(mock).toHaveBeenCalledWith("sk-test", { model: "voyage-3-5" }, {}, "embedding");
|
||||
});
|
||||
|
||||
it("shows a success indicator per target when the health check passes", async () => {
|
||||
const mock = await getMock();
|
||||
mock.mockResolvedValue({ status: "success" });
|
||||
|
||||
renderWithProviders(<AutoRouterConnectionTest accessToken="sk-test" targets={targets} />);
|
||||
|
||||
await waitFor(() => expect(screen.getAllByTestId("test-status-success")).toHaveLength(3));
|
||||
expect(screen.queryByTestId("test-status-error")).toBeNull();
|
||||
expect(screen.getByText("MEDIUM, COMPLEX")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("renders the provider error message for a failing target while others pass", async () => {
|
||||
const mock = await getMock();
|
||||
mock.mockImplementation((_token, litellmParams) =>
|
||||
litellmParams.model === "claude-sonnet-4"
|
||||
? Promise.resolve({ status: "error", result: { error: "litellm.AuthenticationError: invalid api key" } })
|
||||
: Promise.resolve({ status: "success" }),
|
||||
);
|
||||
|
||||
renderWithProviders(<AutoRouterConnectionTest accessToken="sk-test" targets={targets} />);
|
||||
|
||||
await waitFor(() => expect(screen.getByTestId("test-error-message")).toBeInTheDocument());
|
||||
expect(screen.getByTestId("test-error-message")).toHaveTextContent("invalid api key");
|
||||
expect(screen.getByTestId("test-error-message")).not.toHaveTextContent("litellm.AuthenticationError");
|
||||
expect(screen.getAllByTestId("test-status-success")).toHaveLength(2);
|
||||
});
|
||||
|
||||
it("surfaces a thrown network error as a failing row", async () => {
|
||||
const mock = await getMock();
|
||||
mock.mockRejectedValue(new Error("Network request failed"));
|
||||
|
||||
renderWithProviders(
|
||||
<AutoRouterConnectionTest
|
||||
accessToken="sk-test"
|
||||
targets={[{ labels: ["SIMPLE"], modelGroup: "gpt-4o-mini", mode: "chat" }]}
|
||||
/>,
|
||||
);
|
||||
|
||||
await waitFor(() => expect(screen.getByTestId("test-error-message")).toHaveTextContent("Network request failed"));
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,132 @@
|
|||
import React from "react";
|
||||
import { Typography } from "antd";
|
||||
import { CheckCircleTwoTone, CloseCircleTwoTone, LoadingOutlined } from "@ant-design/icons";
|
||||
import { testConnectionRequest } from "../networking";
|
||||
import { AutoRouterTestTarget } from "./build_auto_router_test_targets";
|
||||
|
||||
const { Text } = Typography;
|
||||
|
||||
interface AutoRouterConnectionTestProps {
|
||||
accessToken: string;
|
||||
targets: AutoRouterTestTarget[];
|
||||
onTestComplete?: () => void;
|
||||
}
|
||||
|
||||
type TargetResult = { status: "pending" } | { status: "success" } | { status: "error"; error: string };
|
||||
|
||||
interface NormalizedResponse {
|
||||
ok: boolean;
|
||||
error?: string;
|
||||
}
|
||||
|
||||
const normalizeTestConnectionResponse = (response: unknown): NormalizedResponse => {
|
||||
if (typeof response !== "object" || response === null) {
|
||||
return { ok: false, error: "Unexpected response from connection test" };
|
||||
}
|
||||
const record = response as Record<string, unknown>;
|
||||
if (record.status === "success") {
|
||||
return { ok: true };
|
||||
}
|
||||
const result =
|
||||
typeof record.result === "object" && record.result !== null ? (record.result as Record<string, unknown>) : {};
|
||||
const resultError = typeof result.error === "string" ? result.error : undefined;
|
||||
const recordMessage = typeof record.message === "string" ? record.message : undefined;
|
||||
return { ok: false, error: resultError ?? recordMessage ?? "Unknown error" };
|
||||
};
|
||||
|
||||
const cleanErrorMessage = (error: string): string => {
|
||||
const mainError = error.split("stack trace:")[0].trim();
|
||||
return mainError.replace(/^litellm\.(.*?)Error: /, "");
|
||||
};
|
||||
|
||||
const runTarget = async (accessToken: string, target: AutoRouterTestTarget): Promise<TargetResult> => {
|
||||
try {
|
||||
const response = await testConnectionRequest(accessToken, { model: target.modelGroup }, {}, target.mode);
|
||||
const normalized = normalizeTestConnectionResponse(response);
|
||||
return normalized.ok
|
||||
? { status: "success" }
|
||||
: { status: "error", error: cleanErrorMessage(normalized.error ?? "Unknown error") };
|
||||
} catch (error) {
|
||||
return { status: "error", error: cleanErrorMessage(error instanceof Error ? error.message : String(error)) };
|
||||
}
|
||||
};
|
||||
|
||||
const AutoRouterConnectionTest: React.FC<AutoRouterConnectionTestProps> = ({
|
||||
accessToken,
|
||||
targets,
|
||||
onTestComplete,
|
||||
}) => {
|
||||
const [results, setResults] = React.useState<TargetResult[]>(() => targets.map(() => ({ status: "pending" })));
|
||||
|
||||
React.useEffect(() => {
|
||||
let cancelled = false;
|
||||
const run = async () => {
|
||||
const settled = await Promise.all(targets.map((target) => runTarget(accessToken, target)));
|
||||
if (cancelled) return;
|
||||
setResults(settled);
|
||||
if (onTestComplete) onTestComplete();
|
||||
};
|
||||
run();
|
||||
return () => {
|
||||
cancelled = true;
|
||||
};
|
||||
// eslint-disable-next-line react-hooks/exhaustive-deps -- probes run once per mount; the parent remounts via `key` to start a fresh test, and re-running on prop identity changes would refire paid health checks
|
||||
}, []);
|
||||
|
||||
if (targets.length === 0) {
|
||||
return <Text type="secondary">No complexity tiers are configured yet, so there is nothing to test.</Text>;
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="space-y-3">
|
||||
<Text type="secondary" style={{ display: "block", marginBottom: 8 }}>
|
||||
Each configured tier routes to a saved model group. Test Connection runs a live health check against each one.
|
||||
</Text>
|
||||
{targets.map((target, index) => {
|
||||
const result = results[index] ?? { status: "pending" };
|
||||
return (
|
||||
<div
|
||||
key={`${target.modelGroup}-${target.mode}`}
|
||||
data-testid="auto-router-test-row"
|
||||
style={{
|
||||
border: "1px solid #e5e7eb",
|
||||
borderRadius: 8,
|
||||
padding: "12px 16px",
|
||||
display: "flex",
|
||||
alignItems: "flex-start",
|
||||
gap: 12,
|
||||
}}
|
||||
>
|
||||
<div style={{ fontSize: 18, lineHeight: "24px" }}>
|
||||
{result.status === "pending" && <LoadingOutlined data-testid="test-status-pending" />}
|
||||
{result.status === "success" && (
|
||||
<CheckCircleTwoTone twoToneColor="#52c41a" data-testid="test-status-success" />
|
||||
)}
|
||||
{result.status === "error" && (
|
||||
<CloseCircleTwoTone twoToneColor="#ff4d4f" data-testid="test-status-error" />
|
||||
)}
|
||||
</div>
|
||||
<div style={{ flex: 1 }}>
|
||||
<Text strong>{target.labels.join(", ")}</Text>{" "}
|
||||
<Text type="secondary">
|
||||
{"->"} {target.modelGroup}
|
||||
{target.mode === "embedding" ? " (embedding)" : ""}
|
||||
</Text>
|
||||
{result.status === "error" && (
|
||||
<Text
|
||||
type="danger"
|
||||
data-testid="test-error-message"
|
||||
style={{ display: "block", marginTop: 4, fontSize: 13 }}
|
||||
>
|
||||
{result.error}
|
||||
</Text>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
})}
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
export default AutoRouterConnectionTest;
|
||||
|
|
@ -0,0 +1,68 @@
|
|||
import { buildAutoRouterTestTargets } from "./build_auto_router_test_targets";
|
||||
|
||||
const tiers = {
|
||||
SIMPLE: "gpt-4o-mini",
|
||||
MEDIUM: "claude-sonnet-4",
|
||||
COMPLEX: "claude-sonnet-4",
|
||||
REASONING: "o3",
|
||||
};
|
||||
|
||||
describe("buildAutoRouterTestTargets", () => {
|
||||
it("dedups tiers that share a model group into one chat target carrying both labels", () => {
|
||||
const targets = buildAutoRouterTestTargets({ tiers, semanticMatchingEnabled: false, embeddingModel: undefined });
|
||||
expect(targets).toEqual([
|
||||
{ labels: ["SIMPLE"], modelGroup: "gpt-4o-mini", mode: "chat" },
|
||||
{ labels: ["MEDIUM", "COMPLEX"], modelGroup: "claude-sonnet-4", mode: "chat" },
|
||||
{ labels: ["REASONING"], modelGroup: "o3", mode: "chat" },
|
||||
]);
|
||||
});
|
||||
|
||||
it("drops empty/whitespace tiers", () => {
|
||||
const targets = buildAutoRouterTestTargets({
|
||||
tiers: { SIMPLE: "gpt-4o-mini", MEDIUM: "", COMPLEX: " ", REASONING: "" },
|
||||
semanticMatchingEnabled: false,
|
||||
embeddingModel: undefined,
|
||||
});
|
||||
expect(targets).toEqual([{ labels: ["SIMPLE"], modelGroup: "gpt-4o-mini", mode: "chat" }]);
|
||||
});
|
||||
|
||||
it("returns [] when no tier is configured", () => {
|
||||
expect(
|
||||
buildAutoRouterTestTargets({
|
||||
tiers: { SIMPLE: "", MEDIUM: "", COMPLEX: "", REASONING: "" },
|
||||
semanticMatchingEnabled: false,
|
||||
embeddingModel: undefined,
|
||||
}),
|
||||
).toEqual([]);
|
||||
});
|
||||
|
||||
it("appends an embedding target only when semantic matching is on and a model is set", () => {
|
||||
const targets = buildAutoRouterTestTargets({
|
||||
tiers: { SIMPLE: "gpt-4o-mini", MEDIUM: "", COMPLEX: "", REASONING: "" },
|
||||
semanticMatchingEnabled: true,
|
||||
embeddingModel: "voyage-3-5",
|
||||
});
|
||||
expect(targets).toEqual([
|
||||
{ labels: ["SIMPLE"], modelGroup: "gpt-4o-mini", mode: "chat" },
|
||||
{ labels: ["Embedding"], modelGroup: "voyage-3-5", mode: "embedding" },
|
||||
]);
|
||||
});
|
||||
|
||||
it("omits the embedding target when semantic matching is on but no model is chosen", () => {
|
||||
const targets = buildAutoRouterTestTargets({
|
||||
tiers: { SIMPLE: "gpt-4o-mini", MEDIUM: "", COMPLEX: "", REASONING: "" },
|
||||
semanticMatchingEnabled: true,
|
||||
embeddingModel: undefined,
|
||||
});
|
||||
expect(targets).toEqual([{ labels: ["SIMPLE"], modelGroup: "gpt-4o-mini", mode: "chat" }]);
|
||||
});
|
||||
|
||||
it("omits the embedding target when a model is set but semantic matching is off", () => {
|
||||
const targets = buildAutoRouterTestTargets({
|
||||
tiers: { SIMPLE: "gpt-4o-mini", MEDIUM: "", COMPLEX: "", REASONING: "" },
|
||||
semanticMatchingEnabled: false,
|
||||
embeddingModel: "voyage-3-5",
|
||||
});
|
||||
expect(targets).toEqual([{ labels: ["SIMPLE"], modelGroup: "gpt-4o-mini", mode: "chat" }]);
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,42 @@
|
|||
import { ComplexityTiers } from "./ComplexityRouterConfig";
|
||||
|
||||
export type AutoRouterTestMode = "chat" | "embedding";
|
||||
|
||||
export interface AutoRouterTestTarget {
|
||||
labels: string[];
|
||||
modelGroup: string;
|
||||
mode: AutoRouterTestMode;
|
||||
}
|
||||
|
||||
export interface BuildAutoRouterTestTargetsParams {
|
||||
tiers: ComplexityTiers;
|
||||
semanticMatchingEnabled: boolean;
|
||||
embeddingModel: string | undefined;
|
||||
}
|
||||
|
||||
const TIER_ORDER: (keyof ComplexityTiers)[] = ["SIMPLE", "MEDIUM", "COMPLEX", "REASONING"];
|
||||
|
||||
export const buildAutoRouterTestTargets = ({
|
||||
tiers,
|
||||
semanticMatchingEnabled,
|
||||
embeddingModel,
|
||||
}: BuildAutoRouterTestTargetsParams): AutoRouterTestTarget[] => {
|
||||
const groupedByModel = TIER_ORDER.reduce<Record<string, string[]>>((acc, tier) => {
|
||||
const modelGroup = tiers[tier]?.trim();
|
||||
if (!modelGroup) return acc;
|
||||
return { ...acc, [modelGroup]: [...(acc[modelGroup] ?? []), tier] };
|
||||
}, {});
|
||||
|
||||
const tierTargets: AutoRouterTestTarget[] = Object.entries(groupedByModel).map(([modelGroup, labels]) => ({
|
||||
labels,
|
||||
modelGroup,
|
||||
mode: "chat" as const,
|
||||
}));
|
||||
|
||||
const embeddingTarget: AutoRouterTestTarget[] =
|
||||
semanticMatchingEnabled && embeddingModel?.trim()
|
||||
? [{ labels: ["Embedding"], modelGroup: embeddingModel.trim(), mode: "embedding" as const }]
|
||||
: [];
|
||||
|
||||
return [...tierTargets, ...embeddingTarget];
|
||||
};
|
||||
Loading…
Add table
Reference in a new issue