From b6be165457b637c645290008f98ec37d8edf46bf Mon Sep 17 00:00:00 2001 From: Rob Gordon Date: Wed, 9 Oct 2024 15:32:37 -0400 Subject: [PATCH] Ability to stop AI Process (#750) --- app/src/components/AiToolbar.tsx | 48 +++++++++++++------- app/src/components/ConvertOnPasteOverlay.tsx | 2 +- app/src/lib/runAi.ts | 43 ++++++++++-------- app/src/lib/usePromptStore.ts | 40 +++++++++++----- 4 files changed, 83 insertions(+), 50 deletions(-) diff --git a/app/src/components/AiToolbar.tsx b/app/src/components/AiToolbar.tsx index 2c1fae57..742ac0e8 100644 --- a/app/src/components/AiToolbar.tsx +++ b/app/src/components/AiToolbar.tsx @@ -1,5 +1,5 @@ import { Button2, IconButton2, Textarea } from "../ui/Shared"; -import { CaretDown, CaretUp, MagicWand } from "phosphor-react"; +import { CaretDown, CaretUp, MagicWand, Stop } from "phosphor-react"; import cx from "classnames"; import { t, Trans } from "@lingui/macro"; import { createExamples } from "../pages/createExamples"; @@ -44,7 +44,7 @@ export function AiToolbar() { const isOpen = usePromptStore((state) => state.isOpen); const currentMode = usePromptStore((state) => state.mode); const isRunning = usePromptStore((state) => state.isRunning); - const runAiWithStore = useRunAiWithStore(); + const { runAi, cancelAi } = useRunAiWithStore(); const diff = usePromptStore((state) => state.diff); const toggleOpen = () => setIsOpen(!isOpen); @@ -109,19 +109,33 @@ export function AiToolbar() { )} {!showAcceptDiffButton ? ( - - {!isOpen ? ( - - ) : ( - - )} - + <> +
+ + {!isOpen ? ( + + ) : ( + + )} + + {isRunning ? ( + + + + ) : null} +
+ ) : (
@@ -150,7 +164,7 @@ export function AiToolbar() { onKeyDown={(e) => { if (e.key === "Enter" && !e.shiftKey) { e.preventDefault(); - runAiWithStore(); + runAi(); } }} /> @@ -160,7 +174,7 @@ export function AiToolbar() { size="xs" className="dark:bg-purple-700 dark:hover:bg-purple-600 dark:text-purple-100" disabled={isRunning} - onClick={runAiWithStore} + onClick={runAi} data-session-activity={`Run AI: ${currentMode}`} > {!isRunning ? t`Submit` : "..."} diff --git a/app/src/components/ConvertOnPasteOverlay.tsx b/app/src/components/ConvertOnPasteOverlay.tsx index b02ad270..7b5e9f3c 100644 --- a/app/src/components/ConvertOnPasteOverlay.tsx +++ b/app/src/components/ConvertOnPasteOverlay.tsx @@ -50,7 +50,7 @@ export function ConvertOnPasteOverlay() { * This is positioned at the bottom across the whole screen */ function Overlay() { - const runAiWithStore = useRunAiWithStore(); + const { runAi: runAiWithStore } = useRunAiWithStore(); const isRunning = usePromptStore((s) => s.isRunning); const pasted = useEditorStore((s) => s.userPasted); return ( diff --git a/app/src/lib/runAi.ts b/app/src/lib/runAi.ts index 97125b19..8505deb5 100644 --- a/app/src/lib/runAi.ts +++ b/app/src/lib/runAi.ts @@ -39,13 +39,15 @@ function writeToEditor(text: string) { * Runs an AI endpoint and streams the response back into the editor */ export async function runAi({ - prompt, endpoint, + prompt, sid, + signal, }: { + endpoint: string; prompt: string; - endpoint: "prompt" | "convert" | "edit"; sid?: string; + signal?: AbortSignal; }) { let accumulated = ""; @@ -65,6 +67,7 @@ export async function runAi({ Authorization: sid ? `Bearer ${sid}` : "", }, body: JSON.stringify({ prompt, document: useDoc.getState().text }), + signal, }) .then((response) => { if (response.ok && response.body) { @@ -81,8 +84,8 @@ export async function runAi({ if (endpoint === "edit") { // No setup... } else { - editor.pushUndoStop(); // Make sure you can undo the changes - writeToEditor(""); // Clear the editor content + editor.pushUndoStop(); + writeToEditor(""); } const processText = ({ @@ -91,42 +94,42 @@ export async function runAi({ }: ReadableStreamReadResult): Promise => { if (done) { setLastResult(accumulated); + resolve(accumulated); return Promise.resolve(); } const text = decoder.decode(value, { stream: true }); accumulated += text; - // If we are editing, we want to set the diff if (endpoint === "edit") { setDiff(accumulated); } else { - // If we are not editing, we want to write the text to the editor writeToEditor(accumulated); } - return reader.read().then(processText); + return reader.read().then(processText).catch(reject); }; - reader - .read() - .then(processText) - .finally(() => { - editor.pushUndoStop(); - resolve(accumulated); - }); + reader.read().then(processText).catch(reject); } else { if (response.status === 429) { reject(new Error(RATE_LIMIT_EXCEEDED)); + } else { + reject( + new Error( + t`Sorry, there was an error converting the text to a flowchart. Try again later.` + ) + ); } - reject( - new Error( - t`Sorry, there was an error converting the text to a flowchart. Try again later.` - ) - ); } }) - .catch(reject); + .catch((error) => { + if (error.name === "AbortError") { + reject(new Error("Operation canceled")); + } else { + reject(error); + } + }); }); } diff --git a/app/src/lib/usePromptStore.ts b/app/src/lib/usePromptStore.ts index 945e6d2d..bdf6b947 100644 --- a/app/src/lib/usePromptStore.ts +++ b/app/src/lib/usePromptStore.ts @@ -1,7 +1,6 @@ import { create } from "zustand"; import { RATE_LIMIT_EXCEEDED, runAi } from "./runAi"; -import { isError } from "./helpers"; -import { useCallback, useContext } from "react"; +import { useCallback, useContext, useState } from "react"; import { useHasProAccess } from "./hooks"; import { showPaywall } from "./usePaywallModalStore"; import { t } from "@lingui/macro"; @@ -90,11 +89,14 @@ export function useRunAiWithStore() { const hasProAccess = useHasProAccess(); const customer = useContext(AppContext).customer; const sid = customer?.subscription?.id; + const [abortController, setAbortController] = + useState(null); const handleError = useCallback( (error: Error) => { - if (!hasProAccess && error.message === RATE_LIMIT_EXCEEDED) { - // Show paywall + if (error.name === "AbortError") { + setError(t`Operation canceled`); + } else if (!hasProAccess && error.message === RATE_LIMIT_EXCEEDED) { showPaywall({ title: t`Get Unlimited AI Requests`, content: t`You've used all your free AI conversions. Upgrade to Pro for unlimited AI use, custom themes, private sharing, and more. Keep creating amazing flowcharts effortlessly!`, @@ -113,25 +115,28 @@ export function useRunAiWithStore() { [hasProAccess] ); - return useCallback(() => { + const runAiCallback = useCallback(() => { const store = usePromptStore.getState(); if (store.isRunning) return; - // close the toolbar setIsOpen(false); startConvert(); - // If we're creating, we need to unfreeze the editor if (store.mode === "convert" || store.mode === "prompt") { unfreezeDoc(); } - runAi({ endpoint: store.mode, prompt: store.currentText, sid }) - .catch((err) => { - if (isError(err)) handleError(err); - }) + const newAbortController = new AbortController(); + setAbortController(newAbortController); + + runAi({ + endpoint: store.mode, + prompt: store.currentText, + sid, + signal: newAbortController.signal, + }) + .catch(handleError) .then((result) => { - // Just in case there is an error, run repair text on the result if (result) { const text = repairText(result); if (text) { @@ -142,6 +147,17 @@ export function useRunAiWithStore() { .finally(() => { stopConvert(); useEditorStore.setState({ userPasted: "" }); + setAbortController(null); }); }, [handleError, sid]); + + const cancelAi = useCallback(() => { + if (abortController) { + abortController.abort(); + stopConvert(); + setAbortController(null); + } + }, [abortController]); + + return { runAi: runAiCallback, cancelAi }; }