feat: add streaming

2026-04-29 13:35:21 +02:00
parent d04fbbe73c
commit 18f293ce67
2 changed files with 230 additions and 63 deletions
@@ -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,32 +76,41 @@ 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',
-        {
+      if (settings.stream) {
+        await handleStream(endpoint, body, newMessages);
+      } else {
+        await handleNonStream(endpoint, body);
+      }
+    } catch (err: any) {
+      console.error(err);
+      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;
-        }
-
+      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);
    }

@@ -95,29 +118,45 @@ export default function ChatUI({
    if (assistant) {
      setMessages((prev) => [...prev, assistant]);
    }
+  };

-    } catch (err: any) {
-      console.error(err);
+  // ── Streaming path ────────────────────────────────────────────────────────

-      const message =
-        err?.message || 'Something went wrong while sending request';
+  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: '' }]);

-      setMessages((prev) => [
-        ...prev,
-        {
-          role: 'assistant',
-          content: `❌ ${message}`,
-        },
-      ]);
+    const response = await fetchChatStream(
+      endpoint,
+      body,
+      session?.accessToken as string,
+    );

-      setNotification(message);
+    for await (const { content, done } of readStreamChunks(response)) {
+      if (done) break;
+      if (!content) continue;

-      setTimeout(() => setNotification(null), 3000);
-    } finally {
-      setLoading(false);
+      setMessages((prev) => {
+        const next = [...prev];
+        const target = next[placeholderIndex];
+        if (target) {
+          next[placeholderIndex] = {
+            ...target,
+            content: target.content + content,
+          };
+        }
+        return next;
+      });
    }
  };

+  // ── Render ────────────────────────────────────────────────────────────────
+
  return (
    <div className="flex flex-1 bg-primary text-primary transition-theme relative">
      {notification && (
@@ -146,9 +185,18 @@ export default function ChatUI({
            {messages.map((msg, i) => (
              <div key={i} className="text-sm">
                <b>{msg.role === 'user' ? 'You' : 'AI'}:</b> {msg.content}
+                {/* Blinking cursor shown on the last message while streaming */}
+                {loading &&
+                  settings.stream &&
+                  i === messages.length - 1 &&
+                  msg.role === ASSISTANT_ROLE && (
+                    <span className="inline-block w-[2px] h-[1em] bg-current ml-0.5 align-middle animate-pulse" />
+                  )}
              </div>
            ))}
-            {loading && <div className="text-xs text-muted">Thinking...</div>}
+            {loading && !settings.stream && (
+              <div className="text-xs text-muted">Thinking...</div>
+            )}
          </div>
        </main>

@@ -156,17 +204,14 @@ export default function ChatUI({
        <div
          className={`
            transition-all duration-500 ease-in-out
-            ${messages.length === 0
+            ${
+              messages.length === 0
                ? 'absolute inset-0 flex items-center justify-center pointer-events-none'
-              : 'p-4'}
-          `}
-        >
-          <div
-            className={`
-              w-full max-w-2xl mx-auto pointer-events-auto
-              transition-all duration-500 ease-in-out
+                : 'p-4'
+            }
          `}
        >
+          <div className="w-full max-w-2xl mx-auto pointer-events-auto transition-all duration-500 ease-in-out">
            {messages.length === 0 && (
              <p className="text-center text-primary text-sm mb-3 opacity-60">
                Start a conversation…
@@ -175,8 +220,9 @@ export default function ChatUI({
            <input
              value={input}
              onChange={(e) => 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..."
            />
          </div>
@@ -0,0 +1,121 @@
+/**
+ * chat.stream.ts
+ *
+ * Helpers for consuming a streaming chat-completions response (SSE format).
+ * The server must send chunks as:
+ *   data: {"choices":[{"delta":{"content":"..."}}]}\n\n
+ *   data: [DONE]\n\n
+ */
+
+export type StreamChunk = {
+  content: string;
+  done: boolean;
+};
+
+/**
+ * Async-generator that yields text deltas from a streaming fetch response.
+ * Handles both pure SSE (`data: ...`) and raw JSON-lines.
+ *
+ * @example
+ * for await (const { content, done } of readStreamChunks(response)) {
+ *   if (done) break;
+ *   setPartial(prev => prev + content);
+ * }
+ */
+export async function* readStreamChunks(
+  response: Response,
+): AsyncGenerator<StreamChunk> {
+  if (!response.body) {
+    throw new Error('Response body is null — cannot stream.');
+  }
+
+  const reader = response.body.getReader();
+  const decoder = new TextDecoder();
+  let buffer = '';
+
+  while (true) {
+    const { value, done } = await reader.read();
+
+    if (done) {
+      // Flush any remaining buffer content
+      if (buffer.trim()) {
+        yield* parseBuffer(buffer);
+      }
+      break;
+    }
+
+    buffer += decoder.decode(value, { stream: true });
+
+    // Split on double-newline (SSE event boundary) or single newline (JSON-lines)
+    const parts = buffer.split(/\n\n|\n/);
+
+    // Keep the last (potentially incomplete) chunk in the buffer
+    buffer = parts.pop() ?? '';
+
+    for (const part of parts) {
+      yield* parseBuffer(part);
+    }
+  }
+}
+
+/**
+ * Parses a single SSE line or JSON line and yields a StreamChunk.
+ */
+function* parseBuffer(raw: string): Generator<StreamChunk> {
+  const line = raw.trim();
+  if (!line) return;
+
+  // SSE lines start with "data: "
+  const jsonStr = line.startsWith('data: ') ? line.slice(6) : line;
+
+  if (jsonStr === '[DONE]') {
+    yield { content: '', done: true };
+    return;
+  }
+
+  try {
+    const parsed = JSON.parse(jsonStr);
+    const delta = parsed?.choices?.[0]?.delta?.content;
+    if (typeof delta === 'string') {
+      yield { content: delta, done: false };
+    }
+  } catch {
+    // Not valid JSON — silently skip (could be a comment or keep-alive ping)
+  }
+}
+
+/**
+ * Builds the streaming fetch request and returns the raw Response.
+ * Throws on non-2xx status with a descriptive message.
+ */
+export async function fetchChatStream(
+  url: string,
+  body: object,
+  token: string,
+): Promise<Response> {
+  const res = await fetch(url, {
+    method: 'POST',
+    headers: {
+      'Content-Type': 'application/json',
+      Authorization: `Bearer ${token}`,
+    },
+    body: JSON.stringify(body),
+  });
+
+  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 {
+      const data = await res.json().catch(() => null);
+      message = data?.message ?? data?.error ?? message;
+    }
+
+    throw new Error(message);
+  }
+
+  return res;
+}