From 852a1fee89754d5f095e2a6e4915bfbc58a1eda3 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Fri, 5 Dec 2025 15:51:56 -0800
Subject: [PATCH] Support images in compare UI

---
 .../playground/compareUI/CompareUI.test.tsx   |  69 ++++++++-
 .../playground/compareUI/CompareUI.tsx        | 136 ++++++++++++++----
 .../components/MessageDisplay.test.tsx        |  29 ++++
 .../compareUI/components/MessageDisplay.tsx   |   4 +-
 .../components/MessageInput.test.tsx          |  12 ++
 .../compareUI/components/MessageInput.tsx     |  11 +-
 6 files changed, 230 insertions(+), 31 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/CompareUI.test.tsx b/ui/litellm-dashboard/src/components/playground/compareUI/CompareUI.test.tsx
index 1941943a58..963b752774 100644
--- a/ui/litellm-dashboard/src/components/playground/compareUI/CompareUI.test.tsx
+++ b/ui/litellm-dashboard/src/components/playground/compareUI/CompareUI.test.tsx
@@ -2,6 +2,7 @@ import { render, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import CompareUI from "./CompareUI";
+import { makeOpenAIChatCompletionRequest } from "../llm_calls/chat_completion";
 
 vi.mock("../llm_calls/fetch_models", () => ({
   fetchAvailableModels: vi.fn().mockResolvedValue([{ model_group: "gpt-4" }, { model_group: "gpt-3.5-turbo" }]),
@@ -11,6 +12,34 @@ vi.mock("../llm_calls/chat_completion", () => ({
   makeOpenAIChatCompletionRequest: vi.fn().mockResolvedValue(undefined),
 }));
 
+let capturedOnImageUpload: ((file: File) => false) | null = null;
+
+vi.mock("../chat_ui/ChatImageUpload", () => ({
+  default: ({ onImageUpload }: { onImageUpload: (file: File) => false }) => {
+    capturedOnImageUpload = onImageUpload;
+    return (
+      <div data-testid="chat-image-upload">
+        <button data-testid="trigger-upload">Upload</button>
+      </div>
+    );
+  },
+}));
+
+vi.mock("../chat_ui/ChatImageUtils", () => ({
+  createChatMultimodalMessage: vi.fn().mockResolvedValue({
+    role: "user",
+    content: [
+      { type: "text", text: "test message" },
+      { type: "image_url", image_url: { url: "data:image/png;base64,test" } },
+    ],
+  }),
+  createChatDisplayMessage: vi.fn().mockReturnValue({
+    role: "user",
+    content: "test message [Image attached]",
+    imagePreviewUrl: "blob:test-url",
+  }),
+}));
+
 vi.mock("./components/ComparisonPanel", () => ({
   ComparisonPanel: ({ comparison, onRemove }: { comparison: any; onRemove: () => void }) => (
     <div data-testid={`comparison-panel-${comparison.id}`}>
@@ -22,8 +51,9 @@ vi.mock("./components/ComparisonPanel", () => ({
 }));
 
 vi.mock("./components/MessageInput", () => ({
-  MessageInput: ({ value, onChange, onSend, disabled }: any) => (
+  MessageInput: ({ value, onChange, onSend, disabled, hasAttachment, uploadComponent }: any) => (
     <div data-testid="message-input">
+      {uploadComponent && <div data-testid="upload-component">{uploadComponent}</div>}
       <textarea
         data-testid="message-textarea"
         value={value}
@@ -33,6 +63,7 @@ vi.mock("./components/MessageInput", () => ({
       <button data-testid="send-button" onClick={onSend} disabled={disabled}>
         Send
       </button>
+      {hasAttachment && <div data-testid="has-attachment">Attachment</div>}
     </div>
   ),
 }));
@@ -51,6 +82,10 @@ beforeEach(() => {
       dispatchEvent: () => false,
     }),
   });
+  global.URL.createObjectURL = vi.fn().mockReturnValue("blob:test-url");
+  global.URL.revokeObjectURL = vi.fn();
+  capturedOnImageUpload = null;
+  vi.clearAllMocks();
 });
 
 describe("CompareUI", () => {
@@ -88,4 +123,36 @@ describe("CompareUI", () => {
     expect(getByTestId("comparison-panel-1")).toBeInTheDocument();
     expect(getByTestId("comparison-panel-2")).toBeInTheDocument();
   });
+
+  it("should handle image upload and send message with attachment", async () => {
+    const user = userEvent.setup();
+    const { getByTestId, queryByTestId } = render(
+      <CompareUI accessToken="test-token" disabledPersonalKeyCreation={false} />,
+    );
+
+    const file = new File(["test content"], "test-image.png", { type: "image/png" });
+    
+    await waitFor(() => {
+      expect(capturedOnImageUpload).not.toBeNull();
+    });
+
+    if (capturedOnImageUpload) {
+      capturedOnImageUpload(file);
+    }
+
+    await waitFor(() => {
+      expect(queryByTestId("has-attachment")).toBeInTheDocument();
+    });
+
+    const textarea = getByTestId("message-textarea");
+    await user.type(textarea, "Describe this image");
+
+    const sendButton = getByTestId("send-button");
+    expect(sendButton).not.toBeDisabled();
+    await user.click(sendButton);
+
+    await waitFor(() => {
+      expect(makeOpenAIChatCompletionRequest).toHaveBeenCalled();
+    });
+  });
 });
diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/CompareUI.tsx b/ui/litellm-dashboard/src/components/playground/compareUI/CompareUI.tsx
index ef4d65d409..9ba0953505 100644
--- a/ui/litellm-dashboard/src/components/playground/compareUI/CompareUI.tsx
+++ b/ui/litellm-dashboard/src/components/playground/compareUI/CompareUI.tsx
@@ -1,14 +1,16 @@
 "use client";
 
-import React, { useEffect, useMemo, useState } from "react";
-import { v4 as uuidv4 } from "uuid";
-import { Select, Input, Tooltip, Button } from "antd";
-import { ClearOutlined, PlusOutlined } from "@ant-design/icons";
 import NotificationsManager from "@/components/molecules/notifications_manager";
-import { fetchAvailableModels } from "../llm_calls/fetch_models";
-import { makeOpenAIChatCompletionRequest } from "../llm_calls/chat_completion";
+import { ClearOutlined, DeleteOutlined, FilePdfOutlined, PlusOutlined } from "@ant-design/icons";
+import { Button, Input, Select, Tooltip } from "antd";
+import { useEffect, useMemo, useState } from "react";
+import { v4 as uuidv4 } from "uuid";
+import ChatImageUpload from "../chat_ui/ChatImageUpload";
+import { createChatDisplayMessage, createChatMultimodalMessage } from "../chat_ui/ChatImageUtils";
 import type { TokenUsage } from "../chat_ui/ResponseMetrics";
 import type { MessageType, VectorStoreSearchResponse } from "../chat_ui/types";
+import { makeOpenAIChatCompletionRequest } from "../llm_calls/chat_completion";
+import { fetchAvailableModels } from "../llm_calls/fetch_models";
 import { ComparisonPanel } from "./components/ComparisonPanel";
 import { MessageInput } from "./components/MessageInput";
 export interface ComparisonInstance {
@@ -71,6 +73,8 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
   const [modelOptions, setModelOptions] = useState<string[]>([]);
   const [isLoadingModels, setIsLoadingModels] = useState(false);
   const [inputValue, setInputValue] = useState("");
+  const [uploadedFile, setUploadedFile] = useState<File | null>(null);
+  const [uploadedFilePreviewUrl, setUploadedFilePreviewUrl] = useState<string | null>(null);
   const [apiKeySource, setApiKeySource] = useState<"session" | "custom">(
     disabledPersonalKeyCreation ? "custom" : "session",
   );
@@ -82,6 +86,13 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
     }, 300);
     return () => clearTimeout(timer);
   }, [customApiKey]);
+  useEffect(() => {
+    return () => {
+      if (uploadedFilePreviewUrl) {
+        URL.revokeObjectURL(uploadedFilePreviewUrl);
+      }
+    };
+  }, [uploadedFilePreviewUrl]);
   const effectiveApiKey = useMemo(
     () => (apiKeySource === "session" ? accessToken || "" : debouncedCustomApiKey.trim()),
     [apiKeySource, accessToken, debouncedCustomApiKey],
@@ -215,6 +226,21 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
       );
     });
   };
+  const handleFileUpload = (file: File): false => {
+    if (uploadedFilePreviewUrl) {
+      URL.revokeObjectURL(uploadedFilePreviewUrl);
+    }
+    setUploadedFile(file);
+    setUploadedFilePreviewUrl(URL.createObjectURL(file));
+    return false;
+  };
+  const handleRemoveFile = () => {
+    if (uploadedFilePreviewUrl) {
+      URL.revokeObjectURL(uploadedFilePreviewUrl);
+    }
+    setUploadedFile(null);
+    setUploadedFilePreviewUrl(null);
+  };
   const clearAllChats = () => {
     setComparisons((prev) =>
       prev.map((comparison) => ({
@@ -225,6 +251,7 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
       })),
     );
     setInputValue("");
+    handleRemoveFile();
   };
   const appendAssistantChunk = (comparisonId: string, chunk: string, model?: string) => {
     if (!chunk) {
@@ -389,9 +416,10 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
     );
   };
   const canUseSessionKey = Boolean(accessToken);
-  const handleSendMessage = (input: string) => {
+  const handleSendMessage = async (input: string) => {
     const trimmed = input.trim();
-    if (!trimmed) {
+    const hasAttachment = Boolean(uploadedFile);
+    if (!trimmed && !hasAttachment) {
       return;
     }
     if (!effectiveApiKey) {
@@ -406,6 +434,17 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
       NotificationsManager.fromBackend("Select a model before sending a message.");
       return;
     }
+
+    const apiUserMessage = hasAttachment
+      ? await createChatMultimodalMessage(trimmed, uploadedFile as File)
+      : { role: "user", content: trimmed };
+    const displayUserMessage = createChatDisplayMessage(
+      trimmed,
+      hasAttachment,
+      uploadedFilePreviewUrl || undefined,
+      uploadedFile?.name,
+    );
+
     const preparedTargets = new Map<
       string,
       {
@@ -417,15 +456,19 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
         guardrails: string[];
         temperature: number;
         maxTokens: number;
-        messages: MessageType[];
+        displayMessages: MessageType[];
+        apiChatHistory: Array<{ role: string; content: string | any[] }>;
       }
     >();
     targetComparisons.forEach((comparison) => {
       const traceId = comparison.traceId ?? uuidv4();
-      const userMessage: MessageType = {
-        role: "user",
-        content: trimmed,
-      };
+      const apiChatHistory = [
+        ...comparison.messages.map(({ role, content }) => ({
+          role,
+          content: Array.isArray(content) ? content : typeof content === "string" ? content : "",
+        })),
+        apiUserMessage,
+      ];
       preparedTargets.set(comparison.id, {
         id: comparison.id,
         model: comparison.model,
@@ -435,7 +478,8 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
         guardrails: comparison.guardrails,
         temperature: comparison.temperature,
         maxTokens: comparison.maxTokens,
-        messages: [...comparison.messages, userMessage],
+        displayMessages: [...comparison.messages, displayUserMessage],
+        apiChatHistory,
       });
     });
     if (preparedTargets.size === 0) {
@@ -450,23 +494,22 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
         return {
           ...comparison,
           traceId: prepared.traceId,
-          messages: prepared.messages,
+          messages: prepared.displayMessages,
           isLoading: true,
         };
       }),
     );
+    setInputValue("");
+    handleRemoveFile();
+
     preparedTargets.forEach((prepared) => {
-      const apiChatHistory = prepared.messages.map(({ role, content }) => ({
-        role,
-        content: typeof content === "string" ? content : "",
-      }));
       const tags = prepared.tags.length > 0 ? prepared.tags : undefined;
       const vectorStoreIds = prepared.vectorStores.length > 0 ? prepared.vectorStores : undefined;
       const guardrails = prepared.guardrails.length > 0 ? prepared.guardrails : undefined;
       const comparison = comparisons.find((c) => c.id === prepared.id);
       const useAdvancedParams = comparison?.useAdvancedParams ?? false;
       makeOpenAIChatCompletionRequest(
-        apiChatHistory,
+        prepared.apiChatHistory,
         (chunk, model) => appendAssistantChunk(prepared.id, chunk, model),
         prepared.model,
         effectiveApiKey,
@@ -536,18 +579,19 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
     setInputValue(value);
   };
   const handleSubmit = () => {
-    handleSendMessage(inputValue);
-    setInputValue("");
+    void handleSendMessage(inputValue);
   };
   const handleFollowUpSelect = (question: string) => {
     setInputValue(question);
   };
   const hasMessages = comparisons.some((comparison) => comparison.messages.length > 0);
   const isAnyComparisonLoading = comparisons.some((comparison) => comparison.isLoading);
-  const showSuggestedPrompts = !hasMessages && !isAnyComparisonLoading;
+  const hasAttachment = Boolean(uploadedFile);
+  const isUploadedFilePdf = Boolean(uploadedFile?.name.toLowerCase().endsWith(".pdf"));
+  const showSuggestedPrompts = !hasMessages && !isAnyComparisonLoading && !hasAttachment;
   return (
     <div className="w-full h-full p-4 bg-white">
-      <div className="rounded-2xl border border-gray-200 bg-white shadow-sm min-h-[calc(100vh-140px)] flex flex-col">
+      <div className="rounded-2xl border border-gray-200 bg-white shadow-sm min-h-[calc(100vh-160px)] flex flex-col">
         <div className="border-b px-4 py-2">
           <div className="flex flex-wrap items-center justify-between gap-3">
             <div className="flex items-center gap-2">
@@ -620,7 +664,9 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
           <div className="w-full max-w-3xl px-4">
             <div className="border border-gray-200 shadow-lg rounded-xl bg-white p-4">
               <div className="flex items-center justify-between gap-4 mb-3 min-h-8">
-                {showSuggestedPrompts ? (
+                {hasAttachment ? (
+                  <span className="text-sm text-gray-500">Attachment ready to send</span>
+                ) : showSuggestedPrompts ? (
                   <div className="flex items-center gap-2 overflow-x-auto">
                     {SUGGESTED_PROMPTS.map((prompt) => (
                       <button
@@ -633,7 +679,7 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
                       </button>
                     ))}
                   </div>
-                ) : haveAllResponses ? (
+                ) : haveAllResponses && !hasAttachment ? (
                   <div className="flex items-center gap-2 overflow-x-auto">
                     {GENERIC_FOLLOW_UPS.map((question) => (
                       <button
@@ -655,11 +701,49 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
                   <span className="text-sm text-gray-500">Send a prompt to compare models</span>
                 )}
               </div>
+              {uploadedFile && (
+                <div className="mb-3">
+                  <div className="flex items-center gap-3 p-3 bg-gray-50 rounded-lg border border-gray-200">
+                    <div className="relative inline-block">
+                      {isUploadedFilePdf ? (
+                        <div className="w-10 h-10 rounded-md bg-red-500 flex items-center justify-center">
+                          <FilePdfOutlined style={{ fontSize: "16px", color: "white" }} />
+                        </div>
+                      ) : (
+                        <img
+                          src={uploadedFilePreviewUrl || ""}
+                          alt="Upload preview"
+                          className="w-10 h-10 rounded-md border border-gray-200 object-cover"
+                        />
+                      )}
+                    </div>
+                    <div className="flex-1 min-w-0">
+                      <div className="text-sm font-medium text-gray-900 truncate">{uploadedFile.name}</div>
+                      <div className="text-xs text-gray-500">{isUploadedFilePdf ? "PDF" : "Image"}</div>
+                    </div>
+                    <button
+                      className="flex items-center justify-center w-6 h-6 text-gray-400 hover:text-gray-600 hover:bg-gray-200 rounded-full transition-colors"
+                      onClick={handleRemoveFile}
+                    >
+                      <DeleteOutlined style={{ fontSize: "12px" }} />
+                    </button>
+                  </div>
+                </div>
+              )}
               <MessageInput
                 value={inputValue}
                 onChange={handleInputChange}
                 onSend={handleSubmit}
                 disabled={comparisons.length === 0 || comparisons.every((comparison) => comparison.isLoading)}
+                hasAttachment={hasAttachment}
+                uploadComponent={
+                  <ChatImageUpload
+                    chatUploadedImage={uploadedFile}
+                    chatImagePreviewUrl={uploadedFilePreviewUrl}
+                    onImageUpload={handleFileUpload}
+                    onRemoveImage={handleRemoveFile}
+                  />
+                }
               />
             </div>
           </div>
diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageDisplay.test.tsx b/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageDisplay.test.tsx
index 5e6fa5f739..c635e9555f 100644
--- a/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageDisplay.test.tsx
+++ b/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageDisplay.test.tsx
@@ -17,6 +17,15 @@ vi.mock("../../chat_ui/SearchResultsDisplay", () => ({
   SearchResultsDisplay: () => <div data-testid="search-results">SearchResultsDisplay</div>,
 }));
 
+vi.mock("../../chat_ui/ChatImageRenderer", () => ({
+  default: ({ message }: { message: any }) =>
+    message.imagePreviewUrl ? (
+      <div data-testid="chat-image-renderer">
+        <img src={message.imagePreviewUrl} alt="User uploaded image" />
+      </div>
+    ) : null,
+}));
+
 describe("MessageDisplay", () => {
   it("should render", () => {
     const messages: MessageType[] = [
@@ -63,4 +72,24 @@ describe("MessageDisplay", () => {
     expect(getByText("2+2 equals 4")).toBeInTheDocument();
     expect(getByTestId("response-metrics")).toBeInTheDocument();
   });
+
+  it("should display image attachment in user message", () => {
+    const messages: MessageType[] = [
+      {
+        role: "user",
+        content: "What is in this image? [Image attached]",
+        imagePreviewUrl: "blob:test-image-url",
+      },
+      {
+        role: "assistant",
+        content: "This is a test image",
+        model: "gpt-4",
+      },
+    ];
+    const { getByTestId, getByText } = render(<MessageDisplay messages={messages} isLoading={false} />);
+    expect(getByText("What is in this image? [Image attached]")).toBeInTheDocument();
+    expect(getByTestId("chat-image-renderer")).toBeInTheDocument();
+    const image = getByTestId("chat-image-renderer").querySelector("img");
+    expect(image).toHaveAttribute("src", "blob:test-image-url");
+  });
 });
diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageDisplay.tsx b/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageDisplay.tsx
index 734ddf26da..088e45b91b 100644
--- a/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageDisplay.tsx
+++ b/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageDisplay.tsx
@@ -1,8 +1,9 @@
+import { Bot, Loader2, UserRound } from "lucide-react";
 import React from "react";
 import ReactMarkdown from "react-markdown";
 import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
 import { coy } from "react-syntax-highlighter/dist/esm/styles/prism";
-import { Bot, Loader2, UserRound } from "lucide-react";
+import ChatImageRenderer from "../../chat_ui/ChatImageRenderer";
 import ReasoningContent from "../../chat_ui/ReasoningContent";
 import ResponseMetrics from "../../chat_ui/ResponseMetrics";
 import { SearchResultsDisplay } from "../../chat_ui/SearchResultsDisplay";
@@ -56,6 +57,7 @@ export function MessageDisplay({ messages, isLoading }: MessageDisplayProps) {
         hyphens: "auto",
       }}
     >
+      <ChatImageRenderer message={message} />
       <ReactMarkdown
         components={{
           code({
diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageInput.test.tsx b/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageInput.test.tsx
index b9ca9187f7..4ec45fa2df 100644
--- a/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageInput.test.tsx
+++ b/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageInput.test.tsx
@@ -21,4 +21,16 @@ describe("MessageInput", () => {
 
     expect(button).toBeDisabled();
   });
+
+  it("should enable send button when hasAttachment is true even with empty value", () => {
+    const onChange = vi.fn();
+    const onSend = vi.fn();
+    const uploadComponent = <div data-testid="upload-component">Upload</div>;
+    const { container, getByTestId } = render(
+      <MessageInput value="" onChange={onChange} onSend={onSend} hasAttachment={true} uploadComponent={uploadComponent} />,
+    );
+    const button = container.querySelector("button") as HTMLButtonElement;
+    expect(getByTestId("upload-component")).toBeInTheDocument();
+    expect(button).not.toBeDisabled();
+  });
 });
diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageInput.tsx b/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageInput.tsx
index 23c8e89718..67036282b3 100644
--- a/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageInput.tsx
+++ b/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageInput.tsx
@@ -9,13 +9,17 @@ interface MessageInputProps {
   onChange: (value: string) => void;
   onSend: () => void;
   disabled?: boolean;
+  hasAttachment?: boolean;
+  uploadComponent?: React.ReactNode;
 }
 
-export function MessageInput({ value, onChange, onSend, disabled }: MessageInputProps) {
+export function MessageInput({ value, onChange, onSend, disabled, hasAttachment, uploadComponent }: MessageInputProps) {
+  const canSend = !disabled && (value.trim().length > 0 || Boolean(hasAttachment));
+
   const handleKeyDown = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
     if (e.key === "Enter" && !e.shiftKey) {
       e.preventDefault();
-      if (!disabled && value.trim()) {
+      if (canSend) {
         onSend();
       }
     }
@@ -24,6 +28,7 @@ export function MessageInput({ value, onChange, onSend, disabled }: MessageInput
   return (
     <div className="flex items-center gap-2">
       <div className="flex items-center flex-1 bg-white border border-gray-300 rounded-xl px-3 py-1 min-h-[44px]">
+        {uploadComponent && <div className="flex-shrink-0 mr-2">{uploadComponent}</div>}
         <TextArea
           value={value}
           onChange={(e) => onChange(e.target.value)}
@@ -42,7 +47,7 @@ export function MessageInput({ value, onChange, onSend, disabled }: MessageInput
             lineHeight: "20px",
           }}
         />
-        <Button onClick={onSend} disabled={disabled || !value.trim()} icon={<ArrowUpOutlined />} shape="circle" />
+        <Button onClick={onSend} disabled={!canSend} icon={<ArrowUpOutlined />} shape="circle" />
       </div>
     </div>
   );