diff --git a/docs/my-website/docs/proxy/prod.md b/docs/my-website/docs/proxy/prod.md index 3a24a3427bd..fb2acf230c1 100644 --- a/docs/my-website/docs/proxy/prod.md +++ b/docs/my-website/docs/proxy/prod.md @@ -199,7 +199,7 @@ USE_PRISMA_MIGRATE="True" ```bash -litellm --use_prisma_migrate +litellm ``` diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 21c9131887b..c194c49765c 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -243,7 +243,6 @@ class ProxyExtrasDBManager: bool: True if setup was successful, False otherwise """ schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma" - use_migrate = str_to_bool(os.getenv("USE_PRISMA_MIGRATE")) or use_migrate for attempt in range(4): original_dir = os.getcwd() migrations_dir = ProxyExtrasDBManager._get_prisma_dir() diff --git a/litellm/llms/litellm_proxy/chat/transformation.py b/litellm/llms/litellm_proxy/chat/transformation.py index ea89c4c3bc7..cf6a6ed7a54 100644 --- a/litellm/llms/litellm_proxy/chat/transformation.py +++ b/litellm/llms/litellm_proxy/chat/transformation.py @@ -4,6 +4,7 @@ Translate from OpenAI's `/v1/chat/completions` to VLLM's `/v1/chat/completions` from typing import TYPE_CHECKING, List, Optional, Tuple +from litellm.constants import OPENAI_CHAT_COMPLETION_PARAMS from litellm.secret_managers.main import get_secret_bool, get_secret_str from litellm.types.router import LiteLLM_Params @@ -16,8 +17,7 @@ if TYPE_CHECKING: class LiteLLMProxyChatConfig(OpenAIGPTConfig): def get_supported_openai_params(self, model: str) -> List: params_list = super().get_supported_openai_params(model) - params_list.append("thinking") - params_list.append("reasoning_effort") + params_list.extend(OPENAI_CHAT_COMPLETION_PARAMS) return params_list def _map_openai_params( diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index b48fe5be1c3..99460a0547a 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -6,6 +6,7 @@ model_list: api_base: https://exampleopenaiendpoint-production.up.railway.app/ litellm_settings: + callbacks: ["otel"] cache: true cache_params: type: redis diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 02a75c0b0f7..406ddceabf5 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -87,12 +87,11 @@ class PrismaWrapper: self, new_db_url: str, http_client: Optional[Any] = None ): from prisma import Prisma # type: ignore + try: await self._original_prisma.disconnect() except Exception as e: - verbose_proxy_logger.warning( - f"Failed to disconnect Prisma client: {e}" - ) + verbose_proxy_logger.warning(f"Failed to disconnect Prisma client: {e}") if http_client is not None: self._original_prisma = Prisma(http=http_client) @@ -139,7 +138,6 @@ class PrismaManager: bool: True if setup was successful, False otherwise """ - use_migrate = str_to_bool(os.getenv("USE_PRISMA_MIGRATE")) or use_migrate for attempt in range(4): original_dir = os.getcwd() prisma_dir = PrismaManager._get_prisma_dir() @@ -185,7 +183,7 @@ class PrismaManager: def should_update_prisma_schema( - disable_updates: Optional[Union[bool, str]] = None + disable_updates: Optional[Union[bool, str]] = None, ) -> bool: """ Determines if Prisma Schema updates should be applied during startup. diff --git a/litellm/proxy/prisma_migration.py b/litellm/proxy/prisma_migration.py index 35142736ab3..251d1e56287 100644 --- a/litellm/proxy/prisma_migration.py +++ b/litellm/proxy/prisma_migration.py @@ -13,7 +13,7 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy.proxy_cli import run_server # Call the Click command with standalone_mode=False -run_server(["--use_prisma_migrate", "--skip_server_startup"], standalone_mode=False) +run_server(["--skip_server_startup"], standalone_mode=False) # run prisma generate verbose_proxy_logger.info("Running 'prisma generate'...") diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 80b2f69e67e..cc4b1652d08 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -465,10 +465,10 @@ class ProxyInitializationHelpers: help="Ciphers to use for the SSL setup.", ) @click.option( - "--use_prisma_migrate", + "--use_prisma_db_push", is_flag=True, - default=True, - help="Use prisma migrate instead of prisma db push for database schema updates", + default=False, + help="Use prisma db push instead of prisma migrate for database schema updates", ) @click.option("--local", is_flag=True, default=False, help="for local debugging") @click.option( @@ -519,7 +519,7 @@ def run_server( # noqa: PLR0915 ssl_certfile_path, ciphers, log_config, - use_prisma_migrate, + use_prisma_db_push: bool, skip_server_startup, keepalive_timeout, ): @@ -777,7 +777,7 @@ def run_server( # noqa: PLR0915 ): check_prisma_schema_diff(db_url=None) else: - PrismaManager.setup_database(use_migrate=use_prisma_migrate) + PrismaManager.setup_database(use_migrate=not use_prisma_db_push) else: print( # noqa f"Unable to connect to DB. DATABASE_URL found in environment, but prisma package not found." # noqa diff --git a/tests/llm_translation/test_litellm_proxy_provider.py b/tests/llm_translation/test_litellm_proxy_provider.py index 13cc4bd2637..1ae507a10d5 100644 --- a/tests/llm_translation/test_litellm_proxy_provider.py +++ b/tests/llm_translation/test_litellm_proxy_provider.py @@ -452,7 +452,7 @@ def test_litellm_gateway_from_sdk_with_response_cost_in_additional_headers(): def test_litellm_gateway_from_sdk_with_thinking_param(): - try: + try: response = litellm.completion( model="litellm_proxy/anthropic.claude-3-7-sonnet-20250219-v1:0", messages=[{"role": "user", "content": "Hello world"}], @@ -464,4 +464,3 @@ def test_litellm_gateway_from_sdk_with_thinking_param(): pytest.fail("Expected an error to be raised") except Exception as e: assert "Connection error." in str(e) - diff --git a/tests/test_litellm/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py b/tests/test_litellm/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py index 16d33853a60..3c2f22dca9e 100644 --- a/tests/test_litellm/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py +++ b/tests/test_litellm/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py @@ -30,3 +30,13 @@ def test_litellm_proxy_chat_transformation(): litellm_params={}, headers={}, ) == {"model": "model", "messages": messages} + + +def test_litellm_gateway_from_sdk_with_user_param(): + from litellm.llms.litellm_proxy.chat.transformation import LiteLLMProxyChatConfig + + supported_params = LiteLLMProxyChatConfig().get_supported_openai_params( + "openai/gpt-4o" + ) + print(f"supported_params: {supported_params}") + assert "user" in supported_params diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 48abbc6e003..4235e5d3adb 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -2,18 +2,19 @@ import os import sys from unittest.mock import MagicMock, patch -import pytest import fastapi +import pytest sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system-path -from litellm.proxy.proxy_cli import ProxyInitializationHelpers -from litellm.proxy.health_endpoints.health_app_factory import build_health_app import builtins import types +from litellm.proxy.health_endpoints.health_app_factory import build_health_app +from litellm.proxy.proxy_cli import ProxyInitializationHelpers + class TestProxyInitializationHelpers: @patch("importlib.metadata.version") @@ -308,7 +309,7 @@ class TestProxyInitializationHelpers: keepalive_timeout=30, ) mock_uvicorn_run.assert_called_once() - + # Check that the uvicorn.run was called with the timeout_keep_alive parameter call_args = mock_uvicorn_run.call_args assert call_args[1]["timeout_keep_alive"] == 30 @@ -376,10 +377,12 @@ class TestProxyInitializationHelpers: @patch("builtins.print") def test_run_server_no_config_passed(self, mock_print, mock_uvicorn_run): """Test that run_server properly handles the case when no config is passed""" - from click.testing import CliRunner - from litellm.proxy.proxy_cli import run_server import asyncio + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + runner = CliRunner() mock_app = MagicMock() @@ -389,11 +392,8 @@ class TestProxyInitializationHelpers: # Mock the ProxyConfig.get_config method to return a proper async config async def mock_get_config(config_file_path=None): - return { - "general_settings": {}, - "litellm_settings": {} - } - + return {"general_settings": {}, "litellm_settings": {}} + mock_proxy_config_instance = MagicMock() mock_proxy_config_instance.get_config = mock_get_config mock_proxy_config.return_value = mock_proxy_config_instance @@ -423,7 +423,7 @@ class TestProxyInitializationHelpers: result = runner.invoke(run_server, ["--local"]) assert result.exit_code == 0 - + # Verify that uvicorn.run was called mock_uvicorn_run.assert_called_once() @@ -434,34 +434,34 @@ class TestProxyInitializationHelpers: result = runner.invoke(run_server, ["--local", "--config", "None"]) assert result.exit_code == 0 - + # Verify that uvicorn.run was called again mock_uvicorn_run.assert_called_once() class TestHealthAppFactory: """Test cases for the health app factory module""" - + def test_build_health_app(self): """Test that build_health_app creates a FastAPI app with the correct title and includes the health router""" # Execute health_app = build_health_app() - + # Assert assert health_app.title == "LiteLLM Health Endpoints" assert isinstance(health_app, fastapi.FastAPI) - + # Verify that the app has the expected health endpoints by checking route paths # When a router is included, its routes are flattened into the main app's routes route_paths = [] for route in health_app.routes: - if hasattr(route, 'path'): + if hasattr(route, "path"): route_paths.append(route.path) - + # Check for some expected health endpoints expected_paths = [ "/test", - "/health/services", + "/health/services", "/health", "/health/history", "/health/latest", @@ -470,24 +470,100 @@ class TestHealthAppFactory: "/health/readiness", "/health/liveliness", "/health/liveness", - "/health/test_connection" + "/health/test_connection", ] - + # At least some of the expected health endpoints should be present found_paths = [path for path in expected_paths if path in route_paths] - assert len(found_paths) > 0, f"Expected to find health endpoints, but found: {route_paths}" - + assert ( + len(found_paths) > 0 + ), f"Expected to find health endpoints, but found: {route_paths}" + # Verify that the app has routes (indicating the router was included) - assert len(health_app.routes) > 0, "Health app should have routes from the included router" - + assert ( + len(health_app.routes) > 0 + ), "Health app should have routes from the included router" + def test_build_health_app_returns_different_instances(self): """Test that build_health_app returns different FastAPI instances on each call""" # Execute health_app_1 = build_health_app() health_app_2 = build_health_app() - + # Assert assert health_app_1 is not health_app_2 assert health_app_1.title == health_app_2.title assert isinstance(health_app_1, fastapi.FastAPI) assert isinstance(health_app_2, fastapi.FastAPI) + + @patch("subprocess.run") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") + @patch.dict( + os.environ, {"DATABASE_URL": "postgresql://test:test@localhost:5432/test"} + ) + def test_use_prisma_db_push_flag_behavior( + self, + mock_should_update_schema, + mock_check_schema_diff, + mock_setup_database, + mock_subprocess_run, + ): + """Test that use_prisma_db_push flag correctly controls PrismaManager.setup_database use_migrate parameter""" + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + runner = CliRunner() + + # Mock subprocess.run to simulate prisma being available + mock_subprocess_run.return_value = MagicMock(returncode=0) + + # Mock should_update_prisma_schema to return True (so setup_database gets called) + mock_should_update_schema.return_value = True + + mock_app = MagicMock() + mock_proxy_config = MagicMock() + mock_key_mgmt = MagicMock() + mock_save_worker_config = MagicMock() + + with patch.dict( + "sys.modules", + { + "proxy_server": MagicMock( + app=mock_app, + ProxyConfig=mock_proxy_config, + KeyManagementSettings=mock_key_mgmt, + save_worker_config=mock_save_worker_config, + ) + }, + ), patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" + ) as mock_get_args: + mock_get_args.return_value = { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + + # Test 1: Without --use_prisma_db_push flag (default behavior) + # use_prisma_db_push should be False (default), so use_migrate should be True + result = runner.invoke(run_server, ["--local", "--skip_server_startup"]) + + assert result.exit_code == 0 + mock_setup_database.assert_called_with(use_migrate=True) + + # Reset mocks + mock_setup_database.reset_mock() + mock_should_update_schema.reset_mock() + mock_should_update_schema.return_value = True + + # Test 2: With --use_prisma_db_push flag set + # use_prisma_db_push should be True, so use_migrate should be False + result = runner.invoke( + run_server, ["--local", "--skip_server_startup", "--use_prisma_db_push"] + ) + + assert result.exit_code == 0 + mock_setup_database.assert_called_with(use_migrate=False)