@assistant-ui/react 0.7.65 → 0.7.66

package/src/runtimes/edge/EdgeChatAdapter.ts

@@ -11,6 +11,11 @@ import { streamPartDecoderStream } from "./streams/utils/streamPartDecoderStream
 import { runResultStream } from "./streams/runResultStream";
 import { toolResultStream } from "./streams/toolResultStream";
 import { toLanguageModelMessages } from "./converters";
+import { ThreadMessage } from "../../types";
+import { Tool } from "../../model-context";
+import { z } from "zod";
+import zodToJsonSchema from "zod-to-json-schema";
+import { JSONSchema7 } from "json-schema";
 
 export function asAsyncIterable<T>(
   source: ReadableStream<T>,
@@ -29,23 +34,71 @@ export function asAsyncIterable<T>(
     },
   };
 }
+
 export type EdgeChatAdapterOptions = {
   api: string;
 
+  // experimental_prepareRequestBody?: (options: {
+  //   id: string;
+  //   messages: UIMessage[];
+  //   requestData?: JSONValue;
+  //   requestBody?: object;
+  // }) => unknown;
+
+  // onToolCall?: ({
+  //   toolCall,
+  // }: {
+  //   toolCall: UIMessageToolInvocation;
+  // }) => void | Promise<unknown> | unknown;
+
+  /**
+   * Callback function to be called when the API response is received.
+   */
+  onResponse?: (response: Response) => void | Promise<void>;
+  /**
+   * Optional callback function that is called when the assistant message is finished streaming.
+   */
+  onFinish?: (message: ThreadMessage) => void;
+  /**
+   * Callback function to be called when an error is encountered.
+   */
+  onError?: (error: Error) => void;
+
   credentials?: RequestCredentials;
   headers?: Record<string, string> | Headers;
   body?: object;
 
   /**
-   * When enabled, the adapter will not strip `id` from messages in the messages array.
+   * @deprecated Renamed to `sendExtraMessageFields`.
    */
   unstable_sendMessageIds?: boolean;
 
+  /**
+   * When enabled, the adapter will not strip `id` from messages in the messages array.
+   */
+  sendExtraMessageFields?: boolean;
+
   /**
    * When enabled, the adapter will send messages in the format expected by the Vercel AI SDK Core.
    * This feature will be removed in the future in favor of a better solution.
+   *
+   * `v2` sends frontend tools in a format that can be directly passed to `stremaText`
    */
-  unstable_AISDKInterop?: boolean | undefined;
+  unstable_AISDKInterop?: boolean | "v2" | undefined;
+};
+
+const toAISDKTools = (tools: Record<string, Tool<any, any>>) => {
+  return Object.fromEntries(
+    Object.entries(tools).map(([name, tool]) => [
+      name,
+      {
+        ...(tool.description ? { description: tool.description } : undefined),
+        parameters: (tool.parameters instanceof z.ZodType
+          ? zodToJsonSchema(tool.parameters)
+          : tool.parameters) as JSONSchema7,
+      },
+    ]),
+  );
 };
 
 export class EdgeChatAdapter implements ChatModelAdapter {
@@ -57,6 +110,7 @@ export class EdgeChatAdapter implements ChatModelAdapter {
     abortSignal,
     context,
     unstable_assistantMessageId,
+    unstable_getMessage,
   }: ChatModelRunOptions) {
     const headers = new Headers(this.options.headers);
     headers.set("Content-Type", "application/json");
@@ -69,12 +123,20 @@ export class EdgeChatAdapter implements ChatModelAdapter {
         system: context.system,
         messages: this.options.unstable_AISDKInterop
           ? (toLanguageModelMessages(messages, {
-              unstable_includeId: this.options.unstable_sendMessageIds,
+              unstable_includeId:
+                this.options.unstable_sendMessageIds ||
+                this.options.sendExtraMessageFields,
             }) as EdgeRuntimeRequestOptions["messages"]) // TODO figure out a better way to do this
           : toCoreMessages(messages, {
-              unstable_includeId: this.options.unstable_sendMessageIds,
+              unstable_includeId:
+                this.options.unstable_sendMessageIds ||
+                this.options.sendExtraMessageFields,
             }),
-        tools: context.tools ? toLanguageModelTools(context.tools) : [],
+        tools: context.tools
+          ? this.options.unstable_AISDKInterop === "v2"
+            ? (toAISDKTools(context.tools) as any)
+            : toLanguageModelTools(context.tools)
+          : [],
         unstable_assistantMessageId,
         runConfig,
         ...context.callSettings,
@@ -85,22 +147,31 @@ export class EdgeChatAdapter implements ChatModelAdapter {
       signal: abortSignal,
     });
 
-    if (!result.ok) {
-      throw new Error(`Status ${result.status}: ${await result.text()}`);
-    }
+    await this.options.onResponse?.(result);
 
-    const stream = result
-      .body!.pipeThrough(streamPartDecoderStream())
-      .pipeThrough(assistantDecoderStream())
-      .pipeThrough(toolResultStream(context.tools, abortSignal))
-      .pipeThrough(runResultStream());
+    try {
+      if (!result.ok) {
+        throw new Error(`Status ${result.status}: ${await result.text()}`);
+      }
 
-    let update: ChatModelRunResult | undefined;
-    for await (update of asAsyncIterable(stream)) {
-      yield update;
-    }
+      const stream = result
+        .body!.pipeThrough(streamPartDecoderStream())
+        .pipeThrough(assistantDecoderStream())
+        .pipeThrough(toolResultStream(context.tools, abortSignal))
+        .pipeThrough(runResultStream());
+
+      let update: ChatModelRunResult | undefined;
+      for await (update of asAsyncIterable(stream)) {
+        yield update;
+      }
 
-    if (update === undefined)
-      throw new Error("No data received from Edge Runtime");
+      if (update === undefined)
+        throw new Error("No data received from Edge Runtime");
+
+      this.options.onFinish?.(unstable_getMessage());
+    } catch (error: unknown) {
+      this.options.onError?.(error as Error);
+      throw error;
+    }
   }
 }

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

@@ -167,6 +167,10 @@ export class ExternalStoreThreadListRuntimeCore
     throw new Error("Method not implemented.");
   }
 
+  public generateTitle(): never {
+    throw new Error("Method not implemented.");
+  }
+
   private _subscriptions = new Set<() => void>();
 
   public subscribe(callback: () => void): Unsubscribe {

package/src/runtimes/local/ChatModelAdapter.tsx

@@ -42,6 +42,7 @@ export type ChatModelRunOptions = {
   readonly config: ModelContext;
 
   readonly unstable_assistantMessageId?: string;
+  unstable_getMessage(): ThreadMessage;
 };
 
 export type ChatModelAdapter = {

package/src/runtimes/local/LocalRuntimeOptions.tsx

@@ -4,6 +4,7 @@ import { ThreadMessageLike } from "../external-store";
 import { FeedbackAdapter } from "../adapters/feedback/FeedbackAdapter";
 import { SpeechSynthesisAdapter } from "../adapters/speech/SpeechAdapterTypes";
 import { ChatModelAdapter } from "./ChatModelAdapter";
+import { AssistantCloud } from "../../cloud";
 
 export type LocalRuntimeOptionsBase = {
   maxSteps?: number | undefined;
@@ -24,6 +25,7 @@ export type LocalRuntimeOptionsBase = {
 
 // TODO align LocalRuntimeOptions with LocalRuntimeOptionsBase
 export type LocalRuntimeOptions = Omit<LocalRuntimeOptionsBase, "adapters"> & {
+  cloud?: AssistantCloud | undefined;
   initialMessages?: readonly ThreadMessageLike[] | undefined;
   adapters?: Omit<LocalRuntimeOptionsBase["adapters"], "chatModel"> | undefined;
 };
@@ -32,6 +34,7 @@ export const splitLocalRuntimeOptions = <T extends LocalRuntimeOptions>(
   options: T,
 ) => {
   const {
+    cloud,
     initialMessages,
     maxSteps,
     adapters,
@@ -41,6 +44,7 @@ export const splitLocalRuntimeOptions = <T extends LocalRuntimeOptions>(
 
   return {
     localRuntimeOptions: {
+      cloud,
       initialMessages,
       maxSteps,
       adapters,

package/src/runtimes/local/LocalThreadListRuntimeCore.tsx

@@ -1,235 +1,92 @@
-import type { Unsubscribe } from "../../types";
 import { ThreadListRuntimeCore } from "../core/ThreadListRuntimeCore";
-import { generateId } from "../../utils/idUtils";
+import { BaseSubscribable } from "../remote-thread-list/BaseSubscribable";
 import { LocalThreadRuntimeCore } from "./LocalThreadRuntimeCore";
 
-export type ThreadListAdapter = {
-  subscribe(callback: () => void): Unsubscribe;
-};
-
-export type LocalThreadData = {
-  readonly runtime: LocalThreadRuntimeCore;
-  readonly status: "new" | "regular" | "archived";
-  readonly threadId: string;
-  readonly title?: string | undefined;
-};
-
 export type LocalThreadFactory = () => LocalThreadRuntimeCore;
 
-const RESOLVED_PROMISE = Promise.resolve();
-export class LocalThreadListRuntimeCore implements ThreadListRuntimeCore {
-  private _threadData = new Map<string, LocalThreadData>();
-  private _threadIds: readonly string[] = [];
-  private _archivedThreadIds: readonly string[] = [];
-  private _newThreadId: string | undefined;
+const EMPTY_ARRAY = Object.freeze([]);
+export class LocalThreadListRuntimeCore
+  extends BaseSubscribable
+  implements ThreadListRuntimeCore
+{
+  private _mainThread: LocalThreadRuntimeCore;
+  constructor(_threadFactory: LocalThreadFactory) {
+    super();
 
-  public get newThreadId() {
-    return this._newThreadId;
+    this._mainThread = _threadFactory();
   }
 
-  public get threadIds() {
-    return this._threadIds;
+  public getMainThreadRuntimeCore() {
+    return this._mainThread;
   }
 
-  public get archivedThreadIds() {
-    return this._archivedThreadIds;
+  public get newThreadId(): string {
+    throw new Error("Method not implemented.");
   }
 
-  private _mainThreadId!: string;
-
-  public get mainThreadId(): string {
-    return this._mainThreadId;
+  public get threadIds(): readonly string[] {
+    throw EMPTY_ARRAY;
   }
 
-  constructor(private _threadFactory: LocalThreadFactory) {
-    this.switchToNewThread();
+  public get archivedThreadIds(): readonly string[] {
+    throw EMPTY_ARRAY;
   }
 
-  public getMainThreadRuntimeCore() {
-    const result = this._threadData.get(this._mainThreadId)?.runtime;
-    if (!result)
-      throw new Error("Main thread not found. This is a bug in assistant-ui.");
-    return result;
+  public get mainThreadId(): string {
+    return "__DEFAULT_ID__";
   }
 
-  public getThreadRuntimeCore(threadId: string) {
-    const result = this._threadData.get(threadId)?.runtime;
-    if (!result) throw new Error("Thread not found.");
-    return result;
+  public getThreadRuntimeCore(): never {
+    throw new Error("Method not implemented.");
   }
 
   public getLoadThreadsPromise(): Promise<void> {
-    return RESOLVED_PROMISE;
+    throw new Error("Method not implemented.");
   }
 
   public getItemById(threadId: string) {
-    return this._threadData.get(threadId);
+    if (threadId === this.mainThreadId) {
+      return {
+        status: "regular" as const,
+        threadId: this.mainThreadId,
+        remoteId: this.mainThreadId,
+        externalId: undefined,
+        title: undefined,
+        isMain: true,
+      };
+    }
+    throw new Error("Method not implemented");
   }
 
-  public async switchToThread(threadId: string): Promise<void> {
-    if (this._mainThreadId === threadId) return;
-
-    const data = this._threadData.get(threadId);
-    if (!data) throw new Error("Thread not found");
-
-    if (data.status === "archived") await this.unarchive(threadId);
-
-    this._mainThreadId = data.threadId;
-    this._notifySubscribers();
+  public async switchToThread(): Promise<void> {
+    throw new Error("Method not implemented.");
   }
 
   public switchToNewThread(): Promise<void> {
-    if (this._newThreadId === undefined) {
-      let threadId: string;
-      do {
-        threadId = generateId();
-      } while (this._threadData.has(threadId));
-
-      const runtime = this._threadFactory();
-      const dispose = runtime.unstable_on("initialize", () => {
-        dispose();
-        const data = this._threadData.get(threadId);
-        if (!data) throw new Error("Thread not found");
-
-        this._stateOp(threadId, "regular");
-      });
-      this._threadData.set(threadId, {
-        runtime,
-        status: "new",
-        threadId,
-      });
-      this._newThreadId = threadId;
-    }
-
-    this.switchToThread(this._newThreadId);
-    return Promise.resolve();
-  }
-
-  private async _stateOp(
-    threadId: string,
-    newState: "regular" | "archived" | "deleted",
-  ) {
-    const data = this._threadData.get(threadId);
-    if (!data) throw new Error("Thread not found");
-
-    const { status: lastState } = data;
-    if (lastState === newState) return;
-
-    // lastState
-    switch (lastState) {
-      case "new":
-        this._newThreadId = undefined;
-        break;
-      case "regular":
-        this._threadIds = this._threadIds.filter((t) => t !== threadId);
-        break;
-      case "archived":
-        this._archivedThreadIds = this._archivedThreadIds.filter(
-          (t) => t !== threadId,
-        );
-        break;
-
-      default: {
-        const _exhaustiveCheck: never = lastState;
-        throw new Error(`Unsupported state: ${_exhaustiveCheck}`);
-      }
-    }
-
-    // newState
-    switch (newState) {
-      case "regular":
-        this._threadIds = [data.threadId, ...this._threadIds];
-        break;
-
-      case "archived":
-        this._archivedThreadIds = [data.threadId, ...this._archivedThreadIds];
-        break;
-
-      case "deleted":
-        this._threadData.delete(threadId);
-        break;
-
-      default: {
-        const _exhaustiveCheck: never = newState;
-        throw new Error(`Unsupported state: ${_exhaustiveCheck}`);
-      }
-    }
-
-    if (newState !== "deleted") {
-      this._threadData.set(threadId, {
-        ...data,
-        status: newState,
-      });
-    }
-
-    if (
-      threadId === this._mainThreadId &&
-      (newState === "archived" || newState === "deleted")
-    ) {
-      const lastThreadId = this._threadIds[0];
-      if (lastThreadId) {
-        await this.switchToThread(lastThreadId);
-      } else {
-        await this.switchToNewThread();
-      }
-    } else {
-      this._notifySubscribers();
-    }
+    throw new Error("Method not implemented.");
   }
 
-  public rename(threadId: string, newTitle: string): Promise<void> {
-    const data = this._threadData.get(threadId);
-    if (!data) throw new Error("Thread not found");
-
-    this._threadData.set(threadId, {
-      ...data,
-      title: newTitle,
-    });
-    this._notifySubscribers();
-    return Promise.resolve();
+  public rename(): Promise<void> {
+    throw new Error("Method not implemented.");
   }
 
-  public archive(threadId: string): Promise<void> {
-    const data = this._threadData.get(threadId);
-    if (!data) throw new Error("Thread not found");
-    if (data.status !== "regular")
-      throw new Error("Thread is not yet initialized or already archived");
-
-    this._stateOp(threadId, "archived");
-    return Promise.resolve();
+  public archive(): Promise<void> {
+    throw new Error("Method not implemented.");
   }
 
-  public unarchive(threadId: string): Promise<void> {
-    const data = this._threadData.get(threadId);
-    if (!data) throw new Error("Thread not found");
-    if (data.status !== "archived") throw new Error("Thread is not archived");
-
-    this._stateOp(threadId, "regular");
-    return Promise.resolve();
+  public unarchive(): Promise<void> {
+    throw new Error("Method not implemented.");
   }
 
-  public delete(threadId: string): Promise<void> {
-    const data = this._threadData.get(threadId);
-    if (!data) throw new Error("Thread not found");
-    if (data.status !== "regular" && data.status !== "archived")
-      throw new Error("Thread is not yet initialized");
-
-    this._stateOp(threadId, "deleted");
-    return Promise.resolve();
+  public delete(): Promise<void> {
+    throw new Error("Method not implemented.");
   }
 
   public initialize(): never {
     throw new Error("Method not implemented.");
   }
 
-  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();
+  public generateTitle(): never {
+    throw new Error("Method not implemented.");
   }
 }

package/src/runtimes/local/LocalThreadRuntimeCore.tsx

@@ -249,6 +249,9 @@ export class LocalThreadRuntimeCore
         context,
         config: context,
         unstable_assistantMessageId: message.id,
+        unstable_getMessage() {
+          return message;
+        },
       });
 
       // handle async iterator for streaming results

package/src/runtimes/local/useLocalRuntime.tsx

@@ -10,6 +10,8 @@ import {
 } from "../../api/AssistantRuntime";
 import { ThreadRuntimeImpl } from "../../internal";
 import { useRuntimeAdapters } from "../adapters/RuntimeAdapterProvider";
+import { useRemoteThreadListRuntime } from "../remote-thread-list/useRemoteThreadListRuntime";
+import { useCloudThreadListAdapter } from "../remote-thread-list/adapter/cloud";
 
 export type LocalRuntime = AssistantRuntime & {
   reset: (options?: Parameters<LocalRuntimeCore["reset"]>[0]) => void;
@@ -34,9 +36,9 @@ class LocalRuntimeImpl extends AssistantRuntimeImpl implements LocalRuntime {
   }
 }
 
-export const useLocalRuntime = (
+const useLocalThreadRuntime = (
   adapter: ChatModelAdapter,
-  { initialMessages, ...options }: LocalRuntimeOptions = {},
+  { initialMessages, ...options }: LocalRuntimeOptions,
 ) => {
   const { modelContext, ...threadListAdapters } = useRuntimeAdapters() ?? {};
   const opt = useMemo(
@@ -65,3 +67,14 @@ export const useLocalRuntime = (
 
   return useMemo(() => LocalRuntimeImpl.create(runtime), [runtime]);
 };
+
+export const useLocalRuntime = (
+  adapter: ChatModelAdapter,
+  { cloud, ...options }: LocalRuntimeOptions = {},
+) => {
+  const cloudAdapter = useCloudThreadListAdapter({ cloud });
+  return useRemoteThreadListRuntime({
+    runtimeHook: () => useLocalThreadRuntime(adapter, options),
+    adapter: cloudAdapter,
+  });
+};

package/src/runtimes/remote-thread-list/RemoteThreadListHookInstanceManager.tsx

@@ -121,6 +121,13 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
       return runtime.threads.main.unstable_on("initialize", () => {
         if (threadListItemRuntime.getState().status === "new") {
           threadListItemRuntime.initialize();
+
+          // auto generate a title after first run
+          const dispose = runtime.thread.unstable_on("run-end", () => {
+            dispose();
+
+            threadListItemRuntime.generateTitle();
+          });
         }
       });
     }, [runtime, threadListItemRuntime]);