From 72d7b268116a99eaf2028a501f2706a84b0d644e Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 3 Apr 2025 21:00:42 -0700
Subject: [PATCH] fix allow selecting endpoint on test key page

---
 .../src/components/chat_ui.tsx                | 50 ++++++++++------
 .../chat_ui/llm_calls/fetch_models.tsx        | 59 +++++++------------
 2 files changed, 53 insertions(+), 56 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/chat_ui.tsx b/ui/litellm-dashboard/src/components/chat_ui.tsx
index 547e926559b..bf483a790e2 100644
--- a/ui/litellm-dashboard/src/components/chat_ui.tsx
+++ b/ui/litellm-dashboard/src/components/chat_ui.tsx
@@ -25,7 +25,8 @@ import {
 import { message, Select } from "antd";
 import { makeOpenAIChatCompletionRequest } from "./chat_ui/llm_calls/chat_completion";
 import { makeOpenAIImageGenerationRequest } from "./chat_ui/llm_calls/image_generation";
-import { fetchAvailableModels } from "./chat_ui/llm_calls/fetch_models";
+import { fetchAvailableModels, ModelGroup  } from "./chat_ui/llm_calls/fetch_models";
+import { litellmModeMapping, ModelMode, EndpointType } from "./chat_ui/mode_endpoint_mapping";
 import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
 import { Typography } from "antd";
 import { coy } from 'react-syntax-highlighter/dist/esm/styles/prism';
@@ -55,7 +56,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
     undefined
   );
   const [showCustomModelInput, setShowCustomModelInput] = useState<boolean>(false);
-  const [modelInfo, setModelInfo] = useState<any[]>([]);
+  const [modelInfo, setModelInfo] = useState<ModelGroup[]>([]);
   const customModelTimeout = useRef<NodeJS.Timeout | null>(null);
   const [endpointType, setEndpointType] = useState<'chat' | 'image'>('chat');
 
@@ -74,15 +75,18 @@ const ChatUI: React.FC<ChatUIProps> = ({
       try {
         const uniqueModels = await fetchAvailableModels(
           useApiKey,
-          userID,
-          userRole
         );
   
         console.log("Fetched models:", uniqueModels);
   
         if (uniqueModels.length > 0) {
           setModelInfo(uniqueModels);
-          setSelectedModel(uniqueModels[0].value);
+          setSelectedModel(uniqueModels[0].model_group);
+          // Auto-set endpoint based on the first model's mode if available
+          const firstMode = uniqueModels[0].mode as ModelMode;
+          if (firstMode && litellmModeMapping[firstMode]) {
+            setEndpointType(litellmModeMapping[firstMode] as EndpointType);
+          }
         }
       } catch (error) {
         console.error("Error fetching model info:", error);
@@ -202,6 +206,14 @@ const ChatUI: React.FC<ChatUIProps> = ({
   const onModelChange = (value: string) => {
     console.log(`selected ${value}`);
     setSelectedModel(value);
+    // Look up the selected model to auto-select the endpoint type
+    const selectedOption = modelInfo.find((option) => option.model_group === value);
+    if (selectedOption && selectedOption.mode) {
+      const mode = selectedOption.mode as ModelMode;
+      if (litellmModeMapping[mode]) {
+        setEndpointType(litellmModeMapping[mode] as EndpointType);
+      }
+    }
     setShowCustomModelInput(value === 'custom');
   };
 
@@ -245,23 +257,15 @@ const ChatUI: React.FC<ChatUIProps> = ({
                       )}
                     </Col>
                     <Col className="mx-2">
-                      <Text>Endpoint Type:</Text>
-                      <Select
-                        defaultValue="chat"
-                        style={{ width: "350px", marginBottom: "12px" }}
-                        onChange={handleEndpointChange}
-                        options={[
-                          { value: 'chat', label: '/chat/completions' },
-                          { value: 'image', label: '/images/generations' }
-                        ]}
-                      />
-                      
                       <Text>Select Model:</Text>
                       <Select
                         placeholder="Select a Model"
                         onChange={onModelChange}
                         options={[
-                          ...modelInfo,
+                          ...modelInfo.map((option) => ({
+                            value: option.model_group,
+                            label: option.model_group
+                          })),
                           { value: 'custom', label: 'Enter custom model' }
                         ]}
                         style={{ width: "350px" }}
@@ -283,7 +287,19 @@ const ChatUI: React.FC<ChatUIProps> = ({
                           }}
                         />
                       )}
+                      <Text>Endpoint Type:</Text>
+                      <Select
+                        defaultValue="chat"
+                        style={{ width: "350px", marginBottom: "12px" }}
+                        onChange={handleEndpointChange}
+                        options={[
+                          { value: 'chat', label: '/chat/completions' },
+                          { value: 'image', label: '/images/generations' }
+                        ]}
+                      />
+                      
                     </Col>
+
                   </Grid>
 
                   {/* Clear Chat Button */}
diff --git a/ui/litellm-dashboard/src/components/chat_ui/llm_calls/fetch_models.tsx b/ui/litellm-dashboard/src/components/chat_ui/llm_calls/fetch_models.tsx
index 579af1add2e..4c3e679ce16 100644
--- a/ui/litellm-dashboard/src/components/chat_ui/llm_calls/fetch_models.tsx
+++ b/ui/litellm-dashboard/src/components/chat_ui/llm_calls/fetch_models.tsx
@@ -1,54 +1,35 @@
-import { modelAvailableCall } from "../../networking";
+// fetch_models.ts
 
-interface ModelOption {
-  value: string;
-  label: string;
+import { modelHubCall } from "../../networking";
+
+export interface ModelGroup {
+  model_group: string;
+  mode?: string;
 }
 
 /**
- * Fetches available models for the user and formats them as options
- * for selection dropdowns
+ * Fetches available models using modelHubCall and formats them for the selection dropdown.
  */
 export const fetchAvailableModels = async (
-  apiKey: string | null,
-  userID: string,
-  userRole: string,
-  teamID: string | null = null
-): Promise<ModelOption[]> => {
+  accessToken: string
+): Promise<ModelGroup[]> => {
   try {
-    const fetchedAvailableModels = await modelAvailableCall(
-      apiKey ?? '', // Use empty string if apiKey is null
-      userID,
-      userRole,
-      false,
-      teamID
-    );
+    const fetchedModels = await modelHubCall(accessToken);
+    console.log("model_info:", fetchedModels);
 
-    console.log("model_info:", fetchedAvailableModels);
+    if (fetchedModels?.data.length > 0) {
+      const models: ModelGroup[] = fetchedModels.data.map((item: any) => ({
+        model_group: item.model_group, // Display the model_group to the user
+        mode: item?.model_info?.mode, // Save the mode for auto-selection of endpoint type
+      }));
 
-    if (fetchedAvailableModels?.data.length > 0) {
-      // Create a Map to store unique models using the model ID as key
-      const uniqueModelsMap = new Map();
-      
-      fetchedAvailableModels["data"].forEach((item: { id: string }) => {
-        uniqueModelsMap.set(item.id, {
-          value: item.id,
-          label: item.id
-        });
-      });
-
-      // Convert Map values back to array
-      const uniqueModels = Array.from(uniqueModelsMap.values());
-
-      // Sort models alphabetically
-      uniqueModels.sort((a, b) => a.label.localeCompare(b.label));
-
-      return uniqueModels;
+      // Sort models alphabetically by label
+      models.sort((a, b) => a.model_group.localeCompare(b.model_group));
+      return models;
     }
-    
     return [];
   } catch (error) {
     console.error("Error fetching model info:", error);
     throw error;
   }
-}; 
\ No newline at end of file
+};