@assistant-ui/core 0.1.0

package/src/runtime/base/base-composer-runtime-core.ts

@@ -0,0 +1,418 @@
+import type {
+  Attachment,
+  CompleteAttachment,
+  PendingAttachment,
+  MessageRole,
+  RunConfig,
+  QuoteInfo,
+  AppendMessage,
+  Unsubscribe,
+} from "../../types";
+import { BaseSubscribable } from "../../subscribable";
+import type { AttachmentAdapter } from "../../adapters/attachment";
+import type {
+  ComposerRuntimeCore,
+  ComposerRuntimeEventType,
+  DictationState,
+} from "../interfaces/composer-runtime-core";
+import type { DictationAdapter } from "../../adapters/speech";
+
+const isAttachmentComplete = (a: Attachment): a is CompleteAttachment =>
+  a.status.type === "complete";
+
+export abstract class BaseComposerRuntimeCore
+  extends BaseSubscribable
+  implements ComposerRuntimeCore
+{
+  public readonly isEditing = true;
+
+  protected abstract getAttachmentAdapter(): AttachmentAdapter | undefined;
+  protected abstract getDictationAdapter(): DictationAdapter | undefined;
+
+  public get attachmentAccept(): string {
+    return this.getAttachmentAdapter()?.accept ?? "*";
+  }
+
+  private _attachments: readonly Attachment[] = [];
+  public get attachments() {
+    return this._attachments;
+  }
+
+  protected setAttachments(value: readonly Attachment[]) {
+    this._attachments = value;
+    this._notifySubscribers();
+  }
+
+  public abstract get canCancel(): boolean;
+
+  public get isEmpty() {
+    return !this.text.trim() && !this.attachments.length;
+  }
+
+  private _text = "";
+
+  get text() {
+    return this._text;
+  }
+
+  private _role: MessageRole = "user";
+
+  get role() {
+    return this._role;
+  }
+
+  private _runConfig: RunConfig = {};
+
+  get runConfig() {
+    return this._runConfig;
+  }
+
+  private _quote: QuoteInfo | undefined = undefined;
+
+  get quote() {
+    return this._quote;
+  }
+
+  public setQuote(quote: QuoteInfo | undefined) {
+    if (this._quote === quote) return;
+
+    this._quote = quote;
+    this._notifySubscribers();
+  }
+
+  public setText(value: string) {
+    if (this._text === value) return;
+
+    this._text = value;
+    if (this._dictation) {
+      this._dictationBaseText = value;
+      this._currentInterimText = "";
+      const { status, inputDisabled } = this._dictation;
+      this._dictation = inputDisabled ? { status, inputDisabled } : { status };
+    }
+    this._notifySubscribers();
+  }
+
+  public setRole(role: MessageRole) {
+    if (this._role === role) return;
+
+    this._role = role;
+    this._notifySubscribers();
+  }
+
+  public setRunConfig(runConfig: RunConfig) {
+    if (this._runConfig === runConfig) return;
+
+    this._runConfig = runConfig;
+    this._notifySubscribers();
+  }
+
+  private _emptyTextAndAttachments() {
+    this._attachments = [];
+    this._text = "";
+    this._notifySubscribers();
+  }
+
+  private async _onClearAttachments() {
+    const adapter = this.getAttachmentAdapter();
+    if (adapter) {
+      await Promise.all(this._attachments.map((a) => adapter.remove(a)));
+    }
+  }
+
+  public async reset() {
+    if (
+      this._attachments.length === 0 &&
+      this._text === "" &&
+      this._role === "user" &&
+      Object.keys(this._runConfig).length === 0 &&
+      this._quote === undefined
+    ) {
+      return;
+    }
+
+    this._role = "user";
+    this._runConfig = {};
+    this._quote = undefined;
+
+    const task = this._onClearAttachments();
+    this._emptyTextAndAttachments();
+    await task;
+  }
+
+  public async clearAttachments() {
+    const task = this._onClearAttachments();
+    this.setAttachments([]);
+
+    await task;
+  }
+
+  public async send() {
+    if (this._dictationSession) {
+      this._dictationSession.cancel();
+      this._cleanupDictation();
+    }
+
+    const adapter = this.getAttachmentAdapter();
+    const attachments =
+      adapter && this.attachments.length > 0
+        ? Promise.all(
+            this.attachments.map(async (a) => {
+              if (isAttachmentComplete(a)) return a;
+              const result = await adapter.send(a);
+              return result as CompleteAttachment;
+            }),
+          )
+        : [];
+
+    const text = this.text;
+    const quote = this._quote;
+    this._quote = undefined;
+    this._emptyTextAndAttachments();
+
+    const message: Omit<AppendMessage, "parentId" | "sourceId"> = {
+      createdAt: new Date(),
+      role: this.role,
+      content: text ? [{ type: "text", text }] : [],
+      attachments: await attachments,
+      runConfig: this.runConfig,
+      metadata: { custom: { ...(quote ? { quote } : {}) } },
+    };
+
+    this.handleSend(message);
+    this._notifyEventSubscribers("send");
+  }
+
+  public cancel() {
+    this.handleCancel();
+  }
+
+  protected abstract handleSend(
+    message: Omit<AppendMessage, "parentId" | "sourceId">,
+  ): void;
+  protected abstract handleCancel(): void;
+
+  async addAttachment(file: File) {
+    const adapter = this.getAttachmentAdapter();
+    if (!adapter) throw new Error("Attachments are not supported");
+
+    const upsertAttachment = (a: PendingAttachment) => {
+      const idx = this._attachments.findIndex(
+        (attachment) => attachment.id === a.id,
+      );
+      if (idx !== -1)
+        this._attachments = [
+          ...this._attachments.slice(0, idx),
+          a,
+          ...this._attachments.slice(idx + 1),
+        ];
+      else {
+        this._attachments = [...this._attachments, a];
+      }
+
+      this._notifySubscribers();
+    };
+
+    const promiseOrGenerator = adapter.add({ file });
+    if (Symbol.asyncIterator in promiseOrGenerator) {
+      for await (const r of promiseOrGenerator) {
+        upsertAttachment(r);
+      }
+    } else {
+      upsertAttachment(await promiseOrGenerator);
+    }
+
+    this._notifyEventSubscribers("attachmentAdd");
+    this._notifySubscribers();
+  }
+
+  async removeAttachment(attachmentId: string) {
+    const adapter = this.getAttachmentAdapter();
+    if (!adapter) throw new Error("Attachments are not supported");
+
+    const index = this._attachments.findIndex((a) => a.id === attachmentId);
+    if (index === -1) throw new Error("Attachment not found");
+    const attachment = this._attachments[index]!;
+
+    await adapter.remove(attachment);
+
+    this._attachments = this._attachments.filter((a) => a.id !== attachmentId);
+    this._notifySubscribers();
+  }
+
+  private _dictation: DictationState | undefined;
+  private _dictationSession: DictationAdapter.Session | undefined;
+  private _dictationUnsubscribes: Unsubscribe[] = [];
+  private _dictationBaseText = "";
+  private _currentInterimText = "";
+  private _dictationSessionIdCounter = 0;
+  private _activeDictationSessionId: number | undefined;
+  private _isCleaningDictation = false;
+
+  public get dictation(): DictationState | undefined {
+    return this._dictation;
+  }
+
+  private _isActiveSession(
+    sessionId: number,
+    session: DictationAdapter.Session,
+  ): boolean {
+    return (
+      this._activeDictationSessionId === sessionId &&
+      this._dictationSession === session
+    );
+  }
+
+  public startDictation(): void {
+    const adapter = this.getDictationAdapter();
+    if (!adapter) {
+      throw new Error("Dictation adapter not configured");
+    }
+
+    if (this._dictationSession) {
+      for (const unsub of this._dictationUnsubscribes) {
+        unsub();
+      }
+      this._dictationUnsubscribes = [];
+      const oldSession = this._dictationSession;
+      oldSession.stop().catch(() => {});
+      this._dictationSession = undefined;
+    }
+
+    const inputDisabled = adapter.disableInputDuringDictation ?? false;
+
+    this._dictationBaseText = this._text;
+    this._currentInterimText = "";
+
+    const session = adapter.listen();
+    this._dictationSession = session;
+    const sessionId = ++this._dictationSessionIdCounter;
+    this._activeDictationSessionId = sessionId;
+    this._dictation = { status: session.status, inputDisabled };
+    this._notifySubscribers();
+
+    const unsubSpeech = session.onSpeech((result) => {
+      if (!this._isActiveSession(sessionId, session)) return;
+      const isFinal = result.isFinal !== false;
+
+      const needsSeparator =
+        this._dictationBaseText &&
+        !this._dictationBaseText.endsWith(" ") &&
+        result.transcript;
+      const separator = needsSeparator ? " " : "";
+
+      if (isFinal) {
+        this._dictationBaseText =
+          this._dictationBaseText + separator + result.transcript;
+        this._currentInterimText = "";
+        this._text = this._dictationBaseText;
+
+        if (this._dictation) {
+          const { transcript: _, ...rest } = this._dictation;
+          this._dictation = rest;
+        }
+        this._notifySubscribers();
+      } else {
+        this._currentInterimText = separator + result.transcript;
+        this._text = this._dictationBaseText + this._currentInterimText;
+
+        if (this._dictation) {
+          this._dictation = {
+            ...this._dictation,
+            transcript: result.transcript,
+          };
+        }
+        this._notifySubscribers();
+      }
+    });
+    this._dictationUnsubscribes.push(unsubSpeech);
+
+    const unsubStart = session.onSpeechStart(() => {
+      if (!this._isActiveSession(sessionId, session)) return;
+
+      this._dictation = {
+        status: { type: "running" },
+        inputDisabled,
+        ...(this._dictation?.transcript && {
+          transcript: this._dictation.transcript,
+        }),
+      };
+      this._notifySubscribers();
+    });
+    this._dictationUnsubscribes.push(unsubStart);
+
+    const unsubEnd = session.onSpeechEnd(() => {
+      this._cleanupDictation({ sessionId });
+    });
+    this._dictationUnsubscribes.push(unsubEnd);
+
+    const statusInterval = setInterval(() => {
+      if (!this._isActiveSession(sessionId, session)) return;
+
+      if (session.status.type === "ended") {
+        this._cleanupDictation({ sessionId });
+      }
+    }, 100);
+    this._dictationUnsubscribes.push(() => clearInterval(statusInterval));
+  }
+
+  public stopDictation(): void {
+    if (!this._dictationSession) return;
+
+    const session = this._dictationSession;
+    const sessionId = this._activeDictationSessionId;
+    session.stop().finally(() => {
+      this._cleanupDictation({ sessionId });
+    });
+  }
+
+  private _cleanupDictation(options?: { sessionId: number | undefined }): void {
+    const isStaleSession =
+      options?.sessionId !== undefined &&
+      options.sessionId !== this._activeDictationSessionId;
+    if (isStaleSession || this._isCleaningDictation) return;
+
+    this._isCleaningDictation = true;
+    try {
+      for (const unsub of this._dictationUnsubscribes) {
+        unsub();
+      }
+      this._dictationUnsubscribes = [];
+      this._dictationSession = undefined;
+      this._activeDictationSessionId = undefined;
+      this._dictation = undefined;
+      this._dictationBaseText = "";
+      this._currentInterimText = "";
+      this._notifySubscribers();
+    } finally {
+      this._isCleaningDictation = false;
+    }
+  }
+
+  private _eventSubscribers = new Map<
+    ComposerRuntimeEventType,
+    Set<() => void>
+  >();
+
+  protected _notifyEventSubscribers(event: ComposerRuntimeEventType) {
+    const subscribers = this._eventSubscribers.get(event);
+    if (!subscribers) return;
+
+    for (const callback of subscribers) callback();
+  }
+
+  public unstable_on(event: ComposerRuntimeEventType, callback: () => void) {
+    const subscribers = this._eventSubscribers.get(event);
+    if (!subscribers) {
+      this._eventSubscribers.set(event, new Set([callback]));
+    } else {
+      subscribers.add(callback);
+    }
+
+    return () => {
+      const subscribers = this._eventSubscribers.get(event);
+      if (!subscribers) return;
+      subscribers.delete(callback);
+    };
+  }
+}

package/src/runtime/base/base-thread-runtime-core.ts

@@ -0,0 +1,234 @@
+import type { AppendMessage, ThreadMessage, Unsubscribe } from "../../types";
+import type { ModelContextProvider } from "../../model-context";
+import { getThreadMessageText } from "../../utils/text";
+import {
+  ExportedMessageRepository,
+  MessageRepository,
+} from "../utils/message-repository";
+import { DefaultThreadComposerRuntimeCore } from "./default-thread-composer-runtime-core";
+import type {
+  AddToolResultOptions,
+  ResumeToolCallOptions,
+  ThreadSuggestion,
+  SubmitFeedbackOptions,
+  ThreadRuntimeCore,
+  SpeechState,
+  RuntimeCapabilities,
+  ThreadRuntimeEventType,
+  StartRunConfig,
+  ResumeRunConfig,
+} from "../interfaces/thread-runtime-core";
+import { DefaultEditComposerRuntimeCore } from "./default-edit-composer-runtime-core";
+import type { SpeechSynthesisAdapter } from "../../adapters/speech";
+import type { FeedbackAdapter } from "../../adapters/feedback";
+import type { AttachmentAdapter } from "../../adapters/attachment";
+import type { ThreadMessageLike } from "../utils/thread-message-like";
+
+type BaseThreadAdapters = {
+  speech?: SpeechSynthesisAdapter | undefined;
+  feedback?: FeedbackAdapter | undefined;
+  attachments?: AttachmentAdapter | undefined;
+};
+
+export abstract class BaseThreadRuntimeCore implements ThreadRuntimeCore {
+  private _subscriptions = new Set<() => void>();
+  private _isInitialized = false;
+
+  protected readonly repository = new MessageRepository();
+  public abstract get adapters(): BaseThreadAdapters | undefined;
+  public abstract get isDisabled(): boolean;
+  public abstract get isLoading(): boolean;
+  public abstract get suggestions(): readonly ThreadSuggestion[];
+  public abstract get extras(): unknown;
+
+  public abstract get capabilities(): RuntimeCapabilities;
+  public abstract append(message: AppendMessage): void;
+  public abstract startRun(config: StartRunConfig): void;
+  public abstract resumeRun(config: ResumeRunConfig): void;
+  public abstract addToolResult(options: AddToolResultOptions): void;
+  public abstract resumeToolCall(options: ResumeToolCallOptions): void;
+  public abstract cancelRun(): void;
+  public abstract exportExternalState(): any;
+  public abstract importExternalState(state: any): void;
+  public abstract unstable_loadExternalState(state: any): void;
+
+  public get messages() {
+    return this.repository.getMessages();
+  }
+
+  public get state() {
+    let mostRecentAssistantMessage;
+    for (const message of this.messages) {
+      if (message.role === "assistant") {
+        mostRecentAssistantMessage = message;
+      }
+    }
+
+    return mostRecentAssistantMessage?.metadata.unstable_state ?? null;
+  }
+
+  public readonly composer = new DefaultThreadComposerRuntimeCore(this);
+
+  constructor(private readonly _contextProvider: ModelContextProvider) {}
+
+  public getModelContext() {
+    return this._contextProvider.getModelContext();
+  }
+
+  private _editComposers = new Map<string, DefaultEditComposerRuntimeCore>();
+  public getEditComposer(messageId: string) {
+    return this._editComposers.get(messageId);
+  }
+  public beginEdit(messageId: string) {
+    if (this._editComposers.has(messageId))
+      throw new Error("Edit already in progress");
+
+    this._editComposers.set(
+      messageId,
+      new DefaultEditComposerRuntimeCore(
+        this,
+        () => this._editComposers.delete(messageId),
+        this.repository.getMessage(messageId),
+      ),
+    );
+    this._notifySubscribers();
+  }
+
+  public getMessageById(messageId: string) {
+    try {
+      return this.repository.getMessage(messageId);
+    } catch {
+      return undefined;
+    }
+  }
+
+  public getBranches(messageId: string): string[] {
+    return this.repository.getBranches(messageId);
+  }
+
+  public switchToBranch(branchId: string): void {
+    this.repository.switchToBranch(branchId);
+    this._notifySubscribers();
+  }
+
+  protected _notifySubscribers() {
+    for (const callback of this._subscriptions) callback();
+  }
+
+  public _notifyEventSubscribers(event: ThreadRuntimeEventType) {
+    const subscribers = this._eventSubscribers.get(event);
+    if (!subscribers) return;
+
+    for (const callback of subscribers) callback();
+  }
+
+  public subscribe(callback: () => void): Unsubscribe {
+    this._subscriptions.add(callback);
+    return () => this._subscriptions.delete(callback);
+  }
+
+  public submitFeedback({ messageId, type }: SubmitFeedbackOptions) {
+    const adapter = this.adapters?.feedback;
+    if (!adapter) throw new Error("Feedback adapter not configured");
+
+    const { message, parentId } = this.repository.getMessage(messageId);
+    adapter.submit({ message, type });
+
+    if (message.role === "assistant") {
+      const updatedMessage: ThreadMessage = {
+        ...message,
+        metadata: {
+          ...message.metadata,
+          submittedFeedback: { type },
+        },
+      };
+      this.repository.addOrUpdateMessage(parentId, updatedMessage);
+    }
+
+    this._notifySubscribers();
+  }
+
+  private _stopSpeaking: Unsubscribe | undefined;
+  public speech: SpeechState | undefined;
+
+  public speak(messageId: string) {
+    const adapter = this.adapters?.speech;
+    if (!adapter) throw new Error("Speech adapter not configured");
+
+    const { message } = this.repository.getMessage(messageId);
+
+    this._stopSpeaking?.();
+
+    const utterance = adapter.speak(getThreadMessageText(message));
+    const unsub = utterance.subscribe(() => {
+      if (utterance.status.type === "ended") {
+        this._stopSpeaking = undefined;
+        this.speech = undefined;
+      } else {
+        this.speech = { messageId, status: utterance.status };
+      }
+      this._notifySubscribers();
+    });
+
+    this.speech = { messageId, status: utterance.status };
+    this._notifySubscribers();
+
+    this._stopSpeaking = () => {
+      utterance.cancel();
+      unsub();
+      this.speech = undefined;
+      this._stopSpeaking = undefined;
+    };
+  }
+
+  public stopSpeaking() {
+    if (!this._stopSpeaking) throw new Error("No message is being spoken");
+    this._stopSpeaking();
+    this._notifySubscribers();
+  }
+
+  protected ensureInitialized() {
+    if (!this._isInitialized) {
+      this._isInitialized = true;
+      this._notifyEventSubscribers("initialize");
+    }
+  }
+
+  public export() {
+    return this.repository.export();
+  }
+
+  public import(data: ExportedMessageRepository) {
+    this.ensureInitialized();
+    this.repository.clear();
+    this.repository.import(data);
+    this._notifySubscribers();
+  }
+
+  public reset(initialMessages?: readonly ThreadMessageLike[]) {
+    this.import(ExportedMessageRepository.fromArray(initialMessages ?? []));
+  }
+
+  private _eventSubscribers = new Map<
+    ThreadRuntimeEventType,
+    Set<() => void>
+  >();
+
+  public unstable_on(event: ThreadRuntimeEventType, callback: () => void) {
+    if (event === "modelContextUpdate") {
+      return this._contextProvider.subscribe?.(callback) ?? (() => {});
+    }
+
+    const subscribers = this._eventSubscribers.get(event);
+    if (!subscribers) {
+      this._eventSubscribers.set(event, new Set([callback]));
+    } else {
+      subscribers.add(callback);
+    }
+
+    return () => {
+      const subscribers = this._eventSubscribers.get(event)!;
+      subscribers.delete(callback);
+    };
+  }
+}

package/src/runtime/base/default-edit-composer-runtime-core.ts

@@ -0,0 +1,71 @@
+import type { AppendMessage, ThreadMessage } from "../../types";
+import { getThreadMessageText } from "../../utils/text";
+import type { AttachmentAdapter } from "../../adapters/attachment";
+import type { DictationAdapter } from "../../adapters/speech";
+import type { ThreadRuntimeCore } from "../interfaces/thread-runtime-core";
+import { BaseComposerRuntimeCore } from "./base-composer-runtime-core";
+
+export class DefaultEditComposerRuntimeCore extends BaseComposerRuntimeCore {
+  public get canCancel() {
+    return true;
+  }
+
+  protected getAttachmentAdapter() {
+    return this.runtime.adapters?.attachments;
+  }
+
+  protected getDictationAdapter() {
+    return this.runtime.adapters?.dictation;
+  }
+
+  private _nonTextParts;
+  private _previousText;
+  private _parentId;
+  private _sourceId;
+  constructor(
+    private runtime: ThreadRuntimeCore & {
+      adapters?:
+        | {
+            attachments?: AttachmentAdapter | undefined;
+            dictation?: DictationAdapter | undefined;
+          }
+        | undefined;
+    },
+    private endEditCallback: () => void,
+    { parentId, message }: { parentId: string | null; message: ThreadMessage },
+  ) {
+    super();
+    this._parentId = parentId;
+    this._sourceId = message.id;
+    this._previousText = getThreadMessageText(message);
+    this.setText(this._previousText);
+
+    this.setRole(message.role);
+    this.setAttachments(message.attachments ?? []);
+
+    this._nonTextParts = message.content.filter((part) => part.type !== "text");
+
+    this.setRunConfig({ ...runtime.composer.runConfig });
+  }
+
+  public async handleSend(
+    message: Omit<AppendMessage, "parentId" | "sourceId">,
+  ) {
+    const text = getThreadMessageText(message as AppendMessage);
+    if (text !== this._previousText) {
+      this.runtime.append({
+        ...message,
+        content: [...message.content, ...this._nonTextParts] as any,
+        parentId: this._parentId,
+        sourceId: this._sourceId,
+      });
+    }
+
+    this.handleCancel();
+  }
+
+  public handleCancel() {
+    this.endEditCallback();
+    this._notifySubscribers();
+  }
+}