[RFC/experimental] StreamData API for streaming additional data to the client

[MaxLeiter]

This is an experimental API and is subject to change. Please leave (non-code related) feedback! the code is WIP and is mostly for experimentation

One of the most common requested features has been how to stream additional data to the client. I think we have a pretty good solution now.

There's a lot of duplicate code here because we're going to support the old protocol and new one while experimenting. As a result this functionality is only implemented in ai/react's useChat

On the server (StreamData)

The StreamData class is a new TransformStream wrapper available on the server from the root ai package.

You instantiate it like any other class:

  const data = new experimental_StreamData()

and use it in your API:

+  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'
+        })

        return;
      }
    }, 
+   onFinish() { 
+ 	// IMPORTANT! the user must manually close the stream 
+     data.close();
+	},
+  // IMPORTANT! this is only required while we support both protocols
+  experimental_StreamData: true
  })

+ data.append({ done: true })

  return new StreamingTextResponse(finalStream, {
+ // IMPORTANT! behind the scenes, StreamingTextResponse adds a header if `data` is present. If you aren't using StreamingTextResponse you need to add the header yourself. 
+  {},
+, data);

On the client

The useChat and useCompletion hooks return a new object, data:

  const { messages, input, handleInputChange, handleSubmit, data } = useChat({

data is the array of JS objects you appended to with StreamData. Correlating data to specific messages is not handled by the SDK and will be managed by the user. In the future this will likely change, and we'll maintain the global data and a message.data.

How it works

This requires moving the API from a simple raw text response to a prefixed protocol inspired by React Server Components. Implementation-wise this can later be migrated to Server-Side Events.

The current prefix map looks like this, and is exported so it can be used in userspace:

/**
 * The map of prefixes for data in the stream
 *
 * - 0: Text from the LLM response
 * - 1: (OpenAI) function_call responses
 * - 2: custom JSON added by the user using `Data`
 *
 * Example:
 * ```
 * 0:Vercel
 * 0:'s
 * 0: AI
 * 0: AI
 * 0: SDK
 * 0: is great
 * 0:!
 * 2: { "someJson": "value" }
 * 1: {"function_call": {"name": "get_current_weather", "arguments": "{\\n\\"location\\": \\"Charlottesville, Virginia\\",\\n\\"format\\": \\"celsius\\"\\n}"}}
 *```
 */
export const StreamStringPrefixes = {
  text: 0,
  function_call: 1,
  data: 2
  // user_err: 3?
} as const

Caveats

• Implementation complexity increases. For now, this is only implemented in useChat until it's less experimental/more widely agreed upon
• Client-side parsing is no longer as simple as fetching the request. Non-SDK client's will need to manually parse the prefixes and construct messages, but we can provide helpers to make it simpler.
• More bandwidth is used, although the number of tokens remains the same, so this isn't a very big deal.

What's next?

• this is basically SSE but only consumed by us. In the future we can adjust the transport format.
• data tied to each message if they come in the same response

[changeset-bot]

⚠️ No Changeset found

Latest commit: 3c901a7

Merging this PR will not cause a version bump for any packages. If these changes should not result in a new version, you're good to go. If these changes should result in a version bump, you need to add a changeset.

This PR includes no changesets

When changesets are added to this PR, you'll see the packages that this PR includes changesets for and the associated semver types

Click here to learn what changesets are, and how to add one.

Click here if you're a maintainer who wants to add a changeset to this PR

[dnsosebee]

Thanks for this! This is exactly the API I was imagining to solve #350.

My use case involves per-message data, but I can understand why a global data structure affords more flexibility and will probably cover more use cases. I would just do some zipping for my use case

[newme616]

This is a great idea! Would love to see this in production soon

[shuding]

Looks great in general!

[shuding]

It only updates the data array when a new message begins to stream.

This is something we can improve in a follow up PR as well, thanks for the feedback!

[rajeshdavidbabu]

@MaxLeiter 🔥 it's working now. but the output is coming out strange:

If it helps, happy to hop on a quick call.

I am facing the same issue now, I also added COMPLEX_HEADER and dint fix it for me.

[rajeshdavidbabu]

Am I missing something, I am using stream from Langchain Stream so not entirely sure where to add the experimental_StreamData to true.

[MaxLeiter]

@rajeshdavidbabu ah the new parser hasn't been implemented in LangChainStream yet, sorry.

If you'd like to contribute it, I believe it's mostly a case of removing the pipeThrough(createStreamDataTransformer()) line and adding getStreamString() around the writer.write()s

https://github.com/vercel/ai/blob/main/packages/core/streams/langchain-stream.ts#L33-L68

[rajeshdavidbabu]

Nice. Can take a look !

[rayli09]

@MaxLeiter hi there, we're trying the feature but we got a weird behavior that on the front end it keeps triggering requests.

We're using next-js@13.4.9, ai@2.2.6, and enabled the feature flags.

checking on the console log in browser, we did see custom data being piped, but on the front end, the streamed message first completes, and then vanished, and then it keeps re-rendering the ui and re triggering multiple request with the same user's last input.

  // Convert the response into a friendly text-stream
  const stream = OpenAIStream(response, {
    onStart: async () => {
      // save to db
    },
    onToken: async (_token: string) => {
      // log
    },
    onCompletion: async (completion: string) => {
      // save to db
    },
    onFinal: async (_completion: string) => {
      // close the stream
      await customData.close();
    },
    // experimental flag to stream custom data.
    experimental_streamData: true,
  });
  return new StreamingTextResponse(stream, {}, customData);

[rajeshdavidbabu]

Afaik, its already been implemented https://github.com/vercel/ai/blob/main/packages/core/streams/stream-data.ts#L111.

So I simply have to pass the experimental_streamData flag to be true when I create an instance of LangChainStream

const { stream, handlers } = LangChainStream({
      experimental_streamData: true,
    });

This works for me.

[sebastianhoitz]

@MaxLeiter thank you so much, this is really cool - one question: You said you might be able to get per-message data working. I don't see anything locally at the moment, only the global data.

Is this because I'm using a LangChain agent? Or is it actually missing?

[MaxLeiter]

@sebastianhoitz it's probably only a few line change but it wasn't included in the initial version:

Correlating data to specific messages is not handled by the SDK and will be managed by the user. In the future this will likely change, and we'll maintain the global data and a message.data.

[dnsosebee]

Thanks for the work on this. I'm going to try out function calling and so may not need to stream extra data.

However I've been thinking about DX and have an idea: I think the core data of RAG apps is not a message log, it's an event log containing user messages, function calls, function responses, and assistant messages. A RAG app needs to present views of that event log to both the user and the assistant; the assistant's view is a message log, and the user's view is whatever JSON drives their UI. The current API (pre-experimental_streamData), which thinks of things as OpenAI-style message logs rather than generic event logs, makes it hard to handle those two different views of the data.

For example, in my RAG app I want my search results to be formatted as a plaintext system message for the OpenAI message array, but want to view those same results as JSON client side, so that I can present a "related research" UI section. Under the current functionCalling API that might involve parsing system messages client side, which is unwieldy. With the current streamData API that would involve interleaving the data array back into my message history as system messages, which is also a bit awkward. Either that or duplicating data as system messages AND streamData.

So I think my ideal API would be to have a generic type called "event" that looks something like this:

type Event<T> = {
  name: string // or generic string literal
  data: T
}

This Event type would be part of message history, and I would set up functions on client and server side to determine how to map that event to a system message for OpenAI, and how to interpret the event data as UI client side.

This is a bit different from message-correlated data, since the data event is a neighboring member in the message/event array rather than an additional field of a message. I think that's more sensible, since fundamentally the data event has its own chronological place in the stream, you can even think of a data event as a message from a 3rd party agent, like a 'search agent' who just sent a message to the stream.

Curious to hear people's thoughts on this!

[GorvGoyl]

Client-side parsing is no longer as simple as fetching the request. Non-SDK client's will need to manually parse the prefixes and construct messages, but we can provide helpers to make it simpler.

I'm using Vercel AI streaming response in the Chrome extension and the SDK is not supported. Could you provide the helper to parse data?

[aaron5670]

@MaxLeiter 🔥 it's working now. but the output is coming out strange:
If it helps, happy to hop on a quick call.

I am facing the same issue now, I also added COMPLEX_HEADER and dint fix it for me.

I have the same issue, this is my response:

[
    {
        "content": "Hello",
        "role": "user",
        "createdAt": "2024-01-19T09:07:36.443Z",
        "id": "31kQ0KF"
    },
    {
        "id": "ES6PN2g",
        "createdAt": "2024-01-19T09:07:39.792Z",
        "content": "2:[{\"text\":\"Hello, how are you?\"}]\n0:\"Bonjour\"\n0:\"!\"\n0:\" Comment\"\n0:\" puis\"\n0:\"-\"\n0:\"je\"\n0:\" vous\"\n0:\" aider\"\n0:\" aujourd\"\n0:\"'hui\"\n0:\"?\"\n",
        "role": "assistant"
    }
]

And this is my backend code:

import {experimental_StreamData, OpenAIStream, StreamingTextResponse} from 'ai';
import {ChatCompletionMessage} from 'openai/resources/chat/completions';
import {createOpenAIEmbedding} from "@/helpers/createOpenAIEmbedding";
import {pineconeIndex} from "@/helpers/pinecone";
import {openai} from "@/helpers/openai";

export async function POST(req: Request) {
  if (!process.env.OPENAI_API_KEY) {
    return Response.json({error: "OpenAI is not enabled, check your env variables"}, {status: 400});
  }

  try {
    const body = await req.json();
    const messages: ChatCompletionMessage[] = body.messages;

    const messagesTruncated = messages.slice(-6);

    const embedding = await createOpenAIEmbedding(
      messagesTruncated.map((message) => message.content).join("\n"),
    );

    const vectorQueryResponse = await pineconeIndex.query({
      vector: embedding,
      topK: 4,
      includeMetadata: true,
      filter: {type: 'product'},
    });

    console.log('vectorQueryResponse', vectorQueryResponse.matches.map((document) => document.metadata?.title));

    const systemMessage: ChatCompletionMessage = {
      role: "assistant",
      content:
        `
        You are a helpful assistant of the online e-commerce store.
        It is very important that you always answer in French, you must not deviate from this at any cost!
        Use the following bits of context to answer the question at the end.
        Make sure your answer consists of a maximum of about 45 words.
        If you don't know the answer, just say you don't know, don't go making up answers.
        ----------------
        CONTEXT:
        ${vectorQueryResponse.matches
          .map((document) => {
            return `Title: ${document.metadata?.title}\n\nContent:\n${document.metadata?.text}`
          })
          .join("\n\n")}
         `
    };

    const response = await openai.chat.completions.create({
      model: "gpt-3.5-turbo",
      stream: true,
      messages: [systemMessage, ...messagesTruncated],
    });

    // 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);
  } catch (error) {
    console.error(error);
    return Response.json({error: "Internal server error"}, {status: 500});
  }
}

export async function OPTIONS() {
  return Response.json({status: "OK"});
}

What am I doing wrong and how did you solve this? :)

[aaron5670]

@MaxLeiter I found the issue.

We have currently a Gatsby app where I want to use this. I tried it with a Next.js app and then it works as expected.

Next.js

This is my Next.js front-end code, where it works:

"use client";

import {useChat} from "ai/react";

const Page = () => {
  const { data, handleInputChange, handleSubmit, messages } = useChat({
    api: '/api/chat',
  });

  console.log('data', data);
  console.log('messages', messages);

  return (
    <div>
      <form onSubmit={handleSubmit}>
        <input
          type="text"
          name="input"
          onChange={handleInputChange}
        />
        <button type="submit">Send</button>
      </form>
    </div>
  )
}

export default Page

Next.js response:

[
    {
        "content": "Hello",
        "role": "user",
        "createdAt": "2024-01-19T09:40:26.742Z",
        "id": "cJCvzYF"
    },
    {
        "id": "4W01c2l",
        "role": "assistant",
        "content": "Bonjour! Comment puis-je vous aider aujourd'hui ?",
        "createdAt": "2024-01-19T09:40:29.129Z"
    }
]

Gatsby

And this is my Gatsby front-end code, where I have encoding issues:

import React from 'react';
import { useChat } from 'ai/react';

export const Page = () => {
  const {
    data, handleInputChange, handleSubmit, messages,
  } = useChat({
    api: 'http://localhost:3000/api/chat',
  });

  console.log('data', data);
  console.log('messages', messages);

  return (
    <div>
      <form onSubmit={handleSubmit}>
        <input
          type="text"
          name="input"
          onChange={handleInputChange}
        />
        <button type="submit">Send</button>
      </form>
    </div>
  );
};

Gatsby response:

[
    {
        "content": "Hello",
        "role": "user",
        "createdAt": "2024-01-19T09:41:10.737Z",
        "id": "oK7t7Lh"
    },
    {
        "id": "9F2Pg9s",
        "createdAt": "2024-01-19T09:41:12.195Z",
        "content": "0:\"Bonjour\"\n0:\"!\"\n0:\" Comment\"\n0:\" puis\"\n0:\"-\"\n0:\"je\"\n0:\" vous\"\n0:\" aider\"\n0:\" aujourd\"\n0:\"'hui\"\n0:\"?\"\n",
        "role": "assistant"
    }
]

[zandko]

I also encountered the same problem. It is normal locally. This problem occurs when I call the remote server.

[aaron5670]

I also encountered the same problem. It is normal locally. This problem occurs when I call the remote server.

See #930

[straatrakker]

I am also encountering the same issue. Works fine locally on Next.js but it doesn't seem to work when calling the Next.js API from another React/vite project.

Has anyone found a fix for this?

[Laktus]

How would i send this info not from a next backend but e.g. from a python backend with fast api? What header do i have to add?