@assistant-ui/core 0.1.0

package/src/runtimes/local/local-thread-runtime-core.ts

@@ -0,0 +1,534 @@
+import { fromThreadMessageLike } from "../../runtime/utils/thread-message-like";
+import { generateId } from "../../utils/id";
+import type {
+  ChatModelAdapter,
+  ChatModelRunResult,
+} from "../../runtime/utils/chat-model-adapter";
+import { shouldContinue } from "./should-continue";
+import type { LocalRuntimeOptionsBase } from "./local-runtime-options";
+import type {
+  AddToolResultOptions,
+  ResumeToolCallOptions,
+  ThreadSuggestion,
+  ThreadRuntimeCore,
+  StartRunConfig,
+  ResumeRunConfig,
+} from "../../runtime/interfaces/thread-runtime-core";
+import { BaseThreadRuntimeCore } from "../../runtime/base/base-thread-runtime-core";
+import type {
+  RunConfig,
+  AppendMessage,
+  ThreadAssistantMessage,
+} from "../../types";
+import type { ModelContextProvider } from "../../model-context";
+
+class AbortError extends Error {
+  override name = "AbortError";
+  detach: boolean;
+
+  constructor(detach: boolean, message?: string) {
+    super(message);
+    this.detach = detach;
+  }
+}
+
+export class LocalThreadRuntimeCore
+  extends BaseThreadRuntimeCore
+  implements ThreadRuntimeCore
+{
+  public readonly capabilities = {
+    switchToBranch: true,
+    switchBranchDuringRun: true,
+    edit: true,
+    reload: true,
+    cancel: true,
+    unstable_copy: true,
+    speech: false,
+    dictation: false,
+    attachments: false,
+    feedback: false,
+  };
+
+  private abortController: AbortController | null = null;
+
+  public readonly isDisabled = false;
+
+  private _isLoading = false;
+  public get isLoading() {
+    return this._isLoading;
+  }
+
+  private _suggestions: readonly ThreadSuggestion[] = [];
+  private _suggestionsController: AbortController | null = null;
+  public get suggestions(): readonly ThreadSuggestion[] {
+    return this._suggestions;
+  }
+
+  public get adapters() {
+    return this._options.adapters;
+  }
+
+  constructor(
+    contextProvider: ModelContextProvider,
+    options: LocalRuntimeOptionsBase,
+  ) {
+    super(contextProvider);
+    this.__internal_setOptions(options);
+  }
+
+  private _options!: LocalRuntimeOptionsBase;
+
+  private _lastRunConfig: RunConfig = {};
+
+  private _getThreadId?: () => string | undefined;
+
+  public __internal_setGetThreadId(getThreadId: () => string | undefined) {
+    this._getThreadId = getThreadId;
+  }
+
+  private _getInitializePromise?: () => Promise<unknown> | undefined;
+
+  public __internal_setGetInitializePromise(
+    getPromise: () => Promise<unknown> | undefined,
+  ) {
+    this._getInitializePromise = getPromise;
+  }
+
+  public get extras() {
+    return undefined;
+  }
+
+  public __internal_setOptions(options: LocalRuntimeOptionsBase) {
+    if (this._options === options) return;
+
+    this._options = options;
+
+    let hasUpdates = false;
+
+    const canSpeak = options.adapters?.speech !== undefined;
+    if (this.capabilities.speech !== canSpeak) {
+      this.capabilities.speech = canSpeak;
+      hasUpdates = true;
+    }
+
+    const canDictate = options.adapters?.dictation !== undefined;
+    if (this.capabilities.dictation !== canDictate) {
+      this.capabilities.dictation = canDictate;
+      hasUpdates = true;
+    }
+
+    const canAttach = options.adapters?.attachments !== undefined;
+    if (this.capabilities.attachments !== canAttach) {
+      this.capabilities.attachments = canAttach;
+      hasUpdates = true;
+    }
+
+    const canFeedback = options.adapters?.feedback !== undefined;
+    if (this.capabilities.feedback !== canFeedback) {
+      this.capabilities.feedback = canFeedback;
+      hasUpdates = true;
+    }
+
+    if (hasUpdates) this._notifySubscribers();
+  }
+
+  private _loadPromise: Promise<void> | undefined;
+  public __internal_load() {
+    if (this._loadPromise) return this._loadPromise;
+
+    const promise = this.adapters.history?.load() ?? Promise.resolve(null);
+
+    this._isLoading = true;
+    this._notifySubscribers();
+
+    this._loadPromise = promise
+      .then((repo) => {
+        if (!repo) return;
+        this.repository.import(repo);
+        if (repo.messages.length > 0) {
+          this.ensureInitialized();
+        }
+        this._notifySubscribers();
+
+        const resume = this.adapters.history?.resume?.bind(
+          this.adapters.history,
+        );
+        if (repo.unstable_resume && resume) {
+          this.startRun(
+            {
+              parentId: this.repository.headId,
+              sourceId: this.repository.headId,
+              runConfig: this._lastRunConfig,
+            },
+            resume,
+          ).catch(() => {});
+        }
+      })
+      .finally(() => {
+        this._isLoading = false;
+        this._notifySubscribers();
+      });
+
+    return this._loadPromise;
+  }
+
+  public async append(message: AppendMessage): Promise<void> {
+    this.ensureInitialized();
+
+    const initPromise = this._getInitializePromise?.();
+    if (initPromise) {
+      await initPromise;
+    }
+
+    const newMessage = fromThreadMessageLike(message, generateId(), {
+      type: "complete",
+      reason: "unknown",
+    });
+    this.repository.addOrUpdateMessage(message.parentId, newMessage);
+    this._options.adapters.history?.append({
+      parentId: message.parentId,
+      message: newMessage,
+      ...(message.runConfig !== undefined && { runConfig: message.runConfig }),
+    });
+
+    const startRun = message.startRun ?? message.role === "user";
+    if (startRun) {
+      await this.startRun({
+        parentId: newMessage.id,
+        sourceId: message.sourceId,
+        runConfig: message.runConfig ?? {},
+      });
+    } else {
+      this.repository.resetHead(newMessage.id);
+      this._notifySubscribers();
+    }
+  }
+
+  public resumeRun({ stream, ...startConfig }: ResumeRunConfig): Promise<void> {
+    if (!stream)
+      throw new Error("You must pass a stream parameter to resume runs.");
+    return this.startRun(startConfig, stream);
+  }
+
+  public exportExternalState(): any {
+    throw new Error("Runtime does not support exporting external states.");
+  }
+
+  public importExternalState(): void {
+    throw new Error("Runtime does not support importing external states.");
+  }
+
+  public unstable_loadExternalState(): void {
+    throw new Error("Runtime does not support importing external states.");
+  }
+
+  public async startRun(
+    { parentId, runConfig }: StartRunConfig,
+    runCallback?: ChatModelAdapter["run"],
+  ): Promise<void> {
+    this.ensureInitialized();
+
+    // add assistant message
+    const id = generateId();
+    let message: ThreadAssistantMessage = {
+      id,
+      role: "assistant",
+      status: { type: "running" },
+      content: [],
+      metadata: {
+        unstable_state: this.state,
+        unstable_annotations: [],
+        unstable_data: [],
+        steps: [],
+        custom: {},
+      },
+      createdAt: new Date(),
+    };
+
+    this._notifyEventSubscribers("runStart");
+
+    try {
+      this._suggestions = [];
+      this._suggestionsController?.abort();
+      this._suggestionsController = null;
+      this._notifySubscribers();
+
+      do {
+        message = await this.performRoundtrip(
+          parentId,
+          message,
+          runConfig,
+          runCallback,
+        );
+        runCallback = undefined;
+      } while (shouldContinue(message, this._options.unstable_humanToolNames));
+    } finally {
+      this._notifyEventSubscribers("runEnd");
+    }
+
+    this._suggestionsController = new AbortController();
+    const signal = this._suggestionsController.signal;
+    if (
+      this.adapters.suggestion &&
+      message.status?.type !== "requires-action"
+    ) {
+      const promiseOrGenerator = this.adapters.suggestion?.generate({
+        messages: this.messages,
+      });
+
+      if (Symbol.asyncIterator in promiseOrGenerator) {
+        for await (const r of promiseOrGenerator) {
+          if (signal.aborted) break;
+          this._suggestions = r;
+          this._notifySubscribers();
+        }
+      } else {
+        const result = await promiseOrGenerator;
+        if (signal.aborted) return;
+        this._suggestions = result;
+        this._notifySubscribers();
+      }
+    }
+  }
+
+  private async performRoundtrip(
+    parentId: string | null,
+    message: ThreadAssistantMessage,
+    runConfig: RunConfig | undefined,
+    runCallback?: ChatModelAdapter["run"],
+  ) {
+    const messages = parentId ? this.repository.getMessages(parentId) : [];
+
+    // abort existing run
+    this.abortController?.abort();
+    this.abortController = new AbortController();
+
+    const initialContent = message.content;
+    const initialAnnotations = message.metadata?.unstable_annotations;
+    const initialData = message.metadata?.unstable_data;
+    const initialSteps = message.metadata?.steps;
+    const initialCustom = message.metadata?.custom;
+    const updateMessage = (m: Partial<ChatModelRunResult>) => {
+      const newSteps = m.metadata?.steps;
+      const steps = newSteps
+        ? [...(initialSteps ?? []), ...newSteps]
+        : undefined;
+
+      const newAnnotations = m.metadata?.unstable_annotations;
+      const newData = m.metadata?.unstable_data;
+      const annotations = newAnnotations
+        ? [...(initialAnnotations ?? []), ...newAnnotations]
+        : undefined;
+      const data = newData ? [...(initialData ?? []), ...newData] : undefined;
+
+      message = {
+        ...message,
+        ...(m.content
+          ? { content: [...initialContent, ...(m.content ?? [])] }
+          : undefined),
+        status: m.status ?? message.status,
+        ...(m.metadata
+          ? {
+              metadata: {
+                ...message.metadata,
+                ...(m.metadata.unstable_state
+                  ? { unstable_state: m.metadata.unstable_state }
+                  : undefined),
+                ...(annotations
+                  ? { unstable_annotations: annotations }
+                  : undefined),
+                ...(data ? { unstable_data: data } : undefined),
+                ...(steps ? { steps } : undefined),
+                ...(m.metadata?.timing
+                  ? { timing: m.metadata.timing }
+                  : undefined),
+                ...(m.metadata?.custom
+                  ? {
+                      custom: {
+                        ...(initialCustom ?? {}),
+                        ...m.metadata.custom,
+                      },
+                    }
+                  : undefined),
+              },
+            }
+          : undefined),
+      };
+      this.repository.addOrUpdateMessage(parentId, message);
+      this._notifySubscribers();
+    };
+
+    const maxSteps = this._options.maxSteps ?? 2;
+
+    const steps = message.metadata?.steps?.length ?? 0;
+    if (steps >= maxSteps) {
+      // reached max tool steps
+      updateMessage({
+        status: {
+          type: "incomplete",
+          reason: "tool-calls",
+        },
+      });
+      return message;
+    } else {
+      updateMessage({
+        status: {
+          type: "running",
+        },
+      });
+
+      // Switch to the new message branch right after adding it for the first time
+      this.repository.resetHead(message.id);
+      this._notifySubscribers();
+    }
+
+    try {
+      this._lastRunConfig = runConfig ?? {};
+      const context = this.getModelContext();
+
+      runCallback =
+        runCallback ??
+        this.adapters.chatModel.run.bind(this.adapters.chatModel);
+
+      const abortSignal = this.abortController.signal;
+      const threadId = this._getThreadId?.();
+      const promiseOrGenerator = runCallback({
+        messages,
+        runConfig: this._lastRunConfig,
+        abortSignal,
+        context,
+        config: context,
+        unstable_assistantMessageId: message.id,
+        unstable_threadId: threadId,
+        unstable_parentId: parentId,
+        unstable_getMessage() {
+          return message;
+        },
+      });
+
+      // handle async iterator for streaming results
+      if (Symbol.asyncIterator in promiseOrGenerator) {
+        for await (const r of promiseOrGenerator) {
+          if (abortSignal.aborted) {
+            updateMessage({
+              status: { type: "incomplete", reason: "cancelled" },
+            });
+            break;
+          }
+
+          updateMessage(r);
+        }
+      } else {
+        updateMessage(await promiseOrGenerator);
+      }
+
+      if (message.status.type === "running") {
+        updateMessage({
+          status: { type: "complete", reason: "unknown" },
+        });
+      }
+    } catch (e) {
+      // TODO this should be handled by the run result stream
+      if (e instanceof AbortError) {
+        updateMessage({
+          status: { type: "incomplete", reason: "cancelled" },
+        });
+      } else if (e instanceof Error && e.name === "AbortError") {
+        updateMessage({
+          status: { type: "incomplete", reason: "cancelled" },
+        });
+      } else {
+        updateMessage({
+          status: {
+            type: "incomplete",
+            reason: "error",
+            error:
+              e instanceof Error
+                ? e.message
+                : `[${typeof e}] ${new String(e).toString()}`,
+          },
+        });
+
+        throw e;
+      }
+    } finally {
+      this.abortController = null;
+
+      if (
+        message.status.type === "complete" ||
+        message.status.type === "incomplete"
+      ) {
+        await this._options.adapters.history?.append({
+          parentId,
+          message: message,
+          runConfig: this._lastRunConfig,
+        });
+      }
+    }
+    return message;
+  }
+
+  public detach() {
+    const error = new AbortError(true);
+    this.abortController?.abort(error);
+    this.abortController = null;
+  }
+
+  public cancelRun() {
+    const error = new AbortError(false);
+    this.abortController?.abort(error);
+    this.abortController = null;
+  }
+
+  public addToolResult({
+    messageId,
+    toolCallId,
+    result,
+    isError,
+    artifact,
+  }: AddToolResultOptions) {
+    const messageData = this.repository.getMessage(messageId);
+    const { parentId } = messageData;
+    let { message } = messageData;
+
+    if (message.role !== "assistant")
+      throw new Error("Tried to add tool result to non-assistant message");
+
+    let added = false;
+    let found = false;
+    const newContent = message.content.map((c) => {
+      if (c.type !== "tool-call") return c;
+      if (c.toolCallId !== toolCallId) return c;
+      found = true;
+      if (!c.result) added = true;
+      return {
+        ...c,
+        result,
+        artifact,
+        isError,
+      };
+    });
+
+    if (!found)
+      throw new Error("Tried to add tool result to non-existing tool call");
+
+    message = {
+      ...message,
+      content: newContent,
+    };
+    this.repository.addOrUpdateMessage(parentId, message);
+
+    if (
+      added &&
+      shouldContinue(message, this._options.unstable_humanToolNames)
+    ) {
+      this.performRoundtrip(parentId, message, this._lastRunConfig).catch(
+        () => {},
+      );
+    }
+  }
+
+  public resumeToolCall(_options: ResumeToolCallOptions) {
+    throw new Error("Local runtime does not support resuming tool calls.");
+  }
+}

package/src/runtimes/local/should-continue.ts

@@ -0,0 +1,26 @@
+import type { ThreadAssistantMessage } from "../../types";
+
+export const shouldContinue = (
+  result: ThreadAssistantMessage,
+  humanToolNames: string[] | undefined,
+) => {
+  // TODO legacy behavior -- make specifying human tool names required
+  if (humanToolNames === undefined) {
+    return (
+      result.status?.type === "requires-action" &&
+      result.status.reason === "tool-calls" &&
+      result.content.every((c) => c.type !== "tool-call" || !!c.result)
+    );
+  }
+
+  return (
+    result.status?.type === "requires-action" &&
+    result.status.reason === "tool-calls" &&
+    result.content.every(
+      (c) =>
+        c.type !== "tool-call" ||
+        !!c.result ||
+        !humanToolNames.includes(c.toolName),
+    )
+  );
+};

package/src/runtimes/remote-thread-list/adapter/in-memory.ts

@@ -0,0 +1,43 @@
+import { AssistantStream, AssistantStreamChunk } from "assistant-stream";
+import {
+  RemoteThreadInitializeResponse,
+  RemoteThreadListAdapter,
+  RemoteThreadListResponse,
+  RemoteThreadMetadata,
+} from "../types";
+
+export class InMemoryThreadListAdapter implements RemoteThreadListAdapter {
+  list(): Promise<RemoteThreadListResponse> {
+    return Promise.resolve({
+      threads: [],
+    });
+  }
+
+  rename(): Promise<void> {
+    return Promise.resolve();
+  }
+
+  archive(): Promise<void> {
+    return Promise.resolve();
+  }
+
+  unarchive(): Promise<void> {
+    return Promise.resolve();
+  }
+
+  delete(): Promise<void> {
+    return Promise.resolve();
+  }
+
+  initialize(threadId: string): Promise<RemoteThreadInitializeResponse> {
+    return Promise.resolve({ remoteId: threadId, externalId: undefined });
+  }
+
+  generateTitle(): Promise<AssistantStream> {
+    return Promise.resolve(new ReadableStream<AssistantStreamChunk>());
+  }
+
+  fetch(_threadId: string): Promise<RemoteThreadMetadata> {
+    return Promise.reject(new Error("Thread not found"));
+  }
+}