assistant-stream 0.1.8 → 0.2.1

package/{dist/core/serialization/data-stream/DataStream.mjs → src/core/serialization/data-stream/DataStream.ts}

@@ -1,40 +1,52 @@
-// src/core/serialization/data-stream/DataStream.ts
-import { AssistantTransformStream } from "../../utils/stream/AssistantTransformStream.mjs";
-import { PipeableTransformStream } from "../../utils/stream/PipeableTransformStream.mjs";
-import { DataStreamStreamChunkType } from "./chunk-types.mjs";
-import { LineDecoderStream } from "../../utils/stream/LineDecoderStream.mjs";
+import { AssistantStreamChunk } from "../../AssistantStreamChunk";
+import { ToolCallStreamController } from "../../modules/tool-call";
+import { AssistantTransformStream } from "../../utils/stream/AssistantTransformStream";
+import { PipeableTransformStream } from "../../utils/stream/PipeableTransformStream";
+import { DataStreamChunk, DataStreamStreamChunkType } from "./chunk-types";
+import { LineDecoderStream } from "../../utils/stream/LineDecoderStream";
 import {
   DataStreamChunkDecoder,
-  DataStreamChunkEncoder
-} from "./serialization.mjs";
+  DataStreamChunkEncoder,
+} from "./serialization";
 import {
-  AssistantMetaTransformStream
-} from "../../utils/stream/AssistantMetaTransformStream.mjs";
-var DataStreamEncoder = class extends PipeableTransformStream {
+  AssistantMetaStreamChunk,
+  AssistantMetaTransformStream,
+} from "../../utils/stream/AssistantMetaTransformStream";
+import { TextStreamController } from "../../modules/text";
+import { AssistantStreamEncoder } from "../../AssistantStream";
+
+export class DataStreamEncoder
+  extends PipeableTransformStream<AssistantStreamChunk, Uint8Array>
+  implements AssistantStreamEncoder
+{
   headers = new Headers({
     "Content-Type": "text/plain; charset=utf-8",
-    "x-vercel-ai-data-stream": "v1"
+    "x-vercel-ai-data-stream": "v1",
   });
+
   constructor() {
     super((readable) => {
-      const transform = new TransformStream({
+      const transform = new TransformStream<
+        AssistantMetaStreamChunk,
+        DataStreamChunk
+      >({
         transform(chunk, controller) {
           const type = chunk.type;
           switch (type) {
             case "part-start": {
               const part = chunk.part;
               if (part.type === "tool-call") {
-                const { type: type2, ...value } = part;
+                const { type, ...value } = part;
                 controller.enqueue({
                   type: DataStreamStreamChunkType.StartToolCall,
-                  value
+                  value,
                 });
               }
               if (part.type === "source") {
-                const { type: type2, ...value } = part;
+                const { type, ...value } = part;
                 controller.enqueue({
                   type: DataStreamStreamChunkType.Source,
-                  value
+                  value,
                 });
               }
               break;
@@ -45,14 +57,14 @@ var DataStreamEncoder = class extends PipeableTransformStream {
                 case "text": {
                   controller.enqueue({
                     type: DataStreamStreamChunkType.TextDelta,
-                    value: chunk.textDelta
+                    value: chunk.textDelta,
                   });
                   break;
                 }
                 case "reasoning": {
                   controller.enqueue({
                     type: DataStreamStreamChunkType.ReasoningDelta,
-                    value: chunk.textDelta
+                    value: chunk.textDelta,
                   });
                   break;
                 }
@@ -61,23 +73,24 @@ var DataStreamEncoder = class extends PipeableTransformStream {
                     type: DataStreamStreamChunkType.ToolCallArgsTextDelta,
                     value: {
                       toolCallId: part.toolCallId,
-                      argsTextDelta: chunk.textDelta
-                    }
+                      argsTextDelta: chunk.textDelta,
+                    },
                   });
                   break;
                 }
                 default:
                   throw new Error(
-                    `Unsupported part type for text-delta: ${part.type}`
+                    `Unsupported part type for text-delta: ${part.type}`,
                   );
               }
               break;
             }
             case "result": {
+              // Only tool-call parts can have results.
               const part = chunk.meta;
               if (part.type !== "tool-call") {
                 throw new Error(
-                  `Result chunk on non-tool-call part not supported: ${part.type}`
+                  `Result chunk on non-tool-call part not supported: ${part.type}`,
                 );
               }
               controller.enqueue({
@@ -86,73 +99,81 @@ var DataStreamEncoder = class extends PipeableTransformStream {
                   toolCallId: part.toolCallId,
                   result: chunk.result,
                   artifact: chunk.artifact,
-                  ...chunk.isError ? { isError: chunk.isError } : {}
-                }
+                  ...(chunk.isError ? { isError: chunk.isError } : {}),
+                },
               });
               break;
             }
             case "step-start": {
-              const { type: type2, ...value } = chunk;
+              const { type, ...value } = chunk;
               controller.enqueue({
                 type: DataStreamStreamChunkType.StartStep,
-                value
+                value,
               });
               break;
             }
             case "step-finish": {
-              const { type: type2, ...value } = chunk;
+              const { type, ...value } = chunk;
               controller.enqueue({
                 type: DataStreamStreamChunkType.FinishStep,
-                value
+                value,
               });
               break;
             }
             case "message-finish": {
-              const { type: type2, ...value } = chunk;
+              const { type, ...value } = chunk;
               controller.enqueue({
                 type: DataStreamStreamChunkType.FinishMessage,
-                value
+                value,
               });
               break;
             }
             case "error": {
               controller.enqueue({
                 type: DataStreamStreamChunkType.Error,
-                value: chunk.error
+                value: chunk.error,
               });
               break;
             }
             case "annotations": {
               controller.enqueue({
                 type: DataStreamStreamChunkType.Annotation,
-                value: chunk.annotations
+                value: chunk.annotations,
               });
               break;
             }
             case "data": {
               controller.enqueue({
                 type: DataStreamStreamChunkType.Data,
-                value: chunk.data
+                value: chunk.data,
               });
               break;
             }
+
             // TODO ignore for now
             // in the future, we should create a handler that waits for text parts to finish before continuing
             case "tool-call-args-text-finish":
             case "part-finish":
               break;
+
             default: {
-              const exhaustiveCheck = type;
+              const exhaustiveCheck: never = type;
               throw new Error(`Unsupported chunk type: ${exhaustiveCheck}`);
             }
           }
-        }
+        },
       });
-      return readable.pipeThrough(new AssistantMetaTransformStream()).pipeThrough(transform).pipeThrough(new DataStreamChunkEncoder()).pipeThrough(new TextEncoderStream());
+
+      return readable
+        .pipeThrough(new AssistantMetaTransformStream())
+        .pipeThrough(transform)
+        .pipeThrough(new DataStreamChunkEncoder())
+        .pipeThrough(new TextEncoderStream());
     });
   }
-};
-var TOOL_CALL_ARGS_CLOSING_CHUNKS = [
+}
+
+const TOOL_CALL_ARGS_CLOSING_CHUNKS = [
   DataStreamStreamChunkType.StartToolCall,
   DataStreamStreamChunkType.ToolCall,
   DataStreamStreamChunkType.TextDelta,
@@ -160,63 +181,76 @@ var TOOL_CALL_ARGS_CLOSING_CHUNKS = [
   DataStreamStreamChunkType.Source,
   DataStreamStreamChunkType.Error,
   DataStreamStreamChunkType.FinishStep,
-  DataStreamStreamChunkType.FinishMessage
+  DataStreamStreamChunkType.FinishMessage,
 ];
-var DataStreamDecoder = class extends PipeableTransformStream {
+
+export class DataStreamDecoder extends PipeableTransformStream<
+  Uint8Array,
+  AssistantStreamChunk
+> {
   constructor() {
     super((readable) => {
-      const toolCallControllers = /* @__PURE__ */ new Map();
-      let activeToolCallArgsText;
-      const transform = new AssistantTransformStream({
+      const toolCallControllers = new Map<string, ToolCallStreamController>();
+      let activeToolCallArgsText: TextStreamController | undefined;
+      const transform = new AssistantTransformStream<DataStreamChunk>({
         transform(chunk, controller) {
           const { type, value } = chunk;
+
           if (TOOL_CALL_ARGS_CLOSING_CHUNKS.includes(type)) {
             activeToolCallArgsText?.close();
-            activeToolCallArgsText = void 0;
+            activeToolCallArgsText = undefined;
           }
+
           switch (type) {
             case DataStreamStreamChunkType.ReasoningDelta:
               controller.appendReasoning(value);
               break;
+
             case DataStreamStreamChunkType.TextDelta:
               controller.appendText(value);
               break;
+
             case DataStreamStreamChunkType.StartToolCall: {
               const { toolCallId, toolName } = value;
               const toolCallController = controller.addToolCallPart({
                 toolCallId,
-                toolName
+                toolName,
               });
               toolCallControllers.set(toolCallId, toolCallController);
+
               activeToolCallArgsText = toolCallController.argsText;
               break;
             }
+
             case DataStreamStreamChunkType.ToolCallArgsTextDelta: {
               const { toolCallId, argsTextDelta } = value;
               const toolCallController = toolCallControllers.get(toolCallId);
               if (!toolCallController)
                 throw new Error(
-                  "Encountered tool call with unknown id: " + toolCallId
+                  "Encountered tool call with unknown id: " + toolCallId,
                 );
               toolCallController.argsText.append(argsTextDelta);
               break;
             }
+
             case DataStreamStreamChunkType.ToolCallResult: {
               const { toolCallId, artifact, result, isError } = value;
               const toolCallController = toolCallControllers.get(toolCallId);
               if (!toolCallController)
                 throw new Error(
-                  "Encountered tool call result with unknown id: " + toolCallId
+                  "Encountered tool call result with unknown id: " + toolCallId,
                 );
               toolCallController.setResponse({
                 artifact,
                 result,
-                isError
+                isError,
               });
               break;
             }
+
             case DataStreamStreamChunkType.ToolCall: {
               const { toolCallId, toolName, args } = value;
+
               let toolCallController = toolCallControllers.get(toolCallId);
               if (toolCallController) {
                 toolCallController.argsText.close();
@@ -224,88 +258,98 @@ var DataStreamDecoder = class extends PipeableTransformStream {
                 toolCallController = controller.addToolCallPart({
                   toolCallId,
                   toolName,
-                  args
+                  args,
                 });
                 toolCallControllers.set(toolCallId, toolCallController);
               }
               break;
             }
+
             case DataStreamStreamChunkType.FinishMessage:
               controller.enqueue({
                 type: "message-finish",
                 path: [],
-                ...value
+                ...value,
               });
               break;
+
             case DataStreamStreamChunkType.StartStep:
               controller.enqueue({
                 type: "step-start",
                 path: [],
-                ...value
+                ...value,
               });
               break;
+
             case DataStreamStreamChunkType.FinishStep:
               controller.enqueue({
                 type: "step-finish",
                 path: [],
-                ...value
+                ...value,
               });
               break;
             case DataStreamStreamChunkType.Data:
               controller.enqueue({
                 type: "data",
                 path: [],
-                data: value
+                data: value,
               });
               break;
+
             case DataStreamStreamChunkType.Annotation:
               controller.enqueue({
                 type: "annotations",
                 path: [],
-                annotations: value
+                annotations: value,
               });
               break;
+
             case DataStreamStreamChunkType.Source:
               controller.appendSource({
                 type: "source",
-                ...value
+                ...value,
               });
               break;
+
             case DataStreamStreamChunkType.Error:
               controller.enqueue({
                 type: "error",
                 path: [],
-                error: value
+                error: value,
               });
               break;
+
             case DataStreamStreamChunkType.File:
               controller.appendFile({
                 type: "file",
-                ...value
+                ...value,
               });
               break;
+
             case DataStreamStreamChunkType.ReasoningSignature:
             case DataStreamStreamChunkType.RedactedReasoning:
+              // ignore these for now
               break;
+
             default: {
-              const exhaustiveCheck = type;
+              const exhaustiveCheck: never = type;
               throw new Error(`unsupported chunk type: ${exhaustiveCheck}`);
             }
           }
         },
         flush() {
           activeToolCallArgsText?.close();
-          activeToolCallArgsText = void 0;
+          activeToolCallArgsText = undefined;
           toolCallControllers.forEach((controller) => controller.close());
           toolCallControllers.clear();
-        }
+        },
       });
-      return readable.pipeThrough(new TextDecoderStream()).pipeThrough(new LineDecoderStream()).pipeThrough(new DataStreamChunkDecoder()).pipeThrough(transform);
+
+      return readable
+        .pipeThrough(new TextDecoderStream())
+        .pipeThrough(new LineDecoderStream())
+        .pipeThrough(new DataStreamChunkDecoder())
+        .pipeThrough(transform);
     });
   }
-};
-export {
-  DataStreamDecoder,
-  DataStreamEncoder
-};
-//# sourceMappingURL=DataStream.mjs.map
+}

package/src/core/serialization/data-stream/chunk-types.ts

@@ -0,0 +1,93 @@
+import {
+  ReadonlyJSONObject,
+  ReadonlyJSONValue,
+} from "../../../utils/json/json-value";
+
+export type DataStreamChunk = {
+  [K in DataStreamStreamChunkType]: {
+    type: K;
+    value: DataStreamStreamChunkValue[K];
+  };
+}[DataStreamStreamChunkType];
+
+type LanguageModelV1FinishReason =
+  | "stop"
+  | "length"
+  | "content-filter"
+  | "tool-calls"
+  | "error"
+  | "other"
+  | "unknown";
+
+type LanguageModelV1Usage = {
+  promptTokens: number;
+  completionTokens: number;
+};
+
+export enum DataStreamStreamChunkType {
+  TextDelta = "0",
+  Data = "2",
+  Error = "3",
+  Annotation = "8",
+  ToolCall = "9",
+  ToolCallResult = "a",
+  StartToolCall = "b",
+  ToolCallArgsTextDelta = "c",
+  FinishMessage = "d",
+  FinishStep = "e",
+  StartStep = "f",
+  ReasoningDelta = "g",
+  Source = "h",
+  RedactedReasoning = "i",
+  ReasoningSignature = "j",
+  File = "k",
+}
+type DataStreamStreamChunkValue = {
+  [DataStreamStreamChunkType.TextDelta]: string;
+  [DataStreamStreamChunkType.Data]: ReadonlyJSONValue[];
+  [DataStreamStreamChunkType.Annotation]: ReadonlyJSONValue[];
+  [DataStreamStreamChunkType.ToolCall]: {
+    toolCallId: string;
+    toolName: string;
+    args: ReadonlyJSONObject;
+  };
+  [DataStreamStreamChunkType.StartToolCall]: {
+    toolCallId: string;
+    toolName: string;
+  };
+  [DataStreamStreamChunkType.ToolCallArgsTextDelta]: {
+    toolCallId: string;
+    argsTextDelta: string;
+  };
+  [DataStreamStreamChunkType.ToolCallResult]: {
+    toolCallId: string;
+    result: ReadonlyJSONValue;
+
+    // aui-extensions
+    artifact?: ReadonlyJSONValue | undefined;
+    isError?: boolean;
+  };
+  [DataStreamStreamChunkType.Error]: string;
+  [DataStreamStreamChunkType.FinishStep]: {
+    finishReason: LanguageModelV1FinishReason;
+    usage: LanguageModelV1Usage;
+    isContinued: boolean;
+  };
+  [DataStreamStreamChunkType.FinishMessage]: {
+    finishReason: LanguageModelV1FinishReason;
+    usage: LanguageModelV1Usage;
+  };
+  [DataStreamStreamChunkType.StartStep]: {
+    messageId: string;
+  };
+  [DataStreamStreamChunkType.ReasoningDelta]: string;
+  [DataStreamStreamChunkType.Source]: {
+    sourceType: "url";
+    id: string;
+    url: string;
+    title?: string;
+  };
+  [DataStreamStreamChunkType.RedactedReasoning]: { data: string };
+  [DataStreamStreamChunkType.ReasoningSignature]: { signature: string };
+  [DataStreamStreamChunkType.File]: { data: string; mimeType: string };
+};

package/src/core/serialization/data-stream/serialization.ts

@@ -0,0 +1,32 @@
+import { DataStreamChunk, DataStreamStreamChunkType } from "./chunk-types";
+
+export class DataStreamChunkEncoder extends TransformStream<
+  DataStreamChunk,
+  string
+> {
+  constructor() {
+    super({
+      transform: (chunk, controller) => {
+        controller.enqueue(`${chunk.type}:${JSON.stringify(chunk.value)}\n`);
+      },
+    });
+  }
+}
+
+export class DataStreamChunkDecoder extends TransformStream<
+  string,
+  DataStreamChunk
+> {
+  constructor() {
+    super({
+      transform: (chunk, controller) => {
+        const index = chunk.indexOf(":");
+        if (index === -1) throw new Error("Invalid stream part");
+        controller.enqueue({
+          type: chunk.slice(0, index) as DataStreamStreamChunkType,
+          value: JSON.parse(chunk.slice(index + 1)),
+        });
+      },
+    });
+  }
+}