assistant-stream 0.2.0 → 0.2.2

package/src/core/utils/stream/path-utils.ts

@@ -0,0 +1,71 @@
+import { AssistantStreamChunk } from "../../AssistantStreamChunk";
+import { Counter } from "../Counter";
+
+export class PathAppendEncoder extends TransformStream<
+  AssistantStreamChunk,
+  AssistantStreamChunk
+> {
+  constructor(idx: number) {
+    super({
+      transform(chunk, controller) {
+        controller.enqueue({
+          ...chunk,
+          path: [idx, ...chunk.path],
+        });
+      },
+    });
+  }
+}
+
+export class PathAppendDecoder extends TransformStream<
+  AssistantStreamChunk,
+  AssistantStreamChunk
+> {
+  constructor(idx: number) {
+    super({
+      transform(chunk, controller) {
+        const {
+          path: [idx2, ...path],
+        } = chunk;
+
+        if (idx !== idx2)
+          throw new Error(`Path mismatch: expected ${idx}, got ${idx2}`);
+
+        controller.enqueue({
+          ...chunk,
+          path,
+        });
+      },
+    });
+  }
+}
+
+export class PathMergeEncoder extends TransformStream<
+  AssistantStreamChunk,
+  AssistantStreamChunk
+> {
+  constructor(counter: Counter) {
+    const innerCounter = new Counter();
+    const mapping = new Map<number, number>();
+    super({
+      transform(chunk, controller) {
+        if (chunk.type === "part-start" && chunk.path.length === 0) {
+          mapping.set(innerCounter.up(), counter.up());
+        }
+
+        const [idx, ...path] = chunk.path;
+        if (idx === undefined) {
+          controller.enqueue(chunk);
+          return;
+        }
+        const mappedIdx = mapping.get(idx);
+        if (mappedIdx === undefined) throw new Error("Path not found");
+
+        controller.enqueue({
+          ...chunk,
+          path: [mappedIdx, ...path],
+        });
+      },
+    });
+  }
+}

package/src/core/utils/types.ts

@@ -0,0 +1,150 @@
+import {
+  ReadonlyJSONObject,
+  ReadonlyJSONValue,
+} from "../../utils/json/json-value";
+
+type TextStatus =
+  | {
+      type: "running";
+    }
+  | {
+      type: "complete";
+      reason: "stop" | "unknown";
+    }
+  | {
+      type: "incomplete";
+      reason: "cancelled" | "length" | "content-filter" | "other";
+    };
+
+// export type StepStartPart = {
+//   type: "step-start";
+// };
+
+export type TextPart = {
+  type: "text";
+  text: string;
+  status: TextStatus;
+};
+
+export type ReasoningPart = {
+  type: "reasoning";
+  text: string;
+  status: TextStatus;
+};
+
+type ToolCallStatus =
+  | {
+      type: "running";
+      isArgsComplete: boolean;
+    }
+  | {
+      type: "requires-action";
+      reason: "tool-call-result";
+    }
+  | {
+      type: "complete";
+      reason: "stop" | "unknown";
+    }
+  | {
+      type: "incomplete";
+      reason: "cancelled" | "length" | "content-filter" | "other";
+    };
+
+export type ToolCallPart = {
+  type: "tool-call";
+  state: "partial-call" | "call" | "result";
+  status: ToolCallStatus;
+  toolCallId: string;
+  toolName: string;
+  argsText: string;
+  args: ReadonlyJSONObject;
+  artifact?: unknown;
+  result?: ReadonlyJSONValue;
+  isError?: boolean;
+};
+
+export type SourcePart = {
+  type: "source";
+  sourceType: "url";
+  id: string;
+  url: string;
+  title?: string;
+};
+
+export type FilePart = {
+  type: "file";
+  data: string;
+  mimeType: string;
+};
+
+export type AssistantMessagePart =
+  | TextPart
+  | ReasoningPart
+  | ToolCallPart
+  | SourcePart
+  | FilePart;
+
+type AssistantMessageStepUsage = {
+  promptTokens: number;
+  completionTokens: number;
+};
+
+type AssistantMessageStepMetadata =
+  | {
+      state: "started";
+      messageId: string;
+    }
+  | {
+      state: "finished";
+      messageId: string;
+      finishReason:
+        | "stop"
+        | "length"
+        | "content-filter"
+        | "tool-calls"
+        | "error"
+        | "other"
+        | "unknown";
+      usage?: AssistantMessageStepUsage;
+      isContinued: boolean;
+    };
+
+export type AssistantMessageStatus =
+  | {
+      type: "running";
+    }
+  | {
+      type: "requires-action";
+      reason: "tool-calls";
+    }
+  | {
+      type: "complete";
+      reason: "stop" | "unknown";
+    }
+  | {
+      type: "incomplete";
+      reason:
+        | "cancelled"
+        | "tool-calls"
+        | "length"
+        | "content-filter"
+        | "other"
+        | "error";
+      error?: ReadonlyJSONValue;
+    };
+
+export type AssistantMessage = {
+  role: "assistant";
+  status: AssistantMessageStatus;
+  parts: AssistantMessagePart[];
+  /**
+   * @deprecated Use `parts` instead.
+   */
+  content: AssistantMessagePart[];
+  metadata: {
+    unstable_data: ReadonlyJSONValue[];
+    unstable_annotations: ReadonlyJSONValue[];
+    steps: AssistantMessageStepMetadata[];
+    custom: Record<string, unknown>;
+  };
+};

package/src/core/utils/withPromiseOrValue.ts

@@ -0,0 +1,20 @@
+export function withPromiseOrValue<T>(
+  callback: () => T | PromiseLike<T>,
+  thenHandler: (value: T) => PromiseLike<void> | void,
+  catchHandler: (error: unknown) => PromiseLike<void> | void,
+): PromiseLike<void> | void {
+  try {
+    const promiseOrValue = callback();
+    if (
+      typeof promiseOrValue === "object" &&
+      promiseOrValue !== null &&
+      "then" in promiseOrValue
+    ) {
+      return promiseOrValue.then(thenHandler, catchHandler);
+    } else {
+      thenHandler(promiseOrValue);
+    }
+  } catch (e) {
+    catchHandler(e);
+  }
+}

package/src/index.ts

@@ -0,0 +1 @@
+export * from "./core";

package/src/utils/AsyncIterableStream.ts

@@ -0,0 +1,24 @@
+export type AsyncIterableStream<T> = AsyncIterable<T> & ReadableStream<T>;
+
+async function* streamGeneratorPolyfill<T>(
+  this: ReadableStream<T>,
+): AsyncIterator<T, undefined, unknown> {
+  const reader = this.getReader();
+  try {
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      yield value;
+    }
+  } finally {
+    reader.releaseLock();
+  }
+}
+
+export function asAsyncIterableStream<T>(
+  source: ReadableStream<T>,
+): AsyncIterableStream<T> {
+  (source as AsyncIterableStream<T>)[Symbol.asyncIterator] ??=
+    streamGeneratorPolyfill;
+  return source as AsyncIterableStream<T>;
+}