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

Implement health check backend API and storage functionality (#11678)

Browse Source 
* feat: Add health check functionality and endpoints

- Introduced methods for saving health check results to the database, including validation and cleaning of data.
- Added new health check endpoints to retrieve health check history and latest health statuses for models.
- Updated model prices and context window configuration for new Azure transcription models.

* test: Add unit tests for health check functionality

- Introduced tests for PrismaClient health check methods, including saving results and retrieving health check history.
- Added tests for the _save_health_check_to_db function to ensure proper handling of healthy and unhealthy endpoints.
- Implemented mock objects to simulate database interactions and validate method behaviors.

* Refactor health endpoint model ID handling and improve logging

- Updated health endpoint to use `get_deployment` for retrieving model names based on model IDs, enhancing error handling for missing models.
- Changed health check result saving to the database to be non-blocking by using `asyncio.create_task`.
- Cleaned up code for better readability and maintainability.

* Refactor utility functions in proxy module for improved readability and error handling

- Removed unused imports and simplified exception handling in `_get_redoc_url` and `_get_docs_url` functions to manage circular imports.
- Cleaned up logging statements for consistency and clarity.
- Streamlined error message formatting in `handle_exception_on_proxy` function.

* Enhance type hinting and default values in ProxyUpdateSpend class for improved clarity and robustness

- Added type hints for `_end_user_list_transactions` to specify it as a dictionary mapping end user IDs to spend amounts.
- Updated default values for optional fields in `SpendLogsPayload` to ensure they are initialized properly, enhancing error handling.
- Refactored `_premium_user_check` function to improve model validation logic and error handling.

* Fix disable_spend_updates method to handle None return value gracefully

- Updated the disable_spend_updates method to return False if the environment variable DISABLE_SPEND_UPDATES is not set or is None, improving robustness in configuration handling.

* Refactor join_paths function in utils.py for improved path handling

- Enhanced the join_paths function to better manage leading and trailing slashes, ensuring correct path concatenation.
- Added logic to handle cases where either base_path or route is empty, improving robustness and usability.

* Enhance health check functionality and improve error handling

- Introduced a new method `_save_health_check_to_db` for saving health check results to the database, utilizing safe JSON functions for data integrity.
- Refactored existing health check methods to streamline the process and improve error logging.
- Updated email sending logic to ensure secure connections and better error handling.
- Improved spend update logic with batch processing and retry mechanisms for database operations.
- Added utility functions for projected spend calculations and enhanced validation for team configurations.

* Add health check methods for database interaction

- Introduced `save_health_check_result` method to save health check results with detailed logging and validation.
- Added `get_health_check_history` method for retrieving health check records with optional filtering.
- Implemented `get_all_latest_health_checks` method to fetch the latest health checks for each model.
- Enhanced error handling and logging for all new methods to improve reliability and traceability.

* Refactor health check result saving to use typed arguments

- Updated the `_save_health_check_to_db` function to call `save_health_check_result` with explicitly typed arguments instead of a dictionary spread, enhancing code clarity and type safety.
- Removed unused method bindings in the mock Prisma client tests to streamline the test setup.

* Remove unused `_save_health_check_to_db` function from utils.py to streamline code and improve maintainability.

* Implement response time validation and details cleaning in health check result saving

- Added `_validate_response_time` method to ensure response time values are valid and handle exceptions gracefully.
- Introduced `_clean_details` method to validate and clean details JSON, improving data integrity.
- Refactored `save_health_check_result` to utilize these new methods for optional fields, enhancing code clarity and maintainability.
- Updated tests to bind new methods to the mock Prisma client for comprehensive testing.

* Add health check utility functions and refactor existing endpoints

- Introduced `_convert_health_check_to_dict` to standardize health check record conversion to dictionary format for JSON responses.
- Added `_check_prisma_client` helper function to streamline database availability checks and improve error handling.
- Refactored health check endpoints to utilize the new utility functions, enhancing code clarity and maintainability.

* Refactor health check tests for improved clarity and coverage

- Simplified the mock PrismaClient setup by consolidating method bindings.
- Updated health check result saving tests to use parameterized scenarios for better coverage.
- Added tests for health check history retrieval and graceful handling when no database client is provided.
- Removed redundant mock functions to streamline the test suite.

* Implement helper function for health check and database saving

- Added `_perform_health_check_and_save` to encapsulate health check execution and optional database saving.
- Refactored health endpoint logic to utilize the new helper function, improving code clarity and reducing redundancy.
- Enhanced error handling and streamlined the process of saving health check results to the database.
This commit is contained in:
Cole McIntosh 2025-06-14 15:27:11 -06:00 committed by GitHub
parent b7cb66ee8f
commit 5f34ceea1a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 471 additions and 36 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

255
litellm/proxy/health_endpoints/_health_endpoints.py
View File

@ -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()}

161
litellm/proxy/utils.py
View File

@ -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]:

91
tests/test_litellm/proxy/test_health_check_functions.py Normal file
View File

@ -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__])