Function handleChatEventStream

website/src/lib/chat_stream.ts:15–50

({
  stream,
  onError,
  onPending,
  onToken,
}: ChatStreamHandlerOptions)

}

export async function handleChatEventStream({
  stream,
  onError,
  onPending,
  onToken,
}: ChatStreamHandlerOptions): Promise<InferenceMessage | null> {
  let tokens = "";
  for await (const { event, data } of iteratorSSE(stream)) {
    if (event === "error") {
      await onError(data);
    } else if (event === "ping") {
      continue;
    }
    try {
      const chunk: InferenceEvent = JSON.parse(data);
      if (chunk.event_type === "pending") {
        await onPending({ queuePosition: chunk.queue_position, queueSize: chunk.queue_size });
      } else if (chunk.event_type === "token") {
        tokens += chunk.text;
        await onToken(tokens);
      } else if (chunk.event_type === "message") {
        // final message
        return chunk.message;
      } else if (chunk.event_type === "error") {
        // handle error
        await onError(chunk.error);
        return chunk.message;
      } else {
        console.error("Unexpected event", chunk);
      }
    } catch (e) {
      console.error(`Error parsing data: ${data}, error: ${e}`);
    }
  }
  return null;
}

export async function* iteratorSSE(stream: ReadableStream<Uint8Array>) {
  const reader = stream.pipeThrough(new TextDecoderStream()).getReader();