Merge pull request #19913 from BerriAI/litellm_store_prompt_backend

[Feature] Allow Dynamic Setting of store_prompts_in_spend_logs

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):

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

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(

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"
+

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