diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 4af2d0d8e35..c0a0f788a17 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -1,6 +1,5 @@ # All settings - ```yaml environment_variables: {} @@ -95,6 +94,7 @@ general_settings: allowed_routes: ["route1", "route2"] # list of allowed proxy API routes - a user can access. (currently JWT-Auth only) key_management_system: google_kms # either google_kms or azure_kms master_key: string + maximum_spend_logs_retention_period: 30d # The maximum time to retain spend logs before deletion. # Database Settings database_url: string @@ -211,6 +211,7 @@ general_settings: | enable_oauth2_proxy_auth | boolean | (Enterprise Feature) If true, enables oauth2.0 authentication | | forward_openai_org_id | boolean | If true, forwards the OpenAI Organization ID to the backend LLM call (if it's OpenAI). | | forward_client_headers_to_llm_api | boolean | If true, forwards the client headers (any `x-` headers) to the backend LLM call | +| maximum_spend_logs_retention_period | str | Used to set the max retention time for spend logs in the db, after which they will be auto-purged | ### router_settings - Reference diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py new file mode 100644 index 00000000000..15d8c7b2e4a --- /dev/null +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -0,0 +1,39 @@ +""" +Handles checking if spend logs should be deleted based on maximum retention period +""" + +from typing import Optional, Union + +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.proxy.proxy_server import general_settings + + +def _should_delete_spend_logs() -> bool: + """ + Checks if the Pod should delete spend logs based on maximum retention period + + This setting enables automatic deletion of old spend logs to manage database size. + The maximum_spend_logs_retention_period can be specified in: + - Days (e.g., "30d") + - Hours (e.g., "24h") + - Minutes (e.g., "60m") + - Seconds (e.g., "3600s" or just "3600") + """ + _maximum_spend_logs_retention_period: Optional[Union[int, str]] = general_settings.get( + "maximum_spend_logs_retention_period", None + ) + + if _maximum_spend_logs_retention_period is None: + return False + + try: + if isinstance(_maximum_spend_logs_retention_period, int): + _maximum_spend_logs_retention_period = str(_maximum_spend_logs_retention_period) + duration_in_seconds(_maximum_spend_logs_retention_period) + return True + except ValueError as e: + verbose_proxy_logger.error( + f"Invalid maximum_spend_logs_retention_period value: {_maximum_spend_logs_retention_period}, error: {str(e)}" + ) + return False \ No newline at end of file diff --git a/tests/proxy_unit_tests/test_spend_log_cleanup.py b/tests/proxy_unit_tests/test_spend_log_cleanup.py new file mode 100644 index 00000000000..9c990ba2c86 --- /dev/null +++ b/tests/proxy_unit_tests/test_spend_log_cleanup.py @@ -0,0 +1,56 @@ +""" +Test cases for spend log cleanup functionality +""" + +import pytest +from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import _should_delete_spend_logs +from litellm.proxy.proxy_server import general_settings + + +@pytest.mark.asyncio +async def test_should_delete_spend_logs(): + """ + Test the _should_delete_spend_logs function with various scenarios + """ + # Test case 1: No maximum_spend_logs_retention_period set + general_settings.clear() + assert _should_delete_spend_logs() is False + + # Test case 2: Valid integer maximum_spend_logs_retention_period (in seconds) + general_settings["maximum_spend_logs_retention_period"] = 3600 + assert _should_delete_spend_logs() is True + + # Test case 3: Valid duration string - days + general_settings["maximum_spend_logs_retention_period"] = "30d" + assert _should_delete_spend_logs() is True + + # Test case 4: Valid duration string - hours + general_settings["maximum_spend_logs_retention_period"] = "24h" + assert _should_delete_spend_logs() is True + + # Test case 5: Valid duration string - minutes + general_settings["maximum_spend_logs_retention_period"] = "60m" + assert _should_delete_spend_logs() is True + + # Test case 6: Valid duration string - seconds + general_settings["maximum_spend_logs_retention_period"] = "3600s" + assert _should_delete_spend_logs() is True + + # Test case 7: Valid duration string - weeks + general_settings["maximum_spend_logs_retention_period"] = "1w" + assert _should_delete_spend_logs() is True + + # Test case 8: Valid duration string - months + general_settings["maximum_spend_logs_retention_period"] = "1mo" + assert _should_delete_spend_logs() is True + + # Test case 9: Invalid duration string + general_settings["maximum_spend_logs_retention_period"] = "invalid" + assert _should_delete_spend_logs() is False + + # Test case 10: None value + general_settings["maximum_spend_logs_retention_period"] = None + assert _should_delete_spend_logs() is False + + # Clean up + general_settings.clear() \ No newline at end of file