From 5e03ef73820f0ef64d472bd0f36186f5b91bf1c7 Mon Sep 17 00:00:00 2001 From: = Date: Sat, 4 Oct 2025 14:32:15 -0700 Subject: [PATCH 1/3] fixes bloated key alias network calls with lean endpoint --- .../key_management_endpoints.py | 75 ++++++++++++++++++- .../test_key_generate_prisma.py | 54 ++++++++++++- .../key_team_helpers/filter_helpers.ts | 42 ++--------- .../src/components/networking.tsx | 35 +++++++++ 4 files changed, 168 insertions(+), 38 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 007c0164be4..54e409bba79 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -15,7 +15,7 @@ import json import secrets import traceback from datetime import datetime, timedelta, timezone -from typing import List, Literal, Optional, Tuple, cast +from typing import List, Literal, Optional, Tuple, cast, Set import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status @@ -2614,6 +2614,79 @@ async def list_keys( code=status.HTTP_500_INTERNAL_SERVER_ERROR, ) +@router.get( + "/key/aliases", + tags=["key management"], + dependencies=[Depends(user_api_key_auth)], +) +@management_endpoint_wrapper +async def key_aliases() -> Dict[str, List[str]]: + """ + Lists all key aliases + + Returns: + { + "aliases": List[str] + } + """ + try: + from litellm.proxy.proxy_server import prisma_client + + verbose_proxy_logger.debug("Entering key_aliases function") + + if prisma_client is None: + verbose_proxy_logger.error("Database not connected") + raise Exception("Database not connected") + + where: Dict[str, Any] = {} + try: + where.update(_get_condition_to_filter_out_ui_session_tokens()) + except NameError: + # Helper may not exist in some builds; ignore if missing + pass + + rows = await prisma_client.db.litellm_verificationtoken.find_many( + where=where, + order=[{"key_alias": "asc"}], + ) + + seen = set() + aliases: List[str] = [] + for row in rows: + alias = getattr(row, "key_alias", None) + if alias is None and isinstance(row, dict): + alias = row.get("key_alias") + + if not alias: + continue + + alias_str = str(alias).strip() + if alias_str and alias_str not in seen: + seen.add(alias_str) + aliases.append(alias_str) + + verbose_proxy_logger.debug(f"Returning {len(aliases)} key aliases") + + return {"aliases": aliases} + + except Exception as e: + verbose_proxy_logger.exception(f"Error in key_aliases: {e}") + if isinstance(e, HTTPException): + raise ProxyException( + message=getattr(e, "detail", f"error({str(e)})"), + type=ProxyErrorTypes.internal_server_error, + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), + ) + elif isinstance(e, ProxyException): + raise e + raise ProxyException( + message="Authentication Error, " + str(e), + type=ProxyErrorTypes.internal_server_error, + param=getattr(e, "param", "None"), + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + def _validate_sort_params( sort_by: Optional[str], sort_order: str diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index 6eb4efa7d63..a8df57283ca 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -60,7 +60,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( info_key_fn, list_keys, regenerate_key_fn, - update_key_fn, + update_key_fn, key_aliases, ) from litellm.proxy.management_endpoints.team_endpoints import ( new_team, @@ -3528,6 +3528,58 @@ async def test_list_keys(prisma_client): assert _key in response["keys"] +@pytest.mark.asyncio +async def test_key_aliases(prisma_client): + """ + Test the key_aliases function: + - Returns a list + - Includes alias from a newly created key + - Aliases are unique and sorted + """ + import asyncio + import uuid + import litellm + from litellm.proxy._types import LitellmUserRoles + + # Wire up test prisma client + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + await litellm.proxy.proxy_server.prisma_client.connect() + + # Basic call + response = await key_aliases() + assert "aliases" in response + assert isinstance(response["aliases"], list) + + # Create a new user (and key) with a unique alias + unique_id = str(uuid.uuid4()) + test_alias = f"key-aliases-test-{unique_id}" + test_user_id = f"key-aliases-user-{unique_id}" + + await new_user( + data=NewUserRequest( + user_id=test_user_id, + user_role=LitellmUserRoles.INTERNAL_USER, + key_alias=test_alias, + ), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + # Allow async DB writes to settle + await asyncio.sleep(2) + + # Call again and validate + response_after = await key_aliases() + aliases = response_after["aliases"] + + # Contains the new alias + assert test_alias in aliases + + # Unique & sorted (endpoint dedupes and orders ascending) + assert len(aliases) == len(set(aliases)) + assert aliases == sorted(aliases) + + @pytest.mark.asyncio async def test_auth_vertex_ai_route(prisma_client): """ diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts index 887815596db..a46666cddc9 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts +++ b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts @@ -1,9 +1,9 @@ -import { keyListCall, teamListCall, organizationListCall } from "../networking"; +import { teamListCall, organizationListCall, keyAliasesCall } from "../networking" import { Team } from "./key_list"; import { Organization } from "../networking"; /** - * Fetches all key aliases across all pages + * Fetches all key aliases via the dedicated /key/aliases endpoint * @param accessToken The access token for API authentication * @returns Array of all unique key aliases */ @@ -13,46 +13,16 @@ export const fetchAllKeyAliases = async ( if (!accessToken) return []; try { - // Fetch all pages of keys to extract aliases - let allAliases: string[] = []; - let currentPage = 1; - let hasMorePages = true; - - while (hasMorePages) { - const response = await keyListCall( - accessToken, - null, // organization_id - "", // team_id - null, // selectedKeyAlias - null, // user_id - null, // key_hash - currentPage, - 100 // larger page size to reduce number of requests - ); - - // Extract aliases from this page - const pageAliases = response.keys - .map((key: any) => key.key_alias) - .filter(Boolean) as string[]; - - allAliases = [...allAliases, ...pageAliases]; - - // Check if there are more pages - if (currentPage < response.total_pages) { - currentPage++; - } else { - hasMorePages = false; - } - } - - // Remove duplicates - return Array.from(new Set(allAliases)); + const { aliases } = await keyAliasesCall(accessToken as unknown as String); + // Defensive dedupe & null-guard + return Array.from(new Set((aliases || []).filter(Boolean))); } catch (error) { console.error("Error fetching all key aliases:", error); return []; } }; + /** * Fetches all teams across all pages * @param accessToken The access token for API authentication diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 6b8c7671a10..f674df444a5 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -3305,6 +3305,41 @@ export const keyListCall = async ( } }; +export const keyAliasesCall = async ( + accessToken: String +): Promise<{ aliases: string[] }> => { + /** + * Get all key aliases from proxy + */ + try { + let url = proxyBaseUrl ? `${proxyBaseUrl}/key/aliases` : `/key/aliases`; + console.log("in keyAliasesCall"); + + const response = await fetch(url, { + method: "GET", + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + }); + + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + + const data = await response.json(); + console.log("/key/aliases API Response:", data); + return data; // { aliases: string[] } + } catch (error) { + console.error("Failed to fetch key aliases:", error); + throw error; + } +}; + + export const spendUsersCall = async (accessToken: String, userID: String) => { try { const url = proxyBaseUrl ? `${proxyBaseUrl}/spend/users` : `/spend/users`; From 6ba077593f1601bcadb7f15aa729002d38a49ee9 Mon Sep 17 00:00:00 2001 From: = Date: Sat, 4 Oct 2025 14:36:19 -0700 Subject: [PATCH 2/3] Update test_key_generate_prisma.py --- tests/proxy_unit_tests/test_key_generate_prisma.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index a8df57283ca..1968fde8436 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -60,7 +60,8 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( info_key_fn, list_keys, regenerate_key_fn, - update_key_fn, key_aliases, + update_key_fn, + key_aliases, ) from litellm.proxy.management_endpoints.team_endpoints import ( new_team, @@ -151,7 +152,6 @@ def prisma_client(): @pytest.mark.flaky(retries=6, delay=1) async def test_new_user_response(prisma_client): try: - print("prisma client=", prisma_client) setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) @@ -424,7 +424,6 @@ async def test_call_with_valid_model_using_all_models(prisma_client): setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") try: - await litellm.proxy.proxy_server.prisma_client.connect() team_request = NewTeamRequest( @@ -1786,7 +1785,6 @@ async def test_call_with_key_over_model_budget( litellm.callbacks.append(model_budget_limiter) try: - # set budget for chatgpt-v-3 to 0.000001, expect the next request to fail model_max_budget = { "gpt-4o-mini": { From 1ce21d1f58df807fcb11660f1ce6466545d153b4 Mon Sep 17 00:00:00 2001 From: = Date: Sat, 4 Oct 2025 16:02:15 -0700 Subject: [PATCH 3/3] Update key_management_endpoints.py --- litellm/proxy/management_endpoints/key_management_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 54e409bba79..cf101543011 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -15,7 +15,7 @@ import json import secrets import traceback from datetime import datetime, timedelta, timezone -from typing import List, Literal, Optional, Tuple, cast, Set +from typing import List, Literal, Optional, Tuple, cast import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status