diff --git a/src/components/chat/ChatUI.tsx b/src/components/chat/ChatUI.tsx index 80e2dc4..cac6448 100644 --- a/src/components/chat/ChatUI.tsx +++ b/src/components/chat/ChatUI.tsx @@ -8,6 +8,7 @@ import { Session } from 'next-auth'; import { Message, ModelInfo } from '@/types/chat-api'; import { ChatRequest } from '@/types/chat-api'; import { Settings } from './chat.panels'; +import { fetchChatStream, readStreamChunks } from './stream'; type Props = { session: Session; @@ -16,6 +17,9 @@ type Props = { setSelectedModel: (v: string) => void; }; +// Sentinel role used while the assistant message is being streamed in +const ASSISTANT_ROLE = 'assistant' as const; + export default function ChatUI({ session, models, @@ -39,16 +43,26 @@ export default function ChatUI({ keepAlive: '5m', }); + const showError = (message: string) => { + setMessages((prev) => [ + ...prev, + { role: ASSISTANT_ROLE, content: `❌ ${message}` }, + ]); + setNotification(message); + setTimeout(() => setNotification(null), 3000); + }; + const sendMessage = async () => { if (!input.trim() || !selectedModel) return; - const userMessage = { role: 'user' as const, content: input }; - const newMessages = [...messages, userMessage]; + + const userMessage: Message = { role: 'user', content: input }; + const newMessages: Message[] = [...messages, userMessage]; setMessages(newMessages); setInput(''); setLoading(true); const body: ChatRequest = { - model: selectedModel!, + model: selectedModel, messages: newMessages, temperature: settings.temp, top_p: settings.topP, @@ -62,62 +76,87 @@ export default function ChatUI({ stop: settings.stop, }; + const endpoint = process.env.NEXT_PUBLIC_API_URL + '/chat/completions'; + try { - const res = await fetch( - process.env.NEXT_PUBLIC_API_URL + '/chat/completions', - { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - Authorization: `Bearer ${session?.accessToken}`, - }, - body: JSON.stringify(body), - }, - ); - - const data = await res.json().catch(() => null); - - if (!res.ok) { - let message = "Request failed"; - - if (res.status === 500) { - message = "Server error (likely model too large or OOM)"; - } else if (res.status === 401) { - message = "Unauthorized"; - } else { - message = data?.message || data?.error || message; - } - - throw new Error(message); + if (settings.stream) { + await handleStream(endpoint, body, newMessages); + } else { + await handleNonStream(endpoint, body); } - - const assistant = data?.choices?.[0]?.message; - if (assistant) { - setMessages((prev) => [...prev, assistant]); - } - } catch (err: any) { console.error(err); - - const message = - err?.message || 'Something went wrong while sending request'; - - setMessages((prev) => [ - ...prev, - { - role: 'assistant', - content: `❌ ${message}`, - }, - ]); - - setNotification(message); - - setTimeout(() => setNotification(null), 3000); + showError(err?.message ?? 'Something went wrong while sending request'); } finally { setLoading(false); } }; + // ── Non-streaming path ──────────────────────────────────────────────────── + + const handleNonStream = async (endpoint: string, body: ChatRequest) => { + const res = await fetch(endpoint, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: `Bearer ${session?.accessToken}`, + }, + body: JSON.stringify(body), + }); + + const data = await res.json().catch(() => null); + + if (!res.ok) { + let message = 'Request failed'; + if (res.status === 500) message = 'Server error (likely model too large or OOM)'; + else if (res.status === 401) message = 'Unauthorized'; + else message = data?.message ?? data?.error ?? message; + throw new Error(message); + } + + const assistant = data?.choices?.[0]?.message; + if (assistant) { + setMessages((prev) => [...prev, assistant]); + } + }; + + // ── Streaming path ──────────────────────────────────────────────────────── + + const handleStream = async ( + endpoint: string, + body: ChatRequest, + priorMessages: Message[], + ) => { + // Append an empty assistant message that we'll fill in as tokens arrive + const placeholderIndex = priorMessages.length; // index in the upcoming state + setMessages((prev) => [...prev, { role: ASSISTANT_ROLE, content: '' }]); + + const response = await fetchChatStream( + endpoint, + body, + session?.accessToken as string, + ); + + for await (const { content, done } of readStreamChunks(response)) { + if (done) break; + if (!content) continue; + + setMessages((prev) => { + const next = [...prev]; + const target = next[placeholderIndex]; + if (target) { + next[placeholderIndex] = { + ...target, + content: target.content + content, + }; + } + return next; + }); + } + }; + + // ── Render ──────────────────────────────────────────────────────────────── + return (
Start a conversation… @@ -175,8 +220,9 @@ export default function ChatUI({ setInput(e.target.value)} - onKeyDown={(e) => e.key === 'Enter' && sendMessage()} - className="w-full p-3 rounded-xl bg-black/20 border border-muted outline-none" + onKeyDown={(e) => e.key === 'Enter' && !loading && sendMessage()} + disabled={loading} + className="w-full p-3 rounded-xl bg-black/20 border border-muted outline-none disabled:opacity-50" placeholder="Message..." />