@assistant-ui/core 0.1.0

package/src/model-context/frame/provider.ts

@@ -0,0 +1,193 @@
+import { ModelContextProvider, ModelContext } from "../types";
+import type { Unsubscribe } from "../../types";
+import { Tool } from "assistant-stream";
+import { z } from "zod";
+import {
+  FrameMessage,
+  FRAME_MESSAGE_CHANNEL,
+  SerializedModelContext,
+  SerializedTool,
+} from "./types";
+
+const serializeTool = (tool: Tool<any, any>): SerializedTool => ({
+  ...(tool.description && { description: tool.description }),
+  parameters:
+    tool.parameters instanceof z.ZodType
+      ? ((z as any).toJSONSchema?.(tool.parameters) ?? tool.parameters)
+      : tool.parameters,
+  ...(tool.disabled !== undefined && { disabled: tool.disabled }),
+  ...(tool.type && { type: tool.type }),
+});
+
+const serializeModelContext = (
+  context: ModelContext,
+): SerializedModelContext => ({
+  ...(context.system !== undefined && { system: context.system }),
+  ...(context.tools && {
+    tools: Object.fromEntries(
+      Object.entries(context.tools).map(([name, tool]) => [
+        name,
+        serializeTool(tool),
+      ]),
+    ),
+  }),
+});
+
+export class AssistantFrameProvider {
+  private static _instance: AssistantFrameProvider | null = null;
+
+  private _providers = new Set<ModelContextProvider>();
+  private _providerUnsubscribes = new Map<
+    ModelContextProvider,
+    Unsubscribe | undefined
+  >();
+  private _targetOrigin: string;
+
+  private constructor(targetOrigin: string = "*") {
+    this._targetOrigin = targetOrigin;
+    this.handleMessage = this.handleMessage.bind(this);
+    window.addEventListener("message", this.handleMessage);
+
+    setTimeout(() => this.broadcastUpdate(), 0);
+  }
+
+  private static getInstance(targetOrigin?: string): AssistantFrameProvider {
+    if (!AssistantFrameProvider._instance) {
+      AssistantFrameProvider._instance = new AssistantFrameProvider(
+        targetOrigin,
+      );
+    }
+    return AssistantFrameProvider._instance;
+  }
+
+  private handleMessage(event: MessageEvent) {
+    if (this._targetOrigin !== "*" && event.origin !== this._targetOrigin)
+      return;
+    if (event.data?.channel !== FRAME_MESSAGE_CHANNEL) return;
+
+    const message = event.data.message as FrameMessage;
+
+    switch (message.type) {
+      case "model-context-request":
+        this.sendMessage(event, {
+          type: "model-context-update",
+          context: serializeModelContext(this.getModelContext()),
+        });
+        break;
+
+      case "tool-call":
+        this.handleToolCall(message, event);
+        break;
+    }
+  }
+
+  private async handleToolCall(
+    message: Extract<FrameMessage, { type: "tool-call" }>,
+    event: MessageEvent,
+  ) {
+    const tool = this.getModelContext().tools?.[message.toolName];
+
+    let result: any;
+    let error: string | undefined;
+
+    if (!tool) {
+      error = `Tool "${message.toolName}" not found`;
+    } else {
+      try {
+        result = tool.execute
+          ? await tool.execute(message.args, {
+              toolCallId: message.id,
+              abortSignal: new AbortController().signal,
+              human: async () => {
+                throw new Error(
+                  "Tool human input is not supported in frame context",
+                );
+              },
+            })
+          : undefined;
+      } catch (e) {
+        error = e instanceof Error ? e.message : String(e);
+      }
+    }
+
+    this.sendMessage(event, {
+      type: "tool-result",
+      id: message.id,
+      ...(error ? { error } : { result }),
+    });
+  }
+
+  private sendMessage(event: MessageEvent, message: FrameMessage) {
+    event.source?.postMessage(
+      { channel: FRAME_MESSAGE_CHANNEL, message },
+      { targetOrigin: event.origin },
+    );
+  }
+
+  private getModelContext(): ModelContext {
+    const contexts = Array.from(this._providers).map((p) =>
+      p.getModelContext(),
+    );
+
+    return contexts.reduce(
+      (merged, context) => ({
+        system: context.system
+          ? merged.system
+            ? `${merged.system}\n\n${context.system}`
+            : context.system
+          : merged.system,
+        tools: { ...(merged.tools || {}), ...(context.tools || {}) },
+      }),
+      {} as ModelContext,
+    );
+  }
+
+  private broadcastUpdate() {
+    if (window.parent && window.parent !== window) {
+      const updateMessage: FrameMessage = {
+        type: "model-context-update",
+        context: serializeModelContext(this.getModelContext()),
+      };
+
+      window.parent.postMessage(
+        { channel: FRAME_MESSAGE_CHANNEL, message: updateMessage },
+        this._targetOrigin,
+      );
+    }
+  }
+
+  static addModelContextProvider(
+    provider: ModelContextProvider,
+    targetOrigin?: string,
+  ): Unsubscribe {
+    const instance = AssistantFrameProvider.getInstance(targetOrigin);
+    instance._providers.add(provider);
+
+    const unsubscribe = provider.subscribe?.(() => instance.broadcastUpdate());
+    if (unsubscribe) {
+      instance._providerUnsubscribes.set(provider, unsubscribe);
+    }
+
+    instance.broadcastUpdate();
+
+    return () => {
+      instance._providers.delete(provider);
+      instance._providerUnsubscribes.get(provider)?.();
+      instance._providerUnsubscribes.delete(provider);
+      instance.broadcastUpdate();
+    };
+  }
+
+  static dispose() {
+    if (AssistantFrameProvider._instance) {
+      const instance = AssistantFrameProvider._instance;
+      window.removeEventListener("message", instance.handleMessage);
+
+      instance._providerUnsubscribes.forEach((unsubscribe) => unsubscribe?.());
+      instance._providerUnsubscribes.clear();
+      instance._providers.clear();
+
+      AssistantFrameProvider._instance = null;
+    }
+  }
+}

package/src/model-context/frame/types.ts

@@ -0,0 +1,40 @@
+export type SerializedTool = {
+  description?: string;
+  parameters: any; // JSON Schema
+  disabled?: boolean;
+  type?: string;
+};
+
+export type SerializedModelContext = {
+  system?: string;
+  tools?: Record<string, SerializedTool>;
+};
+
+export type FrameMessageType =
+  | "model-context-request"
+  | "model-context-update"
+  | "tool-call"
+  | "tool-result";
+
+export type FrameMessage =
+  | {
+      type: "model-context-request";
+    }
+  | {
+      type: "model-context-update";
+      context: SerializedModelContext;
+    }
+  | {
+      type: "tool-call";
+      id: string;
+      toolName: string;
+      args: unknown;
+    }
+  | {
+      type: "tool-result";
+      id: string;
+      result?: unknown;
+      error?: string;
+    };
+
+export const FRAME_MESSAGE_CHANNEL = "assistant-ui-frame";

package/src/model-context/index.ts

@@ -0,0 +1,35 @@
+export type {
+  // Language model settings
+  LanguageModelV1CallSettings,
+  LanguageModelConfig,
+  // Model context
+  ModelContext,
+  ModelContextProvider,
+  // Tool & instruction config
+  AssistantToolProps,
+  AssistantInstructionsConfig,
+} from "./types";
+export { mergeModelContexts } from "./types";
+
+export { tool } from "./tool";
+
+export { ModelContextRegistry } from "./registry";
+export type {
+  ModelContextRegistryToolHandle,
+  ModelContextRegistryInstructionHandle,
+  ModelContextRegistryProviderHandle,
+} from "./registry-handles";
+
+export {
+  // Classes
+  AssistantFrameHost,
+  AssistantFrameProvider,
+  // Constants
+  FRAME_MESSAGE_CHANNEL,
+} from "./frame";
+export type {
+  SerializedTool,
+  SerializedModelContext,
+  FrameMessageType,
+  FrameMessage,
+} from "./frame";

package/src/model-context/registry-handles.ts

@@ -0,0 +1,18 @@
+import type { AssistantToolProps, AssistantInstructionsConfig } from "./types";
+
+export interface ModelContextRegistryToolHandle<
+  TArgs extends Record<string, unknown> = any,
+  TResult = any,
+> {
+  update(tool: AssistantToolProps<TArgs, TResult>): void;
+  remove(): void;
+}
+
+export interface ModelContextRegistryInstructionHandle {
+  update(config: string | AssistantInstructionsConfig): void;
+  remove(): void;
+}
+
+export interface ModelContextRegistryProviderHandle {
+  remove(): void;
+}

package/src/model-context/registry.ts

@@ -0,0 +1,159 @@
+import { Tool } from "assistant-stream";
+import {
+  ModelContext,
+  ModelContextProvider,
+  mergeModelContexts,
+  AssistantToolProps,
+  AssistantInstructionsConfig,
+} from "./types";
+import type { Unsubscribe } from "../types";
+import {
+  ModelContextRegistryToolHandle,
+  ModelContextRegistryInstructionHandle,
+  ModelContextRegistryProviderHandle,
+} from "./registry-handles";
+
+export class ModelContextRegistry implements ModelContextProvider {
+  private _tools = new Map<symbol, AssistantToolProps<any, any>>();
+  private _instructions = new Map<symbol, string>();
+  private _providers = new Map<symbol, ModelContextProvider>();
+  private _subscribers = new Set<() => void>();
+  private _providerUnsubscribes = new Map<symbol, Unsubscribe | undefined>();
+
+  getModelContext(): ModelContext {
+    const instructions = Array.from(this._instructions.values()).filter(
+      Boolean,
+    );
+
+    const system =
+      instructions.length > 0 ? instructions.join("\n\n") : undefined;
+
+    const tools: Record<string, Tool<any, any>> = {};
+    for (const toolProps of this._tools.values()) {
+      const { toolName, render, ...tool } = toolProps;
+      tools[toolName] = tool;
+    }
+
+    const providerContexts = mergeModelContexts(
+      new Set(this._providers.values()),
+    );
+
+    const context: ModelContext = {
+      system,
+      tools: Object.keys(tools).length > 0 ? tools : undefined,
+    };
+
+    if (providerContexts.system) {
+      context.system = context.system
+        ? `${context.system}\n\n${providerContexts.system}`
+        : providerContexts.system;
+    }
+
+    if (providerContexts.tools) {
+      context.tools = { ...(context.tools || {}), ...providerContexts.tools };
+    }
+
+    if (providerContexts.callSettings) {
+      context.callSettings = providerContexts.callSettings;
+    }
+
+    if (providerContexts.config) {
+      context.config = providerContexts.config;
+    }
+
+    return context;
+  }
+
+  subscribe(callback: () => void): Unsubscribe {
+    this._subscribers.add(callback);
+    return () => this._subscribers.delete(callback);
+  }
+
+  private notifySubscribers(): void {
+    for (const callback of this._subscribers) {
+      callback();
+    }
+  }
+
+  addTool<TArgs extends Record<string, unknown>, TResult>(
+    tool: AssistantToolProps<TArgs, TResult>,
+  ): ModelContextRegistryToolHandle<TArgs, TResult> {
+    const id = Symbol();
+
+    this._tools.set(id, tool);
+    this.notifySubscribers();
+
+    return {
+      update: (newTool: AssistantToolProps<TArgs, TResult>) => {
+        if (this._tools.has(id)) {
+          this._tools.set(id, newTool);
+          this.notifySubscribers();
+        }
+      },
+      remove: () => {
+        this._tools.delete(id);
+        this.notifySubscribers();
+      },
+    };
+  }
+
+  addInstruction(
+    config: string | AssistantInstructionsConfig,
+  ): ModelContextRegistryInstructionHandle {
+    const id = Symbol();
+
+    const instruction =
+      typeof config === "string" ? config : config.instruction;
+    const disabled = typeof config === "object" ? config.disabled : false;
+
+    if (!disabled) {
+      this._instructions.set(id, instruction);
+      this.notifySubscribers();
+    }
+
+    return {
+      update: (newConfig: string | AssistantInstructionsConfig) => {
+        const newInstruction =
+          typeof newConfig === "string" ? newConfig : newConfig.instruction;
+        const newDisabled =
+          typeof newConfig === "object" ? newConfig.disabled : false;
+
+        if (newDisabled) {
+          this._instructions.delete(id);
+        } else {
+          this._instructions.set(id, newInstruction);
+        }
+        this.notifySubscribers();
+      },
+      remove: () => {
+        this._instructions.delete(id);
+        this.notifySubscribers();
+      },
+    };
+  }
+
+  addProvider(
+    provider: ModelContextProvider,
+  ): ModelContextRegistryProviderHandle {
+    const id = Symbol();
+
+    this._providers.set(id, provider);
+
+    const unsubscribe = provider.subscribe?.(() => {
+      this.notifySubscribers();
+    });
+    this._providerUnsubscribes.set(id, unsubscribe);
+
+    this.notifySubscribers();
+
+    return {
+      remove: () => {
+        this._providers.delete(id);
+        const unsubscribe = this._providerUnsubscribes.get(id);
+        unsubscribe?.();
+        this._providerUnsubscribes.delete(id);
+        this.notifySubscribers();
+      },
+    };
+  }
+}

package/src/model-context/tool.ts

@@ -0,0 +1,9 @@
+import { Tool } from "assistant-stream";
+
+// TODO re-add the inferrence of the parameters
+
+export function tool<TArgs extends Record<string, unknown>, TResult = any>(
+  tool: Tool<TArgs, TResult>,
+): Tool<TArgs, TResult> {
+  return tool;
+}

package/src/model-context/types.ts

@@ -0,0 +1,104 @@
+import type { Unsubscribe } from "../types";
+import { Tool } from "assistant-stream";
+
+// =============================================================================
+// Language Model Settings
+// =============================================================================
+
+export type LanguageModelV1CallSettings = {
+  maxTokens?: number;
+  temperature?: number;
+  topP?: number;
+  presencePenalty?: number;
+  frequencyPenalty?: number;
+  seed?: number;
+  headers?: Record<string, string | undefined>;
+};
+
+export type LanguageModelConfig = {
+  apiKey?: string;
+  baseUrl?: string;
+  modelName?: string;
+};
+
+// =============================================================================
+// Model Context
+// =============================================================================
+
+export type ModelContext = {
+  priority?: number | undefined;
+  system?: string | undefined;
+  tools?: Record<string, Tool<any, any>> | undefined;
+  callSettings?: LanguageModelV1CallSettings | undefined;
+  config?: LanguageModelConfig | undefined;
+};
+
+export type ModelContextProvider = {
+  getModelContext: () => ModelContext;
+  subscribe?: (callback: () => void) => Unsubscribe;
+};
+
+// =============================================================================
+// Tool & Instruction Config
+// =============================================================================
+
+export type AssistantToolProps<
+  TArgs extends Record<string, unknown>,
+  TResult,
+> = Tool<TArgs, TResult> & {
+  toolName: string;
+  render?: unknown;
+};
+
+export type AssistantInstructionsConfig = {
+  disabled?: boolean | undefined;
+  instruction: string;
+};
+
+// =============================================================================
+// Merging
+// =============================================================================
+
+export const mergeModelContexts = (
+  configSet: Set<ModelContextProvider>,
+): ModelContext => {
+  const configs = Array.from(configSet)
+    .map((c) => c.getModelContext())
+    .sort((a, b) => (b.priority ?? 0) - (a.priority ?? 0));
+
+  return configs.reduce((acc, config) => {
+    if (config.system) {
+      if (acc.system) {
+        acc.system += `\n\n${config.system}`;
+      } else {
+        acc.system = config.system;
+      }
+    }
+    if (config.tools) {
+      for (const [name, tool] of Object.entries(config.tools)) {
+        const existing = acc.tools?.[name];
+        if (existing && existing !== tool) {
+          throw new Error(
+            `You tried to define a tool with the name ${name}, but it already exists.`,
+          );
+        }
+
+        if (!acc.tools) acc.tools = {};
+        acc.tools[name] = tool;
+      }
+    }
+    if (config.config) {
+      acc.config = {
+        ...acc.config,
+        ...config.config,
+      };
+    }
+    if (config.callSettings) {
+      acc.callSettings = {
+        ...acc.callSettings,
+        ...config.callSettings,
+      };
+    }
+    return acc;
+  }, {} as ModelContext);
+};

package/src/runtime/api/assistant-runtime.ts

@@ -0,0 +1,115 @@
+import type { Unsubscribe } from "../../types";
+import type { ModelContextProvider } from "../../model-context";
+import type { AssistantRuntimeCore } from "../interfaces/assistant-runtime-core";
+import type { ThreadRuntime } from "./thread-runtime";
+import {
+  ThreadListRuntime,
+  ThreadListRuntimeImpl,
+} from "./thread-list-runtime";
+import { ExportedMessageRepository } from "../utils/message-repository";
+import type { ThreadMessageLike } from "../utils/thread-message-like";
+
+export type AssistantRuntime = {
+  /**
+   * The threads in this assistant.
+   */
+  readonly threads: ThreadListRuntime;
+
+  /**
+   * The currently selected main thread. Equivalent to `threads.main`.
+   */
+  readonly thread: ThreadRuntime;
+
+  /**
+   * @deprecated This field was renamed to `threads`.
+   */
+  readonly threadList: ThreadListRuntime;
+
+  /**
+   * Switch to a new thread.
+   *
+   * @deprecated This method was moved to `threads.switchToNewThread`.
+   */
+  switchToNewThread(): void;
+
+  /**
+   * Switch to a thread.
+   *
+   * @param threadId The thread ID to switch to.
+   * @deprecated This method was moved to `threads.switchToThread`.
+   */
+  switchToThread(threadId: string): void;
+
+  /**
+   * Register a model context provider. Model context providers are configuration such as system message, temperature, etc. that are set in the frontend.
+   *
+   * @param provider The model context provider to register.
+   */
+  registerModelContextProvider(provider: ModelContextProvider): Unsubscribe;
+
+  /**
+   * @deprecated This method was renamed to `registerModelContextProvider`.
+   */
+  registerModelConfigProvider(provider: ModelContextProvider): Unsubscribe;
+
+  /**
+   * @deprecated Use `runtime.thread.reset(initialMessages)`.
+   */
+  reset: unknown; // make it a type error
+};
+
+export class AssistantRuntimeImpl implements AssistantRuntime {
+  public readonly threads;
+  public get threadList() {
+    return this.threads;
+  }
+
+  public readonly _thread: ThreadRuntime;
+
+  public constructor(private readonly _core: AssistantRuntimeCore) {
+    this.threads = new ThreadListRuntimeImpl(_core.threads);
+    this._thread = this.threads.main;
+
+    this.__internal_bindMethods();
+  }
+
+  protected __internal_bindMethods() {
+    this.switchToNewThread = this.switchToNewThread.bind(this);
+    this.switchToThread = this.switchToThread.bind(this);
+    this.registerModelContextProvider =
+      this.registerModelContextProvider.bind(this);
+    this.registerModelConfigProvider =
+      this.registerModelConfigProvider.bind(this);
+    this.reset = this.reset.bind(this);
+  }
+
+  public get thread() {
+    return this._thread;
+  }
+
+  public switchToNewThread() {
+    return this._core.threads.switchToNewThread();
+  }
+
+  public switchToThread(threadId: string) {
+    return this._core.threads.switchToThread(threadId);
+  }
+
+  public registerModelContextProvider(provider: ModelContextProvider) {
+    return this._core.registerModelContextProvider(provider);
+  }
+
+  public registerModelConfigProvider(provider: ModelContextProvider) {
+    return this.registerModelContextProvider(provider);
+  }
+
+  public reset({
+    initialMessages,
+  }: {
+    initialMessages?: ThreadMessageLike[];
+  } = {}) {
+    return this._core.threads
+      .getMainThreadRuntimeCore()
+      .import(ExportedMessageRepository.fromArray(initialMessages ?? []));
+  }
+}