@assistant-ui/react 0.5.92 → 0.5.94

package/src/runtimes/external-store/external-message-converter.tsx

@@ -0,0 +1,205 @@
+import { useMemo } from "react";
+import { ThreadMessageConverter } from "./ThreadMessageConverter";
+import {
+  getExternalStoreMessage,
+  symbolInnerMessage,
+} from "./getExternalStoreMessage";
+import { fromThreadMessageLike, ThreadMessageLike } from "./ThreadMessageLike";
+import { getAutoStatus, isAutoStatus } from "./auto-status";
+import { ToolCallContentPart } from "../../types";
+
+export namespace useExternalMessageConverter {
+  export type Message =
+    | ThreadMessageLike
+    | {
+        role: "tool";
+        toolCallId: string;
+        toolName?: string | undefined;
+        result: any;
+      };
+
+  export type Callback<T> = (message: T) => Message | Message[];
+}
+
+type CallbackResult<T> = {
+  input: T;
+  outputs: useExternalMessageConverter.Message[];
+};
+
+type ChunkResult<T> = {
+  inputs: T[];
+  outputs: useExternalMessageConverter.Message[];
+};
+
+const joinExternalMessages = (
+  messages: useExternalMessageConverter.Message[],
+): ThreadMessageLike => {
+  const assistantMessage: ThreadMessageLike & { content: any[] } = {
+    role: "assistant",
+    content: [],
+  };
+  for (const output of messages) {
+    if (output.role === "tool") {
+      const toolCallIdx = assistantMessage.content.findIndex(
+        (c) => c.type === "tool-call" && c.toolCallId === output.toolCallId,
+      );
+      if (toolCallIdx !== -1) {
+        const toolCall = assistantMessage.content[
+          toolCallIdx
+        ]! as ToolCallContentPart;
+        if (output.toolName) {
+          if (toolCall.toolName !== output.toolName)
+            throw new Error(
+              `Tool call name ${output.toolCallId} ${output.toolName} does not match existing tool call ${toolCall.toolName}`,
+            );
+        }
+        assistantMessage.content[toolCallIdx] = {
+          ...toolCall,
+          result: output.result,
+        };
+      } else {
+        throw new Error(
+          `Tool call ${output.toolCallId} ${output.toolName} not found in assistant message`,
+        );
+      }
+    } else {
+      const role = output.role;
+      switch (role) {
+        case "system":
+        case "user":
+          return output;
+        case "assistant":
+          if (assistantMessage.content.length === 0) {
+            assistantMessage.id = output.id;
+            assistantMessage.createdAt ??= output.createdAt;
+            assistantMessage.status ??= output.status;
+            // TODO keep this in sync
+          }
+          assistantMessage.content.push(...output.content);
+          break;
+        default: {
+          const unsupportedRole: never = role;
+          throw new Error(`Unknown message role: ${unsupportedRole}`);
+        }
+      }
+    }
+  }
+  return assistantMessage;
+};
+
+const chunkExternalMessages = <T,>(callbackResults: CallbackResult<T>[]) => {
+  const results: ChunkResult<T>[] = [];
+  let isAssistant = false;
+  let inputs: T[] = [];
+  let outputs: useExternalMessageConverter.Message[] = [];
+
+  const flush = () => {
+    if (outputs.length) {
+      results.push({
+        inputs,
+        outputs,
+      });
+    }
+    inputs = [];
+    outputs = [];
+  };
+
+  for (const callbackResult of callbackResults) {
+    for (const output of callbackResult.outputs) {
+      if (!isAssistant || output.role === "user" || output.role === "system") {
+        flush();
+      }
+      isAssistant = output.role === "assistant" || output.role === "tool";
+
+      if (inputs.at(-1) !== callbackResult.input) {
+        inputs.push(callbackResult.input);
+      }
+      outputs.push(output);
+    }
+  }
+  flush();
+  return results;
+};
+
+export const useExternalMessageConverter = <T extends WeakKey>({
+  callback,
+  messages,
+  isRunning,
+}: {
+  callback: useExternalMessageConverter.Callback<T>;
+  messages: T[];
+  isRunning: boolean;
+}) => {
+  const state = useMemo(
+    () => ({
+      callback,
+      callbackCache: new WeakMap<T, CallbackResult<T>>(),
+      chunkCache: new WeakMap<
+        useExternalMessageConverter.Message,
+        ChunkResult<T>
+      >(),
+      converterCache: new ThreadMessageConverter(),
+    }),
+    [callback],
+  );
+
+  return useMemo(() => {
+    const callbackResults: CallbackResult<T>[] = [];
+    for (const message of messages) {
+      let result = state.callbackCache.get(message);
+      if (!result) {
+        const output = state.callback(message);
+        const outputs = Array.isArray(output) ? output : [output];
+        result = { input: message, outputs };
+        state.callbackCache.set(message, result);
+      }
+      callbackResults.push(result);
+    }
+
+    const chunks = chunkExternalMessages(callbackResults).map((m) => {
+      const key = m.outputs[0];
+      if (!key) return m;
+
+      const cached = state.chunkCache.get(key);
+      if (cached && shallowArrayEqual(cached.outputs, m.outputs)) return cached;
+      state.chunkCache.set(key, m);
+      return m;
+    });
+
+    return state.converterCache.convertMessages(
+      chunks,
+      (cache, message, idx) => {
+        const isLast = idx === chunks.length - 1;
+        const autoStatus = getAutoStatus(isLast, isRunning);
+
+        if (
+          cache &&
+          (cache.role !== "assistant" ||
+            !isAutoStatus(cache.status) ||
+            cache.status === autoStatus)
+        ) {
+          const inputs = getExternalStoreMessage(cache) as T[];
+          if (shallowArrayEqual(inputs, message.inputs)) {
+            return cache;
+          }
+        }
+
+        const newMessage = fromThreadMessageLike(
+          joinExternalMessages(message.outputs),
+          idx.toString(),
+          autoStatus,
+        );
+        (newMessage as any)[symbolInnerMessage] = message.inputs;
+        return newMessage;
+      },
+    );
+  }, [state, messages, isRunning]);
+};
+
+const shallowArrayEqual = (a: unknown[], b: unknown[]) => {
+  if (a.length !== b.length) return false;
+  for (let i = 0; i < a.length; i++) {
+    if (a[i] !== b[i]) return false;
+  }
+  return true;
+};

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

@@ -0,0 +1,11 @@
+import { ThreadMessage } from "../../types";
+
+export const symbolInnerMessage = Symbol("innerMessage");
+
+export type ExternalStoreThreadMessage<T> = ThreadMessage & {
+  [symbolInnerMessage]?: T;
+};
+
+export const getExternalStoreMessage = <T,>(message: ThreadMessage) => {
+  return (message as ExternalStoreThreadMessage<T>)[symbolInnerMessage];
+};

package/src/runtimes/external-store/index.ts

@@ -0,0 +1,8 @@
+export type {
+  ExternalStoreAdapter,
+  ExternalStoreMessageConverter,
+} from "./ExternalStoreAdapter";
+export type { ThreadMessageLike } from "./ThreadMessageLike";
+export { useExternalStoreRuntime } from "./useExternalStoreRuntime";
+export { getExternalStoreMessage } from "./getExternalStoreMessage";
+export { useExternalMessageConverter } from "./external-message-converter";

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

@@ -0,0 +1,18 @@
+import { useEffect, useMemo, useState } from "react";
+import { ExternalStoreRuntimeCore } from "./ExternalStoreRuntimeCore";
+import { ExternalStoreAdapter } from "./ExternalStoreAdapter";
+import { AssistantRuntimeImpl } from "../../api/AssistantRuntime";
+import { ThreadRuntimeImpl } from "../../api/ThreadRuntime";
+
+export const useExternalStoreRuntime = <T,>(store: ExternalStoreAdapter<T>) => {
+  const [runtime] = useState(() => new ExternalStoreRuntimeCore(store));
+
+  useEffect(() => {
+    runtime.setStore(store);
+  });
+
+  return useMemo(
+    () => AssistantRuntimeImpl.create(runtime, ThreadRuntimeImpl),
+    [runtime],
+  );
+};

package/src/runtimes/feedback/FeedbackAdapter.ts

@@ -0,0 +1,10 @@
+import { ThreadMessage } from "../../types/AssistantTypes";
+
+type FeedbackAdapterFeedback = {
+  message: ThreadMessage;
+  type: "positive" | "negative";
+};
+
+export type FeedbackAdapter = {
+  submit: (feedback: FeedbackAdapterFeedback) => void;
+};

package/src/runtimes/feedback/index.ts

@@ -0,0 +1 @@
+export { type FeedbackAdapter } from "./FeedbackAdapter";

package/src/runtimes/index.ts

@@ -0,0 +1,8 @@
+export * from "./core";
+export * from "./local";
+export * from "./edge";
+export * from "./external-store";
+export * from "./dangerous-in-browser";
+export * from "./speech";
+export * from "./attachment";
+export * from "./feedback";

package/src/runtimes/local/ChatModelAdapter.tsx

@@ -0,0 +1,46 @@
+"use client";
+
+import type {
+  MessageStatus,
+  ThreadAssistantContentPart,
+  ThreadMessage,
+  ThreadStep,
+} from "../../types/AssistantTypes";
+import type { ModelConfig } from "../../types/ModelConfigTypes";
+
+export type ChatModelRunUpdate = {
+  content: ThreadAssistantContentPart[];
+  metadata?: Record<string, unknown>;
+};
+
+export type ChatModelRunResult = {
+  content?: ThreadAssistantContentPart[];
+  status?: MessageStatus;
+  metadata?: {
+    /**
+     * @internal Use `steps` instead. This field will be removed in v0.6.
+     */
+    roundtrips?: ThreadStep[];
+    steps?: ThreadStep[];
+    custom?: Record<string, unknown>;
+  };
+};
+
+export type ChatModelRunOptions = {
+  messages: ThreadMessage[];
+  abortSignal: AbortSignal;
+  config: ModelConfig;
+
+  unstable_assistantMessageId?: string;
+
+  /**
+   * @deprecated Declare the run function as an AsyncGenerator instead. This method will be removed in v0.6
+   */
+  onUpdate: (result: ChatModelRunUpdate) => void;
+};
+
+export type ChatModelAdapter = {
+  run: (
+    options: ChatModelRunOptions,
+  ) => Promise<ChatModelRunResult> | AsyncGenerator<ChatModelRunResult, void>;
+};

package/src/runtimes/local/LocalRuntimeCore.tsx

@@ -0,0 +1,69 @@
+import type { CoreMessage } from "../../types/AssistantTypes";
+import { BaseAssistantRuntimeCore } from "../core/BaseAssistantRuntimeCore";
+import { LocalThreadRuntimeCore } from "./LocalThreadRuntimeCore";
+import { LocalRuntimeOptionsBase } from "./LocalRuntimeOptions";
+import { fromCoreMessages } from "../edge/converters/fromCoreMessage";
+import { LocalThreadManagerRuntimeCore } from "./LocalThreadManagerRuntimeCore";
+import { ExportedMessageRepository } from "../utils/MessageRepository";
+
+const getExportFromInitialMessages = (
+  initialMessages: readonly CoreMessage[],
+): ExportedMessageRepository => {
+  const messages = fromCoreMessages(initialMessages);
+  return {
+    messages: messages.map((m, idx) => ({
+      parentId: messages[idx - 1]?.id ?? null,
+      message: m,
+    })),
+  };
+};
+
+export class LocalRuntimeCore extends BaseAssistantRuntimeCore {
+  public readonly threadManager;
+
+  private _options: LocalRuntimeOptionsBase;
+
+  constructor(
+    options: LocalRuntimeOptionsBase,
+    initialMessages?: CoreMessage[],
+  ) {
+    super();
+
+    this._options = options;
+
+    this.threadManager = new LocalThreadManagerRuntimeCore((threadId, data) => {
+      const thread = new LocalThreadRuntimeCore(
+        this._proxyConfigProvider,
+        threadId,
+        this._options,
+      );
+      thread.import(data);
+      return thread;
+    });
+
+    if (initialMessages) {
+      this.threadManager.mainThread.import(
+        getExportFromInitialMessages(initialMessages),
+      );
+    }
+  }
+
+  public setOptions(options: LocalRuntimeOptionsBase) {
+    this._options = options;
+
+    this.threadManager.mainThread.setOptions(options);
+  }
+
+  public reset({
+    initialMessages,
+  }: {
+    initialMessages?: readonly CoreMessage[] | undefined;
+  } = {}) {
+    this.threadManager.switchToNewThread();
+    if (!initialMessages) return;
+
+    this.threadManager.mainThread.import(
+      getExportFromInitialMessages(initialMessages),
+    );
+  }
+}

package/src/runtimes/local/LocalRuntimeOptions.tsx

@@ -0,0 +1,42 @@
+import type { CoreMessage } from "../../types";
+import { AttachmentAdapter } from "../attachment/AttachmentAdapter";
+import { FeedbackAdapter } from "../feedback/FeedbackAdapter";
+import { SpeechSynthesisAdapter } from "../speech/SpeechAdapterTypes";
+import { ChatModelAdapter } from "./ChatModelAdapter";
+
+export type LocalRuntimeOptionsBase = {
+  maxSteps?: number | undefined;
+  /**
+   * @deprecated Use `maxSteps` (which is `maxToolRoundtrips` + 1; if you set `maxToolRoundtrips` to 2, set `maxSteps` to 3) instead. This field will be removed in v0.6.
+   */
+  maxToolRoundtrips?: number | undefined;
+  adapters: {
+    chatModel: ChatModelAdapter;
+    attachments?: AttachmentAdapter | undefined;
+    speech?: SpeechSynthesisAdapter | undefined;
+    feedback?: FeedbackAdapter | undefined;
+  };
+};
+
+// TODO align LocalRuntimeOptions with LocalRuntimeOptionsBase
+export type LocalRuntimeOptions = Omit<LocalRuntimeOptionsBase, "adapters"> & {
+  initialMessages?: readonly CoreMessage[] | undefined;
+  adapters?: Omit<LocalRuntimeOptionsBase["adapters"], "chatModel"> | undefined;
+};
+
+export const splitLocalRuntimeOptions = <T extends LocalRuntimeOptions>(
+  options: T,
+) => {
+  const { initialMessages, maxToolRoundtrips, maxSteps, adapters, ...rest } =
+    options;
+
+  return {
+    localRuntimeOptions: {
+      initialMessages,
+      maxToolRoundtrips,
+      maxSteps,
+      adapters,
+    },
+    otherOptions: rest,
+  };
+};

package/src/runtimes/local/LocalThreadManagerRuntimeCore.tsx

@@ -0,0 +1,171 @@
+import type { Unsubscribe } from "../../types";
+import {
+  ThreadManagerMetadata,
+  ThreadManagerRuntimeCore,
+} from "../core/ThreadManagerRuntimeCore";
+import { ExportedMessageRepository } from "../utils/MessageRepository";
+import { generateId } from "../../utils/idUtils";
+import { LocalThreadRuntimeCore } from "./LocalThreadRuntimeCore";
+
+export type LocalThreadData = {
+  data: ExportedMessageRepository;
+  metadata: ThreadManagerMetadata;
+  isArchived: boolean;
+};
+
+export type LocalThreadFactory = (
+  threadId: string,
+  data: ExportedMessageRepository,
+) => LocalThreadRuntimeCore;
+
+export class LocalThreadManagerRuntimeCore implements ThreadManagerRuntimeCore {
+  private _threadData = new Map<string, LocalThreadData>();
+
+  private _threads: readonly ThreadManagerMetadata[] = [];
+  private _archivedThreads: readonly ThreadManagerMetadata[] = [];
+
+  public get threads() {
+    return this._threads;
+  }
+
+  public get archivedThreads() {
+    return this._archivedThreads;
+  }
+
+  private _mainThread: LocalThreadRuntimeCore;
+
+  public get mainThread(): LocalThreadRuntimeCore {
+    return this._mainThread;
+  }
+
+  constructor(private _threadFactory: LocalThreadFactory) {
+    const threadId = generateId();
+    this._threadData.set(threadId, {
+      data: { messages: [] },
+      metadata: { threadId },
+      isArchived: false,
+    });
+    this._threads = [{ threadId }];
+    this._mainThread = this._threadFactory(threadId, { messages: [] });
+  }
+
+  public switchToThread(threadId: string): void {
+    if (this._mainThread.threadId === threadId) return;
+
+    const data = this._threadData.get(threadId);
+    if (!data) throw new Error("Thread not found");
+
+    const thread = this._threadFactory(threadId, data.data);
+    this._performThreadSwitch(thread);
+  }
+
+  public switchToNewThread(): void {
+    const threadId = generateId();
+    this._threadData.set(threadId, {
+      data: { messages: [] },
+      metadata: { threadId },
+      isArchived: false,
+    });
+    this._threads = [{ threadId }];
+    const thread = this._threadFactory(threadId, { messages: [] });
+    this._performThreadSwitch(thread);
+  }
+
+  private _performThreadSwitch(newThreadCore: LocalThreadRuntimeCore) {
+    const data = this._threadData.get(this._mainThread.threadId);
+    if (!data) throw new Error("Thread not found");
+
+    const exprt = this._mainThread.export();
+    data.data = exprt;
+
+    this._mainThread._notifyEventSubscribers("switched-away");
+    this._mainThread = newThreadCore;
+    newThreadCore._notifyEventSubscribers("switched-to");
+
+    this._notifySubscribers();
+  }
+
+  private _performMoveOp(
+    threadId: string,
+    operation: "archive" | "unarchive" | "delete",
+  ) {
+    const data = this._threadData.get(threadId);
+    if (!data) throw new Error("Thread not found");
+
+    if (operation === "archive" && data.isArchived) return;
+    if (operation === "unarchive" && !data.isArchived) return;
+
+    if (operation === "archive") {
+      data.isArchived = true;
+      this._archivedThreads = [...this._archivedThreads, data.metadata];
+    }
+    if (operation === "unarchive") {
+      data.isArchived = false;
+      this._threads = [...this._threads, data.metadata];
+    }
+    if (operation === "delete") {
+      this._threadData.delete(threadId);
+    }
+
+    if (
+      operation === "archive" ||
+      (operation === "delete" && data.isArchived)
+    ) {
+      this._archivedThreads = this._archivedThreads.filter(
+        (t) => t.threadId !== threadId,
+      );
+    }
+
+    if (
+      operation === "unarchive" ||
+      (operation === "delete" && !data.isArchived)
+    ) {
+      this._threads = this._threads.filter((t) => t.threadId !== threadId);
+    }
+
+    this._notifySubscribers();
+  }
+
+  public async rename(threadId: string, newTitle: string): Promise<void> {
+    const data = this._threadData.get(threadId);
+    if (!data) throw new Error("Thread not found");
+
+    data.metadata = {
+      ...data.metadata,
+      title: newTitle,
+    };
+
+    const threadList = data.isArchived ? this.archivedThreads : this.threads;
+    const idx = threadList.findIndex((t) => t.threadId === threadId);
+    const updatedThreadList = threadList.toSpliced(idx, 1, data.metadata);
+    if (data.isArchived) {
+      this._archivedThreads = updatedThreadList;
+    } else {
+      this._threads = updatedThreadList;
+    }
+    this._notifySubscribers();
+  }
+
+  public async archive(threadId: string): Promise<void> {
+    this._performMoveOp(threadId, "archive");
+  }
+
+  public async unarchive(threadId: string): Promise<void> {
+    this._performMoveOp(threadId, "unarchive");
+  }
+
+  public async delete(threadId: string): Promise<void> {
+    this._performMoveOp(threadId, "delete");
+  }
+
+  private _subscriptions = new Set<() => void>();
+
+  public subscribe(callback: () => void): Unsubscribe {
+    this._subscriptions.add(callback);
+    return () => this._subscriptions.delete(callback);
+  }
+
+  private _notifySubscribers() {
+    for (const callback of this._subscriptions) callback();
+  }
+}