Reflect key and team update in UI (BerriAI#9825)

* Reflect updates to keys in UI instantly

* Reflect updates to teams in UI instantly

Author: crisshaker

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

@@ -13,9 +13,12 @@ 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";
+import { Setter } from "@/types";
+import { updateExistingKeys } from "@/utils/dataUtils";
 
 interface AllKeysTableProps {
   keys: KeyResponse[];
+  setKeys: Setter<KeyResponse[]>;
   isLoading?: boolean;
   pagination: {
     currentPage: number;
@@ -87,6 +90,7 @@ const TeamFilter = ({
  */
 export function AllKeysTable({ 
   keys, 
+  setKeys,
   isLoading = false,
   pagination,
   onPageChange,
@@ -364,6 +368,23 @@ export function AllKeysTable({
           keyId={selectedKeyId} 
           onClose={() => setSelectedKeyId(null)}
           keyData={keys.find(k => k.token === selectedKeyId)}
+          onKeyDataUpdate={(updatedKeyData) => {
+            setKeys(keys => keys.map(key => {
+              if (key.token === updatedKeyData.token) {
+                // The shape of key is different from that of
+                // updatedKeyData(received from keyUpdateCall in networking.tsx).
+                // Hence, we can't replace key with updatedKeys since it might lead
+                // to unintended bugs/behaviors.
+                // So instead, we only update fields that are present in both.
+                return updateExistingKeys(key, updatedKeyData)
+              }
+              
+              return key
+            }))
+          }}
+          onDelete={() => {
+            setKeys(keys => keys.filter(key => key.token !== selectedKeyId))
+          }}
           accessToken={accessToken}
           userID={userID}
           userRole={userRole}

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

@@ -27,13 +27,15 @@ interface KeyInfoViewProps {
   keyId: string;
   onClose: () => void;
   keyData: KeyResponse | undefined;
+  onKeyDataUpdate?: (data: Partial<KeyResponse>) => void;
+  onDelete?: () => void;
   accessToken: string | null;
   userID: string | null;
   userRole: string | null;
   teams: any[] | null;
 }
 
-export default function KeyInfoView({ keyId, onClose, keyData, accessToken, userID, userRole, teams }: KeyInfoViewProps) {
+export default function KeyInfoView({ keyId, onClose, keyData, accessToken, userID, userRole, teams, onKeyDataUpdate, onDelete }: KeyInfoViewProps) {
   const [isEditing, setIsEditing] = useState(false);
   const [form] = Form.useForm();
   const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false);
@@ -93,6 +95,9 @@ export default function KeyInfoView({ keyId, onClose, keyData, accessToken, user
       }
 
       const newKeyValues = await keyUpdateCall(accessToken, formValues);
+      if (onKeyDataUpdate) {
+        onKeyDataUpdate(newKeyValues)
+      }
       message.success("Key updated successfully");
       setIsEditing(false);
       // Refresh key data here if needed
@@ -107,6 +112,9 @@ export default function KeyInfoView({ keyId, onClose, keyData, accessToken, user
       if (!accessToken) return;
       await keyDeleteCall(accessToken as string, keyData.token);
       message.success("Key deleted successfully");
+      if (onDelete) {
+        onDelete()
+      }
       onClose();
     } catch (error) {
       console.error("Error deleting the key:", error);

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

@@ -1,5 +1,6 @@
 import { useState, useEffect } from 'react';
 import { keyListCall, Organization } from '../networking';
+import { Setter } from '@/types';
 
 export interface Team {
     team_id: string;
@@ -94,13 +95,14 @@ totalPages: number;
 totalCount: number;
 }
 
+
 interface UseKeyListReturn {
 keys: KeyResponse[];
 isLoading: boolean;
 error: Error | null;
 pagination: PaginationData;
 refresh: (params?: Record<string, unknown>) => Promise<void>;
-setKeys: (newKeysOrUpdater: KeyResponse[] | ((prevKeys: KeyResponse[]) => KeyResponse[])) => void;
+setKeys: Setter<KeyResponse[]>;
 }
 
 const useKeyList = ({

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

@@ -4,6 +4,7 @@
 import { all_admin_roles } from "@/utils/roles";
 import { message } from "antd";
 import { TagNewRequest, TagUpdateRequest, TagDeleteRequest, TagInfoRequest, TagListResponse, TagInfoResponse } from "./tag_management/types";
+import { Team } from "./key_team_helpers/key_list";
 
 const isLocal = process.env.NODE_ENV === "development";
 export const proxyBaseUrl = isLocal ? "http://localhost:4000" : null;
@@ -2983,7 +2984,7 @@ export const teamUpdateCall = async (
       console.error("Error response from the server:", errorData);
       throw new Error("Network response was not ok");
     }
-    const data = await response.json();
+    const data = await response.json() as { data: Team, team_id: string };
     console.log("Update Team Response:", data);
     return data;
     // Handle success - you might want to update some state or UI based on the created key

ui/litellm-dashboard/src/components/team/team_info.tsx

@@ -30,6 +30,7 @@ import { PencilAltIcon, PlusIcon, TrashIcon } from "@heroicons/react/outline";
 import MemberModal from "./edit_membership";
 import UserSearchModal from "@/components/common_components/user_search_modal";
 import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key";
+import { Team } from "../key_team_helpers/key_list";
 
 
 interface TeamData {
@@ -69,6 +70,7 @@ interface TeamInfoProps {
   is_proxy_admin: boolean;
   userModels: string[];
   editTeam: boolean;
+  onUpdate?: (team: Team) => void
 }
 
 const TeamInfoView: React.FC<TeamInfoProps> = ({ 
@@ -78,7 +80,8 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
   is_team_admin, 
   is_proxy_admin,
   userModels,
-  editTeam
+  editTeam,
+  onUpdate
 }) => {
   const [teamData, setTeamData] = useState<TeamData | null>(null);
   const [loading, setLoading] = useState(true);
@@ -199,7 +202,10 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
       };
 
       const response = await teamUpdateCall(accessToken, updateData);
-      
+      if (onUpdate) {
+        onUpdate(response.data)
+      }
+    
       message.success("Team settings updated successfully");
       setIsEditing(false);
       fetchTeamInfo();

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

@@ -84,6 +84,7 @@ import {
   modelAvailableCall,
   teamListCall
 } from "./networking";
+import { updateExistingKeys } from "@/utils/dataUtils";
 
 const getOrganizationModels = (organization: Organization | null, userModels: string[]) => {
   let tempModelsToPick = [];
@@ -321,6 +322,22 @@ const Teams: React.FC<TeamProps> = ({
       {selectedTeamId ? (
         <TeamInfoView 
         teamId={selectedTeamId} 
+        onUpdate={data => {
+            setTeams(teams => {
+              if (teams == null) {
+                return teams;
+              }
+            
+              return teams.map(team => {
+                if (data.team_id === team.team_id) {
+                  return updateExistingKeys(team, data)
+                }
+                
+                return team
+              })
+            })
+
+        }}
         onClose={() => {
           setSelectedTeamId(null);
           setEditTeam(false);

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

@@ -418,6 +418,7 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
     <div>
       <AllKeysTable 
         keys={keys}
+        setKeys={setKeys}
         isLoading={isLoading}
         pagination={pagination}
         onPageChange={handlePageChange}

ui/litellm-dashboard/src/types.ts

@@ -0,0 +1 @@
+export type Setter<T> = (newValueOrUpdater: T | ((previousValue: T) => T)) => void

ui/litellm-dashboard/src/utils/dataUtils.ts

@@ -0,0 +1,14 @@
+export function updateExistingKeys<Source extends Object>(
+  target: Source,
+  source: Object
+): Source {
+  const clonedTarget = structuredClone(target);
+  
+  for (const [key, value] of Object.entries(source)) {
+    if (key in clonedTarget) {
+      (clonedTarget as any)[key] = value;
+    }
+  }
+
+  return clonedTarget;
+}