Merge pull request #9683 from BerriAI/litellm_fix_service_account_behavior

[Bug fix] - Service accounts - only apply `service_account_settings.enforced_params` on service accounts

Author: ishaan-jaff

litellm/proxy/auth/service_account_checks.py

@@ -49,7 +49,6 @@
 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

litellm/proxy/auth/user_api_key_auth.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],

litellm/proxy/litellm_pre_call_utils.py

@@ -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: []

litellm/proxy/proxy_config.yaml

@@ -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

tests/litellm/proxy/test_litellm_pre_call_utils.py

@@ -770,56 +770,31 @@ async def test_add_litellm_data_to_request_duplicate_tags(
 
 
 @pytest.mark.parametrize(
-    "general_settings, user_api_key_dict, expected_enforced_params",
+    "general_settings, user_api_key_dict, request_body, expected_error",
     [
         (
             {"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"],
+            {},
+            True,
         ),
-    ],
-)
-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",
-    [
         (
-            {"enforced_params": ["param1", "param2"]},
+            {"service_account_settings": {"enforced_params": ["user"]}},
             UserAPIKeyAuth(
                 api_key="test_api_key", user_id="test_user_id", org_id="test_org_id"
             ),
             {},
-            True,
+            False,
         ),
         (
             {"service_account_settings": {"enforced_params": ["user"]}},
             UserAPIKeyAuth(
-                api_key="test_api_key", user_id="test_user_id", org_id="test_org_id"
+                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,