assistant-stream 0.1.8 → 0.2.1

package/{dist/core/modules/assistant-stream.mjs → src/core/modules/assistant-stream.ts}

@@ -1,141 +1,201 @@
-// src/core/modules/assistant-stream.ts
-import { AssistantStream } from "../AssistantStream.mjs";
-import { createMergeStream } from "../utils/stream/merge.mjs";
-import { createTextStreamController } from "./text.mjs";
+import { AssistantStream } from "../AssistantStream";
+import { AssistantStreamChunk, PartInit } from "../AssistantStreamChunk";
+import { createMergeStream } from "../utils/stream/merge";
+import { createTextStreamController, TextStreamController } from "./text";
 import {
-  createToolCallStreamController
-} from "./tool-call.mjs";
-import { Counter } from "../utils/Counter.mjs";
+  createToolCallStreamController,
+  ToolCallStreamController,
+} from "./tool-call";
+import { Counter } from "../utils/Counter";
 import {
   PathAppendEncoder,
-  PathMergeEncoder
-} from "../utils/stream/path-utils.mjs";
-import { DataStreamEncoder } from "../serialization/data-stream/DataStream.mjs";
-import { generateId } from "../utils/generateId.mjs";
-import { promiseWithResolvers } from "../../utils/promiseWithResolvers.mjs";
-var AssistantStreamControllerImpl = class {
-  _merger = createMergeStream();
-  _append;
-  _contentCounter = new Counter();
+  PathMergeEncoder,
+} from "../utils/stream/path-utils";
+import { DataStreamEncoder } from "../serialization/data-stream/DataStream";
+import { FilePart, SourcePart } from "../utils/types";
+import { generateId } from "../utils/generateId";
+import {
+  ReadonlyJSONObject,
+  ReadonlyJSONValue,
+} from "../../utils/json/json-value";
+import { ToolResponseInit } from "../tool/ToolResponse";
+import { promiseWithResolvers } from "../../utils/promiseWithResolvers";
+
+type ToolCallPartInit = {
+  toolCallId?: string;
+  toolName: string;
+  argsText?: string;
+  args?: ReadonlyJSONObject;
+  response?: ToolResponseInit<ReadonlyJSONValue>;
+};
+
+export type AssistantStreamController = {
+  appendText(textDelta: string): void;
+  appendReasoning(reasoningDelta: string): void;
+  appendSource(options: SourcePart): void;
+  appendFile(options: FilePart): void;
+  addTextPart(): TextStreamController;
+  addToolCallPart(options: string): ToolCallStreamController;
+  addToolCallPart(options: ToolCallPartInit): ToolCallStreamController;
+  enqueue(chunk: AssistantStreamChunk): void;
+  merge(stream: AssistantStream): void;
+  close(): void;
+};
+
+class AssistantStreamControllerImpl implements AssistantStreamController {
+  private _merger = createMergeStream();
+  private _append:
+    | {
+        controller: TextStreamController;
+        kind: "text" | "reasoning";
+      }
+    | undefined;
+  private _contentCounter = new Counter();
+
   get __internal_isClosed() {
     return this._merger.isSealed();
   }
+
   __internal_getReadable() {
     return this._merger.readable;
   }
-  _closeSubscriber;
-  __internal_subscribeToClose(callback) {
+
+  private _closeSubscriber: undefined | (() => void);
+  __internal_subscribeToClose(callback: () => void) {
     this._closeSubscriber = callback;
   }
-  _addPart(part, stream) {
+
+  private _addPart(part: PartInit, stream: AssistantStream) {
     if (this._append) {
       this._append.controller.close();
-      this._append = void 0;
+      this._append = undefined;
     }
+
     this.enqueue({
       type: "part-start",
       part,
-      path: []
+      path: [],
     });
     this._merger.addStream(
-      stream.pipeThrough(new PathAppendEncoder(this._contentCounter.value))
+      stream.pipeThrough(new PathAppendEncoder(this._contentCounter.value)),
     );
   }
-  merge(stream) {
+
+  merge(stream: AssistantStream) {
     this._merger.addStream(
-      stream.pipeThrough(new PathMergeEncoder(this._contentCounter))
+      stream.pipeThrough(new PathMergeEncoder(this._contentCounter)),
     );
   }
-  appendText(textDelta) {
+
+  appendText(textDelta: string) {
     if (this._append?.kind !== "text") {
       this._append = {
         kind: "text",
-        controller: this.addTextPart()
+        controller: this.addTextPart(),
       };
     }
     this._append.controller.append(textDelta);
   }
-  appendReasoning(textDelta) {
+
+  appendReasoning(textDelta: string) {
     if (this._append?.kind !== "reasoning") {
       this._append = {
         kind: "reasoning",
-        controller: this.addReasoningPart()
+        controller: this.addReasoningPart(),
       };
     }
     this._append.controller.append(textDelta);
   }
+
   addTextPart() {
     const [stream, controller] = createTextStreamController();
     this._addPart({ type: "text" }, stream);
     return controller;
   }
+
   addReasoningPart() {
     const [stream, controller] = createTextStreamController();
     this._addPart({ type: "reasoning" }, stream);
     return controller;
   }
-  addToolCallPart(options) {
+
+  addToolCallPart(
+    options: string | ToolCallPartInit,
+  ): ToolCallStreamController {
     const opt = typeof options === "string" ? { toolName: options } : options;
     const toolName = opt.toolName;
     const toolCallId = opt.toolCallId ?? generateId();
+
     const [stream, controller] = createToolCallStreamController();
     this._addPart({ type: "tool-call", toolName, toolCallId }, stream);
-    if (opt.argsText !== void 0) {
+
+    if (opt.argsText !== undefined) {
       controller.argsText.append(opt.argsText);
       controller.argsText.close();
     }
-    if (opt.args !== void 0) {
+    if (opt.args !== undefined) {
       controller.argsText.append(JSON.stringify(opt.args));
       controller.argsText.close();
     }
-    if (opt.response !== void 0) {
+    if (opt.response !== undefined) {
       controller.setResponse(opt.response);
     }
+
     return controller;
   }
-  appendSource(options) {
+
+  appendSource(options: SourcePart) {
     this._addPart(
       options,
       new ReadableStream({
         start(controller) {
           controller.enqueue({
             type: "part-finish",
-            path: []
+            path: [],
           });
           controller.close();
-        }
-      })
+        },
+      }),
     );
   }
-  appendFile(options) {
+
+  appendFile(options: FilePart) {
     this._addPart(
       options,
       new ReadableStream({
         start(controller) {
           controller.enqueue({
             type: "part-finish",
-            path: []
+            path: [],
           });
           controller.close();
-        }
-      })
+        },
+      }),
     );
   }
-  enqueue(chunk) {
+
+  enqueue(chunk: AssistantStreamChunk) {
     this._merger.enqueue(chunk);
+
     if (chunk.type === "part-start" && chunk.path.length === 0) {
       this._contentCounter.up();
     }
   }
+
   close() {
     this._merger.seal();
     this._append?.controller?.close();
+
     this._closeSubscriber?.();
   }
-};
-function createAssistantStream(callback) {
+}
+
+export function createAssistantStream(
+  callback: (controller: AssistantStreamController) => PromiseLike<void> | void,
+): AssistantStream {
   const controller = new AssistantStreamControllerImpl();
-  let promiseOrVoid;
+
+  let promiseOrVoid: PromiseLike<void> | void;
   try {
     promiseOrVoid = callback(controller);
   } catch (e) {
@@ -143,12 +203,13 @@ function createAssistantStream(callback) {
       controller.enqueue({
         type: "error",
         path: [],
-        error: String(e)
+        error: String(e),
       });
       controller.close();
     }
     throw e;
   }
+
   if (promiseOrVoid instanceof Promise) {
     const runTask = async () => {
       try {
@@ -158,7 +219,7 @@ function createAssistantStream(callback) {
           controller.enqueue({
             type: "error",
             path: [],
-            error: String(e)
+            error: String(e),
           });
         }
         throw e;
@@ -174,29 +235,30 @@ function createAssistantStream(callback) {
       controller.close();
     }
   }
+
   return controller.__internal_getReadable();
 }
-function createAssistantStreamController() {
-  const { resolve, promise } = promiseWithResolvers();
-  let controller;
+
+export function createAssistantStreamController() {
+  const { resolve, promise } = promiseWithResolvers<void>();
+  let controller!: AssistantStreamController;
   const stream = createAssistantStream((c) => {
     controller = c;
-    controller.__internal_subscribeToClose(
-      resolve
+
+    (controller as AssistantStreamControllerImpl).__internal_subscribeToClose(
+      resolve,
     );
+
     return promise;
   });
-  return [stream, controller];
+  return [stream, controller] as const;
 }
-function createAssistantStreamResponse(callback) {
+
+export function createAssistantStreamResponse(
+  callback: (controller: AssistantStreamController) => PromiseLike<void> | void,
+) {
   return AssistantStream.toResponse(
     createAssistantStream(callback),
-    new DataStreamEncoder()
+    new DataStreamEncoder(),
   );
 }
-export {
-  createAssistantStream,
-  createAssistantStreamController,
-  createAssistantStreamResponse
-};
-//# sourceMappingURL=assistant-stream.mjs.map

package/src/core/modules/text.ts

@@ -0,0 +1,64 @@
+import { AssistantStream } from "../AssistantStream";
+import { AssistantStreamChunk } from "../AssistantStreamChunk";
+import { UnderlyingReadable } from "../utils/stream/UnderlyingReadable";
+
+export type TextStreamController = {
+  append(textDelta: string): void;
+  close(): void; // TODO reason? error?
+};
+
+class TextStreamControllerImpl implements TextStreamController {
+  private _controller: ReadableStreamDefaultController<AssistantStreamChunk>;
+  private _isClosed = false;
+
+  constructor(
+    controller: ReadableStreamDefaultController<AssistantStreamChunk>,
+  ) {
+    this._controller = controller;
+  }
+
+  append(textDelta: string) {
+    this._controller.enqueue({
+      type: "text-delta",
+      path: [],
+      textDelta,
+    });
+    return this;
+  }
+
+  close() {
+    if (this._isClosed) return;
+    this._isClosed = true;
+    this._controller.enqueue({
+      type: "part-finish",
+      path: [],
+    });
+    this._controller.close();
+  }
+}
+
+export const createTextStream = (
+  readable: UnderlyingReadable<TextStreamController>,
+): AssistantStream => {
+  return new ReadableStream({
+    start(c) {
+      return readable.start?.(new TextStreamControllerImpl(c));
+    },
+    pull(c) {
+      return readable.pull?.(new TextStreamControllerImpl(c));
+    },
+    cancel(c) {
+      return readable.cancel?.(c);
+    },
+  });
+};
+
+export const createTextStreamController = () => {
+  let controller!: TextStreamController;
+  const stream = createTextStream({
+    start(c) {
+      controller = c;
+    },
+  });
+  return [stream, controller] as const;
+};

package/{dist/core/modules/tool-call.mjs → src/core/modules/tool-call.ts}

@@ -1,12 +1,28 @@
-// src/core/modules/tool-call.ts
-import { createTextStream } from "./text.mjs";
-var ToolCallStreamControllerImpl = class {
-  constructor(_controller) {
-    this._controller = _controller;
+import { AssistantStream } from "../AssistantStream";
+import { AssistantStreamChunk } from "../AssistantStreamChunk";
+import { ToolResponseInit } from "../tool/ToolResponse";
+import { ReadonlyJSONValue } from "../../utils/json/json-value";
+import { UnderlyingReadable } from "../utils/stream/UnderlyingReadable";
+import { createTextStream, TextStreamController } from "./text";
+
+export type ToolCallStreamController = {
+  argsText: TextStreamController;
+
+  setResponse(response: ToolResponseInit<ReadonlyJSONValue>): void;
+  close(): void;
+};
+
+class ToolCallStreamControllerImpl implements ToolCallStreamController {
+  private _isClosed = false;
+
+  private _mergeTask: Promise<void>;
+  constructor(
+    private _controller: ReadableStreamDefaultController<AssistantStreamChunk>,
+  ) {
     const stream = createTextStream({
       start: (c) => {
         this._argsTextController = c;
-      }
+      },
     });
     this._mergeTask = stream.pipeTo(
       new WritableStream({
@@ -15,47 +31,56 @@ var ToolCallStreamControllerImpl = class {
             case "text-delta":
               this._controller.enqueue(chunk);
               break;
+
             case "part-finish":
               this._controller.enqueue({
                 type: "tool-call-args-text-finish",
-                path: []
+                path: [],
               });
               break;
+
             default:
               throw new Error(`Unexpected chunk type: ${chunk.type}`);
           }
-        }
-      })
+        },
+      }),
     );
   }
-  _isClosed = false;
-  _mergeTask;
+
   get argsText() {
     return this._argsTextController;
   }
-  _argsTextController;
-  setResponse(response) {
+
+  private _argsTextController!: TextStreamController;
+
+  setResponse(response: ToolResponseInit<ReadonlyJSONValue>) {
     this._controller.enqueue({
       type: "result",
       path: [],
       artifact: response.artifact,
       result: response.result,
-      isError: response.isError ?? false
+      isError: response.isError ?? false,
     });
   }
+
   async close() {
     if (this._isClosed) return;
+
     this._isClosed = true;
     this._argsTextController.close();
     await this._mergeTask;
+
     this._controller.enqueue({
       type: "part-finish",
-      path: []
+      path: [],
     });
     this._controller.close();
   }
-};
-var createToolCallStream = (readable) => {
+}
+
+export const createToolCallStream = (
+  readable: UnderlyingReadable<ToolCallStreamController>,
+): AssistantStream => {
   return new ReadableStream({
     start(c) {
       return readable.start?.(new ToolCallStreamControllerImpl(c));
@@ -65,20 +90,16 @@ var createToolCallStream = (readable) => {
     },
     cancel(c) {
       return readable.cancel?.(c);
-    }
+    },
   });
 };
-var createToolCallStreamController = () => {
-  let controller;
+
+export const createToolCallStreamController = () => {
+  let controller!: ToolCallStreamController;
   const stream = createToolCallStream({
     start(c) {
       controller = c;
-    }
+    },
   });
-  return [stream, controller];
-};
-export {
-  createToolCallStream,
-  createToolCallStreamController
+  return [stream, controller] as const;
 };
-//# sourceMappingURL=tool-call.mjs.map

package/src/core/serialization/PlainText.ts

@@ -0,0 +1,68 @@
+import { AssistantStreamEncoder } from "../AssistantStream";
+import { AssistantStreamChunk } from "../AssistantStreamChunk";
+import { AssistantTransformStream } from "../utils/stream/AssistantTransformStream";
+import { PipeableTransformStream } from "../utils/stream/PipeableTransformStream";
+
+export class PlainTextEncoder
+  extends PipeableTransformStream<AssistantStreamChunk, Uint8Array>
+  implements AssistantStreamEncoder
+{
+  headers = new Headers({
+    "Content-Type": "text/plain; charset=utf-8",
+    "x-vercel-ai-data-stream": "v1",
+  });
+
+  constructor() {
+    super((readable) => {
+      const transform = new TransformStream<AssistantStreamChunk, string>({
+        transform(chunk, controller) {
+          const type = chunk.type;
+          switch (type) {
+            case "text-delta":
+              controller.enqueue(chunk.textDelta);
+              break;
+
+            default:
+              const unsupportedType:
+                | "part-start"
+                | "part-finish"
+                | "tool-call-args-text-finish"
+                | "data"
+                | "step-start"
+                | "step-finish"
+                | "message-finish"
+                | "annotations"
+                | "tool-call-begin"
+                | "tool-call-delta"
+                | "result"
+                | "error" = type;
+              throw new Error(`unsupported chunk type: ${unsupportedType}`);
+          }
+        },
+      });
+
+      return readable
+        .pipeThrough(transform)
+        .pipeThrough(new TextEncoderStream());
+    });
+  }
+}
+
+export class PlainTextDecoder extends PipeableTransformStream<
+  Uint8Array,
+  AssistantStreamChunk
+> {
+  constructor() {
+    super((readable) => {
+      const transform = new AssistantTransformStream<string>({
+        transform(chunk, controller) {
+          controller.appendText(chunk);
+        },
+      });
+
+      return readable
+        .pipeThrough(new TextDecoderStream())
+        .pipeThrough(transform);
+    });
+  }
+}