streams/openai-stream: don't call onStart/onCompletion when recursing (…

…#357)
vercel · Jul 21, 2023 · 0ebc2f0 · 0ebc2f0
1 parent 74e1f35
commit 0ebc2f0
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 5 deletions.
diff --git a/.changeset/blue-cows-laugh.md b/.changeset/blue-cows-laugh.md
@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+streams/openai-stream: don't call onStart/onCompletion when recursing
diff --git a/packages/core/streams/ai-stream.ts b/packages/core/streams/ai-stream.ts
@@ -15,9 +15,9 @@ export interface FunctionCallPayload {
  * @interface
  */
 export interface AIStreamCallbacks {
-  onStart?: () => Promise<void>
-  onCompletion?: (completion: string) => Promise<void>
-  onToken?: (token: string) => Promise<void>
+  onStart?: () => Promise<void> | void
+  onCompletion?: (completion: string) => Promise<void> | void
+  onToken?: (token: string) => Promise<void> | void
 }
 
 /**

diff --git a/packages/core/streams/openai-stream.ts b/packages/core/streams/openai-stream.ts
@@ -293,9 +293,19 @@ function createFunctionCallTransformer(
           return
         }
 
-        // Recursively
-        const openAIStream = OpenAIStream(functionResponse, {
+        // Recursively:
+
+        // We don't want to trigger onStart or onComplete recursively
+        // so we remove them from the callbacks
+        // see https://github.com/vercel-labs/ai/issues/351
+        const filteredCallbacks: OpenAIStreamCallbacks = {
           ...callbacks,
+          onStart: undefined,
+          onCompletion: undefined
+        }
+
+        const openAIStream = OpenAIStream(functionResponse, {
+          ...filteredCallbacks,
           [__internal__OpenAIFnMessagesSymbol]: newFunctionCallMessages
         } as AIStreamCallbacks)