@assistant-ui/react 0.5.91 → 0.5.93

package/src/runtimes/edge/converters/toCoreMessages.ts

@@ -0,0 +1,61 @@
+import { ThreadMessage, CoreMessage } from "../../../types";
+
+type CoreMessageWithConditionalId<T extends boolean> = T extends false
+  ? CoreMessage
+  : CoreMessage & { unstable_id?: string };
+
+export const toCoreMessages = <T extends boolean = false>(
+  messages: ThreadMessage[],
+  options: { unstable_includeId?: T | undefined } = {},
+): CoreMessageWithConditionalId<T>[] => {
+  return messages.map((message) => toCoreMessage(message, options));
+};
+
+export const toCoreMessage = <T extends boolean = false>(
+  message: ThreadMessage,
+  options: { unstable_includeId?: T | undefined } = {},
+): CoreMessageWithConditionalId<T> => {
+  const includeId = options.unstable_includeId ?? false;
+  const role = message.role;
+  switch (role) {
+    case "assistant":
+      return {
+        role,
+        content: message.content.map((part) => {
+          if (part.type === "ui") throw new Error("UI parts are not supported");
+          if (part.type === "tool-call") {
+            const { argsText, ...rest } = part;
+            return rest;
+          }
+          return part;
+        }),
+        ...(includeId ? { unstable_id: message.id } : {}),
+      };
+
+    case "user":
+      return {
+        role,
+        content: [
+          ...message.content.map((part) => {
+            if (part.type === "ui")
+              throw new Error("UI parts are not supported");
+            return part;
+          }),
+          ...message.attachments.map((a) => a.content).flat(),
+        ],
+        ...(includeId ? { unstable_id: message.id } : {}),
+      };
+
+    case "system":
+      return {
+        role,
+        content: message.content,
+        ...(includeId ? { unstable_id: message.id } : {}),
+      };
+
+    default: {
+      const unsupportedRole: never = role;
+      throw new Error(`Unknown message role: ${unsupportedRole}`);
+    }
+  }
+};

package/src/runtimes/edge/converters/toLanguageModelMessages.ts

@@ -0,0 +1,143 @@
+import {
+  LanguageModelV1ImagePart,
+  LanguageModelV1Message,
+  LanguageModelV1TextPart,
+  LanguageModelV1ToolCallPart,
+  LanguageModelV1ToolResultPart,
+} from "@ai-sdk/provider";
+import {
+  CoreMessage,
+  ThreadMessage,
+  TextContentPart,
+  CoreToolCallContentPart,
+} from "../../../types/AssistantTypes";
+
+const assistantMessageSplitter = () => {
+  const stash: LanguageModelV1Message[] = [];
+  let assistantMessage = {
+    role: "assistant" as const,
+    content: [] as (LanguageModelV1TextPart | LanguageModelV1ToolCallPart)[],
+  };
+  let toolMessage = {
+    role: "tool" as const,
+    content: [] as LanguageModelV1ToolResultPart[],
+  };
+
+  return {
+    addTextContentPart: (part: TextContentPart) => {
+      if (toolMessage.content.length > 0) {
+        stash.push(assistantMessage);
+        stash.push(toolMessage);
+
+        assistantMessage = {
+          role: "assistant" as const,
+          content: [] as (
+            | LanguageModelV1TextPart
+            | LanguageModelV1ToolCallPart
+          )[],
+        };
+
+        toolMessage = {
+          role: "tool" as const,
+          content: [] as LanguageModelV1ToolResultPart[],
+        };
+      }
+
+      assistantMessage.content.push(part);
+    },
+    addToolCallPart: (part: CoreToolCallContentPart) => {
+      assistantMessage.content.push({
+        type: "tool-call",
+        toolCallId: part.toolCallId,
+        toolName: part.toolName,
+        args: part.args,
+      });
+
+      toolMessage.content.push({
+        type: "tool-result",
+        toolCallId: part.toolCallId,
+        toolName: part.toolName,
+        result: part.result ?? "<no result>",
+        isError: part.isError ?? false,
+      });
+    },
+    getMessages: () => {
+      if (toolMessage.content.length > 0) {
+        return [...stash, assistantMessage, toolMessage];
+      }
+
+      return [...stash, assistantMessage];
+    },
+  };
+};
+
+export function toLanguageModelMessages(
+  message: readonly CoreMessage[] | readonly ThreadMessage[],
+): LanguageModelV1Message[] {
+  return message.flatMap((message) => {
+    const role = message.role;
+    switch (role) {
+      case "system": {
+        return [{ role: "system", content: message.content[0].text }];
+      }
+
+      case "user": {
+        const msg: LanguageModelV1Message = {
+          role: "user",
+          content: message.content.map(
+            (part): LanguageModelV1TextPart | LanguageModelV1ImagePart => {
+              const type = part.type;
+              switch (type) {
+                case "text": {
+                  return part;
+                }
+
+                case "image": {
+                  return {
+                    type: "image",
+                    image: new URL(part.image),
+                  };
+                }
+
+                default: {
+                  const unhandledType: "ui" | "audio" = type;
+                  throw new Error(
+                    `Unspported content part type: ${unhandledType}`,
+                  );
+                }
+              }
+            },
+          ),
+        };
+        return [msg];
+      }
+
+      case "assistant": {
+        const splitter = assistantMessageSplitter();
+        for (const part of message.content) {
+          const type = part.type;
+          switch (type) {
+            case "text": {
+              splitter.addTextContentPart(part);
+              break;
+            }
+            case "tool-call": {
+              splitter.addToolCallPart(part);
+              break;
+            }
+            default: {
+              const unhandledType: "ui" = type;
+              throw new Error(`Unhandled content part type: ${unhandledType}`);
+            }
+          }
+        }
+        return splitter.getMessages();
+      }
+
+      default: {
+        const unhandledRole: never = role;
+        throw new Error(`Unknown message role: ${unhandledRole}`);
+      }
+    }
+  });
+}

package/src/runtimes/edge/converters/toLanguageModelTools.ts

@@ -0,0 +1,18 @@
+import { LanguageModelV1FunctionTool } from "@ai-sdk/provider";
+import { JSONSchema7 } from "json-schema";
+import { z } from "zod";
+import zodToJsonSchema from "zod-to-json-schema";
+import { Tool } from "../../../types/ModelConfigTypes";
+
+export const toLanguageModelTools = (
+  tools: Record<string, Tool<any, any>>,
+): LanguageModelV1FunctionTool[] => {
+  return Object.entries(tools).map(([name, tool]) => ({
+    type: "function",
+    name,
+    ...(tool.description ? { description: tool.description } : undefined),
+    parameters: (tool.parameters instanceof z.ZodType
+      ? zodToJsonSchema(tool.parameters)
+      : tool.parameters) as JSONSchema7,
+  }));
+};

package/src/runtimes/edge/createEdgeRuntimeAPI.ts

@@ -0,0 +1,253 @@
+import {
+  LanguageModelV1,
+  LanguageModelV1ToolChoice,
+  LanguageModelV1FunctionTool,
+  LanguageModelV1Prompt,
+  LanguageModelV1CallOptions,
+} from "@ai-sdk/provider";
+import {
+  CoreMessage,
+  ThreadMessage,
+  ThreadStep,
+} from "../../types/AssistantTypes";
+import { assistantEncoderStream } from "./streams/assistantEncoderStream";
+import { EdgeRuntimeRequestOptionsSchema } from "./EdgeRuntimeRequestOptions";
+import { toLanguageModelMessages } from "./converters/toLanguageModelMessages";
+import { Tool } from "../../types";
+import { toLanguageModelTools } from "./converters/toLanguageModelTools";
+import {
+  toolResultStream,
+  ToolResultStreamPart,
+} from "./streams/toolResultStream";
+import { runResultStream } from "./streams/runResultStream";
+import {
+  LanguageModelConfig,
+  LanguageModelV1CallSettings,
+  LanguageModelV1CallSettingsSchema,
+} from "../../types/ModelConfigTypes";
+import { ChatModelRunResult } from "../local/ChatModelAdapter";
+import { toCoreMessage } from "./converters/toCoreMessages";
+import { streamPartEncoderStream } from "./streams/utils/streamPartEncoderStream";
+import { z } from "zod";
+
+type FinishResult = {
+  messages: CoreMessage[];
+  metadata: {
+    /**
+     * @deprecated Use `steps` instead. This field will be removed in v0.6.
+     */
+    roundtrips: ThreadStep[];
+    steps: ThreadStep[];
+  };
+};
+
+type LanguageModelCreator = (
+  config: LanguageModelConfig,
+) => Promise<LanguageModelV1> | LanguageModelV1;
+
+export type CreateEdgeRuntimeAPIOptions = LanguageModelV1CallSettings & {
+  model: LanguageModelV1 | LanguageModelCreator;
+  system?: string;
+  tools?: Record<string, Tool<any, any>>;
+  toolChoice?: LanguageModelV1ToolChoice;
+  onFinish?: (result: FinishResult) => void;
+};
+
+const voidStream = () => {
+  return new WritableStream({
+    abort(reason) {
+      console.error("Server stream processing aborted:", reason);
+    },
+  });
+};
+
+type GetEdgeRuntimeStreamOptions = {
+  abortSignal: AbortSignal;
+  requestData: z.infer<typeof EdgeRuntimeRequestOptionsSchema>;
+  options: CreateEdgeRuntimeAPIOptions;
+};
+
+export const getEdgeRuntimeStream = async ({
+  abortSignal,
+  requestData: unsafeRequest,
+  options: {
+    model: modelOrCreator,
+    system: serverSystem,
+    tools: serverTools = {},
+    toolChoice,
+    onFinish,
+    ...unsafeSettings
+  },
+}: GetEdgeRuntimeStreamOptions) => {
+  const settings = LanguageModelV1CallSettingsSchema.parse(unsafeSettings);
+  const lmServerTools = toLanguageModelTools(serverTools);
+  const hasServerTools = Object.values(serverTools).some((v) => !!v.execute);
+
+  const {
+    system: clientSystem,
+    tools: clientTools = [],
+    messages,
+    apiKey,
+    baseUrl,
+    modelName,
+    ...callSettings
+  } = EdgeRuntimeRequestOptionsSchema.parse(unsafeRequest);
+
+  const systemMessages = [];
+  if (serverSystem) systemMessages.push(serverSystem);
+  if (clientSystem) systemMessages.push(clientSystem);
+  const system = systemMessages.join("\n\n");
+
+  for (const clientTool of clientTools) {
+    if (serverTools?.[clientTool.name]) {
+      throw new Error(
+        `Tool ${clientTool.name} was defined in both the client and server tools. This is not allowed.`,
+      );
+    }
+  }
+
+  const model =
+    typeof modelOrCreator === "function"
+      ? await modelOrCreator({ apiKey, baseUrl, modelName })
+      : modelOrCreator;
+
+  let stream: ReadableStream<ToolResultStreamPart>;
+  const streamResult = await streamMessage({
+    ...(settings as Partial<StreamMessageOptions>),
+    ...callSettings,
+
+    model,
+    abortSignal,
+
+    ...(!!system ? { system } : undefined),
+    messages,
+    tools: lmServerTools.concat(clientTools as LanguageModelV1FunctionTool[]),
+    ...(toolChoice ? { toolChoice } : undefined),
+  });
+  stream = streamResult.stream;
+
+  // add tool results if we have server tools
+  const canExecuteTools = hasServerTools && toolChoice?.type !== "none";
+  if (canExecuteTools) {
+    stream = stream.pipeThrough(toolResultStream(serverTools, abortSignal));
+  }
+
+  if (canExecuteTools || onFinish) {
+    // tee the stream to process server tools and onFinish asap
+    const tees = stream.tee();
+    stream = tees[0];
+    let serverStream = tees[1];
+
+    if (onFinish) {
+      let lastChunk: ChatModelRunResult;
+      serverStream = serverStream.pipeThrough(runResultStream()).pipeThrough(
+        new TransformStream({
+          transform(chunk) {
+            lastChunk = chunk;
+          },
+          flush() {
+            if (!lastChunk?.status || lastChunk.status.type === "running")
+              return;
+
+            const resultingMessages = [
+              ...messages,
+              toCoreMessage({
+                role: "assistant",
+                content: lastChunk.content,
+              } as ThreadMessage),
+            ];
+            onFinish({
+              messages: resultingMessages,
+              metadata: {
+                // TODO
+                // eslint-disable-next-line @typescript-eslint/no-non-null-asserted-optional-chain
+                roundtrips: lastChunk.metadata?.steps!,
+                // eslint-disable-next-line @typescript-eslint/no-non-null-asserted-optional-chain
+                steps: lastChunk.metadata?.steps!,
+              },
+            });
+          },
+        }),
+      );
+    }
+
+    // drain the server stream
+    serverStream.pipeTo(voidStream()).catch((e) => {
+      console.error("Server stream processing error:", e);
+    });
+  }
+
+  return stream;
+};
+
+export declare namespace getEdgeRuntimeResponse {
+  export type { GetEdgeRuntimeStreamOptions as Options };
+}
+
+export const getEdgeRuntimeResponse = async (
+  options: getEdgeRuntimeResponse.Options,
+) => {
+  const stream = await getEdgeRuntimeStream(options);
+  return new Response(
+    stream
+      .pipeThrough(assistantEncoderStream())
+      .pipeThrough(streamPartEncoderStream()),
+    {
+      headers: {
+        "Content-Type": "text/plain; charset=utf-8",
+      },
+    },
+  );
+};
+
+export const createEdgeRuntimeAPI = (options: CreateEdgeRuntimeAPIOptions) => ({
+  POST: async (request: Request) =>
+    getEdgeRuntimeResponse({
+      abortSignal: request.signal,
+      requestData: await request.json(),
+      options,
+    }),
+});
+
+type StreamMessageOptions = LanguageModelV1CallSettings & {
+  model: LanguageModelV1;
+  system?: string;
+  messages: CoreMessage[];
+  tools?: LanguageModelV1FunctionTool[];
+  toolChoice?: LanguageModelV1ToolChoice;
+  abortSignal: AbortSignal;
+};
+
+async function streamMessage({
+  model,
+  system,
+  messages,
+  tools,
+  toolChoice,
+  ...options
+}: StreamMessageOptions) {
+  return model.doStream({
+    inputFormat: "messages",
+    mode: {
+      type: "regular",
+      ...(tools ? { tools } : undefined),
+      ...(toolChoice ? { toolChoice } : undefined),
+    },
+    prompt: convertToLanguageModelPrompt(system, messages),
+    ...(options as Partial<LanguageModelV1CallOptions>),
+  });
+}
+
+export function convertToLanguageModelPrompt(
+  system: string | undefined,
+  messages: CoreMessage[],
+): LanguageModelV1Prompt {
+  const languageModelMessages: LanguageModelV1Prompt = [];
+
+  if (system != null) {
+    languageModelMessages.push({ role: "system", content: system });
+  }
+  languageModelMessages.push(...toLanguageModelMessages(messages));
+
+  return languageModelMessages;
+}

package/src/runtimes/edge/index.ts

@@ -0,0 +1,7 @@
+export * from "./converters";
+
+export * from "./streams/utils";
+
+export { useEdgeRuntime, type EdgeRuntimeOptions } from "./useEdgeRuntime";
+export { EdgeChatAdapter } from "./EdgeChatAdapter";
+export type { EdgeRuntimeRequestOptions } from "./EdgeRuntimeRequestOptions";