mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
refactor(ui): simplify Auto Setup model selection
This commit is contained in:
moe-berri
parent
a2e5e7e066
commit
8acb8de997
5 changed files with 112 additions and 314 deletions
|
|
@ -90,9 +90,6 @@ export interface AutoRouterCandidateDeployment {
|
|||
export interface AutoRouterDeployment extends AutoRouterCandidateDeployment {
|
||||
litellm_params?: {
|
||||
model?: string | null;
|
||||
base_model?: string | null;
|
||||
input_cost_per_token?: number | null;
|
||||
output_cost_per_token?: number | null;
|
||||
complexity_router_config?: unknown;
|
||||
complexity_router_default_model?: string | null;
|
||||
auto_router_config?: unknown;
|
||||
|
|
@ -108,8 +105,6 @@ export interface AutoRouterDeployment extends AutoRouterCandidateDeployment {
|
|||
/** False for config.yaml-defined deployments, which the update and delete routes refuse. */
|
||||
db_model?: boolean | null;
|
||||
base_model?: string | null;
|
||||
input_cost_per_token?: number | null;
|
||||
output_cost_per_token?: number | null;
|
||||
created_at?: string | null;
|
||||
updated_at?: string | null;
|
||||
team_id?: string | null;
|
||||
|
|
|
|||
|
|
@ -19,10 +19,6 @@ vi.mock(
|
|||
"@/app/(dashboard)/hooks/autoRouter/useAutoRouterPresets",
|
||||
async () => await import("../../../tests/mocks/autoRouterPresets"),
|
||||
);
|
||||
vi.mock("@/app/(dashboard)/hooks/models/useModelCostMap", () => ({
|
||||
useModelCostMap: () => ({ data: {}, isLoading: false }),
|
||||
}));
|
||||
|
||||
const getAllPresets = (): AutoRouterPreset[] => BUNDLED_PRESETS;
|
||||
const getPresetByKey = (key: string): AutoRouterPreset | undefined => BUNDLED_PRESETS.find((p) => p.key === key);
|
||||
|
||||
|
|
@ -158,56 +154,38 @@ describe("AddAutoRouterTab", () => {
|
|||
expect(screen.getByText("Complexity Tier Configuration")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("configures four unique tiers from the available models with one click", async () => {
|
||||
it("hides automatic setup when no available model is recommended", async () => {
|
||||
mockFetchAvailableModels.mockResolvedValue([
|
||||
{ model_group: "premium", mode: "chat" },
|
||||
{ model_group: "cheap", mode: "chat" },
|
||||
{ model_group: "best", mode: "chat" },
|
||||
{ model_group: "middle", mode: "chat" },
|
||||
]);
|
||||
mockFetchAllModelDeployments.mockResolvedValue([
|
||||
{ model_name: "cheap", litellm_params: { model: "cheap", input_cost_per_token: 1 } },
|
||||
{ model_name: "middle", litellm_params: { model: "middle", input_cost_per_token: 2 } },
|
||||
{ model_name: "premium", litellm_params: { model: "premium", input_cost_per_token: 3 } },
|
||||
{ model_name: "best", litellm_params: { model: "best", input_cost_per_token: 4 } },
|
||||
{ model_group: "unknown-model-a", mode: "chat" },
|
||||
{ model_group: "unknown-model-b", mode: "chat" },
|
||||
]);
|
||||
renderWithProviders(<Harness />);
|
||||
|
||||
const button = screen.getByTestId("configure-automatically-button");
|
||||
await waitFor(() => expect(button).toBeEnabled());
|
||||
await userEvent.click(button);
|
||||
|
||||
expect(screen.getByText(/Simple: cheap.*Medium: middle.*Complex: premium.*Reasoning: best/)).toBeInTheDocument();
|
||||
openTemplateDropdown();
|
||||
await waitFor(() => expect(optionByLabel("Anthropic Family")).toHaveTextContent("Missing:"));
|
||||
expect(screen.queryByTestId("configure-automatically-button")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("prefers the first compatible bundled template over the price fallback", async () => {
|
||||
const firstPreset = getAllPresets()[0];
|
||||
it("mixes preferred tier models even when one complete preset is available", async () => {
|
||||
const anthropicPreset = getPresetByKey("anthropic_family")!;
|
||||
mockFetchAvailableModels.mockResolvedValue(
|
||||
[...getRequiredModelsInPreset(firstPreset)].map((model_group) => ({ model_group, mode: "chat" })),
|
||||
[...getRequiredModelsInPreset(anthropicPreset), "gpt-5.6-luna"].map((model_group) => ({
|
||||
model_group,
|
||||
mode: "chat",
|
||||
})),
|
||||
);
|
||||
mockFetchAllModelDeployments.mockResolvedValue([]);
|
||||
renderWithProviders(<Harness />);
|
||||
|
||||
const button = screen.getByTestId("configure-automatically-button");
|
||||
await waitFor(() => expect(button).toBeEnabled());
|
||||
const button = await screen.findByTestId("configure-automatically-button");
|
||||
await userEvent.click(button);
|
||||
|
||||
expect(toast.success).toHaveBeenCalledWith(`Configured with ${firstPreset.label}`);
|
||||
});
|
||||
|
||||
it("prefers the OpenAI template over Gemini", async () => {
|
||||
const openAiPreset = getPresetByKey("openai_family")!;
|
||||
const geminiPreset = getPresetByKey("gemini_family")!;
|
||||
const available = new Set([...getRequiredModelsInPreset(openAiPreset), ...getRequiredModelsInPreset(geminiPreset)]);
|
||||
mockFetchAvailableModels.mockResolvedValue([...available].map((model_group) => ({ model_group, mode: "chat" })));
|
||||
mockFetchAllModelDeployments.mockResolvedValue([]);
|
||||
renderWithProviders(<Harness />);
|
||||
|
||||
const button = screen.getByTestId("configure-automatically-button");
|
||||
await waitFor(() => expect(button).toBeEnabled());
|
||||
await userEvent.click(button);
|
||||
|
||||
expect(toast.success).toHaveBeenCalledWith(`Configured with ${openAiPreset.label}`);
|
||||
expect(
|
||||
screen.getByText(
|
||||
/Simple: gpt-5.6-luna.*Medium: claude-sonnet-5.*Complex: claude-opus-5.*Reasoning: claude-opus-5/,
|
||||
),
|
||||
).toBeInTheDocument();
|
||||
expect(toast.success).not.toHaveBeenCalledWith(expect.stringContaining("Configured with"));
|
||||
});
|
||||
|
||||
it("mixes available models from the preferred tier catalog when no complete template fits", async () => {
|
||||
|
|
@ -220,8 +198,7 @@ describe("AddAutoRouterTab", () => {
|
|||
mockFetchAllModelDeployments.mockResolvedValue([]);
|
||||
renderWithProviders(<Harness />);
|
||||
|
||||
const button = screen.getByTestId("configure-automatically-button");
|
||||
await waitFor(() => expect(button).toBeEnabled());
|
||||
const button = await screen.findByTestId("configure-automatically-button");
|
||||
await userEvent.click(button);
|
||||
|
||||
expect(
|
||||
|
|
|
|||
|
|
@ -21,7 +21,6 @@ import TeamDropdown from "../common_components/team_dropdown";
|
|||
import { type AddAutoRouterValues, handleAddAutoRouterSubmit } from "./handle_add_auto_router_submit";
|
||||
import { fetchAvailableModels, type ModelGroup } from "@/components/llm_calls/fetch_models";
|
||||
import { autoRouterListKey, fetchAllModelDeployments } from "@/app/(dashboard)/hooks/models/useModels";
|
||||
import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap";
|
||||
import ComplexityRouterConfig, {
|
||||
ComplexityRouterConfigValue,
|
||||
effectiveClassifierType,
|
||||
|
|
@ -106,7 +105,6 @@ const presetDisabledHint = (availability: PresetAvailability): string | null =>
|
|||
const isPresetHintAlarming = (availability: PresetAvailability): boolean => availability.kind === "missing_models";
|
||||
|
||||
const NO_PRESETS: AutoRouterPreset[] = [];
|
||||
const AUTO_SETUP_PRESET_PRIORITY = ["1m_context", "anthropic_family", "openai_family", "gemini_family", "lite"];
|
||||
|
||||
// A one-line summary of what's configured, shown when the detailed section is collapsed so a
|
||||
// caller can see the shape of the config without opening it.
|
||||
|
|
@ -236,7 +234,6 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
|
|||
queryFn: () => fetchAllModelDeployments(accessToken, userId ?? "", userRole),
|
||||
enabled: Boolean(accessToken),
|
||||
});
|
||||
const { data: modelCostMap = {}, isLoading: costsLoading } = useModelCostMap();
|
||||
const modelsLoading = groupsLoading || deploymentsLoading;
|
||||
const modelInfo = React.useMemo(() => data ?? [], [data]);
|
||||
const {
|
||||
|
|
@ -246,7 +243,7 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
|
|||
refetch: refetchPresets,
|
||||
} = useAutoRouterPresets();
|
||||
const presets = presetsData ?? NO_PRESETS;
|
||||
const automaticSetupLoading = modelsLoading || costsLoading || presetsPending;
|
||||
const automaticSetupLoading = modelsLoading || presetsPending;
|
||||
const presetsUnavailable = presetsError && presetsData === undefined;
|
||||
// react-query keeps the last successful list around when a later refetch fails, so isError alone
|
||||
// can't tell "never loaded" apart from "loaded, then a background refetch errored" - only the
|
||||
|
|
@ -271,6 +268,14 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
|
|||
),
|
||||
[modelInfo],
|
||||
);
|
||||
const preferredTierModels = React.useMemo(
|
||||
() => buildPreferredTierModels(presets, availability),
|
||||
[presets, availability],
|
||||
);
|
||||
const automaticRouterConfig = React.useMemo(
|
||||
() => buildAutomaticRouterConfig(modelInfo, deployments ?? [], preferredTierModels),
|
||||
[modelInfo, deployments, preferredTierModels],
|
||||
);
|
||||
|
||||
// A preset's models can only be trusted against a successfully loaded list. Selection and the
|
||||
// greyed-out state derive from this one function, so a preset that cannot be selected can never
|
||||
|
|
@ -319,30 +324,11 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
|
|||
};
|
||||
|
||||
const handleAutomaticSetup = () => {
|
||||
const prioritizedPresets = AUTO_SETUP_PRESET_PRIORITY.flatMap((key) => {
|
||||
const preset = presets.find((candidate) => candidate.key === key);
|
||||
return preset ? [preset] : [];
|
||||
});
|
||||
const matchingPreset = prioritizedPresets.find((preset) => presetAvailability(preset).kind === "available");
|
||||
if (matchingPreset) {
|
||||
const presetState = presetAvailability(matchingPreset);
|
||||
setSelectedPreset(matchingPreset.key);
|
||||
applyPrefill(buildPresetPrefill(matchingPreset.complexity_router_config, availability));
|
||||
setDetailsExpanded(presetState.kind === "available" && presetState.viaDeployments);
|
||||
toast.success(`Configured with ${matchingPreset.label}`);
|
||||
return;
|
||||
}
|
||||
|
||||
const preferredTierModels = buildPreferredTierModels(prioritizedPresets, availability);
|
||||
const generatedConfig = buildAutomaticRouterConfig(modelInfo, deployments ?? [], modelCostMap, preferredTierModels);
|
||||
if (generatedConfig === null) {
|
||||
toast.fromError("Add at least one chat model before configuring an Auto Router");
|
||||
return;
|
||||
}
|
||||
if (automaticRouterConfig === null) return;
|
||||
setSelectedPreset(undefined);
|
||||
applyPrefill({ ...buildEmptyPrefill(), complexityRouterConfig: generatedConfig });
|
||||
applyPrefill({ ...buildEmptyPrefill(), complexityRouterConfig: automaticRouterConfig });
|
||||
setDetailsExpanded(false);
|
||||
toast.success("Automatic setup created", { description: tierConfigSummary(generatedConfig) });
|
||||
toast.success("Automatic setup created", { description: tierConfigSummary(automaticRouterConfig) });
|
||||
};
|
||||
|
||||
const handlePresetChange = (presetKey: string | undefined) => {
|
||||
|
|
@ -540,24 +526,20 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
|
|||
{({ ref, ...field }) => <Input {...field} ref={ref} placeholder="e.g., smart_router, auto_router_1" />}
|
||||
</FormField>
|
||||
|
||||
<div className="flex items-center justify-between gap-4 rounded-lg border border-border p-4">
|
||||
<div>
|
||||
<div className="text-sm font-medium text-foreground">Start with a recommended setup</div>
|
||||
<div className="text-xs text-muted-foreground">
|
||||
Uses your available models and our recommended setups. You can review and edit everything before
|
||||
saving.
|
||||
{!automaticSetupLoading && automaticRouterConfig && (
|
||||
<div className="flex items-center justify-between gap-4 rounded-lg border border-border p-4">
|
||||
<div>
|
||||
<div className="text-sm font-medium text-foreground">Start with a recommended setup</div>
|
||||
<div className="text-xs text-muted-foreground">
|
||||
Uses your available models and our recommended setups. You can review and edit everything before
|
||||
saving.
|
||||
</div>
|
||||
</div>
|
||||
<Button type="button" data-testid="configure-automatically-button" onClick={handleAutomaticSetup}>
|
||||
Configure automatically
|
||||
</Button>
|
||||
</div>
|
||||
<Button
|
||||
type="button"
|
||||
data-testid="configure-automatically-button"
|
||||
disabled={automaticSetupLoading}
|
||||
onClick={handleAutomaticSetup}
|
||||
>
|
||||
{automaticSetupLoading && <UiLoadingSpinner className="size-4" />}
|
||||
Configure automatically
|
||||
</Button>
|
||||
</div>
|
||||
)}
|
||||
|
||||
<div>
|
||||
<label className="block text-sm font-medium text-foreground mb-2">Template</label>
|
||||
|
|
|
|||
|
|
@ -1,184 +1,89 @@
|
|||
import { describe, expect, it } from "vitest";
|
||||
import type { AutoRouterDeployment } from "@/app/(dashboard)/hooks/models/useModels";
|
||||
import { buildAutomaticRouterConfig, type PreferredTierModels } from "./auto_setup";
|
||||
import { buildModelAvailability } from "@/lib/autorouter_presets";
|
||||
import { buildAutomaticRouterConfig, buildPreferredTierModels, type PreferredTierModels } from "./auto_setup";
|
||||
|
||||
const models = (...names: string[]) => names.map((model_group) => ({ model_group, mode: "chat" }));
|
||||
const deployment = (model_name: string, model = model_name): AutoRouterDeployment => ({
|
||||
model_name,
|
||||
litellm_params: { model },
|
||||
});
|
||||
const tierModels = (config: ReturnType<typeof buildAutomaticRouterConfig>) =>
|
||||
config && Object.values(config.tiers).map((tier) => (typeof tier === "string" ? tier : tier[0]));
|
||||
|
||||
const deployment = (name: string, cost: number): AutoRouterDeployment => ({
|
||||
model_name: name,
|
||||
litellm_params: {
|
||||
model: name,
|
||||
input_cost_per_token: cost / 2,
|
||||
output_cost_per_token: cost / 2,
|
||||
},
|
||||
describe("buildPreferredTierModels", () => {
|
||||
it("recognizes curated models that are not in a preset", () => {
|
||||
const available = ["gpt-4o-mini", "claude-sonnet-4-5", "grok-4", "deepseek-reasoner"];
|
||||
const availability = buildModelAvailability(available, []);
|
||||
const preferred = buildPreferredTierModels([], availability);
|
||||
|
||||
expect(preferred).toEqual({
|
||||
SIMPLE: ["gpt-4o-mini"],
|
||||
MEDIUM: ["claude-sonnet-4-5"],
|
||||
COMPLEX: ["grok-4"],
|
||||
REASONING: ["deepseek-reasoner"],
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
const firstModel = (value: string | string[]): string => (typeof value === "string" ? value : value[0]);
|
||||
|
||||
const tierModels = (config: ReturnType<typeof buildAutomaticRouterConfig>) =>
|
||||
config && [
|
||||
firstModel(config.tiers.SIMPLE),
|
||||
firstModel(config.tiers.MEDIUM),
|
||||
firstModel(config.tiers.COMPLEX),
|
||||
firstModel(config.tiers.REASONING),
|
||||
];
|
||||
|
||||
describe("buildAutomaticRouterConfig", () => {
|
||||
it("uses available preferred models before price ranking", () => {
|
||||
it("selects one preferred model for each tier", () => {
|
||||
const preferred: PreferredTierModels = {
|
||||
SIMPLE: ["preferred-simple"],
|
||||
MEDIUM: ["preferred-medium"],
|
||||
COMPLEX: ["preferred-complex"],
|
||||
REASONING: ["preferred-reasoning"],
|
||||
SIMPLE: ["simple"],
|
||||
MEDIUM: ["medium"],
|
||||
COMPLEX: ["complex"],
|
||||
REASONING: ["reasoning"],
|
||||
};
|
||||
const available = [...Object.values(preferred).flat(), "cheap-decoy", "expensive-decoy"];
|
||||
const config = buildAutomaticRouterConfig(
|
||||
models(...available),
|
||||
available.map((name, index) => deployment(name, index + 1)),
|
||||
{},
|
||||
preferred,
|
||||
);
|
||||
|
||||
expect(tierModels(config)).toEqual([
|
||||
"preferred-simple",
|
||||
"preferred-medium",
|
||||
"preferred-complex",
|
||||
"preferred-reasoning",
|
||||
]);
|
||||
expect(
|
||||
tierModels(buildAutomaticRouterConfig(models("simple", "medium", "complex", "reasoning"), [], preferred)),
|
||||
).toEqual(["simple", "medium", "complex", "reasoning"]);
|
||||
});
|
||||
|
||||
it("reuses the nearest preferred model for tiers with no preferred match", () => {
|
||||
it("reuses the closest available tier when a tier has no match", () => {
|
||||
const preferred: PreferredTierModels = {
|
||||
SIMPLE: ["preferred-simple"],
|
||||
SIMPLE: ["simple"],
|
||||
MEDIUM: [],
|
||||
COMPLEX: ["preferred-complex"],
|
||||
COMPLEX: ["complex"],
|
||||
REASONING: [],
|
||||
};
|
||||
const config = buildAutomaticRouterConfig(
|
||||
models("preferred-simple", "preferred-complex", "cheap-decoy"),
|
||||
[deployment("preferred-simple", 4), deployment("preferred-complex", 5), deployment("cheap-decoy", 1)],
|
||||
{},
|
||||
preferred,
|
||||
);
|
||||
|
||||
expect(tierModels(config)).toEqual([
|
||||
"preferred-simple",
|
||||
"preferred-simple",
|
||||
"preferred-complex",
|
||||
"preferred-complex",
|
||||
expect(tierModels(buildAutomaticRouterConfig(models("simple", "complex"), [], preferred))).toEqual([
|
||||
"simple",
|
||||
"simple",
|
||||
"complex",
|
||||
"complex",
|
||||
]);
|
||||
});
|
||||
|
||||
it("uses price ranking when none of the preferred models are available", () => {
|
||||
const unavailablePreferred: PreferredTierModels = {
|
||||
it("returns null when none of the available models are recommended", () => {
|
||||
const preferred: PreferredTierModels = {
|
||||
SIMPLE: ["missing-simple"],
|
||||
MEDIUM: ["missing-medium"],
|
||||
COMPLEX: ["missing-complex"],
|
||||
REASONING: ["missing-reasoning"],
|
||||
};
|
||||
const config = buildAutomaticRouterConfig(
|
||||
models("expensive", "cheap", "premium", "middle"),
|
||||
[deployment("cheap", 1), deployment("middle", 2), deployment("premium", 3), deployment("expensive", 4)],
|
||||
{},
|
||||
unavailablePreferred,
|
||||
);
|
||||
|
||||
expect(tierModels(config)).toEqual(["cheap", "middle", "premium", "expensive"]);
|
||||
expect(buildAutomaticRouterConfig(models("unknown-model"), [], preferred)).toBeNull();
|
||||
});
|
||||
|
||||
it("uses four different models when four are available", () => {
|
||||
const config = buildAutomaticRouterConfig(
|
||||
models("expensive", "cheap", "premium", "middle"),
|
||||
[deployment("cheap", 1), deployment("middle", 2), deployment("premium", 3), deployment("expensive", 4)],
|
||||
{},
|
||||
);
|
||||
|
||||
expect(tierModels(config)).toEqual(["cheap", "middle", "premium", "expensive"]);
|
||||
expect(config?.classifier_type).toBe("heuristic_v2");
|
||||
});
|
||||
|
||||
it("selects one model per tier from a large inventory", () => {
|
||||
const names = Array.from({ length: 100 }, (_, index) => `model-${index.toString().padStart(3, "0")}`);
|
||||
const config = buildAutomaticRouterConfig(
|
||||
models(...names),
|
||||
names.map((name, index) => deployment(name, index + 1)),
|
||||
{},
|
||||
);
|
||||
const expectedTiers = {
|
||||
SIMPLE: ["model-000"],
|
||||
MEDIUM: ["model-033"],
|
||||
COMPLEX: ["model-066"],
|
||||
REASONING: ["model-099"],
|
||||
it("ignores non-chat models and existing auto routers", () => {
|
||||
const preferred: PreferredTierModels = {
|
||||
SIMPLE: ["gpt-4o-mini", "smart-router"],
|
||||
MEDIUM: [],
|
||||
COMPLEX: [],
|
||||
REASONING: [],
|
||||
};
|
||||
const available = [
|
||||
{ model_group: "gpt-4o-mini", mode: "chat" },
|
||||
{ model_group: "image-model", mode: "image_generation" },
|
||||
{ model_group: "smart-router", mode: "chat" },
|
||||
];
|
||||
|
||||
expect(config?.tiers).toEqual(expectedTiers);
|
||||
});
|
||||
|
||||
it("only repeats models when fewer than four are available", () => {
|
||||
expect(
|
||||
tierModels(
|
||||
buildAutomaticRouterConfig(
|
||||
models("cheap", "expensive"),
|
||||
[deployment("cheap", 1), deployment("expensive", 4)],
|
||||
{},
|
||||
),
|
||||
buildAutomaticRouterConfig(available, [deployment("smart-router", "auto_router/complexity_router")], preferred),
|
||||
),
|
||||
).toEqual(["cheap", "cheap", "expensive", "expensive"]);
|
||||
});
|
||||
|
||||
it("uses the published cost map when deployments do not define prices", () => {
|
||||
const config = buildAutomaticRouterConfig(
|
||||
models("premium", "cheap", "middle"),
|
||||
[
|
||||
{ model_name: "premium", litellm_params: { model: "provider/premium" } },
|
||||
{ model_name: "cheap", litellm_params: { model: "provider/cheap" } },
|
||||
{ model_name: "middle", litellm_params: { model: "provider/middle" } },
|
||||
],
|
||||
{
|
||||
"provider/cheap": { input_cost_per_token: 1, output_cost_per_token: 1 },
|
||||
"provider/middle": { input_cost_per_token: 2, output_cost_per_token: 2 },
|
||||
"provider/premium": { input_cost_per_token: 3, output_cost_per_token: 3 },
|
||||
},
|
||||
);
|
||||
|
||||
expect(tierModels(config)).toEqual(["cheap", "middle", "premium", "premium"]);
|
||||
});
|
||||
|
||||
it("uses the most expensive deployment when a group has several", () => {
|
||||
const config = buildAutomaticRouterConfig(
|
||||
models("variable", "steady", "premium", "top"),
|
||||
[
|
||||
deployment("variable", 1),
|
||||
deployment("variable", 8),
|
||||
deployment("steady", 2),
|
||||
deployment("premium", 3),
|
||||
deployment("top", 4),
|
||||
],
|
||||
{},
|
||||
);
|
||||
|
||||
expect(tierModels(config)).toEqual(["steady", "premium", "top", "variable"]);
|
||||
});
|
||||
|
||||
it("ignores non-chat and existing auto-router models", () => {
|
||||
const config = buildAutomaticRouterConfig(
|
||||
[
|
||||
{ model_group: "chat-model", mode: "chat" },
|
||||
{ model_group: "image-model", mode: "image_generation" },
|
||||
{ model_group: "auto_router/existing", mode: "chat" },
|
||||
{ model_group: "smart-router", mode: "chat" },
|
||||
],
|
||||
[
|
||||
deployment("chat-model", 1),
|
||||
{ model_name: "smart-router", litellm_params: { model: "auto_router/complexity_router" } },
|
||||
],
|
||||
{},
|
||||
);
|
||||
|
||||
expect(tierModels(config)).toEqual(["chat-model", "chat-model", "chat-model", "chat-model"]);
|
||||
});
|
||||
|
||||
it("returns null when there are no usable models", () => {
|
||||
expect(buildAutomaticRouterConfig([], [], {})).toBeNull();
|
||||
).toEqual(["gpt-4o-mini", "gpt-4o-mini", "gpt-4o-mini", "gpt-4o-mini"]);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -3,62 +3,15 @@ import type { ModelGroup } from "@/components/llm_calls/fetch_models";
|
|||
import { resolveAvailableModel, type AutoRouterPreset, type ModelAvailability } from "@/lib/autorouter_presets";
|
||||
import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
|
||||
|
||||
type ModelCost = {
|
||||
input_cost_per_token?: number | null;
|
||||
output_cost_per_token?: number | null;
|
||||
};
|
||||
|
||||
export type ModelCostMap = Record<string, ModelCost>;
|
||||
|
||||
const TIER_NAMES = ["SIMPLE", "MEDIUM", "COMPLEX", "REASONING"] as const;
|
||||
type TierName = (typeof TIER_NAMES)[number];
|
||||
export type PreferredTierModels = Record<TierName, string[]>;
|
||||
|
||||
const price = (cost: ModelCost | null | undefined): number | undefined => {
|
||||
const input = cost?.input_cost_per_token;
|
||||
const output = cost?.output_cost_per_token;
|
||||
if (typeof input !== "number" && typeof output !== "number") return undefined;
|
||||
return (input ?? 0) + (output ?? 0);
|
||||
};
|
||||
|
||||
const deploymentPrice = (deployment: AutoRouterDeployment, costMap: ModelCostMap): number | undefined => {
|
||||
const configured = price(deployment.litellm_params) ?? price(deployment.model_info);
|
||||
if (configured !== undefined) return configured;
|
||||
|
||||
const references = [
|
||||
deployment.litellm_params?.model,
|
||||
deployment.litellm_params?.base_model,
|
||||
deployment.model_info?.base_model,
|
||||
];
|
||||
for (const reference of references) {
|
||||
if (reference && costMap[reference]) return price(costMap[reference]);
|
||||
}
|
||||
return undefined;
|
||||
};
|
||||
|
||||
const groupPrice = (
|
||||
modelGroup: string,
|
||||
deployments: AutoRouterDeployment[],
|
||||
costMap: ModelCostMap,
|
||||
): number | undefined => {
|
||||
const groupDeployments = deployments.filter(
|
||||
(deployment) =>
|
||||
deployment.model_name === modelGroup && !deployment.litellm_params?.model?.startsWith("auto_router/"),
|
||||
);
|
||||
if (groupDeployments.length === 0) return price(costMap[modelGroup]);
|
||||
const prices = groupDeployments.map((deployment) => deploymentPrice(deployment, costMap));
|
||||
if (prices.some((value) => value === undefined)) return undefined;
|
||||
const knownPrices = prices.filter((value): value is number => value !== undefined);
|
||||
return Math.max(...knownPrices);
|
||||
};
|
||||
|
||||
const selectTierModels = (ranked: string[]): [string, string, string, string] => {
|
||||
if (ranked.length === 1) return [ranked[0], ranked[0], ranked[0], ranked[0]];
|
||||
if (ranked.length === 2) return [ranked[0], ranked[0], ranked[1], ranked[1]];
|
||||
if (ranked.length === 3) return [ranked[0], ranked[1], ranked[2], ranked[2]];
|
||||
|
||||
const last = ranked.length - 1;
|
||||
return [ranked[0], ranked[Math.floor(last / 3)], ranked[Math.floor((2 * last) / 3)], ranked[last]];
|
||||
const ADDITIONAL_TIER_MODELS: PreferredTierModels = {
|
||||
SIMPLE: ["gpt-4o-mini", "gpt-5-mini", "gemini-2.5-flash", "deepseek-chat"],
|
||||
MEDIUM: ["gpt-5-mini", "gpt-4o", "claude-sonnet-4-5", "gemini-2.5-flash", "deepseek-chat"],
|
||||
COMPLEX: ["gpt-5", "gpt-4o", "claude-sonnet-4-6", "gemini-2.5-pro", "grok-4"],
|
||||
REASONING: ["o3", "deepseek-reasoner", "claude-opus-4-6", "gemini-2.5-pro", "gpt-5"],
|
||||
};
|
||||
|
||||
export const buildPreferredTierModels = (
|
||||
|
|
@ -70,12 +23,13 @@ export const buildPreferredTierModels = (
|
|||
tier,
|
||||
Array.from(
|
||||
new Set(
|
||||
presets.flatMap((preset) =>
|
||||
preset.complexity_router_config.tiers[tier].flatMap((model) => {
|
||||
const resolved = resolveAvailableModel(model, availability);
|
||||
return resolved ? [resolved] : [];
|
||||
}),
|
||||
),
|
||||
[
|
||||
...presets.flatMap((preset) => preset.complexity_router_config.tiers[tier]),
|
||||
...ADDITIONAL_TIER_MODELS[tier],
|
||||
].flatMap((model) => {
|
||||
const resolved = resolveAvailableModel(model, availability);
|
||||
return resolved ? [resolved] : [];
|
||||
}),
|
||||
),
|
||||
),
|
||||
]),
|
||||
|
|
@ -99,8 +53,7 @@ const selectPreferredTierModels = (
|
|||
export const buildAutomaticRouterConfig = (
|
||||
models: ModelGroup[],
|
||||
deployments: AutoRouterDeployment[],
|
||||
costMap: ModelCostMap,
|
||||
preferredByTier?: PreferredTierModels,
|
||||
preferredByTier: PreferredTierModels,
|
||||
): ComplexityRouterConfigValue | null => {
|
||||
const autoRouterNames: ReadonlySet<string> = new Set(
|
||||
deployments
|
||||
|
|
@ -117,22 +70,8 @@ export const buildAutomaticRouterConfig = (
|
|||
);
|
||||
if (names.length === 0) return null;
|
||||
const usableNames: ReadonlySet<string> = new Set(names);
|
||||
|
||||
const ranked = names
|
||||
.map((name) => ({ name, price: groupPrice(name, deployments, costMap) }))
|
||||
.sort((left, right) => {
|
||||
if (left.price === undefined && right.price !== undefined) return 1;
|
||||
if (left.price !== undefined && right.price === undefined) return -1;
|
||||
if (left.price !== undefined && right.price !== undefined && left.price !== right.price) {
|
||||
return left.price - right.price;
|
||||
}
|
||||
return left.name.localeCompare(right.name);
|
||||
})
|
||||
.map(({ name }) => name);
|
||||
|
||||
const selected = preferredByTier
|
||||
? selectPreferredTierModels(preferredByTier, usableNames) ?? selectTierModels(ranked)
|
||||
: selectTierModels(ranked);
|
||||
const selected = selectPreferredTierModels(preferredByTier, usableNames);
|
||||
if (selected === null) return null;
|
||||
|
||||
return {
|
||||
tiers: {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue