Add Health Check Model for Wildcard in UI

This commit is contained in:
yuneng-jiang
2025-12-19 15:13:06 -08:00
parent 5b1fda02fb
commit 40b823af87
2 changed files with 120 additions and 0 deletions
@@ -107,6 +107,41 @@ vi.mock("./networking", () => ({
],
}),
credentialGetCall: vi.fn().mockResolvedValue({}),
getGuardrailsList: vi.fn().mockResolvedValue({
guardrails: [{ guardrail_name: "content_filter" }, { guardrail_name: "toxicity_filter" }],
}),
tagListCall: vi.fn().mockResolvedValue({
test_tag: {
name: "test_tag",
description: "A test tag",
},
production_tag: {
name: "production_tag",
description: "Production ready models",
},
}),
}));
// Mock the useModelsInfo hook since it uses React Query
vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({
useModelsInfo: vi.fn().mockReturnValue({
data: {
data: [
{
model_name: "bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0",
provider: "bedrock",
litellm_model_name: "bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0",
},
{
model_name: "openai/gpt-4",
provider: "openai",
litellm_model_name: "gpt-4",
},
],
},
isLoading: false,
error: null,
}),
}));
describe("ModelInfoView", () => {
@@ -242,6 +277,36 @@ describe("ModelInfoView", () => {
});
});
it("should render health check model field for wildcard routes", async () => {
const wildcardModelData = {
...modelData,
litellm_model_name: "openai/gpt-4*",
};
const WILDCARD_ADMIN_PROPS = {
...DEFAULT_ADMIN_PROPS,
modelData: wildcardModelData,
};
const { getByText } = render(<ModelInfoView {...WILDCARD_ADMIN_PROPS} />);
await waitFor(() => {
expect(getByText("Model Settings")).toBeInTheDocument();
});
await waitFor(() => {
expect(getByText("Health Check Model")).toBeInTheDocument();
});
});
it("should not render health check model field for non-wildcard routes", async () => {
const { queryByText } = render(<ModelInfoView {...DEFAULT_ADMIN_PROPS} />);
await waitFor(() => {
expect(queryByText("Model Settings")).toBeInTheDocument();
});
await waitFor(() => {
expect(queryByText("Health Check Model")).not.toBeInTheDocument();
});
});
describe("View Model", () => {
it("should render the model info view", async () => {
const { getByText } = render(<ModelInfoView {...DEFAULT_ADMIN_PROPS} />);
@@ -37,6 +37,7 @@ import { getProviderLogoAndName } from "./provider_info_helpers";
import NumericalInput from "./shared/numerical_input";
import { Tag } from "./tag_management/types";
import { getDisplayModelName } from "./view_model/model_name_display";
import { useModelsInfo } from "@/app/(dashboard)/hooks/models/useModels";
interface ModelInfoViewProps {
modelId: string;
@@ -83,6 +84,8 @@ export default function ModelInfoView({
const isAdmin = userRole === "Admin";
const isAutoRouter = modelData?.litellm_params?.auto_router_config != null;
const { data: modelsInfoData } = useModelsInfo(accessToken, userID, userRole);
console.log("modelsInfoData, ", modelsInfoData);
const usingExistingCredential =
modelData?.litellm_params?.litellm_credential_name != null &&
modelData?.litellm_params?.litellm_credential_name != undefined;
@@ -226,6 +229,13 @@ export default function ModelInfoView({
access_groups: values.model_access_group,
};
}
// Override health_check_model from the form
if (values.health_check_model !== undefined) {
updatedModelInfo = {
...updatedModelInfo,
health_check_model: values.health_check_model,
};
}
} catch (e) {
NotificationsManager.fromBackend("Invalid JSON in Model Info");
return;
@@ -342,6 +352,7 @@ export default function ModelInfoView({
onModelUpdate(updatedModel);
}
};
const isWildcardModel = modelData.litellm_model_name.includes("*");
return (
<div className="p-4">
@@ -545,6 +556,7 @@ export default function ModelInfoView({
? localModelData.litellm_params.guardrails
: [],
tags: Array.isArray(localModelData.litellm_params?.tags) ? localModelData.litellm_params.tags : [],
health_check_model: isWildcardModel ? localModelData.model_info?.health_check_model : null,
litellm_extra_params: JSON.stringify(localModelData.litellm_params || {}, null, 2),
}}
layout="vertical"
@@ -868,6 +880,49 @@ export default function ModelInfoView({
)}
</div>
{isWildcardModel && (
<div>
<Text className="font-medium">Health Check Model</Text>
{isEditing ? (
<Form.Item name="health_check_model" className="mb-0">
<Select
showSearch
placeholder="Select existing health check model"
optionFilterProp="children"
allowClear
options={(() => {
const seen = new Set();
return modelsInfoData?.data
?.filter((model: any) => {
const modelProvider = model.provider;
const wildcardProvider = modelData.litellm_model_name.split("/")[0];
return (
modelProvider === wildcardProvider &&
model.model_name !== modelData.litellm_model_name
);
})
.filter((model: any) => {
if (seen.has(model.model_name)) {
return false;
}
seen.add(model.model_name);
return true;
})
.map((model: any) => ({
value: model.model_name,
label: model.model_name,
}));
})()}
/>
</Form.Item>
) : (
<div className="mt-1 p-2 bg-gray-50 rounded">
{localModelData.model_info?.health_check_model || "Not Set"}
</div>
)}
</div>
)}
{/* Cache Control Section */}
{isEditing ? (
<CacheControlSettings