tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-17 00:26:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #24172 from BerriAI/litellm_extract_useChatHistory_hook

Browse Source 
[Refactor] Extract useChatHistory hook from ChatUI.tsx
This commit is contained in:
yuneng-jiang
2026-03-20 00:17:07 -07:00
committed by GitHub
parent bdf24757c0 f34fe4758a
commit 8c396e5ca9
4 changed files with 1023 additions and 327 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx
+35 -327
View File
@@ -28,7 +28,6 @@ import ReactMarkdown from "react-markdown";
import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
import { coy } from "react-syntax-highlighter/dist/esm/styles/prism";
import { v4 as uuidv4 } from "uuid";
import { truncateString } from "../../../utils/textUtils";
import GuardrailSelector from "../../guardrails/GuardrailSelector";
import PolicySelector from "../../policies/PolicySelector";
import MCPToolArgumentsForm, { MCPToolArgumentsFormRef } from "../../mcp_tools/MCPToolArgumentsForm";
@@ -62,7 +61,6 @@ import { generateCodeSnippet } from "./CodeSnippets";
import EndpointSelector from "./EndpointSelector";
import FilePreviewCard from "./FilePreviewCard";
import MCPEventsDisplay from "./MCPEventsDisplay";
import type { MCPEvent } from "../../mcp_tools/types";
import { EndpointType, getEndpointType } from "./mode_endpoint_mapping";
import ReasoningContent from "./ReasoningContent";
import ResponseMetrics, { TokenUsage } from "./ResponseMetrics";
@@ -74,6 +72,7 @@ import SessionManagement from "./SessionManagement";
import RealtimePlayground from "./RealtimePlayground";
import { A2ATaskMetadata, MessageType } from "./types";
import { useCodeInterpreter } from "./useCodeInterpreter";
import { useChatHistory } from "./useChatHistory";
const { TextArea } = Input;
const { Dragger } = Upload;
@@ -134,6 +133,34 @@ const ChatUI: React.FC<ChatUIProps> = ({
return {};
}
});
const {
chatHistory,
setChatHistory,
mcpEvents,
setMCPEvents,
messageTraceId,
setMessageTraceId,
responsesSessionId,
setResponsesSessionId,
useApiSessionManagement,
setUseApiSessionManagement,
updateTextUI,
updateReasoningContent,
updateTimingData,
updateUsageData,
updateA2AMetadata,
updateTotalLatency,
updateSearchResults,
handleResponseId,
handleToggleSessionManagement,
handleMCPEvent,
updateImageUI,
updateEmbeddingsUI,
updateAudioUI,
updateChatImageUI,
clearChatHistory: clearChatHistoryHook,
clearMCPEvents,
} = useChatHistory({ simplified });
const [apiKeySource, setApiKeySource] = useState<"session" | "custom">(() => {
const saved = sessionStorage.getItem("apiKeySource");
if (saved) {
@@ -150,16 +177,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
() => sessionStorage.getItem("customProxyBaseUrl") || "",
);
const [inputMessage, setInputMessage] = useState("");
const [chatHistory, setChatHistory] = useState<MessageType[]>(() => {
if (simplified) return [];
try {
const saved = sessionStorage.getItem("chatHistory");
return saved ? JSON.parse(saved) : [];
} catch (error) {
console.error("Error parsing chatHistory from sessionStorage", error);
return [];
}
});
const [selectedModel, setSelectedModel] = useState<string | undefined>(simplified ? fixedModel : undefined);
const [showCustomModelInput, setShowCustomModelInput] = useState<boolean>(false);
const [modelInfo, setModelInfo] = useState<ModelGroup[]>([]);
@@ -217,16 +234,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
return [];
}
});
const [messageTraceId, setMessageTraceId] = useState<string | null>(
() => sessionStorage.getItem("messageTraceId") || null,
);
const [responsesSessionId, setResponsesSessionId] = useState<string | null>(
() => sessionStorage.getItem("responsesSessionId") || null,
);
const [useApiSessionManagement, setUseApiSessionManagement] = useState<boolean>(() => {
const saved = sessionStorage.getItem("useApiSessionManagement");
return saved ? JSON.parse(saved) : true; // Default to API session management
});
const [uploadedImages, setUploadedImages] = useState<File[]>([]);
const [imagePreviewUrls, setImagePreviewUrls] = useState<string[]>([]);
const [responsesUploadedImage, setResponsesUploadedImage] = useState<File | null>(null);
@@ -237,7 +244,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
const [isGetCodeModalVisible, setIsGetCodeModalVisible] = useState(false);
const [generatedCode, setGeneratedCode] = useState("");
const [selectedSdk, setSelectedSdk] = useState<"openai" | "azure">("openai");
const [mcpEvents, setMCPEvents] = useState<MCPEvent[]>([]);
const [temperature, setTemperature] = useState<number>(1.0);
const [maxTokens, setMaxTokens] = useState<number>(2048);
const [useAdvancedParams, setUseAdvancedParams] = useState<boolean>(false);
@@ -331,17 +337,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
proxySettings,
]);
useEffect(() => {
if (simplified) return; // Do not persist chat history in simplified (embedded) mode
const handler = setTimeout(() => {
sessionStorage.setItem("chatHistory", JSON.stringify(chatHistory));
}, 500); // Debounce by 500ms
return () => {
clearTimeout(handler);
};
}, [chatHistory, simplified]);
useEffect(() => {
sessionStorage.setItem("apiKeySource", JSON.stringify(apiKeySource));
sessionStorage.setItem("apiKey", apiKey);
@@ -362,17 +357,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
sessionStorage.removeItem("selectedModel");
}
}
if (messageTraceId) {
sessionStorage.setItem("messageTraceId", messageTraceId);
} else {
sessionStorage.removeItem("messageTraceId");
}
if (responsesSessionId) {
sessionStorage.setItem("responsesSessionId", responsesSessionId);
} else {
sessionStorage.removeItem("responsesSessionId");
}
sessionStorage.setItem("useApiSessionManagement", JSON.stringify(useApiSessionManagement));
// Note: codeInterpreterEnabled and selectedContainerId are persisted by useCodeInterpreter hook
}, [
simplified,
@@ -384,9 +368,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
selectedVectorStores,
selectedGuardrails,
selectedPolicies,
messageTraceId,
responsesSessionId,
useApiSessionManagement,
selectedMCPServers,
mcpServerToolRestrictions,
selectedVoice,
@@ -478,264 +459,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
}
}, [chatHistory]);
const updateTextUI = (role: string, chunk: string, model?: string) => {
console.log("updateTextUI called with:", role, chunk, model);
setChatHistory((prev) => {
const last = prev[prev.length - 1];
// if the last message is already from this same role, append
if (last && last.role === role && !last.isImage && !last.isAudio) {
// build a new object, but only set `model` if it wasn't there already
const updated: MessageType = {
...last,
content: last.content + chunk,
model: last.model ?? model, // ← only use the passed‐in model on the first chunk
};
return [...prev.slice(0, -1), updated];
} else {
// otherwise start a brand new assistant bubble
return [
...prev,
{
role,
content: chunk,
model, // model set exactly once here
},
];
}
});
};
const updateReasoningContent = (chunk: string) => {
setChatHistory((prevHistory) => {
const lastMessage = prevHistory[prevHistory.length - 1];
if (lastMessage && lastMessage.role === "assistant" && !lastMessage.isImage && !lastMessage.isAudio) {
return [
...prevHistory.slice(0, prevHistory.length - 1),
{
...lastMessage,
reasoningContent: (lastMessage.reasoningContent || "") + chunk,
},
];
} else {
// If there's no assistant message yet, we'll create one with empty content
// but with reasoning content
if (prevHistory.length > 0 && prevHistory[prevHistory.length - 1].role === "user") {
return [
...prevHistory,
{
role: "assistant",
content: "",
reasoningContent: chunk,
},
];
}
return prevHistory;
}
});
};
const updateTimingData = (timeToFirstToken: number) => {
console.log("updateTimingData called with:", timeToFirstToken);
setChatHistory((prevHistory) => {
const lastMessage = prevHistory[prevHistory.length - 1];
console.log("Current last message:", lastMessage);
if (lastMessage && lastMessage.role === "assistant") {
console.log("Updating assistant message with timeToFirstToken:", timeToFirstToken);
const updatedHistory = [
...prevHistory.slice(0, prevHistory.length - 1),
{
...lastMessage,
timeToFirstToken,
},
];
console.log("Updated chat history:", updatedHistory);
return updatedHistory;
}
// If the last message is a user message and no assistant message exists yet,
// create a new assistant message with empty content
else if (lastMessage && lastMessage.role === "user") {
console.log("Creating new assistant message with timeToFirstToken:", timeToFirstToken);
return [
...prevHistory,
{
role: "assistant",
content: "",
timeToFirstToken,
},
];
}
console.log("No appropriate message found to update timing");
return prevHistory;
});
};
const updateUsageData = (usage: TokenUsage, toolName?: string) => {
console.log("Received usage data:", usage);
setChatHistory((prevHistory) => {
const lastMessage = prevHistory[prevHistory.length - 1];
if (lastMessage && lastMessage.role === "assistant") {
console.log("Updating message with usage data:", usage);
const updatedMessage = {
...lastMessage,
usage,
toolName,
};
console.log("Updated message:", updatedMessage);
return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
}
return prevHistory;
});
};
const updateA2AMetadata = (a2aMetadata: A2ATaskMetadata) => {
console.log("Received A2A metadata:", a2aMetadata);
setChatHistory((prevHistory) => {
const lastMessage = prevHistory[prevHistory.length - 1];
if (lastMessage && lastMessage.role === "assistant") {
const updatedMessage = {
...lastMessage,
a2aMetadata,
};
return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
}
return prevHistory;
});
};
const updateTotalLatency = (totalLatency: number) => {
setChatHistory((prevHistory) => {
const lastMessage = prevHistory[prevHistory.length - 1];
if (lastMessage && lastMessage.role === "assistant") {
return [
...prevHistory.slice(0, prevHistory.length - 1),
{
...lastMessage,
totalLatency,
},
];
}
return prevHistory;
});
};
const updateSearchResults = (searchResults: any[]) => {
console.log("Received search results:", searchResults);
setChatHistory((prevHistory) => {
const lastMessage = prevHistory[prevHistory.length - 1];
if (lastMessage && lastMessage.role === "assistant") {
console.log("Updating message with search results");
const updatedMessage = {
...lastMessage,
searchResults,
};
return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
}
return prevHistory;
});
};
const handleResponseId = (responseId: string) => {
console.log("Received response ID for session management:", responseId);
if (useApiSessionManagement) {
setResponsesSessionId(responseId);
}
};
const handleToggleSessionManagement = (useApi: boolean) => {
setUseApiSessionManagement(useApi);
if (!useApi) {
// Clear API session when switching to UI mode
setResponsesSessionId(null);
}
};
const handleMCPEvent = (event: MCPEvent) => {
console.log("ChatUI: Received MCP event:", event);
setMCPEvents((prev) => {
// Check if this is a duplicate event (same item_id and type)
// Only check for duplicates if item_id is defined (for mcp_list_tools, item_id is "mcp_list_tools")
const isDuplicate = event.item_id
? prev.some(
(existingEvent) =>
existingEvent.item_id === event.item_id &&
existingEvent.type === event.type &&
(existingEvent.sequence_number === event.sequence_number ||
(existingEvent.sequence_number === undefined && event.sequence_number === undefined)),
)
: false;
if (isDuplicate) {
console.log("ChatUI: Duplicate MCP event, skipping");
return prev;
}
const newEvents = [...prev, event];
console.log("ChatUI: Updated MCP events:", newEvents);
return newEvents;
});
};
const updateImageUI = (imageUrl: string, model: string) => {
setChatHistory((prevHistory) => [...prevHistory, { role: "assistant", content: imageUrl, model, isImage: true }]);
};
const updateEmbeddingsUI = (embeddings: string, model?: string) => {
setChatHistory((prevHistory) => [
...prevHistory,
{ role: "assistant", content: truncateString(embeddings, 100), model, isEmbeddings: true },
]);
};
const updateAudioUI = (audioUrl: string, model: string) => {
setChatHistory((prevHistory) => [...prevHistory, { role: "assistant", content: audioUrl, model, isAudio: true }]);
};
const updateChatImageUI = (imageUrl: string, model?: string) => {
setChatHistory((prev) => {
const last = prev[prev.length - 1];
// If the last message is from assistant and has content, add image to it
if (last && last.role === "assistant" && !last.isImage && !last.isAudio) {
const updated = {
...last,
image: {
url: imageUrl,
detail: "auto",
},
model: last.model ?? model,
};
return [...prev.slice(0, -1), updated];
} else {
// Otherwise create a new assistant message with just the image
return [
...prev,
{
role: "assistant",
content: "",
model,
image: {
url: imageUrl,
detail: "auto",
},
},
];
}
});
};
const handleKeyDown = (event: React.KeyboardEvent<HTMLTextAreaElement>) => {
if (event.key === "Enter" && !event.shiftKey) {
event.preventDefault(); // Prevent default to avoid newline
@@ -966,7 +689,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
}
setChatHistory([...chatHistory, displayMessage]);
setMCPEvents([]); // Clear previous MCP events for new conversation turn
clearMCPEvents(); // Clear previous MCP events for new conversation turn
codeInterpreter.clearResult(); // Clear previous code interpreter results
setIsLoading(true);
@@ -1222,26 +945,11 @@ const ChatUI: React.FC<ChatUIProps> = ({
};
const clearChatHistory = () => {
// Clean up audio object URLs before clearing history
chatHistory.forEach((message) => {
if (message.isAudio && typeof message.content === "string") {
URL.revokeObjectURL(message.content);
}
});
setChatHistory([]);
setMessageTraceId(null);
setResponsesSessionId(null); // Clear responses session ID
setMCPEvents([]); // Clear MCP events
handleRemoveAllImages(); // Clear any uploaded images for image edits
handleRemoveResponsesImage(); // Clear any uploaded images for responses
handleRemoveChatImage(); // Clear any uploaded images for chat completions
handleRemoveAudio(); // Clear any uploaded audio for transcription
if (!simplified) {
sessionStorage.removeItem("chatHistory");
sessionStorage.removeItem("messageTraceId");
sessionStorage.removeItem("responsesSessionId");
}
clearChatHistoryHook();
handleRemoveAllImages();
handleRemoveResponsesImage();
handleRemoveChatImage();
handleRemoveAudio();
NotificationsManager.success("Chat history cleared.");
};
ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.test.ts
+591
View File
@@ -0,0 +1,591 @@
import { renderHook, act } from "@testing-library/react";
import { describe, it, expect, beforeEach, vi } from "vitest";
import { useChatHistory } from "./useChatHistory";
describe("useChatHistory", () => {
beforeEach(() => {
sessionStorage.clear();
});
describe("updateTextUI", () => {
it("should create a new assistant message when chat is empty", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateTextUI("assistant", "Hello", "gpt-4");
});
expect(result.current.chatHistory).toEqual([
{ role: "assistant", content: "Hello", model: "gpt-4" },
]);
});
it("should append to the last assistant message", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateTextUI("assistant", "Hello", "gpt-4");
});
act(() => {
result.current.updateTextUI("assistant", " world");
});
expect(result.current.chatHistory).toEqual([
{ role: "assistant", content: "Hello world", model: "gpt-4" },
]);
});
it("should not overwrite model on subsequent chunks", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateTextUI("assistant", "Hello", "gpt-4");
});
act(() => {
result.current.updateTextUI("assistant", " world", "gpt-3.5");
});
expect(result.current.chatHistory[0].model).toBe("gpt-4");
});
it("should create a new message when role changes", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateTextUI("user", "Hi");
});
act(() => {
result.current.updateTextUI("assistant", "Hello", "gpt-4");
});
expect(result.current.chatHistory).toHaveLength(2);
expect(result.current.chatHistory[0].role).toBe("user");
expect(result.current.chatHistory[1].role).toBe("assistant");
});
it("should not append to image messages", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateImageUI("http://img.png", "dall-e");
});
act(() => {
result.current.updateTextUI("assistant", "description", "gpt-4");
});
expect(result.current.chatHistory).toHaveLength(2);
});
it("should not append to audio messages", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateAudioUI("http://audio.mp3", "tts-1");
});
act(() => {
result.current.updateTextUI("assistant", "text", "gpt-4");
});
expect(result.current.chatHistory).toHaveLength(2);
});
});
describe("updateReasoningContent", () => {
it("should add reasoning content to existing assistant message", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateTextUI("assistant", "Answer", "gpt-4");
});
act(() => {
result.current.updateReasoningContent("thinking...");
});
expect(result.current.chatHistory[0].reasoningContent).toBe("thinking...");
});
it("should append reasoning content across chunks", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateTextUI("assistant", "", "gpt-4");
});
act(() => {
result.current.updateReasoningContent("step 1");
});
act(() => {
result.current.updateReasoningContent(" step 2");
});
expect(result.current.chatHistory[0].reasoningContent).toBe("step 1 step 2");
});
it("should create assistant message with reasoning when last message is user", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.setChatHistory([{ role: "user", content: "question" }]);
});
act(() => {
result.current.updateReasoningContent("thinking...");
});
expect(result.current.chatHistory).toHaveLength(2);
expect(result.current.chatHistory[1]).toEqual({
role: "assistant",
content: "",
reasoningContent: "thinking...",
});
});
it("should not update when chat is empty", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateReasoningContent("thinking...");
});
expect(result.current.chatHistory).toHaveLength(0);
});
});
describe("updateTimingData", () => {
it("should add timeToFirstToken to existing assistant message", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateTextUI("assistant", "Hello", "gpt-4");
});
act(() => {
result.current.updateTimingData(150);
});
expect(result.current.chatHistory[0].timeToFirstToken).toBe(150);
});
it("should create assistant message when last is user", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.setChatHistory([{ role: "user", content: "hi" }]);
});
act(() => {
result.current.updateTimingData(200);
});
expect(result.current.chatHistory).toHaveLength(2);
expect(result.current.chatHistory[1].timeToFirstToken).toBe(200);
});
});
describe("updateUsageData", () => {
it("should add usage data to assistant message", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateTextUI("assistant", "Hello", "gpt-4");
});
const usage = { completionTokens: 10, promptTokens: 5, totalTokens: 15 };
act(() => {
result.current.updateUsageData(usage);
});
expect(result.current.chatHistory[0].usage).toEqual(usage);
});
it("should add toolName when provided", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateTextUI("assistant", "Hello", "gpt-4");
});
const usage = { completionTokens: 10, promptTokens: 5, totalTokens: 15 };
act(() => {
result.current.updateUsageData(usage, "search_tool");
});
expect(result.current.chatHistory[0].toolName).toBe("search_tool");
});
});
describe("updateTotalLatency", () => {
it("should add totalLatency to assistant message", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateTextUI("assistant", "Hello", "gpt-4");
});
act(() => {
result.current.updateTotalLatency(500);
});
expect(result.current.chatHistory[0].totalLatency).toBe(500);
});
});
describe("updateA2AMetadata", () => {
it("should add A2A metadata to assistant message", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateTextUI("assistant", "Hello", "gpt-4");
});
const metadata = { taskId: "task-1", contextId: "ctx-1" };
act(() => {
result.current.updateA2AMetadata(metadata);
});
expect(result.current.chatHistory[0].a2aMetadata).toEqual(metadata);
});
});
describe("updateSearchResults", () => {
it("should add search results to assistant message", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateTextUI("assistant", "Hello", "gpt-4");
});
const searchResults = [{ object: "search", search_query: "test", data: [] }];
act(() => {
result.current.updateSearchResults(searchResults);
});
expect(result.current.chatHistory[0].searchResults).toEqual(searchResults);
});
});
describe("updateImageUI", () => {
it("should add image message to history", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateImageUI("http://img.png", "dall-e-3");
});
expect(result.current.chatHistory).toEqual([
{ role: "assistant", content: "http://img.png", model: "dall-e-3", isImage: true },
]);
});
});
describe("updateEmbeddingsUI", () => {
it("should add truncated embeddings message", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateEmbeddingsUI("[0.1, 0.2, 0.3]", "text-embedding-ada");
});
expect(result.current.chatHistory[0].isEmbeddings).toBe(true);
expect(result.current.chatHistory[0].model).toBe("text-embedding-ada");
});
});
describe("updateAudioUI", () => {
it("should add audio message to history", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateAudioUI("http://audio.mp3", "tts-1");
});
expect(result.current.chatHistory).toEqual([
{ role: "assistant", content: "http://audio.mp3", model: "tts-1", isAudio: true },
]);
});
});
describe("updateChatImageUI", () => {
it("should add image to existing assistant message", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateTextUI("assistant", "Here is the image", "gpt-4");
});
act(() => {
result.current.updateChatImageUI("http://img.png", "gpt-4");
});
expect(result.current.chatHistory[0].image).toEqual({
url: "http://img.png",
detail: "auto",
});
});
it("should create new assistant message with image when no assistant message exists", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateChatImageUI("http://img.png", "gpt-4");
});
expect(result.current.chatHistory[0]).toEqual({
role: "assistant",
content: "",
model: "gpt-4",
image: { url: "http://img.png", detail: "auto" },
});
});
});
describe("handleMCPEvent", () => {
it("should add MCP event", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.handleMCPEvent({ type: "tool_call", item_id: "1" });
});
expect(result.current.mcpEvents).toHaveLength(1);
});
it("should deduplicate events by item_id and type", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
const event = { type: "tool_call", item_id: "1" };
act(() => {
result.current.handleMCPEvent(event);
});
act(() => {
result.current.handleMCPEvent(event);
});
expect(result.current.mcpEvents).toHaveLength(1);
});
it("should allow events without item_id (no dedup)", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.handleMCPEvent({ type: "tool_call" });
});
act(() => {
result.current.handleMCPEvent({ type: "tool_call" });
});
expect(result.current.mcpEvents).toHaveLength(2);
});
it("should allow events with same item_id/type but different sequence_number", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.handleMCPEvent({ type: "tool_call", item_id: "1", sequence_number: 1 });
});
act(() => {
result.current.handleMCPEvent({ type: "tool_call", item_id: "1", sequence_number: 2 });
});
expect(result.current.mcpEvents).toHaveLength(2);
});
});
describe("clearMCPEvents", () => {
it("should clear MCP events without affecting chat history", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateTextUI("assistant", "Hello", "gpt-4");
result.current.handleMCPEvent({ type: "tool_call", item_id: "1" });
});
act(() => {
result.current.clearMCPEvents();
});
expect(result.current.mcpEvents).toEqual([]);
expect(result.current.chatHistory).toHaveLength(1);
});
});
describe("clearChatHistory", () => {
it("should clear all state", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateTextUI("assistant", "Hello", "gpt-4");
result.current.handleMCPEvent({ type: "tool_call", item_id: "1" });
});
act(() => {
result.current.clearChatHistory();
});
expect(result.current.chatHistory).toEqual([]);
expect(result.current.mcpEvents).toEqual([]);
expect(result.current.messageTraceId).toBeNull();
expect(result.current.responsesSessionId).toBeNull();
});
it("should revoke audio object URLs when clearing", () => {
const revokeSpy = vi.spyOn(URL, "revokeObjectURL").mockImplementation(() => {});
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.updateAudioUI("blob:http://localhost/audio-1", "tts-1");
});
act(() => {
result.current.clearChatHistory();
});
expect(revokeSpy).toHaveBeenCalledWith("blob:http://localhost/audio-1");
revokeSpy.mockRestore();
});
it("should clear sessionStorage when not simplified", () => {
vi.useFakeTimers();
const { result } = renderHook(() => useChatHistory({ simplified: false }));
sessionStorage.setItem("chatHistory", "[]");
sessionStorage.setItem("messageTraceId", "trace-1");
sessionStorage.setItem("responsesSessionId", "resp-1");
act(() => {
result.current.clearChatHistory();
});
// Advance past the 500ms debounce to verify it does not re-write the key
act(() => {
vi.advanceTimersByTime(600);
});
expect(sessionStorage.getItem("chatHistory")).toBeNull();
expect(sessionStorage.getItem("messageTraceId")).toBeNull();
expect(sessionStorage.getItem("responsesSessionId")).toBeNull();
vi.useRealTimers();
});
it("should NOT clear sessionStorage when simplified", () => {
sessionStorage.setItem("chatHistory", '[{"role":"user","content":"hi"}]');
const { result } = renderHook(() => useChatHistory({ simplified: true }));
act(() => {
result.current.clearChatHistory();
});
// simplified mode should not touch sessionStorage
expect(sessionStorage.getItem("chatHistory")).toBe('[{"role":"user","content":"hi"}]');
});
it("should not re-write chatHistory to sessionStorage after clear via debounce", () => {
vi.useFakeTimers();
const { result } = renderHook(() => useChatHistory({ simplified: false }));
// Add a message so the debounce has something to persist
act(() => {
result.current.updateTextUI("assistant", "Hello", "gpt-4");
});
// Let the debounce fire so the message is persisted
act(() => {
vi.advanceTimersByTime(600);
});
expect(sessionStorage.getItem("chatHistory")).not.toBeNull();
// Now clear
act(() => {
result.current.clearChatHistory();
});
// Advance past the debounce — the key should stay removed
act(() => {
vi.advanceTimersByTime(600);
});
expect(sessionStorage.getItem("chatHistory")).toBeNull();
vi.useRealTimers();
});
});
describe("simplified mode session isolation", () => {
it("should not hydrate messageTraceId from sessionStorage in simplified mode", () => {
sessionStorage.setItem("messageTraceId", "trace-from-playground");
const { result } = renderHook(() => useChatHistory({ simplified: true }));
expect(result.current.messageTraceId).toBeNull();
});
it("should not hydrate responsesSessionId from sessionStorage in simplified mode", () => {
sessionStorage.setItem("responsesSessionId", "resp-from-playground");
const { result } = renderHook(() => useChatHistory({ simplified: true }));
expect(result.current.responsesSessionId).toBeNull();
});
it("should not hydrate useApiSessionManagement from sessionStorage in simplified mode", () => {
sessionStorage.setItem("useApiSessionManagement", "false");
const { result } = renderHook(() => useChatHistory({ simplified: true }));
// Should get the default (true), not the stored value
expect(result.current.useApiSessionManagement).toBe(true);
});
it("should not persist session state to sessionStorage in simplified mode", () => {
vi.useFakeTimers();
const { result } = renderHook(() => useChatHistory({ simplified: true }));
act(() => {
result.current.setMessageTraceId("trace-embedded");
result.current.setResponsesSessionId("resp-embedded");
});
// Flush effects
act(() => {
vi.advanceTimersByTime(0);
});
expect(sessionStorage.getItem("messageTraceId")).toBeNull();
expect(sessionStorage.getItem("responsesSessionId")).toBeNull();
vi.useRealTimers();
});
});
describe("session management", () => {
it("handleResponseId should set responsesSessionId when useApiSessionManagement is true", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.handleResponseId("resp-123");
});
expect(result.current.responsesSessionId).toBe("resp-123");
});
it("handleResponseId should NOT set responsesSessionId when useApiSessionManagement is false", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.handleToggleSessionManagement(false);
});
act(() => {
result.current.handleResponseId("resp-123");
});
expect(result.current.responsesSessionId).toBeNull();
});
it("handleToggleSessionManagement should clear session when switching to UI mode", () => {
const { result } = renderHook(() => useChatHistory({ simplified: false }));
act(() => {
result.current.handleResponseId("resp-123");
});
act(() => {
result.current.handleToggleSessionManagement(false);
});
expect(result.current.useApiSessionManagement).toBe(false);
expect(result.current.responsesSessionId).toBeNull();
});
});
});
ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts
+392
View File
@@ -0,0 +1,392 @@
import React, { useState, useEffect } from "react";
import { MessageType, A2ATaskMetadata } from "./types";
import { TokenUsage } from "./ResponseMetrics";
import { MCPEvent } from "../../mcp_tools/types";
import { truncateString } from "../../../utils/textUtils";
export interface UseChatHistoryReturn {
// State
chatHistory: MessageType[];
setChatHistory: React.Dispatch<React.SetStateAction<MessageType[]>>;
mcpEvents: MCPEvent[];
setMCPEvents: React.Dispatch<React.SetStateAction<MCPEvent[]>>;
messageTraceId: string | null;
setMessageTraceId: React.Dispatch<React.SetStateAction<string | null>>;
responsesSessionId: string | null;
setResponsesSessionId: React.Dispatch<React.SetStateAction<string | null>>;
useApiSessionManagement: boolean;
setUseApiSessionManagement: React.Dispatch<React.SetStateAction<boolean>>;
// Actions
updateTextUI: (role: string, chunk: string, model?: string) => void;
updateReasoningContent: (chunk: string) => void;
updateTimingData: (timeToFirstToken: number) => void;
updateUsageData: (usage: TokenUsage, toolName?: string) => void;
updateA2AMetadata: (a2aMetadata: A2ATaskMetadata) => void;
updateTotalLatency: (totalLatency: number) => void;
updateSearchResults: (searchResults: any[]) => void;
handleResponseId: (responseId: string) => void;
handleToggleSessionManagement: (useApi: boolean) => void;
handleMCPEvent: (event: MCPEvent) => void;
updateImageUI: (imageUrl: string, model: string) => void;
updateEmbeddingsUI: (embeddings: string, model?: string) => void;
updateAudioUI: (audioUrl: string, model: string) => void;
updateChatImageUI: (imageUrl: string, model?: string) => void;
clearChatHistory: () => void;
clearMCPEvents: () => void;
}
export function useChatHistory({ simplified }: { simplified: boolean }): UseChatHistoryReturn {
const [chatHistory, setChatHistory] = useState<MessageType[]>(() => {
if (simplified) return [];
try {
const saved = sessionStorage.getItem("chatHistory");
return saved ? JSON.parse(saved) : [];
} catch (error) {
console.error("Error parsing chatHistory from sessionStorage", error);
return [];
}
});
const [mcpEvents, setMCPEvents] = useState<MCPEvent[]>([]);
const [messageTraceId, setMessageTraceId] = useState<string | null>(
() => (simplified ? null : sessionStorage.getItem("messageTraceId") || null),
);
const [responsesSessionId, setResponsesSessionId] = useState<string | null>(
() => (simplified ? null : sessionStorage.getItem("responsesSessionId") || null),
);
const [useApiSessionManagement, setUseApiSessionManagement] = useState<boolean>(() => {
if (simplified) return true;
const saved = sessionStorage.getItem("useApiSessionManagement");
return saved ? JSON.parse(saved) : true; // Default to API session management
});
// Debounced chatHistory persistence
useEffect(() => {
if (simplified) return; // Do not persist chat history in simplified (embedded) mode
// When chatHistory is empty (e.g. after clearChatHistory removed the key),
// don't re-write an empty array back into sessionStorage.
if (chatHistory.length === 0) return;
const handler = setTimeout(() => {
sessionStorage.setItem("chatHistory", JSON.stringify(chatHistory));
}, 500); // Debounce by 500ms
return () => {
clearTimeout(handler);
};
}, [chatHistory, simplified]);
// messageTraceId/responsesSessionId/useApiSessionManagement persistence
useEffect(() => {
if (simplified) return;
if (messageTraceId) {
sessionStorage.setItem("messageTraceId", messageTraceId);
} else {
sessionStorage.removeItem("messageTraceId");
}
if (responsesSessionId) {
sessionStorage.setItem("responsesSessionId", responsesSessionId);
} else {
sessionStorage.removeItem("responsesSessionId");
}
sessionStorage.setItem("useApiSessionManagement", JSON.stringify(useApiSessionManagement));
}, [messageTraceId, responsesSessionId, useApiSessionManagement, simplified]);
const updateTextUI = (role: string, chunk: string, model?: string) => {
setChatHistory((prev) => {
const last = prev[prev.length - 1];
// if the last message is already from this same role, append
if (last && last.role === role && !last.isImage && !last.isAudio) {
// build a new object, but only set `model` if it wasn't there already
const updated: MessageType = {
...last,
content: last.content + chunk,
model: last.model ?? model, // ← only use the passed‐in model on the first chunk
};
return [...prev.slice(0, -1), updated];
} else {
// otherwise start a brand new assistant bubble
return [
...prev,
{
role,
content: chunk,
model, // model set exactly once here
},
];
}
});
};
const updateReasoningContent = (chunk: string) => {
setChatHistory((prevHistory) => {
const lastMessage = prevHistory[prevHistory.length - 1];
if (lastMessage && lastMessage.role === "assistant" && !lastMessage.isImage && !lastMessage.isAudio) {
return [
...prevHistory.slice(0, prevHistory.length - 1),
{
...lastMessage,
reasoningContent: (lastMessage.reasoningContent || "") + chunk,
},
];
} else {
// If there's no assistant message yet, we'll create one with empty content
// but with reasoning content
if (prevHistory.length > 0 && prevHistory[prevHistory.length - 1].role === "user") {
return [
...prevHistory,
{
role: "assistant",
content: "",
reasoningContent: chunk,
},
];
}
return prevHistory;
}
});
};
const updateTimingData = (timeToFirstToken: number) => {
setChatHistory((prevHistory) => {
const lastMessage = prevHistory[prevHistory.length - 1];
if (lastMessage && lastMessage.role === "assistant") {
return [
...prevHistory.slice(0, prevHistory.length - 1),
{
...lastMessage,
timeToFirstToken,
},
];
}
// If the last message is a user message and no assistant message exists yet,
// create a new assistant message with empty content
else if (lastMessage && lastMessage.role === "user") {
return [
...prevHistory,
{
role: "assistant",
content: "",
timeToFirstToken,
},
];
}
return prevHistory;
});
};
const updateUsageData = (usage: TokenUsage, toolName?: string) => {
setChatHistory((prevHistory) => {
const lastMessage = prevHistory[prevHistory.length - 1];
if (lastMessage && lastMessage.role === "assistant") {
const updatedMessage = {
...lastMessage,
usage,
toolName,
};
return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
}
return prevHistory;
});
};
const updateA2AMetadata = (a2aMetadata: A2ATaskMetadata) => {
setChatHistory((prevHistory) => {
const lastMessage = prevHistory[prevHistory.length - 1];
if (lastMessage && lastMessage.role === "assistant") {
const updatedMessage = {
...lastMessage,
a2aMetadata,
};
return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
}
return prevHistory;
});
};
const updateTotalLatency = (totalLatency: number) => {
setChatHistory((prevHistory) => {
const lastMessage = prevHistory[prevHistory.length - 1];
if (lastMessage && lastMessage.role === "assistant") {
return [
...prevHistory.slice(0, prevHistory.length - 1),
{
...lastMessage,
totalLatency,
},
];
}
return prevHistory;
});
};
const updateSearchResults = (searchResults: any[]) => {
setChatHistory((prevHistory) => {
const lastMessage = prevHistory[prevHistory.length - 1];
if (lastMessage && lastMessage.role === "assistant") {
const updatedMessage = {
...lastMessage,
searchResults,
};
return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
}
return prevHistory;
});
};
const handleResponseId = (responseId: string) => {
if (useApiSessionManagement) {
setResponsesSessionId(responseId);
}
};
const handleToggleSessionManagement = (useApi: boolean) => {
setUseApiSessionManagement(useApi);
if (!useApi) {
// Clear API session when switching to UI mode
setResponsesSessionId(null);
}
};
const handleMCPEvent = (event: MCPEvent) => {
setMCPEvents((prev) => {
// Check if this is a duplicate event (same item_id and type)
// Only check for duplicates if item_id is defined (for mcp_list_tools, item_id is "mcp_list_tools")
const isDuplicate = event.item_id
? prev.some(
(existingEvent) =>
existingEvent.item_id === event.item_id &&
existingEvent.type === event.type &&
(existingEvent.sequence_number === event.sequence_number ||
(existingEvent.sequence_number === undefined && event.sequence_number === undefined)),
)
: false;
if (isDuplicate) {
return prev;
}
return [...prev, event];
});
};
const updateImageUI = (imageUrl: string, model: string) => {
setChatHistory((prevHistory) => [...prevHistory, { role: "assistant", content: imageUrl, model, isImage: true }]);
};
const updateEmbeddingsUI = (embeddings: string, model?: string) => {
setChatHistory((prevHistory) => [
...prevHistory,
{ role: "assistant", content: truncateString(embeddings, 100), model, isEmbeddings: true },
]);
};
const updateAudioUI = (audioUrl: string, model: string) => {
setChatHistory((prevHistory) => [...prevHistory, { role: "assistant", content: audioUrl, model, isAudio: true }]);
};
const updateChatImageUI = (imageUrl: string, model?: string) => {
setChatHistory((prev) => {
const last = prev[prev.length - 1];
// If the last message is from assistant and has content, add image to it
if (last && last.role === "assistant" && !last.isImage && !last.isAudio) {
const updated = {
...last,
image: {
url: imageUrl,
detail: "auto",
},
model: last.model ?? model,
};
return [...prev.slice(0, -1), updated];
} else {
// Otherwise create a new assistant message with just the image
return [
...prev,
{
role: "assistant",
content: "",
model,
image: {
url: imageUrl,
detail: "auto",
},
},
];
}
});
};
const clearChatHistory = () => {
// Use functional updater to get the latest snapshot — avoids stale-closure
// bugs where audio messages added between the last render and the click
// would leak their blob URLs.
setChatHistory((prev) => {
prev.forEach((message) => {
if (message.isAudio && typeof message.content === "string") {
URL.revokeObjectURL(message.content);
}
});
return [];
});
setMessageTraceId(null);
setResponsesSessionId(null); // Clear responses session ID
setMCPEvents([]); // Clear MCP events
if (!simplified) {
sessionStorage.removeItem("chatHistory");
sessionStorage.removeItem("messageTraceId");
sessionStorage.removeItem("responsesSessionId");
}
};
const clearMCPEvents = () => {
setMCPEvents([]);
};
return {
chatHistory,
setChatHistory,
mcpEvents,
setMCPEvents,
messageTraceId,
setMessageTraceId,
responsesSessionId,
setResponsesSessionId,
useApiSessionManagement,
setUseApiSessionManagement,
updateTextUI,
updateReasoningContent,
updateTimingData,
updateUsageData,
updateA2AMetadata,
updateTotalLatency,
updateSearchResults,
handleResponseId,
handleToggleSessionManagement,
handleMCPEvent,
updateImageUI,
updateEmbeddingsUI,
updateAudioUI,
updateChatImageUI,
clearChatHistory,
clearMCPEvents,
};
}
ui/litellm-dashboard/tests/setupTests.ts
+5
View File
@@ -100,6 +100,11 @@ if (!document.getAnimations) {
document.getAnimations = () => [];
}
// Stub URL.revokeObjectURL so vi.spyOn can intercept it in tests
if (!URL.revokeObjectURL) {
URL.revokeObjectURL = () => {};
}
// Mock ResizeObserver for components that use it (e.g., Tremor UI components)
// This prevents "ResizeObserver is not defined" errors in JSDOM
global.ResizeObserver = class ResizeObserver {