@assistant-ui/react 0.5.91 → 0.5.93

package/src/runtimes/edge/partial-json/fix-json.ts

@@ -0,0 +1,417 @@
+// LICENSE for this file only
+
+// Copyright 2023 Vercel, Inc.
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+
+//     http://www.apache.org/licenses/LICENSE-2.0
+
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+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.
+export function fixJson(input: string): string {
+  const stack: State[] = ["ROOT"];
+  let lastValidIndex = -1;
+  let literalStart: number | null = null;
+
+  function processValueStart(char: string, i: number, swapState: State) {
+    {
+      switch (char) {
+        case '"': {
+          lastValidIndex = i;
+          stack.pop();
+          stack.push(swapState);
+          stack.push("INSIDE_STRING");
+          break;
+        }
+
+        case "f":
+        case "t":
+        case "n": {
+          lastValidIndex = i;
+          literalStart = i;
+          stack.pop();
+          stack.push(swapState);
+          stack.push("INSIDE_LITERAL");
+          break;
+        }
+
+        case "-": {
+          stack.pop();
+          stack.push(swapState);
+          stack.push("INSIDE_NUMBER");
+          break;
+        }
+        case "0":
+        case "1":
+        case "2":
+        case "3":
+        case "4":
+        case "5":
+        case "6":
+        case "7":
+        case "8":
+        case "9": {
+          lastValidIndex = i;
+          stack.pop();
+          stack.push(swapState);
+          stack.push("INSIDE_NUMBER");
+          break;
+        }
+
+        case "{": {
+          lastValidIndex = i;
+          stack.pop();
+          stack.push(swapState);
+          stack.push("INSIDE_OBJECT_START");
+          break;
+        }
+
+        case "[": {
+          lastValidIndex = i;
+          stack.pop();
+          stack.push(swapState);
+          stack.push("INSIDE_ARRAY_START");
+          break;
+        }
+      }
+    }
+  }
+
+  function processAfterObjectValue(char: string, i: number) {
+    switch (char) {
+      case ",": {
+        stack.pop();
+        stack.push("INSIDE_OBJECT_AFTER_COMMA");
+        break;
+      }
+      case "}": {
+        lastValidIndex = i;
+        stack.pop();
+        break;
+      }
+    }
+  }
+
+  function processAfterArrayValue(char: string, i: number) {
+    switch (char) {
+      case ",": {
+        stack.pop();
+        stack.push("INSIDE_ARRAY_AFTER_COMMA");
+        break;
+      }
+      case "]": {
+        lastValidIndex = i;
+        stack.pop();
+        break;
+      }
+    }
+  }
+
+  for (let i = 0; i < input.length; 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 '"': {
+            stack.pop();
+            stack.push("INSIDE_OBJECT_KEY");
+            break;
+          }
+          case "}": {
+            lastValidIndex = i;
+            stack.pop();
+            break;
+          }
+        }
+        break;
+      }
+
+      case "INSIDE_OBJECT_AFTER_COMMA": {
+        switch (char) {
+          case '"': {
+            stack.pop();
+            stack.push("INSIDE_OBJECT_KEY");
+            break;
+          }
+        }
+        break;
+      }
+
+      case "INSIDE_OBJECT_KEY": {
+        switch (char) {
+          case '"': {
+            stack.pop();
+            stack.push("INSIDE_OBJECT_AFTER_KEY");
+            break;
+          }
+        }
+        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;
+            break;
+          }
+
+          case "\\": {
+            stack.push("INSIDE_STRING_ESCAPE");
+            break;
+          }
+
+          default: {
+            lastValidIndex = i;
+          }
+        }
+
+        break;
+      }
+
+      case "INSIDE_ARRAY_START": {
+        switch (char) {
+          case "]": {
+            lastValidIndex = i;
+            stack.pop();
+            break;
+          }
+
+          default: {
+            lastValidIndex = i;
+            processValueStart(char, i, "INSIDE_ARRAY_AFTER_VALUE");
+            break;
+          }
+        }
+        break;
+      }
+
+      case "INSIDE_ARRAY_AFTER_VALUE": {
+        switch (char) {
+          case ",": {
+            stack.pop();
+            stack.push("INSIDE_ARRAY_AFTER_COMMA");
+            break;
+          }
+
+          case "]": {
+            lastValidIndex = i;
+            stack.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();
+        lastValidIndex = i;
+
+        break;
+      }
+
+      case "INSIDE_NUMBER": {
+        switch (char) {
+          case "0":
+          case "1":
+          case "2":
+          case "3":
+          case "4":
+          case "5":
+          case "6":
+          case "7":
+          case "8":
+          case "9": {
+            lastValidIndex = i;
+            break;
+          }
+
+          case "e":
+          case "E":
+          case "-":
+          case ".": {
+            break;
+          }
+
+          case ",": {
+            stack.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();
+
+            if (stack[stack.length - 1] === "INSIDE_OBJECT_AFTER_VALUE") {
+              processAfterObjectValue(char, i);
+            }
+
+            break;
+          }
+
+          case "]": {
+            stack.pop();
+
+            if (stack[stack.length - 1] === "INSIDE_ARRAY_AFTER_VALUE") {
+              processAfterArrayValue(char, i);
+            }
+
+            break;
+          }
+
+          default: {
+            stack.pop();
+            break;
+          }
+        }
+
+        break;
+      }
+
+      case "INSIDE_LITERAL": {
+        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") {
+            processAfterArrayValue(char, i);
+          }
+        } 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":
+      case "INSIDE_OBJECT_START":
+      case "INSIDE_OBJECT_BEFORE_VALUE":
+      case "INSIDE_OBJECT_AFTER_VALUE": {
+        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);
+
+        if ("true".startsWith(partialLiteral)) {
+          result += "true".slice(partialLiteral.length);
+        } else if ("false".startsWith(partialLiteral)) {
+          result += "false".slice(partialLiteral.length);
+        } else if ("null".startsWith(partialLiteral)) {
+          result += "null".slice(partialLiteral.length);
+        }
+      }
+    }
+  }
+
+  return result;
+}

package/src/runtimes/edge/partial-json/parse-partial-json.ts

@@ -0,0 +1,14 @@
+import sjson from "secure-json-parse";
+import { fixJson } from "./fix-json";
+
+export const parsePartialJson = (json: string) => {
+  try {
+    return sjson.parse(json);
+  } catch {
+    try {
+      return sjson.parse(fixJson(json));
+    } catch {
+      return undefined;
+    }
+  }
+};

package/src/runtimes/edge/streams/AssistantStreamChunkType.ts

@@ -0,0 +1,58 @@
+import { LanguageModelV1StreamPart } from "@ai-sdk/provider";
+
+export enum AssistantStreamChunkType {
+  TextDelta = "0",
+  Data = "2",
+  Error = "3",
+  ToolCall = "9",
+  ToolCallResult = "a",
+  ToolCallBegin = "b",
+  ToolCallDelta = "c",
+  Finish = "d",
+  StepFinish = "e",
+}
+
+export type AssistantStreamChunk = {
+  [AssistantStreamChunkType.TextDelta]: string;
+  [AssistantStreamChunkType.Data]: unknown;
+  [AssistantStreamChunkType.ToolCall]: {
+    toolCallId: string;
+    toolName: string;
+    args: unknown;
+  };
+  [AssistantStreamChunkType.ToolCallBegin]: {
+    toolCallId: string;
+    toolName: string;
+  };
+  [AssistantStreamChunkType.ToolCallDelta]: {
+    toolCallId: string;
+    argsTextDelta: string;
+  };
+  [AssistantStreamChunkType.ToolCallResult]: {
+    toolCallId: string;
+    result: any;
+  };
+  [AssistantStreamChunkType.Error]: unknown;
+  [AssistantStreamChunkType.StepFinish]: {
+    finishReason:
+      | "stop"
+      | "length"
+      | "content-filter"
+      | "tool-calls"
+      | "error"
+      | "other"
+      | "unknown";
+    usage: {
+      promptTokens: number;
+      completionTokens: number;
+    };
+    isContinued: boolean;
+  };
+  [AssistantStreamChunkType.Finish]: Omit<
+    LanguageModelV1StreamPart & {
+      type: "finish";
+    },
+    "type"
+  >;
+};
+ 

package/src/runtimes/edge/streams/assistantDecoderStream.ts

@@ -0,0 +1,139 @@
+import {
+  AssistantStreamChunk,
+  AssistantStreamChunkType,
+} from "./AssistantStreamChunkType";
+import { StreamPart } from "./utils/StreamPart";
+import { ToolResultStreamPart } from "./toolResultStream";
+
+export function assistantDecoderStream() {
+  const toolCallNames = new Map<string, string>();
+  let currentToolCall:
+    | { id: string; name: string; argsText: string }
+    | undefined;
+
+  return new TransformStream<
+    StreamPart<AssistantStreamChunk>,
+    ToolResultStreamPart
+  >({
+    transform({ type, value }, controller) {
+      if (
+        currentToolCall &&
+        type !== AssistantStreamChunkType.ToolCallDelta &&
+        type !== AssistantStreamChunkType.Error
+      ) {
+        controller.enqueue({
+          type: "tool-call",
+          toolCallType: "function",
+          toolCallId: currentToolCall.id,
+          toolName: currentToolCall.name,
+          args: currentToolCall.argsText,
+        });
+        currentToolCall = undefined;
+      }
+
+      switch (type) {
+        case AssistantStreamChunkType.TextDelta: {
+          controller.enqueue({
+            type: "text-delta",
+            textDelta: value,
+          });
+          break;
+        }
+        case AssistantStreamChunkType.ToolCallBegin: {
+          const { toolCallId: id, toolName: name } = value;
+          toolCallNames.set(id, name);
+
+          currentToolCall = { id, name, argsText: "" };
+
+          controller.enqueue({
+            type: "tool-call-delta",
+            toolCallType: "function",
+            toolCallId: id,
+            toolName: name,
+            argsTextDelta: "",
+          });
+          break;
+        }
+        case AssistantStreamChunkType.ToolCallDelta: {
+          const { toolCallId, argsTextDelta } = value;
+
+          const toolName = toolCallNames.get(toolCallId)!;
+          if (currentToolCall?.id === toolCallId) {
+            currentToolCall.argsText += argsTextDelta;
+          }
+          controller.enqueue({
+            type: "tool-call-delta",
+            toolCallType: "function",
+            toolCallId,
+            toolName,
+            argsTextDelta: argsTextDelta,
+          });
+          break;
+        }
+        case AssistantStreamChunkType.ToolCallResult: {
+          controller.enqueue({
+            type: "tool-result",
+            toolCallType: "function",
+            toolCallId: value.toolCallId,
+            toolName: toolCallNames.get(value.toolCallId)!,
+            result: value.result,
+          });
+          break;
+        }
+        case AssistantStreamChunkType.Finish: {
+          controller.enqueue({
+            type: "finish",
+            ...value,
+          });
+          break;
+        }
+        case AssistantStreamChunkType.Error: {
+          controller.enqueue({
+            type: "error",
+            error: value,
+          });
+          break;
+        }
+
+        case AssistantStreamChunkType.ToolCall: {
+          const { toolCallId, toolName, args } = value;
+          toolCallNames.set(toolCallId, toolName);
+
+          const argsText = JSON.stringify(args);
+          controller.enqueue({
+            type: "tool-call-delta",
+            toolCallType: "function",
+            toolCallId,
+            toolName,
+            argsTextDelta: argsText,
+          });
+          controller.enqueue({
+            type: "tool-call",
+            toolCallType: "function",
+            toolCallId: toolCallId,
+            toolName: toolName,
+            args: argsText,
+          });
+          break;
+        }
+
+        case AssistantStreamChunkType.StepFinish: {
+          controller.enqueue({
+            type: "step-finish",
+            ...value,
+          });
+          break;
+        }
+
+        // TODO
+        case AssistantStreamChunkType.Data:
+          break;
+
+        default: {
+          const unhandledType: never = type;
+          throw new Error(`Unhandled chunk type: ${unhandledType}`);
+        }
+      }
+    },
+  });
+}

package/src/runtimes/edge/streams/assistantEncoderStream.ts

@@ -0,0 +1,95 @@
+import {
+  AssistantStreamChunk,
+  AssistantStreamChunkType,
+} from "./AssistantStreamChunkType";
+import { StreamPart } from "./utils/StreamPart";
+import { ToolResultStreamPart } from "./toolResultStream";
+
+export function assistantEncoderStream() {
+  const toolCalls = new Set<string>();
+  return new TransformStream<
+    ToolResultStreamPart,
+    StreamPart<AssistantStreamChunk>
+  >({
+    transform(chunk, controller) {
+      const chunkType = chunk.type;
+      switch (chunkType) {
+        case "text-delta": {
+          if (!chunk.textDelta) break; // ignore empty text deltas
+          controller.enqueue({
+            type: AssistantStreamChunkType.TextDelta,
+            value: chunk.textDelta,
+          });
+          break;
+        }
+        case "tool-call-delta": {
+          if (!toolCalls.has(chunk.toolCallId)) {
+            toolCalls.add(chunk.toolCallId);
+            controller.enqueue({
+              type: AssistantStreamChunkType.ToolCallBegin,
+              value: {
+                toolCallId: chunk.toolCallId,
+                toolName: chunk.toolName,
+              },
+            });
+          }
+
+          controller.enqueue({
+            type: AssistantStreamChunkType.ToolCallDelta,
+            value: {
+              toolCallId: chunk.toolCallId,
+              argsTextDelta: chunk.argsTextDelta,
+            },
+          });
+          break;
+        }
+
+        // ignore
+        case "tool-call":
+        case "response-metadata":
+          break;
+
+        case "tool-result": {
+          controller.enqueue({
+            type: AssistantStreamChunkType.ToolCallResult,
+            value: {
+              toolCallId: chunk.toolCallId,
+              result: chunk.result,
+            },
+          });
+          break;
+        }
+
+        case "step-finish": {
+          const { type, ...rest } = chunk;
+          controller.enqueue({
+            type: AssistantStreamChunkType.StepFinish,
+            value: rest,
+          });
+          break;
+        }
+
+        case "finish": {
+          const { type, ...rest } = chunk;
+          controller.enqueue({
+            type: AssistantStreamChunkType.Finish,
+            value: rest,
+          });
+          break;
+        }
+
+        case "error": {
+          controller.enqueue({
+            type: AssistantStreamChunkType.Error,
+            value: chunk.error,
+          });
+          break;
+        }
+        default: {
+          const unhandledType: never = chunkType;
+          throw new Error(`Unhandled chunk type: ${unhandledType}`);
+        }
+      }
+    },
+  });
+}