From 28ca9912967093f6ee1c41a93c31370c85fc7e83 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 27 Jan 2026 20:52:07 -0800 Subject: [PATCH] Allow dynamic setting of store_prompts_in_spend_logs --- litellm/proxy/_types.py | 8 + litellm/proxy/proxy_server.py | 101 ++++++++- .../spend_tracking/spend_tracking_utils.py | 17 +- .../test_spend_tracking_utils.py | 46 +++++ tests/test_litellm/proxy/test_proxy_server.py | 194 ++++++++++++++++++ 5 files changed, 361 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c854d81ec71..e15179a8e29 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2073,6 +2073,14 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="Controls how non-admin users interact with MCP servers in the dashboard. 'restricted' shows only accessible servers, 'view_all' lists every server in read-only mode.", ) + store_prompts_in_spend_logs: Optional[bool] = Field( + None, + description="If True, stores request messages and responses in spend logs. Default is False.", + ) + maximum_spend_logs_retention_period: Optional[str] = Field( + None, + description="Maximum retention period for spend logs (e.g., '7d' for 7 days). Logs older than this will be deleted.", + ) class ConfigYAML(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 12e34e221d7..e7fcd5f4422 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3540,6 +3540,79 @@ class ProxyConfig: llm_router=llm_router, ) + async def _reschedule_spend_log_cleanup_job(self): + """ + Reschedule the spend log cleanup job based on current general_settings. + This is called when maximum_spend_logs_retention_period is updated dynamically. + If the retention period is None, the job will be removed. + """ + global scheduler, general_settings, prisma_client + if scheduler is None: + return + + # Remove existing job if it exists + try: + scheduler.remove_job("spend_log_cleanup_job") + verbose_proxy_logger.info("Removed existing spend log cleanup job") + except Exception: + pass # Job might not exist, which is fine + + # Schedule new job if retention period is set (not None) + retention_period = general_settings.get("maximum_spend_logs_retention_period") + if retention_period is not None: + from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import ( + SpendLogCleanup, + ) + + spend_log_cleanup = SpendLogCleanup() + cleanup_cron = general_settings.get("maximum_spend_logs_cleanup_cron") + + if cleanup_cron: + from apscheduler.triggers.cron import CronTrigger + + try: + cron_trigger = CronTrigger.from_crontab(cleanup_cron) + scheduler.add_job( + spend_log_cleanup.cleanup_old_spend_logs, + cron_trigger, + args=[prisma_client], + id="spend_log_cleanup_job", + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + verbose_proxy_logger.info( + f"Spend log cleanup rescheduled with cron: {cleanup_cron}" + ) + except ValueError: + verbose_proxy_logger.error( + f"Invalid maximum_spend_logs_cleanup_cron value: {cleanup_cron}" + ) + else: + # Interval-based scheduling (existing behavior) + from litellm.litellm_core_utils.duration_parser import duration_in_seconds + + retention_interval = general_settings.get( + "maximum_spend_logs_retention_interval", "1d" + ) + try: + interval_seconds = duration_in_seconds(retention_interval) + scheduler.add_job( + spend_log_cleanup.cleanup_old_spend_logs, + "interval", + seconds=interval_seconds + random.randint(0, 60), + args=[prisma_client], + id="spend_log_cleanup_job", + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + verbose_proxy_logger.info( + f"Spend log cleanup rescheduled with interval: {retention_interval}" + ) + except ValueError: + verbose_proxy_logger.error( + "Invalid maximum_spend_logs_retention_interval value" + ) + async def _update_general_settings(self, db_general_settings: Optional[Json]): """ Pull from DB, read general settings value @@ -3579,6 +3652,32 @@ class ProxyConfig: if "ui_access_mode" in _general_settings: general_settings["ui_access_mode"] = _general_settings["ui_access_mode"] + ## STORE PROMPTS IN SPEND LOGS ## + if "store_prompts_in_spend_logs" in _general_settings: + value = _general_settings["store_prompts_in_spend_logs"] + # Normalize case: handle True/true/TRUE, False/false/FALSE, None/null + if value is None: + general_settings["store_prompts_in_spend_logs"] = None + elif isinstance(value, bool): + general_settings["store_prompts_in_spend_logs"] = value + elif isinstance(value, str): + # Case-insensitive string comparison + general_settings["store_prompts_in_spend_logs"] = ( + value.lower() == "true" + ) + else: + # For other types, convert to bool + general_settings["store_prompts_in_spend_logs"] = bool(value) + + ## MAXIMUM SPEND LOGS RETENTION PERIOD ## + if "maximum_spend_logs_retention_period" in _general_settings: + old_value = general_settings.get("maximum_spend_logs_retention_period") + new_value = _general_settings["maximum_spend_logs_retention_period"] + general_settings["maximum_spend_logs_retention_period"] = new_value + # Reschedule cleanup job if value changed (including when set to None) + if old_value != new_value: + await self._reschedule_spend_log_cleanup_job() + def _update_config_fields( self, current_config: dict, @@ -4688,7 +4787,7 @@ class ProxyStartupEvent: proxy_logging_obj: ProxyLogging, ): """Initializes scheduled background jobs""" - global store_model_in_db + global store_model_in_db, scheduler # MEMORY LEAK FIX: Configure scheduler with optimized settings # Memray analysis showed APScheduler's normalize() and _apply_jitter() causing diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index db4e4beec21..bd148ecb481 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -759,10 +759,19 @@ def _should_store_prompts_and_responses_in_spend_logs() -> bool: from litellm.proxy.proxy_server import general_settings from litellm.secret_managers.main import get_secret_bool - return ( - general_settings.get("store_prompts_in_spend_logs") is True - or get_secret_bool("STORE_PROMPTS_IN_SPEND_LOGS") is True - ) + # Check general_settings (from DB or proxy_config.yaml) + store_prompts_value = general_settings.get("store_prompts_in_spend_logs") + + # Normalize case: handle True/true/TRUE, False/false/FALSE, None/null + if store_prompts_value is True: + return True + elif isinstance(store_prompts_value, str): + # Case-insensitive string comparison + if store_prompts_value.lower() == "true": + return True + + # Also check environment variable + return get_secret_bool("STORE_PROMPTS_IN_SPEND_LOGS") is True def _get_status_for_spend_log( diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index dd4cbfa2e7f..1972103c3d2 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -23,6 +23,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( _get_response_for_spend_logs_payload, _get_vector_store_request_for_spend_logs_payload, _sanitize_request_body_for_spend_logs_payload, + _should_store_prompts_and_responses_in_spend_logs, get_logging_payload, ) from litellm.types.utils import ( @@ -911,3 +912,48 @@ def test_spend_logs_redacts_request_and_response_when_turn_off_message_logging_e parsed_response = json.loads(response_result) assert parsed_response == {"text": "redacted-by-litellm"} + +@patch("litellm.secret_managers.main.get_secret_bool") +def test_should_store_prompts_and_responses_in_spend_logs_case_insensitive_string( + mock_get_secret_bool, +): + """ + Test that _should_store_prompts_and_responses_in_spend_logs handles + case-insensitive string values for store_prompts_in_spend_logs in general_settings. + """ + # Test case-insensitive string "true" variations + for true_value in ["true", "TRUE", "True", "TrUe"]: + with patch("litellm.proxy.proxy_server.general_settings", {"store_prompts_in_spend_logs": true_value}): + mock_get_secret_bool.return_value = False # Ensure env var is False + result = _should_store_prompts_and_responses_in_spend_logs() + assert result is True, f"Expected True for '{true_value}', got {result}" + + # Test boolean True + with patch("litellm.proxy.proxy_server.general_settings", {"store_prompts_in_spend_logs": True}): + mock_get_secret_bool.return_value = False + result = _should_store_prompts_and_responses_in_spend_logs() + assert result is True, f"Expected True for boolean True, got {result}" + + # Test that non-true values fall back to environment variable + for false_value in [False, None, "false", "FALSE", "False", "anything"]: + with patch("litellm.proxy.proxy_server.general_settings", {"store_prompts_in_spend_logs": false_value}): + # When env var is True, should return True + mock_get_secret_bool.return_value = True + result = _should_store_prompts_and_responses_in_spend_logs() + assert result is True, f"Expected True (from env var) for '{false_value}', got {result}" + + # When env var is False, should return False + mock_get_secret_bool.return_value = False + result = _should_store_prompts_and_responses_in_spend_logs() + assert result is False, f"Expected False (from env var) for '{false_value}', got {result}" + + # Test when general_settings doesn't have the key at all + with patch("litellm.proxy.proxy_server.general_settings", {}): + mock_get_secret_bool.return_value = True + result = _should_store_prompts_and_responses_in_spend_logs() + assert result is True, "Expected True (from env var) when key missing, got False" + + mock_get_secret_bool.return_value = False + result = _should_store_prompts_and_responses_in_spend_logs() + assert result is False, "Expected False (from env var) when key missing, got True" + diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 22a60f220ee..9069a57643b 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5121,3 +5121,197 @@ async def test_model_list_no_scope_parameter(monkeypatch): # Verify router methods were NOT called (normal path) mock_router.get_model_names.assert_not_called() mock_router.get_model_access_groups.assert_not_called() + + +@pytest.mark.asyncio +async def test_update_general_settings_store_prompts_in_spend_logs(monkeypatch): + """ + Test that _update_general_settings correctly normalizes store_prompts_in_spend_logs + values (handles bool, string, None, and other types). + """ + from unittest.mock import patch + + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + # Test Case 1: None value + mock_gs = {} + with patch("litellm.proxy.proxy_server.general_settings", mock_gs): + await proxy_config._update_general_settings( + {"store_prompts_in_spend_logs": None} + ) + assert mock_gs.get("store_prompts_in_spend_logs") is None + + # Test Case 2: bool True + mock_gs = {} + with patch("litellm.proxy.proxy_server.general_settings", mock_gs): + await proxy_config._update_general_settings( + {"store_prompts_in_spend_logs": True} + ) + assert mock_gs.get("store_prompts_in_spend_logs") is True + + # Test Case 3: bool False + mock_gs = {} + with patch("litellm.proxy.proxy_server.general_settings", mock_gs): + await proxy_config._update_general_settings( + {"store_prompts_in_spend_logs": False} + ) + assert mock_gs.get("store_prompts_in_spend_logs") is False + + # Test Case 4: string "true" (lowercase) + mock_gs = {} + with patch("litellm.proxy.proxy_server.general_settings", mock_gs): + await proxy_config._update_general_settings( + {"store_prompts_in_spend_logs": "true"} + ) + assert mock_gs.get("store_prompts_in_spend_logs") is True + + # Test Case 5: string "True" (capitalized) + mock_gs = {} + with patch("litellm.proxy.proxy_server.general_settings", mock_gs): + await proxy_config._update_general_settings( + {"store_prompts_in_spend_logs": "True"} + ) + assert mock_gs.get("store_prompts_in_spend_logs") is True + + # Test Case 6: string "TRUE" (uppercase) + mock_gs = {} + with patch("litellm.proxy.proxy_server.general_settings", mock_gs): + await proxy_config._update_general_settings( + {"store_prompts_in_spend_logs": "TRUE"} + ) + assert mock_gs.get("store_prompts_in_spend_logs") is True + + # Test Case 7: string "false" (lowercase) + mock_gs = {} + with patch("litellm.proxy.proxy_server.general_settings", mock_gs): + await proxy_config._update_general_settings( + {"store_prompts_in_spend_logs": "false"} + ) + assert mock_gs.get("store_prompts_in_spend_logs") is False + + # Test Case 8: string "False" (capitalized) + mock_gs = {} + with patch("litellm.proxy.proxy_server.general_settings", mock_gs): + await proxy_config._update_general_settings( + {"store_prompts_in_spend_logs": "False"} + ) + assert mock_gs.get("store_prompts_in_spend_logs") is False + + # Test Case 9: string "FALSE" (uppercase) + mock_gs = {} + with patch("litellm.proxy.proxy_server.general_settings", mock_gs): + await proxy_config._update_general_settings( + {"store_prompts_in_spend_logs": "FALSE"} + ) + assert mock_gs.get("store_prompts_in_spend_logs") is False + + # Test Case 10: other string value (should be False) + mock_gs = {} + with patch("litellm.proxy.proxy_server.general_settings", mock_gs): + await proxy_config._update_general_settings( + {"store_prompts_in_spend_logs": "invalid"} + ) + assert mock_gs.get("store_prompts_in_spend_logs") is False + + # Test Case 11: integer 1 (should be True) + mock_gs = {} + with patch("litellm.proxy.proxy_server.general_settings", mock_gs): + await proxy_config._update_general_settings( + {"store_prompts_in_spend_logs": 1} + ) + assert mock_gs.get("store_prompts_in_spend_logs") is True + + # Test Case 12: integer 0 (should be False) + mock_gs = {} + with patch("litellm.proxy.proxy_server.general_settings", mock_gs): + await proxy_config._update_general_settings( + {"store_prompts_in_spend_logs": 0} + ) + assert mock_gs.get("store_prompts_in_spend_logs") is False + + +@pytest.mark.asyncio +async def test_update_general_settings_maximum_spend_logs_retention_period(monkeypatch): + """ + Test that _update_general_settings correctly handles maximum_spend_logs_retention_period + and reschedules cleanup job when value changes. + """ + from unittest.mock import AsyncMock, patch + + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + # Test Case 1: Setting a new value should reschedule cleanup job + mock_reschedule = AsyncMock() + mock_gs = {} + with patch("litellm.proxy.proxy_server.general_settings", mock_gs), patch.object( + proxy_config, "_reschedule_spend_log_cleanup_job", mock_reschedule + ): + await proxy_config._update_general_settings( + {"maximum_spend_logs_retention_period": "7d"} + ) + assert mock_gs.get("maximum_spend_logs_retention_period") == "7d" + mock_reschedule.assert_called_once() + + # Test Case 2: Setting the same value should not reschedule + mock_reschedule.reset_mock() + mock_gs = {"maximum_spend_logs_retention_period": "7d"} + with patch("litellm.proxy.proxy_server.general_settings", mock_gs), patch.object( + proxy_config, "_reschedule_spend_log_cleanup_job", mock_reschedule + ): + await proxy_config._update_general_settings( + {"maximum_spend_logs_retention_period": "7d"} + ) + assert mock_gs.get("maximum_spend_logs_retention_period") == "7d" + mock_reschedule.assert_not_called() + + # Test Case 3: Changing value should reschedule + mock_reschedule.reset_mock() + mock_gs = {"maximum_spend_logs_retention_period": "7d"} + with patch("litellm.proxy.proxy_server.general_settings", mock_gs), patch.object( + proxy_config, "_reschedule_spend_log_cleanup_job", mock_reschedule + ): + await proxy_config._update_general_settings( + {"maximum_spend_logs_retention_period": "30d"} + ) + assert mock_gs.get("maximum_spend_logs_retention_period") == "30d" + mock_reschedule.assert_called_once() + + # Test Case 4: Setting to None should reschedule + mock_reschedule.reset_mock() + mock_gs = {"maximum_spend_logs_retention_period": "7d"} + with patch("litellm.proxy.proxy_server.general_settings", mock_gs), patch.object( + proxy_config, "_reschedule_spend_log_cleanup_job", mock_reschedule + ): + await proxy_config._update_general_settings( + {"maximum_spend_logs_retention_period": None} + ) + assert mock_gs.get("maximum_spend_logs_retention_period") is None + mock_reschedule.assert_called_once() + + # Test Case 5: Changing from None to a value should reschedule + mock_reschedule.reset_mock() + mock_gs = {} + with patch("litellm.proxy.proxy_server.general_settings", mock_gs), patch.object( + proxy_config, "_reschedule_spend_log_cleanup_job", mock_reschedule + ): + await proxy_config._update_general_settings( + {"maximum_spend_logs_retention_period": "24h"} + ) + assert mock_gs.get("maximum_spend_logs_retention_period") == "24h" + mock_reschedule.assert_called_once() + + # Test Case 6: Setting None when already None should not reschedule + mock_reschedule.reset_mock() + mock_gs = {} + with patch("litellm.proxy.proxy_server.general_settings", mock_gs), patch.object( + proxy_config, "_reschedule_spend_log_cleanup_job", mock_reschedule + ): + await proxy_config._update_general_settings( + {"maximum_spend_logs_retention_period": None} + ) + assert mock_gs.get("maximum_spend_logs_retention_period") is None + mock_reschedule.assert_not_called()