[Feature] Stop model generating

mlc-ai · May 15, 2024 · 7325284 · 7325284
1 parent aa44981
commit 7325284
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 24 deletions.
diff --git a/app/client/api.ts b/app/client/api.ts
@@ -59,5 +59,6 @@ export abstract class LLMApi {
   abstract chat(options: ChatOptions): Promise<void>;
   abstract usage(): Promise<LLMUsage>;
   abstract models(): Promise<LLMModel[]>;
+  abstract abort(): Promise<void>;
   abstract clear(): void;
 }
diff --git a/app/client/webllm.ts b/app/client/webllm.ts
@@ -92,6 +92,10 @@ export class WebLLMApi implements LLMApi {
     }
   }
 
+  async abort() {
+    await this.engine?.interruptGenerate();
+  }
+
   async usage() {
     return {
       used: 0,

diff --git a/app/components/chat.tsx b/app/components/chat.tsx
@@ -441,10 +441,6 @@ export function ChatActions(props: {
     config.update((config) => (config.theme = nextTheme));
   }
 
-  // stop all responses
-  const couldStop = ChatControllerPool.hasPending();
-  const stopAll = () => ChatControllerPool.stopAll();
-
   // switch model
   const currentModel = chatStore.currentSession().mask.modelConfig.model;
   const allModels = useAllModels();
@@ -489,13 +485,6 @@ export function ChatActions(props: {
 
   return (
     <div className={styles["chat-input-actions"]}>
-      {couldStop && (
-        <ChatAction
-          onClick={stopAll}
-          text={Locale.Chat.InputActions.Stop}
-          icon={<StopIcon />}
-        />
-      )}
       {!props.hitBottom && (
         <ChatAction
           onClick={props.scrollToBottom}
@@ -669,7 +658,7 @@ function _Chat() {
   const config = useAppConfig();
   const fontSize = config.fontSize;
 
-  const isGenerating = session.isGenerating;
+  const isStreaming = session.messages.some((m) => m.streaming);
 
   const [showExport, setShowExport] = useState(false);
 
@@ -772,7 +761,7 @@ function _Chat() {
       return;
     }
 
-    if (isGenerating) return;
+    if (isStreaming) return;
     setIsLoading(true);
     chatStore
       .onUserInput(userInput, webllm!, attachImages)
@@ -803,8 +792,8 @@ function _Chat() {
   };
 
   // stop response
-  const onUserStop = (messageId: string) => {
-    ChatControllerPool.stop(session.id, messageId);
+  const onUserStop = () => {
+    webllm?.abort();
   };
 
   // Reset session status on initial loading
@@ -1337,7 +1326,7 @@ function _Chat() {
                               <ChatAction
                                 text={Locale.Chat.Actions.Stop}
                                 icon={<StopIcon />}
-                                onClick={() => onUserStop(message.id ?? i)}
+                                onClick={() => onUserStop()}
                               />
                             ) : (
                               <>
@@ -1512,14 +1501,23 @@ function _Chat() {
               })}
             </div>
           )}
-          <IconButton
-            icon={<SendWhiteIcon />}
-            text={Locale.Chat.Send}
-            className={styles["chat-input-send"]}
-            type="primary"
-            onClick={() => doSubmit(userInput)}
-            disabled={isGenerating}
-          />
+          {isStreaming ? (
+            <IconButton
+              icon={<StopIcon />}
+              text={Locale.Chat.InputActions.Stop}
+              className={styles["chat-input-send"]}
+              type="primary"
+              onClick={() => onUserStop()}
+            />
+          ) : (
+            <IconButton
+              icon={<SendWhiteIcon />}
+              text={Locale.Chat.Send}
+              className={styles["chat-input-send"]}
+              type="primary"
+              onClick={() => doSubmit(userInput)}
+            />
+          )}
         </label>
       </div>