From 35b5a89d82e9e9148f89a8cf043e5acd2149744d 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         | 139 +++-------------
 .../components/common_components/filter.tsx   |   9 +-
 .../key_team_helpers/filter_helpers.ts        |   2 +-
 .../key_team_helpers/filter_logic.tsx         | 148 ++++++++++++++++++
 4 files changed, 176 insertions(+), 122 deletions(-)
 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 844ecbc9e84..bc0d1135ea7 100644
--- a/ui/litellm-dashboard/src/components/all_keys_table.tsx
+++ b/ui/litellm-dashboard/src/components/all_keys_table.tsx
@@ -12,8 +12,7 @@ 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 { keyListCall } from "./networking";
-import { fetchAllKeyAliases, fetchAllTeams, fetchAllOrganizations } from './key_team_helpers/filter_helpers';
+import { useFilterLogic } from "./key_team_helpers/filter_logic";
 
 interface AllKeysTableProps {
   keys: KeyResponse[];
@@ -101,63 +100,25 @@ export function AllKeysTable({
   setCurrentOrg,
 }: AllKeysTableProps) {
   const [selectedKeyId, setSelectedKeyId] = useState<string | null>(null);
-  const [filters, setFilters] = useState<{
-    'Team ID': string;
-    'Organization ID': string;
-    'Key Alias': string;
-  }>({
-    'Team ID': '',
-    'Organization ID': '',
-    'Key Alias': ''
-  });
   const [userList, setUserList] = useState<UserResponse[]>([]);
-  const [allKeyAliases, setAllKeyAliases] = useState<string[]>([]);
-  const [allTeams, setAllTeams] = useState<Team[]>(teams || []);
-  const [allOrganizations, setAllOrganizations] = useState<Organization[]>(organizations || []);
-
-  // Fetch all data for filters when component mounts
-  useEffect(() => {
-    const loadAllFilterData = async () => {
-      // Load all key aliases
-      const aliases = await fetchAllKeyAliases(accessToken);
-      setAllKeyAliases(aliases);
-      
-      // 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]);
+  
+  // Use the filter logic hook
+  const {
+    filters,
+    filteredKeys,
+    allKeyAliases,
+    allTeams,
+    allOrganizations,
+    handleFilterChange,
+    handleFilterReset
+  } = useFilterLogic({
+    keys,
+    teams,
+    organizations,
+    accessToken,
+    setSelectedTeam,
+    setCurrentOrg
+  });
 
   useEffect(() => {
     if (accessToken) {
@@ -170,45 +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'] || '',
-      'Key Alias': newFilters['Key Alias'] || ''
-    });
-  
-    // 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': '',
-      'Key Alias': ''
-    });
-    
-    // Reset team and org selections
-    setSelectedTeam(null);
-    setCurrentOrg(null);
-  };
-  
-
   const columns: ColumnDef<KeyResponse>[] = [
     {
       id: "expander",
@@ -259,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";
       },
     },
@@ -413,23 +335,6 @@ export function AllKeysTable({
         }));
       }
     },
-    {
-      name: 'Key Alias',
-      label: 'Key Alias',
-      isSearchable: true,
-      searchFn: async (searchText: string) => {
-        if (!allKeyAliases || allKeyAliases.length === 0) return [];
-        
-        const filteredAliases = allKeyAliases.filter(alias => 
-          alias && alias.toLowerCase().includes(searchText.toLowerCase())
-        );
-        
-        return filteredAliases.map(alias => ({
-          label: alias,
-          value: alias
-        }));
-      }
-    }
   ];
   
   
@@ -443,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">
@@ -481,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 d68b7f98312..ebe6ab1eb74 100644
--- a/ui/litellm-dashboard/src/components/common_components/filter.tsx
+++ b/ui/litellm-dashboard/src/components/common_components/filter.tsx
@@ -178,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>
@@ -186,13 +186,14 @@ const FilterComponent: React.FC<FilterComponentProps> = ({
               <Dropdown
                 menu={{
                   items: dropdownItems,
-                  onClick: ({ key }) => handleFilterSelect(key)
+                  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" />
@@ -219,7 +220,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
index 94fc738a204..4639c9291fd 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
@@ -124,4 +124,4 @@ export const fetchAllOrganizations = async (accessToken: string | null): Promise
     console.error("Error fetching all organizations:", error);
     return [];
   }
-}; 
\ No newline at end of file
+}; 
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