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

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)

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"]
+      model: mistral/*
+      api_key: os.environ/MISTRAL_API_KEY
+      access_groups: ["beta-models"]

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

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)]
 )

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,
+        }
+    )

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

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,