From 439472f80063fa24bd46b09298d4213d73ac1afe Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 16 Jan 2026 15:36:56 -0800 Subject: [PATCH 1/2] /public/model_hub health information --- .../public_endpoints/public_endpoints.py | 25 +- .../model_management_endpoints.py | 5 +- .../public_endpoints/test_public_endpoints.py | 256 ++++++++++++++++++ 3 files changed, 284 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index abb69050464..aebea8616b4 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -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 as e: + 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 diff --git a/litellm/types/proxy/management_endpoints/model_management_endpoints.py b/litellm/types/proxy/management_endpoints/model_management_endpoints.py index a8ff3971305..c488c46ecc2 100644 --- a/litellm/types/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/model_management_endpoints.py @@ -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): diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 213297fc80a..5f5e2cf1ff8 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -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() + From 0c98efe045e3d50395854084a07d3e42f4fcd5c2 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 16 Jan 2026 18:14:33 -0800 Subject: [PATCH 2/2] fixing lint --- litellm/proxy/public_endpoints/public_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index aebea8616b4..6d60a218fd1 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -57,7 +57,7 @@ async def public_model_hub(): health_checks_map[key] = health_check_dict if check.model_name: health_checks_map[check.model_name] = health_check_dict - except Exception as e: + except Exception: pass for model_group in model_groups: