From c13f6514bf11d59f8f33971b2dabda84b3f65313 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 19 Mar 2025 11:35:14 -0700 Subject: [PATCH 1/5] fix(proxy_server.py): move db credential check to only run behind feature flag --- litellm/proxy/proxy_server.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) 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 From 858da57b3cadddebc5d43ee2dcd2f567553af158 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 19 Mar 2025 11:44:00 -0700 Subject: [PATCH 2/5] test(test_proxy_server.py): add unit test to ensure get credentials only called behind feature flag --- tests/litellm/proxy/test_proxy_server.py | 74 ++++++++++++++++++++++++ 1 file changed, 74 insertions(+) create mode 100644 tests/litellm/proxy/test_proxy_server.py diff --git a/tests/litellm/proxy/test_proxy_server.py b/tests/litellm/proxy/test_proxy_server.py new file mode 100644 index 0000000000..3f0a57e284 --- /dev/null +++ b/tests/litellm/proxy/test_proxy_server.py @@ -0,0 +1,74 @@ +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(): + """ + Test that get_credentials is only called when store_model_in_db is True + """ + 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 From 79db96dd13dd5f7714d0c0ee3ef42c0fbe43da79 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 19 Mar 2025 11:47:14 -0700 Subject: [PATCH 3/5] docs(config_settings.md): update env var information --- docs/my-website/docs/proxy/config_settings.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index cbd0706970..0cad66c6e3 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 From 041d5391eb96b5173a97567cba954288367335b2 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 19 Mar 2025 12:01:37 -0700 Subject: [PATCH 4/5] test(test_proxy_server.py): make test work on ci/cd --- tests/litellm/proxy/test_proxy_server.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/litellm/proxy/test_proxy_server.py b/tests/litellm/proxy/test_proxy_server.py index 3f0a57e284..77b521e836 100644 --- a/tests/litellm/proxy/test_proxy_server.py +++ b/tests/litellm/proxy/test_proxy_server.py @@ -20,10 +20,12 @@ import litellm @pytest.mark.asyncio -async def test_initialize_scheduled_jobs_credentials(): +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 From 8464f06df79eac00d4eabd5d2bdd9864f4848d87 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 19 Mar 2025 12:06:07 -0700 Subject: [PATCH 5/5] docs(config_settings.md): cleanup --- docs/my-website/docs/proxy/config_settings.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 0cad66c6e3..71ea045fb1 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -503,7 +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. | +| 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 @@ -515,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 -