-
-
Notifications
You must be signed in to change notification settings - Fork 1.9k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
(feat) Allow disabling ErrorLogs written to the DB (#6940)
* fix - allow disabling logging error logs * docs on disabling error logs * doc string for _PROXY_failure_handler * test_disable_error_logs * rename file * fix rename file * increase test coverage for test_enable_error_logs
- Loading branch information
1 parent
0ac2d8b
commit 05f8109
Showing
5 changed files
with
218 additions
and
79 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,87 @@ | ||
| """ | ||
| Runs when LLM Exceptions occur on LiteLLM Proxy | ||
| """ | ||
|
|
||
| import copy | ||
| import json | ||
| import uuid | ||
|
|
||
| import litellm | ||
| from litellm.proxy._types import LiteLLM_ErrorLogs | ||
|
|
||
|
|
||
| async def _PROXY_failure_handler( | ||
| kwargs, # kwargs to completion | ||
| completion_response: litellm.ModelResponse, # response from completion | ||
| start_time=None, | ||
| end_time=None, # start/end time for completion | ||
| ): | ||
| """ | ||
| Async Failure Handler - runs when LLM Exceptions occur on LiteLLM Proxy. | ||
| This function logs the errors to the Prisma DB | ||
| Can be disabled by setting the following on proxy_config.yaml: | ||
| ```yaml | ||
| general_settings: | ||
| disable_error_logs: True | ||
| ``` | ||
| """ | ||
| from litellm._logging import verbose_proxy_logger | ||
| from litellm.proxy.proxy_server import general_settings, prisma_client | ||
|
|
||
| if general_settings.get("disable_error_logs") is True: | ||
| return | ||
|
|
||
| if prisma_client is not None: | ||
| verbose_proxy_logger.debug( | ||
| "inside _PROXY_failure_handler kwargs=", extra=kwargs | ||
| ) | ||
|
|
||
| _exception = kwargs.get("exception") | ||
| _exception_type = _exception.__class__.__name__ | ||
| _model = kwargs.get("model", None) | ||
|
|
||
| _optional_params = kwargs.get("optional_params", {}) | ||
| _optional_params = copy.deepcopy(_optional_params) | ||
|
|
||
| for k, v in _optional_params.items(): | ||
| v = str(v) | ||
| v = v[:100] | ||
|
|
||
| _status_code = "500" | ||
| try: | ||
| _status_code = str(_exception.status_code) | ||
| except Exception: | ||
| # Don't let this fail logging the exception to the dB | ||
| pass | ||
|
|
||
| _litellm_params = kwargs.get("litellm_params", {}) or {} | ||
| _metadata = _litellm_params.get("metadata", {}) or {} | ||
| _model_id = _metadata.get("model_info", {}).get("id", "") | ||
| _model_group = _metadata.get("model_group", "") | ||
| api_base = litellm.get_api_base(model=_model, optional_params=_litellm_params) | ||
| _exception_string = str(_exception) | ||
|
|
||
| error_log = LiteLLM_ErrorLogs( | ||
| request_id=str(uuid.uuid4()), | ||
| model_group=_model_group, | ||
| model_id=_model_id, | ||
| litellm_model_name=kwargs.get("model"), | ||
| request_kwargs=_optional_params, | ||
| api_base=api_base, | ||
| exception_type=_exception_type, | ||
| status_code=_status_code, | ||
| exception_string=_exception_string, | ||
| startTime=kwargs.get("start_time"), | ||
| endTime=kwargs.get("end_time"), | ||
| ) | ||
|
|
||
| error_log_dict = error_log.model_dump() | ||
| error_log_dict["request_kwargs"] = json.dumps(error_log_dict["request_kwargs"]) | ||
|
|
||
| await prisma_client.db.litellm_errorlogs.create( | ||
| data=error_log_dict # type: ignore | ||
| ) | ||
|
|
||
| pass |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,111 @@ | ||
| import asyncio | ||
| import os | ||
| import sys | ||
| from unittest.mock import Mock, patch, AsyncMock | ||
| import pytest | ||
| from fastapi import Request | ||
| from litellm.proxy.utils import _get_redoc_url, _get_docs_url | ||
|
|
||
| sys.path.insert(0, os.path.abspath("../..")) | ||
| import litellm | ||
|
|
||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_disable_error_logs(): | ||
| """ | ||
| Test that the error logs are not written to the database when disable_error_logs is True | ||
| """ | ||
| # Mock the necessary components | ||
| mock_prisma_client = AsyncMock() | ||
| mock_general_settings = {"disable_error_logs": True} | ||
|
|
||
| with patch( | ||
| "litellm.proxy.proxy_server.general_settings", mock_general_settings | ||
| ), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): | ||
|
|
||
| # Create a test exception | ||
| test_exception = Exception("Test error") | ||
| test_kwargs = { | ||
| "model": "gpt-4", | ||
| "exception": test_exception, | ||
| "optional_params": {}, | ||
| "litellm_params": {"metadata": {}}, | ||
| } | ||
|
|
||
| # Call the failure handler | ||
| from litellm.proxy.proxy_server import _PROXY_failure_handler | ||
|
|
||
| await _PROXY_failure_handler( | ||
| kwargs=test_kwargs, | ||
| completion_response=None, | ||
| start_time="2024-01-01", | ||
| end_time="2024-01-01", | ||
| ) | ||
|
|
||
| # Verify prisma client was not called to create error logs | ||
| if hasattr(mock_prisma_client, "db"): | ||
| assert not mock_prisma_client.db.litellm_errorlogs.create.called | ||
|
|
||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_disable_spend_logs(): | ||
| """ | ||
| Test that the spend logs are not written to the database when disable_spend_logs is True | ||
| """ | ||
| # Mock the necessary components | ||
| mock_prisma_client = Mock() | ||
| mock_prisma_client.spend_log_transactions = [] | ||
|
|
||
| with patch("litellm.proxy.proxy_server.disable_spend_logs", True), patch( | ||
| "litellm.proxy.proxy_server.prisma_client", mock_prisma_client | ||
| ): | ||
| from litellm.proxy.proxy_server import update_database | ||
|
|
||
| # Call update_database with disable_spend_logs=True | ||
| await update_database( | ||
| token="fake-token", | ||
| response_cost=0.1, | ||
| user_id="user123", | ||
| completion_response=None, | ||
| start_time="2024-01-01", | ||
| end_time="2024-01-01", | ||
| ) | ||
| # Verify no spend logs were added | ||
| assert len(mock_prisma_client.spend_log_transactions) == 0 | ||
|
|
||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_enable_error_logs(): | ||
| """ | ||
| Test that the error logs are written to the database when disable_error_logs is False | ||
| """ | ||
| # Mock the necessary components | ||
| mock_prisma_client = AsyncMock() | ||
| mock_general_settings = {"disable_error_logs": False} | ||
|
|
||
| with patch( | ||
| "litellm.proxy.proxy_server.general_settings", mock_general_settings | ||
| ), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): | ||
|
|
||
| # Create a test exception | ||
| test_exception = Exception("Test error") | ||
| test_kwargs = { | ||
| "model": "gpt-4", | ||
| "exception": test_exception, | ||
| "optional_params": {}, | ||
| "litellm_params": {"metadata": {}}, | ||
| } | ||
|
|
||
| # Call the failure handler | ||
| from litellm.proxy.proxy_server import _PROXY_failure_handler | ||
|
|
||
| await _PROXY_failure_handler( | ||
| kwargs=test_kwargs, | ||
| completion_response=None, | ||
| start_time="2024-01-01", | ||
| end_time="2024-01-01", | ||
| ) | ||
|
|
||
| # Verify prisma client was called to create error logs | ||
| if hasattr(mock_prisma_client, "db"): | ||
| assert mock_prisma_client.db.litellm_errorlogs.create.called |