lib/openai-stream.ts

import type { ChatCompletionChunk } from "openai/resources";

export function createOpenAIStream(
  completion: AsyncIterable<ChatCompletionChunk>,
  callbacks: { onFinal: () => void }
): ReadableStream {
  const { readable, writable } = new TransformStream();
  const writer = writable.getWriter();
  const encoder = new TextEncoder();

  let isFinished = false;

  async function forwardCompletion() {
    try {
      for await (const chunk of completion) {
        const { content } = chunk.choices[0].delta;
        if (content) {
          await writer.write(encoder.encode(`0:${JSON.stringify(content)}\n`));
        }
      }
    } catch (error) {
      console.error("Error forwarding completion:", error);
      await writer.abort(error);
    } finally {
      isFinished = true;
      await writer.close();
      callbacks.onFinal();
    }
  }

  forwardCompletion();

  return readable;
}