From 6356db560d0a234fdd178eba2c19020b6b4012f9 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 18 Feb 2026 17:25:39 -0800 Subject: [PATCH] [Feature] Allow store_model_in_db to be set via database Users had to set store_model_in_db in the config YAML and restart the proxy, causing service downtime. This change allows the value to be written to the LiteLLM_Config table and read from the database at runtime, with DB values overriding config file values. Co-Authored-By: Claude Opus 4.6 (1M context) --- litellm/proxy/_types.py | 4 + litellm/proxy/proxy_server.py | 42 ++- tests/test_litellm/proxy/test_proxy_server.py | 294 ++++++++++++++++++ 3 files changed, 339 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c327dd130a..71b376ea28 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2125,6 +2125,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="CIDR ranges of trusted reverse proxies. When set, X-Forwarded-For headers are only trusted from these IPs.", ) + store_model_in_db: Optional[bool] = Field( + None, + description="If True, models and config are stored in and loaded from the database. Default is False.", + ) class ConfigYAML(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 318c64e4a0..56d1d139c2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2804,6 +2804,7 @@ class ProxyConfig: store_model_in_db = general_settings.get("store_model_in_db", False) if store_model_in_db is None: store_model_in_db = False + general_settings["store_model_in_db"] = store_model_in_db ### CUSTOM API KEY AUTH ### ## pass filepath custom_auth = general_settings.get("custom_auth", None) @@ -3841,7 +3842,7 @@ class ProxyConfig: """ Pull from DB, read general settings value """ - global general_settings + global general_settings, store_model_in_db if db_general_settings is None: return _general_settings = dict(db_general_settings) @@ -3893,6 +3894,19 @@ class ProxyConfig: # For other types, convert to bool general_settings["store_prompts_in_spend_logs"] = bool(value) + ## STORE MODEL IN DB ## + if "store_model_in_db" in _general_settings: + value = _general_settings["store_model_in_db"] + if value is None: + pass # Don't change store_model_in_db to None; keep current value + elif isinstance(value, bool): + store_model_in_db = value + elif isinstance(value, str): + store_model_in_db = value.lower() == "true" + else: + store_model_in_db = bool(value) + general_settings["store_model_in_db"] = store_model_in_db + ## MAXIMUM SPEND LOGS RETENTION PERIOD ## if "maximum_spend_logs_retention_period" in _general_settings: old_value = general_settings.get("maximum_spend_logs_retention_period") @@ -5429,6 +5443,31 @@ class ProxyStartupEvent: get_secret_bool("STORE_MODEL_IN_DB", store_model_in_db) or store_model_in_db ) + # If store_model_in_db is still False, check DB for override. + # This breaks the chicken-and-egg where DB has store_model_in_db=True + # but YAML config has False. + if store_model_in_db is not True and prisma_client is not None: + try: + _db_gs_record = await prisma_client.db.litellm_config.find_first( + where={"param_name": "general_settings"} + ) + if _db_gs_record is not None and isinstance( + _db_gs_record.param_value, dict + ): + _db_val = _db_gs_record.param_value.get("store_model_in_db") + if _db_val is True or ( + isinstance(_db_val, str) + and _db_val.lower() == "true" + ): + store_model_in_db = True + verbose_proxy_logger.info( + "store_model_in_db=True loaded from DB, overriding config/env" + ) + except Exception as e: + verbose_proxy_logger.debug( + "Failed to check DB for store_model_in_db: %s", str(e) + ) + if store_model_in_db is True: # MEMORY LEAK FIX: Increase interval from 10s to 30s minimum # Frequent polling was causing excessive memory allocations @@ -11317,6 +11356,7 @@ async def get_config_list( "max_request_size_mb": {"type": "Integer"}, "max_response_size_mb": {"type": "Integer"}, "pass_through_endpoints": {"type": "PydanticModel"}, + "store_model_in_db": {"type": "Boolean"}, "store_prompts_in_spend_logs": {"type": "Boolean"}, "maximum_spend_logs_retention_period": {"type": "String"}, "mcp_internal_ip_ranges": {"type": "List"}, diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index aefd19ef3c..14a2b77841 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -94,6 +94,7 @@ def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "") + monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None) client = TestClient(app) response = client.post( @@ -3337,3 +3338,296 @@ class TestInvitationEndpoints: # ProxyException handler returns {"error": {...}}, HTTPException returns {"detail": {...}} error_content = body.get("error", body.get("detail", body)) assert "not allowed" in str(error_content).lower() + + +# ============================================================================ +# store_model_in_db DB Config Override Tests +# ============================================================================ + + +def test_store_model_in_db_in_config_general_settings(): + """ + Verify store_model_in_db is a valid field in ConfigGeneralSettings + and validates correctly for True/False values. + """ + from litellm.proxy._types import ConfigGeneralSettings + + assert "store_model_in_db" in ConfigGeneralSettings.model_fields + + # Should validate with True + config = ConfigGeneralSettings(store_model_in_db=True) + assert config.store_model_in_db is True + + # Should validate with False + config = ConfigGeneralSettings(store_model_in_db=False) + assert config.store_model_in_db is False + + # Should validate with None (default) + config = ConfigGeneralSettings(store_model_in_db=None) + assert config.store_model_in_db is None + + # Should validate with no value + config = ConfigGeneralSettings() + assert config.store_model_in_db is None + + +@pytest.mark.asyncio +async def test_update_general_settings_store_model_in_db_true(): + """ + Verify _update_general_settings sets global store_model_in_db to True + when DB general_settings has store_model_in_db=True. + """ + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + with patch( + "litellm.proxy.proxy_server.store_model_in_db", False + ) as mock_store, patch( + "litellm.proxy.proxy_server.general_settings", {} + ) as mock_gs: + await proxy_config._update_general_settings( + db_general_settings={"store_model_in_db": True} + ) + + import litellm.proxy.proxy_server as ps + + assert ps.store_model_in_db is True + assert ps.general_settings["store_model_in_db"] is True + + +@pytest.mark.asyncio +async def test_update_general_settings_store_model_in_db_false(): + """ + Verify _update_general_settings sets global store_model_in_db to False + when DB general_settings has store_model_in_db=False. + """ + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + with patch( + "litellm.proxy.proxy_server.store_model_in_db", True + ), patch("litellm.proxy.proxy_server.general_settings", {}): + await proxy_config._update_general_settings( + db_general_settings={"store_model_in_db": False} + ) + + import litellm.proxy.proxy_server as ps + + assert ps.store_model_in_db is False + assert ps.general_settings["store_model_in_db"] is False + + +@pytest.mark.asyncio +async def test_update_general_settings_store_model_in_db_string_normalization(): + """ + Verify _update_general_settings normalizes string values for store_model_in_db. + """ + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + # Test "true" string + with patch( + "litellm.proxy.proxy_server.store_model_in_db", False + ), patch("litellm.proxy.proxy_server.general_settings", {}): + await proxy_config._update_general_settings( + db_general_settings={"store_model_in_db": "true"} + ) + import litellm.proxy.proxy_server as ps + + assert ps.store_model_in_db is True + + # Test "True" string + with patch( + "litellm.proxy.proxy_server.store_model_in_db", False + ), patch("litellm.proxy.proxy_server.general_settings", {}): + await proxy_config._update_general_settings( + db_general_settings={"store_model_in_db": "True"} + ) + import litellm.proxy.proxy_server as ps + + assert ps.store_model_in_db is True + + # Test "false" string + with patch( + "litellm.proxy.proxy_server.store_model_in_db", True + ), patch("litellm.proxy.proxy_server.general_settings", {}): + await proxy_config._update_general_settings( + db_general_settings={"store_model_in_db": "false"} + ) + import litellm.proxy.proxy_server as ps + + assert ps.store_model_in_db is False + + +@pytest.mark.asyncio +async def test_update_general_settings_store_model_in_db_none_keeps_current(): + """ + Verify _update_general_settings does not change store_model_in_db + when DB value is None. + """ + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + # When current is True and DB sends None, should stay True + with patch( + "litellm.proxy.proxy_server.store_model_in_db", True + ), patch("litellm.proxy.proxy_server.general_settings", {}): + await proxy_config._update_general_settings( + db_general_settings={"store_model_in_db": None} + ) + import litellm.proxy.proxy_server as ps + + assert ps.store_model_in_db is True + + # When current is False and DB sends None, should stay False + with patch( + "litellm.proxy.proxy_server.store_model_in_db", False + ), patch("litellm.proxy.proxy_server.general_settings", {}): + await proxy_config._update_general_settings( + db_general_settings={"store_model_in_db": None} + ) + import litellm.proxy.proxy_server as ps + + assert ps.store_model_in_db is False + + +@pytest.mark.asyncio +async def test_store_model_in_db_db_override_when_config_false(): + """ + Verify the early DB check in initialize_scheduled_background_jobs + overrides store_model_in_db=False when DB has True. + """ + from litellm.proxy.proxy_server import ProxyStartupEvent + from litellm.proxy.utils import ProxyLogging + + mock_prisma_client = MagicMock() + + # Mock DB returning store_model_in_db=True in general_settings + mock_db_record = MagicMock() + mock_db_record.param_value = {"store_model_in_db": True} + mock_prisma_client.db.litellm_config.find_first = AsyncMock( + return_value=mock_db_record + ) + + mock_proxy_logging = MagicMock(spec=ProxyLogging) + mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_config = AsyncMock() + + with patch( + "litellm.proxy.proxy_server.proxy_config", mock_proxy_config + ), patch( + "litellm.proxy.proxy_server.store_model_in_db", False + ), patch( + "litellm.proxy.proxy_server.get_secret_bool", return_value=False + ): + await ProxyStartupEvent.initialize_scheduled_background_jobs( + general_settings={}, + prisma_client=mock_prisma_client, + proxy_budget_rescheduler_min_time=1, + proxy_budget_rescheduler_max_time=2, + proxy_batch_write_at=5, + proxy_logging_obj=mock_proxy_logging, + ) + + import litellm.proxy.proxy_server as ps + + # store_model_in_db should now be True (overridden by DB) + assert ps.store_model_in_db is True + + # add_deployment and get_credentials should have been called + # since store_model_in_db is now True + assert mock_proxy_config.add_deployment.call_count == 1 + assert mock_proxy_config.get_credentials.call_count == 1 + + +@pytest.mark.asyncio +async def test_store_model_in_db_db_check_skipped_when_already_true(monkeypatch): + """ + Verify the early DB check is skipped when store_model_in_db is already True. + The DB query for the early check should not be called. + """ + monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) + from litellm.proxy.proxy_server import ProxyStartupEvent + from litellm.proxy.utils import ProxyLogging + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=None) + + mock_proxy_logging = MagicMock(spec=ProxyLogging) + mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_config = AsyncMock() + + with patch( + "litellm.proxy.proxy_server.proxy_config", mock_proxy_config + ), patch( + "litellm.proxy.proxy_server.store_model_in_db", True + ), patch( + "litellm.proxy.proxy_server.get_secret_bool", return_value=True + ): + await ProxyStartupEvent.initialize_scheduled_background_jobs( + general_settings={}, + prisma_client=mock_prisma_client, + proxy_budget_rescheduler_min_time=1, + proxy_budget_rescheduler_max_time=2, + proxy_batch_write_at=5, + proxy_logging_obj=mock_proxy_logging, + ) + + # The early DB check uses find_first with param_name="general_settings". + # When store_model_in_db is already True, the early check should be skipped. + # However, add_deployment may also call find_first. + # We just verify that store_model_in_db stays True and jobs are scheduled. + import litellm.proxy.proxy_server as ps + + assert ps.store_model_in_db is True + assert mock_proxy_config.add_deployment.call_count == 1 + + +@pytest.mark.asyncio +async def test_store_model_in_db_db_failure_graceful(monkeypatch): + """ + Verify the early DB check handles DB failures gracefully + without crashing and keeps store_model_in_db as False. + """ + monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) + from litellm.proxy.proxy_server import ProxyStartupEvent + from litellm.proxy.utils import ProxyLogging + + mock_prisma_client = MagicMock() + # Simulate DB failure + mock_prisma_client.db.litellm_config.find_first = AsyncMock( + side_effect=Exception("DB connection error") + ) + + mock_proxy_logging = MagicMock(spec=ProxyLogging) + mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_config = AsyncMock() + + with patch( + "litellm.proxy.proxy_server.proxy_config", mock_proxy_config + ), patch( + "litellm.proxy.proxy_server.store_model_in_db", False + ), patch( + "litellm.proxy.proxy_server.get_secret_bool", return_value=False + ): + # Should not raise an exception + await ProxyStartupEvent.initialize_scheduled_background_jobs( + general_settings={}, + prisma_client=mock_prisma_client, + proxy_budget_rescheduler_min_time=1, + proxy_budget_rescheduler_max_time=2, + proxy_batch_write_at=5, + proxy_logging_obj=mock_proxy_logging, + ) + + import litellm.proxy.proxy_server as ps + + # store_model_in_db should remain False + assert ps.store_model_in_db is False + + # add_deployment should NOT have been called since store_model_in_db is False + mock_proxy_config.add_deployment.assert_not_called()