assistant-stream 0.2.0 → 0.2.2

package/src/core/modules/assistant-stream.ts

@@ -0,0 +1,264 @@
+import { AssistantStream } from "../AssistantStream";
+import { AssistantStreamChunk, PartInit } from "../AssistantStreamChunk";
+import { createMergeStream } from "../utils/stream/merge";
+import { createTextStreamController, TextStreamController } from "./text";
+import {
+  createToolCallStreamController,
+  ToolCallStreamController,
+} from "./tool-call";
+import { Counter } from "../utils/Counter";
+import {
+  PathAppendEncoder,
+  PathMergeEncoder,
+} from "../utils/stream/path-utils";
+import { DataStreamEncoder } from "../serialization/data-stream/DataStream";
+import { FilePart, SourcePart } from "../utils/types";
+import { generateId } from "../utils/generateId";
+import {
+  ReadonlyJSONObject,
+  ReadonlyJSONValue,
+} from "../../utils/json/json-value";
+import { ToolResponseInit } from "../tool/ToolResponse";
+import { promiseWithResolvers } from "../../utils/promiseWithResolvers";
+
+type ToolCallPartInit = {
+  toolCallId?: string;
+  toolName: string;
+  argsText?: string;
+  args?: ReadonlyJSONObject;
+  response?: ToolResponseInit<ReadonlyJSONValue>;
+};
+
+export type AssistantStreamController = {
+  appendText(textDelta: string): void;
+  appendReasoning(reasoningDelta: string): void;
+  appendSource(options: SourcePart): void;
+  appendFile(options: FilePart): void;
+  addTextPart(): TextStreamController;
+  addToolCallPart(options: string): ToolCallStreamController;
+  addToolCallPart(options: ToolCallPartInit): ToolCallStreamController;
+  enqueue(chunk: AssistantStreamChunk): void;
+  merge(stream: AssistantStream): void;
+  close(): void;
+};
+
+class AssistantStreamControllerImpl implements AssistantStreamController {
+  private _merger = createMergeStream();
+  private _append:
+    | {
+        controller: TextStreamController;
+        kind: "text" | "reasoning";
+      }
+    | undefined;
+  private _contentCounter = new Counter();
+
+  get __internal_isClosed() {
+    return this._merger.isSealed();
+  }
+
+  __internal_getReadable() {
+    return this._merger.readable;
+  }
+
+  private _closeSubscriber: undefined | (() => void);
+  __internal_subscribeToClose(callback: () => void) {
+    this._closeSubscriber = callback;
+  }
+
+  private _addPart(part: PartInit, stream: AssistantStream) {
+    if (this._append) {
+      this._append.controller.close();
+      this._append = undefined;
+    }
+
+    this.enqueue({
+      type: "part-start",
+      part,
+      path: [],
+    });
+    this._merger.addStream(
+      stream.pipeThrough(new PathAppendEncoder(this._contentCounter.value)),
+    );
+  }
+
+  merge(stream: AssistantStream) {
+    this._merger.addStream(
+      stream.pipeThrough(new PathMergeEncoder(this._contentCounter)),
+    );
+  }
+
+  appendText(textDelta: string) {
+    if (this._append?.kind !== "text") {
+      this._append = {
+        kind: "text",
+        controller: this.addTextPart(),
+      };
+    }
+    this._append.controller.append(textDelta);
+  }
+
+  appendReasoning(textDelta: string) {
+    if (this._append?.kind !== "reasoning") {
+      this._append = {
+        kind: "reasoning",
+        controller: this.addReasoningPart(),
+      };
+    }
+    this._append.controller.append(textDelta);
+  }
+
+  addTextPart() {
+    const [stream, controller] = createTextStreamController();
+    this._addPart({ type: "text" }, stream);
+    return controller;
+  }
+
+  addReasoningPart() {
+    const [stream, controller] = createTextStreamController();
+    this._addPart({ type: "reasoning" }, stream);
+    return controller;
+  }
+
+  addToolCallPart(
+    options: string | ToolCallPartInit,
+  ): ToolCallStreamController {
+    const opt = typeof options === "string" ? { toolName: options } : options;
+    const toolName = opt.toolName;
+    const toolCallId = opt.toolCallId ?? generateId();
+
+    const [stream, controller] = createToolCallStreamController();
+    this._addPart({ type: "tool-call", toolName, toolCallId }, stream);
+
+    if (opt.argsText !== undefined) {
+      controller.argsText.append(opt.argsText);
+      controller.argsText.close();
+    }
+    if (opt.args !== undefined) {
+      controller.argsText.append(JSON.stringify(opt.args));
+      controller.argsText.close();
+    }
+    if (opt.response !== undefined) {
+      controller.setResponse(opt.response);
+    }
+
+    return controller;
+  }
+
+  appendSource(options: SourcePart) {
+    this._addPart(
+      options,
+      new ReadableStream({
+        start(controller) {
+          controller.enqueue({
+            type: "part-finish",
+            path: [],
+          });
+          controller.close();
+        },
+      }),
+    );
+  }
+
+  appendFile(options: FilePart) {
+    this._addPart(
+      options,
+      new ReadableStream({
+        start(controller) {
+          controller.enqueue({
+            type: "part-finish",
+            path: [],
+          });
+          controller.close();
+        },
+      }),
+    );
+  }
+
+  enqueue(chunk: AssistantStreamChunk) {
+    this._merger.enqueue(chunk);
+
+    if (chunk.type === "part-start" && chunk.path.length === 0) {
+      this._contentCounter.up();
+    }
+  }
+
+  close() {
+    this._merger.seal();
+    this._append?.controller?.close();
+
+    this._closeSubscriber?.();
+  }
+}
+
+export function createAssistantStream(
+  callback: (controller: AssistantStreamController) => PromiseLike<void> | void,
+): AssistantStream {
+  const controller = new AssistantStreamControllerImpl();
+
+  let promiseOrVoid: PromiseLike<void> | void;
+  try {
+    promiseOrVoid = callback(controller);
+  } catch (e) {
+    if (!controller.__internal_isClosed) {
+      controller.enqueue({
+        type: "error",
+        path: [],
+        error: String(e),
+      });
+      controller.close();
+    }
+    throw e;
+  }
+
+  if (promiseOrVoid instanceof Promise) {
+    const runTask = async () => {
+      try {
+        await promiseOrVoid;
+      } catch (e) {
+        if (!controller.__internal_isClosed) {
+          controller.enqueue({
+            type: "error",
+            path: [],
+            error: String(e),
+          });
+        }
+        throw e;
+      } finally {
+        if (!controller.__internal_isClosed) {
+          controller.close();
+        }
+      }
+    };
+    runTask();
+  } else {
+    if (!controller.__internal_isClosed) {
+      controller.close();
+    }
+  }
+
+  return controller.__internal_getReadable();
+}
+
+export function createAssistantStreamController() {
+  const { resolve, promise } = promiseWithResolvers<void>();
+  let controller!: AssistantStreamController;
+  const stream = createAssistantStream((c) => {
+    controller = c;
+
+    (controller as AssistantStreamControllerImpl).__internal_subscribeToClose(
+      resolve,
+    );
+
+    return promise;
+  });
+  return [stream, controller] as const;
+}
+
+export function createAssistantStreamResponse(
+  callback: (controller: AssistantStreamController) => PromiseLike<void> | void,
+) {
+  return AssistantStream.toResponse(
+    createAssistantStream(callback),
+    new DataStreamEncoder(),
+  );
+}

package/src/core/modules/text.ts

@@ -0,0 +1,64 @@
+import { AssistantStream } from "../AssistantStream";
+import { AssistantStreamChunk } from "../AssistantStreamChunk";
+import { UnderlyingReadable } from "../utils/stream/UnderlyingReadable";
+
+export type TextStreamController = {
+  append(textDelta: string): void;
+  close(): void; // TODO reason? error?
+};
+
+class TextStreamControllerImpl implements TextStreamController {
+  private _controller: ReadableStreamDefaultController<AssistantStreamChunk>;
+  private _isClosed = false;
+
+  constructor(
+    controller: ReadableStreamDefaultController<AssistantStreamChunk>,
+  ) {
+    this._controller = controller;
+  }
+
+  append(textDelta: string) {
+    this._controller.enqueue({
+      type: "text-delta",
+      path: [],
+      textDelta,
+    });
+    return this;
+  }
+
+  close() {
+    if (this._isClosed) return;
+    this._isClosed = true;
+    this._controller.enqueue({
+      type: "part-finish",
+      path: [],
+    });
+    this._controller.close();
+  }
+}
+
+export const createTextStream = (
+  readable: UnderlyingReadable<TextStreamController>,
+): AssistantStream => {
+  return new ReadableStream({
+    start(c) {
+      return readable.start?.(new TextStreamControllerImpl(c));
+    },
+    pull(c) {
+      return readable.pull?.(new TextStreamControllerImpl(c));
+    },
+    cancel(c) {
+      return readable.cancel?.(c);
+    },
+  });
+};
+
+export const createTextStreamController = () => {
+  let controller!: TextStreamController;
+  const stream = createTextStream({
+    start(c) {
+      controller = c;
+    },
+  });
+  return [stream, controller] as const;
+};

package/src/core/modules/tool-call.ts

@@ -0,0 +1,105 @@
+import { AssistantStream } from "../AssistantStream";
+import { AssistantStreamChunk } from "../AssistantStreamChunk";
+import { ToolResponseInit } from "../tool/ToolResponse";
+import { ReadonlyJSONValue } from "../../utils/json/json-value";
+import { UnderlyingReadable } from "../utils/stream/UnderlyingReadable";
+import { createTextStream, TextStreamController } from "./text";
+
+export type ToolCallStreamController = {
+  argsText: TextStreamController;
+
+  setResponse(response: ToolResponseInit<ReadonlyJSONValue>): void;
+  close(): void;
+};
+
+class ToolCallStreamControllerImpl implements ToolCallStreamController {
+  private _isClosed = false;
+
+  private _mergeTask: Promise<void>;
+  constructor(
+    private _controller: ReadableStreamDefaultController<AssistantStreamChunk>,
+  ) {
+    const stream = createTextStream({
+      start: (c) => {
+        this._argsTextController = c;
+      },
+    });
+    this._mergeTask = stream.pipeTo(
+      new WritableStream({
+        write: (chunk) => {
+          switch (chunk.type) {
+            case "text-delta":
+              this._controller.enqueue(chunk);
+              break;
+
+            case "part-finish":
+              this._controller.enqueue({
+                type: "tool-call-args-text-finish",
+                path: [],
+              });
+              break;
+
+            default:
+              throw new Error(`Unexpected chunk type: ${chunk.type}`);
+          }
+        },
+      }),
+    );
+  }
+
+  get argsText() {
+    return this._argsTextController;
+  }
+
+  private _argsTextController!: TextStreamController;
+
+  setResponse(response: ToolResponseInit<ReadonlyJSONValue>) {
+    this._controller.enqueue({
+      type: "result",
+      path: [],
+      artifact: response.artifact,
+      result: response.result,
+      isError: response.isError ?? false,
+    });
+  }
+
+  async close() {
+    if (this._isClosed) return;
+
+    this._isClosed = true;
+    this._argsTextController.close();
+    await this._mergeTask;
+
+    this._controller.enqueue({
+      type: "part-finish",
+      path: [],
+    });
+    this._controller.close();
+  }
+}
+
+export const createToolCallStream = (
+  readable: UnderlyingReadable<ToolCallStreamController>,
+): AssistantStream => {
+  return new ReadableStream({
+    start(c) {
+      return readable.start?.(new ToolCallStreamControllerImpl(c));
+    },
+    pull(c) {
+      return readable.pull?.(new ToolCallStreamControllerImpl(c));
+    },
+    cancel(c) {
+      return readable.cancel?.(c);
+    },
+  });
+};
+
+export const createToolCallStreamController = () => {
+  let controller!: ToolCallStreamController;
+  const stream = createToolCallStream({
+    start(c) {
+      controller = c;
+    },
+  });
+  return [stream, controller] as const;
+};

package/src/core/serialization/PlainText.ts

@@ -0,0 +1,68 @@
+import { AssistantStreamEncoder } from "../AssistantStream";
+import { AssistantStreamChunk } from "../AssistantStreamChunk";
+import { AssistantTransformStream } from "../utils/stream/AssistantTransformStream";
+import { PipeableTransformStream } from "../utils/stream/PipeableTransformStream";
+
+export class PlainTextEncoder
+  extends PipeableTransformStream<AssistantStreamChunk, Uint8Array>
+  implements AssistantStreamEncoder
+{
+  headers = new Headers({
+    "Content-Type": "text/plain; charset=utf-8",
+    "x-vercel-ai-data-stream": "v1",
+  });
+
+  constructor() {
+    super((readable) => {
+      const transform = new TransformStream<AssistantStreamChunk, string>({
+        transform(chunk, controller) {
+          const type = chunk.type;
+          switch (type) {
+            case "text-delta":
+              controller.enqueue(chunk.textDelta);
+              break;
+
+            default:
+              const unsupportedType:
+                | "part-start"
+                | "part-finish"
+                | "tool-call-args-text-finish"
+                | "data"
+                | "step-start"
+                | "step-finish"
+                | "message-finish"
+                | "annotations"
+                | "tool-call-begin"
+                | "tool-call-delta"
+                | "result"
+                | "error" = type;
+              throw new Error(`unsupported chunk type: ${unsupportedType}`);
+          }
+        },
+      });
+
+      return readable
+        .pipeThrough(transform)
+        .pipeThrough(new TextEncoderStream());
+    });
+  }
+}
+
+export class PlainTextDecoder extends PipeableTransformStream<
+  Uint8Array,
+  AssistantStreamChunk
+> {
+  constructor() {
+    super((readable) => {
+      const transform = new AssistantTransformStream<string>({
+        transform(chunk, controller) {
+          controller.appendText(chunk);
+        },
+      });
+
+      return readable
+        .pipeThrough(new TextDecoderStream())
+        .pipeThrough(transform);
+    });
+  }
+}