ui new chat ui endpoints

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

@@ -20,18 +20,28 @@ import {
   SelectItem,
   TextInput,
   Button,
+  Divider,
 } from "@tremor/react";
 
-import { message, Select } from "antd";
+import { message, Select, Spin, Typography, Tooltip } from "antd";
 import { makeOpenAIChatCompletionRequest } from "./chat_ui/llm_calls/chat_completion";
 import { makeOpenAIImageGenerationRequest } from "./chat_ui/llm_calls/image_generation";
 import { fetchAvailableModels, ModelGroup  } from "./chat_ui/llm_calls/fetch_models";
 import { litellmModeMapping, ModelMode, EndpointType, getEndpointType } 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';
 import EndpointSelector from "./chat_ui/EndpointSelector";
 import { determineEndpointType } from "./chat_ui/EndpointUtils";
+import { 
+  SendOutlined, 
+  ApiOutlined, 
+  KeyOutlined, 
+  ClearOutlined, 
+  RobotOutlined, 
+  UserOutlined,
+  DeleteOutlined,
+  LoadingOutlined
+} from "@ant-design/icons";
 
 interface ChatUIProps {
   accessToken: string | null;
@@ -61,6 +71,8 @@ const ChatUI: React.FC<ChatUIProps> = ({
   const [modelInfo, setModelInfo] = useState<ModelGroup[]>([]);
   const customModelTimeout = useRef<NodeJS.Timeout | null>(null);
   const [endpointType, setEndpointType] = useState<string>(EndpointType.CHAT);
+  const [isLoading, setIsLoading] = useState<boolean>(false);
+  const abortControllerRef = useRef<AbortController | null>(null);
 
   const chatEndRef = useRef<HTMLDivElement>(null);
 
@@ -141,6 +153,15 @@ const ChatUI: React.FC<ChatUIProps> = ({
     }
   };
 
+  const handleCancelRequest = () => {
+    if (abortControllerRef.current) {
+      abortControllerRef.current.abort();
+      abortControllerRef.current = null;
+      setIsLoading(false);
+      message.info("Request cancelled");
+    }
+  };
+
   const handleSendMessage = async () => {
     if (inputMessage.trim() === "") return;
 
@@ -155,11 +176,16 @@ const ChatUI: React.FC<ChatUIProps> = ({
       return;
     }
 
+    // Create new abort controller for this request
+    abortControllerRef.current = new AbortController();
+    const signal = abortControllerRef.current.signal;
+
     // Create message object without model field for API call
     const newUserMessage = { role: "user", content: inputMessage };
     
     // Update UI with full message object
     setChatHistory([...chatHistory, newUserMessage]);
+    setIsLoading(true);
 
     try {
       if (selectedModel) {
@@ -172,7 +198,8 @@ const ChatUI: React.FC<ChatUIProps> = ({
             apiChatHistory,
             (chunk, model) => updateTextUI("assistant", chunk, model),
             selectedModel,
-            effectiveApiKey
+            effectiveApiKey,
+            signal
           );
         } else if (endpointType === EndpointType.IMAGE) {
           // For image generation
@@ -180,13 +207,21 @@ const ChatUI: React.FC<ChatUIProps> = ({
             inputMessage,
             (imageUrl, model) => updateImageUI(imageUrl, model),
             selectedModel,
-            effectiveApiKey
+            effectiveApiKey,
+            signal
           );
         }
       }
     } catch (error) {
-      console.error("Error fetching response", error);
+      if (signal.aborted) {
-      updateTextUI("assistant", "Error fetching response");
+        console.log("Request was cancelled");
+      } else {
+        console.error("Error fetching response", error);
+        updateTextUI("assistant", "Error fetching response");
+      }
+    } finally {
+      setIsLoading(false);
+      abortControllerRef.current = null;
     }
 
     setInputMessage("");
@@ -224,128 +259,144 @@ const ChatUI: React.FC<ChatUIProps> = ({
     setEndpointType(value);
   };
 
+  const antIcon = <LoadingOutlined style={{ fontSize: 24 }} spin />;
+
   return (
-    <div style={{ width: "100%", position: "relative" }}>
+    <div className="w-full h-screen p-4 bg-white">
-      <div className="flex h-[80vh] w-full mt-2">
+    <Card className="w-full rounded-xl shadow-md overflow-hidden">
+      <div className="flex h-[80vh] w-full">
         {/* Left Sidebar with Controls */}
-        <div className="w-1/4 p-4 border-r">
+        <div className="w-1/4 p-4 border-r border-gray-200 bg-gray-50">
           <div className="mb-6">
-            <Text className="font-medium block mb-2">API Key Source</Text>
+            <div className="space-y-6">
-            <Select
+              <div>
-              disabled={disabledPersonalKeyCreation}
+                <Text className="font-medium block mb-2 text-gray-700 flex items-center">
-              defaultValue="session"
+                  <KeyOutlined className="mr-2" /> API Key Source
-              style={{ width: "100%" }}
+                </Text>
-              onChange={(value) => setApiKeySource(value as "session" | "custom")}
+                <Select
-              options={[
+                  disabled={disabledPersonalKeyCreation}
-                { value: 'session', label: 'Current UI Session' },
+                  defaultValue="session"
-                { value: 'custom', label: 'Virtual Key' },
+                  style={{ width: "100%" }}
-              ]}
+                  onChange={(value) => setApiKeySource(value as "session" | "custom")}
-            />
+                  options={[
-            {apiKeySource === 'custom' && (
+                    { value: 'session', label: 'Current UI Session' },
-              <TextInput
+                    { value: 'custom', label: 'Virtual Key' },
-                className="mt-2"
+                  ]}
-                placeholder="Enter custom API key"
+                  className="rounded-md"
-                type="password"
+                />
-                onValueChange={setApiKey}
+                {apiKeySource === 'custom' && (
-                value={apiKey}
+                  <TextInput
-              />
+                    className="mt-2"
-            )}
+                    placeholder="Enter custom API key"
+                    type="password"
+                    onValueChange={setApiKey}
+                    value={apiKey}
+                    icon={KeyOutlined}
+                  />
+                )}
+              </div>
+              
+              <div>
+                <Text className="font-medium block mb-2 text-gray-700 flex items-center">
+                  <RobotOutlined className="mr-2" /> Select Model
+                </Text>
+                <Select
+                  placeholder="Select a Model"
+                  onChange={onModelChange}
+                  options={[
+                    ...modelInfo.map((option) => ({
+                      value: option.model_group,
+                      label: option.model_group
+                    })),
+                    { value: 'custom', label: 'Enter custom model' }
+                  ]}
+                  style={{ width: "100%" }}
+                  showSearch={true}
+                  className="rounded-md"
+                />
+                {showCustomModelInput && (
+                  <TextInput
+                    className="mt-2"
+                    placeholder="Enter custom model name"
+                    onValueChange={(value) => {
+                      // Using setTimeout to create a simple debounce effect
+                      if (customModelTimeout.current) {
+                        clearTimeout(customModelTimeout.current);
+                      }
+                      
+                      customModelTimeout.current = setTimeout(() => {
+                        setSelectedModel(value);
+                      }, 500); // 500ms delay after typing stops
+                    }}
+                  />
+                )}
+              </div>
+              
+              <div>
+                <Text className="font-medium block mb-2 text-gray-700 flex items-center">
+                  <ApiOutlined className="mr-2" /> Endpoint Type
+                </Text>
+                <EndpointSelector 
+                  endpointType={endpointType}
+                  onEndpointChange={handleEndpointChange}
+                  className="mb-4"
+                />
+              </div>
+              
+              <Button
+                onClick={clearChatHistory}
+                className="w-full bg-gray-100 hover:bg-gray-200 text-gray-700 border-gray-300 mt-4"
+                icon={ClearOutlined}
+              >
+                Clear Chat
+              </Button>
+            </div>
           </div>
-          
-          <div className="mb-6">
-            <Text className="font-medium block mb-2">Select Model:</Text>
-            <Select
-              placeholder="Select a Model"
-              onChange={onModelChange}
-              options={[
-                ...modelInfo.map((option) => ({
-                  value: option.model_group,
-                  label: option.model_group
-                })),
-                { value: 'custom', label: 'Enter custom model' }
-              ]}
-              style={{ width: "100%" }}
-              showSearch={true}
-            />
-            {showCustomModelInput && (
-              <TextInput
-                className="mt-2"
-                placeholder="Enter custom model name"
-                onValueChange={(value) => {
-                  // Using setTimeout to create a simple debounce effect
-                  if (customModelTimeout.current) {
-                    clearTimeout(customModelTimeout.current);
-                  }
-                  
-                  customModelTimeout.current = setTimeout(() => {
-                    setSelectedModel(value);
-                  }, 500); // 500ms delay after typing stops
-                }}
-              />
-            )}
-          </div>
-          
-          <div className="mb-6">
-            <EndpointSelector 
-              endpointType={endpointType}
-              onEndpointChange={handleEndpointChange}
-            />
-          </div>
-          
-          <Button
-            onClick={clearChatHistory}
-            className="w-full"
-          >
-            Clear Chat
-          </Button>
         </div>
         
         {/* Main Chat Area */}
-        <div className="w-3/4 flex flex-col">
+        <div className="w-3/4 flex flex-col bg-white">
-          <div className="flex-1 overflow-auto p-4">
+          <div className="flex-1 overflow-auto p-4 pb-0">
+            {chatHistory.length === 0 && (
+              <div className="h-full flex flex-col items-center justify-center text-gray-400">
+                <RobotOutlined style={{ fontSize: '48px', marginBottom: '16px' }} />
+                <Text>Start a conversation or generate an image</Text>
+              </div>
+            )}
+            
             {chatHistory.map((message, index) => (
               <div 
                 key={index} 
                 className={`mb-4 ${message.role === "user" ? "text-right" : "text-left"}`}
               >
-                <div style={{ 
+                <div className="inline-block max-w-[80%] rounded-lg shadow-sm p-3.5 px-4" style={{
-                  display: 'inline-block',
+                  backgroundColor: message.role === "user" ? '#f0f8ff' : '#ffffff',
-                  maxWidth: '80%',
+                  border: message.role === "user" ? '1px solid #e6f0fa' : '1px solid #f0f0f0',
-                  borderRadius: '8px',
-                  padding: '10px 14px',
-                  backgroundColor: message.role === "user" ? '#f0f0f0' : '#f9f9f9',
                   textAlign: 'left'
                 }}>
-                  <div style={{ 
+                  <div className="flex items-center gap-2 mb-1.5">
-                    display: 'flex',
+                    <div className="flex items-center justify-center w-6 h-6 rounded-full mr-1" style={{
-                    alignItems: 'center',
+                      backgroundColor: message.role === "user" ? '#e6f0fa' : '#f5f5f5',
-                    gap: '8px',
+                    }}>
-                    marginBottom: '4px'
+                      {message.role === "user" ? 
-                  }}>
+                        <UserOutlined style={{ fontSize: '12px', color: '#2563eb' }} /> : 
-                    <strong>{message.role}</strong>
+                        <RobotOutlined style={{ fontSize: '12px', color: '#4b5563' }} />
+                      }
+                    </div>
+                    <strong className="text-sm capitalize">{message.role}</strong>
                     {message.role === "assistant" && message.model && (
-                      <span style={{
+                      <span className="text-xs px-2 py-0.5 rounded bg-gray-100 text-gray-600 font-normal">
-                        fontSize: '12px',
-                        color: '#666',
-                        backgroundColor: '#f5f5f5',
-                        padding: '2px 6px',
-                        borderRadius: '4px',
-                        fontWeight: 'normal'
-                      }}>
                         {message.model}
                       </span>
                     )}
                   </div>
-                  <div style={{ 
+                  <div className="whitespace-pre-wrap break-words max-w-full message-content">
-                    whiteSpace: "pre-wrap", 
-                    wordBreak: "break-word",
-                    maxWidth: "100%"
-                  }}>
                     {message.isImage ? (
                       <img 
                         src={message.content} 
                         alt="Generated image" 
-                        style={{ maxWidth: '100%', maxHeight: '500px' }} 
+                        className="max-w-full rounded-md border border-gray-200 shadow-sm" 
+                        style={{ maxHeight: '500px' }} 
                       />
                     ) : (
                       <ReactMarkdown
@@ -360,12 +411,13 @@ const ChatUI: React.FC<ChatUIProps> = ({
                                 style={coy as any}
                                 language={match[1]}
                                 PreTag="div"
+                                className="rounded-md my-2"
                                 {...props}
                               >
                                 {String(children).replace(/\n$/, '')}
                               </SyntaxHighlighter>
                             ) : (
-                              <code className={className} {...props}>
+                              <code className={`${className} px-1.5 py-0.5 rounded bg-gray-100 text-sm font-mono`} {...props}>
                                 {children}
                               </code>
                             );
@@ -379,11 +431,16 @@ const ChatUI: React.FC<ChatUIProps> = ({
                 </div>
               </div>
             ))}
+            {isLoading && (
+              <div className="flex justify-center items-center my-4">
+                <Spin indicator={antIcon} />
+              </div>
+            )}
             <div ref={chatEndRef} style={{ height: "1px" }} />
           </div>
           
-          <div className="p-4 border-t">
+          <div className="p-4 border-t border-gray-200 bg-white">
-            <div className="flex">
+            <div className="flex items-center">
               <TextInput
                 type="text"
                 value={inputMessage}
@@ -394,17 +451,31 @@ const ChatUI: React.FC<ChatUIProps> = ({
                     ? "Type your message..." 
                     : "Describe the image you want to generate..."
                 }
+                disabled={isLoading}
+                className="flex-1"
               />
-              <Button
+              {isLoading ? (
-                onClick={handleSendMessage}
+                <Button
-                className="ml-2"
+                  onClick={handleCancelRequest}
-              >
+                  className="ml-2 bg-red-50 hover:bg-red-100 text-red-600 border-red-200"
-                {endpointType === EndpointType.CHAT ? "Send" : "Generate"}
+                  icon={DeleteOutlined}
-              </Button>
+                >
+                  Cancel
+                </Button>
+              ) : (
+                <Button
+                  onClick={handleSendMessage}
+                  className="ml-2 text-white"
+                  icon={endpointType === EndpointType.CHAT ? SendOutlined : RobotOutlined}
+                >
+                  {endpointType === EndpointType.CHAT ? "Send" : "Generate"}
+                </Button>
+              )}
             </div>
           </div>
         </div>
       </div>
+    </Card>
     </div>
   );
 };

ui/litellm-dashboard/src/components/chat_ui/EndpointSelector.tsx

@@ -28,9 +28,10 @@ const EndpointSelector: React.FC<EndpointSelectorProps> = ({
       <Text>Endpoint Type:</Text>
       <Select
         value={endpointType}
-        style={{ width: "100%", marginBottom: "12px" }}
+        style={{ width: "100%" }}
         onChange={onEndpointChange}
         options={endpointOptions}
+        className="rounded-md"
       />
     </div>
   );

ui/litellm-dashboard/src/components/chat_ui/llm_calls/chat_completion.tsx

@@ -1,4 +1,3 @@
-
 import openai from "openai";
 import { ChatCompletionMessageParam } from "openai/resources/chat/completions";
 import { message } from "antd";
@@ -7,7 +6,8 @@ export async function makeOpenAIChatCompletionRequest(
     chatHistory: { role: string; content: string }[],
     updateUI: (chunk: string, model: string) => void,
     selectedModel: string,
-    accessToken: string
+    accessToken: string,
+    signal?: AbortSignal
   ) {
     // base url should be the current base_url
     const isLocal = process.env.NODE_ENV === "development";
@@ -29,7 +29,7 @@ export async function makeOpenAIChatCompletionRequest(
         model: selectedModel,
         stream: true,
         messages: chatHistory as ChatCompletionMessageParam[],
-      });
+      }, { signal });
   
       for await (const chunk of response) {
         console.log(chunk);
@@ -38,7 +38,12 @@ export async function makeOpenAIChatCompletionRequest(
         }
       }
     } catch (error) {
-      message.error(`Error occurred while generating model response. Please try again. Error: ${error}`, 20);
+      if (signal?.aborted) {
+        console.log("Chat completion request was cancelled");
+      } else {
+        message.error(`Error occurred while generating model response. Please try again. Error: ${error}`, 20);
+      }
+      throw error; // Re-throw to allow the caller to handle the error
     }
   }
   

ui/litellm-dashboard/src/components/chat_ui/llm_calls/image_generation.tsx

@@ -5,7 +5,8 @@ export async function makeOpenAIImageGenerationRequest(
   prompt: string,
   updateUI: (imageUrl: string, model: string) => void,
   selectedModel: string,
-  accessToken: string
+  accessToken: string,
+  signal?: AbortSignal
 ) {
   // base url should be the current base_url
   const isLocal = process.env.NODE_ENV === "development";
@@ -26,7 +27,7 @@ export async function makeOpenAIImageGenerationRequest(
     const response = await client.images.generate({
       model: selectedModel,
       prompt: prompt,
-    });
+    }, { signal });
 
     console.log(response.data);
     
@@ -46,6 +47,11 @@ export async function makeOpenAIImageGenerationRequest(
       throw new Error("Invalid response format");
     }
   } catch (error) {
-    message.error(`Error occurred while generating image. Please try again. Error: ${error}`, 20);
+    if (signal?.aborted) {
+      console.log("Image generation request was cancelled");
+    } else {
+      message.error(`Error occurred while generating image. Please try again. Error: ${error}`, 20);
+    }
+    throw error; // Re-throw to allow the caller to handle the error
   }
 }