feat(thread): cancel using prompt spinning #1243 (#1251)

Opla · Sep 12, 2024 · 8a739f4 · 8a739f4
1 parent a12ae9f
commit 8a739f4
Show file tree

Hide file tree

Showing 3 changed files with 95 additions and 44 deletions.
diff --git a/webapp/features/Threads/Conversation/ConversationContext.tsx b/webapp/features/Threads/Conversation/ConversationContext.tsx
@@ -67,7 +67,7 @@ type Context = {
     submit: boolean,
   ) => Promise<void>;
   handleStartMessageEdit: (messageId: string, index: number) => void;
-  handleCancelSending: (messageId: string) => void;
+  handleCancelSending: (conversationId: string, messageId: string | undefined) => void;
 };
 
 type ConversationProviderProps = {
@@ -110,6 +110,7 @@ function ConversationProvider({
   const {
     conversations,
     messages: messagesCache,
+    updateConversationMessages,
     updateMessagesAndConversation,
   } = useThreadStore();
   const { parseAndValidatePrompt, clearPrompt } = useContext(PromptContext) || {};
@@ -580,37 +581,61 @@ function ConversationProvider({
   );
 
   const handleCancelSending = useCallback(
-    async (messageId: string) => {
-      if (selectedConversation && selectedModelId) {
-        try {
-          await cancelSending(
-            messageId,
-            selectedConversation,
-            selectedModelId,
-            assistant,
-            modelStorage,
-            activeService,
-          );
-        } catch (e) {
-          logger.error(e);
-          const conversationMessages = messagesCache[selectedConversation.id];
-          const previousMessage = conversationMessages.find((m) => m.id === messageId);
-          if (!previousMessage) {
-            logger.error(
-              "Can't find previous message",
-              messageId,
+    async (cId: string, messageId: string | undefined) => {
+      const cancelMessageSending = async (mId: string) => {
+        if (selectedConversation && selectedModelId) {
+          try {
+            await cancelSending(
+              mId,
               selectedConversation,
-              conversationMessages,
+              selectedModelId,
+              assistant,
+              modelStorage,
+              activeService,
+            );
+          } catch (e) {
+            logger.error(e);
+            const conversationMessages = messagesCache[selectedConversation.id];
+            const previousMessage = conversationMessages.find((m) => m.id === mId);
+            if (!previousMessage) {
+              logger.error(
+                "Can't find previous message",
+                mId,
+                selectedConversation,
+                conversationMessages,
+              );
+              return;
+            }
+            const updatedMessage = changeMessageContent(
+              previousMessage,
+              t('Cancelled'),
+              t('Cancelled'),
+              MessageStatus.Delivered,
+            );
+            await updateConversationMessages(
+              cId,
+              messagesCache[cId].map((m) => (m.id === mId ? updatedMessage : m)),
             );
-            return;
           }
-          changeMessageContent(
-            previousMessage,
-            t('Cancelled'),
-            t('Cancelled'),
-            MessageStatus.Delivered,
-          );
         }
+      };
+
+      if (messageId) {
+        await cancelMessageSending(messageId);
+      }
+      if (messagesCache[cId]) {
+        const promises = messagesCache[cId]
+          .map((message) => {
+            if (
+              messageId !== message.id &&
+              (message.status === 'pending' || message.status === 'stream')
+            ) {
+              return cancelMessageSending(message.id);
+            }
+            return undefined;
+          })
+          .filter((p) => !!p);
+        await Promise.all(promises);
       }
     },
     [
@@ -621,6 +646,7 @@ function ConversationProvider({
       selectedModelId,
       t,
       messagesCache,
+      updateConversationMessages,
     ],
   );
 

diff --git a/webapp/features/Threads/Conversation/ConversationList.tsx b/webapp/features/Threads/Conversation/ConversationList.tsx
@@ -36,7 +36,7 @@ type ConversationListProps = {
   onChangeMessageContent: (m: Message, newContent: string, submit: boolean) => void;
   onStartMessageEdit: (messageId: string, index: number) => void;
   onCopyMessage: (messageId: string, state: boolean) => void;
-  onCancelSending: (messageId: string) => void;
+  onCancelSending: (conversationId: string, messageId: string) => void;
 };
 
 function ConversationList({
@@ -102,7 +102,7 @@ function ConversationList({
               }}
               onCopyMessage={onCopyMessage}
               onCancelSending={() => {
-                onCancelSending(m.id);
+                onCancelSending(conversation.id, m.id);
               }}
             />
           ))}

diff --git a/webapp/features/Threads/Prompt/index.tsx b/webapp/features/Threads/Prompt/index.tsx
@@ -56,7 +56,12 @@ export default function Prompt({
   isModelLoading,
 }: PromptProps) {
   const { t } = useTranslation();
-  const { isProcessing, errorMessages, handleSendMessage: sendMessage } = useConversationContext();
+  const {
+    isProcessing,
+    errorMessages,
+    handleSendMessage: sendMessage,
+    handleCancelSending,
+  } = useConversationContext();
   const textareaRef = useRef<HTMLTextAreaElement>(null);
   const [needFocus, setNeedFocus] = useState(false);
   const { usage, changedPrompt, conversationPrompt, setChangedPrompt, tokenValidator } =
@@ -67,6 +72,11 @@ export default function Prompt({
   const { conversations, updateConversations, messages, updateConversationMessages } =
     useThreadStore();
 
+  const handleCancel = (e: MouseEvent) => {
+    e.preventDefault();
+    handleCancelSending(conversationId, undefined);
+  };
+
   const handleSendMessage = (e: MouseEvent) => {
     e.preventDefault();
     if (prompt) {
@@ -172,7 +182,8 @@ export default function Prompt({
     return undefined;
   }
 
-  let isLoading = conversationId ? isProcessing[conversationId] || false : false;
+  const isConnected = conversationId ? isProcessing[conversationId] || false : false;
+  let isLoading = false;
   let placeholder;
   if (isModelLoading || isModelLoading === undefined) {
     isLoading = true;
@@ -228,36 +239,50 @@ export default function Prompt({
               onValueChange={handleValueChange}
               onFocus={handleFocus}
               onKeyDown={handleKeypress}
-              disabled={isLoading || disabled}
+              disabled={isLoading || isConnected || disabled}
             />
           </PromptCommands>
           <Tooltip>
             <TooltipTrigger asChild>
               <Button
-                disabled={isLoading || prompt?.raw?.length === 0}
+                disabled={isLoading || (prompt?.raw?.length === 0 && !isConnected)}
                 type="button"
                 aria-label={t('Send')}
-                onClick={handleSendMessage}
+                onClick={isConnected ? handleCancel : handleSendMessage}
                 className="ml-2"
                 size="icon"
                 variant="outline"
               >
-                {isLoading ? (
+                {isLoading || isConnected ? (
                   <Loader2 strokeWidth={1.5} className="loading-icon h-4 w-4 animate-spin" />
                 ) : (
                   <SendHorizontal className="strokeWidth={1.5} h-4 w-4" />
                 )}
               </Button>
             </TooltipTrigger>
             <TooltipContent side="right" sideOffset={12} className="mt-1">
-              <div className="flex w-full flex-row items-center justify-between gap-2 pb-2">
-                <p>{t(shortcutSend.description)}</p>
-                <ShortcutBadge command={shortcutSend.command} />
-              </div>
-              <div className="flex w-full flex-row items-center justify-between gap-2">
-                <p>{t(shortcutNewLine.description)}</p>
-                <ShortcutBadge command={shortcutNewLine.command} />
-              </div>
+              {!isLoading && !isConnected && (
+                <>
+                  <div className="flex w-full flex-row items-center justify-between gap-2 pb-2">
+                    <p>{t(shortcutSend.description)}</p>
+                    <ShortcutBadge command={shortcutSend.command} />
+                  </div>
+                  <div className="flex w-full flex-row items-center justify-between gap-2">
+                    <p>{t(shortcutNewLine.description)}</p>
+                    <ShortcutBadge command={shortcutNewLine.command} />
+                  </div>
+                </>
+              )}
+              {isLoading && (
+                <div className="flex w-full flex-row items-center justify-between gap-2 pb-2">
+                  <p>{t('Please wait...')}</p>
+                </div>
+              )}
+              {isProcessing && (
+                <div className="flex w-full flex-row items-center justify-between gap-2 pb-2">
+                  <p>{t('Cancel')}</p>
+                </div>
+              )}
             </TooltipContent>
           </Tooltip>
         </div>