From a7f6a95dd9bd99b041b935567613e4c09ca68fcf Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 8 Mar 2025 15:58:29 -0800
Subject: [PATCH] (UI) - Fix, Allow Filter Keys by Team Alias, Key Alias and
 Org (#9083)

* working team, key, org filters

* fix size of filter

* all keys filter

* 1 file for filter logic

* remove key alias check
---
 .../src/components/all_keys_table.tsx         | 110 +++++++------
 .../components/common_components/filter.tsx   |  54 +++++--
 .../key_team_helpers/filter_helpers.ts        | 127 +++++++++++++++
 .../key_team_helpers/filter_logic.tsx         | 148 ++++++++++++++++++
 4 files changed, 380 insertions(+), 59 deletions(-)
 create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts
 create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/filter_logic.tsx

diff --git a/ui/litellm-dashboard/src/components/all_keys_table.tsx b/ui/litellm-dashboard/src/components/all_keys_table.tsx
index 1f05c243ee2..bc0d1135ea7 100644
--- a/ui/litellm-dashboard/src/components/all_keys_table.tsx
+++ b/ui/litellm-dashboard/src/components/all_keys_table.tsx
@@ -12,6 +12,8 @@ import { FilterOption } from "./common_components/filter";
 import { Organization, userListCall } from "./networking";
 import { createTeamSearchFunction } from "./key_team_helpers/team_search_fn";
 import { createOrgSearchFunction } from "./key_team_helpers/organization_search_fn";
+import { useFilterLogic } from "./key_team_helpers/filter_logic";
+
 interface AllKeysTableProps {
   keys: KeyResponse[];
   isLoading?: boolean;
@@ -95,17 +97,28 @@ export function AllKeysTable({
   userID,
   userRole,
   organizations,
-  setCurrentOrg
+  setCurrentOrg,
 }: AllKeysTableProps) {
   const [selectedKeyId, setSelectedKeyId] = useState<string | null>(null);
-  const [filters, setFilters] = useState<{
-    'Team ID': string;
-    'Organization ID': string;
-  }>({
-    'Team ID': '',
-    'Organization ID': ''
-  });
   const [userList, setUserList] = useState<UserResponse[]>([]);
+  
+  // Use the filter logic hook
+  const {
+    filters,
+    filteredKeys,
+    allKeyAliases,
+    allTeams,
+    allOrganizations,
+    handleFilterChange,
+    handleFilterReset
+  } = useFilterLogic({
+    keys,
+    teams,
+    organizations,
+    accessToken,
+    setSelectedTeam,
+    setCurrentOrg
+  });
 
   useEffect(() => {
     if (accessToken) {
@@ -118,43 +131,6 @@ export function AllKeysTable({
     }
   }, [accessToken, keys]);
 
-  const handleFilterChange = (newFilters: Record<string, string>) => {
-    // Update filters state
-    setFilters({
-      'Team ID': newFilters['Team ID'] || '',
-      'Organization ID': newFilters['Organization ID'] || ''
-    });
-  
-    // Handle Team change
-    if (newFilters['Team ID']) {
-      const selectedTeamData = teams?.find(team => team.team_id === newFilters['Team ID']);
-      if (selectedTeamData) {
-        setSelectedTeam(selectedTeamData);
-      }
-    }
-  
-    // Handle Org change
-    if (newFilters['Organization ID']) {
-      const selectedOrg = organizations?.find(org => org.organization_id === newFilters['Organization ID']);
-      if (selectedOrg) {
-        setCurrentOrg(selectedOrg);
-      }
-    }
-  };
-
-  const handleFilterReset = () => {
-    // Reset filters state
-    setFilters({
-      'Team ID': '',
-      'Organization ID': ''
-    });
-    
-    // Reset team and org selections
-    setSelectedTeam(null); // or whatever your default value should be
-    setCurrentOrg(null); // or whatever your default value should be
-  };
-  
-
   const columns: ColumnDef<KeyResponse>[] = [
     {
       id: "expander",
@@ -205,7 +181,7 @@ export function AllKeysTable({
       accessorKey: "team_id", // Change to access the team_id
       cell: ({ row, getValue }) => {
         const teamId = getValue() as string;
-        const team = teams?.find(t => t.team_id === teamId);
+        const team = allTeams?.find(t => t.team_id === teamId);
         return team?.team_alias || "Unknown";
       },
     },
@@ -323,8 +299,42 @@ export function AllKeysTable({
   ];
 
   const filterOptions: FilterOption[] = [
-    { name: 'Team ID', label: 'Team ID', isSearchable: true, searchFn: createTeamSearchFunction(teams) },
-    { name: 'Organization ID', label: 'Organization ID', isSearchable: true, searchFn: createOrgSearchFunction(organizations) }
+    { 
+      name: 'Team ID', 
+      label: 'Team ID', 
+      isSearchable: true, 
+      searchFn: async (searchText: string) => {
+        if (!allTeams || allTeams.length === 0) return [];
+        
+        const filteredTeams = allTeams.filter(team => 
+          team.team_id.toLowerCase().includes(searchText.toLowerCase()) || 
+          (team.team_alias && team.team_alias.toLowerCase().includes(searchText.toLowerCase()))
+        );
+        
+        return filteredTeams.map(team => ({
+          label: `${team.team_alias || team.team_id} (${team.team_id})`,
+          value: team.team_id
+        }));
+      }
+    },
+    { 
+      name: 'Organization ID', 
+      label: 'Organization ID', 
+      isSearchable: true, 
+      searchFn: async (searchText: string) => {
+        if (!allOrganizations || allOrganizations.length === 0) return [];
+        
+        const filteredOrgs = allOrganizations.filter(org => 
+          org.organization_id.toLowerCase().includes(searchText.toLowerCase()) || 
+          (org.organization_name && org.organization_name.toLowerCase().includes(searchText.toLowerCase()))
+        );
+        
+        return filteredOrgs.map(org => ({
+          label: `${org.organization_name || 'Unknown'} (${org.organization_id})`,
+          value: org.organization_id
+        }));
+      }
+    },
   ];
   
   
@@ -338,7 +348,7 @@ export function AllKeysTable({
           accessToken={accessToken}
           userID={userID}
           userRole={userRole}
-          teams={teams}
+          teams={allTeams}
         />
       ) : (
         <div className="border-b py-4 flex-1 overflow-hidden">
@@ -376,7 +386,7 @@ export function AllKeysTable({
             
             <DataTable
               columns={columns.filter(col => col.id !== 'expander')}
-              data={keys}
+              data={filteredKeys}
               isLoading={isLoading}
               getRowCanExpand={() => false}
               renderSubComponent={() => <></>}
diff --git a/ui/litellm-dashboard/src/components/common_components/filter.tsx b/ui/litellm-dashboard/src/components/common_components/filter.tsx
index 7c3f3c2d524..cca4f70b001 100644
--- a/ui/litellm-dashboard/src/components/common_components/filter.tsx
+++ b/ui/litellm-dashboard/src/components/common_components/filter.tsx
@@ -62,9 +62,48 @@ const FilterComponent: React.FC<FilterComponentProps> = ({
     return () => document.removeEventListener('mousedown', handleClickOutside);
   }, []);
 
+  useEffect(() => {
+    if (options.length > 0 && options[0].isSearchable && options[0].searchFn) {
+      loadInitialOptions(options[0]);
+    }
+  }, []);
+
+  const loadInitialOptions = async (option: FilterOption) => {
+    if (!option.isSearchable || !option.searchFn) return;
+    
+    setSearchLoading(true);
+    try {
+      const results = await option.searchFn('');
+      setSearchOptions(results);
+    } catch (error) {
+      console.error('Error loading initial options:', error);
+      setSearchOptions([]);
+    } finally {
+      setSearchLoading(false);
+    }
+  };
+  
+  useEffect(() => {
+    if (showFilters && currentOption?.isSearchable && currentOption?.searchFn) {
+      loadInitialOptions(currentOption);
+    }
+  }, [showFilters, selectedFilter]);
+
+  const handleFilterSelect = (key: string) => {
+    setSelectedFilter(key);
+    setDropdownOpen(false);
+    
+    const newOption = options.find(opt => opt.name === key);
+    if (newOption?.isSearchable && newOption?.searchFn) {
+      loadInitialOptions(newOption);
+    } else {
+      setSearchOptions([]);
+    }
+  };
+
   const debouncedSearch = useCallback(
     debounce(async (value: string, option: FilterOption) => {
-      if (!value || !option.isSearchable || !option.searchFn) return;
+      if (!option.isSearchable || !option.searchFn) return;
       
       setSearchLoading(true);
       try {
@@ -139,7 +178,7 @@ const FilterComponent: React.FC<FilterComponentProps> = ({
         {buttonLabel}
       </TremorButton>
       {showFilters && (
-        <Card className="absolute left-0 mt-2 w-96 z-50 border border-gray-200 shadow-lg">
+        <Card className="absolute left-0 mt-2 w-[500px] z-50 border border-gray-200 shadow-lg">
           <div className="flex flex-col gap-4">
             <div className="flex items-center gap-2">
               <span className="text-sm font-medium">Where</span>
@@ -147,17 +186,14 @@ const FilterComponent: React.FC<FilterComponentProps> = ({
               <Dropdown
                 menu={{
                   items: dropdownItems,
-                  onClick: ({ key }) => {
-                    setSelectedFilter(key);
-                    setDropdownOpen(false);
-                    setSearchOptions([]);
-                  }
+                  onClick: ({ key }) => handleFilterSelect(key),
+                  style: { minWidth: '200px' }
                 }}
                 onOpenChange={setDropdownOpen}
                 open={dropdownOpen}
                 trigger={['click']}
               >
-                <Button className="min-w-32 text-left flex justify-between items-center">
+                <Button className="min-w-40 text-left flex justify-between items-center">
                   {currentOption?.label || selectedFilter}
                   {dropdownOpen ? (
                     <ChevronUpIcon className="h-4 w-4" />
@@ -185,7 +221,7 @@ const FilterComponent: React.FC<FilterComponentProps> = ({
                   }
                 }}
                 filterOption={false}
-                className="flex-1 w-full max-w-full truncate"
+                className="flex-1 w-full max-w-full truncate min-w-100"
                 loading={searchLoading}
                 options={searchOptions}
                 allowClear
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
new file mode 100644
index 00000000000..70a78ffb03f
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts
@@ -0,0 +1,127 @@
+import { keyListCall, teamListCall, organizationListCall } from '../networking';
+import { Team } from './key_list';
+import { Organization } from '../networking';
+
+/**
+ * Fetches all key aliases across all pages
+ * @param accessToken The access token for API authentication
+ * @returns Array of all unique key aliases
+ */
+export const fetchAllKeyAliases = async (accessToken: string | null): Promise<string[]> => {
+  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
+        currentPage,
+        100 // larger page size to reduce number of requests
+      );
+      
+      // Extract aliases from this page
+      const pageAliases = response.keys
+        .map(key => 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 [...new Set(allAliases)];
+  } 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
+ * @param organizationId Optional organization ID to filter teams
+ * @returns Array of all teams
+ */
+export const fetchAllTeams = async (accessToken: string | null, organizationId?: string | null): Promise<Team[]> => {
+  if (!accessToken) return [];
+  
+  try {
+    let allTeams: Team[] = [];
+    let currentPage = 1;
+    let hasMorePages = true;
+    
+    while (hasMorePages) {
+      const response = await teamListCall(
+        accessToken,
+        organizationId || null,
+        currentPage,
+        100 // larger page size to reduce number of requests
+      );
+      
+      // Add teams from this page
+      allTeams = [...allTeams, ...response.teams];
+      
+      // Check if there are more pages
+      if (currentPage < response.total_pages) {
+        currentPage++;
+      } else {
+        hasMorePages = false;
+      }
+    }
+    
+    return allTeams;
+  } catch (error) {
+    console.error("Error fetching all teams:", error);
+    return [];
+  }
+};
+
+/**
+ * Fetches all organizations across all pages
+ * @param accessToken The access token for API authentication
+ * @returns Array of all organizations
+ */
+export const fetchAllOrganizations = async (accessToken: string | null): Promise<Organization[]> => {
+  if (!accessToken) return [];
+  
+  try {
+    let allOrganizations: Organization[] = [];
+    let currentPage = 1;
+    let hasMorePages = true;
+    
+    while (hasMorePages) {
+      const response = await organizationListCall(
+        accessToken,
+        currentPage,
+        100 // larger page size to reduce number of requests
+      );
+      
+      // Add organizations from this page
+      allOrganizations = [...allOrganizations, ...response.organizations];
+      
+      // Check if there are more pages
+      if (currentPage < response.total_pages) {
+        currentPage++;
+      } else {
+        hasMorePages = false;
+      }
+    }
+    
+    return allOrganizations;
+  } catch (error) {
+    console.error("Error fetching all organizations:", error);
+    return [];
+  }
+};
diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/filter_logic.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/filter_logic.tsx
new file mode 100644
index 00000000000..6333f20b240
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/key_team_helpers/filter_logic.tsx
@@ -0,0 +1,148 @@
+import { useEffect, useState } from "react";
+import { KeyResponse } from "../key_team_helpers/key_list";
+import { Team, Organization } from "../networking";
+
+export interface FilterState {
+  'Team ID': string;
+  'Organization ID': string;
+  [key: string]: string;
+}
+
+export function useFilterLogic({
+  keys,
+  teams,
+  organizations,
+  accessToken,
+  setSelectedTeam,
+  setCurrentOrg
+}: {
+  keys: KeyResponse[];
+  teams: Team[] | null;
+  organizations: Organization[] | null;
+  accessToken: string | null;
+  setSelectedTeam: (team: Team | null) => void;
+  setCurrentOrg: React.Dispatch<React.SetStateAction<Organization | null>>;
+}) {
+  const [filters, setFilters] = useState<FilterState>({
+    'Team ID': '',
+    'Organization ID': '',
+  });
+  const [allKeyAliases, setAllKeyAliases] = useState<string[]>([]);
+  const [allTeams, setAllTeams] = useState<Team[]>(teams || []);
+  const [allOrganizations, setAllOrganizations] = useState<Organization[]>(organizations || []);
+  const [filteredKeys, setFilteredKeys] = useState<KeyResponse[]>(keys);
+
+  // Apply filters to keys whenever keys or filters change
+  useEffect(() => {
+    if (!keys) {
+      setFilteredKeys([]);
+      return;
+    }
+
+    let result = [...keys];
+
+    // Apply Team ID filter
+    if (filters['Team ID']) {
+      result = result.filter(key => key.team_id === filters['Team ID']);
+    }
+
+    // Apply Organization ID filter
+    if (filters['Organization ID']) {
+      result = result.filter(key => key.organization_id === filters['Organization ID']);
+    }
+
+
+    setFilteredKeys(result);
+  }, [keys, filters]);
+
+  // Fetch all data for filters when component mounts
+  useEffect(() => {
+    const loadAllFilterData = async () => {
+
+      
+      // Load all teams - no organization filter needed here
+      const teamsData = await fetchAllTeams(accessToken);
+      if (teamsData.length > 0) {
+        setAllTeams(teamsData);
+      }
+      
+      // Load all organizations
+      const orgsData = await fetchAllOrganizations(accessToken);
+      if (orgsData.length > 0) {
+        setAllOrganizations(orgsData);
+      }
+    };
+    
+    if (accessToken) {
+      loadAllFilterData();
+    }
+  }, [accessToken]);
+
+  // Update teams and organizations when props change
+  useEffect(() => {
+    if (teams && teams.length > 0) {
+      setAllTeams(prevTeams => {
+        // Only update if we don't already have a larger set of teams
+        return prevTeams.length < teams.length ? teams : prevTeams;
+      });
+    }
+  }, [teams]);
+
+  useEffect(() => {
+    if (organizations && organizations.length > 0) {
+      setAllOrganizations(prevOrgs => {
+        // Only update if we don't already have a larger set of organizations
+        return prevOrgs.length < organizations.length ? organizations : prevOrgs;
+      });
+    }
+  }, [organizations]);
+
+  const handleFilterChange = (newFilters: Record<string, string>) => {
+    // Update filters state
+    setFilters({
+      'Team ID': newFilters['Team ID'] || '',
+      'Organization ID': newFilters['Organization ID'] || '',
+    });
+  
+    // Handle Team change
+    if (newFilters['Team ID']) {
+      const selectedTeamData = allTeams?.find(team => team.team_id === newFilters['Team ID']);
+      if (selectedTeamData) {
+        setSelectedTeam(selectedTeamData);
+      }
+    }
+  
+    // Handle Org change
+    if (newFilters['Organization ID']) {
+      const selectedOrg = allOrganizations?.find(org => org.organization_id === newFilters['Organization ID']);
+      if (selectedOrg) {
+        setCurrentOrg(selectedOrg);
+      }
+    }
+  };
+
+  const handleFilterReset = () => {
+    // Reset filters state
+    setFilters({
+      'Team ID': '',
+      'Organization ID': '',
+    });
+    
+    // Reset team and org selections
+    setSelectedTeam(null);
+    setCurrentOrg(null);
+  };
+
+  return {
+    filters,
+    filteredKeys,
+    allKeyAliases,
+    allTeams,
+    allOrganizations,
+    handleFilterChange,
+    handleFilterReset
+  };
+}
+
+// These functions are imported from key_team_helpers/filter_helpers.ts
+import { fetchAllKeyAliases, fetchAllTeams, fetchAllOrganizations } from './filter_helpers'; 
\ No newline at end of file