@assistant-ui/react 0.10.44 → 0.10.45

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

@@ -0,0 +1,165 @@
+import { Tool } from "assistant-stream";
+import {
+  ModelContext,
+  ModelContextProvider,
+  mergeModelContexts,
+} from "../../model-context/ModelContextTypes";
+import { Unsubscribe } from "../../types/Unsubscribe";
+import {
+  ModelContextRegistryToolHandle,
+  ModelContextRegistryInstructionHandle,
+  ModelContextRegistryProviderHandle,
+} from "./ModelContextRegistryHandles";
+import type { AssistantToolProps } from "../../model-context/useAssistantTool";
+import type { AssistantInstructionsConfig } from "../../model-context/useAssistantInstructions";
+
+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 {
+    // Merge instructions
+    const instructions = Array.from(this._instructions.values()).filter(
+      Boolean,
+    );
+
+    const system =
+      instructions.length > 0 ? instructions.join("\n\n") : undefined;
+
+    // Collect tools
+    const tools: Record<string, Tool<any, any>> = {};
+    for (const toolProps of this._tools.values()) {
+      const { toolName, render, ...tool } = toolProps;
+      tools[toolName] = tool;
+    }
+
+    // Merge provider contexts
+    const providerContexts = mergeModelContexts(
+      new Set(this._providers.values()),
+    );
+
+    // Combine everything
+    const context: ModelContext = {
+      system,
+      tools: Object.keys(tools).length > 0 ? tools : undefined,
+    };
+
+    // Merge with provider contexts
+    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);
+
+    // Subscribe to provider changes
+    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/registry/ModelContextRegistryHandles.ts

@@ -0,0 +1,19 @@
+import type { AssistantToolProps } from "../../model-context/useAssistantTool";
+import type { AssistantInstructionsConfig } from "../../model-context/useAssistantInstructions";
+
+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/SPEC_ModelContextRegistry.md

@@ -0,0 +1,40 @@
+# ModelContextRegistry
+
+An imperative API for registering tools and instructions to a model context provider.
+
+```typescript
+const registry = new ModelContextRegistry();
+
+const handle = registry.addTool({
+  toolName: "search",
+  description: "Search the web",
+  parameters: z.object({ query: z.string() }),
+  execute: async (args) => {
+    return { results: ["..."] };
+  },
+});
+```
+
+## API
+
+- addTool(tool: Tool & { toolName: string }): ModelContextRegistryToolHandle
+- addInstruction(instruction: string): ModelContextRegistryInstructionHandle
+- addProvider(provider: ModelContextProvider): ModelContextRegistryProviderHandle
+
+## ModelContextRegistryToolHandle
+
+- update(tool: Tool & { toolName: string }): void;
+- remove(): void;
+
+## ModelContextRegistryInstructionHandle
+
+- update(instruction: string): void;
+- remove(): void;
+
+## ModelContextRegistryProviderHandle
+
+- remove(): void;
+
+## ModelContextProvider
+
+The registry is a ModelContextProvider.

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

@@ -0,0 +1,2 @@
+export * from "./ModelContextRegistry";
+export * from "./ModelContextRegistryHandles";

package/src/model-context/useAssistantInstructions.tsx

@@ -3,7 +3,7 @@
 import { useEffect } from "react";
 import { useAssistantRuntime } from "../context";
 
-type AssistantInstructionsConfig = {
+export type AssistantInstructionsConfig = {
   disabled?: boolean | undefined;
   instruction: string;
 };