assistant-stream 0.2.0 → 0.2.1

package/src/core/tool/tool-types.ts

@@ -0,0 +1,107 @@
+import { JSONSchema7 } from "json-schema";
+import { TypeAtPath, TypePath } from "./type-path-utils";
+import { DeepPartial } from "ai";
+import { AsyncIterableStream } from "../../utils";
+import type { StandardSchemaV1 } from "@standard-schema/spec";
+import { ToolResponse } from "./ToolResponse";
+
+/**
+ * Interface for reading tool call arguments from a stream, which are
+ * generated by a language learning model (LLM). Provides methods to
+ * retrieve specific values, partial streams, or complete items from
+ * an array, based on a specified path.
+ *
+ * @template TArgs The type of arguments being read.
+ */
+export interface ToolCallArgsReader<TArgs> {
+  /**
+   * Returns a promise that will resolve to the value at the given path,
+   * as soon as that path is generated by the LLM.
+   *
+   * @param fieldPath An array of object keys or array indices.
+   */
+  get<PathT extends TypePath<TArgs>>(
+    ...fieldPath: PathT
+  ): Promise<TypeAtPath<TArgs, PathT>>;
+
+  /**
+   * Returns a stream that will emit partial values at the given path,
+   * as they are generated by the LLM.
+   *
+   * @param fieldPath An array of object keys or array indices.
+   */
+  streamValues<PathT extends TypePath<TArgs>>(
+    ...fieldPath: PathT
+  ): AsyncIterableStream<DeepPartial<TypeAtPath<TArgs, PathT>>>;
+
+  /**
+   * Returns a stream that will emit partial text at the given path,
+   * as they are generated by the LLM.
+   *
+   * @param fieldPath An array of object keys or array indices.
+   */
+  streamText<PathT extends TypePath<TArgs>>(
+    ...fieldPath: PathT
+  ): TypeAtPath<TArgs, PathT> extends string & infer U
+    ? AsyncIterableStream<U>
+    : never;
+
+  /**
+   * Returns a stream that will emit complete items in the array
+   * at the given path, as they are generated by the LLM.
+   *
+   * @param fieldPath An array of object keys or array indices.
+   */
+  forEach<PathT extends TypePath<TArgs>>(
+    ...fieldPath: PathT
+  ): TypeAtPath<TArgs, PathT> extends Array<infer U>
+    ? AsyncIterableStream<U>
+    : never;
+}
+
+export interface ToolCallResponseReader<TResult> {
+  get: () => Promise<ToolResponse<TResult>>;
+}
+
+export interface ToolCallReader<TArgs, TResult> {
+  args: ToolCallArgsReader<TArgs>;
+  response: ToolCallResponseReader<TResult>;
+
+  /**
+   * @deprecated Deprecated. Use `response.get().result` instead.
+   */
+  result: {
+    get: () => Promise<TResult>;
+  };
+}
+
+type ToolExecutionContext = {
+  toolCallId: string;
+  abortSignal: AbortSignal;
+};
+
+export type ToolExecuteFunction<TArgs, TResult> = (
+  args: TArgs,
+  context: ToolExecutionContext,
+) => TResult | Promise<TResult>;
+
+export type ToolStreamCallFunction<TArgs, TResult> = (
+  reader: ToolCallReader<TArgs, TResult>,
+  context: ToolExecutionContext,
+) => void;
+
+type OnSchemaValidationErrorFunction<TResult> = ToolExecuteFunction<
+  unknown,
+  TResult
+>;
+
+export type Tool<TArgs = unknown, TResult = unknown> = {
+  description?: string | undefined;
+  parameters: StandardSchemaV1<TArgs> | JSONSchema7;
+  execute?: ToolExecuteFunction<TArgs, TResult>;
+  /**
+   * @deprecated Experimental, API may change.
+   */
+  streamCall?: ToolStreamCallFunction<TArgs, TResult>;
+  experimental_onSchemaValidationError?: OnSchemaValidationErrorFunction<TResult>;
+};

package/src/core/tool/toolResultStream.ts

@@ -0,0 +1,119 @@
+import { Tool, ToolCallReader, ToolExecuteFunction } from "./tool-types";
+import { StandardSchemaV1 } from "@standard-schema/spec";
+import { ToolResponse } from "./ToolResponse";
+import { ToolExecutionStream } from "./ToolExecutionStream";
+import { AssistantMessage } from "../utils/types";
+
+const isStandardSchemaV1 = (
+  schema: unknown,
+): schema is StandardSchemaV1<unknown> => {
+  return (
+    typeof schema === "object" &&
+    schema !== null &&
+    "~standard" in schema &&
+    (schema as StandardSchemaV1<unknown>)["~standard"].version === 1
+  );
+};
+
+function getToolResponse(
+  tools: Record<string, Tool<any, any>> | undefined,
+  abortSignal: AbortSignal,
+  toolCall: {
+    toolCallId: string;
+    toolName: string;
+    args: unknown;
+  },
+) {
+  const tool = tools?.[toolCall.toolName];
+  if (!tool || !tool.execute) return undefined;
+
+  const getResult = async (toolExecute: ToolExecuteFunction<any, any>) => {
+    let executeFn = toolExecute;
+
+    if (isStandardSchemaV1(tool.parameters)) {
+      let result = tool.parameters["~standard"].validate(toolCall.args);
+      if (result instanceof Promise) result = await result;
+
+      if (result.issues) {
+        executeFn =
+          tool.experimental_onSchemaValidationError ??
+          (() => {
+            throw new Error(
+              `Function parameter validation failed. ${JSON.stringify(result.issues)}`,
+            );
+          });
+      }
+    }
+
+    const result = await executeFn(toolCall.args, {
+      toolCallId: toolCall.toolCallId,
+      abortSignal,
+    });
+    if (result instanceof ToolResponse) return result;
+    return new ToolResponse({
+      result: result === undefined ? "<no result>" : result,
+    });
+  };
+
+  return getResult(tool.execute);
+}
+
+function getToolStreamResponse<TArgs, TResult>(
+  tools: Record<string, Tool<any, any>> | undefined,
+  abortSignal: AbortSignal,
+  reader: ToolCallReader<TArgs, TResult>,
+  context: {
+    toolCallId: string;
+    toolName: string;
+  },
+) {
+  tools?.[context.toolName]?.streamCall?.(reader, {
+    toolCallId: context.toolCallId,
+    abortSignal,
+  });
+}
+
+export async function unstable_runPendingTools(
+  message: AssistantMessage,
+  tools: Record<string, Tool<any, any>> | undefined,
+  abortSignal: AbortSignal,
+) {
+  // TODO parallel tool calling
+  for (const part of message.parts) {
+    if (part.type === "tool-call") {
+      const promiseOrUndefined = getToolResponse(tools, abortSignal, part);
+      if (promiseOrUndefined) {
+        const result = await promiseOrUndefined;
+        const updatedParts = message.parts.map((p) => {
+          if (p.type === "tool-call" && p.toolCallId === part.toolCallId) {
+            return {
+              ...p,
+              state: "result" as const,
+              artifact: result.artifact,
+              result: result.result,
+              isError: result.isError,
+            };
+          }
+          return p;
+        });
+        message = {
+          ...message,
+          parts: updatedParts,
+          content: updatedParts,
+        };
+      }
+    }
+  }
+  return message;
+}
+
+export function toolResultStream(
+  tools: Record<string, Tool<any, any>> | undefined,
+  abortSignal: AbortSignal,
+) {
+  return new ToolExecutionStream({
+    execute: (toolCall) => getToolResponse(tools, abortSignal, toolCall),
+    streamCall: ({ reader, ...context }) =>
+      getToolStreamResponse(tools, abortSignal, reader, context),
+  });
+}

package/src/core/tool/type-path-utils.ts

@@ -0,0 +1,36 @@
+type AsNumber<K> = K extends `${infer N extends number}` ? N | K : never;
+type TupleIndex<T extends readonly any[]> = Exclude<keyof T, keyof any[]>;
+type ObjectKey<T> = keyof T & (string | number);
+
+export type TypePath<T> =
+  | []
+  | (0 extends 1 & T // IsAny
+      ? any[]
+      : T extends object // IsObjectOrArrayOrTuple
+        ? T extends readonly any[] // IsArrayOrTuple
+          ? number extends T["length"] // IsTuple
+            ? // Tuple: make union of [index, ...TypePath<element>]
+              {
+                [K in TupleIndex<T>]: [AsNumber<K>, ...TypePath<T[K]>];
+              }[TupleIndex<T>]
+            : // Array: use number index
+              [number, ...TypePath<T[number]>]
+          : // Object: make union of [key, ...TypePath<value>]
+            { [K in ObjectKey<T>]: [K, ...TypePath<T[K]>] }[ObjectKey<T>]
+        : // Base case: primitive values have no path
+          []);
+
+export type TypeAtPath<T, P extends readonly any[]> = P extends [
+  infer Head,
+  ...infer Rest,
+]
+  ? Head extends keyof T
+    ? TypeAtPath<T[Head], Rest>
+    : never
+  : T;
+
+export type DeepPartial<T> = T extends readonly any[]
+  ? readonly DeepPartial<T[number]>[]
+  : T extends { [key: string]: any }
+    ? { readonly [K in keyof T]?: DeepPartial<T[K]> }
+    : T;

package/src/core/utils/Counter.ts

@@ -0,0 +1,7 @@
+export class Counter {
+  public value = -1;
+
+  up() {
+    return ++this.value;
+  }
+}

package/src/core/utils/generateId.tsx

@@ -0,0 +1,6 @@
+import { customAlphabet } from "nanoid/non-secure";
+
+export const generateId = customAlphabet(
+  "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz",
+  7,
+);

package/src/core/utils/stream/AssistantMetaTransformStream.ts

@@ -0,0 +1,74 @@
+import { AssistantStreamChunk, PartInit } from "../../AssistantStreamChunk";
+
+/**
+ * For chunk types that are associated with a part,
+ * we require a non‑nullable meta field.
+ */
+export type AssistantMetaStreamChunk =
+  | (AssistantStreamChunk & {
+      type: "text-delta" | "part-finish";
+      meta: PartInit;
+    })
+  | (AssistantStreamChunk & {
+      type: "result" | "tool-call-args-text-finish";
+      meta: PartInit & { type: "tool-call" };
+    })
+  | (AssistantStreamChunk & {
+      type: Exclude<
+        AssistantStreamChunk["type"],
+        "text-delta" | "result" | "tool-call-args-text-finish" | "part-finish"
+      >;
+    });
+export class AssistantMetaTransformStream extends TransformStream<
+  AssistantStreamChunk,
+  AssistantMetaStreamChunk
+> {
+  constructor() {
+    // We use an array to record parts as they are introduced.
+    const parts: PartInit[] = [];
+
+    super({
+      transform(chunk, controller) {
+        // For chunks that introduce a new part.
+        if (chunk.type === "part-start") {
+          if (chunk.path.length !== 0) {
+            controller.error(new Error("Nested parts are not supported"));
+            return;
+          }
+          parts.push(chunk.part);
+          controller.enqueue(chunk);
+          return;
+        }
+
+        // For chunks that expect an associated part.
+        if (
+          chunk.type === "text-delta" ||
+          chunk.type === "result" ||
+          chunk.type === "part-finish" ||
+          chunk.type === "tool-call-args-text-finish"
+        ) {
+          if (chunk.path.length !== 1) {
+            controller.error(
+              new Error(`${chunk.type} chunks must have a path of length 1`),
+            );
+            return;
+          }
+          const idx = chunk.path[0]!;
+          if (idx < 0 || idx >= parts.length) {
+            controller.error(new Error(`Invalid path index: ${idx}`));
+            return;
+          }
+          const part = parts[idx]!;
+
+          controller.enqueue({
+            ...chunk,
+            meta: part as any, // TODO
+          });
+          return;
+        }
+
+        controller.enqueue(chunk);
+      },
+    });
+  }
+}

package/src/core/utils/stream/AssistantTransformStream.ts

@@ -0,0 +1,74 @@
+import { AssistantStreamChunk } from "../../AssistantStreamChunk";
+import {
+  AssistantStreamController,
+  createAssistantStreamController,
+} from "../../modules/assistant-stream";
+
+type AssistantTransformerFlushCallback = (
+  controller: AssistantStreamController,
+) => void | PromiseLike<void>;
+
+type AssistantTransformerStartCallback = (
+  controller: AssistantStreamController,
+) => void | PromiseLike<void>;
+
+type AssistantTransformerTransformCallback<I> = (
+  chunk: I,
+  controller: AssistantStreamController,
+) => void | PromiseLike<void>;
+
+type AssistantTransformer<I> = {
+  flush?: AssistantTransformerFlushCallback;
+  start?: AssistantTransformerStartCallback;
+  transform?: AssistantTransformerTransformCallback<I>;
+};
+
+export class AssistantTransformStream<I> extends TransformStream<
+  I,
+  AssistantStreamChunk
+> {
+  constructor(
+    transformer: AssistantTransformer<I>,
+    writableStrategy?: QueuingStrategy<I>,
+    readableStrategy?: QueuingStrategy<AssistantStreamChunk>,
+  ) {
+    const [stream, runController] = createAssistantStreamController();
+
+    let runPipeTask: Promise<void>;
+    super(
+      {
+        start(controller) {
+          runPipeTask = stream
+            .pipeTo(
+              new WritableStream({
+                write(chunk) {
+                  controller.enqueue(chunk);
+                },
+                abort(reason?: any) {
+                  controller.error(reason);
+                },
+                close() {
+                  controller.terminate();
+                },
+              }),
+            )
+            .catch((error) => {
+              controller.error(error);
+            });
+
+          return transformer.start?.(runController);
+        },
+        transform(chunk) {
+          return transformer.transform?.(chunk, runController);
+        },
+        async flush() {
+          await transformer.flush?.(runController);
+          runController.close();
+          await runPipeTask;
+        },
+      },
+      writableStrategy,
+      readableStrategy,
+    );
+  }
+}

package/src/core/utils/stream/LineDecoderStream.ts

@@ -0,0 +1,29 @@
+export class LineDecoderStream extends TransformStream<string, string> {
+  private buffer = "";
+
+  constructor() {
+    super({
+      transform: (chunk, controller) => {
+        this.buffer += chunk;
+        const lines = this.buffer.split("\n");
+
+        // Process all complete lines
+        for (let i = 0; i < lines.length - 1; i++) {
+          controller.enqueue(lines[i]);
+        }
+
+        // Keep the last incomplete line in the buffer
+        this.buffer = lines[lines.length - 1] || "";
+      },
+      flush: () => {
+        // If there's content in the buffer when the stream ends, it means
+        // the stream ended with an incomplete line (no trailing newline)
+        if (this.buffer) {
+          throw new Error(
+            `Stream ended with an incomplete line: "${this.buffer}"`,
+          );
+        }
+      },
+    });
+  }
+}

package/src/core/utils/stream/PipeableTransformStream.ts

@@ -0,0 +1,10 @@
+export class PipeableTransformStream<I, O> extends TransformStream<I, O> {
+  constructor(transform: (readable: ReadableStream<I>) => ReadableStream<O>) {
+    super();
+    const readable = transform(super.readable as unknown as ReadableStream<I>);
+    Object.defineProperty(this, "readable", {
+      value: readable,
+      writable: false,
+    });
+  }
+}

package/src/core/utils/stream/UnderlyingReadable.ts

@@ -0,0 +1,5 @@
+export type UnderlyingReadable<TController> = {
+  start?: (controller: TController) => void;
+  pull?: (controller: TController) => void | PromiseLike<void>;
+  cancel?: UnderlyingSourceCancelCallback;
+};

package/src/core/utils/stream/merge.ts

@@ -0,0 +1,212 @@
+import { AssistantStreamChunk } from "../../AssistantStreamChunk";
+import { promiseWithResolvers } from "../../../utils/promiseWithResolvers";
+
+type MergeStreamItem = {
+  reader: ReadableStreamDefaultReader<AssistantStreamChunk>;
+  promise?: Promise<unknown> | undefined;
+};
+
+export const createMergeStream = () => {
+  const list: MergeStreamItem[] = [];
+  let sealed = false;
+  let controller: ReadableStreamDefaultController<AssistantStreamChunk>;
+  let currentPull: ReturnType<typeof promiseWithResolvers<void>> | undefined;
+
+  const handlePull = (item: MergeStreamItem) => {
+    if (!item.promise) {
+      // TODO for most streams, we can directly pipeTo to avoid the microTask queue
+      // add an option to eagerly pipe the stream to the merge stream
+      // ideally, using assitant-stream w sync run method + piping to a sync WritableStream runs in the same microtask
+      // this is useful because we often use AssistantStreams internally as a serialization utility, e. g. AssistantTransformStream
+      // idea: avoid reader.read() by instead using a WritableStream & if (!hasPendingPull) await waitForPull()?
+      item.promise = item.reader
+        .read()
+        .then(({ done, value }) => {
+          item.promise = undefined;
+          if (done) {
+            list.splice(list.indexOf(item), 1);
+            if (sealed && list.length === 0) {
+              controller.close();
+            }
+          } else {
+            controller.enqueue(value);
+          }
+
+          currentPull?.resolve();
+          currentPull = undefined;
+        })
+        .catch((e) => {
+          console.error(e);
+
+          list.forEach((item) => {
+            item.reader.cancel();
+          });
+          list.length = 0;
+
+          controller.error(e);
+
+          currentPull?.reject(e);
+          currentPull = undefined;
+        });
+    }
+  };
+
+  const readable = new ReadableStream<AssistantStreamChunk>({
+    start(c) {
+      controller = c;
+    },
+    pull() {
+      currentPull = promiseWithResolvers();
+      list.forEach((item) => {
+        handlePull(item);
+      });
+
+      return currentPull.promise;
+    },
+    cancel() {
+      list.forEach((item) => {
+        item.reader.cancel();
+      });
+      list.length = 0;
+    },
+  });
+
+  return {
+    readable,
+    isSealed() {
+      return sealed;
+    },
+    seal() {
+      sealed = true;
+      if (list.length === 0) controller.close();
+    },
+    addStream(stream: ReadableStream<AssistantStreamChunk>) {
+      if (sealed)
+        throw new Error(
+          "Cannot add streams after the run callback has settled.",
+        );
+
+      const item = { reader: stream.getReader() };
+      list.push(item);
+      handlePull(item);
+    },
+    enqueue(chunk: AssistantStreamChunk) {
+      this.addStream(
+        new ReadableStream({
+          start(c) {
+            c.enqueue(chunk);
+            c.close();
+          },
+        }),
+      );
+    },
+  };
+};
+
+// TODO
+// export class SpanContainerMerger {
+//   public get isSealed() {
+//     return this.mergeStream.isSealed();
+//   }
+
+//   public get readable() {
+//     return this.mergeStream.readable;
+//   }
+
+//   private subAllocator = new Counter();
+//   private mergeStream = createMergeStream();
+
+//   constructor() {
+//     // id 0 is auto allocated
+//     this.subAllocator.up();
+//   }
+
+//   add(stream: ReadableStream<AssistantStreamChunk>) {
+//     this.mergeStream.addStream(
+//       stream.pipeThrough(new SpanParentEncoder(this.subAllocator)),
+//     );
+//   }
+
+//   enqueue(chunk: AssistantStreamChunk & { parentId: 0 }) {
+//     this.mergeStream.addStream(
+//       new ReadableStream({
+//         start(c) {
+//           c.enqueue(chunk);
+//           c.close();
+//         },
+//       }),
+//     );
+//   }
+
+//   seal() {
+//     this.mergeStream.seal();
+//   }
+// }
+
+// export class SpanContainerSplitter {
+//   public writable;
+
+//   private isSealed = false;
+//   private writers = new Map<
+//     number,
+//     WritableStreamDefaultWriter<AssistantStreamChunk>
+//   >();
+
+//   private closeTasks: Promise<void>[] = [];
+
+//   private allocator = new Counter();
+//   private subAllocator = new Counter();
+
+//   constructor() {
+//     // id 0 is auto-allocated
+//     this.allocator.up();
+
+//     this.writable = new WritableStream({
+//       write: (chunk) => {
+//         const { type, parentId } = chunk;
+
+//         const writer = this.writers.get(parentId);
+//         if (writer === undefined) throw new Error("Parent id not found");
+
+//         writer.write(chunk);
+
+//         if (type === "span") {
+//           // allocate a new span id
+//           this.writers.set(this.allocator.up(), writer);
+//         }
+//         if (type === "finish") {
+//           this.writers.delete(parentId);
+//           writer.close();
+
+//           if (this.writers.size === 0) {
+//             const closeTask = this.writable.close();
+//             this.closeTasks.push(closeTask);
+//             closeTask.then(() => {
+//               this.closeTasks.splice(this.closeTasks.indexOf(closeTask), 1);
+//             });
+//           }
+//         }
+//       },
+//       close: async () => {
+//         if (this.writers.size > 0) throw new Error("Not all writers closed");
+
+//         // await and throw on any errors
+//         await Promise.all(this.closeTasks);
+//       },
+//     });
+//   }
+
+//   add(stream: WritableStream<AssistantStreamChunk>) {
+//     if (this.isSealed) throw new Error("Cannot add streams after sealing");
+
+//     const decoder = new SpanParentDecoder(this.subAllocator);
+//     decoder.readable.pipeTo(stream);
+
+//     this.writers.set(this.allocator.up(), decoder.writable.getWriter());
+//   }
+
+//   seal() {
+//     this.isSealed = true;
+//     if (this.writers.size === 0) this.writable.close();
+//   }
+// }