assistant-stream 0.2.0 → 0.2.2

package/src/core/tool/ToolCallReader.ts

@@ -0,0 +1,432 @@
+import { promiseWithResolvers } from "../../utils/promiseWithResolvers";
+import {
+  parsePartialJsonObject,
+  getPartialJsonObjectFieldState,
+} from "../../utils/json/parse-partial-json-object";
+import {
+  ToolCallArgsReader,
+  ToolCallReader,
+  ToolCallResponseReader,
+} from "./tool-types";
+import { TypeAtPath, TypePath } from "./type-path-utils";
+import { ToolResponse } from "./ToolResponse";
+
+// TODO: remove dispose
+
+function getField<T>(obj: T, fieldPath: (string | number)[]): any {
+  let current: any = obj;
+  for (const key of fieldPath) {
+    if (current === undefined || current === null) {
+      return undefined;
+    }
+    current = current[key as string | number];
+  }
+  return current;
+}
+
+interface Handle {
+  update(args: unknown): void;
+  dispose(): void;
+}
+
+class GetHandle<T> implements Handle {
+  private resolve: (value: any) => void;
+  private reject: (reason: unknown) => void;
+  private disposed = false;
+  private fieldPath: (string | number)[];
+
+  constructor(
+    resolve: (value: any) => void,
+    reject: (reason: unknown) => void,
+    fieldPath: (string | number)[],
+  ) {
+    this.resolve = resolve;
+    this.reject = reject;
+    this.fieldPath = fieldPath;
+  }
+
+  update(args: unknown): void {
+    if (this.disposed) return;
+
+    try {
+      // Check if the field is complete
+      if (
+        getPartialJsonObjectFieldState(
+          args as Record<string, unknown>,
+          this.fieldPath,
+        ) === "complete"
+      ) {
+        const value = getField(args as T, this.fieldPath);
+        if (value !== undefined) {
+          this.resolve(value);
+          this.dispose();
+        }
+      }
+    } catch (e) {
+      this.reject(e);
+      this.dispose();
+    }
+  }
+
+  dispose(): void {
+    this.disposed = true;
+  }
+}
+
+class StreamValuesHandle<T> implements Handle {
+  private controller: ReadableStreamDefaultController<any>;
+  private disposed = false;
+  private fieldPath: (string | number)[];
+
+  constructor(
+    controller: ReadableStreamDefaultController<any>,
+    fieldPath: (string | number)[],
+  ) {
+    this.controller = controller;
+    this.fieldPath = fieldPath;
+  }
+
+  update(args: unknown): void {
+    if (this.disposed) return;
+
+    try {
+      const value = getField(args as T, this.fieldPath);
+
+      if (value !== undefined) {
+        this.controller.enqueue(value);
+      }
+
+      // Check if the field is complete, if so close the stream
+      if (
+        getPartialJsonObjectFieldState(
+          args as Record<string, unknown>,
+          this.fieldPath,
+        ) === "complete"
+      ) {
+        this.controller.close();
+        this.dispose();
+      }
+    } catch (e) {
+      this.controller.error(e);
+      this.dispose();
+    }
+  }
+
+  dispose(): void {
+    this.disposed = true;
+  }
+}
+
+class StreamTextHandle<T> implements Handle {
+  private controller: ReadableStreamDefaultController<any>;
+  private disposed = false;
+  private fieldPath: (string | number)[];
+  private lastValue: any = undefined;
+
+  constructor(
+    controller: ReadableStreamDefaultController<any>,
+    fieldPath: (string | number)[],
+  ) {
+    this.controller = controller;
+    this.fieldPath = fieldPath;
+  }
+
+  update(args: unknown): void {
+    if (this.disposed) return;
+
+    try {
+      const value = getField(args as T, this.fieldPath);
+
+      if (value !== undefined && typeof value === "string") {
+        const delta = value.substring(this.lastValue?.length || 0);
+        this.lastValue = value;
+        this.controller.enqueue(delta);
+      }
+
+      // Check if the field is complete, if so close the stream
+      if (
+        getPartialJsonObjectFieldState(
+          args as Record<string, unknown>,
+          this.fieldPath,
+        ) === "complete"
+      ) {
+        this.controller.close();
+        this.dispose();
+      }
+    } catch (e) {
+      this.controller.error(e);
+      this.dispose();
+    }
+  }
+
+  dispose(): void {
+    this.disposed = true;
+  }
+}
+
+class ForEachHandle<T> implements Handle {
+  private controller: ReadableStreamDefaultController<any>;
+  private disposed = false;
+  private fieldPath: (string | number)[];
+  private processedIndexes = new Set<number>();
+
+  constructor(
+    controller: ReadableStreamDefaultController<any>,
+    fieldPath: (string | number)[],
+  ) {
+    this.controller = controller;
+    this.fieldPath = fieldPath;
+  }
+
+  update(args: unknown): void {
+    if (this.disposed) return;
+
+    try {
+      const array = getField(args as T, this.fieldPath) as unknown as any[];
+
+      if (!Array.isArray(array)) {
+        return;
+      }
+
+      // Check each array element and emit completed ones that haven't been processed
+      for (let i = 0; i < array.length; i++) {
+        if (!this.processedIndexes.has(i)) {
+          const elementPath = [...this.fieldPath, i];
+          if (
+            getPartialJsonObjectFieldState(
+              args as Record<string, unknown>,
+              elementPath,
+            ) === "complete"
+          ) {
+            this.controller.enqueue(array[i]);
+            this.processedIndexes.add(i);
+          }
+        }
+      }
+
+      // Check if the entire array is complete
+      if (
+        getPartialJsonObjectFieldState(
+          args as Record<string, unknown>,
+          this.fieldPath,
+        ) === "complete"
+      ) {
+        this.controller.close();
+        this.dispose();
+      }
+    } catch (e) {
+      this.controller.error(e);
+      this.dispose();
+    }
+  }
+
+  dispose(): void {
+    this.disposed = true;
+  }
+}
+
+// Implementation of ToolCallReader that uses stream of partial JSON
+export class ToolCallArgsReaderImpl<T> implements ToolCallArgsReader<T> {
+  private argTextDeltas: ReadableStream<string>;
+  private handles: Set<Handle> = new Set();
+  private args: any = parsePartialJsonObject("");
+
+  constructor(argTextDeltas: ReadableStream<string>) {
+    this.argTextDeltas = argTextDeltas;
+    this.processStream();
+  }
+
+  private async processStream(): Promise<void> {
+    try {
+      let accumulatedText = "";
+      const reader = this.argTextDeltas.getReader();
+
+      while (true) {
+        const { value, done } = await reader.read();
+        if (done) break;
+
+        accumulatedText += value;
+        const parsedArgs = parsePartialJsonObject(accumulatedText);
+
+        if (parsedArgs !== undefined) {
+          this.args = parsedArgs;
+          // Notify all handles of the updated args
+          for (const handle of this.handles) {
+            handle.update(parsedArgs);
+          }
+        }
+      }
+    } catch (error) {
+      console.error("Error processing argument stream:", error);
+      // Notify handles of the error
+      for (const handle of this.handles) {
+        handle.dispose();
+      }
+    }
+  }
+
+  get<PathT extends TypePath<T>>(
+    ...fieldPath: PathT
+  ): Promise<TypeAtPath<T, PathT>> {
+    return new Promise<any>((resolve, reject) => {
+      const handle = new GetHandle<T>(resolve, reject, fieldPath);
+
+      // Check if the field is already complete in current args
+      if (
+        this.args &&
+        getPartialJsonObjectFieldState(
+          this.args as Record<string, unknown>,
+          fieldPath,
+        ) === "complete"
+      ) {
+        const value = getField(this.args as T, fieldPath);
+        if (value !== undefined) {
+          resolve(value);
+          return;
+        }
+      }
+
+      this.handles.add(handle);
+      handle.update(this.args);
+    });
+  }
+
+  streamValues<PathT extends TypePath<T>>(...fieldPath: PathT): any {
+    // Use a type assertion to convert the complex TypePath to a simple array
+    const simplePath = fieldPath as unknown as (string | number)[];
+
+    const stream = new ReadableStream<any>({
+      start: (controller) => {
+        const handle = new StreamValuesHandle<T>(controller, simplePath);
+        this.handles.add(handle);
+
+        // Check current args immediately
+        handle.update(this.args);
+      },
+      cancel: () => {
+        // Find and dispose the corresponding handle
+        for (const handle of this.handles) {
+          if (handle instanceof StreamValuesHandle) {
+            handle.dispose();
+            this.handles.delete(handle);
+            break;
+          }
+        }
+      },
+    });
+
+    // For type compatibility, cast the stream to the required type
+    return stream as any;
+  }
+
+  streamText<PathT extends TypePath<T>>(...fieldPath: PathT): any {
+    // Use a type assertion to convert the complex TypePath to a simple array
+    const simplePath = fieldPath as unknown as (string | number)[];
+
+    const stream = new ReadableStream<any>({
+      start: (controller) => {
+        const handle = new StreamTextHandle<T>(controller, simplePath);
+        this.handles.add(handle);
+
+        // Check current args immediately
+        handle.update(this.args);
+      },
+      cancel: () => {
+        // Find and dispose the corresponding handle
+        for (const handle of this.handles) {
+          if (handle instanceof StreamTextHandle) {
+            handle.dispose();
+            this.handles.delete(handle);
+            break;
+          }
+        }
+      },
+    });
+
+    // For type compatibility, cast the stream to the required type
+    return stream as any;
+  }
+
+  forEach<PathT extends TypePath<T>>(...fieldPath: PathT): any {
+    // Use a type assertion to convert the complex TypePath to a simple array
+    const simplePath = fieldPath as unknown as (string | number)[];
+
+    const stream = new ReadableStream<any>({
+      start: (controller) => {
+        const handle = new ForEachHandle<T>(controller, simplePath);
+        this.handles.add(handle);
+
+        // Check current args immediately
+        handle.update(this.args);
+      },
+      cancel: () => {
+        // Find and dispose the corresponding handle
+        for (const handle of this.handles) {
+          if (handle instanceof ForEachHandle) {
+            handle.dispose();
+            this.handles.delete(handle);
+            break;
+          }
+        }
+      },
+    });
+
+    // For type compatibility, cast the stream to the required type
+    return stream as any;
+  }
+}
+
+export class ToolCallResponseReaderImpl<TResult>
+  implements ToolCallResponseReader<TResult>
+{
+  constructor(private readonly promise: Promise<ToolResponse<TResult>>) {}
+
+  public get() {
+    return this.promise;
+  }
+}
+
+export class ToolCallReaderImpl<TArgs, TResult>
+  implements ToolCallReader<TArgs, TResult>
+{
+  public readonly args: ToolCallArgsReaderImpl<TArgs>;
+  public readonly response: ToolCallResponseReaderImpl<TResult>;
+  private readonly writable: WritableStream<string>;
+  private readonly resolve: (value: ToolResponse<TResult>) => void;
+
+  public argsText: string = "";
+
+  constructor() {
+    const stream = new TransformStream<string, string>();
+    this.writable = stream.writable;
+    this.args = new ToolCallArgsReaderImpl<TArgs>(stream.readable);
+
+    const { promise, resolve } = promiseWithResolvers<ToolResponse<TResult>>();
+    this.resolve = resolve;
+    this.response = new ToolCallResponseReaderImpl<TResult>(promise);
+  }
+
+  async appendArgsTextDelta(text: string): Promise<void> {
+    const writer = this.writable.getWriter();
+    try {
+      await writer.write(text);
+    } catch (err) {
+      console.warn(err);
+    } finally {
+      writer.releaseLock();
+    }
+
+    this.argsText += text;
+  }
+
+  setResponse(value: ToolResponse<TResult>): void {
+    this.resolve(value);
+  }
+
+  result = {
+    get: async () => {
+      const response = await this.response.get();
+      return response.result;
+    },
+  };
+}

package/src/core/tool/ToolExecutionStream.ts

@@ -0,0 +1,180 @@
+import sjson from "secure-json-parse";
+import { AssistantStreamChunk } from "../AssistantStreamChunk";
+import {
+  AssistantMetaStreamChunk,
+  AssistantMetaTransformStream,
+} from "../utils/stream/AssistantMetaTransformStream";
+import { PipeableTransformStream } from "../utils/stream/PipeableTransformStream";
+import { ReadonlyJSONValue } from "../../utils/json/json-value";
+import { ToolResponse } from "./ToolResponse";
+import { withPromiseOrValue } from "../utils/withPromiseOrValue";
+import { ToolCallReaderImpl } from "./ToolCallReader";
+import { ToolCallReader } from "./tool-types";
+
+type ToolCallback = (toolCall: {
+  toolCallId: string;
+  toolName: string;
+  args: unknown;
+}) =>
+  | Promise<ToolResponse<ReadonlyJSONValue>>
+  | ToolResponse<ReadonlyJSONValue>
+  | undefined;
+
+type ToolStreamCallback = <TArgs, TResult>(toolCall: {
+  reader: ToolCallReader<TArgs, TResult>;
+  toolCallId: string;
+  toolName: string;
+}) => void;
+
+type ToolExecutionOptions = {
+  execute: ToolCallback;
+  streamCall: ToolStreamCallback;
+};
+
+export class ToolExecutionStream extends PipeableTransformStream<
+  AssistantStreamChunk,
+  AssistantStreamChunk
+> {
+  constructor(options: ToolExecutionOptions) {
+    const toolCallPromises = new Map<string, PromiseLike<void>>();
+    const toolCallControllers = new Map<
+      string,
+      ToolCallReaderImpl<unknown, unknown>
+    >();
+
+    super((readable) => {
+      const transform = new TransformStream<
+        AssistantMetaStreamChunk,
+        AssistantStreamChunk
+      >({
+        transform(chunk, controller) {
+          // forward everything
+          if (chunk.type !== "part-finish" || chunk.meta.type !== "tool-call") {
+            controller.enqueue(chunk);
+          }
+
+          const type = chunk.type;
+
+          switch (type) {
+            case "part-start":
+              if (chunk.part.type === "tool-call") {
+                const reader = new ToolCallReaderImpl<unknown, unknown>();
+                toolCallControllers.set(chunk.part.toolCallId, reader);
+
+                options.streamCall({
+                  reader,
+                  toolCallId: chunk.part.toolCallId,
+                  toolName: chunk.part.toolName,
+                });
+              }
+              break;
+            case "text-delta": {
+              if (chunk.meta.type === "tool-call") {
+                const toolCallId = chunk.meta.toolCallId;
+
+                const controller = toolCallControllers.get(toolCallId);
+                if (!controller)
+                  throw new Error("No controller found for tool call");
+                controller.appendArgsTextDelta(chunk.textDelta);
+              }
+              break;
+            }
+            case "tool-call-args-text-finish": {
+              if (chunk.meta.type !== "tool-call") break;
+
+              const { toolCallId, toolName } = chunk.meta;
+              const streamController = toolCallControllers.get(toolCallId)!;
+              if (!streamController)
+                throw new Error("No controller found for tool call");
+
+              const promise = withPromiseOrValue(
+                () => {
+                  if (!streamController.argsText) {
+                    console.log(
+                      "Encountered tool call without args, this should never happen",
+                    );
+                    throw new Error(
+                      "Encountered tool call without args, this is unexpected.",
+                    );
+                  }
+
+                  let args;
+                  try {
+                    args = sjson.parse(streamController.argsText);
+                  } catch (e) {
+                    throw new Error(
+                      `Function parameter parsing failed. ${JSON.stringify((e as Error).message)}`,
+                    );
+                  }
+
+                  return options.execute({
+                    toolCallId,
+                    toolName,
+                    args,
+                  });
+                },
+                (c) => {
+                  if (c === undefined) return;
+
+                  // TODO how to handle new ToolResult({ result: undefined })?
+                  const result = new ToolResponse({
+                    artifact: c.artifact,
+                    result: c.result,
+                    isError: c.isError,
+                  });
+                  streamController.setResponse(result);
+                  controller.enqueue({
+                    type: "result",
+                    path: chunk.path,
+                    ...result,
+                  });
+                },
+                (e) => {
+                  const result = new ToolResponse({
+                    result: String(e),
+                    isError: true,
+                  });
+
+                  streamController.setResponse(result);
+                  controller.enqueue({
+                    type: "result",
+                    path: chunk.path,
+                    ...result,
+                  });
+                },
+              );
+              if (promise) {
+                toolCallPromises.set(toolCallId, promise);
+              }
+              break;
+            }
+
+            case "part-finish": {
+              if (chunk.meta.type !== "tool-call") break;
+
+              const { toolCallId } = chunk.meta;
+              const toolCallPromise = toolCallPromises.get(toolCallId);
+              if (toolCallPromise) {
+                toolCallPromise.then(() => {
+                  toolCallPromises.delete(toolCallId);
+                  toolCallControllers.delete(toolCallId);
+
+                  controller.enqueue(chunk);
+                });
+              } else {
+                controller.enqueue(chunk);
+              }
+            }
+          }
+        },
+        async flush() {
+          await Promise.all(toolCallPromises.values());
+        },
+      });
+
+      return readable
+        .pipeThrough(new AssistantMetaTransformStream())
+        .pipeThrough(transform);
+    });
+  }
+}

package/src/core/tool/ToolResponse.ts

@@ -0,0 +1,31 @@
+import { ReadonlyJSONValue } from "../../utils/json/json-value";
+
+const TOOL_RESPONSE_SYMBOL = Symbol.for("aui.tool-response");
+
+export type ToolResponseInit<TResult> = {
+  result: TResult;
+  artifact?: ReadonlyJSONValue | undefined;
+  isError?: boolean | undefined;
+};
+
+export class ToolResponse<TResult> {
+  get [TOOL_RESPONSE_SYMBOL]() {
+    return true;
+  }
+
+  readonly artifact?: ReadonlyJSONValue | undefined;
+  readonly result: TResult;
+  readonly isError: boolean;
+
+  constructor(options: ToolResponseInit<TResult>) {
+    this.artifact = options.artifact;
+    this.result = options.result;
+    this.isError = options.isError ?? false;
+  }
+
+  static [Symbol.hasInstance](obj: unknown): obj is ToolResponse<unknown> {
+    return (
+      typeof obj === "object" && obj !== null && TOOL_RESPONSE_SYMBOL in obj
+    );
+  }
+}

package/src/core/tool/index.ts

@@ -0,0 +1,8 @@
+export type { Tool } from "./tool-types";
+export { ToolResponse } from "./ToolResponse";
+export { ToolExecutionStream } from "./ToolExecutionStream";
+export type { ToolCallReader } from "./tool-types";
+export {
+  toolResultStream as unstable_toolResultStream,
+  unstable_runPendingTools,
+} from "./toolResultStream";