Merge pull request #19256 from BerriAI/litellm_model_hub_health

[Feature] /public/model_hub Health Information

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:
+            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
 
 

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):

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()
+