Merge pull request #19258 from BerriAI/litellm_ui_model_hub_health_1

[Feature] UI - Public Model Hub: Health Checks

ui/litellm-dashboard/src/components/public_model_hub.test.tsx

@@ -1,5 +1,5 @@
 import { describe, it, expect, vi, beforeAll, beforeEach } from "vitest";
-import { render } from "@testing-library/react";
+import { render, screen, waitFor } from "@testing-library/react";
 import PublicModelHub from "./public_model_hub";
 
 vi.mock("next/navigation", () => ({
@@ -38,10 +38,10 @@ beforeAll(() => {
       matches: false,
       media: query,
       onchange: null,
-      addListener: () => {},
-      removeListener: () => {},
-      addEventListener: () => {},
-      removeEventListener: () => {},
+      addListener: () => { },
+      removeListener: () => { },
+      addEventListener: () => { },
+      removeEventListener: () => { },
       dispatchEvent: () => false,
     }),
   });
@@ -64,4 +64,111 @@ describe("PublicModelHub", () => {
     const { container } = render(<PublicModelHub />);
     expect(container).toBeInTheDocument();
   });
+
+  it("displays health status correctly for models with health check information", async () => {
+    const mockModelsWithHealthChecks = [
+      {
+        model_group: "gpt-4",
+        providers: ["openai"],
+        mode: "chat",
+        health_status: "healthy",
+        health_response_time: 150.5,
+        health_checked_at: "2024-01-15T10:30:00Z",
+        supports_function_calling: true,
+        supports_vision: false,
+        supports_parallel_function_calling: false,
+      },
+      {
+        model_group: "claude-3",
+        providers: ["anthropic"],
+        mode: "chat",
+        health_status: "unhealthy",
+        health_response_time: 5000.0,
+        health_checked_at: "2024-01-15T10:25:00Z",
+        supports_function_calling: true,
+        supports_vision: false,
+        supports_parallel_function_calling: false,
+      },
+      {
+        model_group: "gpt-3.5-turbo",
+        providers: ["openai"],
+        mode: "chat",
+        health_status: undefined,
+        health_response_time: undefined,
+        health_checked_at: undefined,
+        supports_function_calling: false,
+        supports_vision: false,
+        supports_parallel_function_calling: false,
+      },
+    ];
+
+    const networkingModule = await import("./networking");
+    vi.mocked(networkingModule.modelHubPublicModelsCall).mockResolvedValue(mockModelsWithHealthChecks);
+
+    render(<PublicModelHub />);
+
+    // Wait for the component to load and render the table
+    await waitFor(() => {
+      expect(screen.getByText("gpt-4")).toBeInTheDocument();
+    });
+
+    // Check that health status is displayed for healthy model (gpt-4)
+    // Find the row containing "gpt-4" and verify it has "healthy" status
+    await waitFor(() => {
+      const gpt4Cell = screen.getByText("gpt-4");
+      const gpt4Row = gpt4Cell.closest("tr");
+      expect(gpt4Row).toBeInTheDocument();
+
+      // Find all cells in the row
+      const cells = gpt4Row?.querySelectorAll("td");
+      expect(cells).toBeTruthy();
+
+      // Find the cell containing "healthy" text (health status column)
+      // The health status is in a Tag component, so look for a Tag containing "healthy"
+      const healthyStatus = Array.from(cells || []).find((cell) => {
+        const tag = cell.querySelector('[class*="ant-tag"]');
+        const text = tag?.textContent?.toLowerCase();
+        return text === "healthy";
+      });
+      expect(healthyStatus).toBeInTheDocument();
+    });
+
+    // Check that health status is displayed for unhealthy model (claude-3)
+    await waitFor(() => {
+      const claude3Cell = screen.getByText("claude-3");
+      const claude3Row = claude3Cell.closest("tr");
+      expect(claude3Row).toBeInTheDocument();
+
+      // Find all cells in the row
+      const cells = claude3Row?.querySelectorAll("td");
+      expect(cells).toBeTruthy();
+
+      // Find the cell containing "unhealthy" text (health status column)
+      const unhealthyStatus = Array.from(cells || []).find((cell) => {
+        const tag = cell.querySelector('[class*="ant-tag"]');
+        const text = tag?.textContent?.toLowerCase();
+        return text === "unhealthy";
+      });
+      expect(unhealthyStatus).toBeInTheDocument();
+    });
+
+    // Check that "Unknown" is displayed for model without health status (gpt-3.5-turbo)
+    await waitFor(() => {
+      const gpt35Cell = screen.getByText("gpt-3.5-turbo");
+      const gpt35Row = gpt35Cell.closest("tr");
+      expect(gpt35Row).toBeInTheDocument();
+
+      // Find all cells in the row
+      const cells = gpt35Row?.querySelectorAll("td");
+      expect(cells).toBeTruthy();
+
+      // Find the cell containing "Unknown" text (health status column)
+      const unknownStatus = Array.from(cells || []).find((cell) => {
+        const tag = cell.querySelector('[class*="ant-tag"]');
+        const text = tag?.textContent;
+        return text === "Unknown";
+      });
+      expect(unknownStatus).toBeInTheDocument();
+    });
+  });
 });

ui/litellm-dashboard/src/components/public_model_hub.tsx

@@ -36,6 +36,9 @@ interface ModelGroupInfo {
   supports_vision: boolean;
   supports_function_calling: boolean;
   supported_openai_params?: string[];
+  health_status?: string;
+  health_response_time?: number;
+  health_checked_at?: string;
   [key: string]: any;
 }
 
@@ -687,6 +690,27 @@ const PublicModelHub: React.FC<PublicModelHubProps> = ({ accessToken, isEmbedded
       },
       size: 120,
     },
+    {
+      header: "Health Status",
+      accessorKey: "health_status",
+      enableSorting: true,
+      cell: ({ row }) => {
+        const original = row.original;
+        const tagColor = original.health_status === "healthy" ? "green" : original.health_status === "unhealthy" ? "red" : "default";
+        const responseTimeLabel = original.health_response_time ? `Response Time: ${Number(original.health_response_time).toFixed(2)}ms` : "N/A";
+        const lastCheckedLabel = original.health_checked_at ? `Last Checked: ${new Date(original.health_checked_at).toLocaleString()}` : "N/A";
+
+        return <Tooltip title={<>
+          <div>
+            {responseTimeLabel}
+          </div>
+          <div>
+            {lastCheckedLabel}
+          </div>
+        </>}><Tag key={original.model_group} color={tagColor}><span className="capitalize">{original.health_status ?? "Unknown"}</span></Tag></Tooltip>;
+      },
+      size: 100,
+    },
     {
       header: "Limits",
       accessorKey: "rpm",