@assistant-ui/react-ai-sdk 0.10.14 → 0.10.16

package/src/dangerous-in-browser/createEdgeRuntimeAPI.ts

@@ -0,0 +1,145 @@
+import {
+  LanguageModelV1,
+  LanguageModelV1ToolChoice,
+  LanguageModelV1FunctionTool,
+  LanguageModelV1Prompt,
+  LanguageModelV1CallOptions,
+} from "@ai-sdk/provider";
+import { toLanguageModelMessages } from "../converters/toLanguageModelMessages";
+import { toLanguageModelTools } from "../converters/toLanguageModelTools";
+import { AssistantStreamChunk } from "assistant-stream";
+import { LanguageModelV1StreamDecoder } from "assistant-stream/ai-sdk";
+import { ThreadMessage, Tool } from "@assistant-ui/react";
+
+export type LanguageModelV1CallSettings = {
+  maxTokens?: number;
+  temperature?: number;
+  topP?: number;
+  presencePenalty?: number;
+  frequencyPenalty?: number;
+  seed?: number;
+  headers?: Record<string, string>;
+};
+
+export type LanguageModelConfig = {
+  apiKey?: string;
+  baseUrl?: string;
+  modelName?: string;
+};
+
+export type EdgeRuntimeRequestOptions = {
+  system?: string | undefined;
+  messages: readonly ThreadMessage[];
+  runConfig?: {
+    custom?: Record<string, unknown>;
+  };
+  tools?: any[];
+  unstable_assistantMessageId?: string;
+  state?: unknown;
+} & LanguageModelV1CallSettings &
+  LanguageModelConfig;
+
+type LanguageModelCreator = (
+  config: LanguageModelConfig,
+) => Promise<LanguageModelV1> | LanguageModelV1;
+
+type ThreadStep = {
+  readonly messageId?: string;
+  readonly usage?:
+    | {
+        readonly promptTokens: number;
+        readonly completionTokens: number;
+      }
+    | undefined;
+};
+
+type FinishResult = {
+  messages: readonly ThreadMessage[];
+  metadata: {
+    steps: readonly ThreadStep[];
+  };
+};
+
+export type CreateEdgeRuntimeAPIOptions = LanguageModelV1CallSettings & {
+  model: LanguageModelV1 | LanguageModelCreator;
+  system?: string;
+  tools?: Record<string, Tool<any, any>>;
+  toolChoice?: LanguageModelV1ToolChoice;
+  onFinish?: (result: FinishResult) => void;
+};
+
+type GetEdgeRuntimeStreamOptions = {
+  abortSignal: AbortSignal;
+  requestData: EdgeRuntimeRequestOptions;
+  options: CreateEdgeRuntimeAPIOptions;
+};
+
+export const getEdgeRuntimeStream = async ({
+  abortSignal,
+  requestData: request,
+  options: {
+    model: modelOrCreator,
+    system: serverSystem,
+    tools: serverTools = {},
+    toolChoice,
+    onFinish,
+    ...unsafeSettings
+  },
+}: GetEdgeRuntimeStreamOptions): Promise<
+  ReadableStream<AssistantStreamChunk>
+> => {
+  const model =
+    typeof modelOrCreator === "function"
+      ? await modelOrCreator(request)
+      : modelOrCreator;
+
+  const messages: LanguageModelV1Prompt = toLanguageModelMessages(
+    request.messages,
+  );
+
+  if (serverSystem || request.system) {
+    messages.unshift({
+      role: "system",
+      content: serverSystem || request.system!,
+    });
+  }
+
+  const allTools = serverTools || {};
+  const tools: LanguageModelV1FunctionTool[] = toLanguageModelTools(allTools);
+  const callOptions: LanguageModelV1CallOptions = {
+    inputFormat: "messages",
+    mode:
+      tools.length > 0
+        ? {
+            type: "regular",
+            tools,
+            ...(toolChoice && { toolChoice }),
+          }
+        : { type: "regular" },
+    prompt: messages,
+    ...(request.maxTokens !== undefined && { maxTokens: request.maxTokens }),
+    ...(request.temperature !== undefined && {
+      temperature: request.temperature,
+    }),
+    ...(request.topP !== undefined && { topP: request.topP }),
+    ...(request.presencePenalty !== undefined && {
+      presencePenalty: request.presencePenalty,
+    }),
+    ...(request.frequencyPenalty !== undefined && {
+      frequencyPenalty: request.frequencyPenalty,
+    }),
+    ...(request.seed !== undefined && { seed: request.seed }),
+    ...(request.headers && {
+      headers: Object.fromEntries(
+        Object.entries(request.headers).filter(([_, v]) => v !== undefined),
+      ) as Record<string, string>,
+    }),
+    abortSignal,
+    ...Object.fromEntries(
+      Object.entries(unsafeSettings).filter(([_, v]) => v !== undefined),
+    ),
+  };
+
+  const result = await model.doStream(callOptions);
+  return result.stream.pipeThrough(new LanguageModelV1StreamDecoder());
+};

package/src/dangerous-in-browser/index.ts

@@ -0,0 +1,4 @@
+export {
+  useDangerousInBrowserRuntime,
+  type DangerousInBrowserRuntimeOptions,
+} from "./useDangerousInBrowserRuntime";

package/src/dangerous-in-browser/useDangerousInBrowserRuntime.ts

@@ -0,0 +1,26 @@
+"use client";
+
+import { useState } from "react";
+import {
+  DangerousInBrowserAdapter,
+  DangerousInBrowserAdapterOptions,
+} from "./DangerousInBrowserAdapter";
+import {
+  INTERNAL,
+  LocalRuntimeOptions,
+  useLocalRuntime,
+} from "@assistant-ui/react";
+
+const { splitLocalRuntimeOptions } = INTERNAL;
+
+export type DangerousInBrowserRuntimeOptions =
+  DangerousInBrowserAdapterOptions & LocalRuntimeOptions;
+
+export const useDangerousInBrowserRuntime = (
+  options: DangerousInBrowserRuntimeOptions,
+) => {
+  const { localRuntimeOptions, otherOptions } =
+    splitLocalRuntimeOptions(options);
+  const [adapter] = useState(() => new DangerousInBrowserAdapter(otherOptions));
+  return useLocalRuntime(adapter, localRuntimeOptions);
+};

package/src/index.ts

@@ -2,12 +2,6 @@ export * from "./rsc";
 export * from "./ui";
 export * from "./useChatRuntime";
 export * from "./useCloudRuntime";
-
-export {
-  toLanguageModelMessages,
-  toLanguageModelTools,
-  fromLanguageModelMessages,
-  fromLanguageModelTools,
-  useDangerousInBrowserRuntime,
-} from "@assistant-ui/react-edge";
+export * from "./converters";
+export * from "./dangerous-in-browser";
 export { frontendTools } from "./frontendTools";

package/src/ui/utils/convertMessage.ts

@@ -1,12 +1,12 @@
 import { Message } from "@ai-sdk/ui-utils";
 import {
   unstable_createMessageConverter,
-  type ReasoningContentPart,
-  type ToolCallContentPart,
-  type TextContentPart,
+  type ReasoningMessagePart,
+  type ToolCallMessagePart,
+  type TextMessagePart,
   type CompleteAttachment,
-  type SourceContentPart,
-  type FileContentPart,
+  type SourceMessagePart,
+  type FileMessagePart,
 } from "@assistant-ui/react";
 
 const convertParts = (message: Message) => {
@@ -20,7 +20,7 @@ const convertParts = (message: Message) => {
             return {
               type: "text",
               text: part.text,
-            } satisfies TextContentPart;
+            } satisfies TextMessagePart;
           case "tool-invocation":
             return {
               type: "tool-call",
@@ -31,23 +31,23 @@ const convertParts = (message: Message) => {
               result:
                 part.toolInvocation.state === "result" &&
                 part.toolInvocation.result,
-            } satisfies ToolCallContentPart;
+            } satisfies ToolCallMessagePart;
           case "reasoning":
             return {
               type: "reasoning",
               text: part.reasoning,
-            } satisfies ReasoningContentPart;
+            } satisfies ReasoningMessagePart;
           case "source":
             return {
               type: "source",
               ...part.source,
-            } satisfies SourceContentPart;
+            } satisfies SourceMessagePart;
           case "file":
             return {
               type: "file",
               data: part.data,
               mimeType: part.mimeType,
-            } satisfies FileContentPart;
+            } satisfies FileMessagePart;
           default: {
             const _unsupported: never = type;
             throw new Error(
@@ -63,7 +63,7 @@ const convertParts = (message: Message) => {
         {
           type: "text",
           text: message.content,
-        } satisfies TextContentPart,
+        } satisfies TextMessagePart,
       ]
     : [];
 };
@@ -117,7 +117,7 @@ export const AISDKMessageConverter = unstable_createMessageConverter(
       case "data": {
         type MaybeSupportedDataMessage =
           | { type?: "unsafe_other" }
-          | ToolCallContentPart
+          | ToolCallMessagePart
           | {
               type: "tool-result";
               toolCallId: string;

package/src/ui/utils/sliceMessagesUntil.tsx

@@ -9,7 +9,7 @@ export const sliceMessagesUntil = (
   let messageIdx = messages.findIndex((m) => m.id === messageId);
   if (messageIdx === -1)
     throw new Error(
-      "useVercelAIThreadState: Message not found. This is liekly an internal bug in assistant-ui.",
+      "useVercelAIThreadState: Message not found. This is likely an internal bug in assistant-ui.",
     );
 
   while (messages[messageIdx + 1]?.role === "assistant") {

package/src/useChatRuntime.ts

@@ -1,13 +1,157 @@
-import { EdgeRuntimeOptions, useEdgeRuntime } from "@assistant-ui/react-edge";
-
-export type UseChatRuntimeOptions = Omit<
-  EdgeRuntimeOptions,
-  "unstable_AISDKInterop"
->;
-
-export const useChatRuntime = (options: UseChatRuntimeOptions) => {
-  return useEdgeRuntime({
-    ...options,
-    unstable_AISDKInterop: "v2",
-  });
+"use client";
+
+import { toLanguageModelMessages } from "./converters";
+import {
+  AssistantRuntime,
+  ChatModelAdapter,
+  ChatModelRunOptions,
+  INTERNAL,
+  LocalRuntimeOptions,
+  ThreadMessage,
+  Tool,
+  useLocalRuntime,
+} from "@assistant-ui/react";
+import { z } from "zod";
+import zodToJsonSchema from "zod-to-json-schema";
+import { JSONSchema7 } from "json-schema";
+import {
+  AssistantMessageAccumulator,
+  DataStreamDecoder,
+  unstable_toolResultStream,
+} from "assistant-stream";
+import { asAsyncIterableStream } from "assistant-stream/utils";
+
+const { splitLocalRuntimeOptions } = INTERNAL;
+
+type HeadersValue = Record<string, string> | Headers;
+
+export type UseChatRuntimeOptions = {
+  api: string;
+  onResponse?: (response: Response) => void | Promise<void>;
+  onFinish?: (message: ThreadMessage) => void;
+  onError?: (error: Error) => void;
+  onCancel?: () => void;
+  credentials?: RequestCredentials;
+  headers?: HeadersValue | (() => Promise<HeadersValue>);
+  body?: object;
+  sendExtraMessageFields?: boolean;
+} & LocalRuntimeOptions;
+
+type ChatRuntimeRequestOptions = {
+  messages: any[];
+  tools: any;
+  system?: string | undefined;
+  runConfig?: any;
+  unstable_assistantMessageId?: string;
+  state?: any;
+};
+
+const toAISDKTools = (tools: Record<string, Tool>) => {
+  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,
+      },
+    ]),
+  );
+};
+
+const getEnabledTools = (tools: Record<string, Tool>) => {
+  return Object.fromEntries(
+    Object.entries(tools).filter(
+      ([, tool]) => !tool.disabled && tool.type !== "backend",
+    ),
+  );
+};
+
+class ChatRuntimeAdapter implements ChatModelAdapter {
+  constructor(
+    private options: Omit<UseChatRuntimeOptions, keyof LocalRuntimeOptions>,
+  ) {}
+
+  async *run({
+    messages,
+    runConfig,
+    abortSignal,
+    context,
+    unstable_assistantMessageId,
+    unstable_getMessage,
+  }: ChatModelRunOptions) {
+    const headersValue =
+      typeof this.options.headers === "function"
+        ? await this.options.headers()
+        : this.options.headers;
+
+    abortSignal.addEventListener(
+      "abort",
+      () => {
+        if (!abortSignal.reason?.detach) this.options.onCancel?.();
+      },
+      { once: true },
+    );
+
+    const headers = new Headers(headersValue);
+    headers.set("Content-Type", "application/json");
+
+    const result = await fetch(this.options.api, {
+      method: "POST",
+      headers,
+      credentials: this.options.credentials ?? "same-origin",
+      body: JSON.stringify({
+        system: context.system,
+        messages: toLanguageModelMessages(messages, {
+          unstable_includeId: this.options.sendExtraMessageFields,
+        }) as ChatRuntimeRequestOptions["messages"],
+        tools: toAISDKTools(
+          getEnabledTools(context.tools ?? {}),
+        ) as unknown as ChatRuntimeRequestOptions["tools"],
+        ...(unstable_assistantMessageId ? { unstable_assistantMessageId } : {}),
+        runConfig,
+        state: unstable_getMessage().metadata.unstable_state || undefined,
+        ...context.callSettings,
+        ...context.config,
+        ...this.options.body,
+      } satisfies ChatRuntimeRequestOptions),
+      signal: abortSignal,
+    });
+
+    await this.options.onResponse?.(result);
+
+    try {
+      if (!result.ok) {
+        throw new Error(`Status ${result.status}: ${await result.text()}`);
+      }
+      if (!result.body) {
+        throw new Error("Response body is null");
+      }
+
+      const stream = result.body
+        .pipeThrough(new DataStreamDecoder())
+        .pipeThrough(unstable_toolResultStream(context.tools, abortSignal))
+        .pipeThrough(new AssistantMessageAccumulator());
+
+      yield* asAsyncIterableStream(stream);
+
+      this.options.onFinish?.(unstable_getMessage());
+    } catch (error: unknown) {
+      this.options.onError?.(error as Error);
+      throw error;
+    }
+  }
+}
+
+export const useChatRuntime = (
+  options: UseChatRuntimeOptions,
+): AssistantRuntime => {
+  const { localRuntimeOptions, otherOptions } =
+    splitLocalRuntimeOptions(options);
+
+  return useLocalRuntime(
+    new ChatRuntimeAdapter(otherOptions),
+    localRuntimeOptions,
+  );
 };