From 3e50a06fb4006838e6854e20ba5995ce79f486fc Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 27 Nov 2024 16:22:04 -0800 Subject: [PATCH] fix _enforce_unique_key_alias --- .../management_endpoints/key_management_endpoints.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index b44e233d9..d0670d062 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -593,6 +593,7 @@ async def update_key_fn( await _enforce_unique_key_alias( key_alias=non_default_values.get("key_alias", None), prisma_client=prisma_client, + existing_key_id=key, ) response = await prisma_client.update_data( @@ -1898,6 +1899,7 @@ async def test_key_logging( async def _enforce_unique_key_alias( key_alias: Optional[str], prisma_client: Any, + existing_key_id: Optional[str] = None, ) -> None: """ Helper to enforce unique key aliases across all keys. @@ -1905,13 +1907,19 @@ async def _enforce_unique_key_alias( Args: key_alias (Optional[str]): The key alias to check prisma_client (Any): Prisma client instance + existing_key_id (Optional[str]): ID of existing key being updated, to exclude from uniqueness check Raises: - HTTPException: If key alias already exists + HTTPException: If key alias already exists on a different key """ if key_alias is not None and prisma_client is not None: + where_clause: dict[str, Any] = {"key_alias": key_alias} + if existing_key_id: + # Exclude the current key from the uniqueness check + where_clause["NOT"] = {"token": existing_key_id} + existing_key = await prisma_client.db.litellm_verificationtoken.find_first( - where={"key_alias": key_alias} + where=where_clause ) if existing_key is not None: raise HTTPException(