assistant-stream 0.1.8 → 0.2.1

package/src/core/utils/stream/path-utils.ts

@@ -0,0 +1,71 @@
+import { AssistantStreamChunk } from "../../AssistantStreamChunk";
+import { Counter } from "../Counter";
+
+export class PathAppendEncoder extends TransformStream<
+  AssistantStreamChunk,
+  AssistantStreamChunk
+> {
+  constructor(idx: number) {
+    super({
+      transform(chunk, controller) {
+        controller.enqueue({
+          ...chunk,
+          path: [idx, ...chunk.path],
+        });
+      },
+    });
+  }
+}
+
+export class PathAppendDecoder extends TransformStream<
+  AssistantStreamChunk,
+  AssistantStreamChunk
+> {
+  constructor(idx: number) {
+    super({
+      transform(chunk, controller) {
+        const {
+          path: [idx2, ...path],
+        } = chunk;
+
+        if (idx !== idx2)
+          throw new Error(`Path mismatch: expected ${idx}, got ${idx2}`);
+
+        controller.enqueue({
+          ...chunk,
+          path,
+        });
+      },
+    });
+  }
+}
+
+export class PathMergeEncoder extends TransformStream<
+  AssistantStreamChunk,
+  AssistantStreamChunk
+> {
+  constructor(counter: Counter) {
+    const innerCounter = new Counter();
+    const mapping = new Map<number, number>();
+    super({
+      transform(chunk, controller) {
+        if (chunk.type === "part-start" && chunk.path.length === 0) {
+          mapping.set(innerCounter.up(), counter.up());
+        }
+
+        const [idx, ...path] = chunk.path;
+        if (idx === undefined) {
+          controller.enqueue(chunk);
+          return;
+        }
+        const mappedIdx = mapping.get(idx);
+        if (mappedIdx === undefined) throw new Error("Path not found");
+
+        controller.enqueue({
+          ...chunk,
+          path: [mappedIdx, ...path],
+        });
+      },
+    });
+  }
+}

package/src/core/utils/types.ts

@@ -0,0 +1,150 @@
+import {
+  ReadonlyJSONObject,
+  ReadonlyJSONValue,
+} from "../../utils/json/json-value";
+
+type TextStatus =
+  | {
+      type: "running";
+    }
+  | {
+      type: "complete";
+      reason: "stop" | "unknown";
+    }
+  | {
+      type: "incomplete";
+      reason: "cancelled" | "length" | "content-filter" | "other";
+    };
+
+// export type StepStartPart = {
+//   type: "step-start";
+// };
+
+export type TextPart = {
+  type: "text";
+  text: string;
+  status: TextStatus;
+};
+
+export type ReasoningPart = {
+  type: "reasoning";
+  text: string;
+  status: TextStatus;
+};
+
+type ToolCallStatus =
+  | {
+      type: "running";
+      isArgsComplete: boolean;
+    }
+  | {
+      type: "requires-action";
+      reason: "tool-call-result";
+    }
+  | {
+      type: "complete";
+      reason: "stop" | "unknown";
+    }
+  | {
+      type: "incomplete";
+      reason: "cancelled" | "length" | "content-filter" | "other";
+    };
+
+export type ToolCallPart = {
+  type: "tool-call";
+  state: "partial-call" | "call" | "result";
+  status: ToolCallStatus;
+  toolCallId: string;
+  toolName: string;
+  argsText: string;
+  args: ReadonlyJSONObject;
+  artifact?: unknown;
+  result?: ReadonlyJSONValue;
+  isError?: boolean;
+};
+
+export type SourcePart = {
+  type: "source";
+  sourceType: "url";
+  id: string;
+  url: string;
+  title?: string;
+};
+
+export type FilePart = {
+  type: "file";
+  data: string;
+  mimeType: string;
+};
+
+export type AssistantMessagePart =
+  | TextPart
+  | ReasoningPart
+  | ToolCallPart
+  | SourcePart
+  | FilePart;
+
+type AssistantMessageStepUsage = {
+  promptTokens: number;
+  completionTokens: number;
+};
+
+type AssistantMessageStepMetadata =
+  | {
+      state: "started";
+      messageId: string;
+    }
+  | {
+      state: "finished";
+      messageId: string;
+      finishReason:
+        | "stop"
+        | "length"
+        | "content-filter"
+        | "tool-calls"
+        | "error"
+        | "other"
+        | "unknown";
+      usage?: AssistantMessageStepUsage;
+      isContinued: boolean;
+    };
+
+export type AssistantMessageStatus =
+  | {
+      type: "running";
+    }
+  | {
+      type: "requires-action";
+      reason: "tool-calls";
+    }
+  | {
+      type: "complete";
+      reason: "stop" | "unknown";
+    }
+  | {
+      type: "incomplete";
+      reason:
+        | "cancelled"
+        | "tool-calls"
+        | "length"
+        | "content-filter"
+        | "other"
+        | "error";
+      error?: ReadonlyJSONValue;
+    };
+
+export type AssistantMessage = {
+  role: "assistant";
+  status: AssistantMessageStatus;
+  parts: AssistantMessagePart[];
+  /**
+   * @deprecated Use `parts` instead.
+   */
+  content: AssistantMessagePart[];
+  metadata: {
+    unstable_data: ReadonlyJSONValue[];
+    unstable_annotations: ReadonlyJSONValue[];
+    steps: AssistantMessageStepMetadata[];
+    custom: Record<string, unknown>;
+  };
+};

package/src/core/utils/withPromiseOrValue.ts

@@ -0,0 +1,20 @@
+export function withPromiseOrValue<T>(
+  callback: () => T | PromiseLike<T>,
+  thenHandler: (value: T) => PromiseLike<void> | void,
+  catchHandler: (error: unknown) => PromiseLike<void> | void,
+): PromiseLike<void> | void {
+  try {
+    const promiseOrValue = callback();
+    if (
+      typeof promiseOrValue === "object" &&
+      promiseOrValue !== null &&
+      "then" in promiseOrValue
+    ) {
+      return promiseOrValue.then(thenHandler, catchHandler);
+    } else {
+      thenHandler(promiseOrValue);
+    }
+  } catch (e) {
+    catchHandler(e);
+  }
+}

package/src/index.ts

@@ -0,0 +1 @@
+export * from "./core";

package/src/utils/AsyncIterableStream.ts

@@ -0,0 +1,24 @@
+export type AsyncIterableStream<T> = AsyncIterable<T> & ReadableStream<T>;
+
+async function* streamGeneratorPolyfill<T>(
+  this: ReadableStream<T>,
+): AsyncIterator<T, undefined, unknown> {
+  const reader = this.getReader();
+  try {
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      yield value;
+    }
+  } finally {
+    reader.releaseLock();
+  }
+}
+
+export function asAsyncIterableStream<T>(
+  source: ReadableStream<T>,
+): AsyncIterableStream<T> {
+  (source as AsyncIterableStream<T>)[Symbol.asyncIterator] ??=
+    streamGeneratorPolyfill;
+  return source as AsyncIterableStream<T>;
+}

package/{dist/utils/json/fix-json.mjs → src/utils/json/fix-json.ts}

@@ -1,17 +1,81 @@
-// src/utils/json/fix-json.ts
-function fixJson(input) {
-  const stack = ["ROOT"];
+// LICENSE for this file only
+
+// MIT License
+
+// Copyright (c) 2025 AgentbaseAI Inc.
+// Copyright (c) 2023 Lars Grammel
+
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+type State =
+  | "ROOT"
+  | "FINISH"
+  | "INSIDE_STRING"
+  | "INSIDE_STRING_ESCAPE"
+  | "INSIDE_LITERAL"
+  | "INSIDE_NUMBER"
+  | "INSIDE_OBJECT_START"
+  | "INSIDE_OBJECT_KEY"
+  | "INSIDE_OBJECT_AFTER_KEY"
+  | "INSIDE_OBJECT_BEFORE_VALUE"
+  | "INSIDE_OBJECT_AFTER_VALUE"
+  | "INSIDE_OBJECT_AFTER_COMMA"
+  | "INSIDE_ARRAY_START"
+  | "INSIDE_ARRAY_AFTER_VALUE"
+  | "INSIDE_ARRAY_AFTER_COMMA";
+
+// Implemented as a scanner with additional fixing
+// that performs a single linear time scan pass over the partial JSON.
+//
+// The states should ideally match relevant states from the JSON spec:
+// https://www.json.org/json-en.html
+//
+// Please note that invalid JSON is not considered/covered, because it
+// is assumed that the resulting JSON will be processed by a standard
+// JSON parser that will detect any invalid JSON.
+
+// Returns a tuple of [fixedJson, partialPath]
+// partialPath is an array of object/array keys that represent
+// the currently partial values. An object is considered partial
+// if through appending extra characters to the JSON string, its
+// value could change.
+
+// Example input: '{"foo":[{"a":f'
+// Example output: ['{"foo":[{"a":false}]}', ['foo', '0']]
+// Example input: '{"foo":'
+// Example output: ['{}', []]
+
+export function fixJson(input: string): [string, string[]] {
+  const stack: State[] = ["ROOT"];
   let lastValidIndex = -1;
-  let literalStart = null;
-  const path = [];
-  let currentKey;
-  function pushCurrentKeyToPath() {
-    if (currentKey !== void 0) {
+  let literalStart: number | null = null;
+  const path: string[] = [];
+  let currentKey: string | undefined;
+
+  function pushCurrentKeyToPath(): void {
+    if (currentKey !== undefined) {
       path.push(JSON.parse('"' + currentKey + '"'));
-      currentKey = void 0;
+      currentKey = undefined;
     }
   }
-  function processValueStart(char, i, swapState) {
+
+  function processValueStart(char: string, i: number, swapState: State) {
     {
       switch (char) {
         case '"': {
@@ -19,9 +83,11 @@ function fixJson(input) {
           stack.pop();
           stack.push(swapState);
           stack.push("INSIDE_STRING");
+
           pushCurrentKeyToPath();
           break;
         }
+
         case "f":
         case "t":
         case "n": {
@@ -32,10 +98,12 @@ function fixJson(input) {
           stack.push("INSIDE_LITERAL");
           break;
         }
+
         case "-": {
           stack.pop();
           stack.push(swapState);
           stack.push("INSIDE_NUMBER");
+
           pushCurrentKeyToPath();
           break;
         }
@@ -53,29 +121,35 @@ function fixJson(input) {
           stack.pop();
           stack.push(swapState);
           stack.push("INSIDE_NUMBER");
+
           pushCurrentKeyToPath();
           break;
         }
+
         case "{": {
           lastValidIndex = i;
           stack.pop();
           stack.push(swapState);
           stack.push("INSIDE_OBJECT_START");
+
           pushCurrentKeyToPath();
           break;
         }
+
         case "[": {
           lastValidIndex = i;
           stack.pop();
           stack.push(swapState);
           stack.push("INSIDE_ARRAY_START");
+
           pushCurrentKeyToPath();
           break;
         }
       }
     }
   }
-  function processAfterObjectValue(char, i) {
+
+  function processAfterObjectValue(char: string, i: number) {
     switch (char) {
       case ",": {
         stack.pop();
@@ -90,7 +164,8 @@ function fixJson(input) {
       }
     }
   }
-  function processAfterArrayValue(char, i) {
+
+  function processAfterArrayValue(char: string, i: number) {
     switch (char) {
       case ",": {
         stack.pop();
@@ -106,13 +181,16 @@ function fixJson(input) {
       }
     }
   }
+
   for (let i = 0; i < input.length; i++) {
-    const char = input[i];
+    const char = input[i]!;
     const currentState = stack[stack.length - 1];
+
     switch (currentState) {
       case "ROOT":
         processValueStart(char, i, "FINISH");
         break;
+
       case "INSIDE_OBJECT_START": {
         switch (char) {
           case '"': {
@@ -130,6 +208,7 @@ function fixJson(input) {
         }
         break;
       }
+
       case "INSIDE_OBJECT_AFTER_COMMA": {
         switch (char) {
           case '"': {
@@ -141,6 +220,7 @@ function fixJson(input) {
         }
         break;
       }
+
       case "INSIDE_OBJECT_KEY": {
         switch (char) {
           case '"': {
@@ -160,42 +240,52 @@ function fixJson(input) {
         }
         break;
       }
+
       case "INSIDE_OBJECT_AFTER_KEY": {
         switch (char) {
           case ":": {
             stack.pop();
             stack.push("INSIDE_OBJECT_BEFORE_VALUE");
+
             break;
           }
         }
         break;
       }
+
       case "INSIDE_OBJECT_BEFORE_VALUE": {
         processValueStart(char, i, "INSIDE_OBJECT_AFTER_VALUE");
         break;
       }
+
       case "INSIDE_OBJECT_AFTER_VALUE": {
         processAfterObjectValue(char, i);
         break;
       }
+
       case "INSIDE_STRING": {
         switch (char) {
           case '"': {
             stack.pop();
             lastValidIndex = i;
+
             currentKey = path.pop();
             break;
           }
+
           case "\\": {
             stack.push("INSIDE_STRING_ESCAPE");
             break;
           }
+
           default: {
             lastValidIndex = i;
           }
         }
+
         break;
       }
+
       case "INSIDE_ARRAY_START": {
         switch (char) {
           case "]": {
@@ -204,6 +294,7 @@ function fixJson(input) {
             currentKey = path.pop();
             break;
           }
+
           default: {
             lastValidIndex = i;
             currentKey = "0";
@@ -213,40 +304,50 @@ function fixJson(input) {
         }
         break;
       }
+
       case "INSIDE_ARRAY_AFTER_VALUE": {
         switch (char) {
           case ",": {
             stack.pop();
             stack.push("INSIDE_ARRAY_AFTER_COMMA");
+
             currentKey = (Number(currentKey) + 1).toString();
             break;
           }
+
           case "]": {
             lastValidIndex = i;
             stack.pop();
             currentKey = path.pop();
             break;
           }
+
           default: {
             lastValidIndex = i;
             break;
           }
         }
+
         break;
       }
+
       case "INSIDE_ARRAY_AFTER_COMMA": {
         processValueStart(char, i, "INSIDE_ARRAY_AFTER_VALUE");
         break;
       }
+
       case "INSIDE_STRING_ESCAPE": {
         stack.pop();
+
         if (stack[stack.length - 1] === "INSIDE_STRING") {
           lastValidIndex = i;
         } else if (stack[stack.length - 1] === "INSIDE_OBJECT_KEY") {
           currentKey += char;
         }
+
         break;
       }
+
       case "INSIDE_NUMBER": {
         switch (char) {
           case "0":
@@ -262,51 +363,71 @@ function fixJson(input) {
             lastValidIndex = i;
             break;
           }
+
           case "e":
           case "E":
           case "-":
           case ".": {
             break;
           }
+
           case ",": {
             stack.pop();
             currentKey = path.pop();
+
             if (stack[stack.length - 1] === "INSIDE_ARRAY_AFTER_VALUE") {
               processAfterArrayValue(char, i);
             }
+
             if (stack[stack.length - 1] === "INSIDE_OBJECT_AFTER_VALUE") {
               processAfterObjectValue(char, i);
             }
+
             break;
           }
+
           case "}": {
             stack.pop();
             currentKey = path.pop();
+
             if (stack[stack.length - 1] === "INSIDE_OBJECT_AFTER_VALUE") {
               processAfterObjectValue(char, i);
             }
+
             break;
           }
+
           case "]": {
             stack.pop();
             currentKey = path.pop();
+
             if (stack[stack.length - 1] === "INSIDE_ARRAY_AFTER_VALUE") {
               processAfterArrayValue(char, i);
             }
+
             break;
           }
+
           default: {
             stack.pop();
             currentKey = path.pop();
             break;
           }
         }
+
         break;
       }
+
       case "INSIDE_LITERAL": {
-        const partialLiteral = input.substring(literalStart, i + 1);
-        if (!"false".startsWith(partialLiteral) && !"true".startsWith(partialLiteral) && !"null".startsWith(partialLiteral)) {
+        const partialLiteral = input.substring(literalStart!, i + 1);
+
+        if (
+          !"false".startsWith(partialLiteral) &&
+          !"true".startsWith(partialLiteral) &&
+          !"null".startsWith(partialLiteral)
+        ) {
           stack.pop();
+
           if (stack[stack.length - 1] === "INSIDE_OBJECT_AFTER_VALUE") {
             processAfterObjectValue(char, i);
           } else if (stack[stack.length - 1] === "INSIDE_ARRAY_AFTER_VALUE") {
@@ -315,18 +436,23 @@ function fixJson(input) {
         } else {
           lastValidIndex = i;
         }
+
         break;
       }
     }
   }
+
   let result = input.slice(0, lastValidIndex + 1);
+
   for (let i = stack.length - 1; i >= 0; i--) {
     const state = stack[i];
+
     switch (state) {
       case "INSIDE_STRING": {
         result += '"';
         break;
       }
+
       case "INSIDE_OBJECT_KEY":
       case "INSIDE_OBJECT_AFTER_KEY":
       case "INSIDE_OBJECT_AFTER_COMMA":
@@ -336,14 +462,17 @@ function fixJson(input) {
         result += "}";
         break;
       }
+
       case "INSIDE_ARRAY_START":
       case "INSIDE_ARRAY_AFTER_COMMA":
       case "INSIDE_ARRAY_AFTER_VALUE": {
         result += "]";
         break;
       }
+
       case "INSIDE_LITERAL": {
-        const partialLiteral = input.substring(literalStart, input.length);
+        const partialLiteral = input.substring(literalStart!, input.length);
+
         if ("true".startsWith(partialLiteral)) {
           result += "true".slice(partialLiteral.length);
         } else if ("false".startsWith(partialLiteral)) {
@@ -354,9 +483,6 @@ function fixJson(input) {
       }
     }
   }
+
   return [result, path];
 }
-export {
-  fixJson
-};
-//# sourceMappingURL=fix-json.mjs.map