@assistant-ui/react 0.5.92 → 0.5.94

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

@@ -0,0 +1,254 @@
+import { ChatModelRunResult } from "../../local/ChatModelAdapter";
+import { parsePartialJson } from "../partial-json/parse-partial-json";
+import { LanguageModelV1StreamPart } from "@ai-sdk/provider";
+import { ToolResultStreamPart } from "./toolResultStream";
+import { MessageStatus, ToolCallContentPart } from "../../../types";
+
+export function runResultStream() {
+  let message: ChatModelRunResult = {
+    content: [],
+    status: { type: "running" },
+  };
+
+  return new TransformStream<ToolResultStreamPart, ChatModelRunResult>({
+    transform(chunk, controller) {
+      const chunkType = chunk.type;
+      switch (chunkType) {
+        case "text-delta": {
+          message = appendOrUpdateText(message, chunk.textDelta);
+          controller.enqueue(message);
+          break;
+        }
+
+        case "tool-call-delta": {
+          const { toolCallId, toolName, argsTextDelta } = chunk;
+
+          message = appendOrUpdateToolCall(
+            message,
+            toolCallId,
+            toolName,
+            argsTextDelta,
+          );
+          controller.enqueue(message);
+          break;
+        }
+
+        case "tool-call":
+        case "response-metadata":
+          break;
+
+        case "tool-result": {
+          message = appendOrUpdateToolResult(
+            message,
+            chunk.toolCallId,
+            chunk.toolName,
+            chunk.result,
+          );
+          controller.enqueue(message);
+          break;
+        }
+        case "step-finish": {
+          message = appendOrUpdateStepFinish(message, chunk);
+          controller.enqueue(message);
+          break;
+        }
+        case "finish": {
+          message = appendOrUpdateFinish(message, chunk);
+          controller.enqueue(message);
+          break;
+        }
+        case "error": {
+          if (
+            chunk.error instanceof Error &&
+            chunk.error.name === "AbortError"
+          ) {
+            message = appendOrUpdateCancel(message);
+            controller.enqueue(message);
+            break;
+          } else {
+            throw chunk.error;
+          }
+        }
+        default: {
+          const unhandledType: never = chunkType;
+          throw new Error(`Unhandled chunk type: ${unhandledType}`);
+        }
+      }
+    },
+  });
+}
+
+const appendOrUpdateText = (message: ChatModelRunResult, textDelta: string) => {
+  let contentParts = message.content ?? [];
+  let contentPart = message.content?.at(-1);
+  if (contentPart?.type !== "text") {
+    contentPart = { type: "text", text: textDelta };
+  } else {
+    contentParts = contentParts.slice(0, -1);
+    contentPart = { type: "text", text: contentPart.text + textDelta };
+  }
+  return {
+    ...message,
+    content: contentParts.concat([contentPart]),
+  };
+};
+
+const appendOrUpdateToolCall = (
+  message: ChatModelRunResult,
+  toolCallId: string,
+  toolName: string,
+  argsTextDelta: string,
+): ChatModelRunResult => {
+  let contentParts = message.content ?? [];
+  const contentPartIdx = contentParts.findIndex(
+    (c) => c.type === "tool-call" && c.toolCallId === toolCallId,
+  );
+  let contentPart =
+    contentPartIdx === -1
+      ? null
+      : (contentParts[contentPartIdx] as ToolCallContentPart);
+
+  if (contentPart == null) {
+    contentPart = {
+      type: "tool-call",
+      toolCallId,
+      toolName,
+      argsText: argsTextDelta,
+      args: parsePartialJson(argsTextDelta),
+    };
+    contentParts = [...contentParts, contentPart];
+  } else {
+    const argsText = contentPart.argsText + argsTextDelta;
+    contentPart = {
+      ...contentPart,
+      argsText,
+      args: parsePartialJson(argsText),
+    };
+    contentParts = [
+      ...contentParts.slice(0, contentPartIdx),
+      contentPart,
+      ...contentParts.slice(contentPartIdx + 1),
+    ];
+  }
+
+  return {
+    ...message,
+    content: contentParts,
+  };
+};
+
+const appendOrUpdateToolResult = (
+  message: ChatModelRunResult,
+  toolCallId: string,
+  toolName: string,
+  result: any,
+) => {
+  let found = false;
+  const newContentParts = message.content?.map((part) => {
+    if (part.type !== "tool-call" || part.toolCallId !== toolCallId)
+      return part;
+    found = true;
+
+    if (part.toolName !== toolName)
+      throw new Error(
+        `Tool call ${toolCallId} found with tool name ${part.toolName}, but expected ${toolName}`,
+      );
+
+    return {
+      ...part,
+      result,
+    };
+  });
+  if (!found)
+    throw new Error(
+      `Received tool result for unknown tool call "${toolName}" / "${toolCallId}". This is likely an internal bug in assistant-ui.`,
+    );
+
+  return {
+    ...message,
+    content: newContentParts!,
+  };
+};
+
+const appendOrUpdateStepFinish = (
+  message: ChatModelRunResult,
+  chunk: ToolResultStreamPart & { type: "step-finish" },
+): ChatModelRunResult => {
+  const { type, ...rest } = chunk;
+  const steps = [
+    ...(message.metadata?.steps ?? []),
+    {
+      usage: rest.usage,
+    },
+  ];
+  return {
+    ...message,
+    metadata: {
+      ...message.metadata,
+      roundtrips: steps,
+      steps,
+    },
+  };
+};
+
+const appendOrUpdateFinish = (
+  message: ChatModelRunResult,
+  chunk: LanguageModelV1StreamPart & { type: "finish" },
+): ChatModelRunResult => {
+  const { type, ...rest } = chunk;
+
+  const steps = [
+    ...(message.metadata?.steps ?? []),
+    {
+      logprobs: rest.logprobs,
+      usage: rest.usage,
+    },
+  ];
+  return {
+    ...message,
+    status: getStatus(chunk),
+    metadata: {
+      ...message.metadata,
+      roundtrips: steps,
+      steps,
+    },
+  };
+};
+
+const getStatus = (
+  chunk:
+    | (LanguageModelV1StreamPart & { type: "finish" })
+    | (ToolResultStreamPart & { type: "step-finish" }),
+): MessageStatus => {
+  if (chunk.finishReason === "tool-calls") {
+    return {
+      type: "requires-action",
+      reason: "tool-calls",
+    };
+  } else if (
+    chunk.finishReason === "stop" ||
+    chunk.finishReason === "unknown"
+  ) {
+    return {
+      type: "complete",
+      reason: chunk.finishReason,
+    };
+  } else {
+    return {
+      type: "incomplete",
+      reason: chunk.finishReason,
+    };
+  }
+};
+
+const appendOrUpdateCancel = (
+  message: ChatModelRunResult,
+): ChatModelRunResult => {
+  return {
+    ...message,
+    status: {
+      type: "incomplete",
+      reason: "cancelled",
+    },
+  };
+};

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

@@ -0,0 +1,123 @@
+import { Tool } from "../../../types/ModelConfigTypes";
+import { LanguageModelV1StreamPart } from "@ai-sdk/provider";
+import { z } from "zod";
+import sjson from "secure-json-parse";
+
+export type ToolResultStreamPart =
+  | LanguageModelV1StreamPart
+  | {
+      type: "tool-result";
+      toolCallType: "function";
+      toolCallId: string;
+      toolName: string;
+      result: unknown;
+      isError?: boolean;
+    }
+  | {
+      type: "step-finish";
+      finishReason:
+        | "stop"
+        | "length"
+        | "content-filter"
+        | "tool-calls"
+        | "error"
+        | "other"
+        | "unknown";
+      usage: {
+        promptTokens: number;
+        completionTokens: number;
+      };
+      isContinued: boolean;
+    };
+
+export function toolResultStream(
+  tools: Record<string, Tool> | undefined,
+  abortSignal: AbortSignal,
+) {
+  const toolCallExecutions = new Map<string, Promise<any>>();
+
+  return new TransformStream<ToolResultStreamPart, ToolResultStreamPart>({
+    transform(chunk, controller) {
+      // forward everything
+      controller.enqueue(chunk);
+
+      // handle tool calls
+      const chunkType = chunk.type;
+      switch (chunkType) {
+        case "tool-call": {
+          const { toolCallId, toolCallType, toolName, args: argsText } = chunk;
+          const tool = tools?.[toolName];
+          if (!tool || !tool.execute) return;
+
+          const args = sjson.parse(argsText);
+          if (tool.parameters instanceof z.ZodType) {
+            const result = tool.parameters.safeParse(args);
+            if (!result.success) {
+              controller.enqueue({
+                type: "tool-result",
+                toolCallType,
+                toolCallId,
+                toolName,
+                result:
+                  "Function parameter validation failed. " +
+                  JSON.stringify(result.error.issues),
+                isError: true,
+              });
+              return;
+            } else {
+              toolCallExecutions.set(
+                toolCallId,
+                (async () => {
+                  if (!tool.execute) return;
+
+                  try {
+                    const result = await tool.execute(args, { abortSignal });
+
+                    controller.enqueue({
+                      type: "tool-result",
+                      toolCallType,
+                      toolCallId,
+                      toolName,
+                      result,
+                    });
+                  } catch (error) {
+                    controller.enqueue({
+                      type: "tool-result",
+                      toolCallType,
+                      toolCallId,
+                      toolName,
+                      result: "Error: " + error,
+                      isError: true,
+                    });
+                  } finally {
+                    toolCallExecutions.delete(toolCallId);
+                  }
+                })(),
+              );
+            }
+          }
+          break;
+        }
+
+        // ignore other parts
+        case "text-delta":
+        case "tool-call-delta":
+        case "tool-result":
+        case "step-finish":
+        case "finish":
+        case "error":
+        case "response-metadata":
+          break;
+
+        default: {
+          const unhandledType: never = chunkType;
+          throw new Error(`Unhandled chunk type: ${unhandledType}`);
+        }
+      }
+    },
+
+    async flush() {
+      await Promise.all(toolCallExecutions.values());
+    },
+  });
+}

package/src/runtimes/edge/streams/utils/PipeableTransformStream.ts

@@ -0,0 +1,10 @@
+export class PipeableTransformStream<I, O> extends TransformStream<I, O> {
+  constructor(transform: (readable: ReadableStream<I>) => ReadableStream<O>) {
+    super();
+    const readable = transform(super.readable as any);
+    Object.defineProperty(this, "readable", {
+      value: readable,
+      writable: false,
+    });
+  }
+}

package/src/runtimes/edge/streams/utils/StreamPart.ts

@@ -0,0 +1,3 @@
+export type StreamPart<T extends Record<string, unknown>> = {
+  [K in keyof T]: { type: K; value: T[K] };
+}[keyof T];

package/src/runtimes/edge/streams/utils/chunkByLineStream.ts

@@ -0,0 +1,24 @@
+export function chunkByLineStream() {
+  let buffer = "";
+
+  return new TransformStream({
+    transform(chunk, controller) {
+      buffer += chunk;
+      const lines = buffer.split("\n");
+
+      // Process all complete lines
+      for (let i = 0; i < lines.length - 1; i++) {
+        controller.enqueue(lines[i]);
+      }
+
+      // Keep the last incomplete line in the buffer
+      buffer = lines[lines.length - 1]!;
+    },
+    flush(controller) {
+      // flush any remaining content in the buffer
+      if (buffer) {
+        controller.enqueue(buffer);
+      }
+    },
+  });
+}

package/src/runtimes/edge/streams/utils/index.ts

@@ -0,0 +1,15 @@
+import { streamPartDecoderStream } from "./streamPartDecoderStream";
+import { streamPartEncoderStream } from "./streamPartEncoderStream";
+import { StreamPart } from "./StreamPart";
+
+export declare namespace StreamUtils {
+  export { StreamPart };
+}
+
+/**
+ * @deprecated `streamUtils` will be replaced with `assistant-stream` once it is ready.
+ */
+export const streamUtils = {
+  streamPartEncoderStream,
+  streamPartDecoderStream,
+};

package/src/runtimes/edge/streams/utils/streamPartDecoderStream.ts

@@ -0,0 +1,29 @@
+import { chunkByLineStream } from "./chunkByLineStream";
+import { PipeableTransformStream } from "./PipeableTransformStream";
+import { StreamPart } from "./StreamPart";
+
+const decodeStreamPart = <T extends Record<string, unknown>>(
+  part: string,
+): StreamPart<T> => {
+  const index = part.indexOf(":");
+  if (index === -1) throw new Error("Invalid stream part");
+  return {
+    type: part.slice(0, index),
+    value: JSON.parse(part.slice(index + 1)),
+  };
+};
+
+export function streamPartDecoderStream<T extends Record<string, unknown>>() {
+  const decodeStream = new TransformStream<string, StreamPart<T>>({
+    transform(chunk, controller) {
+      controller.enqueue(decodeStreamPart<T>(chunk));
+    },
+  });
+
+  return new PipeableTransformStream((readable) => {
+    return readable
+      .pipeThrough(new TextDecoderStream())
+      .pipeThrough(chunkByLineStream())
+      .pipeThrough(decodeStream);
+  });
+}

package/src/runtimes/edge/streams/utils/streamPartEncoderStream.ts

@@ -0,0 +1,23 @@
+import { PipeableTransformStream } from "./PipeableTransformStream";
+import { StreamPart } from "./StreamPart";
+
+function encodeStreamPart<T extends Record<string, unknown>>({
+  type,
+  value,
+}: StreamPart<T>): string {
+  return `${type as string}:${JSON.stringify(value)}\n`;
+}
+
+export function streamPartEncoderStream<T extends Record<string, unknown>>() {
+  const encodeStream = new TransformStream<StreamPart<T>, string>({
+    transform(chunk, controller) {
+      controller.enqueue(encodeStreamPart<T>(chunk));
+    },
+  });
+
+  return new PipeableTransformStream((readable) => {
+    return readable
+      .pipeThrough(encodeStream)
+      .pipeThrough(new TextEncoderStream());
+  });
+}

package/src/runtimes/edge/useEdgeRuntime.ts

@@ -0,0 +1,15 @@
+import { LocalRuntimeOptions, useLocalRuntime } from "..";
+import { EdgeChatAdapterOptions, EdgeChatAdapter } from "./EdgeChatAdapter";
+import { splitLocalRuntimeOptions } from "../local/LocalRuntimeOptions";
+
+export type EdgeRuntimeOptions = EdgeChatAdapterOptions & LocalRuntimeOptions;
+
+export const useEdgeRuntime = (options: EdgeRuntimeOptions) => {
+  const { localRuntimeOptions, otherOptions } =
+    splitLocalRuntimeOptions(options);
+
+  return useLocalRuntime(
+    new EdgeChatAdapter(otherOptions),
+    localRuntimeOptions,
+  );
+};

package/src/runtimes/external-store/ExternalStoreAdapter.tsx

@@ -0,0 +1,85 @@
+import { AppendMessage, ThreadMessage } from "../../types";
+import { AttachmentAdapter } from "../attachment";
+import {
+  AddToolResultOptions,
+  ThreadSuggestion,
+} from "../core/ThreadRuntimeCore";
+import { FeedbackAdapter } from "../feedback/FeedbackAdapter";
+import { SpeechSynthesisAdapter } from "../speech/SpeechAdapterTypes";
+import { ThreadManagerMetadata } from "../core/ThreadManagerRuntimeCore";
+import { ThreadMessageLike } from "./ThreadMessageLike";
+
+export type ExternalStoreThreadManagerAdapter = {
+  threadId?: string | undefined;
+  threads?: readonly ThreadManagerMetadata[] | undefined;
+  archivedThreads?: readonly ThreadManagerMetadata[] | undefined;
+  onSwitchToNewThread?: (() => Promise<void> | void) | undefined;
+  onSwitchToThread?: ((threadId: string) => Promise<void> | void) | undefined;
+  onRename?: (
+    threadId: string,
+    newTitle: string,
+  ) => (Promise<void> | void) | undefined;
+  onArchive?: ((threadId: string) => Promise<void> | void) | undefined;
+  onUnarchive?: ((threadId: string) => Promise<void> | void) | undefined;
+  onDelete?: ((threadId: string) => Promise<void> | void) | undefined;
+};
+
+export type ExternalStoreMessageConverter<T> = (
+  message: T,
+  idx: number,
+) => ThreadMessageLike;
+
+type ExternalStoreMessageConverterAdapter<T> = {
+  convertMessage: ExternalStoreMessageConverter<T>;
+};
+
+type ExternalStoreAdapterBase<T> = {
+  /**
+   * @deprecated Use `adapters.threadManager.threadId` instead. This will be removed in 0.6.0.
+   */
+  threadId?: string | undefined;
+
+  /**
+   * @deprecated Use `adapters.threadManager.onSwitchToThread` instead. This will be removed in 0.6.0.
+   */
+  onSwitchToThread?: ((threadId: string) => Promise<void> | void) | undefined;
+  /**
+   * @deprecated Use `adapters.threadManager.onSwitchToNewThread` instead. This will be removed in 0.6.0.
+   */
+  onSwitchToNewThread?: (() => Promise<void> | void) | undefined;
+
+  isDisabled?: boolean | undefined;
+  isRunning?: boolean | undefined;
+  messages: T[];
+  suggestions?: readonly ThreadSuggestion[] | undefined;
+  extras?: unknown;
+
+  setMessages?: ((messages: T[]) => void) | undefined;
+  onNew: (message: AppendMessage) => Promise<void>;
+  onEdit?: ((message: AppendMessage) => Promise<void>) | undefined;
+  onReload?: ((parentId: string | null) => Promise<void>) | undefined;
+  onCancel?: (() => Promise<void>) | undefined;
+  onAddToolResult?:
+    | ((options: AddToolResultOptions) => Promise<void> | void)
+    | undefined;
+  convertMessage?: ExternalStoreMessageConverter<T> | undefined;
+  adapters?:
+    | {
+        attachments?: AttachmentAdapter | undefined;
+        speech?: SpeechSynthesisAdapter | undefined;
+        feedback?: FeedbackAdapter | undefined;
+        threadManager?: ExternalStoreThreadManagerAdapter | undefined;
+      }
+    | undefined;
+  unstable_capabilities?:
+    | {
+        copy?: boolean | undefined;
+      }
+    | undefined;
+};
+
+export type ExternalStoreAdapter<T = ThreadMessage> =
+  ExternalStoreAdapterBase<T> &
+    (T extends ThreadMessage
+      ? object
+      : ExternalStoreMessageConverterAdapter<T>);

package/src/runtimes/external-store/ExternalStoreRuntimeCore.tsx

@@ -0,0 +1,39 @@
+import { BaseAssistantRuntimeCore } from "../../internal";
+import { ExternalStoreThreadManagerRuntimeCore } from "./ExternalStoreThreadManagementAdapter";
+import { ExternalStoreAdapter } from "./ExternalStoreAdapter";
+import { ExternalStoreThreadRuntimeCore } from "./ExternalStoreThreadRuntimeCore";
+
+const getThreadManagerAdapter = (store: ExternalStoreAdapter<any>) => {
+  return {
+    threadId: store.threadId,
+    onSwitchToNewThread: store.onSwitchToNewThread,
+    onSwitchToThread: store.onSwitchToThread,
+    ...store.adapters?.threadManager,
+  };
+};
+
+export class ExternalStoreRuntimeCore extends BaseAssistantRuntimeCore {
+  public readonly threadManager;
+
+  private _store: ExternalStoreAdapter<any>;
+
+  constructor(store: ExternalStoreAdapter<any>) {
+    super();
+    this._store = store;
+    this.threadManager = new ExternalStoreThreadManagerRuntimeCore(
+      getThreadManagerAdapter(store),
+      (threadId) =>
+        new ExternalStoreThreadRuntimeCore(
+          this._proxyConfigProvider,
+          threadId,
+          this._store,
+        ),
+    );
+  }
+
+  public setStore(store: ExternalStoreAdapter<any>) {
+    this._store = store;
+    this.threadManager.setAdapter(getThreadManagerAdapter(store));
+    this.threadManager.mainThread.setStore(store);
+  }
+}