From 40b823af87b2c0d8f0b4e7da4b4dce03bb599f70 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Fri, 19 Dec 2025 15:13:06 -0800
Subject: [PATCH] Add Health Check Model for Wildcard in UI

---
 .../src/components/model_info_view.test.tsx   | 65 +++++++++++++++++++
 .../src/components/model_info_view.tsx        | 55 ++++++++++++++++
 2 files changed, 120 insertions(+)

diff --git a/ui/litellm-dashboard/src/components/model_info_view.test.tsx b/ui/litellm-dashboard/src/components/model_info_view.test.tsx
index d63d3dd6eb..21402ad467 100644
--- a/ui/litellm-dashboard/src/components/model_info_view.test.tsx
+++ b/ui/litellm-dashboard/src/components/model_info_view.test.tsx
@@ -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} />);
diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx
index 64f96ac915..37aa68bcc7 100644
--- a/ui/litellm-dashboard/src/components/model_info_view.tsx
+++ b/ui/litellm-dashboard/src/components/model_info_view.tsx
@@ -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