feat: switch to DataStream transfer protocol for edge runtime (#953)

Yonom · Oct 8, 2024 · 0ff22a7 · 0ff22a7
1 parent 8a9c654
commit 0ff22a7
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 22 deletions.
diff --git a/.changeset/rich-trees-care.md b/.changeset/rich-trees-care.md
@@ -0,0 +1,5 @@
+---
+"@assistant-ui/react": patch
+---
+
+feat: switch to DataStream transfer protocol for edge runtime
diff --git a/packages/react/src/runtimes/edge/streams/AssistantStreamChunkType.ts b/packages/react/src/runtimes/edge/streams/AssistantStreamChunkType.ts
@@ -2,22 +2,25 @@ import { LanguageModelV1StreamPart } from "@ai-sdk/provider";
 
 export enum AssistantStreamChunkType {
   TextDelta = "0",
-  ToolCallBegin = "1",
-  ToolCallArgsTextDelta = "2",
-  ToolCallResult = "3",
-  Error = "E",
-  Finish = "F",
+  ToolCallBegin = "b",
+  ToolCallDelta = "c",
+  ToolCallResult = "a",
+  Error = "3",
+  Finish = "d",
 }
 
 export type AssistantStreamChunk = {
   [AssistantStreamChunkType.TextDelta]: string;
   [AssistantStreamChunkType.ToolCallBegin]: {
-    id: string;
-    name: string;
+    toolCallId: string;
+    toolName: string;
+  };
+  [AssistantStreamChunkType.ToolCallDelta]: {
+    toolCallId: string;
+    argsTextDelta: string;
   };
-  [AssistantStreamChunkType.ToolCallArgsTextDelta]: string;
   [AssistantStreamChunkType.ToolCallResult]: {
-    id: string;
+    toolCallId: string;
     result: any;
   };
   [AssistantStreamChunkType.Error]: unknown;

diff --git a/packages/react/src/runtimes/edge/streams/assistantDecoderStream.ts b/packages/react/src/runtimes/edge/streams/assistantDecoderStream.ts
@@ -18,7 +18,7 @@ export function assistantDecoderStream() {
     transform({ type, value }, controller) {
       if (
         currentToolCall &&
-        type !== AssistantStreamChunkType.ToolCallArgsTextDelta &&
+        type !== AssistantStreamChunkType.ToolCallDelta &&
         type !== AssistantStreamChunkType.Error
       ) {
         controller.enqueue({
@@ -40,29 +40,35 @@ export function assistantDecoderStream() {
           break;
         }
         case AssistantStreamChunkType.ToolCallBegin: {
-          const { id, name } = value;
+          const { toolCallId: id, toolName: name } = value;
           toolCallNames.set(id, name);
           currentToolCall = { id, name, argsText: "" };
           break;
         }
-        case AssistantStreamChunkType.ToolCallArgsTextDelta: {
-          const delta = value;
-          currentToolCall!.argsText += delta;
+        case AssistantStreamChunkType.ToolCallDelta: {
+          const { toolCallId, argsTextDelta } = value;
+          if (currentToolCall?.id !== toolCallId) {
+            throw new Error(
+              `Received tool call delta for unknown tool call "${toolCallId}".`,
+            );
+          }
+
+          currentToolCall!.argsText += argsTextDelta;
           controller.enqueue({
             type: "tool-call-delta",
             toolCallType: "function",
             toolCallId: currentToolCall!.id,
             toolName: currentToolCall!.name,
-            argsTextDelta: delta,
+            argsTextDelta: argsTextDelta,
           });
           break;
         }
         case AssistantStreamChunkType.ToolCallResult: {
           controller.enqueue({
             type: "tool-result",
             toolCallType: "function",
-            toolCallId: value.id,
-            toolName: toolCallNames.get(value.id)!,
+            toolCallId: value.toolCallId,
+            toolName: toolCallNames.get(value.toolCallId)!,
             result: value.result,
           });
           break;

diff --git a/packages/react/src/runtimes/edge/streams/assistantEncoderStream.ts b/packages/react/src/runtimes/edge/streams/assistantEncoderStream.ts
@@ -27,15 +27,18 @@ export function assistantEncoderStream() {
             controller.enqueue({
               type: AssistantStreamChunkType.ToolCallBegin,
               value: {
-                id: chunk.toolCallId,
-                name: chunk.toolName,
+                toolCallId: chunk.toolCallId,
+                toolName: chunk.toolName,
               },
             });
           }
 
           controller.enqueue({
-            type: AssistantStreamChunkType.ToolCallArgsTextDelta,
-            value: chunk.argsTextDelta,
+            type: AssistantStreamChunkType.ToolCallDelta,
+            value: {
+              toolCallId: chunk.toolCallId,
+              argsTextDelta: chunk.argsTextDelta,
+            },
           });
           break;
         }
@@ -49,7 +52,7 @@ export function assistantEncoderStream() {
           controller.enqueue({
             type: AssistantStreamChunkType.ToolCallResult,
             value: {
-              id: chunk.toolCallId,
+              toolCallId: chunk.toolCallId,
               result: chunk.result,
             },
           });