diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 5996f4055f..e1286d0aca 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -1,6 +1,7 @@ import asyncio import copy import os +import time import traceback from datetime import datetime, timedelta from typing import Dict, Literal, Optional, Union @@ -302,12 +303,120 @@ async def health_services_endpoint( # noqa: PLR0915 ) +def _convert_health_check_to_dict(check) -> dict: + """Convert health check database record to dictionary format""" + return { + "health_check_id": check.health_check_id, + "model_name": check.model_name, + "model_id": check.model_id, + "status": check.status, + "healthy_count": check.healthy_count, + "unhealthy_count": check.unhealthy_count, + "error_message": check.error_message, + "response_time_ms": check.response_time_ms, + "details": check.details, + "checked_by": check.checked_by, + "checked_at": check.checked_at.isoformat() if check.checked_at else None, + "created_at": check.created_at.isoformat() if check.created_at else None, + } + + +def _check_prisma_client(): + """Helper to check if prisma_client is available and raise appropriate error""" + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={"error": "Database not initialized"}, + ) + return prisma_client + + +async def _save_health_check_to_db( + prisma_client, + model_name: str, + healthy_endpoints: list, + unhealthy_endpoints: list, + start_time: float, + user_id: Optional[str], + model_id: Optional[str] = None +): + """Helper function to save health check results to database""" + try: + # Extract error message from first unhealthy endpoint if available + error_message = ( + str(unhealthy_endpoints[0]["error"])[:500] + if unhealthy_endpoints and unhealthy_endpoints[0].get("error") + else None + ) + + await prisma_client.save_health_check_result( + model_name=model_name, + model_id=model_id, + status="healthy" if healthy_endpoints else "unhealthy", + healthy_count=len(healthy_endpoints), + unhealthy_count=len(unhealthy_endpoints), + error_message=error_message, + response_time_ms=(time.time() - start_time) * 1000, + details=None, # Skip details for now to avoid JSON serialization issues + checked_by=user_id, + ) + except Exception as db_error: + verbose_proxy_logger.warning(f"Failed to save health check to database for model {model_name}: {db_error}") + # Continue execution - don't let database save failure break health checks + + +async def _perform_health_check_and_save( + model_list, + target_model, + cli_model, + details, + prisma_client, + start_time, + user_id, + model_id=None +): + """Helper function to perform health check and save results to database""" + healthy_endpoints, unhealthy_endpoints = await perform_health_check( + model_list=model_list, + cli_model=cli_model, + target_model=target_model, + details=details + ) + + # Optionally save health check result to database (non-blocking) + if prisma_client is not None: + # For CLI model, use cli_model name; for router models, use target_model + model_name_for_db = cli_model if cli_model is not None else target_model + if model_name_for_db is not None: + asyncio.create_task(_save_health_check_to_db( + prisma_client, + model_name_for_db, + healthy_endpoints, + unhealthy_endpoints, + start_time, + user_id, + model_id=model_id + )) + + return { + "healthy_endpoints": healthy_endpoints, + "unhealthy_endpoints": unhealthy_endpoints, + "healthy_count": len(healthy_endpoints), + "unhealthy_count": len(unhealthy_endpoints), + } + + @router.get("/health", tags=["health"], dependencies=[Depends(user_api_key_auth)]) async def health_endpoint( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), model: Optional[str] = fastapi.Query( None, description="Specify the model name (optional)" ), + model_id: Optional[str] = fastapi.Query( + None, description="Specify the model ID (optional)" + ), ): """ 🚨 USE `/health/liveliness` to health check the proxy 🚨 @@ -329,23 +438,55 @@ async def health_endpoint( health_check_details, health_check_results, llm_model_list, + llm_router, use_background_health_checks, user_model, + prisma_client, ) + import time + start_time = time.time() + + # Handle model_id parameter - convert to model name for health check + target_model = model + if model_id and not model: + # Use get_deployment from router to find the model name + if llm_router is not None: + try: + deployment = llm_router.get_deployment(model_id=model_id) + if deployment is not None: + target_model = deployment.model_name + else: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": f"Model with ID {model_id} not found"}, + ) + except Exception as e: + verbose_proxy_logger.error(f"Error getting deployment for model_id {model_id}: {e}") + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": f"Model with ID {model_id} not found"}, + ) + else: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": f"Model with ID {model_id} not found"}, + ) + try: if llm_model_list is None: # if no router set, check if user set a model using litellm --model ollama/llama2 if user_model is not None: - healthy_endpoints, unhealthy_endpoints = await perform_health_check( - model_list=[], cli_model=user_model, details=health_check_details + return await _perform_health_check_and_save( + model_list=[], + target_model=None, + cli_model=user_model, + details=health_check_details, + prisma_client=prisma_client, + start_time=start_time, + user_id=user_api_key_dict.user_id, + model_id=None # CLI model doesn't have model_id ) - return { - "healthy_endpoints": healthy_endpoints, - "unhealthy_endpoints": unhealthy_endpoints, - "healthy_count": len(healthy_endpoints), - "unhealthy_count": len(unhealthy_endpoints), - } raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail={"error": "Model list not initialized"}, @@ -359,16 +500,16 @@ async def health_endpoint( if use_background_health_checks: return health_check_results else: - healthy_endpoints, unhealthy_endpoints = await perform_health_check( - _llm_model_list, model, details=health_check_details + return await _perform_health_check_and_save( + model_list=_llm_model_list, + target_model=target_model, + cli_model=None, + details=health_check_details, + prisma_client=prisma_client, + start_time=start_time, + user_id=user_api_key_dict.user_id, + model_id=model_id ) - - return { - "healthy_endpoints": healthy_endpoints, - "unhealthy_endpoints": unhealthy_endpoints, - "healthy_count": len(healthy_endpoints), - "unhealthy_count": len(unhealthy_endpoints), - } except Exception as e: verbose_proxy_logger.error( "litellm.proxy.proxy_server.py::health_endpoint(): Exception occured - {}".format( @@ -379,6 +520,86 @@ async def health_endpoint( raise e +@router.get("/health/history", tags=["health"], dependencies=[Depends(user_api_key_auth)]) +async def health_check_history_endpoint( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + model: Optional[str] = fastapi.Query( + None, description="Filter by specific model name" + ), + status_filter: Optional[str] = fastapi.Query( + None, description="Filter by status (healthy/unhealthy)" + ), + limit: int = fastapi.Query( + 100, description="Number of records to return", ge=1, le=1000 + ), + offset: int = fastapi.Query( + 0, description="Number of records to skip", ge=0 + ), +): + """ + Get health check history for models + + Returns historical health check data with optional filtering. + """ + prisma_client = _check_prisma_client() + + try: + history = await prisma_client.get_health_check_history( + model_name=model, + limit=limit, + offset=offset, + status_filter=status_filter, + ) + + # Convert to dict format for JSON response using helper function + history_data = [_convert_health_check_to_dict(check) for check in history] + + return { + "health_checks": history_data, + "total_records": len(history_data), + "limit": limit, + "offset": offset, + } + except Exception as e: + verbose_proxy_logger.error(f"Error getting health check history: {e}") + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={"error": f"Failed to retrieve health check history: {str(e)}"}, + ) + + +@router.get("/health/latest", tags=["health"], dependencies=[Depends(user_api_key_auth)]) +async def latest_health_checks_endpoint( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Get the latest health check status for all models + + Returns the most recent health check result for each model. + """ + prisma_client = _check_prisma_client() + + try: + latest_checks = await prisma_client.get_all_latest_health_checks() + + # Convert to dict format for JSON response using helper function + checks_data = { + (check.model_id if check.model_id else check.model_name): _convert_health_check_to_dict(check) + for check in latest_checks + } + + return { + "latest_health_checks": checks_data, + "total_models": len(checks_data), + } + except Exception as e: + verbose_proxy_logger.error(f"Error getting latest health checks: {e}") + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={"error": f"Failed to retrieve latest health checks: {str(e)}"}, + ) + + db_health_cache = {"status": "unknown", "last_updated": datetime.now()} diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 091fe0a173..fdf21dcba2 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -54,6 +54,8 @@ from litellm import ( ) from litellm._logging import verbose_proxy_logger from litellm._service_logger import ServiceLogging, ServiceTypes +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.caching.caching import DualCache, RedisCache from litellm.exceptions import RejectedRequestError from litellm.integrations.custom_guardrail import CustomGuardrail @@ -2458,8 +2460,121 @@ class PrismaClient: value=_num_spend_logs_rows, ) + # Health Check Database Methods + def _validate_response_time(self, response_time_ms: Optional[float]) -> Optional[float]: + """Validate and clean response time value""" + if response_time_ms is None: + return None + try: + value = float(response_time_ms) + return value if value == value and value not in (float('inf'), float('-inf')) else None + except (ValueError, TypeError): + verbose_proxy_logger.warning(f"Invalid response_time_ms value: {response_time_ms}") + return None + + def _clean_details(self, details: Optional[dict]) -> Optional[dict]: + """Clean and validate details JSON""" + if not isinstance(details, dict): + return None + try: + return safe_json_loads(safe_dumps(details)) + except Exception as e: + verbose_proxy_logger.warning(f"Failed to clean details JSON: {e}") + return None + + async def save_health_check_result( + self, + model_name: str, + status: str, + healthy_count: int = 0, + unhealthy_count: int = 0, + error_message: Optional[str] = None, + response_time_ms: Optional[float] = None, + details: Optional[dict] = None, + checked_by: Optional[str] = None, + model_id: Optional[str] = None, + ): + """Save health check result to database""" + try: + # Build base data with required fields + health_check_data = { + "model_name": str(model_name), + "status": str(status), + "healthy_count": int(healthy_count), + "unhealthy_count": int(unhealthy_count), + } + + # Add optional fields using dict comprehension and helper methods + optional_fields = { + "error_message": str(error_message)[:500] if error_message else None, + "response_time_ms": self._validate_response_time(response_time_ms), + "details": self._clean_details(details), + "checked_by": str(checked_by) if checked_by else None, + "model_id": str(model_id) if model_id else None, + } + + # Add only non-None optional fields + health_check_data.update({k: v for k, v in optional_fields.items() if v is not None}) + + verbose_proxy_logger.debug(f"Saving health check data: {health_check_data}") + return await self.db.litellm_healthchecktable.create(data=health_check_data) + + except Exception as e: + verbose_proxy_logger.error(f"Error saving health check result for model {model_name}: {e}") + return None + + async def get_health_check_history( + self, + model_name: Optional[str] = None, + limit: int = 100, + offset: int = 0, + status_filter: Optional[str] = None, + ): + """ + Get health check history with optional filtering + """ + try: + where_clause = {} + if model_name: + where_clause["model_name"] = model_name + if status_filter: + where_clause["status"] = status_filter + + results = await self.db.litellm_healthchecktable.find_many( + where=where_clause, + order={"checked_at": "desc"}, + take=limit, + skip=offset, + ) + return results + except Exception as e: + verbose_proxy_logger.error(f"Error getting health check history: {e}") + return [] + + async def get_all_latest_health_checks(self): + """ + Get the latest health check for each model + """ + try: + # Get all unique model names first + all_checks = await self.db.litellm_healthchecktable.find_many( + order={"checked_at": "desc"} + ) + + # Group by model_name and get the latest for each + latest_checks = {} + for check in all_checks: + if check.model_name not in latest_checks: + latest_checks[check.model_name] = check + + return list(latest_checks.values()) + except Exception as e: + verbose_proxy_logger.error(f"Error getting all latest health checks: {e}") + return [] + ### HELPER FUNCTIONS ### + async def _cache_user_row(user_id: str, cache: DualCache, db: PrismaClient): """ Check if a user_id exists in cache, @@ -2957,32 +3072,40 @@ def is_known_model(model: Optional[str], llm_router: Optional[Router]) -> bool: def join_paths(base_path: str, route: str) -> str: - # Remove trailing/leading slashes + # Remove trailing slashes from base_path and leading slashes from route base_path = base_path.rstrip("/") route = route.lstrip("/") - - # Join with a single slash + + # If base_path is empty, return route with leading slash + if not base_path: + return f"/{route}" if route else "/" + + # If route is empty, return just base_path + if not route: + return base_path + + # Join with single slash return f"{base_path}/{route}" def get_custom_url(request_base_url: str, route: Optional[str] = None) -> str: - """ - Use proxy base url, if set. - - Else, use request base url. - """ - from httpx import URL - - proxy_base_url = os.getenv("PROXY_BASE_URL") - server_root_path = os.getenv("SERVER_ROOT_PATH") or "" - if route is not None: - server_root_path = join_paths(base_path=server_root_path, route=route) - if proxy_base_url: - ui_link = str(URL(proxy_base_url).join(server_root_path)) + # Use environment variable value, otherwise use URL from request + server_base_url = get_proxy_base_url() + if server_base_url is not None: + base_url = server_base_url else: - ui_link = str(URL(request_base_url).join(server_root_path)) - - return ui_link + base_url = request_base_url + + server_root_path = get_server_root_path() + if route is not None: + if server_root_path != "": + # First join base_url with server_root_path, then with route + intermediate_url = join_paths(base_url, server_root_path) + return join_paths(intermediate_url, route) + else: + return join_paths(base_url, route) + else: + return join_paths(base_url, server_root_path) def get_proxy_base_url() -> Optional[str]: diff --git a/tests/test_litellm/proxy/test_health_check_functions.py b/tests/test_litellm/proxy/test_health_check_functions.py new file mode 100644 index 0000000000..4f014ce1be --- /dev/null +++ b/tests/test_litellm/proxy/test_health_check_functions.py @@ -0,0 +1,91 @@ +import asyncio +import pytest +from unittest.mock import AsyncMock, MagicMock +import sys +import os + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.utils import PrismaClient +from litellm.proxy.health_endpoints._health_endpoints import _save_health_check_to_db + + +@pytest.fixture +def mock_prisma(): + """Simplified mock PrismaClient with bound methods""" + client = MagicMock() + client.db.litellm_healthchecktable.create = AsyncMock(return_value={"id": "test-id"}) + client.db.litellm_healthchecktable.find_many = AsyncMock(return_value=[{"id": "1", "model_name": "test"}]) + + # Bind actual methods + import types + for method in ['save_health_check_result', '_validate_response_time', '_clean_details', + 'get_health_check_history', 'get_all_latest_health_checks']: + setattr(client, method, types.MethodType(getattr(PrismaClient, method), client)) + + return client + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status,healthy,unhealthy,should_succeed", [ + ("healthy", 1, 0, True), + ("unhealthy", 0, 1, True), + ("healthy", 1, 0, False), # Database error case +]) +async def test_save_health_check_result(mock_prisma, status, healthy, unhealthy, should_succeed): + """Test health check result saving with various scenarios""" + if not should_succeed: + mock_prisma.db.litellm_healthchecktable.create.side_effect = Exception("DB Error") + + result = await mock_prisma.save_health_check_result( + model_name="test-model", status=status, healthy_count=healthy, unhealthy_count=unhealthy + ) + + if should_succeed: + mock_prisma.db.litellm_healthchecktable.create.assert_called_once() + else: + assert result is None + + +@pytest.mark.asyncio +async def test_get_health_check_history(mock_prisma): + """Test health check history retrieval""" + result = await mock_prisma.get_health_check_history(model_name="test", limit=50) + mock_prisma.db.litellm_healthchecktable.find_many.assert_called_once() + assert len(result) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("healthy_count,unhealthy_count,expected_status", [ + (1, 0, "healthy"), + (0, 1, "unhealthy"), + (2, 1, "healthy"), +]) +async def test_save_health_check_to_db(healthy_count, unhealthy_count, expected_status): + """Test _save_health_check_to_db function with different endpoint counts""" + mock_client = MagicMock() + mock_client.save_health_check_result = AsyncMock() + + healthy_endpoints = [{"model": "test"}] * healthy_count + unhealthy_endpoints = [{"error": "test error"}] * unhealthy_count + + await _save_health_check_to_db( + mock_client, "test-model", healthy_endpoints, unhealthy_endpoints, + 1234567890.0, "test-user" + ) + + call_args = mock_client.save_health_check_result.call_args[1] + assert call_args["status"] == expected_status + assert call_args["healthy_count"] == healthy_count + assert call_args["unhealthy_count"] == unhealthy_count + + +@pytest.mark.asyncio +async def test_save_health_check_to_db_no_client(): + """Test graceful handling when no database client""" + result = await _save_health_check_to_db(None, "test", [], [], 0.0, "user") + assert result is None + + +if __name__ == "__main__": + pytest.main([__file__]) \ No newline at end of file