mirror of
https://github.com/tiennm99/litellm.git
synced 2026-06-29 19:04:28 +00:00
Cole McIntosh
7d1e0651d8
* 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. * refactor: rename target_model parameter to model in health check function
91 lines
3.2 KiB
Python
91 lines
3.2 KiB
Python
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__]) |