feat: assistant-stream (#952)

packages/assistant-stream/package.json

@@ -1,6 +1,62 @@
 {
   "name": "assistant-stream",
-  "version": "0.0.0",
-  "description": "",
-  "author": "Simon Farshid"
+  "version": "0.0.0-rc.2",
+  "license": "MIT",
+  "exports": {
+    ".": {
+      "import": {
+        "types": "./dist/index.d.mts",
+        "default": "./dist/index.mjs"
+      },
+      "require": {
+        "types": "./dist/index.d.ts",
+        "default": "./dist/index.js"
+      }
+    },
+    "./ai-sdk": {
+      "import": {
+        "types": "./dist/ai-sdk.d.mts",
+        "default": "./dist/ai-sdk.mjs"
+      },
+      "require": {
+        "types": "./dist/ai-sdk.d.ts",
+        "default": "./dist/ai-sdk.js"
+      }
+    }
+  },
+  "source": "./src/index.ts",
+  "main": "./dist/index.js",
+  "module": "./dist/index.mjs",
+  "types": "./dist/index.d.ts",
+  "files": [
+    "dist",
+    "README.md"
+  ],
+  "sideEffects": false,
+  "scripts": {
+    "build": "tsx scripts/build.mts"
+  },
+  "devDependencies": {
+    "@assistant-ui/tsconfig": "workspace:*",
+    "ai": "^3.4.9",
+    "eslint": "^8",
+    "eslint-config-next": "14.2.14",
+    "tsup": "8.3.0",
+    "tsx": "^4.19.1"
+  },
+  "publishConfig": {
+    "access": "public"
+  },
+  "homepage": "https://assistant-ui.com/",
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/Yonom/assistant-ui.git"
+  },
+  "bugs": {
+    "url": "https://github.com/Yonom/assistant-ui/issues"
+  },
+  "dependencies": {
+    "nanoid": "^5.0.7",
+    "secure-json-parse": "^2.7.0"
+  }
 }

packages/assistant-stream/scripts/build.mts

@@ -0,0 +1,10 @@
+import { build } from "tsup";
+
+// JS
+await build({
+  entry: ["src/index.ts", "src/ai-sdk.ts"],
+  format: ["cjs", "esm"],
+  dts: true,
+  sourcemap: true,
+  clean: true,
+});

packages/assistant-stream/src/ai-sdk.ts

@@ -0,0 +1 @@
+export * from "./ai-sdk/index";

packages/assistant-stream/src/ai-sdk/index.ts

@@ -0,0 +1,135 @@
+import type { TextStreamPart, CoreTool, ObjectStreamPart } from "ai";
+import { AssistantStream, AssistantStreamChunk } from "../core";
+import { generateId } from "../core/utils/generateId";
+
+export const fromAISDKStreamText = (
+  stream: ReadableStream<TextStreamPart<Record<string, CoreTool>>>,
+): AssistantStream => {
+  const transformer = new TransformStream<
+    TextStreamPart<Record<string, CoreTool>>,
+    AssistantStreamChunk
+  >({
+    transform(chunk, controller) {
+      const { type } = chunk;
+      switch (type) {
+        case "text-delta": {
+          const { textDelta } = chunk;
+          controller.enqueue({
+            type: "text-delta",
+            textDelta,
+          });
+          break;
+        }
+        case "tool-call-streaming-start": {
+          const { toolCallId, toolName } = chunk;
+          controller.enqueue({
+            type: "tool-call-begin",
+            toolCallId,
+            toolName,
+          });
+          break;
+        }
+        case "tool-call-delta": {
+          const { toolCallId, argsTextDelta } = chunk;
+          controller.enqueue({
+            type: "tool-call-delta",
+            toolCallId,
+            argsTextDelta,
+          });
+          break;
+        }
+        case "tool-result" as string: {
+          const { toolCallId, result } = chunk as unknown as {
+            toolCallId: string;
+            result: unknown;
+          };
+          controller.enqueue({
+            type: "tool-result",
+            toolCallId,
+            result,
+          });
+          break;
+        }
+        case "tool-call": {
+          const { toolCallId, toolName, args } = chunk;
+          controller.enqueue({
+            type: "tool-call-begin",
+            toolCallId,
+            toolName,
+          });
+          controller.enqueue({
+            type: "tool-call-delta",
+            toolCallId,
+            argsTextDelta: JSON.stringify(args),
+          });
+          break;
+        }
+        case "step-finish":
+        case "error":
+        case "finish": {
+          break;
+        }
+
+        default: {
+          const unhandledType: never = type;
+          throw new Error(`Unhandled chunk type: ${unhandledType}`);
+        }
+      }
+    },
+  });
+
+  return new AssistantStream(stream.pipeThrough(transformer));
+};
+
+export const fromAISDKStreamObject = (
+  stream: ReadableStream<ObjectStreamPart<unknown>>,
+  toolName: string,
+): AssistantStream => {
+  const toolCallId = generateId();
+  const transformer = new TransformStream<
+    ObjectStreamPart<unknown>,
+    AssistantStreamChunk
+  >({
+    start(controller) {
+      controller.enqueue({
+        type: "tool-call-begin",
+        toolName,
+        toolCallId,
+      });
+    },
+    transform(chunk, controller) {
+      const { type } = chunk;
+      switch (type) {
+        case "text-delta": {
+          const { textDelta } = chunk;
+          controller.enqueue({
+            type: "tool-call-delta",
+            toolCallId,
+            argsTextDelta: textDelta,
+          });
+          break;
+        }
+        case "finish": {
+          controller.enqueue({
+            type: "tool-result",
+            toolCallId,
+            result: "",
+          });
+          break;
+        }
+
+        case "object":
+        case "error": {
+          break;
+        }
+
+        default: {
+          const unhandledType: never = type;
+          throw new Error(`Unhandled chunk type: ${unhandledType}`);
+        }
+      }
+    },
+  });
+
+  return new AssistantStream(stream.pipeThrough(transformer));
+};

packages/assistant-stream/src/core/AssistantStream.ts

@@ -0,0 +1,69 @@
+export type AssistantStream2 = {
+  // kind: AssistantStreamKind;
+  readable: ReadableStream<AssistantStreamChunk>;
+  toResponse(
+    format: ReadableWritablePair<Uint8Array, AssistantStreamChunk>,
+  ): Response;
+  toByteStream(
+    format: ReadableWritablePair<Uint8Array, AssistantStreamChunk>,
+  ): ReadableStream<Uint8Array>;
+  tee(): [AssistantStream, AssistantStream];
+};
+
+export type AssistantStreamChunk =
+  | {
+      type: "text-delta";
+      textDelta: string;
+    }
+  | {
+      type: "tool-call-begin";
+      toolCallId: string;
+      toolName: string;
+    }
+  | {
+      type: "tool-call-delta";
+      toolCallId: string;
+      argsTextDelta: string;
+    }
+  | {
+      type: "tool-result";
+      toolCallId: string;
+      result: any;
+    };
+
+export class AssistantStream {
+  constructor(public readonly readable: ReadableStream<AssistantStreamChunk>) {
+    this.readable = readable;
+  }
+
+  toResponse(
+    transformer: ReadableWritablePair<Uint8Array, AssistantStreamChunk>,
+  ) {
+    return new Response(this.toByteStream(transformer));
+  }
+
+  static fromResponse(
+    response: Response,
+    transformer: ReadableWritablePair<AssistantStreamChunk, Uint8Array>,
+  ) {
+    return AssistantStream.fromByteStream(response.body!, transformer);
+  }
+
+  toByteStream(
+    transformer: ReadableWritablePair<Uint8Array, AssistantStreamChunk>,
+  ) {
+    return this.readable.pipeThrough(transformer);
+  }
+
+  static fromByteStream(
+    readable: ReadableStream<Uint8Array>,
+    transformer: ReadableWritablePair<AssistantStreamChunk, Uint8Array>,
+  ) {
+    return new AssistantStream(readable.pipeThrough(transformer));
+  }
+
+  tee(): [AssistantStream, AssistantStream] {
+    const [readable1, readable2] = this.readable.tee();
+    return [new AssistantStream(readable1), new AssistantStream(readable2)];
+  }
+}