From d282babd3eabe5260fe22177dd5d9dd17ade8e72 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Mon, 19 May 2025 21:01:53 -0700 Subject: [PATCH] Validate migrating keys to teams + Fix mistral image url on async translation (#10966) * feat(key_management_endpoints.py): add validation checks for migrating key to team Ensures requests with migrated key can actually succeed Prevent migrated keys from failing in prod due to team missing required permissions * fix(mistral/): fix image url handling for mistral on async call * fix(key_management_endpoints.py): improve check for running team validation on key update --- .../mistral/mistral_chat_transformation.py | 5 +- litellm/proxy/_new_secret_config.yaml | 12 +-- litellm/proxy/auth/auth_checks.py | 4 +- .../key_management_endpoints.py | 96 +++++++++++++++++++ .../test_mistral_chat_transformation.py | 34 +++++++ .../test_team_member_permissions.py | 1 - tests/proxy_unit_tests/test_auth_checks.py | 2 +- 7 files changed, 141 insertions(+), 13 deletions(-) create mode 100644 tests/litellm/llms/mistral/test_mistral_chat_transformation.py diff --git a/litellm/llms/mistral/mistral_chat_transformation.py b/litellm/llms/mistral/mistral_chat_transformation.py index f580414df8..a675beebbd 100644 --- a/litellm/llms/mistral/mistral_chat_transformation.py +++ b/litellm/llms/mistral/mistral_chat_transformation.py @@ -184,7 +184,10 @@ class MistralConfig(OpenAIGPTConfig): if _content_block and isinstance(_content_block, list): for c in _content_block: if c.get("type") == "image_url": - return messages + if is_async: + return super()._transform_messages(messages, model, True) + else: + return super()._transform_messages(messages, model, False) ## 2. If content is list, then convert to string messages = handle_messages_with_content_list_to_str_conversion(messages) diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index a356818c85..a6fb25f3a7 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -66,12 +66,8 @@ model_list: model_info: id: my-general-azure-deployment mode: batch - - model_name: "gpt-4o-batch" + - model_name: mistral/* litellm_params: - model: azure/gpt-4o-mini - api_base: https://krris-m2f9a9i7-eastus2.openai.azure.com - api_key: 04d22fb7e9ad4d9c8afe7c6abf97a6fc - model_info: - id: my-unique-azure-deployment - mode: batch - access_groups: ["beta-models"] \ No newline at end of file + model: mistral/* + api_key: os.environ/MISTRAL_API_KEY + access_groups: ["beta-models"] diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b95d16ebbb..e18b358fa6 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -103,7 +103,7 @@ async def common_checks( # 2. If team can call model if _model and team_object: - if not await can_team_access_model( + if not can_team_access_model( model=_model, team_object=team_object, llm_router=llm_router, @@ -1284,7 +1284,7 @@ def can_org_access_model( ) -async def can_team_access_model( +def can_team_access_model( model: Union[str, List[str]], team_object: Optional[LiteLLM_TeamTable], llm_router: Optional[Router], diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 69daa2debb..4c0e9466e5 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -30,6 +30,7 @@ from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import ( _cache_key_object, _delete_cache_key_object, + can_team_access_model, get_key_object, get_team_object, ) @@ -683,6 +684,16 @@ def prepare_key_update_data( return non_default_values +def is_different_team( + data: UpdateKeyRequest, existing_key_row: LiteLLM_VerificationToken +) -> bool: + if data.team_id is None: + return False + if existing_key_row.team_id is None: + return True + return data.team_id != existing_key_row.team_id + + @router.post( "/key/update", tags=["key management"], dependencies=[Depends(user_api_key_auth)] ) @@ -788,6 +799,27 @@ async def update_key_fn( user_api_key_cache=user_api_key_cache, ) + # if team change - check if this is possible + if is_different_team(data=data, existing_key_row=existing_key_row): + team_obj = await get_team_object( + team_id=cast(str, data.team_id), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + if llm_router is None: + raise HTTPException( + status_code=400, + detail={ + "error": "LLM router not found. Please set it up by passing in a valid config.yaml or adding models via the UI." + }, + ) + validate_key_team_change( + key=existing_key_row, + team=team_obj, + change_initiated_by=user_api_key_dict, + llm_router=llm_router, + ) non_default_values = prepare_key_update_data( data=data, existing_key_row=existing_key_row ) @@ -847,6 +879,70 @@ async def update_key_fn( ) +def validate_key_team_change( + key: LiteLLM_VerificationToken, + team: LiteLLM_TeamTable, + change_initiated_by: UserAPIKeyAuth, + llm_router: Router, +): + """ + Validate that a key can be moved to a new team. + + - The team must have access to the key's models + - The key's user_id must be a member of the team + - The key's tpm/rpm limit must be less than the team's tpm/rpm limit + - The person initiating the change must be either Proxy Admin or Team Admin + """ + # Check if the team has access to the key's models + if len(key.models) > 0: + for model in key.models: + can_team_access_model( + model=model, + team_object=team, + llm_router=llm_router, + ) + + # Check if the key's user_id is a member of the team + if key.user_id is not None: + is_member = False + for member in team.members_with_roles: + if member.user_id == key.user_id: + is_member = True + break + if not is_member: + raise HTTPException( + status_code=403, + detail=f"User={key.user_id} is not a member of the team={team.team_id}. Check team members via `/team/info`.", + ) + + # Check if the key's tpm/rpm limit is less than the team's tpm/rpm limit + if key.tpm_limit is not None: + if team.tpm_limit and key.tpm_limit > team.tpm_limit: + raise HTTPException( + status_code=403, + detail=f"Key={key.token} has a tpm_limit={key.tpm_limit} which is greater than the team's tpm_limit={team.tpm_limit}.", + ) + if team.rpm_limit and key.rpm_limit and key.rpm_limit > team.rpm_limit: + raise HTTPException( + status_code=403, + detail=f"Key={key.token} has a rpm_limit={key.rpm_limit} which is greater than the team's rpm_limit={team.rpm_limit}.", + ) + + # Check if the person initiating the change is a Proxy Admin or Team Admin + if change_initiated_by.user_role == LitellmUserRoles.PROXY_ADMIN.value: + return + elif _is_user_team_admin( + user_api_key_dict=change_initiated_by, + team_obj=team, + ): + return + else: + raise HTTPException( + status_code=403, + detail=f"User={change_initiated_by.user_id} is not a Proxy Admin or Team Admin for team={team.team_id}.", + ) + + @router.post( "/key/delete", tags=["key management"], dependencies=[Depends(user_api_key_auth)] ) diff --git a/tests/litellm/llms/mistral/test_mistral_chat_transformation.py b/tests/litellm/llms/mistral/test_mistral_chat_transformation.py new file mode 100644 index 0000000000..e8e50890e1 --- /dev/null +++ b/tests/litellm/llms/mistral/test_mistral_chat_transformation.py @@ -0,0 +1,34 @@ +import os +import sys +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + +from litellm.llms.mistral.mistral_chat_transformation import MistralConfig + + +@pytest.mark.asyncio +async def test_mistral_chat_transformation(): + mistral_config = MistralConfig() + result = mistral_config._transform_messages( + **{ + "messages": [ + { + "content": [ + {"type": "text", "text": "Here is a representation of text"}, + { + "type": "image_url", + "image_url": "https://images.pexels.com/photos/13268478/pexels-photo-13268478.jpeg", + }, + ], + "role": "user", + } + ], + "model": "mistral-medium-latest", + "is_async": True, + } + ) diff --git a/tests/otel_tests/test_team_member_permissions.py b/tests/otel_tests/test_team_member_permissions.py index 8c98fa3b92..3b266a9e68 100644 --- a/tests/otel_tests/test_team_member_permissions.py +++ b/tests/otel_tests/test_team_member_permissions.py @@ -358,7 +358,6 @@ async def test_edit_delete_permissions(): # Generate an admin key for the team admin_key_data = await generate_key(session, master_key, team_id) - admin_key = admin_key_data["key"] key_id = admin_key_data["key"] # Create a user key diff --git a/tests/proxy_unit_tests/test_auth_checks.py b/tests/proxy_unit_tests/test_auth_checks.py index 7695306c87..b78543f2c8 100644 --- a/tests/proxy_unit_tests/test_auth_checks.py +++ b/tests/proxy_unit_tests/test_auth_checks.py @@ -410,7 +410,7 @@ async def test_can_team_access_model(model, team_models, expect_to_work): team_id="test-team", models=team_models, ) - result = await can_team_access_model( + result = can_team_access_model( model=model, team_object=team_object, llm_router=None,