From eb24398ee17c99149d45916df397a32fe8815f0b Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 3 Apr 2024 20:29:25 -0700
Subject: [PATCH] (ui) fix bug when editing

---
 .../src/components/user_dashboard.tsx         |  2 +-
 .../src/components/view_key_table.tsx         | 39 ++++++++++++++-----
 2 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx
index 8f9977e08..eb13faaae 100644
--- a/ui/litellm-dashboard/src/components/user_dashboard.tsx
+++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx
@@ -217,10 +217,10 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
 
           <ViewKeyTable
             userID={userID}
+            userRole={userRole}
             accessToken={accessToken}
             selectedTeam={selectedTeam ? selectedTeam : null}
             data={keys}
-            userModels={userModels}
             setData={setKeys}
           />
           <CreateKey
diff --git a/ui/litellm-dashboard/src/components/view_key_table.tsx b/ui/litellm-dashboard/src/components/view_key_table.tsx
index c1e9f0d56..a4b49bd48 100644
--- a/ui/litellm-dashboard/src/components/view_key_table.tsx
+++ b/ui/litellm-dashboard/src/components/view_key_table.tsx
@@ -1,6 +1,6 @@
 "use client";
 import React, { useEffect, useState } from "react";
-import { keyDeleteCall } from "./networking";
+import { keyDeleteCall, modelAvailableCall } from "./networking";
 import { InformationCircleIcon, StatusOnlineIcon, TrashIcon, PencilAltIcon } from "@heroicons/react/outline";
 import { keySpendLogsCall, PredictedSpendLogsCall, keyUpdateCall } from "./networking";
 import {
@@ -47,9 +47,9 @@ interface EditKeyModalProps {
 // Define the props type
 interface ViewKeyTableProps {
   userID: string;
+  userRole: string | null;
   accessToken: string;
   selectedTeam: any | null;
-  userModels: string[];
   data: any[] | null;
   setData: React.Dispatch<React.SetStateAction<any[] | null>>;
 }
@@ -73,9 +73,9 @@ interface ItemData {
 
 const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
   userID,
+  userRole,
   accessToken,
   selectedTeam,
-  userModels,
   data,
   setData,
 }) => {
@@ -91,18 +91,39 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
 
   const [editModalVisible, setEditModalVisible] = useState(false);
   const [selectedToken, setSelectedToken] = useState<ItemData | null>(null);
+  const [userModels, setUserModels] = useState([]);
+
+  useEffect(() => {
+    const fetchUserModels = async () => {
+      try {
+        if (userID === null) {
+          return;
+        }
+
+        if (accessToken !== null && userRole !== null) {
+          const model_available = await modelAvailableCall(accessToken, userID, userRole);
+          let available_model_names = model_available["data"].map(
+            (element: { id: string }) => element.id
+          );
+          console.log("available_model_names:", available_model_names);
+          setUserModels(available_model_names);
+        }
+      } catch (error) {
+        console.error("Error fetching user models:", error);
+      }
+    };
+  
+    fetchUserModels();
+  }, [accessToken, userID, userRole]);
 
   const EditKeyModal: React.FC<EditKeyModalProps> = ({ visible, onCancel, token, onSubmit }) => {
     const [form] = Form.useForm();
 
-    // check token.models length == 0
-    if (token.models.length == 0 && selectedTeam) {
-      token.models = selectedTeam.models;
-      
-    }
+    console.log("in edit key modal:", token);
+    console.log("in edit key modal, team:", selectedTeam);
 
+    
 
-  
     const handleOk = () => {
       form
         .validateFields()