diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index cbd0706970..71ea045fb1 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/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 - diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f10b4ef86b..ae1c8d18af 100644 --- a/litellm/proxy/proxy_server.py +++ b/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 diff --git a/tests/litellm/proxy/test_proxy_server.py b/tests/litellm/proxy/test_proxy_server.py new file mode 100644 index 0000000000..77b521e836 --- /dev/null +++ b/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