assistant-stream 0.2.0 → 0.2.2

package/src/core/AssistantStream.ts

@@ -0,0 +1,39 @@
+import { AssistantStreamChunk } from "./AssistantStreamChunk";
+
+export type AssistantStream = ReadableStream<AssistantStreamChunk>;
+
+export type AssistantStreamEncoder = ReadableWritablePair<
+  Uint8Array,
+  AssistantStreamChunk
+> & {
+  headers?: Headers;
+};
+
+export const AssistantStream = {
+  toResponse(stream: AssistantStream, transformer: AssistantStreamEncoder) {
+    return new Response(AssistantStream.toByteStream(stream, transformer), {
+      headers: transformer.headers ?? {},
+    });
+  },
+
+  fromResponse(
+    response: Response,
+    transformer: ReadableWritablePair<AssistantStreamChunk, Uint8Array>,
+  ) {
+    return AssistantStream.fromByteStream(response.body!, transformer);
+  },
+
+  toByteStream(
+    stream: AssistantStream,
+    transformer: ReadableWritablePair<Uint8Array, AssistantStreamChunk>,
+  ) {
+    return stream.pipeThrough(transformer);
+  },
+
+  fromByteStream(
+    readable: ReadableStream<Uint8Array>,
+    transformer: ReadableWritablePair<AssistantStreamChunk, Uint8Array>,
+  ) {
+    return readable.pipeThrough(transformer);
+  },
+};

package/src/core/AssistantStreamChunk.ts

@@ -0,0 +1,93 @@
+import { ReadonlyJSONValue } from "../utils/json/json-value";
+
+export type PartInit =
+  | {
+      readonly type: "text" | "reasoning";
+    }
+  | {
+      readonly type: "tool-call";
+      readonly toolCallId: string;
+      readonly toolName: string;
+    }
+  | {
+      readonly type: "source";
+      readonly sourceType: "url";
+      readonly id: string;
+      readonly url: string;
+      readonly title?: string;
+    }
+  | {
+      readonly type: "file";
+      readonly data: string;
+      readonly mimeType: string;
+    };
+
+export type AssistantStreamChunk = { readonly path: readonly number[] } & (
+  | {
+      readonly type: "part-start";
+      readonly part: PartInit;
+    }
+  | {
+      readonly type: "part-finish";
+    }
+  | {
+      readonly type: "tool-call-args-text-finish";
+    }
+  | {
+      readonly type: "text-delta";
+      readonly textDelta: string;
+    }
+  | {
+      readonly type: "annotations";
+      readonly annotations: ReadonlyJSONValue[];
+    }
+  | {
+      readonly type: "data";
+      readonly data: ReadonlyJSONValue[];
+    }
+  | {
+      readonly type: "step-start";
+      readonly messageId: string;
+    }
+  | {
+      readonly type: "step-finish";
+      readonly finishReason:
+        | "stop"
+        | "length"
+        | "content-filter"
+        | "tool-calls"
+        | "error"
+        | "other"
+        | "unknown";
+      readonly usage: {
+        readonly promptTokens: number;
+        readonly completionTokens: number;
+      };
+      readonly isContinued: boolean;
+    }
+  | {
+      readonly type: "message-finish";
+      readonly finishReason:
+        | "stop"
+        | "length"
+        | "content-filter"
+        | "tool-calls"
+        | "error"
+        | "other"
+        | "unknown";
+      readonly usage: {
+        readonly promptTokens: number;
+        readonly completionTokens: number;
+      };
+    }
+  | {
+      readonly type: "result";
+      readonly artifact?: ReadonlyJSONValue | undefined;
+      readonly result: ReadonlyJSONValue;
+      readonly isError: boolean;
+    }
+  | {
+      readonly type: "error";
+      readonly error: string;
+    }
+);

package/src/core/accumulators/AssistantMessageStream.ts

@@ -0,0 +1,56 @@
+import { AssistantStream } from "../AssistantStream";
+import { AssistantMessage } from "../utils/types";
+import { AssistantMessageAccumulator } from "./assistant-message-accumulator";
+
+export class AssistantMessageStream {
+  constructor(public readonly readable: ReadableStream<AssistantMessage>) {
+    this.readable = readable;
+  }
+
+  static fromAssistantStream(stream: AssistantStream) {
+    return new AssistantMessageStream(
+      stream.pipeThrough(new AssistantMessageAccumulator()),
+    );
+  }
+
+  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_data: [],
+          unstable_annotations: [],
+          steps: [],
+          custom: {},
+        },
+      };
+    }
+    return last;
+  }
+
+  [Symbol.asyncIterator]() {
+    const reader = this.readable.getReader();
+    return {
+      async next(): Promise<IteratorResult<AssistantMessage, undefined>> {
+        const { done, value } = await reader.read();
+        return done ? { done: true, value: undefined } : { done: false, value };
+      },
+    };
+  }
+
+  tee(): [AssistantMessageStream, AssistantMessageStream] {
+    const [readable1, readable2] = this.readable.tee();
+    return [
+      new AssistantMessageStream(readable1),
+      new AssistantMessageStream(readable2),
+    ];
+  }
+}

package/src/core/accumulators/assistant-message-accumulator.ts

@@ -0,0 +1,394 @@
+import { AssistantStreamChunk } from "../AssistantStreamChunk";
+import { generateId } from "../utils/generateId";
+import { parsePartialJsonObject } from "../../utils/json/parse-partial-json-object";
+import {
+  AssistantMessage,
+  AssistantMessageStatus,
+  TextPart,
+  ToolCallPart,
+  SourcePart,
+  AssistantMessagePart,
+  ReasoningPart,
+  FilePart,
+} from "../utils/types";
+
+const createInitialMessage = (): AssistantMessage => ({
+  role: "assistant",
+  status: { type: "running" },
+  parts: [],
+  get content() {
+    return this.parts;
+  },
+  metadata: {
+    unstable_data: [],
+    unstable_annotations: [],
+    steps: [],
+    custom: {},
+  },
+});
+
+const updatePartForPath = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk,
+  updater: (part: AssistantMessagePart) => AssistantMessagePart,
+): AssistantMessage => {
+  if (message.parts.length === 0) {
+    throw new Error("No parts available to update.");
+  }
+
+  if (chunk.path.length !== 1)
+    throw new Error("Nested paths are not supported yet.");
+
+  const partIndex = chunk.path[0]!;
+  const updatedPart = updater(message.parts[partIndex]!);
+  return {
+    ...message,
+    parts: [
+      ...message.parts.slice(0, partIndex),
+      updatedPart,
+      ...message.parts.slice(partIndex + 1),
+    ],
+    get content() {
+      return this.parts;
+    },
+  };
+};
+
+const handlePartStart = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & { readonly type: "part-start" },
+): AssistantMessage => {
+  const partInit = chunk.part;
+  if (partInit.type === "text" || partInit.type === "reasoning") {
+    const newTextPart: TextPart | ReasoningPart = {
+      type: partInit.type,
+      text: "",
+      status: { type: "running" },
+    };
+    return {
+      ...message,
+      parts: [...message.parts, newTextPart],
+      get content() {
+        return this.parts;
+      },
+    };
+  } else if (partInit.type === "tool-call") {
+    const newToolCallPart: ToolCallPart = {
+      type: "tool-call",
+      state: "partial-call",
+      status: { type: "running", isArgsComplete: false },
+      toolCallId: partInit.toolCallId,
+      toolName: partInit.toolName,
+      argsText: "",
+      args: {},
+    };
+    return {
+      ...message,
+      parts: [...message.parts, newToolCallPart],
+      get content() {
+        return this.parts;
+      },
+    };
+  } else if (partInit.type === "source") {
+    const newSourcePart: SourcePart = {
+      type: "source",
+      sourceType: partInit.sourceType,
+      id: partInit.id,
+      url: partInit.url,
+      ...(partInit.title ? { title: partInit.title } : undefined),
+    };
+    return {
+      ...message,
+      parts: [...message.parts, newSourcePart],
+      get content() {
+        return this.parts;
+      },
+    };
+  } else if (partInit.type === "file") {
+    const newFilePart: FilePart = {
+      type: "file",
+      mimeType: partInit.mimeType,
+      data: partInit.data,
+    };
+    return {
+      ...message,
+      parts: [...message.parts, newFilePart],
+      get content() {
+        return this.parts;
+      },
+    };
+  } else {
+    throw new Error(`Unsupported part type: ${partInit.type}`);
+  }
+};
+
+const handleToolCallArgsTextFinish = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & {
+    readonly type: "tool-call-args-text-finish";
+  },
+): AssistantMessage => {
+  return updatePartForPath(message, chunk, (part) => {
+    if (part.type !== "tool-call") {
+      throw new Error("Last is not a tool call");
+    }
+    return {
+      ...part,
+      state: "call",
+    };
+  });
+};
+
+const handlePartFinish = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & { readonly type: "part-finish" },
+): AssistantMessage => {
+  return updatePartForPath(message, chunk, (part) => ({
+    ...part,
+    status: { type: "complete", reason: "unknown" },
+  }));
+};
+
+const handleTextDelta = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & { type: "text-delta" },
+): AssistantMessage => {
+  return updatePartForPath(message, chunk, (part) => {
+    if (part.type === "text" || part.type === "reasoning") {
+      return { ...part, text: part.text + chunk.textDelta };
+    } else if (part.type === "tool-call") {
+      const newArgsText = part.argsText + chunk.textDelta;
+
+      // Fall back to existing args if parsing fails
+      const newArgs = parsePartialJsonObject(newArgsText) ?? part.args;
+
+      return { ...part, argsText: newArgsText, args: newArgs };
+    } else {
+      throw new Error(
+        "text-delta received but part is neither text nor tool-call",
+      );
+    }
+  });
+};
+
+const handleResult = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & { type: "result" },
+): AssistantMessage => {
+  return updatePartForPath(message, chunk, (part) => {
+    if (part.type === "tool-call") {
+      return {
+        ...part,
+        state: "result",
+        artifact: chunk.artifact,
+        result: chunk.result,
+        isError: chunk.isError ?? false,
+        status: { type: "complete", reason: "stop" },
+      };
+    } else {
+      throw new Error("Result chunk received but part is not a tool-call");
+    }
+  });
+};
+
+const handleMessageFinish = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & { type: "message-finish" },
+): AssistantMessage => {
+  const newStatus = getStatus(chunk);
+  return { ...message, status: newStatus };
+};
+
+const getStatus = (
+  chunk:
+    | (AssistantStreamChunk & { type: "message-finish" })
+    | (AssistantStreamChunk & { type: "step-finish" }),
+): AssistantMessageStatus => {
+  if (chunk.finishReason === "tool-calls") {
+    return {
+      type: "requires-action",
+      reason: "tool-calls",
+    };
+  } else if (
+    chunk.finishReason === "stop" ||
+    chunk.finishReason === "unknown"
+  ) {
+    return {
+      type: "complete",
+      reason: chunk.finishReason,
+    };
+  } else {
+    return {
+      type: "incomplete",
+      reason: chunk.finishReason,
+    };
+  }
+};
+
+const handleAnnotations = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & { type: "annotations" },
+): AssistantMessage => {
+  return {
+    ...message,
+    metadata: {
+      ...message.metadata,
+      unstable_annotations: [
+        ...message.metadata.unstable_annotations,
+        ...chunk.annotations,
+      ],
+    },
+  };
+};
+
+const handleData = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & { type: "data" },
+): AssistantMessage => {
+  return {
+    ...message,
+    metadata: {
+      ...message.metadata,
+      unstable_data: [...message.metadata.unstable_data, ...chunk.data],
+    },
+  };
+};
+
+const handleStepStart = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & { type: "step-start" },
+): AssistantMessage => {
+  return {
+    ...message,
+    metadata: {
+      ...message.metadata,
+      steps: [
+        ...message.metadata.steps,
+        { state: "started", messageId: chunk.messageId },
+      ],
+    },
+  };
+};
+
+const handleStepFinish = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & { type: "step-finish" },
+): AssistantMessage => {
+  const steps = message.metadata.steps.slice();
+  const lastIndex = steps.length - 1;
+
+  // Check if the previous step is a step-start (has state "started")
+  if (steps.length > 0 && steps[lastIndex]?.state === "started") {
+    steps[lastIndex] = {
+      ...steps[lastIndex],
+      state: "finished",
+      finishReason: chunk.finishReason,
+      usage: chunk.usage,
+      isContinued: chunk.isContinued,
+    };
+  } else {
+    // If no previous step-start exists, append a finished step
+    steps.push({
+      state: "finished",
+      messageId: generateId(),
+      finishReason: chunk.finishReason,
+      usage: chunk.usage,
+      isContinued: chunk.isContinued,
+    });
+  }
+
+  return {
+    ...message,
+    metadata: {
+      ...message.metadata,
+      steps,
+    },
+  };
+};
+
+const handleErrorChunk = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & { type: "error" },
+): AssistantMessage => {
+  return {
+    ...message,
+    status: { type: "incomplete", reason: "error", error: chunk.error },
+  };
+};
+
+export class AssistantMessageAccumulator extends TransformStream<
+  AssistantStreamChunk,
+  AssistantMessage
+> {
+  constructor({
+    initialMessage,
+  }: {
+    initialMessage?: AssistantMessage;
+  } = {}) {
+    let message = initialMessage ?? createInitialMessage();
+    super({
+      transform(chunk, controller) {
+        const type = chunk.type;
+        switch (type) {
+          case "part-start":
+            message = handlePartStart(message, chunk);
+            break;
+
+          case "tool-call-args-text-finish":
+            message = handleToolCallArgsTextFinish(message, chunk);
+            break;
+
+          case "part-finish":
+            message = handlePartFinish(message, chunk);
+            break;
+
+          case "text-delta":
+            message = handleTextDelta(message, chunk);
+            break;
+          case "result":
+            message = handleResult(message, chunk);
+            break;
+          case "message-finish":
+            message = handleMessageFinish(message, chunk);
+            break;
+          case "annotations":
+            message = handleAnnotations(message, chunk);
+            break;
+          case "data":
+            message = handleData(message, chunk);
+            break;
+          case "step-start":
+            message = handleStepStart(message, chunk);
+            break;
+          case "step-finish":
+            message = handleStepFinish(message, chunk);
+            break;
+          case "error":
+            message = handleErrorChunk(message, chunk);
+            break;
+          default: {
+            const unhandledType: never = type;
+            throw new Error(`Unsupported chunk type: ${unhandledType}`);
+          }
+        }
+        controller.enqueue(message);
+      },
+      flush(controller) {
+        if (message.status?.type === "running") {
+          // TODO this test isn't sound
+          const requiresAction = message.parts?.at(-1)?.type === "tool-call";
+          message = handleMessageFinish(message, {
+            type: "message-finish",
+            path: [],
+            finishReason: requiresAction ? "tool-calls" : "unknown",
+            usage: {
+              promptTokens: 0,
+              completionTokens: 0,
+            },
+          });
+          controller.enqueue(message);
+        }
+      },
+    });
+  }
+}

package/src/core/index.ts

@@ -0,0 +1,17 @@
+export {
+  createAssistantStream,
+  createAssistantStreamResponse,
+} from "./modules/assistant-stream";
+export { AssistantMessageAccumulator } from "./accumulators/assistant-message-accumulator";
+export { AssistantStream } from "./AssistantStream";
+export type { AssistantStreamController } from "./modules/assistant-stream";
+export type { AssistantStreamChunk } from "./AssistantStreamChunk";
+export {
+  DataStreamDecoder,
+  DataStreamEncoder,
+} from "./serialization/data-stream/DataStream";
+export { PlainTextDecoder, PlainTextEncoder } from "./serialization/PlainText";
+export { AssistantMessageStream } from "./accumulators/AssistantMessageStream";
+export type { AssistantMessage } from "./utils/types";
+
+export * from "./tool";