assistant-stream 0.1.8 → 0.2.1

package/src/core/accumulators/AssistantMessageStream.ts

@@ -0,0 +1,56 @@
+import { AssistantStream } from "../AssistantStream";
+import { AssistantMessage } from "../utils/types";
+import { AssistantMessageAccumulator } from "./assistant-message-accumulator";
+
+export class AssistantMessageStream {
+  constructor(public readonly readable: ReadableStream<AssistantMessage>) {
+    this.readable = readable;
+  }
+
+  static fromAssistantStream(stream: AssistantStream) {
+    return new AssistantMessageStream(
+      stream.pipeThrough(new AssistantMessageAccumulator()),
+    );
+  }
+
+  async unstable_result(): Promise<AssistantMessage> {
+    let last: AssistantMessage | undefined;
+    for await (const chunk of this) {
+      last = chunk;
+    }
+
+    if (!last) {
+      return {
+        role: "assistant",
+        status: { type: "complete", reason: "unknown" },
+        parts: [],
+        content: [],
+        metadata: {
+          unstable_data: [],
+          unstable_annotations: [],
+          steps: [],
+          custom: {},
+        },
+      };
+    }
+    return last;
+  }
+
+  [Symbol.asyncIterator]() {
+    const reader = this.readable.getReader();
+    return {
+      async next(): Promise<IteratorResult<AssistantMessage, undefined>> {
+        const { done, value } = await reader.read();
+        return done ? { done: true, value: undefined } : { done: false, value };
+      },
+    };
+  }
+
+  tee(): [AssistantMessageStream, AssistantMessageStream] {
+    const [readable1, readable2] = this.readable.tee();
+    return [
+      new AssistantMessageStream(readable1),
+      new AssistantMessageStream(readable2),
+    ];
+  }
+}

package/{dist/core/accumulators/assistant-message-accumulator.mjs → src/core/accumulators/assistant-message-accumulator.ts}

@@ -1,7 +1,18 @@
-// src/core/accumulators/assistant-message-accumulator.ts
-import { generateId } from "../utils/generateId.mjs";
-import { parsePartialJsonObject } from "../../utils/json/parse-partial-json-object.mjs";
-var createInitialMessage = () => ({
+import { AssistantStreamChunk } from "../AssistantStreamChunk";
+import { generateId } from "../utils/generateId";
+import { parsePartialJsonObject } from "../../utils/json/parse-partial-json-object";
+import {
+  AssistantMessage,
+  AssistantMessageStatus,
+  TextPart,
+  ToolCallPart,
+  SourcePart,
+  AssistantMessagePart,
+  ReasoningPart,
+  FilePart,
+} from "../utils/types";
+
+const createInitialMessage = (): AssistantMessage => ({
   role: "assistant",
   status: { type: "running" },
   parts: [],
@@ -12,126 +23,158 @@ var createInitialMessage = () => ({
     unstable_data: [],
     unstable_annotations: [],
     steps: [],
-    custom: {}
-  }
+    custom: {},
+  },
 });
-var updatePartForPath = (message, chunk, updater) => {
+
+const updatePartForPath = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk,
+  updater: (part: AssistantMessagePart) => AssistantMessagePart,
+): AssistantMessage => {
   if (message.parts.length === 0) {
     throw new Error("No parts available to update.");
   }
+
   if (chunk.path.length !== 1)
     throw new Error("Nested paths are not supported yet.");
-  const partIndex = chunk.path[0];
-  const updatedPart = updater(message.parts[partIndex]);
+
+  const partIndex = chunk.path[0]!;
+  const updatedPart = updater(message.parts[partIndex]!);
   return {
     ...message,
     parts: [
       ...message.parts.slice(0, partIndex),
       updatedPart,
-      ...message.parts.slice(partIndex + 1)
+      ...message.parts.slice(partIndex + 1),
     ],
     get content() {
       return this.parts;
-    }
+    },
   };
 };
-var handlePartStart = (message, chunk) => {
+
+const handlePartStart = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & { readonly type: "part-start" },
+): AssistantMessage => {
   const partInit = chunk.part;
   if (partInit.type === "text" || partInit.type === "reasoning") {
-    const newTextPart = {
+    const newTextPart: TextPart | ReasoningPart = {
       type: partInit.type,
       text: "",
-      status: { type: "running" }
+      status: { type: "running" },
     };
     return {
       ...message,
       parts: [...message.parts, newTextPart],
       get content() {
         return this.parts;
-      }
+      },
     };
   } else if (partInit.type === "tool-call") {
-    const newToolCallPart = {
+    const newToolCallPart: ToolCallPart = {
       type: "tool-call",
       state: "partial-call",
       status: { type: "running", isArgsComplete: false },
       toolCallId: partInit.toolCallId,
       toolName: partInit.toolName,
       argsText: "",
-      args: {}
+      args: {},
     };
     return {
       ...message,
       parts: [...message.parts, newToolCallPart],
       get content() {
         return this.parts;
-      }
+      },
     };
   } else if (partInit.type === "source") {
-    const newSourcePart = {
+    const newSourcePart: SourcePart = {
       type: "source",
       sourceType: partInit.sourceType,
       id: partInit.id,
       url: partInit.url,
-      ...partInit.title ? { title: partInit.title } : void 0
+      ...(partInit.title ? { title: partInit.title } : undefined),
     };
     return {
       ...message,
       parts: [...message.parts, newSourcePart],
       get content() {
         return this.parts;
-      }
+      },
     };
   } else if (partInit.type === "file") {
-    const newFilePart = {
+    const newFilePart: FilePart = {
       type: "file",
       mimeType: partInit.mimeType,
-      data: partInit.data
+      data: partInit.data,
     };
     return {
       ...message,
       parts: [...message.parts, newFilePart],
       get content() {
         return this.parts;
-      }
+      },
     };
   } else {
     throw new Error(`Unsupported part type: ${partInit.type}`);
   }
 };
-var handleToolCallArgsTextFinish = (message, chunk) => {
+
+const handleToolCallArgsTextFinish = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & {
+    readonly type: "tool-call-args-text-finish";
+  },
+): AssistantMessage => {
   return updatePartForPath(message, chunk, (part) => {
     if (part.type !== "tool-call") {
       throw new Error("Last is not a tool call");
     }
     return {
       ...part,
-      state: "call"
+      state: "call",
     };
   });
 };
-var handlePartFinish = (message, chunk) => {
+
+const handlePartFinish = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & { readonly type: "part-finish" },
+): AssistantMessage => {
   return updatePartForPath(message, chunk, (part) => ({
     ...part,
-    status: { type: "complete", reason: "unknown" }
+    status: { type: "complete", reason: "unknown" },
   }));
 };
-var handleTextDelta = (message, chunk) => {
+
+const handleTextDelta = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & { type: "text-delta" },
+): AssistantMessage => {
   return updatePartForPath(message, chunk, (part) => {
     if (part.type === "text" || part.type === "reasoning") {
       return { ...part, text: part.text + chunk.textDelta };
     } else if (part.type === "tool-call") {
       const newArgsText = part.argsText + chunk.textDelta;
+
+      // Fall back to existing args if parsing fails
       const newArgs = parsePartialJsonObject(newArgsText) ?? part.args;
+
       return { ...part, argsText: newArgsText, args: newArgs };
     } else {
       throw new Error(
-        "text-delta received but part is neither text nor tool-call"
+        "text-delta received but part is neither text nor tool-call",
       );
     }
   });
 };
-var handleResult = (message, chunk) => {
+
+const handleResult = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & { type: "result" },
+): AssistantMessage => {
   return updatePartForPath(message, chunk, (part) => {
     if (part.type === "tool-call") {
       return {
@@ -140,105 +183,147 @@ var handleResult = (message, chunk) => {
         artifact: chunk.artifact,
         result: chunk.result,
         isError: chunk.isError ?? false,
-        status: { type: "complete", reason: "stop" }
+        status: { type: "complete", reason: "stop" },
       };
     } else {
       throw new Error("Result chunk received but part is not a tool-call");
     }
   });
 };
-var handleMessageFinish = (message, chunk) => {
+
+const handleMessageFinish = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & { type: "message-finish" },
+): AssistantMessage => {
   const newStatus = getStatus(chunk);
   return { ...message, status: newStatus };
 };
-var getStatus = (chunk) => {
+
+const getStatus = (
+  chunk:
+    | (AssistantStreamChunk & { type: "message-finish" })
+    | (AssistantStreamChunk & { type: "step-finish" }),
+): AssistantMessageStatus => {
   if (chunk.finishReason === "tool-calls") {
     return {
       type: "requires-action",
-      reason: "tool-calls"
+      reason: "tool-calls",
     };
-  } else if (chunk.finishReason === "stop" || chunk.finishReason === "unknown") {
+  } else if (
+    chunk.finishReason === "stop" ||
+    chunk.finishReason === "unknown"
+  ) {
     return {
       type: "complete",
-      reason: chunk.finishReason
+      reason: chunk.finishReason,
     };
   } else {
     return {
       type: "incomplete",
-      reason: chunk.finishReason
+      reason: chunk.finishReason,
     };
   }
 };
-var handleAnnotations = (message, chunk) => {
+
+const handleAnnotations = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & { type: "annotations" },
+): AssistantMessage => {
   return {
     ...message,
     metadata: {
       ...message.metadata,
       unstable_annotations: [
         ...message.metadata.unstable_annotations,
-        ...chunk.annotations
-      ]
-    }
+        ...chunk.annotations,
+      ],
+    },
   };
 };
-var handleData = (message, chunk) => {
+
+const handleData = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & { type: "data" },
+): AssistantMessage => {
   return {
     ...message,
     metadata: {
       ...message.metadata,
-      unstable_data: [...message.metadata.unstable_data, ...chunk.data]
-    }
+      unstable_data: [...message.metadata.unstable_data, ...chunk.data],
+    },
   };
 };
-var handleStepStart = (message, chunk) => {
+
+const handleStepStart = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & { type: "step-start" },
+): AssistantMessage => {
   return {
     ...message,
     metadata: {
       ...message.metadata,
       steps: [
         ...message.metadata.steps,
-        { state: "started", messageId: chunk.messageId }
-      ]
-    }
+        { state: "started", messageId: chunk.messageId },
+      ],
+    },
   };
 };
-var handleStepFinish = (message, chunk) => {
+
+const handleStepFinish = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & { type: "step-finish" },
+): AssistantMessage => {
   const steps = message.metadata.steps.slice();
   const lastIndex = steps.length - 1;
+
+  // Check if the previous step is a step-start (has state "started")
   if (steps.length > 0 && steps[lastIndex]?.state === "started") {
     steps[lastIndex] = {
       ...steps[lastIndex],
       state: "finished",
       finishReason: chunk.finishReason,
       usage: chunk.usage,
-      isContinued: chunk.isContinued
+      isContinued: chunk.isContinued,
     };
   } else {
+    // If no previous step-start exists, append a finished step
     steps.push({
       state: "finished",
       messageId: generateId(),
       finishReason: chunk.finishReason,
       usage: chunk.usage,
-      isContinued: chunk.isContinued
+      isContinued: chunk.isContinued,
     });
   }
+
   return {
     ...message,
     metadata: {
       ...message.metadata,
-      steps
-    }
+      steps,
+    },
   };
 };
-var handleErrorChunk = (message, chunk) => {
+
+const handleErrorChunk = (
+  message: AssistantMessage,
+  chunk: AssistantStreamChunk & { type: "error" },
+): AssistantMessage => {
   return {
     ...message,
-    status: { type: "incomplete", reason: "error", error: chunk.error }
+    status: { type: "incomplete", reason: "error", error: chunk.error },
   };
 };
-var AssistantMessageAccumulator = class extends TransformStream {
+
+export class AssistantMessageAccumulator extends TransformStream<
+  AssistantStreamChunk,
+  AssistantMessage
+> {
   constructor({
-    initialMessage
+    initialMessage,
+  }: {
+    initialMessage?: AssistantMessage;
   } = {}) {
     let message = initialMessage ?? createInitialMessage();
     super({
@@ -248,12 +333,15 @@ var AssistantMessageAccumulator = class extends TransformStream {
           case "part-start":
             message = handlePartStart(message, chunk);
             break;
+
           case "tool-call-args-text-finish":
             message = handleToolCallArgsTextFinish(message, chunk);
             break;
+
           case "part-finish":
             message = handlePartFinish(message, chunk);
             break;
+
           case "text-delta":
             message = handleTextDelta(message, chunk);
             break;
@@ -279,7 +367,7 @@ var AssistantMessageAccumulator = class extends TransformStream {
             message = handleErrorChunk(message, chunk);
             break;
           default: {
-            const unhandledType = type;
+            const unhandledType: never = type;
             throw new Error(`Unsupported chunk type: ${unhandledType}`);
           }
         }
@@ -287,6 +375,7 @@ var AssistantMessageAccumulator = class extends TransformStream {
       },
       flush(controller) {
         if (message.status?.type === "running") {
+          // TODO this test isn't sound
           const requiresAction = message.parts?.at(-1)?.type === "tool-call";
           message = handleMessageFinish(message, {
             type: "message-finish",
@@ -294,16 +383,12 @@ var AssistantMessageAccumulator = class extends TransformStream {
             finishReason: requiresAction ? "tool-calls" : "unknown",
             usage: {
               promptTokens: 0,
-              completionTokens: 0
-            }
+              completionTokens: 0,
+            },
           });
           controller.enqueue(message);
         }
-      }
+      },
     });
   }
-};
-export {
-  AssistantMessageAccumulator
-};
-//# sourceMappingURL=assistant-message-accumulator.mjs.map
+}

package/src/core/index.ts

@@ -0,0 +1,17 @@
+export {
+  createAssistantStream,
+  createAssistantStreamResponse,
+} from "./modules/assistant-stream";
+export { AssistantMessageAccumulator } from "./accumulators/assistant-message-accumulator";
+export { AssistantStream } from "./AssistantStream";
+export type { AssistantStreamController } from "./modules/assistant-stream";
+export type { AssistantStreamChunk } from "./AssistantStreamChunk";
+export {
+  DataStreamDecoder,
+  DataStreamEncoder,
+} from "./serialization/data-stream/DataStream";
+export { PlainTextDecoder, PlainTextEncoder } from "./serialization/PlainText";
+export { AssistantMessageStream } from "./accumulators/AssistantMessageStream";
+export type { AssistantMessage } from "./utils/types";
+
+export * from "./tool";