Merge pull request #45 from Yidadaa/bugfix-0326

v1.3 Stop and Retry Button

app/components/home.tsx

@@ -27,6 +27,7 @@ import Locale from "../locales";
 
 import dynamic from "next/dynamic";
 import { REPO_URL } from "../constant";
+import { ControllerPool } from "../requests";
 
 export function Loading(props: { noLogo?: boolean }) {
   return (
@@ -146,28 +147,67 @@ function useSubmitHandler() {
 export function Chat(props: { showSideBar?: () => void }) {
   type RenderMessage = Message & { preview?: boolean };
 
-  const session = useChatStore((state) => state.currentSession());
+  const [session, sessionIndex] = useChatStore((state) => [
+    state.currentSession(),
+    state.currentSessionIndex,
+  ]);
   const [userInput, setUserInput] = useState("");
   const [isLoading, setIsLoading] = useState(false);
   const { submitKey, shouldSubmit } = useSubmitHandler();
 
   const onUserInput = useChatStore((state) => state.onUserInput);
+
+  // submit user input
   const onUserSubmit = () => {
     if (userInput.length <= 0) return;
     setIsLoading(true);
     onUserInput(userInput).then(() => setIsLoading(false));
     setUserInput("");
   };
+
+  // stop response
+  const onUserStop = (messageIndex: number) => {
+    console.log(ControllerPool, sessionIndex, messageIndex);
+    ControllerPool.stop(sessionIndex, messageIndex);
+  };
+
+  // check if should send message
   const onInputKeyDown = (e: KeyboardEvent) => {
     if (shouldSubmit(e)) {
       onUserSubmit();
       e.preventDefault();
     }
   };
+  const onRightClick = (e: any, message: Message) => {
+    // auto fill user input
+    if (message.role === "user") {
+      setUserInput(message.content);
+    }
+
+    // copy to clipboard
+    if (selectOrCopy(e.currentTarget, message.content)) {
+      e.preventDefault();
+    }
+  };
+
+  const onResend = (botIndex: number) => {
+    // find last user input message and resend
+    for (let i = botIndex; i >= 0; i -= 1) {
+      if (messages[i].role === "user") {
+        setIsLoading(true);
+        onUserInput(messages[i].content).then(() => setIsLoading(false));
+        return;
+      }
+    }
+  };
+
+  // for auto-scroll
   const latestMessageRef = useRef<HTMLDivElement>(null);
 
-  const [hoveringMessage, setHoveringMessage] = useState(false);
+  // wont scroll while hovering messages
+  const [autoScroll, setAutoScroll] = useState(false);
 
+  // preview messages
   const messages = (session.messages as RenderMessage[])
     .concat(
       isLoading
@@ -194,10 +234,11 @@ export function Chat(props: { showSideBar?: () => void }) {
         : []
     );
 
+  // auto scroll
   useLayoutEffect(() => {
     setTimeout(() => {
       const dom = latestMessageRef.current;
-      if (dom && !isIOS() && !hoveringMessage) {
+      if (dom && !isIOS() && autoScroll) {
         dom.scrollIntoView({
           behavior: "smooth",
           block: "end",
@@ -252,15 +293,7 @@ export function Chat(props: { showSideBar?: () => void }) {
         </div>
       </div>
 
-      <div
+      <div className={styles["chat-body"]}>
-        className={styles["chat-body"]}
-        onMouseOver={() => {
-          setHoveringMessage(true);
-        }}
-        onMouseOut={() => {
-          setHoveringMessage(false);
-        }}
-      >
         {messages.map((message, i) => {
           const isUser = message.role === "user";
 
@@ -283,13 +316,20 @@ export function Chat(props: { showSideBar?: () => void }) {
                 <div className={styles["chat-message-item"]}>
                   {!isUser && (
                     <div className={styles["chat-message-top-actions"]}>
-                      {message.streaming && (
+                      {message.streaming ? (
                         <div
                           className={styles["chat-message-top-action"]}
-                          onClick={() => showToast(Locale.WIP)}
+                          onClick={() => onUserStop(i)}
                         >
                           {Locale.Chat.Actions.Stop}
                         </div>
+                      ) : (
+                        <div
+                          className={styles["chat-message-top-action"]}
+                          onClick={() => onResend(i)}
+                        >
+                          {Locale.Chat.Actions.Retry}
+                        </div>
                       )}
 
                       <div
@@ -306,11 +346,7 @@ export function Chat(props: { showSideBar?: () => void }) {
                   ) : (
                     <div
                       className="markdown-body"
-                      onContextMenu={(e) => {
+                      onContextMenu={(e) => onRightClick(e, message)}
-                        if (selectOrCopy(e.currentTarget, message.content)) {
-                          e.preventDefault();
-                        }
-                      }}
                     >
                       <Markdown content={message.content} />
                     </div>
@@ -341,6 +377,9 @@ export function Chat(props: { showSideBar?: () => void }) {
             onInput={(e) => setUserInput(e.currentTarget.value)}
             value={userInput}
             onKeyDown={(e) => onInputKeyDown(e as any)}
+            onFocus={() => setAutoScroll(true)}
+            onBlur={() => setAutoScroll(false)}
+            autoFocus
           />
           <IconButton
             icon={<SendWhiteIcon />}

app/locales/cn.ts

@@ -14,6 +14,7 @@ const cn = {
       Export: "导出聊天记录",
       Copy: "复制",
       Stop: "停止",
+      Retry: "重试",
     },
     Typing: "正在输入…",
     Input: (submitKey: string) => `输入消息，${submitKey} 发送`,

app/locales/en.ts

@@ -17,6 +17,7 @@ const en: LocaleType = {
       Export: "Export All Messages as Markdown",
       Copy: "Copy",
       Stop: "Stop",
+      Retry: "Retry",
     },
     Typing: "Typing…",
     Input: (submitKey: string) =>

app/requests.ts

@@ -60,6 +60,7 @@ export async function requestChatStream(
     modelConfig?: ModelConfig;
     onMessage: (message: string, done: boolean) => void;
     onError: (error: Error) => void;
+    onController?: (controller: AbortController) => void;
   }
 ) {
   const req = makeRequestParam(messages, {
@@ -96,12 +97,12 @@ export async function requestChatStream(
       controller.abort();
     };
 
-    console.log(res);
-
     if (res.ok) {
       const reader = res.body?.getReader();
       const decoder = new TextDecoder();
 
+      options?.onController?.(controller);
+
       while (true) {
         // handle time out, will stop if no response in 10 secs
         const resTimeoutId = setTimeout(() => finish(), TIME_OUT_MS);
@@ -146,3 +147,34 @@ export async function requestWithPrompt(messages: Message[], prompt: string) {
 
   return res.choices.at(0)?.message?.content ?? "";
 }
+
+// To store message streaming controller
+export const ControllerPool = {
+  controllers: {} as Record<string, AbortController>,
+
+  addController(
+    sessionIndex: number,
+    messageIndex: number,
+    controller: AbortController
+  ) {
+    const key = this.key(sessionIndex, messageIndex);
+    this.controllers[key] = controller;
+    return key;
+  },
+
+  stop(sessionIndex: number, messageIndex: number) {
+    const key = this.key(sessionIndex, messageIndex);
+    const controller = this.controllers[key];
+    console.log(controller);
+    controller?.abort();
+  },
+
+  remove(sessionIndex: number, messageIndex: number) {
+    const key = this.key(sessionIndex, messageIndex);
+    delete this.controllers[key];
+  },
+
+  key(sessionIndex: number, messageIndex: number) {
+    return `${sessionIndex},${messageIndex}`;
+  },
+};

app/store/app.ts

@@ -2,7 +2,11 @@ import { create } from "zustand";
 import { persist } from "zustand/middleware";
 
 import { type ChatCompletionResponseMessage } from "openai";
-import { requestChatStream, requestWithPrompt } from "../requests";
+import {
+  ControllerPool,
+  requestChatStream,
+  requestWithPrompt,
+} from "../requests";
 import { trimTopic } from "../utils";
 
 import Locale from "../locales";
@@ -296,6 +300,8 @@ export const useChatStore = create<ChatStore>()(
         // get recent messages
         const recentMessages = get().getMessagesWithMemory();
         const sendMessages = recentMessages.concat(userMessage);
+        const sessionIndex = get().currentSessionIndex;
+        const messageIndex = get().currentSession().messages.length + 1;
 
         // save user's and bot's message
         get().updateCurrentSession((session) => {
@@ -303,13 +309,16 @@ export const useChatStore = create<ChatStore>()(
           session.messages.push(botMessage);
         });
 
+        // make request
         console.log("[User Input] ", sendMessages);
         requestChatStream(sendMessages, {
           onMessage(content, done) {
+            // stream response
             if (done) {
               botMessage.streaming = false;
               botMessage.content = content;
               get().onNewMessage(botMessage);
+              ControllerPool.remove(sessionIndex, messageIndex);
             } else {
               botMessage.content = content;
               set(() => ({}));
@@ -319,6 +328,15 @@ export const useChatStore = create<ChatStore>()(
             botMessage.content += "\n\n" + Locale.Store.Error;
             botMessage.streaming = false;
             set(() => ({}));
+            ControllerPool.remove(sessionIndex, messageIndex);
+          },
+          onController(controller) {
+            // collect controller for stop/retry
+            ControllerPool.addController(
+              sessionIndex,
+              messageIndex,
+              controller
+            );
           },
           filterBot: !get().config.sendBotMessages,
           modelConfig: get().config.modelConfig,