AssistantResponse: specify forwardStream return type. (#1164)

vercel · Mar 15, 2024 · 42209be · 42209be
1 parent efff0d4
commit 42209be
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 5 deletions.
diff --git a/.changeset/giant-dodos-glow.md b/.changeset/giant-dodos-glow.md
@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+AssistantResponse: specify forwardStream return type.
diff --git a/docs/pages/docs/api-reference/providers/assistant-response.mdx b/docs/pages/docs/api-reference/providers/assistant-response.mdx
@@ -35,7 +35,7 @@ The process parameter is a callback in which you can run the assistant on thread
 
 It gets invoked with the following functions that you can use to send messages and data messages to the client:
 
-- `forwardStream: (stream: AssistantStream) => void`: Forwards the assistant response stream to the client.
+- `forwardStream: (stream: AssistantStream) => Run | undefined`: Forwards the assistant response stream to the client. Returns the `Run` object after it completes, or when it requires an action.
 - `sendDataMessage: (message: DataMessage) => void`: Send a data message to the client. You can use this to provide information for rendering custom UIs while the assistant is processing the thread.
 
 ## Example
@@ -100,7 +100,7 @@ export async function POST(req: Request) {
 
       // status can be: queued, in_progress, requires_action, cancelling, cancelled, failed, completed, or expired
       while (
-        runResult.status === 'requires_action' &&
+        runResult?.status === 'requires_action' &&
         runResult.required_action?.type === 'submit_tool_outputs'
       ) {
         const tool_outputs =

diff --git a/examples/next-openai/app/api/assistant/route.ts b/examples/next-openai/app/api/assistant/route.ts
@@ -50,7 +50,7 @@ export async function POST(req: Request) {
 
       // status can be: queued, in_progress, requires_action, cancelling, cancelled, failed, completed, or expired
       while (
-        runResult.status === 'requires_action' &&
+        runResult?.status === 'requires_action' &&
         runResult.required_action?.type === 'submit_tool_outputs'
       ) {
         const tool_outputs =

diff --git a/packages/core/streams/assistant-response.ts b/packages/core/streams/assistant-response.ts
@@ -1,6 +1,7 @@
 import { AssistantStream } from 'openai/lib/AssistantStream';
 import { formatStreamPart } from '../shared/stream-parts';
 import { AssistantMessage, DataMessage } from '../shared/types';
+import { Run } from 'openai/resources/beta/threads/runs/runs';
 
 type AssistantResponseSettings = {
   threadId: string;
@@ -12,7 +13,7 @@ type AssistantResponseCallback = (options: {
   messageId: string;
   sendMessage: (message: AssistantMessage) => void;
   sendDataMessage: (message: DataMessage) => void;
-  forwardStream: (stream: AssistantStream) => Promise<any>;
+  forwardStream: (stream: AssistantStream) => Promise<Run | undefined>;
 }) => Promise<void>;
 
 export function experimental_AssistantResponse(
@@ -42,7 +43,7 @@ export function experimental_AssistantResponse(
       };
 
       const forwardStream = async (stream: AssistantStream) => {
-        let result: any = undefined;
+        let result: Run | undefined = undefined;
 
         for await (const value of stream) {
           switch (value.event) {