[RFC/experimental] StreamData API for streaming additional data to th…

…e client, add onFinal callback (#425)

Co-authored-by: Shu Ding <g@shud.in>

.changeset/green-masks-repair.md

@@ -0,0 +1,7 @@
+---
+'ai': patch
+---
+
+Add experimental_StreamData and new opt-in wire protocol to enable streaming additional data. See https://github.com/vercel/ai/pull/425.
+
+Changes `onCompletion` back to run every completion, including recursive function calls. Adds an `onFinish` callback that runs once everything has streamed.

docs/pages/docs/api-reference.mdx

@@ -4,7 +4,7 @@ title: API Reference
 
 # API Reference
 
-## React Hooks
+## Hooks
 
 - [`useChat`](./api-reference/use-chat)
 - [`useCompletion`](./api-reference/use-completion)
@@ -22,8 +22,10 @@ title: API Reference
 - [`StreamingTextResponse`](./api-reference/streaming-text-response)
 - [`AIStream`](./api-reference/ai-stream)
 - [`streamToResponse`](./api-reference/stream-to-response)
+- [`experimental_StreamData`](./api-reference/stream-data)
 
 ## Prompt Construction Helpers
 
 - [`experimental_buildOpenAssistantPrompt`](./api-reference/prompts#experimental_buildopenassistantprompt)
 - [`experimental_buildStarChatBetaPrompt`](./api-reference/prompts#experimental_buildstarchatbetaprompt)
+- [`experimental_buildLlama2Prompt`](./api-reference/prompts#experimental_buildllama2prompt)

docs/pages/docs/api-reference/_meta.json

@@ -7,6 +7,7 @@
   "langchain-stream": "LangChainStream",
   "openai-stream": "OpenAIStream",
   "replicate-stream": "ReplicateStream",
+  "stream-data": "experimental_StreamData",
   "streaming-text-response": "StreamingTextResponse",
   "stream-to-response": "streamToResponse",
   "tokens": "<Tokens />"

docs/pages/docs/api-reference/ai-stream.mdx

@@ -50,7 +50,12 @@ This is an object that contains the following properties:
     [
       'onCompletion',
       '(completion: string) => Promise<void>',
-      "An optional function that is called when the stream processing is complete. It's passed the completion as a string."
+      "An optional function that is called for every completion. It's passed the completion as a string."
+    ],
+    [
+      'onFinal',
+      '(completion: string) => Promise<void>',
+      "An optional function that is called once for every request. It's passed the completion as a string. Differs from onCompletion when function calls are present."
     ],
     [
       'onToken',
@@ -106,8 +111,11 @@ const anthropicStream = AnthropicStream(fetchResponse, {
     console.log('Stream started')
   },
   onCompletion: async completion => {
-    console.log('Stream completed', completion)
+    console.log('Completion completed', completion)
   },
+  onFinal: async completion => {
+    console.log("Stream completed", completion)
+  }
   onToken: async token => {
     console.log('Token received', token)
   }

docs/pages/docs/api-reference/langchain-stream.mdx

@@ -14,7 +14,8 @@ The `AIStreamCallbacks` object has the following properties:
 
 - `onStart?: () => Promise<void>`: A function that is called at the start of the process.
 - `onToken?: (token: string) => Promise<void>`: A function that is called for each new token. The token is passed as a parameter.
-- `onCompletion?: (fullResponse: string) => Promise<void>`: A function that is called when the process is complete. The full response is passed as a parameter.
+- `onCompletion?: (completion: string) => Promise<void>`: A function that is called when a completion is complete. The full completion is passed as a parameter.
+- `onFinal?: (completion: string) => Promise<void>`: A function that is called when a response is complete. The full completion is passed as a parameter.
 
 ## Returns
 

docs/pages/docs/api-reference/stream-data.mdx

@@ -0,0 +1,96 @@
+---
+title: experimental_StreamData
+layout:
+  toc: false
+---
+
+import { Callout } from 'nextra-theme-docs'
+
+# `experimental_StreamData`
+
+The `experimental_StreamData` class allows you to stream arbitrary data to the client alongside your LLM response.
+For information on the implementation, see the associated [pull request](https://github.com/vercel/ai/pull/425).
+
+<Callout>
+  The `experimental_` prefix indicates that the API is not yet stable and may
+  change in the future without a major version bump.
+
+It is currently only implemented from `ai/react`'s `useChat` hook.
+
+</Callout>
+
+## Usage
+
+### On the Server
+
+```jsx filename="app/api/chat/route.ts" {11-12,26-28,45-46,49-51,53-56}
+export async function POST(req: Request) {
+  const { messages } = await req.json()
+
+  const response = await openai.chat.completions.create({
+    model: 'gpt-3.5-turbo-0613',
+    stream: true,
+    messages,
+    functions
+  })
+
+  // Instantiate the StreamData. It works with all API providers.
+  const data = new experimental_StreamData()
+
+  const stream = OpenAIStream(response, {
+    experimental_onFunctionCall: async (
+      { name, arguments: args },
+      createFunctionCallMessages
+    ) => {
+      if (name === 'get_current_weather') {
+        // Call a weather API here
+        const weatherData = {
+          temperature: 20,
+          unit: args.format === 'celsius' ? 'C' : 'F'
+        }
+
+        data.append({
+          text: 'Some custom data'
+        })
+
+        const newMessages = createFunctionCallMessages(weatherData)
+        return openai.chat.completions.create({
+          messages: [...messages, ...newMessages],
+          stream: true,
+          model: 'gpt-3.5-turbo-0613'
+        })
+      }
+    },
+    onCompletion(completion) {
+      console.log('completion', completion)
+    },
+    onFinal(completion) {
+      // IMPORTANT! you must close StreamData manually or the response will never finish.
+      data.close()
+    },
+    // IMPORTANT! until this is stable, you must explicitly opt in to supporting streamData.
+    experimental_streamData: true
+  })
+
+  data.append({
+    text: 'Hello, how are you?'
+  })
+
+  // IMPORTANT! If you aren't using StreamingTextResponse, you MUST have the `X-Experimental-Stream-Data: 'true'` header
+  // in your response so the client uses the correct parsing logic.
+  return new StreamingTextResponse(stream, {}, data)
+}
+```
+
+### On the client
+
+In future versions, each `Message` will have a `data` object attached to it. For the initial implementation, the SDK only supports a global `data`
+returned by the `useChat` hook:
+
+```jsx
+const { data } = useChat({
+  api: '/api/chat'
+})
+```
+
+And `data` is of the type `JSONValue[]`

docs/pages/docs/api-reference/use-chat.mdx

@@ -153,7 +153,8 @@ The `useChat` hook returns an object containing several helper methods and varia
       'isLoading',
       'boolean',
       'Boolean flag indicating whether a request is currently in progress.'
-    ]
+    ],
+    ['data', 'JSONValue[]', 'Data returned from experimental_StreamData']
   ]}
 />
 

docs/pages/docs/concepts/caching.mdx

@@ -60,8 +60,8 @@ export async function POST(req: Request) {
 
   // Convert the response into a friendly text-stream
   const stream = OpenAIStream(response, {
-    async onCompletion(completion) {
-      // Cache the response
+    async onFinal(completion) {
+      // Cache the response. Note that this will also cache function calls.
       await kv.set(key, completion)
       await kv.expire(key, 60 * 60)
     }

docs/pages/docs/guides/providers/anthropic.mdx

@@ -250,7 +250,7 @@ export async function POST(req: Request) {
       console.log(token)
     },
     onCompletion: async (completion: string) => {
-      // This callback is called when the stream completes
+      // This callback is called when the completion is ready
       // You can use this to save the final completion to your database
       await saveCompletionToDatabase(completion)
     }

examples/next-openai/app/api/chat-with-functions/route.ts

@@ -1,4 +1,8 @@
-import { OpenAIStream, StreamingTextResponse } from 'ai'
+import {
+  OpenAIStream,
+  StreamingTextResponse,
+  experimental_StreamData
+} from 'ai'
 import OpenAI from 'openai'
 import { CompletionCreateParams } from 'openai/resources/chat'
 // Create an OpenAI API client (that's edge friendly!)
@@ -12,15 +16,14 @@ export const runtime = 'edge'
 const functions: CompletionCreateParams.Function[] = [
   {
     name: 'get_current_weather',
-    description: 'Get the current weather',
+    description: 'Get the current weather.',
     parameters: {
       type: 'object',
       properties: {
         format: {
           type: 'string',
           enum: ['celsius', 'fahrenheit'],
-          description:
-            'The temperature unit to use. Infer this from the users location.'
+          description: 'The temperature unit to use.'
         }
       },
       required: ['format']
@@ -54,6 +57,7 @@ export async function POST(req: Request) {
     functions
   })
 
+  const data = new experimental_StreamData()
   const stream = OpenAIStream(response, {
     experimental_onFunctionCall: async (
       { name, arguments: args },
@@ -65,16 +69,31 @@ export async function POST(req: Request) {
           temperature: 20,
           unit: args.format === 'celsius' ? 'C' : 'F'
         }
+
+        data.append({
+          text: 'Some custom data'
+        })
+
         const newMessages = createFunctionCallMessages(weatherData)
         return openai.chat.completions.create({
           messages: [...messages, ...newMessages],
           stream: true,
-          model: 'gpt-3.5-turbo-0613',
-          functions
+          model: 'gpt-3.5-turbo-0613'
         })
       }
-    }
+    },
+    onCompletion(completion) {
+      console.log('completion', completion)
+    },
+    onFinal(completion) {
+      data.close()
+    },
+    experimental_streamData: true
+  })
+
+  data.append({
+    text: 'Hello, how are you?'
   })
 
-  return new StreamingTextResponse(stream)
+  return new StreamingTextResponse(stream, {}, data)
 }

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

@@ -18,10 +18,7 @@ export async function POST(req: Request) {
   const response = await openai.chat.completions.create({
     model: 'gpt-3.5-turbo',
     stream: true,
-    messages: messages.map((message: any) => ({
-      content: message.content,
-      role: message.role
-    }))
+    messages: messages
   })
 
   // Convert the response into a friendly text-stream

examples/next-openai/app/function-calling/page.tsx

@@ -33,7 +33,7 @@ export default function Chat() {
     }
   }
 
-  const { messages, input, handleInputChange, handleSubmit } = useChat({
+  const { messages, input, handleInputChange, handleSubmit, data } = useChat({
     api: '/api/chat-with-functions',
     experimental_onFunctionCall: functionCallHandler
   })
@@ -63,7 +63,6 @@ export default function Chat() {
           ))
         : null}
       <div id="chart-goes-here"></div>
-
       <form onSubmit={handleSubmit}>
         <input
           className="fixed bottom-0 w-full max-w-md p-2 mb-8 border border-gray-300 rounded shadow-xl"

examples/next-openai/app/page.tsx

@@ -3,8 +3,7 @@
 import { useChat } from 'ai/react'
 
 export default function Chat() {
-  const { messages, input, handleInputChange, handleSubmit } = useChat()
-
+  const { messages, input, handleInputChange, handleSubmit, data } = useChat()
   return (
     <div className="flex flex-col w-full max-w-md py-24 mx-auto stretch">
       {messages.length > 0