diff --git a/litellm/proxy/auth/service_account_checks.py b/litellm/proxy/auth/service_account_checks.py deleted file mode 100644 index 87d7d66854e..00000000000 --- a/litellm/proxy/auth/service_account_checks.py +++ /dev/null @@ -1,53 +0,0 @@ -""" -Checks for LiteLLM service account keys - -""" - -from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth - - -def check_if_token_is_service_account(valid_token: UserAPIKeyAuth) -> bool: - """ - Checks if the token is a service account - - Returns: - bool: True if token is a service account - - """ - if valid_token.metadata: - if "service_account_id" in valid_token.metadata: - return True - return False - - -async def service_account_checks( - valid_token: UserAPIKeyAuth, request_data: dict -) -> bool: - """ - If a virtual key is a service account, checks it's a valid service account - - A token is a service account if it has a service_account_id in its metadata - - Service Account Specific Checks: - - Check if required_params is set - """ - - if check_if_token_is_service_account(valid_token) is not True: - return True - - from litellm.proxy.proxy_server import general_settings - - if "service_account_settings" in general_settings: - service_account_settings = general_settings["service_account_settings"] - if "enforced_params" in service_account_settings: - _enforced_params = service_account_settings["enforced_params"] - for param in _enforced_params: - if param not in request_data: - raise ProxyException( - type=ProxyErrorTypes.bad_request_error.value, - code=400, - param=param, - message=f"BadRequest please pass param={param} in request body. This is a required param for service account", - ) - - return True diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index eddbf4e0d97..1b140bdaaba 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -49,7 +49,6 @@ from litellm.proxy.auth.auth_utils import ( from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler from litellm.proxy.auth.oauth2_check import check_oauth2_token from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request -from litellm.proxy.auth.service_account_checks import service_account_checks from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.types.services import ServiceTypes @@ -905,12 +904,6 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 else: _team_obj = None - # Check 7: Check if key is a service account key - await service_account_checks( - valid_token=valid_token, - request_data=request_data, - ) - user_api_key_cache.set_cache( key=valid_token.team_id, value=_team_obj ) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 6427be5a6e2..dade6c933e6 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -747,7 +747,10 @@ def _get_enforced_params( enforced_params: Optional[list] = None if general_settings is not None: enforced_params = general_settings.get("enforced_params") - if "service_account_settings" in general_settings: + if ( + "service_account_settings" in general_settings + and check_if_token_is_service_account(user_api_key_dict) is True + ): service_account_settings = general_settings["service_account_settings"] if "enforced_params" in service_account_settings: if enforced_params is None: @@ -760,6 +763,20 @@ def _get_enforced_params( return enforced_params +def check_if_token_is_service_account(valid_token: UserAPIKeyAuth) -> bool: + """ + Checks if the token is a service account + + Returns: + bool: True if token is a service account + + """ + if valid_token.metadata: + if "service_account_id" in valid_token.metadata: + return True + return False + + def _enforced_params_check( request_body: dict, general_settings: Optional[dict], diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index fe8d73d26aa..17658df9031 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -4,12 +4,3 @@ model_list: model: openai/fake api_key: fake-key api_base: https://exampleopenaiendpoint-production.up.railway.app/ - -general_settings: - use_redis_transaction_buffer: true - -litellm_settings: - cache: True - cache_params: - type: redis - supported_call_types: [] \ No newline at end of file diff --git a/tests/litellm/proxy/test_litellm_pre_call_utils.py b/tests/litellm/proxy/test_litellm_pre_call_utils.py new file mode 100644 index 00000000000..94f2c512ea9 --- /dev/null +++ b/tests/litellm/proxy/test_litellm_pre_call_utils.py @@ -0,0 +1,105 @@ +import json +import os +import sys +from unittest.mock import MagicMock, patch + +import pytest + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.litellm_pre_call_utils import ( + _get_enforced_params, + check_if_token_is_service_account, +) + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + + +def test_check_if_token_is_service_account(): + """ + Test that only keys with `service_account_id` in metadata are considered service accounts + """ + # Test case 1: Service account token + service_account_token = UserAPIKeyAuth( + api_key="test-key", metadata={"service_account_id": "test-service-account"} + ) + assert check_if_token_is_service_account(service_account_token) == True + + # Test case 2: Regular user token + regular_token = UserAPIKeyAuth(api_key="test-key", metadata={}) + assert check_if_token_is_service_account(regular_token) == False + + # Test case 3: Token with other metadata + other_metadata_token = UserAPIKeyAuth( + api_key="test-key", metadata={"user_id": "test-user"} + ) + assert check_if_token_is_service_account(other_metadata_token) == False + + +def test_get_enforced_params_for_service_account_settings(): + """ + Test that service account enforced params are only added to service account keys + """ + service_account_token = UserAPIKeyAuth( + api_key="test-key", metadata={"service_account_id": "test-service-account"} + ) + general_settings_with_service_account_settings = { + "service_account_settings": {"enforced_params": ["metadata.service"]}, + } + result = _get_enforced_params( + general_settings=general_settings_with_service_account_settings, + user_api_key_dict=service_account_token, + ) + assert result == ["metadata.service"] + + regular_token = UserAPIKeyAuth( + api_key="test-key", metadata={"enforced_params": ["user"]} + ) + result = _get_enforced_params( + general_settings=general_settings_with_service_account_settings, + user_api_key_dict=regular_token, + ) + assert result == ["user"] + + +@pytest.mark.parametrize( + "general_settings, user_api_key_dict, expected_enforced_params", + [ + ( + {"enforced_params": ["param1", "param2"]}, + UserAPIKeyAuth( + api_key="test_api_key", user_id="test_user_id", org_id="test_org_id" + ), + ["param1", "param2"], + ), + ( + {"service_account_settings": {"enforced_params": ["param1", "param2"]}}, + UserAPIKeyAuth( + api_key="test_api_key", + user_id="test_user_id", + org_id="test_org_id", + metadata={"service_account_id": "test_service_account_id"}, + ), + ["param1", "param2"], + ), + ( + {"service_account_settings": {"enforced_params": ["param1", "param2"]}}, + UserAPIKeyAuth( + api_key="test_api_key", + metadata={ + "enforced_params": ["param3", "param4"], + "service_account_id": "test_service_account_id", + }, + ), + ["param1", "param2", "param3", "param4"], + ), + ], +) +def test_get_enforced_params( + general_settings, user_api_key_dict, expected_enforced_params +): + from litellm.proxy.litellm_pre_call_utils import _get_enforced_params + + enforced_params = _get_enforced_params(general_settings, user_api_key_dict) + assert enforced_params == expected_enforced_params diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index a8c0ebe2ffa..1281d508638 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -769,42 +769,6 @@ async def test_add_litellm_data_to_request_duplicate_tags( ), f"Expected {expected_tags}, got {result['metadata']['tags']}" -@pytest.mark.parametrize( - "general_settings, user_api_key_dict, expected_enforced_params", - [ - ( - {"enforced_params": ["param1", "param2"]}, - UserAPIKeyAuth( - api_key="test_api_key", user_id="test_user_id", org_id="test_org_id" - ), - ["param1", "param2"], - ), - ( - {"service_account_settings": {"enforced_params": ["param1", "param2"]}}, - UserAPIKeyAuth( - api_key="test_api_key", user_id="test_user_id", org_id="test_org_id" - ), - ["param1", "param2"], - ), - ( - {"service_account_settings": {"enforced_params": ["param1", "param2"]}}, - UserAPIKeyAuth( - api_key="test_api_key", - metadata={"enforced_params": ["param3", "param4"]}, - ), - ["param1", "param2", "param3", "param4"], - ), - ], -) -def test_get_enforced_params( - general_settings, user_api_key_dict, expected_enforced_params -): - from litellm.proxy.litellm_pre_call_utils import _get_enforced_params - - enforced_params = _get_enforced_params(general_settings, user_api_key_dict) - assert enforced_params == expected_enforced_params - - @pytest.mark.parametrize( "general_settings, user_api_key_dict, request_body, expected_error", [ @@ -822,6 +786,17 @@ def test_get_enforced_params( api_key="test_api_key", user_id="test_user_id", org_id="test_org_id" ), {}, + False, + ), + ( + {"service_account_settings": {"enforced_params": ["user"]}}, + UserAPIKeyAuth( + api_key="test_api_key", + user_id="test_user_id", + org_id="test_org_id", + metadata={"service_account_id": "test_service_account_id"} + ), + {}, True, ), ( @@ -854,6 +829,7 @@ def test_get_enforced_params( {"service_account_settings": {"enforced_params": ["user"]}}, UserAPIKeyAuth( api_key="test_api_key", + metadata={"service_account_id": "test_service_account_id"} ), {"user": "test_user"}, False,