assistant-stream/src/core/accumulators/AssistantMessageStream.ts

Streaming utilities for AI assistants

import type { AssistantStream } from "../AssistantStream";
import type { AssistantMessage } from "../utils/types";
import { asAsyncIterableStream } from "../../utils/AsyncIterableStream";
import { AssistantMessageAccumulator } from "./assistant-message-accumulator";

export class AssistantMessageStream {
  public readonly readable: ReadableStream<AssistantMessage>;

  constructor(readable: ReadableStream<AssistantMessage>) {
    this.readable = readable;
  }

  static fromAssistantStream(stream: AssistantStream) {
    return new AssistantMessageStream(
      stream.pipeThrough(new AssistantMessageAccumulator()),
    );
  }

  /** @deprecated Experimental since 2025-01-14. Not scheduled for removal; the API may change in any release. */
  async unstable_result(): Promise<AssistantMessage> {
    let last: AssistantMessage | undefined;
    for await (const chunk of this) {
      last = chunk;
    }

    if (!last) {
      return {
        role: "assistant",
        status: { type: "complete", reason: "unknown" },
        parts: [],
        content: [],
        metadata: {
          unstable_state: null,
          unstable_data: [],
          unstable_annotations: [],
          steps: [],
          custom: {},
        },
      };
    }
    return last;
  }

  [Symbol.asyncIterator]() {
    return asAsyncIterableStream(this.readable)[Symbol.asyncIterator]();
  }

  tee(): [AssistantMessageStream, AssistantMessageStream] {
    const [readable1, readable2] = this.readable.tee();
    return [
      new AssistantMessageStream(readable1),
      new AssistantMessageStream(readable2),
    ];
  }
}