@assistant-ui/react 0.5.91 → 0.5.93

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

@@ -0,0 +1,124 @@
+import type { Unsubscribe } from "../../types";
+import { ExternalStoreThreadRuntimeCore } from "./ExternalStoreThreadRuntimeCore";
+import { ThreadManagerRuntimeCore } from "../core/ThreadManagerRuntimeCore";
+import { ExternalStoreThreadManagerAdapter } from "./ExternalStoreAdapter";
+
+export type ExternalStoreThreadFactory = (
+  threadId: string,
+) => ExternalStoreThreadRuntimeCore;
+
+const EMPTY_ARRAY = Object.freeze([]);
+const DEFAULT_THREAD_ID = "DEFAULT_THREAD_ID";
+
+export class ExternalStoreThreadManagerRuntimeCore
+  implements ThreadManagerRuntimeCore
+{
+  public get threads() {
+    return this.adapter.threads ?? EMPTY_ARRAY;
+  }
+
+  public get archivedThreads() {
+    return this.adapter.archivedThreads ?? EMPTY_ARRAY;
+  }
+
+  private _mainThread: ExternalStoreThreadRuntimeCore;
+
+  public get mainThread() {
+    return this._mainThread;
+  }
+
+  constructor(
+    private adapter: ExternalStoreThreadManagerAdapter = {},
+    private threadFactory: ExternalStoreThreadFactory,
+  ) {
+    this._mainThread = this.threadFactory(DEFAULT_THREAD_ID);
+  }
+
+  public setAdapter(adapter: ExternalStoreThreadManagerAdapter) {
+    const previousAdapter = this.adapter;
+    this.adapter = adapter;
+
+    const newThreadId = adapter.threadId ?? DEFAULT_THREAD_ID;
+    const newThreads = adapter.threads ?? EMPTY_ARRAY;
+    const newArchivedThreads = adapter.archivedThreads ?? EMPTY_ARRAY;
+
+    if (
+      previousAdapter.threadId === newThreadId &&
+      previousAdapter.threads === newThreads &&
+      previousAdapter.archivedThreads === newArchivedThreads
+    ) {
+      return;
+    }
+
+    if (previousAdapter.threadId !== newThreadId) {
+      this._mainThread._notifyEventSubscribers("switched-away");
+      this._mainThread = this.threadFactory(newThreadId);
+      this._mainThread._notifyEventSubscribers("switched-to");
+    }
+
+    this._notifySubscribers();
+  }
+
+  public switchToThread(threadId: string): void {
+    if (this._mainThread?.threadId === threadId) return;
+    const onSwitchToThread = this.adapter.onSwitchToThread;
+    if (!onSwitchToThread)
+      throw new Error(
+        "External store adapter does not support switching to thread",
+      );
+    onSwitchToThread(threadId);
+  }
+
+  public switchToNewThread(): void {
+    const onSwitchToNewThread = this.adapter.onSwitchToNewThread;
+    if (!onSwitchToNewThread)
+      throw new Error(
+        "External store adapter does not support switching to new thread",
+      );
+
+    onSwitchToNewThread();
+  }
+
+  public async rename(threadId: string, newTitle: string): Promise<void> {
+    const onRename = this.adapter.onRename;
+    if (!onRename)
+      throw new Error("External store adapter does not support renaming");
+
+    onRename(threadId, newTitle);
+  }
+
+  public async archive(threadId: string): Promise<void> {
+    const onArchive = this.adapter.onArchive;
+    if (!onArchive)
+      throw new Error("External store adapter does not support archiving");
+
+    onArchive(threadId);
+  }
+
+  public async unarchive(threadId: string): Promise<void> {
+    const onUnarchive = this.adapter.onUnarchive;
+    if (!onUnarchive)
+      throw new Error("External store adapter does not support unarchiving");
+
+    onUnarchive(threadId);
+  }
+
+  public async delete(threadId: string): Promise<void> {
+    const onDelete = this.adapter.onDelete;
+    if (!onDelete)
+      throw new Error("External store adapter does not support deleting");
+
+    onDelete(threadId);
+  }
+
+  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();
+  }
+}

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

@@ -0,0 +1,246 @@
+import {
+  AddToolResultOptions,
+  ThreadSuggestion,
+} from "../core/ThreadRuntimeCore";
+
+import { AppendMessage, ModelConfigProvider, ThreadMessage } from "../../types";
+import { ExternalStoreAdapter } from "./ExternalStoreAdapter";
+import {
+  getExternalStoreMessage,
+  symbolInnerMessage,
+} from "./getExternalStoreMessage";
+import { ThreadMessageConverter } from "./ThreadMessageConverter";
+import { getAutoStatus, isAutoStatus } from "./auto-status";
+import { fromThreadMessageLike } from "./ThreadMessageLike";
+import { getThreadMessageText } from "../../utils/getThreadMessageText";
+import {
+  RuntimeCapabilities,
+  ThreadRuntimeCore,
+} from "../core/ThreadRuntimeCore";
+import { BaseThreadRuntimeCore } from "../core/BaseThreadRuntimeCore";
+
+const EMPTY_ARRAY = Object.freeze([]);
+
+export const hasUpcomingMessage = (
+  isRunning: boolean,
+  messages: ThreadMessage[],
+) => {
+  return isRunning && messages[messages.length - 1]?.role !== "assistant";
+};
+
+export class ExternalStoreThreadRuntimeCore
+  extends BaseThreadRuntimeCore
+  implements ThreadRuntimeCore
+{
+  private assistantOptimisticId: string | null = null;
+
+  private _capabilities: RuntimeCapabilities = {
+    switchToBranch: false,
+    edit: false,
+    reload: false,
+    cancel: false,
+    unstable_copy: false,
+    speech: false,
+    attachments: false,
+    feedback: false,
+  };
+
+  public get capabilities() {
+    return this._capabilities;
+  }
+
+  public threadId!: string;
+  private _messages!: ThreadMessage[];
+  public isDisabled!: boolean;
+
+  public override get messages() {
+    return this._messages;
+  }
+
+  public get adapters() {
+    return this._store.adapters;
+  }
+
+  public suggestions: readonly ThreadSuggestion[] = [];
+  public extras: unknown = undefined;
+
+  private _converter = new ThreadMessageConverter();
+
+  private _store!: ExternalStoreAdapter<any>;
+
+  public override beginEdit(messageId: string) {
+    if (!this._store.onEdit)
+      throw new Error("Runtime does not support editing.");
+
+    super.beginEdit(messageId);
+  }
+
+  constructor(
+    configProvider: ModelConfigProvider,
+    threadId: string,
+    store: ExternalStoreAdapter<any>,
+  ) {
+    super(configProvider);
+    this.threadId = threadId;
+    this.setStore(store);
+  }
+
+  public setStore(store: ExternalStoreAdapter<any>) {
+    if (this._store === store) return;
+
+    const isRunning = store.isRunning ?? false;
+    this.isDisabled = store.isDisabled ?? false;
+
+    const oldStore = this._store as ExternalStoreAdapter<any> | undefined;
+    this._store = store;
+    this.extras = store.extras;
+    this.suggestions = store.suggestions ?? EMPTY_ARRAY;
+    this._capabilities = {
+      switchToBranch: this._store.setMessages !== undefined,
+      edit: this._store.onEdit !== undefined,
+      reload: this._store.onReload !== undefined,
+      cancel: this._store.onCancel !== undefined,
+      speech: this._store.adapters?.speech !== undefined,
+      unstable_copy: this._store.unstable_capabilities?.copy !== false, // default true
+      attachments: !!this._store.adapters?.attachments,
+      feedback: !!this._store.adapters?.feedback,
+    };
+
+    if (oldStore) {
+      // flush the converter cache when the convertMessage prop changes
+      if (oldStore.convertMessage !== store.convertMessage) {
+        this._converter = new ThreadMessageConverter();
+      } else if (
+        oldStore.isRunning === store.isRunning &&
+        oldStore.messages === store.messages
+      ) {
+        this._notifySubscribers();
+        // no conversion update
+        return;
+      }
+    }
+
+    const messages = !store.convertMessage
+      ? store.messages
+      : this._converter.convertMessages(store.messages, (cache, m, idx) => {
+          if (!store.convertMessage) return m;
+
+          const isLast = idx === store.messages.length - 1;
+          const autoStatus = getAutoStatus(isLast, isRunning);
+
+          if (
+            cache &&
+            (cache.role !== "assistant" ||
+              !isAutoStatus(cache.status) ||
+              cache.status === autoStatus)
+          )
+            return cache;
+
+          const newMessage = fromThreadMessageLike(
+            store.convertMessage(m, idx),
+            idx.toString(),
+            autoStatus,
+          );
+          (newMessage as any)[symbolInnerMessage] = m;
+          return newMessage;
+        });
+
+    for (let i = 0; i < messages.length; i++) {
+      const message = messages[i]!;
+      const parent = messages[i - 1];
+      this.repository.addOrUpdateMessage(parent?.id ?? null, message);
+    }
+
+    if (this.assistantOptimisticId) {
+      this.repository.deleteMessage(this.assistantOptimisticId);
+      this.assistantOptimisticId = null;
+    }
+
+    if (hasUpcomingMessage(isRunning, messages)) {
+      this.assistantOptimisticId = this.repository.appendOptimisticMessage(
+        messages.at(-1)?.id ?? null,
+        {
+          role: "assistant",
+          content: [],
+        },
+      );
+    }
+
+    this.repository.resetHead(
+      this.assistantOptimisticId ?? messages.at(-1)?.id ?? null,
+    );
+
+    this._messages = this.repository.getMessages();
+    this._notifySubscribers();
+  }
+
+  public override switchToBranch(branchId: string): void {
+    if (!this._store.setMessages)
+      throw new Error("Runtime does not support switching branches.");
+
+    this.repository.switchToBranch(branchId);
+    this.updateMessages(this.repository.getMessages());
+  }
+
+  public async append(message: AppendMessage): Promise<void> {
+    if (message.parentId !== (this.messages.at(-1)?.id ?? null)) {
+      if (!this._store.onEdit)
+        throw new Error("Runtime does not support editing messages.");
+      await this._store.onEdit(message);
+    } else {
+      await this._store.onNew(message);
+    }
+  }
+
+  public async startRun(parentId: string | null): Promise<void> {
+    if (!this._store.onReload)
+      throw new Error("Runtime does not support reloading messages.");
+
+    await this._store.onReload(parentId);
+  }
+
+  public cancelRun(): void {
+    if (!this._store.onCancel)
+      throw new Error("Runtime does not support cancelling runs.");
+
+    this._store.onCancel();
+
+    if (this.assistantOptimisticId) {
+      this.repository.deleteMessage(this.assistantOptimisticId);
+      this.assistantOptimisticId = null;
+    }
+
+    let messages = this.repository.getMessages();
+    const previousMessage = messages[messages.length - 1];
+    if (
+      previousMessage?.role === "user" &&
+      previousMessage.id === messages.at(-1)?.id // ensure the previous message is a leaf node
+    ) {
+      this.repository.deleteMessage(previousMessage.id);
+      if (!this.composer.text.trim()) {
+        this.composer.setText(getThreadMessageText(previousMessage));
+      }
+
+      messages = this.repository.getMessages();
+    } else {
+      this._notifySubscribers();
+    }
+
+    // resync messages (for reloading, to restore the previous branch)
+    setTimeout(() => {
+      this.updateMessages(messages);
+    }, 0);
+  }
+
+  public addToolResult(options: AddToolResultOptions) {
+    if (!this._store.onAddToolResult)
+      throw new Error("Runtime does not support tool results.");
+    this._store.onAddToolResult(options);
+  }
+
+  private updateMessages = (messages: ThreadMessage[]) => {
+    this._store.setMessages?.(
+      messages.flatMap(getExternalStoreMessage).filter((m) => m != null),
+    );
+  };
+}

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

@@ -0,0 +1,23 @@
+import { ThreadMessage } from "../../types";
+
+export type ConverterCallback<TIn> = (
+  cache: ThreadMessage | undefined,
+  message: TIn,
+  idx: number,
+) => ThreadMessage;
+
+export class ThreadMessageConverter {
+  private readonly cache = new WeakMap<WeakKey, ThreadMessage>();
+
+  convertMessages<TIn extends WeakKey>(
+    messages: TIn[],
+    converter: ConverterCallback<TIn>,
+  ): ThreadMessage[] {
+    return messages.map((m, idx) => {
+      const cached = this.cache.get(m);
+      const newMessage = converter(cached, m, idx);
+      this.cache.set(m, newMessage);
+      return newMessage;
+    });
+  }
+}

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

@@ -0,0 +1,133 @@
+import {
+  MessageStatus,
+  TextContentPart,
+  ImageContentPart,
+  ToolCallContentPart,
+  UIContentPart,
+  ThreadMessage,
+  ThreadAssistantContentPart,
+  ThreadAssistantMessage,
+  ThreadUserContentPart,
+  ThreadUserMessage,
+  ThreadSystemMessage,
+  CompleteAttachment,
+} from "../../types";
+import {
+  CoreToolCallContentPart,
+  Unstable_AudioContentPart,
+} from "../../types/AssistantTypes";
+
+export type ThreadMessageLike = {
+  role: "assistant" | "user" | "system";
+  content:
+    | string
+    | (
+        | TextContentPart
+        | ImageContentPart
+        | Unstable_AudioContentPart
+        | ToolCallContentPart<any, any>
+        | CoreToolCallContentPart<any, any>
+        | UIContentPart
+      )[];
+  id?: string | undefined;
+  createdAt?: Date | undefined;
+  status?: MessageStatus | undefined;
+  attachments?: CompleteAttachment[] | undefined;
+};
+
+export const fromThreadMessageLike = (
+  like: ThreadMessageLike,
+  fallbackId: string,
+  fallbackStatus: MessageStatus,
+): ThreadMessage => {
+  const { role, id, createdAt, attachments, status } = like;
+  const common = {
+    id: id ?? fallbackId,
+    createdAt: createdAt ?? new Date(),
+  };
+
+  const content =
+    typeof like.content === "string"
+      ? [{ type: "text" as const, text: like.content }]
+      : like.content;
+
+  if (role !== "user" && attachments)
+    throw new Error("Attachments are only supported for user messages");
+  // TODO add in 0.6
+  // if (role !== "assistant" && status)
+  //   throw new Error("Status is only supported for assistant messages");
+
+  switch (role) {
+    case "assistant":
+      return {
+        ...common,
+        role,
+        content: content
+          .map((part): ThreadAssistantContentPart | null => {
+            const type = part.type;
+            switch (type) {
+              case "text":
+                if (part.text.trim().length === 0) return null;
+                return part;
+
+              case "ui":
+                return part;
+
+              case "tool-call": {
+                if ("argsText" in part) return part;
+                return {
+                  ...part,
+                  argsText: JSON.stringify(part.args),
+                };
+              }
+
+              default: {
+                const unhandledType: "image" | "audio" = type;
+                throw new Error(`Unknown content part type: ${unhandledType}`);
+              }
+            }
+          })
+          .filter((c) => !!c),
+        status: status ?? fallbackStatus,
+      } satisfies ThreadAssistantMessage;
+
+    case "user":
+      return {
+        ...common,
+        role,
+        content: content.map((part): ThreadUserContentPart => {
+          const type = part.type;
+          switch (type) {
+            case "text":
+            case "ui":
+            case "image":
+            case "audio":
+              return part;
+
+            default: {
+              const unhandledType: "tool-call" = type;
+              throw new Error(`Unknown content part type: ${unhandledType}`);
+            }
+          }
+        }),
+        attachments: attachments ?? [],
+      } satisfies ThreadUserMessage;
+
+    case "system":
+      if (content.length !== 1 || content[0]!.type !== "text")
+        throw new Error(
+          "System messages must have exactly one text content part.",
+        );
+
+      return {
+        ...common,
+        role,
+        content: content as [TextContentPart],
+      } satisfies ThreadSystemMessage;
+
+    default: {
+      const unsupportedRole: never = role;
+      throw new Error(`Unknown message role: ${unsupportedRole}`);
+    }
+  }
+};

package/src/runtimes/external-store/auto-status.tsx

@@ -0,0 +1,13 @@
+import { MessageStatus } from "../../types";
+
+const AUTO_STATUS_RUNNING = Object.freeze({ type: "running" });
+const AUTO_STATUS_COMPLETE = Object.freeze({
+  type: "complete",
+  reason: "unknown",
+});
+
+export const isAutoStatus = (status: MessageStatus) =>
+  status === AUTO_STATUS_RUNNING || status === AUTO_STATUS_COMPLETE;
+
+export const getAutoStatus = (isLast: boolean, isRunning: boolean) =>
+  isLast && isRunning ? AUTO_STATUS_RUNNING : AUTO_STATUS_COMPLETE;