Use experimental_StreamData to send OpenAI finish_reason to client (response to issue #616)

[peterdresslar]

Hello, I have been thinking about the discussion on issue #616 about sending metadata to the client from the OpenAI API Chat Completion Chunk object. It turns out that this is possible, but a little clunky, since it requires reading the stream for different reasons at the same time. At some point I am confident that the AIStream helpers will evolve to have more data richness, but in the meantime, this hack (and it is a hack!) does work.

We take advantage of the ability to tee a response into streams. The happy news there is that the signature for OpenAIStream accepts a raw Response, so we just have to convert type and everything runs smoothly with the tee'd stream. The other branch of the stream is a little rockier. In this implementation, we are taking advantage of the fact that the last chunk from OpenAI always(?) has the finish_reason and all the others have a null finish_reason. There are almost certainly more elegant ways to do this, but hopefully you get the idea.

This example just fetches against the OpenAI API rather than using the Typescript library, but I think it would be straightforward to switch that.

//route.ts
import {
  OpenAIStream,
  StreamingTextResponse,
  experimental_StreamData,
} from 'ai';

const apiKey = process.env.OPENAI_API_KEY;

export async function POST(req: Request) {
  let resolveFinishPromise: (finishReason: string) => void;
  const finishPromise = new Promise(resolve => {
    resolveFinishPromise = resolve;
  });

  const { messages } = await req.json();
  // Instantiate the StreamData.
  const data = new experimental_StreamData();

  // construct our request to the OpenAI API with fetch
  const response = await fetch('https://api.openai.com/v1/chat/completions', {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      Authorization: `Bearer ${apiKey}`,
    },
    body: JSON.stringify({
      model: 'gpt-3.5-turbo',
      stream: true,
      messages,
    }),
  });
  if (response.body) {
    const [stream1, stream2] = response.body.tee();

    const newResponse = new Response(stream1, {
      headers: response.headers,
      status: response.status,
      statusText: response.statusText,
    });

    const stream = OpenAIStream(newResponse, {
      onToken(token) {
        console.log('token', token);
      },
      onCompletion(completion) {
        console.log('completion', completion);
      },
      onFinal(completion) {
        console.log('final', completion);
        //hang on and wait to read a value to finish:
        let finishReason = 'no finish reason';

        finishPromise.then(finishReason => {
          console.log('finishing with reason', finishReason);
          data.append({ text: `finish reason ${finishReason}` });
          data.close();
        });
      },
      experimental_streamData: true,
    });

    // Use stream2 to get the finish_reason
    const reader2 = stream2.getReader();
    let finishReason = 'no finish reason';
    reader2.read().then(async function processText({
      done,
      value,
    }): Promise<void> {
      const regex = /"finish_reason"\s*:\s*"(\w+)"/;

      if (value === undefined) {
        return;
      }

      const text = new TextDecoder('utf-8').decode(value);
      const match = regex.exec(text);

      if (match && match[1]) {
        finishReason = match[1];
        resolveFinishPromise(finishReason);
        return;
      }

      if (done) {
        //this does not seem to be firing, would be nice to resolve.
        resolveFinishPromise(finishReason);
        return;
      }

      const result_1 = await reader2.read();
      return processText(result_1);
    });

    return new StreamingTextResponse(stream, {}, data);
  } else {
    return new Response('no response body', { status: 500 });
  }
}
// yes, there is much error handling needed

This will always send a text message into your data array on the client, for you to process however you like. For an example client consumer you can check out the code here. (The example route is openai-experimental-streamdata)

(Edited out many log statements)

[gtokman]

@peterdresslar thanks for posting this. I tried it and the result includes more data but doesn't match an OpenAI response that you'd get with streaming.

0:" market"
0:"."
2:"[{\"text\":\"finish reason stop\"}]"

[peterdresslar]

@gtokman, just to clarify, are you saying that the data are different if you procure the stream by fetching the API using the the OpenAI node sdk instead of the bare fetch in the example? That would be somewhat interesting, though straightforward to adjust to, I think?

[gtokman]

@peterdresslar

When I use the SDK directly in the response looks like this:

"{\"id\":\"chatcmpl-8CC61E5Jto5axQyjArXSVYyxxCNib\",\"object\":\"chat.completion.chunk\",\"created\":1697917125,\"model\":\"gpt-3.5-turbo-0613\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"\"},\"finish_reason\":null}]}"

[peterdresslar]

It may be obvious from the code, but in case it isn't, the regex in the code above is looking at the API response from the fetch and just saying, give me any match where the 'finish_response:' has a value in it. Every chunk returned by the OpenAI API includes a finish_response: key, but only the final one has a value (apparently.) When we match a non-null value, that match is getting returned back through the promise closure and through data. (The code is a little updated today on the front end and works a bit more as I would expect it to.)

IIRC, the main difference with the OpenAI Node SDK is that Choice is nested in the ChatCompletionChunks (or, as I like to call them, Chunks!) rather than flat. I could be misremembering that.

It would be a fun exercise to swap in the SDK chat.completions.create() call for the fetch... I will see if I can get to that tonight or tomorrow. But, really, the regex in the example could be adjusted to do whatever you need, most likely.

[gtokman]

Why is it so hard for the streaming response to be identical to what OpenAI gives you if you use the SDK directly? I mainly am looking for a solution like this so that I don't hard code API keys in my client and us Vercel as a proxy server.

[peterdresslar]

I can't speak for Vercel, and I suspect they have a pretty busy week this week, but IMO it is worth considering that the underlying model serves a number of AI/LLM providers, and it also delivers a very elegant close coupling of streaming with the client side, particularly React. It is pretty cool that you can use the same chat client with a number of AI APIs and interchange almost immediately. The varied (and varying!) models of the APIs makes this a bit of work, which is why there is so much going on under the hood in ‘streams‘.

[gtokman]

I was able to make it work with this:

  const stream = new ReadableStream({
    async start(controller) {
      const encoder = new TextEncoder();

      for await (const part of response) {
        const chunk = encoder.encode(JSON.stringify(part) + "\n");
        controller.enqueue(chunk);
      }
      controller.close();
    },
  });

Agreed on your points. I think I assumed this library was the status quo for streaming a response back to the client but I was simply misundering the use case.