From f8ce30763ea7155bf8383dd6632f907a377ba77d Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 1 Apr 2025 07:36:35 -0700 Subject: [PATCH 1/6] fix check_if_token_is_service_account --- ...odel_prices_and_context_window_backup.json | 36 +++++++++++++ litellm/proxy/auth/service_account_checks.py | 53 ------------------- litellm/proxy/litellm_pre_call_utils.py | 19 ++++++- litellm/proxy/proxy_config.yaml | 12 ++--- 4 files changed, 57 insertions(+), 63 deletions(-) delete mode 100644 litellm/proxy/auth/service_account_checks.py diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 42bddca955c..28f8acd21ca 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -4453,6 +4453,42 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", "supports_tool_choice": true }, + "gemini-2.5-pro-exp-03-25": { + "max_tokens": 65536, + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_images_per_prompt": 3000, + "max_videos_per_prompt": 10, + "max_video_length": 1, + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_pdf_size_mb": 30, + "input_cost_per_image": 0, + "input_cost_per_video_per_second": 0, + "input_cost_per_audio_per_second": 0, + "input_cost_per_token": 0, + "input_cost_per_character": 0, + "input_cost_per_token_above_128k_tokens": 0, + "input_cost_per_character_above_128k_tokens": 0, + "input_cost_per_image_above_128k_tokens": 0, + "input_cost_per_video_per_second_above_128k_tokens": 0, + "input_cost_per_audio_per_second_above_128k_tokens": 0, + "output_cost_per_token": 0, + "output_cost_per_character": 0, + "output_cost_per_token_above_128k_tokens": 0, + "output_cost_per_character_above_128k_tokens": 0, + "litellm_provider": "vertex_ai-language-models", + "mode": "chat", + "supports_system_messages": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true, + "supports_pdf_input": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, "gemini-2.0-pro-exp-02-05": { "max_tokens": 8192, "max_input_tokens": 2097152, 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/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 c0107d7ea6f..9cdd1e72360 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -4,12 +4,6 @@ model_list: model: openai/gpt-4o api_key: sk-xxxxxxx -mcp_servers: - { - "zapier_mcp": { - "url": "https://actions.zapier.com/mcp/sk-akxxxxx/sse" - }, - "fetch": { - "url": "http://localhost:8000/sse" - } - } +general_settings: + service_account_settings: + enforced_params: ["user"] # this means the "user" param is enforced for all requests made through any service account keys \ No newline at end of file From 13aa7f75f6dff49a036a71c5db012a68d8338752 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 1 Apr 2025 07:40:31 -0700 Subject: [PATCH 2/6] test_enforced_params_check --- tests/proxy_unit_tests/test_proxy_utils.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index d613118fc89..9fb00c094fe 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -822,6 +822,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 +865,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, From e5f6529c426148071f295bfbe5ca9e81ff2f4e7b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 1 Apr 2025 07:46:38 -0700 Subject: [PATCH 3/6] test_get_enforced_params_for_service_account_settings --- .../proxy/test_litellm_pre_call_utils.py | 61 +++++++++++++++++++ 1 file changed, 61 insertions(+) create mode 100644 tests/litellm/proxy/test_litellm_pre_call_utils.py 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..5124b6d2595 --- /dev/null +++ b/tests/litellm/proxy/test_litellm_pre_call_utils.py @@ -0,0 +1,61 @@ +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 = { + "enforced_params": ["user"], + "service_account_settings": {"enforced_params": ["metadata.service"]}, + } + result = _get_enforced_params( + general_settings=general_settings_with_service, + user_api_key_dict=service_account_token, + ) + assert result == ["user", "metadata.service"] + + regular_token = UserAPIKeyAuth(api_key="test-key", metadata={}) + result = _get_enforced_params( + general_settings=general_settings_with_service, user_api_key_dict=regular_token + ) + assert result == ["user"] From 01bd418e0b99b957c913e1b4caf24be3ccc08d3f Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 1 Apr 2025 07:55:28 -0700 Subject: [PATCH 4/6] remove deprecated logic --- litellm/proxy/auth/user_api_key_auth.py | 7 ------- 1 file changed, 7 deletions(-) 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 From f805e15f7b40ee0b1879ffdc94acb3a9f3bfa218 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 1 Apr 2025 08:39:41 -0700 Subject: [PATCH 5/6] test_get_enforced_params_for_service_account_settings --- tests/litellm/proxy/test_litellm_pre_call_utils.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/tests/litellm/proxy/test_litellm_pre_call_utils.py b/tests/litellm/proxy/test_litellm_pre_call_utils.py index 5124b6d2595..b671f715067 100644 --- a/tests/litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/litellm/proxy/test_litellm_pre_call_utils.py @@ -44,18 +44,20 @@ def test_get_enforced_params_for_service_account_settings(): service_account_token = UserAPIKeyAuth( api_key="test-key", metadata={"service_account_id": "test-service-account"} ) - general_settings_with_service = { - "enforced_params": ["user"], + general_settings_with_service_account_settings = { "service_account_settings": {"enforced_params": ["metadata.service"]}, } result = _get_enforced_params( - general_settings=general_settings_with_service, + general_settings=general_settings_with_service_account_settings, user_api_key_dict=service_account_token, ) - assert result == ["user", "metadata.service"] + assert result == ["metadata.service"] - regular_token = UserAPIKeyAuth(api_key="test-key", metadata={}) + regular_token = UserAPIKeyAuth( + api_key="test-key", metadata={"enforced_params": ["user"]} + ) result = _get_enforced_params( - general_settings=general_settings_with_service, user_api_key_dict=regular_token + general_settings=general_settings_with_service_account_settings, + user_api_key_dict=regular_token, ) assert result == ["user"] From c2c5dbf24f565f1ac1e6895caebac43bbc675a4d Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 1 Apr 2025 08:41:53 -0700 Subject: [PATCH 6/6] test_get_enforced_params --- .../proxy/test_litellm_pre_call_utils.py | 42 +++++++++++++++++++ tests/proxy_unit_tests/test_proxy_utils.py | 36 ---------------- 2 files changed, 42 insertions(+), 36 deletions(-) diff --git a/tests/litellm/proxy/test_litellm_pre_call_utils.py b/tests/litellm/proxy/test_litellm_pre_call_utils.py index b671f715067..94f2c512ea9 100644 --- a/tests/litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/litellm/proxy/test_litellm_pre_call_utils.py @@ -61,3 +61,45 @@ def test_get_enforced_params_for_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 9fb00c094fe..b28948094ec 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", [