Skip to content

Commit

Permalink
feat (ai/ui): add stop() helper to useAssistant (#1524)
Browse files Browse the repository at this point in the history
  • Loading branch information
lgrammel committed May 10, 2024
1 parent 1b9e2fb commit ceb44bc
Show file tree
Hide file tree
Showing 8 changed files with 90 additions and 29 deletions.
5 changes: 5 additions & 0 deletions .changeset/light-chairs-clap.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
'ai': patch
---

feat (ai/ui): add stop() helper to useAssistant (important: AssistantResponse now requires OpenAI SDK 4.42+)
31 changes: 20 additions & 11 deletions examples/next-openai/app/api/assistant/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -108,6 +116,7 @@ export async function POST(req: Request) {
threadId,
runResult.id,
{ tool_outputs },
{ signal: req.signal },
),
);
}
Expand Down
20 changes: 17 additions & 3 deletions examples/next-openai/app/assistant/page.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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>
);
}
2 changes: 1 addition & 1 deletion examples/next-openai/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
6 changes: 5 additions & 1 deletion packages/core/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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",
Expand All @@ -148,6 +149,9 @@
},
"zod": {
"optional": true
},
"openai": {
"optional": true
}
},
"engines": {
Expand Down
37 changes: 33 additions & 4 deletions packages/core/react/use-assistant.ts
Original file line number Diff line number Diff line change
@@ -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';

Expand Down Expand Up @@ -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.
*/
Expand Down Expand Up @@ -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?: {
Expand All @@ -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,
Expand Down Expand Up @@ -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 (
Expand Down Expand Up @@ -276,6 +304,7 @@ export function useAssistant({
submitMessage,
status,
error,
stop,
};
}

Expand Down
2 changes: 1 addition & 1 deletion packages/core/streams/assistant-response.ts
Original file line number Diff line number Diff line change
@@ -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';
Expand Down
16 changes: 8 additions & 8 deletions pnpm-lock.yaml

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit ceb44bc

Please sign in to comment.