assistant-stream 0.1.8 → 0.2.1

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/{dist/core/tool/ToolExecutionStream.mjs → src/core/tool/ToolExecutionStream.ts}

@@ -1,115 +1,164 @@
-// src/core/tool/ToolExecutionStream.ts
 import sjson from "secure-json-parse";
+import { AssistantStreamChunk } from "../AssistantStreamChunk";
 import {
-  AssistantMetaTransformStream
-} from "../utils/stream/AssistantMetaTransformStream.mjs";
-import { PipeableTransformStream } from "../utils/stream/PipeableTransformStream.mjs";
-import { ToolResponse } from "./ToolResponse.mjs";
-import { withPromiseOrValue } from "../utils/withPromiseOrValue.mjs";
-import { ToolCallReaderImpl } from "./ToolCallReader.mjs";
-var ToolExecutionStream = class extends PipeableTransformStream {
-  constructor(options) {
-    const toolCallPromises = /* @__PURE__ */ new Map();
-    const toolCallControllers = /* @__PURE__ */ new Map();
+  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({
+      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();
+                const reader = new ToolCallReaderImpl<unknown, unknown>();
                 toolCallControllers.set(chunk.part.toolCallId, reader);
+
                 options.streamCall({
                   reader,
                   toolCallId: chunk.part.toolCallId,
-                  toolName: chunk.part.toolName
+                  toolName: chunk.part.toolName,
                 });
               }
               break;
             case "text-delta": {
               if (chunk.meta.type === "tool-call") {
                 const toolCallId = chunk.meta.toolCallId;
-                const controller2 = toolCallControllers.get(toolCallId);
-                if (!controller2)
+
+                const controller = toolCallControllers.get(toolCallId);
+                if (!controller)
                   throw new Error("No controller found for tool call");
-                controller2.appendArgsTextDelta(chunk.textDelta);
+                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);
+              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"
+                      "Encountered tool call without args, this should never happen",
                     );
                     throw new Error(
-                      "Encountered tool call without args, this is unexpected."
+                      "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.message)}`
+                      `Function parameter parsing failed. ${JSON.stringify((e as Error).message)}`,
                     );
                   }
+
                   return options.execute({
                     toolCallId,
                     toolName,
-                    args
+                    args,
                   });
                 },
                 (c) => {
-                  if (c === void 0) return;
+                  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
+                    isError: c.isError,
                   });
                   streamController.setResponse(result);
                   controller.enqueue({
                     type: "result",
                     path: chunk.path,
-                    ...result
+                    ...result,
                   });
                 },
                 (e) => {
                   const result = new ToolResponse({
                     result: String(e),
-                    isError: true
+                    isError: true,
                   });
+
                   streamController.setResponse(result);
                   controller.enqueue({
                     type: "result",
                     path: chunk.path,
-                    ...result
+                    ...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 {
@@ -120,13 +169,12 @@ var ToolExecutionStream = class extends PipeableTransformStream {
         },
         async flush() {
           await Promise.all(toolCallPromises.values());
-        }
+        },
       });
-      return readable.pipeThrough(new AssistantMetaTransformStream()).pipeThrough(transform);
+
+      return readable
+        .pipeThrough(new AssistantMetaTransformStream())
+        .pipeThrough(transform);
     });
   }
-};
-export {
-  ToolExecutionStream
-};
-//# sourceMappingURL=ToolExecutionStream.mjs.map
+}

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";