feat: implement streaming response

components/ChatView/index.tsx

@@ -1,7 +1,7 @@
-import axios from "axios";
 import { useEffect, useRef, useState } from "react";
+import { toast } from "react-hot-toast";
 import { getAssistantById, getPromptGeneratorOfAssistant, useChatStore, useMessageStore, useConnectionStore } from "@/store";
-import { CreatorRole } from "@/types";
+import { CreatorRole, Message } from "@/types";
 import { generateUUID } from "@/utils";
 import Icon from "../Icon";
 import Header from "./Header";
@@ -42,27 +42,58 @@ const ChatView = () => {
       const promptGenerator = getPromptGeneratorOfAssistant(getAssistantById(currentChat.assistantId)!);
       prompt = promptGenerator(tables.map((table) => table.structure).join("/n"));
     }
-    const { data } = await axios.post<string>("/api/chat", {
-      messages: [
-        {
-          role: CreatorRole.System,
-          content: prompt,
-        },
-        ...messageList.map((message) => ({
-          role: message.creatorRole,
-          content: message.content,
-        })),
-      ],
+    const rawRes = await fetch("/api/chat", {
+      method: "POST",
+      body: JSON.stringify({
+        messages: [
+          {
+            role: CreatorRole.System,
+            content: prompt,
+          },
+          ...messageList.map((message) => ({
+            role: message.creatorRole,
+            content: message.content,
+          })),
+        ],
+      }),
     });
-    messageStore.addMessage({
+    setIsRequesting(false);
+
+    if (!rawRes.ok) {
+      const res = await rawRes.json();
+      toast.error(res.error.message);
+      return;
+    }
+    const data = rawRes.body;
+    if (!data) {
+      toast.error("No data return");
+      return;
+    }
+
+    const message: Message = {
       id: generateUUID(),
       chatId: currentChat.id,
       creatorId: currentChat.assistantId,
       creatorRole: CreatorRole.Assistant,
       createdAt: Date.now(),
-      content: data,
-    });
-    setIsRequesting(false);
+      content: "",
+    };
+    messageStore.addMessage(message);
+
+    const reader = data.getReader();
+    const decoder = new TextDecoder("utf-8");
+    let done = false;
+    while (!done) {
+      const { value, done: readerDone } = await reader.read();
+      if (value) {
+        const char = decoder.decode(value);
+        if (char) {
+          message.content = message.content + char;
+          messageStore.updateMessageContent(message.id, message.content);
+        }
+      }
+      done = readerDone;
+    }
   };
 
   return (

package.json

@@ -13,6 +13,7 @@
     "csstype": "^3.1.1",
     "daisyui": "^2.51.5",
     "dayjs": "^1.11.7",
+    "eventsource-parser": "^1.0.0",
     "highlight.js": "^11.7.0",
     "lodash-es": "^4.17.21",
     "marked": "^4.2.12",

pages/api/chat.ts

@@ -1,4 +1,5 @@
 import { openAIApiKey } from "@/utils/openai";
+import { createParser, ParsedEvent, ReconnectInterval } from "eventsource-parser";
 import { NextRequest } from "next/server";
 
 export const config = {
@@ -7,7 +8,7 @@ export const config = {
 
 const handler = async (req: NextRequest) => {
   const reqBody = await req.json();
-  const res = await fetch(`https://api.openai.com/v1/chat/completions`, {
+  const rawRes = await fetch(`https://api.openai.com/v1/chat/completions`, {
     headers: {
       "Content-Type": "application/json",
       Authorization: `Bearer ${openAIApiKey}`,
@@ -20,17 +21,44 @@ const handler = async (req: NextRequest) => {
       temperature: 0,
       frequency_penalty: 0.0,
       presence_penalty: 0.0,
+      stream: true,
     }),
   });
+  if (!rawRes.ok) {
+    return new Response(rawRes.body, {
+      status: rawRes.status,
+      statusText: rawRes.statusText,
+    });
+  }
 
-  const data = await res.json();
-  return new Response(JSON.stringify(data.choices[0].message?.content || ""), {
-    status: 200,
-    headers: {
-      "content-type": "application/json",
-      "cache-control": "public, s-maxage=1200, stale-while-revalidate=600",
+  const encoder = new TextEncoder();
+  const decoder = new TextDecoder();
+  const stream = new ReadableStream({
+    async start(controller) {
+      const streamParser = (event: ParsedEvent | ReconnectInterval) => {
+        if (event.type === "event") {
+          const data = event.data;
+          if (data === "[DONE]") {
+            controller.close();
+            return;
+          }
+          try {
+            const json = JSON.parse(data);
+            const text = json.choices[0].delta?.content;
+            const queue = encoder.encode(text);
+            controller.enqueue(queue);
+          } catch (e) {
+            controller.error(e);
+          }
+        }
+      };
+      const parser = createParser(streamParser);
+      for await (const chunk of rawRes.body as any) {
+        parser.feed(decoder.decode(chunk));
+      }
     },
   });
+  return new Response(stream);
 };
 
 export default handler;

pnpm-lock.yaml

@@ -18,6 +18,7 @@ specifiers:
   dayjs: ^1.11.7
   eslint: 8.20.0
   eslint-config-next: 12.2.3
+  eventsource-parser: ^1.0.0
   highlight.js: ^11.7.0
   lodash-es: ^4.17.21
   marked: ^4.2.12
@@ -42,6 +43,7 @@ dependencies:
   csstype: 3.1.1
   daisyui: 2.51.5_j7yt3jd32cwenjqavrrga47yr4
   dayjs: 1.11.7
+  eventsource-parser: 1.0.0
   highlight.js: 11.7.0
   lodash-es: 4.17.21
   marked: 4.2.12
@@ -1230,6 +1232,11 @@ packages:
     engines: {node: '>=0.10.0'}
     dev: true
 
+  /eventsource-parser/1.0.0:
+    resolution: {integrity: sha512-9jgfSCa3dmEme2ES3mPByGXfgZ87VbP97tng1G2nWwWx6bV2nYxm2AWCrbQjXToSe+yYlqaZNtxffR9IeQr95g==}
+    engines: {node: '>=14.18'}
+    dev: false
+
   /fast-deep-equal/3.1.3:
     resolution: {integrity: sha512-f3qQ9oQy9j2AhBe/H9VC91wLmKBCCU/gDOnKNAYG5hswO7BLKj09Hc5HYNz9cGI++xlpDCIgDaitVs03ATR84Q==}
     dev: true

store/message.ts

@@ -1,11 +1,13 @@
+import { uniqBy } from "lodash-es";
 import { create } from "zustand";
 import { persist } from "zustand/middleware";
-import { Message } from "@/types";
+import { Id, Message } from "@/types";
 
 interface MessageState {
   messageList: Message[];
   getState: () => MessageState;
   addMessage: (message: Message) => void;
+  updateMessageContent: (messageId: Id, content: string) => void;
   clearMessage: (filter: (message: Message) => boolean) => void;
 }
 
@@ -15,6 +17,14 @@ export const useMessageStore = create<MessageState>()(
       messageList: [],
       getState: () => get(),
       addMessage: (message: Message) => set((state) => ({ messageList: [...state.messageList, message] })),
+      updateMessageContent: (messageId: Id, content: string) => {
+        const message = get().messageList.find((message) => message.id === messageId);
+        if (!message) {
+          return;
+        }
+        message.content = content;
+        set((state) => ({ messageList: uniqBy([...state.messageList, message], (message) => message.id) }));
+      },
       clearMessage: (filter: (message: Message) => boolean) => set((state) => ({ messageList: state.messageList.filter(filter) })),
     }),
     {