Virtual Keys: Filter by key alias (#10035) (#10085)

Co-authored-by: Christian Owusu <36159205+crisshaker@users.noreply.github.com>

ui/litellm-dashboard/src/components/all_keys_table.tsx

@@ -30,6 +30,8 @@ interface AllKeysTableProps {
   teams: Team[] | null;
   selectedTeam: Team | null;
   setSelectedTeam: (team: Team | null) => void;
+  selectedKeyAlias: string | null;
+  setSelectedKeyAlias: Setter<string | null>;
   accessToken: string | null;
   userID: string | null;
   userRole: string | null;
@@ -98,6 +100,8 @@ export function AllKeysTable({
   teams,
   selectedTeam,
   setSelectedTeam,
+  selectedKeyAlias,
+  setSelectedKeyAlias,
   accessToken,
   userID,
   userRole,
@@ -123,7 +127,8 @@ export function AllKeysTable({
     organizations,
     accessToken,
     setSelectedTeam,
-    setCurrentOrg
+    setCurrentOrg,
+    setSelectedKeyAlias
   });
 
   useEffect(() => {
@@ -358,6 +363,23 @@ export function AllKeysTable({
           }));
       }
     },
+    {
+      name: "Key Alias",
+      label: "Key Alias",
+      isSearchable: true,
+      searchFn: async (searchText) => {
+        const filteredKeyAliases = allKeyAliases.filter(key => {
+          return key.toLowerCase().includes(searchText.toLowerCase())
+        });
+
+        return filteredKeyAliases.map((key) => {
+          return {
+            label: key,
+            value: key
+          }
+        });
+      }
+    }
   ];
   
   

ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts

@@ -21,6 +21,7 @@ export const fetchAllKeyAliases = async (accessToken: string | null): Promise<st
         accessToken,
         null, // organization_id
         "", // team_id
+        null, // selectedKeyAlias
         currentPage,
         100 // larger page size to reduce number of requests
       );

ui/litellm-dashboard/src/components/key_team_helpers/filter_logic.tsx

@@ -2,10 +2,15 @@ import { useEffect, useState } from "react";
 import { KeyResponse } from "../key_team_helpers/key_list";
 import { Organization } from "../networking";
 import { Team } from "../key_team_helpers/key_list";
+import { useQuery } from "@tanstack/react-query";
+import { fetchAllKeyAliases, fetchAllOrganizations, fetchAllTeams } from "./filter_helpers";
+import { Setter } from "@/types";
+
 
 export interface FilterState {
   'Team ID': string;
   'Organization ID': string;
+  'Key Alias': string;
   [key: string]: string;
 }
 
@@ -15,7 +20,8 @@ export function useFilterLogic({
   organizations,
   accessToken,
   setSelectedTeam,
-  setCurrentOrg
+  setCurrentOrg,
+  setSelectedKeyAlias
 }: {
   keys: KeyResponse[];
   teams: Team[] | null;
@@ -23,12 +29,13 @@ export function useFilterLogic({
   accessToken: string | null;
   setSelectedTeam: (team: Team | null) => void;
   setCurrentOrg: React.Dispatch<React.SetStateAction<Organization | null>>;
+  setSelectedKeyAlias: Setter<string | null>
 }) {
   const [filters, setFilters] = useState<FilterState>({
     'Team ID': '',
     'Organization ID': '',
+    'Key Alias': ''
   });
-  const [allKeyAliases, setAllKeyAliases] = useState<string[]>([]);
   const [allTeams, setAllTeams] = useState<Team[]>(teams || []);
   const [allOrganizations, setAllOrganizations] = useState<Organization[]>(organizations || []);
   const [filteredKeys, setFilteredKeys] = useState<KeyResponse[]>(keys);
@@ -79,6 +86,16 @@ export function useFilterLogic({
     }
   }, [accessToken]);
 
+  const queryAllKeysQuery = useQuery({
+    queryKey: ['allKeys'],
+    queryFn: async () => {
+      if (!accessToken) throw new Error('Access token required');
+      return await fetchAllKeyAliases(accessToken);
+    },
+    enabled: !!accessToken
+  });
+  const allKeyAliases = queryAllKeysQuery.data || []
+
   // Update teams and organizations when props change
   useEffect(() => {
     if (teams && teams.length > 0) {
@@ -103,6 +120,7 @@ export function useFilterLogic({
     setFilters({
       'Team ID': newFilters['Team ID'] || '',
       'Organization ID': newFilters['Organization ID'] || '',
+      'Key Alias': newFilters['Key Alias'] || ''
     });
   
     // Handle Team change
@@ -120,6 +138,12 @@ export function useFilterLogic({
         setCurrentOrg(selectedOrg);
       }
     }
+
+    const keyAlias = newFilters['Key Alias'];
+    const selectedKeyAlias = keyAlias
+      ? allKeyAliases.find((k) => k === keyAlias) || null
+      : null;
+    setSelectedKeyAlias(selectedKeyAlias)
   };
 
   const handleFilterReset = () => {
@@ -127,6 +151,7 @@ export function useFilterLogic({
     setFilters({
       'Team ID': '',
       'Organization ID': '',
+      'Key Alias': ''
     });
     
     // Reset team and org selections
@@ -144,6 +169,3 @@ export function useFilterLogic({
     handleFilterReset
   };
 }
-
-// These functions are imported from key_team_helpers/filter_helpers.ts
-import { fetchAllKeyAliases, fetchAllTeams, fetchAllOrganizations } from './filter_helpers'; 

ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx

@@ -85,6 +85,7 @@ total_pages: number;
 interface UseKeyListProps {
 selectedTeam?: Team;
 currentOrg: Organization | null;
+selectedKeyAlias: string | null;
 accessToken: string;
 currentPage?: number;
 }
@@ -108,9 +109,9 @@ setKeys: Setter<KeyResponse[]>;
 const useKeyList = ({
     selectedTeam,
     currentOrg,
+    selectedKeyAlias,
     accessToken,
     currentPage = 1,
-
 }: UseKeyListProps): UseKeyListReturn => {
     const [keyData, setKeyData] = useState<KeyListResponse>({ 
         keys: [], 
@@ -134,8 +135,9 @@ const useKeyList = ({
                 accessToken,
                 currentOrg?.organization_id || null,
                 selectedTeam?.team_id || "",
+                selectedKeyAlias,
                 params.page as number || 1,
-                50
+                50,
             );
             console.log("data", data);
             setKeyData(data);
@@ -149,8 +151,17 @@ const useKeyList = ({
 
     useEffect(() => {
         fetchKeys();
-        console.log("selectedTeam", selectedTeam, "currentOrg", currentOrg, "accessToken", accessToken);
-    }, [selectedTeam, currentOrg, accessToken]);
+        console.log(
+          'selectedTeam',
+          selectedTeam,
+          'currentOrg',
+          currentOrg,
+          'accessToken',
+          accessToken,
+          'selectedKeyAlias',
+          selectedKeyAlias
+        );
+    }, [selectedTeam, currentOrg, accessToken, selectedKeyAlias]);
 
     const setKeys = (newKeysOrUpdater: KeyResponse[] | ((prevKeys: KeyResponse[]) => KeyResponse[])) => {
         setKeyData(prevData => {

ui/litellm-dashboard/src/components/networking.tsx

@@ -2569,8 +2569,9 @@ export const keyListCall = async (
   accessToken: String, 
   organizationID: string | null,
   teamID: string | null,
+  selectedKeyAlias: string | null,
   page: number,
-  pageSize: number
+  pageSize: number,
 ) => {
   /**
    * Get all available teams on proxy
@@ -2588,6 +2589,10 @@ export const keyListCall = async (
       queryParams.append('organization_id', organizationID.toString());
     }
 
+    if (selectedKeyAlias) {
+      queryParams.append('key_alias', selectedKeyAlias)
+    }
+
     if (page) {
       queryParams.append('page', page.toString());
     }

ui/litellm-dashboard/src/components/user_dashboard.tsx

@@ -108,6 +108,7 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
     team_id: null,
   };
   const [selectedTeam, setSelectedTeam] = useState<any | null>(null);
+  const [selectedKeyAlias, setSelectedKeyAlias] = useState<string | null>(null);
   // check if window is not undefined
   if (typeof window !== "undefined") {
     window.addEventListener("beforeunload", function () {
@@ -350,6 +351,8 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
             accessToken={accessToken}
             selectedTeam={selectedTeam ? selectedTeam : null}
             setSelectedTeam={setSelectedTeam}
+            selectedKeyAlias={selectedKeyAlias}
+            setSelectedKeyAlias={setSelectedKeyAlias}
             data={keys}
             setData={setKeys}
             premiumUser={premiumUser}

ui/litellm-dashboard/src/components/view_key_table.tsx

@@ -71,6 +71,7 @@ import useKeyList from "./key_team_helpers/key_list";
 import { KeyResponse } from "./key_team_helpers/key_list";
 import { AllKeysTable } from "./all_keys_table";
 import { Team } from "./key_team_helpers/key_list";
+import { Setter } from "@/types";
 
 const isLocal = process.env.NODE_ENV === "development";
 const proxyBaseUrl = isLocal ? "http://localhost:4000" : null;
@@ -107,6 +108,8 @@ interface ViewKeyTableProps {
   currentOrg: Organization | null;
   organizations: Organization[] | null;
   setCurrentOrg: React.Dispatch<React.SetStateAction<Organization | null>>;
+  selectedKeyAlias: string | null;
+  setSelectedKeyAlias: Setter<string | null>;
 }
 
 interface ItemData {
@@ -154,7 +157,9 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
   premiumUser,
   currentOrg,
   organizations,
-  setCurrentOrg
+  setCurrentOrg,
+  selectedKeyAlias,
+  setSelectedKeyAlias
 }) => {
   const [isButtonClicked, setIsButtonClicked] = useState(false);
   const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false);
@@ -179,6 +184,7 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
   const { keys, isLoading, error, pagination, refresh, setKeys } = useKeyList({
     selectedTeam,
     currentOrg,
+    selectedKeyAlias,
     accessToken,
   });
 
@@ -432,6 +438,8 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
         organizations={organizations}
         setCurrentOrg={setCurrentOrg}
         refresh={refresh}
+        selectedKeyAlias={selectedKeyAlias}
+        setSelectedKeyAlias={setSelectedKeyAlias}
       />
 
       {isDeleteModalOpen && (