ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin' into litellm_deleted_keys_endpoint

Browse Source
This commit is contained in:
yuneng-jiang 2026-01-16 19:25:03 -08:00
commit 90dc0e9120
5 changed files with 350 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
litellm/proxy/public_endpoints/public_endpoints.py
View File

@ -29,7 +29,8 @@ router = APIRouter()
)
async def public_model_hub():
import litellm
from litellm.proxy.proxy_server import _get_model_group_info, llm_router
from litellm.proxy.proxy_server import _get_model_group_info, llm_router, prisma_client
from litellm.proxy.health_endpoints._health_endpoints import _convert_health_check_to_dict
if llm_router is None:
raise HTTPException(
@ -44,6 +45,28 @@ async def public_model_hub():
model_group=None,
)
# Fetch health check information if available
health_checks_map = {}
if prisma_client is not None:
try:
latest_checks = await prisma_client.get_all_latest_health_checks()
for check in latest_checks:
key = check.model_id if check.model_id else check.model_name
if key:
health_check_dict = _convert_health_check_to_dict(check)
health_checks_map[key] = health_check_dict
if check.model_name:
health_checks_map[check.model_name] = health_check_dict
except Exception:
pass
for model_group in model_groups:
health_info = health_checks_map.get(model_group.model_group)
if health_info:
model_group.health_status = health_info.get("status")
model_group.health_response_time = health_info.get("response_time_ms")
model_group.health_checked_at = health_info.get("checked_at")
return model_groups

5
litellm/types/proxy/management_endpoints/model_management_endpoints.py
View File

@ -1,4 +1,4 @@
from typing import Dict, List, Union, Any
from typing import Dict, List, Union, Any, Optional
from pydantic import BaseModel, Field
@ -7,6 +7,9 @@ from ...router import ModelGroupInfo
class ModelGroupInfoProxy(ModelGroupInfo):
is_public_model_group: bool = Field(default=False)
health_status: Optional[str] = Field(default=None)
health_response_time: Optional[float] = Field(default=None)
health_checked_at: Optional[str] = Field(default=None)
class UpdateUsefulLinksRequest(BaseModel):

256
tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
View File

@ -1,5 +1,9 @@
import os
import sys
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
sys.path.insert(
0, os.path.abspath("../../..")
@ -8,7 +12,11 @@ sys.path.insert(
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.public_endpoints import router
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
ModelGroupInfoProxy,
)
from litellm.types.utils import LlmProviders
@ -101,3 +109,251 @@ def test_watsonx_provider_fields():
assert "token" in field_keys
assert "zen_api_key" in field_keys
def test_public_model_hub_with_healthy_model():
"""Test that health information is populated for a healthy model"""
app = FastAPI()
app.include_router(router)
# Override auth dependency
app.dependency_overrides[user_api_key_auth] = lambda: MagicMock()
client = TestClient(app)
# Create mock model groups
mock_model_group = ModelGroupInfoProxy(
model_group="gpt-3.5-turbo",
providers=["openai"],
is_public_model_group=True,
)
# Create mock health check
mock_health_check = MagicMock()
mock_health_check.model_id = None
mock_health_check.model_name = "gpt-3.5-turbo"
mock_health_check.status = "healthy"
mock_health_check.response_time_ms = 150.5
mock_health_check.checked_at = datetime.now(timezone.utc)
mock_llm_router = MagicMock()
mock_prisma = MagicMock()
mock_prisma.get_all_latest_health_checks = AsyncMock(
return_value=[mock_health_check]
)
with patch("litellm.public_model_groups", ["gpt-3.5-turbo"]), \
patch("litellm.proxy.proxy_server._get_model_group_info") as mock_get_info, \
patch("litellm.proxy.proxy_server.llm_router", mock_llm_router), \
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), \
patch("litellm.proxy.health_endpoints._health_endpoints._convert_health_check_to_dict") as mock_convert:
mock_get_info.return_value = [mock_model_group]
mock_convert.return_value = {
"status": "healthy",
"response_time_ms": 150.5,
"checked_at": mock_health_check.checked_at.isoformat(),
}
response = client.get(
"/public/model_hub",
headers={"Authorization": "Bearer test-key"},
)
assert response.status_code == 200
data = response.json()
assert len(data) == 1
assert data[0]["model_group"] == "gpt-3.5-turbo"
assert data[0]["health_status"] == "healthy"
assert data[0]["health_response_time"] == 150.5
assert data[0]["health_checked_at"] is not None
app.dependency_overrides.clear()
def test_public_model_hub_with_unhealthy_model():
"""Test that health information is populated for an unhealthy model"""
app = FastAPI()
app.include_router(router)
app.dependency_overrides[user_api_key_auth] = lambda: MagicMock()
client = TestClient(app)
mock_model_group = ModelGroupInfoProxy(
model_group="gpt-4",
providers=["openai"],
is_public_model_group=True,
)
mock_health_check = MagicMock()
mock_health_check.model_id = None
mock_health_check.model_name = "gpt-4"
mock_health_check.status = "unhealthy"
mock_health_check.response_time_ms = None
mock_health_check.checked_at = datetime.now(timezone.utc)
mock_llm_router = MagicMock()
mock_prisma = MagicMock()
mock_prisma.get_all_latest_health_checks = AsyncMock(
return_value=[mock_health_check]
)
with patch("litellm.public_model_groups", ["gpt-4"]), \
patch("litellm.proxy.proxy_server._get_model_group_info") as mock_get_info, \
patch("litellm.proxy.proxy_server.llm_router", mock_llm_router), \
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), \
patch("litellm.proxy.health_endpoints._health_endpoints._convert_health_check_to_dict") as mock_convert:
mock_get_info.return_value = [mock_model_group]
mock_convert.return_value = {
"status": "unhealthy",
"response_time_ms": None,
"checked_at": mock_health_check.checked_at.isoformat(),
}
response = client.get(
"/public/model_hub",
headers={"Authorization": "Bearer test-key"},
)
assert response.status_code == 200
data = response.json()
assert len(data) == 1
assert data[0]["model_group"] == "gpt-4"
assert data[0]["health_status"] == "unhealthy"
assert data[0]["health_response_time"] is None
assert data[0]["health_checked_at"] is not None
app.dependency_overrides.clear()
def test_public_model_hub_without_health_check():
"""Test that health information is null when no health check exists"""
app = FastAPI()
app.include_router(router)
app.dependency_overrides[user_api_key_auth] = lambda: MagicMock()
client = TestClient(app)
mock_model_group = ModelGroupInfoProxy(
model_group="claude-3",
providers=["anthropic"],
is_public_model_group=True,
)
mock_llm_router = MagicMock()
mock_prisma = MagicMock()
mock_prisma.get_all_latest_health_checks = AsyncMock(return_value=[])
with patch("litellm.public_model_groups", ["claude-3"]), \
patch("litellm.proxy.proxy_server._get_model_group_info") as mock_get_info, \
patch("litellm.proxy.proxy_server.llm_router", mock_llm_router), \
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
mock_get_info.return_value = [mock_model_group]
response = client.get(
"/public/model_hub",
headers={"Authorization": "Bearer test-key"},
)
assert response.status_code == 200
data = response.json()
assert len(data) == 1
assert data[0]["model_group"] == "claude-3"
assert data[0]["health_status"] is None
assert data[0]["health_response_time"] is None
assert data[0]["health_checked_at"] is None
app.dependency_overrides.clear()
def test_public_model_hub_mixed_health_statuses():
"""Test multiple models with different health statuses"""
app = FastAPI()
app.include_router(router)
app.dependency_overrides[user_api_key_auth] = lambda: MagicMock()
client = TestClient(app)
healthy_model = ModelGroupInfoProxy(
model_group="gpt-3.5-turbo",
providers=["openai"],
is_public_model_group=True,
)
unhealthy_model = ModelGroupInfoProxy(
model_group="gpt-4",
providers=["openai"],
is_public_model_group=True,
)
no_health_model = ModelGroupInfoProxy(
model_group="claude-3",
providers=["anthropic"],
is_public_model_group=True,
)
healthy_check = MagicMock()
healthy_check.model_id = None
healthy_check.model_name = "gpt-3.5-turbo"
healthy_check.status = "healthy"
healthy_check.response_time_ms = 120.0
healthy_check.checked_at = datetime.now(timezone.utc)
unhealthy_check = MagicMock()
unhealthy_check.model_id = None
unhealthy_check.model_name = "gpt-4"
unhealthy_check.status = "unhealthy"
unhealthy_check.response_time_ms = None
unhealthy_check.checked_at = datetime.now(timezone.utc)
mock_llm_router = MagicMock()
mock_prisma = MagicMock()
mock_prisma.get_all_latest_health_checks = AsyncMock(
return_value=[healthy_check, unhealthy_check]
)
def convert_side_effect(check):
if check.model_name == "gpt-3.5-turbo":
return {
"status": "healthy",
"response_time_ms": 120.0,
"checked_at": check.checked_at.isoformat(),
}
elif check.model_name == "gpt-4":
return {
"status": "unhealthy",
"response_time_ms": None,
"checked_at": check.checked_at.isoformat(),
}
return {}
with patch("litellm.public_model_groups", ["gpt-3.5-turbo", "gpt-4", "claude-3"]), \
patch("litellm.proxy.proxy_server._get_model_group_info") as mock_get_info, \
patch("litellm.proxy.proxy_server.llm_router", mock_llm_router), \
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), \
patch("litellm.proxy.health_endpoints._health_endpoints._convert_health_check_to_dict") as mock_convert:
mock_get_info.return_value = [
healthy_model,
unhealthy_model,
no_health_model,
]
mock_convert.side_effect = convert_side_effect
response = client.get(
"/public/model_hub",
headers={"Authorization": "Bearer test-key"},
)
assert response.status_code == 200
data = response.json()
assert len(data) == 3
# Find each model and verify health status
gpt35 = next(m for m in data if m["model_group"] == "gpt-3.5-turbo")
assert gpt35["health_status"] == "healthy"
assert gpt35["health_response_time"] == 120.0
assert gpt35["health_checked_at"] is not None
gpt4 = next(m for m in data if m["model_group"] == "gpt-4")
assert gpt4["health_status"] == "unhealthy"
assert gpt4["health_response_time"] is None
assert gpt4["health_checked_at"] is not None
claude = next(m for m in data if m["model_group"] == "claude-3")
assert claude["health_status"] is None
assert claude["health_response_time"] is None
assert claude["health_checked_at"] is None
app.dependency_overrides.clear()

63
ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx
View File

@ -6,6 +6,7 @@ import { KeyResponse, Team } from "../key_team_helpers/key_list";
import { Organization } from "../networking";
import { KeysResponse, useKeys } from "@/app/(dashboard)/hooks/keys/useKeys";
import { useFilterLogic } from "../key_team_helpers/filter_logic";
import useTeams from "@/app/(dashboard)/hooks/useTeams";
// Mock network calls
vi.mock("./networking", async (importOriginal) => {
@ -21,6 +22,7 @@ vi.mock("./networking", async (importOriginal) => {
},
],
}),
teamListCall: vi.fn().mockResolvedValue([]),
};
});
@ -51,6 +53,20 @@ vi.mock("../key_team_helpers/filter_logic", () => ({
useFilterLogic: vi.fn(),
}));
// Mock useTeams hook (used by KeyInfoView)
vi.mock("@/app/(dashboard)/hooks/useTeams", () => ({
default: vi.fn(),
}));
// Mock fetchTeams to prevent network calls
vi.mock("@/app/(dashboard)/networking", async (importOriginal) => {
const actual = await importOriginal<typeof import("@/app/(dashboard)/networking")>();
return {
...actual,
fetchTeams: vi.fn().mockResolvedValue([]),
};
});
const mockKey: KeyResponse = {
token: "sk-1234567890abcdef",
token_id: "key-1",
@ -146,6 +162,7 @@ const mockOrganization: Organization = {
// Mock hook implementations
const mockUseKeys = useKeys as MockedFunction<typeof useKeys>;
const mockUseFilterLogic = useFilterLogic as MockedFunction<typeof useFilterLogic>;
const mockUseTeams = useTeams as MockedFunction<typeof useTeams>;
beforeEach(() => {
// Reset mocks before each test
@ -181,6 +198,12 @@ beforeEach(() => {
handleFilterChange: vi.fn(),
handleFilterReset: vi.fn(),
});
// Mock useTeams hook (used by KeyInfoView)
mockUseTeams.mockReturnValue({
teams: [mockTeam],
setTeams: vi.fn(),
});
});
it("should render VirtualKeysTable component", () => {
@ -394,3 +417,43 @@ it("should handle column resizing hover events", () => {
fireEvent.mouseLeave(headerCell);
expect(resizer.style.opacity).toBe("0");
});
it("should open KeyInfoView when clicking on a key ID button", async () => {
const mockProps = {
teams: [mockTeam],
organizations: [mockOrganization],
onSortChange: vi.fn(),
currentSort: {
sortBy: "created_at",
sortOrder: "desc" as const,
},
};
renderWithProviders(<VirtualKeysTable {...mockProps} />);
// Wait for the table to render
await waitFor(() => {
expect(screen.getByText("Test Key Alias")).toBeInTheDocument();
});
// Verify table is visible before clicking - check for table-specific text
expect(screen.getByText(/Showing.*results/)).toBeInTheDocument();
// Find the key ID button (it should show the truncated token)
const keyIdButton = screen.getByText("sk-1234...");
expect(keyIdButton).toBeInTheDocument();
// Click on the key ID button
fireEvent.click(keyIdButton);
// Wait for KeyInfoView to appear - check for unique elements that only exist in KeyInfoView
await waitFor(() => {
expect(screen.getByText("Back to Keys")).toBeInTheDocument();
// KeyInfoView shows "Created:" or "Updated:" which is unique to it
expect(screen.getByText(/Created:|Updated:/)).toBeInTheDocument();
});
// Verify that table-specific elements are no longer visible
// The "Showing X of Y results" text should not be visible when KeyInfoView is open
expect(screen.queryByText(/Showing.*results/)).not.toBeInTheDocument();
});

6
ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx
View File

@ -533,6 +533,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
onClose={() => setSelectedKey(null)}
keyData={selectedKey}
teams={allTeams}
onDelete={refetch}
/>
) : (
<div className="border-b py-4 flex-1 overflow-hidden">
@ -599,11 +600,10 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
<TableHeaderCell
key={header.id}
data-header-id={header.id}
className={`py-1 h-8 relative hover:bg-gray-50 ${
header.id === "actions"
className={`py-1 h-8 relative hover:bg-gray-50 ${header.id === "actions"
? "sticky right-0 bg-white shadow-[-4px_0_8px_-6px_rgba(0,0,0,0.1)]"
: ""
}`}
}`}
style={{
width: header.getSize(),
position: "relative",