Merge pull request #9376 from BerriAI/litellm_fix_cred_check

Feature flag checking LiteLLM_CredentialsTable

docs/my-website/docs/proxy/config_settings.md

@@ -177,7 +177,7 @@ general_settings:
 | use_x_forwarded_for | str | If true, uses the X-Forwarded-For header to get the client IP address |
 | service_account_settings | List[Dict[str, Any]] | Set `service_account_settings` if you want to create settings that only apply to service account keys (Doc on service accounts)[./service_accounts.md] | 
 | image_generation_model | str | The default model to use for image generation - ignores model set in request |
-| store_model_in_db | boolean | If true, allows `/model/new` endpoint to store model information in db. Endpoint disabled by default. [Doc on `/model/new` endpoint](./model_management.md#create-a-new-model) |
+| store_model_in_db | boolean | If true, enables storing model + credential information in the DB. |
 | store_prompts_in_spend_logs | boolean | If true, allows prompts and responses to be stored in the spend logs table. |
 | max_request_size_mb | int | The maximum size for requests in MB. Requests above this size will be rejected. |
 | max_response_size_mb | int | The maximum size for responses in MB. LLM Responses above this size will not be sent. |
@@ -503,6 +503,7 @@ router_settings:
 | SSL_VERIFY | Flag to enable or disable SSL certificate verification
 | SUPABASE_KEY | API key for Supabase service
 | SUPABASE_URL | Base URL for Supabase instance
+| STORE_MODEL_IN_DB | If true, enables storing model + credential information in the DB. 
 | TEST_EMAIL_ADDRESS | Email address used for testing purposes
 | UI_LOGO_PATH | Path to the logo image used in the UI
 | UI_PASSWORD | Password for accessing the UI
@@ -514,4 +515,3 @@ router_settings:
 | UPSTREAM_LANGFUSE_SECRET_KEY | Secret key for upstream Langfuse authentication
 | USE_AWS_KMS | Flag to enable AWS Key Management Service for encryption
 | WEBHOOK_URL | URL for receiving webhooks from external services
-

litellm/proxy/proxy_server.py

@@ -3282,14 +3282,14 @@ class ProxyStartupEvent:
                 prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj
             )
 
-        ### GET STORED CREDENTIALS ###
-        scheduler.add_job(
-            proxy_config.get_credentials,
-            "interval",
-            seconds=10,
-            args=[prisma_client],
-        )
-        await proxy_config.get_credentials(prisma_client=prisma_client)
+            ### GET STORED CREDENTIALS ###
+            scheduler.add_job(
+                proxy_config.get_credentials,
+                "interval",
+                seconds=10,
+                args=[prisma_client],
+            )
+            await proxy_config.get_credentials(prisma_client=prisma_client)
         if (
             proxy_logging_obj is not None
             and proxy_logging_obj.slack_alerting_instance.alerting is not None

tests/litellm/proxy/test_proxy_server.py

@@ -0,0 +1,76 @@
+import importlib
+import json
+import os
+import socket
+import subprocess
+import sys
+from unittest.mock import AsyncMock, MagicMock, mock_open, patch
+
+import click
+import httpx
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+sys.path.insert(
+    0, os.path.abspath("../../..")
+)  # Adds the parent directory to the system-path
+
+import litellm
+
+
+@pytest.mark.asyncio
+async def test_initialize_scheduled_jobs_credentials(monkeypatch):
+    """
+    Test that get_credentials is only called when store_model_in_db is True
+    """
+    monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False)
+    monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False)
+    from litellm.proxy.proxy_server import ProxyStartupEvent
+    from litellm.proxy.utils import ProxyLogging
+
+    # Mock dependencies
+    mock_prisma_client = MagicMock()
+    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
+    ):  # set store_model_in_db to False
+
+        # Test when store_model_in_db is 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,
+        )
+
+        # Verify get_credentials was not called
+        mock_proxy_config.get_credentials.assert_not_called()
+
+    # Now test with store_model_in_db = True
+    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,
+        )
+
+        # Verify get_credentials was called both directly and scheduled
+        assert mock_proxy_config.get_credentials.call_count == 1  # Direct call
+
+        # Verify a scheduled job was added for get_credentials
+        mock_scheduler_calls = [
+            call[0] for call in mock_proxy_config.get_credentials.mock_calls
+        ]
+        assert len(mock_scheduler_calls) > 0