feat (ai/ui): add stop() helper to useAssistant (#1524)

Author: lgrammel

.changeset/light-chairs-clap.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+feat (ai/ui): add stop() helper to useAssistant (important: AssistantResponse now requires OpenAI SDK 4.42+)

examples/next-openai/app/api/assistant/route.ts

@@ -27,22 +27,30 @@ export async function POST(req: Request) {
   const threadId = input.threadId ?? (await openai.beta.threads.create({})).id;
 
   // Add a message to the thread
-  const createdMessage = await openai.beta.threads.messages.create(threadId, {
-    role: 'user',
-    content: input.message,
-  });
+  const createdMessage = await openai.beta.threads.messages.create(
+    threadId,
+    {
+      role: 'user',
+      content: input.message,
+    },
+    { signal: req.signal },
+  );
 
   return AssistantResponse(
     { threadId, messageId: createdMessage.id },
     async ({ forwardStream, sendDataMessage }) => {
       // Run the assistant on the thread
-      const runStream = openai.beta.threads.runs.createAndStream(threadId, {
-        assistant_id:
-          process.env.ASSISTANT_ID ??
-          (() => {
-            throw new Error('ASSISTANT_ID is not set');
-          })(),
-      });
+      const runStream = openai.beta.threads.runs.stream(
+        threadId,
+        {
+          assistant_id:
+            process.env.ASSISTANT_ID ??
+            (() => {
+              throw new Error('ASSISTANT_ID is not set');
+            })(),
+        },
+        { signal: req.signal },
+      );
 
       // forward run status would stream message deltas
       let runResult = await forwardStream(runStream);
@@ -108,6 +116,7 @@ export async function POST(req: Request) {
             threadId,
             runResult.id,
             { tool_outputs },
+            { signal: req.signal },
           ),
         );
       }

examples/next-openai/app/assistant/page.tsx

@@ -13,8 +13,15 @@ const roleToColorMap: Record<Message['role'], string> = {
 };
 
 export default function Chat() {
-  const { status, messages, input, submitMessage, handleInputChange, error } =
-    useAssistant({ api: '/api/assistant' });
+  const {
+    status,
+    messages,
+    input,
+    submitMessage,
+    handleInputChange,
+    error,
+    stop,
+  } = useAssistant({ api: '/api/assistant' });
 
   // When status changes to accepting messages, focus the input:
   const inputRef = useRef<HTMLInputElement>(null);
@@ -64,12 +71,19 @@ export default function Chat() {
         <input
           ref={inputRef}
           disabled={status !== 'awaiting_message'}
-          className="fixed bottom-0 w-full max-w-md p-2 mb-8 border border-gray-300 rounded shadow-xl"
+          className="fixed w-full max-w-md p-2 mb-8 border border-gray-300 rounded shadow-xl bottom-14 ax-w-md"
           value={input}
           placeholder="What is the temperature in the living room?"
           onChange={handleInputChange}
         />
       </form>
+
+      <button
+        className="fixed bottom-0 w-full max-w-md p-2 mb-8 text-white bg-red-500 rounded-lg"
+        onClick={stop}
+      >
+        Stop
+      </button>
     </div>
   );
 }

examples/next-openai/package.json

@@ -12,7 +12,7 @@
     "@ai-sdk/openai": "latest",
     "ai": "latest",
     "next": "latest",
-    "openai": "4.29.0",
+    "openai": "4.42.0",
     "react": "18.2.0",
     "react-dom": "^18.2.0",
     "zod": "3.23.4"

packages/core/package.json

@@ -117,7 +117,7 @@
     "jsdom": "^23.0.0",
     "langchain": "0.0.196",
     "msw": "2.0.9",
-    "openai": "4.29.0",
+    "openai": "4.42.0",
     "react-dom": "^18.2.0",
     "react-server-dom-webpack": "18.3.0-canary-eb33bd747-20240312",
     "solid-js": "^1.8.7",
@@ -127,6 +127,7 @@
     "zod": "3.22.4"
   },
   "peerDependencies": {
+    "openai": "^4.42.0",
     "react": "^18.2.0",
     "solid-js": "^1.7.7",
     "svelte": "^3.0.0 || ^4.0.0",
@@ -148,6 +149,9 @@
     },
     "zod": {
       "optional": true
+    },
+    "openai": {
+      "optional": true
     }
   },
   "engines": {

packages/core/react/use-assistant.ts

@@ -1,10 +1,11 @@
 /* eslint-disable react-hooks/rules-of-hooks */
 
-import { useState } from 'react';
-
+import { isAbortError } from '@ai-sdk/provider-utils';
+import { useCallback, useRef, useState } from 'react';
 import { generateId } from '../shared/generate-id';
 import { readDataStream } from '../shared/read-data-stream';
 import { CreateMessage, Message } from '../shared/types';
+import { abort } from 'node:process';
 
 export type AssistantStatus = 'in_progress' | 'awaiting_message';
 
@@ -42,6 +43,11 @@ export type UseAssistantHelpers = {
     },
   ) => Promise<void>;
 
+  /**
+Abort the current request immediately, keep the generated tokens if any.
+   */
+  stop: () => void;
+
   /**
    * setState-powered method to update the input value.
    */
@@ -135,6 +141,16 @@ export function useAssistant({
     setInput(event.target.value);
   };
 
+  // Abort controller to cancel the current API call.
+  const abortControllerRef = useRef<AbortController | null>(null);
+
+  const stop = useCallback(() => {
+    if (abortControllerRef.current) {
+      abortControllerRef.current.abort();
+      abortControllerRef.current = null;
+    }
+  }, []);
+
   const append = async (
     message: Message | CreateMessage,
     requestOptions?: {
@@ -153,10 +169,15 @@ export function useAssistant({
 
     setInput('');
 
+    const abortController = new AbortController();
+
     try {
+      abortControllerRef.current = abortController;
+
       const result = await fetch(api, {
         method: 'POST',
         credentials,
+        signal: abortController.signal,
         headers: { 'Content-Type': 'application/json', ...headers },
         body: JSON.stringify({
           ...body,
@@ -240,14 +261,21 @@ export function useAssistant({
         }
       }
     } catch (error) {
+      // Ignore abort errors as they are expected when the user cancels the request:
+      if (isAbortError(error) && abortController.signal.aborted) {
+        abortControllerRef.current = null;
+        return;
+      }
+
       if (onError && error instanceof Error) {
         onError(error);
       }
 
       setError(error as Error);
+    } finally {
+      abortControllerRef.current = null;
+      setStatus('awaiting_message');
     }
-
-    setStatus('awaiting_message');
   };
 
   const submitMessage = async (
@@ -276,6 +304,7 @@ export function useAssistant({
     submitMessage,
     status,
     error,
+    stop,
   };
 }
 

packages/core/streams/assistant-response.ts

@@ -1,4 +1,4 @@
-import { AssistantStream } from 'openai/lib/AssistantStream';
+import { type AssistantStream } from 'openai/lib/AssistantStream';
 import { Run } from 'openai/resources/beta/threads/runs/runs';
 import { formatStreamPart } from '../shared/stream-parts';
 import { AssistantMessage, DataMessage } from '../shared/types';