diff --git a/llama_stack/ui/app/chat-playground/page.tsx b/llama_stack/ui/app/chat-playground/page.tsx index b8651aca0..30439554d 100644 --- a/llama_stack/ui/app/chat-playground/page.tsx +++ b/llama_stack/ui/app/chat-playground/page.tsx @@ -1,6 +1,6 @@ "use client"; -import { useState, useEffect } from "react"; +import { useState, useEffect, useCallback, useRef } from "react"; import { flushSync } from "react-dom"; import { Button } from "@/components/ui/button"; import { @@ -15,9 +15,24 @@ import { type Message } from "@/components/chat-playground/chat-message"; import { useAuthClient } from "@/hooks/use-auth-client"; import type { CompletionCreateParams } from "llama-stack-client/resources/chat/completions"; import type { Model } from "llama-stack-client/resources/models"; +import { + SessionManager, + SessionUtils, +} from "@/components/chat-playground/session-manager"; +interface ChatSession { + id: string; + name: string; + messages: Message[]; + selectedModel: string; + systemMessage: string; + createdAt: number; + updatedAt: number; +} export default function ChatPlaygroundPage() { - const [messages, setMessages] = useState([]); + const [currentSession, setCurrentSession] = useState( + null + ); const [input, setInput] = useState(""); const [isGenerating, setIsGenerating] = useState(false); const [error, setError] = useState(null); @@ -26,9 +41,52 @@ export default function ChatPlaygroundPage() { const [modelsLoading, setModelsLoading] = useState(true); const [modelsError, setModelsError] = useState(null); const client = useAuthClient(); + const abortControllerRef = useRef(null); const isModelsLoading = modelsLoading ?? true; + useEffect(() => { + const saved = SessionUtils.loadCurrentSession(); + if (saved) { + setCurrentSession(saved); + } else { + const def = SessionUtils.createDefaultSession(); + const defaultSession: ChatSession = { + ...def, + selectedModel: "", + systemMessage: def.systemMessage || "You are a helpful assistant.", + }; + setCurrentSession(defaultSession); + SessionUtils.saveCurrentSession(defaultSession); + } + }, []); + + const handleModelChange = useCallback((newModel: string) => { + setSelectedModel(newModel); + setCurrentSession(prev => + prev + ? { + ...prev, + selectedModel: newModel, + updatedAt: Date.now(), + } + : prev + ); + }, []); + + useEffect(() => { + if (currentSession) { + if (abortControllerRef.current) { + abortControllerRef.current.abort(); + abortControllerRef.current = null; + setIsGenerating(false); + } + + SessionUtils.saveCurrentSession(currentSession); + setSelectedModel(currentSession.selectedModel); + } + }, [currentSession]); + useEffect(() => { const fetchModels = async () => { try { @@ -38,7 +96,7 @@ export default function ChatPlaygroundPage() { const llmModels = modelList.filter(model => model.model_type === "llm"); setModels(llmModels); if (llmModels.length > 0) { - setSelectedModel(llmModels[0].identifier); + handleModelChange(llmModels[0].identifier); } } catch (err) { console.error("Error fetching models:", err); @@ -49,7 +107,7 @@ export default function ChatPlaygroundPage() { }; fetchModels(); - }, [client]); + }, [client, handleModelChange]); const extractTextContent = (content: unknown): string => { if (typeof content === "string") { @@ -91,7 +149,6 @@ export default function ChatPlaygroundPage() { event?.preventDefault?.(); if (!input.trim()) return; - // Add user message to chat const userMessage: Message = { id: Date.now().toString(), role: "user", @@ -99,10 +156,17 @@ export default function ChatPlaygroundPage() { createdAt: new Date(), }; - setMessages(prev => [...prev, userMessage]); + setCurrentSession(prev => + prev + ? { + ...prev, + messages: [...prev.messages, userMessage], + updatedAt: Date.now(), + } + : prev + ); setInput(""); - // Use the helper function with the content await handleSubmitWithContent(userMessage.content); }; @@ -110,9 +174,19 @@ export default function ChatPlaygroundPage() { setIsGenerating(true); setError(null); + if (abortControllerRef.current) { + abortControllerRef.current.abort(); + } + + const abortController = new AbortController(); + abortControllerRef.current = abortController; + try { const messageParams: CompletionCreateParams["messages"] = [ - ...messages.map(msg => { + ...(currentSession?.systemMessage + ? [{ role: "system" as const, content: currentSession.systemMessage }] + : []), + ...(currentSession?.messages || []).map(msg => { const msgContent = typeof msg.content === "string" ? msg.content @@ -128,11 +202,16 @@ export default function ChatPlaygroundPage() { { role: "user" as const, content }, ]; - const response = await client.chat.completions.create({ - model: selectedModel, - messages: messageParams, - stream: true, - }); + const response = await client.chat.completions.create( + { + model: selectedModel || "", + messages: messageParams, + stream: true, + }, + { + signal: abortController.signal, + } as { signal: AbortSignal } + ); const assistantMessage: Message = { id: (Date.now() + 1).toString(), @@ -141,7 +220,15 @@ export default function ChatPlaygroundPage() { createdAt: new Date(), }; - setMessages(prev => [...prev, assistantMessage]); + setCurrentSession(prev => + prev + ? { + ...prev, + messages: [...prev.messages, assistantMessage], + updatedAt: Date.now(), + } + : null + ); let fullContent = ""; for await (const chunk of response) { if (chunk.choices && chunk.choices[0]?.delta?.content) { @@ -149,23 +236,37 @@ export default function ChatPlaygroundPage() { fullContent += deltaContent; flushSync(() => { - setMessages(prev => { - const newMessages = [...prev]; - const lastMessage = newMessages[newMessages.length - 1]; - if (lastMessage.role === "assistant") { - lastMessage.content = fullContent; - } - return newMessages; + setCurrentSession(prev => { + if (!prev) return null; + const newMessages = [...prev.messages]; + const last = newMessages[newMessages.length - 1]; + if (last.role === "assistant") last.content = fullContent; + return { ...prev, messages: newMessages, updatedAt: Date.now() }; }); }); } } } catch (err) { + // don't show error if request was aborted + if (err instanceof Error && err.name === "AbortError") { + console.log("Request aborted"); + return; + } + console.error("Error sending message:", err); setError("Failed to send message. Please try again."); - setMessages(prev => prev.slice(0, -1)); + setCurrentSession(prev => + prev + ? { + ...prev, + messages: prev.messages.slice(0, -1), + updatedAt: Date.now(), + } + : prev + ); } finally { setIsGenerating(false); + abortControllerRef.current = null; } }; const suggestions = [ @@ -181,69 +282,165 @@ export default function ChatPlaygroundPage() { content: message.content, createdAt: new Date(), }; - setMessages(prev => [...prev, newMessage]); + setCurrentSession(prev => + prev + ? { + ...prev, + messages: [...prev.messages, newMessage], + updatedAt: Date.now(), + } + : prev + ); handleSubmitWithContent(newMessage.content); }; const clearChat = () => { - setMessages([]); + setCurrentSession(prev => + prev ? { ...prev, messages: [], updatedAt: Date.now() } : prev + ); setError(null); }; + const handleSessionChange = (session: ChatSession) => { + setCurrentSession(session); + setError(null); + }; + + const handleNewSession = () => { + const defaultModel = + currentSession?.selectedModel || + (models.length > 0 ? models[0].identifier : ""); + + const newSession: ChatSession = { + ...SessionUtils.createDefaultSession(), + selectedModel: defaultModel, + systemMessage: + currentSession?.systemMessage || "You are a helpful assistant.", + messages: [], + updatedAt: Date.now(), + createdAt: Date.now(), + }; + setCurrentSession(newSession); + SessionUtils.saveCurrentSession(newSession); + }; + return ( -
-
-

Chat Playground (Completions)

-
- - +
+ {/* Header */} +
+
+

Chat Playground

+
+ + +
+ {/* Main Two-Column Layout */} +
+ {/* Left Column - Configuration Panel */} +
+

+ Settings +

- {modelsError && ( -
-

{modelsError}

+ {/* Model Configuration */} +
+

+ Model Configuration +

+
+
+ + + {modelsError && ( +

{modelsError}

+ )} +
+ +
+ +