diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx index df04eece289..80e25f18cce 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx @@ -69,6 +69,7 @@ import { createDisplayMessage, createMultimodalMessage } from "./ResponsesImageU import { SearchResultsDisplay } from "./SearchResultsDisplay"; import SessionManagement from "./SessionManagement"; import RealtimePlayground from "./RealtimePlayground"; +import MessageActions from "./MessageActions"; import { A2ATaskMetadata, MessageType } from "./types"; import { useCodeInterpreter } from "./useCodeInterpreter"; @@ -242,6 +243,8 @@ const ChatUI: React.FC = ({ // Code Interpreter state (using custom hook) const codeInterpreter = useCodeInterpreter(); + const [pendingResend, setPendingResend] = useState(false); + const chatEndRef = useRef(null); // Fetch MCP servers @@ -1217,6 +1220,54 @@ const ChatUI: React.FC = ({ setInputMessage(""); }; + useEffect(() => { + if (pendingResend && inputMessage.trim() !== "" && !isLoading) { + setPendingResend(false); + handleSendMessage(); + } + }, [pendingResend, inputMessage]); + + const handleRetry = () => { + if (isLoading || chatHistory.length === 0) return; + + // Find the last assistant message index + let lastAssistantIdx = -1; + for (let i = chatHistory.length - 1; i >= 0; i--) { + if (chatHistory[i].role === "assistant") { + lastAssistantIdx = i; + break; + } + } + if (lastAssistantIdx === -1) return; + + // Find the user message right before the assistant message + let userMsgIdx = -1; + for (let i = lastAssistantIdx - 1; i >= 0; i--) { + if (chatHistory[i].role === "user") { + userMsgIdx = i; + break; + } + } + if (userMsgIdx === -1) return; + + const userMessage = chatHistory[userMsgIdx]; + const userContent = typeof userMessage.content === "string" ? userMessage.content : ""; + + // Truncate history: keep everything before the user message that triggered the response + setChatHistory(chatHistory.slice(0, userMsgIdx)); + setInputMessage(userContent); + setPendingResend(true); + }; + + const handleEditSubmit = (messageIndex: number, newContent: string) => { + if (isLoading) return; + + // Truncate history at the edited message (remove it and everything after) + setChatHistory(chatHistory.slice(0, messageIndex)); + setInputMessage(newContent); + setPendingResend(true); + }; + const clearChatHistory = () => { // Clean up audio object URLs before clearing history chatHistory.forEach((message) => { @@ -1867,7 +1918,7 @@ const ChatUI: React.FC = ({ {chatHistory.map((message, index) => (
-
+
= ({ )}
+ m.role).lastIndexOf("assistant") + } + isLoading={isLoading} + isImage={message.isImage} + isAudio={message.isAudio} + isEmbeddings={message.isEmbeddings} + onRetry={handleRetry} + onEditSubmit={handleEditSubmit} + />
))} diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/MessageActions.test.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/MessageActions.test.tsx new file mode 100644 index 00000000000..86c135bad4d --- /dev/null +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/MessageActions.test.tsx @@ -0,0 +1,218 @@ +import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import MessageActions from "./MessageActions"; + +describe("MessageActions", () => { + const defaultProps = { + role: "user", + content: "Hello world", + messageIndex: 0, + isLastAssistantMessage: false, + isLoading: false, + onRetry: vi.fn(), + onEditSubmit: vi.fn(), + }; + + it("should render edit button for user messages on hover", () => { + render(); + const editButton = screen.getByTestId("edit-message-button"); + expect(editButton).toBeInTheDocument(); + }); + + it("should not render any buttons for user messages when loading", () => { + render(); + expect(screen.queryByTestId("edit-message-button")).toBeNull(); + expect(screen.queryByTestId("retry-message-button")).toBeNull(); + }); + + it("should render retry button for the last assistant message", () => { + render( + , + ); + const retryButton = screen.getByTestId("retry-message-button"); + expect(retryButton).toBeInTheDocument(); + }); + + it("should not render retry button for non-last assistant messages", () => { + render( + , + ); + expect(screen.queryByTestId("retry-message-button")).toBeNull(); + }); + + it("should not render edit button for assistant messages", () => { + render( + , + ); + expect(screen.queryByTestId("edit-message-button")).toBeNull(); + }); + + it("should not render any buttons for image messages", () => { + render(); + expect(screen.queryByTestId("edit-message-button")).toBeNull(); + expect(screen.queryByTestId("message-actions")).toBeNull(); + }); + + it("should not render any buttons for audio messages", () => { + render(); + expect(screen.queryByTestId("message-actions")).toBeNull(); + }); + + it("should call onRetry when retry button is clicked", () => { + const onRetry = vi.fn(); + render( + , + ); + + act(() => { + fireEvent.click(screen.getByTestId("retry-message-button")); + }); + + expect(onRetry).toHaveBeenCalledTimes(1); + }); + + it("should show edit textarea when edit button is clicked", async () => { + render(); + + act(() => { + fireEvent.click(screen.getByTestId("edit-message-button")); + }); + + await waitFor(() => { + expect(screen.getByTestId("edit-message-container")).toBeInTheDocument(); + }); + + const textarea = screen.getByRole("textbox"); + expect(textarea).toBeInTheDocument(); + expect(textarea).toHaveValue("Hello world"); + }); + + it("should call onEditSubmit when confirming edit", async () => { + const onEditSubmit = vi.fn(); + render(); + + act(() => { + fireEvent.click(screen.getByTestId("edit-message-button")); + }); + + await waitFor(() => { + expect(screen.getByTestId("edit-message-container")).toBeInTheDocument(); + }); + + const textarea = screen.getByRole("textbox"); + act(() => { + fireEvent.change(textarea, { target: { value: "Updated message" } }); + }); + + act(() => { + fireEvent.click(screen.getByTestId("confirm-edit-button")); + }); + + expect(onEditSubmit).toHaveBeenCalledWith(0, "Updated message"); + }); + + it("should cancel edit when cancel button is clicked", async () => { + render(); + + act(() => { + fireEvent.click(screen.getByTestId("edit-message-button")); + }); + + await waitFor(() => { + expect(screen.getByTestId("edit-message-container")).toBeInTheDocument(); + }); + + act(() => { + fireEvent.click(screen.getByTestId("cancel-edit-button")); + }); + + await waitFor(() => { + expect(screen.queryByTestId("edit-message-container")).toBeNull(); + }); + + expect(screen.getByTestId("edit-message-button")).toBeInTheDocument(); + }); + + it("should cancel edit when Escape key is pressed", async () => { + render(); + + act(() => { + fireEvent.click(screen.getByTestId("edit-message-button")); + }); + + await waitFor(() => { + expect(screen.getByTestId("edit-message-container")).toBeInTheDocument(); + }); + + const textarea = screen.getByRole("textbox"); + act(() => { + fireEvent.keyDown(textarea, { key: "Escape" }); + }); + + await waitFor(() => { + expect(screen.queryByTestId("edit-message-container")).toBeNull(); + }); + }); + + it("should submit edit when Enter key is pressed without Shift", async () => { + const onEditSubmit = vi.fn(); + render(); + + act(() => { + fireEvent.click(screen.getByTestId("edit-message-button")); + }); + + await waitFor(() => { + expect(screen.getByTestId("edit-message-container")).toBeInTheDocument(); + }); + + const textarea = screen.getByRole("textbox"); + act(() => { + fireEvent.change(textarea, { target: { value: "New content" } }); + }); + + act(() => { + fireEvent.keyDown(textarea, { key: "Enter", shiftKey: false }); + }); + + expect(onEditSubmit).toHaveBeenCalledWith(0, "New content"); + }); + + it("should not submit edit with empty content", async () => { + const onEditSubmit = vi.fn(); + render(); + + act(() => { + fireEvent.click(screen.getByTestId("edit-message-button")); + }); + + await waitFor(() => { + expect(screen.getByTestId("edit-message-container")).toBeInTheDocument(); + }); + + const textarea = screen.getByRole("textbox"); + act(() => { + fireEvent.change(textarea, { target: { value: " " } }); + }); + + const confirmButton = screen.getByTestId("confirm-edit-button"); + expect(confirmButton).toBeDisabled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/MessageActions.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/MessageActions.tsx new file mode 100644 index 00000000000..5b2f614ef43 --- /dev/null +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/MessageActions.tsx @@ -0,0 +1,150 @@ +import { + CheckOutlined, + CloseOutlined, + EditOutlined, + ReloadOutlined, +} from "@ant-design/icons"; +import { Button, Input, Tooltip } from "antd"; +import React, { useEffect, useRef, useState } from "react"; + +const { TextArea } = Input; + +interface MessageActionsProps { + role: string; + content: string; + messageIndex: number; + isLastAssistantMessage: boolean; + isLoading: boolean; + isImage?: boolean; + isAudio?: boolean; + isEmbeddings?: boolean; + onRetry: () => void; + onEditSubmit: (messageIndex: number, newContent: string) => void; +} + +const MessageActions: React.FC = ({ + role, + content, + messageIndex, + isLastAssistantMessage, + isLoading, + isImage, + isAudio, + isEmbeddings, + onRetry, + onEditSubmit, +}) => { + const [isEditing, setIsEditing] = useState(false); + const [editValue, setEditValue] = useState(content); + const textAreaRef = useRef(null); + + useEffect(() => { + if (isEditing && textAreaRef.current) { + const ta = textAreaRef.current?.resizableTextArea?.textArea; + if (ta) { + ta.focus(); + ta.setSelectionRange(ta.value.length, ta.value.length); + } + } + }, [isEditing]); + + const handleStartEdit = () => { + setEditValue(content); + setIsEditing(true); + }; + + const handleCancelEdit = () => { + setIsEditing(false); + setEditValue(content); + }; + + const handleConfirmEdit = () => { + if (editValue.trim() === "") return; + setIsEditing(false); + onEditSubmit(messageIndex, editValue); + }; + + const handleEditKeyDown = (e: React.KeyboardEvent) => { + if (e.key === "Enter" && !e.shiftKey) { + e.preventDefault(); + handleConfirmEdit(); + } + if (e.key === "Escape") { + handleCancelEdit(); + } + }; + + const isSpecialMessage = isImage || isAudio || isEmbeddings; + const canEdit = role === "user" && !isLoading && !isSpecialMessage; + const canRetry = role === "assistant" && isLastAssistantMessage && !isLoading && !isSpecialMessage; + + if (isEditing) { + return ( +
+