streamfold 0.1.1

package/src/internal/wasm-runtime.js

@@ -0,0 +1,227 @@
+import { wasmBinaryBase64 } from "./wasm-binary.js";
+
+const DEPTH_MASK = 0x00ff_ffff;
+const COMPLETE_FLAG = 1 << 24;
+const IN_STRING_FLAG = 1 << 25;
+const ERROR_SHIFT = 28;
+const PATCH_SET_OBJECT = 0;
+const PATCH_SET_ARRAY = 1;
+const PATCH_SET_STRING = 2;
+const PATCH_APPEND_STRING = 3;
+const PATCH_SET_NUMBER = 4;
+const PATCH_SET_TRUE = 5;
+const PATCH_SET_FALSE = 6;
+const PATCH_SET_NULL = 7;
+const PATCH_COMPLETE = 8;
+
+const encoder = new TextEncoder();
+let exports;
+
+const decodeBase64 = (base64) => {
+  const binary = globalThis.atob(base64);
+  const bytes = new Uint8Array(binary.length);
+  for (let index = 0; index < binary.length; index++) {
+    bytes[index] = binary.charCodeAt(index);
+  }
+  return bytes;
+};
+
+const getExports = () => {
+  if (exports !== undefined) return exports;
+  if (typeof WebAssembly !== "object") {
+    throw new Error("Streamfold requires WebAssembly support");
+  }
+  const module = new WebAssembly.Module(decodeBase64(wasmBinaryBase64));
+  exports = new WebAssembly.Instance(module, {}).exports;
+  return exports;
+};
+
+const parserError = ({ wasm, handle, maxBytes, maxDepth }, code) => {
+  const offset = wasm.streamfold_parser_error_offset(handle) >>> 0;
+  const byte = wasm.streamfold_parser_error_byte(handle);
+  if (code === 1) {
+    const kind = byte === 44 || byte === 58 ? "separator" : "closing";
+    return new SyntaxError(`Unexpected ${kind} at ${offset}`);
+  }
+  if (code === 2) return new SyntaxError(`Mismatched closing at ${offset}`);
+  if (code === 3) return new SyntaxError(`Trailing data at ${offset}`);
+  if (code === 4) return new SyntaxError("Empty JSON input");
+  if (code === 5) return new SyntaxError(`Incomplete JSON at ${offset}`);
+  if (code === 6) return new SyntaxError(`Invalid JSON at ${offset}`);
+  if (code === 7) {
+    return new RangeError(
+      `Structured stream exceeds maxBytes (${maxBytes}) at ${offset}`,
+    );
+  }
+  if (code === 8) {
+    return new RangeError(
+      `Structured stream exceeds maxDepth (${maxDepth}) at ${offset}`,
+    );
+  }
+  return new SyntaxError(`Rust parser failed with error code ${code}`);
+};
+
+const readState = (parser, encoded) => {
+  const { wasm, handle } = parser;
+  const error = encoded >>> ERROR_SHIFT;
+  if (error !== 0) throw parserError(parser, error);
+  return {
+    bytesSeen: wasm.streamfold_parser_bytes_seen(handle) >>> 0,
+    depth: encoded & DEPTH_MASK,
+    complete: (encoded & COMPLETE_FLAG) !== 0,
+    inString: (encoded & IN_STRING_FLAG) !== 0,
+  };
+};
+
+const readString = (view, cursor, length) => {
+  let value = "";
+  const batch = [];
+  for (let index = 0; index < length; index++) {
+    batch.push(view.getUint16(cursor.offset, true));
+    cursor.offset += 2;
+    if (batch.length === 4096) {
+      value += String.fromCharCode(...batch);
+      batch.length = 0;
+    }
+  }
+  if (batch.length > 0) value += String.fromCharCode(...batch);
+  return value;
+};
+
+const readPath = (view, cursor) => {
+  const length = view.getUint16(cursor.offset, true);
+  cursor.offset += 2;
+  const path = [];
+  for (let index = 0; index < length; index++) {
+    const kind = view.getUint8(cursor.offset++);
+    if (kind === 0) {
+      const units = view.getUint32(cursor.offset, true);
+      cursor.offset += 4;
+      path.push(readString(view, cursor, units));
+    } else {
+      path.push(view.getUint32(cursor.offset, true));
+      cursor.offset += 4;
+    }
+  }
+  return path;
+};
+
+const readNumber = (view, cursor) => {
+  const length = view.getUint32(cursor.offset, true);
+  cursor.offset += 4;
+  let text = "";
+  for (let index = 0; index < length; index++) {
+    text += String.fromCharCode(view.getUint8(cursor.offset++));
+  }
+  return Number(text);
+};
+
+const readPatches = (wasm, handle) => {
+  const length = wasm.streamfold_parser_output_len(handle) >>> 0;
+  if (length === 0) return [];
+  const pointer = wasm.streamfold_parser_output(handle) >>> 0;
+  const view = new DataView(wasm.memory.buffer, pointer, length);
+  const cursor = { offset: 0 };
+  const patches = [];
+
+  while (cursor.offset < length) {
+    const operation = view.getUint8(cursor.offset++);
+    const path = readPath(view, cursor);
+    if (operation === PATCH_SET_OBJECT) {
+      patches.push({ op: "set", path, value: {} });
+    } else if (operation === PATCH_SET_ARRAY) {
+      patches.push({ op: "set", path, value: [] });
+    } else if (operation === PATCH_SET_STRING) {
+      patches.push({ op: "set", path, value: "" });
+    } else if (operation === PATCH_APPEND_STRING) {
+      const units = view.getUint32(cursor.offset, true);
+      cursor.offset += 4;
+      patches.push({
+        op: "append",
+        path,
+        value: readString(view, cursor, units),
+      });
+    } else if (operation === PATCH_SET_NUMBER) {
+      patches.push({ op: "set", path, value: readNumber(view, cursor) });
+    } else if (operation === PATCH_SET_TRUE) {
+      patches.push({ op: "set", path, value: true });
+    } else if (operation === PATCH_SET_FALSE) {
+      patches.push({ op: "set", path, value: false });
+    } else if (operation === PATCH_SET_NULL) {
+      patches.push({ op: "set", path, value: null });
+    } else if (operation === PATCH_COMPLETE) {
+      patches.push({ op: "complete", path });
+    } else {
+      throw new Error(`Unknown Rust patch operation: ${operation}`);
+    }
+  }
+
+  return patches;
+};
+
+export const createWasmParser = ({ maxBytes, maxDepth }) => {
+  const wasm = getExports();
+  const handle = wasm.streamfold_parser_new(maxDepth, maxBytes);
+  if (handle === 0) throw new Error("Unable to allocate the Rust parser");
+  return {
+    wasm,
+    handle,
+    maxBytes,
+    maxDepth,
+    pendingHighSurrogate: "",
+  };
+};
+
+const pushChunk = (parser, chunk) => {
+  const { wasm, handle } = parser;
+  const capacity = chunk.length * 3;
+  const pointer = wasm.streamfold_parser_input(handle, capacity);
+  let length = 0;
+  if (capacity > 0) {
+    const input = new Uint8Array(wasm.memory.buffer, pointer, capacity);
+    const result = encoder.encodeInto(chunk, input);
+    if (result.read !== chunk.length) {
+      throw new Error("Unable to encode the complete stream fragment");
+    }
+    length = result.written;
+  }
+  const state = readState(parser, wasm.streamfold_parser_push(handle, length));
+  return { state, changes: readPatches(wasm, handle) };
+};
+
+export const pushWasmParser = (parser, chunk) => {
+  let input = parser.pendingHighSurrogate + chunk;
+  parser.pendingHighSurrogate = "";
+
+  if (input.length > 0) {
+    const lastUnit = input.charCodeAt(input.length - 1);
+    if (lastUnit >= 0xd800 && lastUnit <= 0xdbff) {
+      parser.pendingHighSurrogate = input.at(-1);
+      input = input.slice(0, -1);
+    }
+  }
+
+  return pushChunk(parser, input);
+};
+
+export const finishWasmParser = (parser) => {
+  const { wasm, handle } = parser;
+  const changes = [];
+  if (parser.pendingHighSurrogate.length > 0) {
+    changes.push(...pushChunk(parser, parser.pendingHighSurrogate).changes);
+    parser.pendingHighSurrogate = "";
+  }
+  const state = readState(parser, wasm.streamfold_parser_finish(handle));
+  changes.push(...readPatches(wasm, handle));
+  return { state, changes };
+};
+
+export const readWasmParser = (parser) =>
+  readState(
+    parser,
+    parser.wasm.streamfold_parser_state(parser.handle),
+  );
+
+export const freeWasmParser = ({ wasm, handle }) => {
+  wasm.streamfold_parser_free(handle);
+};

package/src/langchain.d.ts

@@ -0,0 +1,26 @@
+import type {
+  EventStructuredStream,
+  StructuredStreamIntegration,
+  StructuredStreamPool,
+} from "./index.js";
+
+export interface LangChainToolCallChunk {
+  readonly type?: string;
+  readonly index: number | string;
+  readonly id?: string;
+  readonly name?: string;
+  readonly args?: string;
+}
+
+export interface LangChainToolCallMessage {
+  readonly tool_call_chunks?: readonly LangChainToolCallChunk[];
+}
+
+export const langchain: StructuredStreamIntegration<
+  LangChainToolCallMessage,
+  string
+>;
+
+export function createStructuredStream(
+  pool?: StructuredStreamPool<string>,
+): EventStructuredStream<LangChainToolCallMessage, string>;

package/src/langchain.js

@@ -0,0 +1,31 @@
+import { createStructuredStreamPool } from "./index.js";
+import { createIntegration } from "./internal/integration.js";
+
+export const langchain = (pool = createStructuredStreamPool()) => {
+  const idsByIndex = new Map();
+
+  return createIntegration(
+    pool,
+    (message) => {
+      let update;
+      for (const chunk of message.tool_call_chunks ?? []) {
+        let id = idsByIndex.get(chunk.index);
+        if (id === undefined && chunk.id) {
+          id = chunk.id;
+          idsByIndex.set(chunk.index, id);
+          update = pool.start(id);
+        }
+        if (id !== undefined && chunk.args) {
+          update = pool.push(id, chunk.args);
+        }
+      }
+      return update;
+    },
+    (complete) => {
+      for (const id of idsByIndex.values()) complete(id);
+      idsByIndex.clear();
+    },
+  );
+};
+
+export { langchain as createStructuredStream };

package/src/openai.d.ts

@@ -0,0 +1,37 @@
+import type {
+  EventStructuredStream,
+  StructuredStreamIntegration,
+  StructuredStreamPool,
+} from "./index.js";
+
+export type OpenAiResponsesToolEvent =
+  | {
+      readonly type: "response.output_item.added";
+      readonly item: {
+        readonly type: string;
+        readonly id: string;
+      };
+    }
+  | {
+      readonly type: "response.function_call_arguments.delta";
+      readonly item_id: string;
+      readonly delta: string;
+    }
+  | {
+      readonly type: "response.function_call_arguments.done";
+      readonly item_id: string;
+      readonly arguments: string;
+    }
+  | {
+      readonly type: string;
+      readonly [key: string]: unknown;
+    };
+
+export const openAI: StructuredStreamIntegration<
+  OpenAiResponsesToolEvent,
+  string
+>;
+
+export function createStructuredStream(
+  pool?: StructuredStreamPool<string>,
+): EventStructuredStream<OpenAiResponsesToolEvent, string>;

package/src/openai.js

@@ -0,0 +1,20 @@
+import { createStructuredStreamPool } from "./index.js";
+import { append, createIntegration } from "./internal/integration.js";
+
+export const openAI = (pool = createStructuredStreamPool()) =>
+  createIntegration(pool, (event, complete) => {
+    if (event.type === "response.output_item.added") {
+      if (event.item.type !== "function_call") return undefined;
+      return pool.start(event.item.id);
+    }
+    if (event.type === "response.function_call_arguments.delta") {
+      return append(pool, event.item_id, event.delta);
+    }
+    if (event.type === "response.function_call_arguments.done") {
+      if (!pool.has(event.item_id)) pool.start(event.item_id, event.arguments);
+      return complete(event.item_id);
+    }
+    return undefined;
+  });
+
+export { openAI as createStructuredStream };

package/src/vercel-ai.d.ts

@@ -0,0 +1,52 @@
+import type {
+  EventStructuredStream,
+  StructuredStreamIntegration,
+  StructuredStreamPool,
+} from "./index.js";
+
+export type VercelAiToolInputEvent =
+  | {
+      readonly type: "tool-input-start";
+      readonly id: string;
+      readonly toolName: string;
+    }
+  | {
+      readonly type: "tool-input-start";
+      readonly toolCallId: string;
+      readonly toolName: string;
+    }
+  | {
+      readonly type: "tool-input-delta";
+      readonly id: string;
+      readonly delta: string;
+    }
+  | {
+      readonly type: "tool-input-delta";
+      readonly toolCallId: string;
+      readonly inputTextDelta: string;
+    }
+  | {
+      readonly type: "tool-input-end";
+      readonly id: string;
+    }
+  | {
+      readonly type: "tool-input-available";
+      readonly toolCallId: string;
+      readonly toolName: string;
+      readonly input: unknown;
+    }
+  | {
+      readonly type: string;
+      readonly id?: string;
+      readonly toolCallId?: string;
+      readonly [key: string]: unknown;
+    };
+
+export const vercelAI: StructuredStreamIntegration<
+  VercelAiToolInputEvent,
+  string
+>;
+
+export function createStructuredStream(
+  pool?: StructuredStreamPool<string>,
+): EventStructuredStream<VercelAiToolInputEvent, string>;

package/src/vercel-ai.js

@@ -0,0 +1,22 @@
+import { createStructuredStreamPool } from "./index.js";
+import { createIntegration } from "./internal/integration.js";
+
+export const vercelAI = (pool = createStructuredStreamPool()) =>
+  createIntegration(pool, (event, complete) => {
+    const id = event.id ?? event.toolCallId;
+    if (id === undefined) return undefined;
+
+    if (event.type === "tool-input-start") return pool.start(id);
+    if (event.type === "tool-input-delta") {
+      return pool.push(id, event.delta ?? event.inputTextDelta);
+    }
+    if (
+      event.type === "tool-input-end" ||
+      event.type === "tool-input-available"
+    ) {
+      return complete(id);
+    }
+    return undefined;
+  });
+
+export { vercelAI as createStructuredStream };