feat (core): add streamObject onFinish callback. (#1864)

Author: lgrammel

.changeset/blue-pets-punch.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+feat (core): add streamObject onFinish callback

content/docs/07-reference/ai-sdk-core/04-stream-object.mdx

@@ -305,7 +305,84 @@ for await (const partialObject of partialObjectStream) {
       isOptional: true,
       description:
         'An optional abort signal that can be used to cancel the call.'
-    }
+    },
+    {
+      name: 'onFinish',
+      type: '(result: OnFinishResult) => void',
+      isOptional: true,
+      description:
+        'Callback that is called when the LLM response and all request tool executions (for tools that have an `execute` function) are finished.',
+      properties: [
+        {
+          type: 'OnFinishResult',
+          parameters: [
+            {
+              name: 'usage',
+              type: 'TokenUsage',
+              description: 'The token usage of the generated text.',
+              properties: [
+                {
+                  type: 'TokenUsage',
+                  parameters: [
+                    {
+                      name: 'promptTokens',
+                      type: 'number',
+                      description: 'The total number of tokens in the prompt.',
+                    },
+                    {
+                      name: 'completionTokens',
+                      type: 'number',
+                      description:
+                        'The total number of tokens in the completion.',
+                    },
+                    {
+                      name: 'totalTokens',
+                      type: 'number',
+                      description: 'The total number of tokens generated.',
+                    },
+                  ],
+                },
+              ],
+            },
+            {
+              name: 'object',
+              type: 'T | undefined',
+              description:
+                'The generated object (typed according to the schema). Can be undefined if the final object does not match the schema.',
+            },
+            {
+              name:"error",
+              type:"unknown | undefined",
+              description:"Optional error object. This is e.g. a TypeValidationError when the final object does not match the schema."
+            },
+            {
+              name: 'warnings',
+              type: 'Warning[] | undefined',
+              description:
+                'Warnings from the model provider (e.g. unsupported settings).',
+            },
+            {
+              name: 'rawResponse',
+              type: 'RawResponse',
+              description: 'Optional raw response data.',
+              properties: [
+                {
+                  type: 'RawResponse',
+                  parameters: [
+                    {
+                      name: 'header',
+                      optional: true,
+                      type: 'Record<string, string>',
+                      description: 'Response headers.',
+                    },
+                  ],
+                },
+              ],
+            },
+          ],
+        },
+      ],
+    },
 
 ]}
 />

content/examples/03-node/02-streaming-structured-data/10-token-usage.mdx

@@ -5,7 +5,36 @@ description: Examples of how to record token usage when streaming structured dat
 
 # Recording Token Usage
 
-When you're streaming structured data, you may want to record the token usage for billing purposes.
+When you're streaming structured data with [`streamObject`](/docs/reference/ai-sdk-core/stream-object),
+you may want to record the token usage for billing purposes.
+
+## `onFinish` Callback
+
+You can use the `onFinish` callback to record token usage.
+It is called when the stream is finished.
+
+```ts file='index.ts' highlight={"15-17"}
+import { openai } from '@ai-sdk/openai';
+import { streamObject } from 'ai';
+import { z } from 'zod';
+
+const result = await streamObject({
+  model: openai('gpt-4-turbo'),
+  schema: z.object({
+    recipe: z.object({
+      name: z.string(),
+      ingredients: z.array(z.string()),
+      steps: z.array(z.string()),
+    }),
+  }),
+  prompt: 'Generate a lasagna recipe.',
+  onFinish({ usage }) {
+    console.log('Token usage:', usage);
+  },
+});
+```
+
+## `usage` Promise
 
 The [`streamObject`](/docs/reference/ai-sdk-core/stream-object) result contains a `usage` promise that resolves to the total token usage.
 

content/examples/03-node/02-streaming-structured-data/12-object.mdx

@@ -7,6 +7,43 @@ description: Examples of how to record the final object when streaming structure
 
 When you're streaming structured data, you may want to record the final object for logging or other purposes.
 
+## `onFinish` Callback
+
+You can use the `onFinish` callback to record the final object.
+It is called when the stream is finished.
+
+The `object` parameter contains the final object, or `undefined` if the type validation fails.
+There is also an `error` parameter that contains error when e.g. the object does not match the schema.
+
+```ts file='index.ts' highlight={"15-23"}
+import { openai } from '@ai-sdk/openai';
+import { streamObject } from 'ai';
+import { z } from 'zod';
+
+const result = await streamObject({
+  model: openai('gpt-4-turbo'),
+  schema: z.object({
+    recipe: z.object({
+      name: z.string(),
+      ingredients: z.array(z.string()),
+      steps: z.array(z.string()),
+    }),
+  }),
+  prompt: 'Generate a lasagna recipe.',
+  onFinish({ object, error }) {
+    // handle type validation failure (when the object does not match the schema):
+    if (object === undefined) {
+      console.error('Error:', error);
+      return;
+    }
+
+    console.log('Final object:', JSON.stringify(object, null, 2));
+  },
+});
+```
+
+## `object` Promise
+
 The [`streamObject`](/docs/reference/ai-sdk-core/stream-object) result contains an `object` promise that resolves to the final object.
 The object is fully typed. When the type validation according to the schema fails, the promise will be rejected with a `TypeValidationError`.
 

examples/ai-core/src/stream-object/openai-on-finish.ts

@@ -0,0 +1,43 @@
+import { openai } from '@ai-sdk/openai';
+import { streamObject } from 'ai';
+import dotenv from 'dotenv';
+import { z } from 'zod';
+
+dotenv.config();
+
+async function main() {
+  const result = await streamObject({
+    model: openai('gpt-4-turbo'),
+    schema: z.object({
+      characters: z.array(
+        z.object({
+          name: z.string(),
+          class: z
+            .string()
+            .describe('Character class, e.g. warrior, mage, or thief.'),
+          description: z.string(),
+        }),
+      ),
+    }),
+    prompt:
+      'Generate 3 character descriptions for a fantasy role playing game.',
+    onFinish({ usage, object, error }) {
+      console.log();
+      console.log('onFinish');
+      console.log('Token usage:', usage);
+
+      // handle type validation failure (when the object does not match the schema):
+      if (object === undefined) {
+        console.error('Error:', error);
+      } else {
+        console.log('Final object:', JSON.stringify(object, null, 2));
+      }
+    },
+  });
+
+  // consume the partialObjectStream:
+  for await (const partialObject of result.partialObjectStream) {
+  }
+}
+
+main().catch(console.error);

examples/ai-core/src/stream-object/openai.ts

@@ -8,7 +8,6 @@ dotenv.config();
 async function main() {
   const result = await streamObject({
     model: openai('gpt-4-turbo'),
-    maxTokens: 2000,
     schema: z.object({
       characters: z.array(
         z.object({

examples/ai-core/src/stream-text/openai-on-finish.ts

@@ -1,16 +1,12 @@
 import { openai } from '@ai-sdk/openai';
 import { streamText } from 'ai';
 import dotenv from 'dotenv';
-import { weatherTool } from '../tools/weather-tool';
 
 dotenv.config();
 
 async function main() {
   const result = await streamText({
     model: openai('gpt-4-turbo'),
-    maxTokens: 128,
-    temperature: 0.3,
-    maxRetries: 5,
     prompt: 'Invent a new holiday and describe its traditions.',
     onFinish({
       usage,

packages/core/core/generate-object/stream-object.test.ts

@@ -356,3 +356,149 @@ describe('result.object', () => {
       });
   });
 });
+
+describe('onFinish callback', () => {
+  describe('with successfully validated object', () => {
+    let result: Parameters<
+      Required<Parameters<typeof streamObject>[0]>['onFinish']
+    >[0];
+
+    beforeEach(async () => {
+      const { partialObjectStream } = await streamObject({
+        model: new MockLanguageModelV1({
+          doStream: async ({ prompt, mode }) => {
+            assert.deepStrictEqual(mode, { type: 'object-json' });
+            assert.deepStrictEqual(prompt, [
+              {
+                role: 'system',
+                content:
+                  'JSON schema:\n' +
+                  '{"type":"object","properties":{"content":{"type":"string"}},"required":["content"],"additionalProperties":false,"$schema":"http://json-schema.org/draft-07/schema#"}\n' +
+                  'You MUST answer with a JSON object that matches the JSON schema above.',
+              },
+              { role: 'user', content: [{ type: 'text', text: 'prompt' }] },
+            ]);
+
+            return {
+              stream: convertArrayToReadableStream([
+                { type: 'text-delta', textDelta: '{ ' },
+                { type: 'text-delta', textDelta: '"content": ' },
+                { type: 'text-delta', textDelta: `"Hello, ` },
+                { type: 'text-delta', textDelta: `world` },
+                { type: 'text-delta', textDelta: `!"` },
+                { type: 'text-delta', textDelta: ' }' },
+                {
+                  type: 'finish',
+                  finishReason: 'stop',
+                  usage: { completionTokens: 10, promptTokens: 3 },
+                },
+              ]),
+              rawCall: { rawPrompt: 'prompt', rawSettings: {} },
+            };
+          },
+        }),
+        schema: z.object({ content: z.string() }),
+        mode: 'json',
+        prompt: 'prompt',
+        onFinish: async event => {
+          result = event as unknown as typeof result;
+        },
+      });
+
+      // consume stream
+      await convertAsyncIterableToArray(partialObjectStream);
+    });
+
+    it('should contain token usage', async () => {
+      assert.deepStrictEqual(result.usage, {
+        completionTokens: 10,
+        promptTokens: 3,
+        totalTokens: 13,
+      });
+    });
+
+    it('should contain the full object', async () => {
+      assert.deepStrictEqual(result.object, {
+        content: 'Hello, world!',
+      });
+    });
+
+    it('should not contain an error object', async () => {
+      assert.deepStrictEqual(result.error, undefined);
+    });
+  });
+
+  describe("with object that doesn't match the schema", () => {
+    let result: Parameters<
+      Required<Parameters<typeof streamObject>[0]>['onFinish']
+    >[0];
+
+    beforeEach(async () => {
+      const { partialObjectStream, object } = await streamObject({
+        model: new MockLanguageModelV1({
+          doStream: async ({ prompt, mode }) => {
+            assert.deepStrictEqual(mode, { type: 'object-json' });
+            assert.deepStrictEqual(prompt, [
+              {
+                role: 'system',
+                content:
+                  'JSON schema:\n' +
+                  '{"type":"object","properties":{"content":{"type":"string"}},"required":["content"],"additionalProperties":false,"$schema":"http://json-schema.org/draft-07/schema#"}\n' +
+                  'You MUST answer with a JSON object that matches the JSON schema above.',
+              },
+              { role: 'user', content: [{ type: 'text', text: 'prompt' }] },
+            ]);
+
+            return {
+              stream: convertArrayToReadableStream([
+                { type: 'text-delta', textDelta: '{ ' },
+                { type: 'text-delta', textDelta: '"invalid": ' },
+                { type: 'text-delta', textDelta: `"Hello, ` },
+                { type: 'text-delta', textDelta: `world` },
+                { type: 'text-delta', textDelta: `!"` },
+                { type: 'text-delta', textDelta: ' }' },
+                {
+                  type: 'finish',
+                  finishReason: 'stop',
+                  usage: { completionTokens: 10, promptTokens: 3 },
+                },
+              ]),
+              rawCall: { rawPrompt: 'prompt', rawSettings: {} },
+            };
+          },
+        }),
+        schema: z.object({ content: z.string() }),
+        mode: 'json',
+        prompt: 'prompt',
+        onFinish: async event => {
+          result = event as unknown as typeof result;
+        },
+      });
+
+      // consume stream
+      await convertAsyncIterableToArray(partialObjectStream);
+
+      // consume expected error rejection
+      await object.catch(() => {});
+    });
+
+    it('should contain token usage', async () => {
+      assert.deepStrictEqual(result.usage, {
+        completionTokens: 10,
+        promptTokens: 3,
+        totalTokens: 13,
+      });
+    });
+
+    it('should not contain a full object', async () => {
+      assert.deepStrictEqual(result.object, undefined);
+    });
+
+    it('should contain an error object', async () => {
+      assert.deepStrictEqual(
+        TypeValidationError.isTypeValidationError(result.error),
+        true,
+      );
+    });
+  });
+});

packages/core/core/generate-object/stream-object.ts

@@ -67,6 +67,7 @@ export async function streamObject<T>({
   messages,
   maxRetries,
   abortSignal,
+  onFinish,
   ...settings
 }: CallSettings &
   Prompt & {
@@ -95,6 +96,41 @@ Please note that most providers do not support all modes.
 Default and recommended: 'auto' (best mode for the model).
      */
     mode?: 'auto' | 'json' | 'tool' | 'grammar';
+
+    /**
+Callback that is called when the LLM response and the final object validation are finished.
+     */
+    onFinish?: (event: {
+      /**
+The token usage of the generated response.
+*/
+      usage: TokenUsage;
+
+      /**
+The generated object (typed according to the schema). Can be undefined if the final object does not match the schema.
+   */
+      object: T | undefined;
+
+      /**
+Optional error object. This is e.g. a TypeValidationError when the final object does not match the schema.
+   */
+      error: unknown | undefined;
+
+      /**
+Optional raw response data.
+   */
+      rawResponse?: {
+        /**
+Response headers.
+     */
+        headers?: Record<string, string>;
+      };
+
+      /**
+Warnings from the model provider (e.g. unsupported settings).
+       */
+      warnings?: CallWarning[];
+    }) => Promise<void> | void;
   }): Promise<StreamObjectResult<T>> {
   const retry = retryWithExponentialBackoff({ maxRetries });
   const jsonSchema = convertZodToJSONSchema(schema);
@@ -229,6 +265,7 @@ Default and recommended: 'auto' (best mode for the model).
     warnings: result.warnings,
     rawResponse: result.rawResponse,
     schema,
+    onFinish,
   });
 }
 
@@ -291,13 +328,15 @@ Response headers.
     warnings,
     rawResponse,
     schema,
+    onFinish,
   }: {
     stream: ReadableStream<string | ObjectStreamInputPart>;
     warnings: CallWarning[] | undefined;
     rawResponse?: {
       headers?: Record<string, string>;
     };
     schema: z.Schema<T>;
+    onFinish: Parameters<typeof streamObject<T>>[0]['onFinish'];
   }) {
     this.warnings = warnings;
     this.rawResponse = rawResponse;
@@ -318,6 +357,8 @@ Response headers.
 
     // store information for onFinish callback:
     let usage: TokenUsage | undefined;
+    let object: T | undefined;
+    let error: unknown | undefined;
 
     // pipe chunks through a transformation stream that extracts metadata:
     let accumulatedText = '';
@@ -360,9 +401,11 @@ Response headers.
               });
 
               if (validationResult.success) {
-                resolveObject(validationResult.value);
+                object = validationResult.value;
+                resolveObject(object);
               } else {
-                rejectObject(validationResult.error);
+                error = validationResult.error;
+                rejectObject(error);
               }
 
               break;
@@ -374,6 +417,26 @@ Response headers.
             }
           }
         },
+
+        // invoke onFinish callback and resolve toolResults promise when the stream is about to close:
+        async flush(controller) {
+          try {
+            // call onFinish callback:
+            await onFinish?.({
+              usage: usage ?? {
+                promptTokens: NaN,
+                completionTokens: NaN,
+                totalTokens: NaN,
+              },
+              object,
+              error,
+              rawResponse,
+              warnings,
+            });
+          } catch (error) {
+            controller.error(error);
+          }
+        },
       }),
     );
   }