scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 137

Merge pull request #9683 from BerriAI/litellm_fix_service_account_behavior

Browse source 
[Bug fix] - Service accounts - only apply `service_account_settings.enforced_params` on service accounts
This commit is contained in:
Ishaan Jaff 2025-04-01 12:04:38 -07:00 • committed by GitHub
commit 655ce2e745
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 135 additions and 106 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

53
litellm/proxy/auth/service_account_checks.py
View file

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

7
litellm/proxy/auth/user_api_key_auth.py
View file

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

19
litellm/proxy/litellm_pre_call_utils.py
View file

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

9
litellm/proxy/proxy_config.yaml
View file

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

105
tests/litellm/proxy/test_litellm_pre_call_utils.py Normal file
View file

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

48
tests/proxy_unit_tests/test_proxy_utils.py
View file

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