diff --git a/.circleci/config.yml b/.circleci/config.yml index 78bdf3d8e..d33f62cf3 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1373,6 +1373,7 @@ jobs: name: Install Dependencies command: | npm install -D @playwright/test + npm install @google-cloud/vertexai pip install "pytest==7.3.1" pip install "pytest-retry==1.6.3" pip install "pytest-asyncio==0.21.1" @@ -1434,7 +1435,7 @@ jobs: - run: name: Run Playwright Tests command: | - npx playwright test --reporter=html --output=test-results + npx playwright test e2e_ui_tests/ --reporter=html --output=test-results no_output_timeout: 120m - store_test_results: path: test-results diff --git a/docs/my-website/docs/proxy/virtual_keys.md b/docs/my-website/docs/proxy/virtual_keys.md index 98b06d33b..5bbb6b2a0 100644 --- a/docs/my-website/docs/proxy/virtual_keys.md +++ b/docs/my-website/docs/proxy/virtual_keys.md @@ -820,6 +820,7 @@ litellm_settings: key_generation_settings: team_key_generation: allowed_team_member_roles: ["admin"] + required_params: ["tags"] # require team admins to set tags for cost-tracking when generating a team key personal_key_generation: # maps to 'Default Team' on UI allowed_user_roles: ["proxy_admin"] ``` @@ -829,10 +830,12 @@ litellm_settings: ```python class TeamUIKeyGenerationConfig(TypedDict): allowed_team_member_roles: List[str] + required_params: List[str] # require params on `/key/generate` to be set if a team key (team_id in request) is being generated class PersonalUIKeyGenerationConfig(TypedDict): allowed_user_roles: List[LitellmUserRoles] + required_params: List[str] # require params on `/key/generate` to be set if a personal key (no team_id in request) is being generated class StandardKeyGenerationConfig(TypedDict, total=False): diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 7baf2224c..7ff209094 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -11,28 +11,4 @@ model_list: model: vertex_ai/claude-3-5-sonnet-v2 vertex_ai_project: "adroit-crow-413218" vertex_ai_location: "us-east5" - - model_name: fake-openai-endpoint - litellm_params: - model: openai/fake - api_key: fake-key - api_base: https://exampleopenaiendpoint-production.up.railway.app/ - -router_settings: - model_group_alias: - "gpt-4-turbo": # Aliased model name - model: "gpt-4" # Actual model name in 'model_list' - hidden: true -litellm_settings: - default_team_settings: - - team_id: team-1 - success_callback: ["langfuse"] - failure_callback: ["langfuse"] - langfuse_public_key: os.environ/LANGFUSE_PROJECT1_PUBLIC # Project 1 - langfuse_secret: os.environ/LANGFUSE_PROJECT1_SECRET # Project 1 - - team_id: team-2 - success_callback: ["langfuse"] - failure_callback: ["langfuse"] - langfuse_public_key: os.environ/LANGFUSE_PROJECT2_PUBLIC # Project 2 - langfuse_secret: os.environ/LANGFUSE_PROJECT2_SECRET # Project 2 - langfuse_host: https://us.cloud.langfuse.com diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ab13616d5..511e5a940 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -39,16 +39,20 @@ from litellm.proxy.utils import ( handle_exception_on_proxy, ) from litellm.secret_managers.main import get_secret +from litellm.types.utils import PersonalUIKeyGenerationConfig, TeamUIKeyGenerationConfig def _is_team_key(data: GenerateKeyRequest): return data.team_id is not None -def _team_key_generation_check(user_api_key_dict: UserAPIKeyAuth): +def _team_key_generation_team_member_check( + user_api_key_dict: UserAPIKeyAuth, + team_key_generation: Optional[TeamUIKeyGenerationConfig], +): if ( - litellm.key_generation_settings is None - or litellm.key_generation_settings.get("team_key_generation") is None + team_key_generation is None + or "allowed_team_member_roles" not in team_key_generation ): return True @@ -59,12 +63,7 @@ def _team_key_generation_check(user_api_key_dict: UserAPIKeyAuth): ) team_member_role = user_api_key_dict.team_member.role - if ( - team_member_role - not in litellm.key_generation_settings["team_key_generation"][ # type: ignore - "allowed_team_member_roles" - ] - ): + if team_member_role not in team_key_generation["allowed_team_member_roles"]: raise HTTPException( status_code=400, detail=f"Team member role {team_member_role} not in allowed_team_member_roles={litellm.key_generation_settings['team_key_generation']['allowed_team_member_roles']}", # type: ignore @@ -72,7 +71,67 @@ def _team_key_generation_check(user_api_key_dict: UserAPIKeyAuth): return True -def _personal_key_generation_check(user_api_key_dict: UserAPIKeyAuth): +def _key_generation_required_param_check( + data: GenerateKeyRequest, required_params: Optional[List[str]] +): + if required_params is None: + return True + + data_dict = data.model_dump(exclude_unset=True) + for param in required_params: + if param not in data_dict: + raise HTTPException( + status_code=400, + detail=f"Required param {param} not in data", + ) + return True + + +def _team_key_generation_check( + user_api_key_dict: UserAPIKeyAuth, data: GenerateKeyRequest +): + if ( + litellm.key_generation_settings is None + or litellm.key_generation_settings.get("team_key_generation") is None + ): + return True + + _team_key_generation = litellm.key_generation_settings["team_key_generation"] # type: ignore + + _team_key_generation_team_member_check( + user_api_key_dict, + team_key_generation=_team_key_generation, + ) + _key_generation_required_param_check( + data, + _team_key_generation.get("required_params"), + ) + + return True + + +def _personal_key_membership_check( + user_api_key_dict: UserAPIKeyAuth, + personal_key_generation: Optional[PersonalUIKeyGenerationConfig], +): + if ( + personal_key_generation is None + or "allowed_user_roles" not in personal_key_generation + ): + return True + + if user_api_key_dict.user_role not in personal_key_generation["allowed_user_roles"]: + raise HTTPException( + status_code=400, + detail=f"Personal key creation has been restricted by admin. Allowed roles={litellm.key_generation_settings['personal_key_generation']['allowed_user_roles']}. Your role={user_api_key_dict.user_role}", # type: ignore + ) + + return True + + +def _personal_key_generation_check( + user_api_key_dict: UserAPIKeyAuth, data: GenerateKeyRequest +): if ( litellm.key_generation_settings is None @@ -80,16 +139,18 @@ def _personal_key_generation_check(user_api_key_dict: UserAPIKeyAuth): ): return True - if ( - user_api_key_dict.user_role - not in litellm.key_generation_settings["personal_key_generation"][ # type: ignore - "allowed_user_roles" - ] - ): - raise HTTPException( - status_code=400, - detail=f"Personal key creation has been restricted by admin. Allowed roles={litellm.key_generation_settings['personal_key_generation']['allowed_user_roles']}. Your role={user_api_key_dict.user_role}", # type: ignore - ) + _personal_key_generation = litellm.key_generation_settings["personal_key_generation"] # type: ignore + + _personal_key_membership_check( + user_api_key_dict, + personal_key_generation=_personal_key_generation, + ) + + _key_generation_required_param_check( + data, + _personal_key_generation.get("required_params"), + ) + return True @@ -99,16 +160,23 @@ def key_generation_check( """ Check if admin has restricted key creation to certain roles for teams or individuals """ - if litellm.key_generation_settings is None: + if ( + litellm.key_generation_settings is None + or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + ): return True ## check if key is for team or individual is_team_key = _is_team_key(data=data) if is_team_key: - return _team_key_generation_check(user_api_key_dict) + return _team_key_generation_check( + user_api_key_dict=user_api_key_dict, data=data + ) else: - return _personal_key_generation_check(user_api_key_dict=user_api_key_dict) + return _personal_key_generation_check( + user_api_key_dict=user_api_key_dict, data=data + ) router = APIRouter() diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 334894320..9fc58dff6 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1604,11 +1604,17 @@ class StandardCallbackDynamicParams(TypedDict, total=False): langsmith_base_url: Optional[str] -class TeamUIKeyGenerationConfig(TypedDict): +class KeyGenerationConfig(TypedDict, total=False): + required_params: List[ + str + ] # specify params that must be present in the key generation request + + +class TeamUIKeyGenerationConfig(KeyGenerationConfig): allowed_team_member_roles: List[str] -class PersonalUIKeyGenerationConfig(TypedDict): +class PersonalUIKeyGenerationConfig(KeyGenerationConfig): allowed_user_roles: List[str] diff --git a/tests/proxy_admin_ui_tests/playwright.config.ts b/tests/proxy_admin_ui_tests/playwright.config.ts index c77897a02..3be77a319 100644 --- a/tests/proxy_admin_ui_tests/playwright.config.ts +++ b/tests/proxy_admin_ui_tests/playwright.config.ts @@ -13,6 +13,8 @@ import { defineConfig, devices } from '@playwright/test'; */ export default defineConfig({ testDir: './e2e_ui_tests', + testIgnore: ['**/tests/pass_through_tests/**', '../pass_through_tests/**/*'], + testMatch: '**/*.spec.ts', // Only run files ending in .spec.ts /* Run tests in files in parallel */ fullyParallel: true, /* Fail the build on CI if you accidentally left test.only in the source code. */ diff --git a/tests/proxy_admin_ui_tests/test_key_management.py b/tests/proxy_admin_ui_tests/test_key_management.py index 81d9fb676..0b392a268 100644 --- a/tests/proxy_admin_ui_tests/test_key_management.py +++ b/tests/proxy_admin_ui_tests/test_key_management.py @@ -551,7 +551,7 @@ def test_is_team_key(): assert not _is_team_key(GenerateKeyRequest(user_id="test_user_id")) -def test_team_key_generation_check(): +def test_team_key_generation_team_member_check(): from litellm.proxy.management_endpoints.key_management_endpoints import ( _team_key_generation_check, ) @@ -562,22 +562,86 @@ def test_team_key_generation_check(): } assert _team_key_generation_check( - UserAPIKeyAuth( + user_api_key_dict=UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-1234", team_member=Member(role="admin", user_id="test_user_id"), - ) + ), + data=GenerateKeyRequest(), ) with pytest.raises(HTTPException): _team_key_generation_check( - UserAPIKeyAuth( + user_api_key_dict=UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-1234", user_id="test_user_id", team_member=Member(role="user", user_id="test_user_id"), + ), + data=GenerateKeyRequest(), + ) + + +@pytest.mark.parametrize( + "team_key_generation_settings, input_data, expected_result", + [ + ({"required_params": ["tags"]}, GenerateKeyRequest(tags=["test_tags"]), True), + ({}, GenerateKeyRequest(), True), + ( + {"required_params": ["models"]}, + GenerateKeyRequest(tags=["test_tags"]), + False, + ), + ], +) +@pytest.mark.parametrize("key_type", ["team_key", "personal_key"]) +def test_key_generation_required_params_check( + team_key_generation_settings, input_data, expected_result, key_type +): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _team_key_generation_check, + _personal_key_generation_check, + ) + from litellm.types.utils import ( + TeamUIKeyGenerationConfig, + StandardKeyGenerationConfig, + PersonalUIKeyGenerationConfig, + ) + from fastapi import HTTPException + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-1234", + user_id="test_user_id", + team_id="test_team_id", + team_member=Member(role="admin", user_id="test_user_id"), + ) + + if key_type == "team_key": + litellm.key_generation_settings = StandardKeyGenerationConfig( + team_key_generation=TeamUIKeyGenerationConfig( + **team_key_generation_settings ) ) + elif key_type == "personal_key": + litellm.key_generation_settings = StandardKeyGenerationConfig( + personal_key_generation=PersonalUIKeyGenerationConfig( + **team_key_generation_settings + ) + ) + + if expected_result: + if key_type == "team_key": + assert _team_key_generation_check(user_api_key_dict, input_data) + elif key_type == "personal_key": + assert _personal_key_generation_check(user_api_key_dict, input_data) + else: + if key_type == "team_key": + with pytest.raises(HTTPException): + _team_key_generation_check(user_api_key_dict, input_data) + elif key_type == "personal_key": + with pytest.raises(HTTPException): + _personal_key_generation_check(user_api_key_dict, input_data) def test_personal_key_generation_check(): @@ -591,16 +655,18 @@ def test_personal_key_generation_check(): } assert _personal_key_generation_check( - UserAPIKeyAuth( + user_api_key_dict=UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="admin" - ) + ), + data=GenerateKeyRequest(), ) with pytest.raises(HTTPException): _personal_key_generation_check( - UserAPIKeyAuth( + user_api_key_dict=UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-1234", user_id="admin", - ) + ), + data=GenerateKeyRequest(), ) diff --git a/ui/litellm-dashboard/src/components/create_key_button.tsx b/ui/litellm-dashboard/src/components/create_key_button.tsx index 0af3a064c..4f771b111 100644 --- a/ui/litellm-dashboard/src/components/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/create_key_button.tsx @@ -40,6 +40,31 @@ interface CreateKeyProps { setData: React.Dispatch>; } +const getPredefinedTags = (data: any[] | null) => { + let allTags = []; + + console.log("data:", JSON.stringify(data)); + + if (data) { + for (let key of data) { + if (key["metadata"] && key["metadata"]["tags"]) { + allTags.push(...key["metadata"]["tags"]); + } + } + } + + // Deduplicate using Set + const uniqueTags = Array.from(new Set(allTags)).map(tag => ({ + value: tag, + label: tag, + })); + + + console.log("uniqueTags:", uniqueTags); + return uniqueTags; +} + + const CreateKey: React.FC = ({ userID, team, @@ -55,6 +80,8 @@ const CreateKey: React.FC = ({ const [userModels, setUserModels] = useState([]); const [modelsToPick, setModelsToPick] = useState([]); const [keyOwner, setKeyOwner] = useState("you"); + const [predefinedTags, setPredefinedTags] = useState(getPredefinedTags(data)); + const handleOk = () => { setIsModalVisible(false); @@ -355,6 +382,15 @@ const CreateKey: React.FC = ({ placeholder="Enter metadata as JSON" /> + +