From a93ef15b5531deb2a3026daca50498cdcc5d009e Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Fri, 14 Mar 2025 20:48:30 -0700
Subject: [PATCH] fix(key_edit_view.tsx): fix showing available models on key
 edit for non-team key

---
 .../src/components/key_edit_view.tsx          | 29 +++++++++++--------
 1 file changed, 17 insertions(+), 12 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/key_edit_view.tsx b/ui/litellm-dashboard/src/components/key_edit_view.tsx
index e6fcadf3fc0..bbe75e0d0eb 100644
--- a/ui/litellm-dashboard/src/components/key_edit_view.tsx
+++ b/ui/litellm-dashboard/src/components/key_edit_view.tsx
@@ -2,7 +2,7 @@ import React, { useState, useEffect } from "react";
 import { Form, Input, InputNumber, Select } from "antd";
 import { Button, TextInput } from "@tremor/react";
 import { KeyResponse } from "./key_team_helpers/key_list";
-import { getTeamModels } from "../components/create_key_button";
+import { fetchTeamModels } from "../components/create_key_button";
 import { modelAvailableCall } from "./networking";
 
 interface KeyEditViewProps {
@@ -45,31 +45,36 @@ export function KeyEditView({
   const [form] = Form.useForm();
   const [userModels, setUserModels] = useState<string[]>([]);
   const team = teams?.find(team => team.team_id === keyData.team_id);
-  const availableModels = getTeamModels(team, userModels);
-  
+  const [availableModels, setAvailableModels] = useState<string[]>([]);
 
   useEffect(() => {
-    const fetchUserModels = async () => {
+    const fetchModels = async () => {
+      if (!userID || !userRole || !accessToken) return;
+
       try {
-        if (accessToken && userID && userRole) {
+        if (keyData.team_id === null) {
+          // Fetch user models if no team
           const model_available = await modelAvailableCall(
             accessToken,
-            userID,
+            userID, 
             userRole
           );
-          let available_model_names = model_available["data"].map(
+          const available_model_names = model_available["data"].map(
             (element: { id: string }) => element.id
           );
-          console.log("available_model_names:", available_model_names);
-          setUserModels(available_model_names);
+          setAvailableModels(available_model_names);
+        } else if (team?.team_id) {
+          // Fetch team models if team exists
+          const models = await fetchTeamModels(userID, userRole, accessToken, team.team_id);
+          setAvailableModels(Array.from(new Set([...team.models, ...models])));
         }
       } catch (error) {
-        console.error("Error fetching user models:", error);
+        console.error("Error fetching models:", error);
       }
     };
 
-    fetchUserModels();
-  }, []);
+    fetchModels();
+  }, [userID, userRole, accessToken, team, keyData.team_id]);
 
   // Convert API budget duration to form format
   const getBudgetDuration = (duration: string | null) => {