@convex-dev/agent 0.1.1 → 0.1.2-alpha.0

package/src/client/index.ts

@@ -4,32 +4,20 @@ import type {
   DeepPartial,
   GenerateObjectResult,
   GenerateTextResult,
-  JSONValue,
-  RepairTextFunction,
-  Schema,
   StepResult,
   StreamObjectResult,
   StreamTextResult,
-  TelemetrySettings,
-  Tool,
-  ToolChoice,
-  ToolExecutionOptions,
   ToolSet,
-  UIMessage,
-} from "ai";
-import type { ToolInvocationUIPart } from "@ai-sdk/ui-utils";
-import {
-  generateObject,
-  generateText,
-  streamObject,
-  streamText,
-  tool,
 } from "ai";
+import { generateObject, generateText, streamObject, streamText } from "ai";
 import { assert } from "convex-helpers";
-import { internalActionGeneric, internalMutationGeneric } from "convex/server";
+import {
+  internalActionGeneric,
+  internalMutationGeneric,
+  PaginationOptions,
+  PaginationResult,
+} from "convex/server";
 import { v } from "convex/values";
-import { z } from "zod";
-import { Mounts } from "../component/_generated/api.js";
 import {
   validateVectorDimension,
   type VectorDimension,
@@ -42,7 +30,6 @@ import {
   serializeNewMessagesInStep,
   serializeObjectResult,
   serializeStep,
-  toUIFilePart,
 } from "../mapping.js";
 import {
   DEFAULT_MESSAGE_RANGE,
@@ -51,132 +38,71 @@ import {
   isTool,
 } from "../shared.js";
 import {
-  type CallSettings,
   type MessageWithMetadata as InnerMessageWithMetadata,
+  type MessageStatus,
   type ProviderMetadata,
-  type ProviderOptions,
   type SearchOptions,
+  type StreamArgs,
   type Usage,
   vMessageWithMetadata,
   vSafeObjectArgs,
   vTextArgs,
 } from "../validators.js";
 import type {
+  AgentComponent,
+  ContextOptions,
+  GenerationOutputMetadata,
   OpaqueIds,
+  Options,
+  OurObjectArgs,
+  OurStreamObjectArgs,
   RunActionCtx,
   RunMutationCtx,
   RunQueryCtx,
-  UseApi,
+  StorageOptions,
+  StreamingTextArgs,
+  SyncStreamsReturnValue,
+  TextArgs,
+  Thread,
+  UsageHandler,
 } from "./types.js";
 
 import type { MessageDoc, ThreadDoc } from "../component/schema.js";
+import { createTool, wrapTools } from "./createTool.js";
+import {
+  DeltaStreamer,
+  mergeTransforms,
+  StreamingOptions,
+} from "./streaming.js";
 
 export { vMessageDoc, vThreadDoc } from "../component/schema.js";
-export { extractText, isTool };
-export type { Usage, ProviderMetadata, MessageDoc, ThreadDoc };
 export {
-  /** @deprecated Use vPaginationResult instead. */
-  paginationResultValidator,
-  paginationResultValidator as vPaginationResult,
+  vAssistantMessage,
   vContextOptions,
-  vUsage,
+  vMessage,
+  vPaginationResult,
   vProviderMetadata,
-  vUserMessage,
-  vAssistantMessage,
-  vToolMessage,
   vStorageOptions,
+  vStreamArgs,
   vSystemMessage,
-  vMessage,
+  vToolMessage,
+  vUsage,
+  vUserMessage,
 } from "../validators.js";
-
-/**
- * Options to configure what messages are fetched as context,
- * automatically with thread.generateText, or directly via search.
- */
-export type ContextOptions = {
-  /** @deprecated Use excludeToolMessages instead. */
-  includeToolCalls?: boolean;
-  /**
-   * Whether to include tool messages in the context.
-   * By default, tool calls and results are not included.
-   */
-  excludeToolMessages?: boolean;
-  /**
-   * How many recent messages to include. These are added after the search
-   * messages, and do not count against the search limit.
-   * Default: 100
-   */
-  recentMessages?: number;
-  /**
-   * Options for searching messages.
-   */
-  searchOptions?: {
-    /**
-     * The maximum number of messages to fetch. Default is 10.
-     */
-    limit: number;
-    /**
-     * Whether to use text search to find messages. Default is false.
-     */
-    textSearch?: boolean;
-    /**
-     * Whether to use vector search to find messages. Default is false.
-     * At least one of textSearch or vectorSearch must be true.
-     */
-    vectorSearch?: boolean;
-    /**
-     * What messages around the search results to include.
-     * Default: { before: 2, after: 1 }
-     * (two before, and one after each message found in the search)
-     * Note, this is after the limit is applied.
-     * By default this will quadruple the number of messages fetched.
-     */
-    messageRange?: { before: number; after: number };
-  };
-  /**
-   * Whether to search across other threads for relevant messages.
-   * By default, only the current thread is searched.
-   */
-  searchOtherThreads?: boolean;
+export { createTool, extractText, isTool };
+export type {
+  AgentComponent,
+  ContextOptions,
+  MessageDoc,
+  ProviderMetadata,
+  StorageOptions,
+  SyncStreamsReturnValue,
+  Thread,
+  ThreadDoc,
+  Usage,
+  UsageHandler,
 };
 
-/**
- * Options to configure the automatic saving of messages
- * when generating text / objects in a thread.
- */
-export type StorageOptions = {
-  /**
-   * Defaults to false, allowing you to pass in arbitrary context that will
-   * be in addition to automatically fetched content.
-   * Pass true to have all input messages saved to the thread history.
-   */
-  saveAllInputMessages?: boolean;
-  /** Defaults to true, saving the prompt, or last message passed to generateText. */
-  saveAnyInputMessages?: boolean;
-  /** Defaults to true. Whether to save messages generated while chatting. */
-  saveOutputMessages?: boolean;
-};
-
-export type GenerationOutputMetadata = { messageId?: string };
-
-type CoreMessageMaybeWithId = CoreMessage & { id?: string | undefined };
-
-export type UsageHandler = (
-  ctx: RunActionCtx,
-  args: {
-    userId: string | undefined;
-    threadId: string | undefined;
-    agentName: string | undefined;
-    usage: Usage;
-    // Often has more information, like cached token usage in the case of openai.
-    providerMetadata: ProviderMetadata | undefined;
-    model: string;
-    provider: string;
-  }
-) => void | Promise<void>;
-
-export type AgentComponent = UseApi<Mounts>;
-
 export class Agent<AgentTools extends ToolSet> {
   constructor(
     public component: AgentComponent,
@@ -406,6 +332,72 @@ export class Agent<AgentTools extends ToolSet> {
     };
   }
 
+  /**
+   * List messages from a thread.
+   * @param ctx A ctx object from a query, mutation, or action.
+   * @param args.threadId The thread to list messages from.
+   * @param args.paginationOpts Pagination options (e.g. via usePaginatedQuery).
+   * @param args.excludeToolMessages Whether to exclude tool messages.
+   *   False by default.
+   * @param args.statuses What statuses to include. All by default.
+   * @returns The MessageDoc's in a format compatible with usePaginatedQuery.
+   */
+  async listMessages(
+    ctx: RunQueryCtx,
+    args: {
+      threadId: string;
+      paginationOpts: PaginationOptions;
+      excludeToolMessages?: boolean;
+      statuses?: MessageStatus[];
+    }
+  ): Promise<PaginationResult<MessageDoc>> {
+    if (args.paginationOpts.numItems === 0) {
+      return {
+        page: [],
+        isDone: true,
+        continueCursor: args.paginationOpts.cursor ?? "",
+      };
+    }
+    return ctx.runQuery(this.component.messages.listMessagesByThreadId, {
+      order: "desc",
+      ...args,
+    });
+  }
+
+  /**
+   * A function that handles fetching stream deltas, used with the React hooks
+   * `useThreadMessages` or `useStreamingThreadMessages`.
+   * @param ctx A ctx object from a query, mutation, or action.
+   * @param args.threadId The thread to sync streams for.
+   * @param args.streamArgs The stream arguments with per-stream cursors.
+   * @returns The deltas for each stream from their existing cursor.
+   */
+  async syncStreams(
+    ctx: RunQueryCtx,
+    args: {
+      threadId: string;
+      streamArgs: StreamArgs | undefined;
+    }
+  ): Promise<SyncStreamsReturnValue | undefined> {
+    if (!args.streamArgs) return undefined;
+    if (args.streamArgs.kind === "list") {
+      return {
+        kind: "list",
+        messages: await ctx.runQuery(this.component.streams.list, {
+          threadId: args.threadId,
+        }),
+      };
+    } else {
+      return {
+        kind: "deltas",
+        deltas: await ctx.runQuery(this.component.streams.listDeltas, {
+          threadId: args.threadId,
+          cursors: args.streamArgs.cursors,
+        }),
+      };
+    }
+  }
+
   /**
    * Fetch the context messages for a thread.
    * @param ctx Either a query, mutation, or action ctx.
@@ -432,7 +424,7 @@ export class Agent<AgentTools extends ToolSet> {
     assert(args.userId || args.threadId, "Specify userId or threadId");
     // Fetch the latest messages from the thread
     let included: Set<string> | undefined;
-    const opts = this.mergedContextOptions(args.contextOptions);
+    const opts = this._mergedContextOptions(args.contextOptions);
     const contextMessages: MessageDoc[] = [];
     if (
       args.threadId &&
@@ -475,7 +467,7 @@ export class Agent<AgentTools extends ToolSet> {
           userId: opts?.searchOtherThreads ? args.userId : undefined,
           threadId: args.threadId,
           beforeMessageId: args.upToAndIncludingMessageId,
-          ...(await this.searchOptionsWithDefaults(opts, messagesToSearch)),
+          ...(await this._searchOptionsWithDefaults(opts, messagesToSearch)),
         }
       );
       // TODO: track what messages we used for context
@@ -537,6 +529,35 @@ export class Agent<AgentTools extends ToolSet> {
     return embeddings;
   }
 
+  async saveMessage(
+    ctx: RunMutationCtx,
+    args: {
+      threadId: string;
+      userId?: string;
+      metadata?: Omit<MessageWithMetadata, "message">;
+    } & (
+      | {
+          prompt?: undefined;
+          message: CoreMessage;
+        }
+      | {
+          prompt: string;
+          message?: undefined;
+        }
+    )
+  ) {
+    const { lastMessageId } = await this.saveMessages(ctx, {
+      threadId: args.threadId,
+      userId: args.userId,
+      messages:
+        args.prompt !== undefined
+          ? [{ role: "user", content: args.prompt }]
+          : [args.message],
+      metadata: args.metadata ? [args.metadata] : undefined,
+    });
+    return { messageId: lastMessageId };
+  }
+
   /**
    * Explicitly save messages associated with the thread (& user if provided)
    * @param ctx The ctx parameter to a mutation or action.
@@ -576,7 +597,7 @@ export class Agent<AgentTools extends ToolSet> {
     }
   ): Promise<{
     lastMessageId: string;
-    messageIds: string[];
+    messages: MessageDoc[];
   }> {
     const embeddings = await this.generateEmbeddings(args.messages);
     const result = await ctx.runMutation(this.component.messages.addMessages, {
@@ -597,7 +618,7 @@ export class Agent<AgentTools extends ToolSet> {
     });
     return {
       lastMessageId: result.messages.at(-1)!._id,
-      messageIds: result.messages.map((m) => m._id),
+      messages: result.messages,
     };
   }
 
@@ -622,11 +643,15 @@ export class Agent<AgentTools extends ToolSet> {
         failPendingSteps: v.optional(v.boolean()),
       },
       handler: async (ctx, args) => {
-        return this.saveMessages(ctx, {
+        const { lastMessageId, messages } = await this.saveMessages(ctx, {
           ...args,
           messages: args.messages.map((m) => m.message),
           metadata: args.messages.map(({ message: _, ...m }) => m),
         });
+        return {
+          lastMessageId,
+          messageIds: messages.map((m) => m._id),
+        };
       },
     });
   }
@@ -660,7 +685,7 @@ export class Agent<AgentTools extends ToolSet> {
        */
       provider?: string;
     }
-  ): Promise<void> {
+  ): Promise<MessageDoc[]> {
     const step = serializeStep(args.step as StepResult<ToolSet>);
     const messages = serializeNewMessagesInStep(args.step, {
       provider: args.provider ?? this.options.chat.provider,
@@ -669,13 +694,14 @@ export class Agent<AgentTools extends ToolSet> {
     const embeddings = await this.generateEmbeddings(
       messages.map((m) => m.message)
     );
-    await ctx.runMutation(this.component.messages.addStep, {
+    const saved = await ctx.runMutation(this.component.messages.addStep, {
       userId: args.userId,
       threadId: args.threadId,
       promptMessageId: args.promptMessageId,
       step: { step, messages, embeddings },
       failPendingSteps: false,
     });
+    return saved;
   }
 
   /**
@@ -752,7 +778,7 @@ export class Agent<AgentTools extends ToolSet> {
       args,
       { userId, threadId, ...options }
     );
-    const toolCtx = { ...ctx, userId, threadId, messageId };
+    const toolCtx = { ...ctx, userId, threadId, messageId, agent: this };
     const tools = wrapTools(
       toolCtx,
       args.tools ?? threadTools ?? this.options.tools
@@ -844,7 +870,19 @@ export class Agent<AgentTools extends ToolSet> {
      * The {@link ContextOptions} and {@link StorageOptions}
      * options to use for fetching contextual messages and saving input/output messages.
      */
-    options?: Options
+    options?: Options & {
+      /**
+       * Whether to save incremental data (deltas) from streaming responses.
+       * Defaults to false.
+       * If false, it will not save any deltas to the database.
+       * If true, it will save deltas with {@link DEFAULT_STREAMING_OPTIONS}.
+       *
+       * Regardless of this option, when streaming you are able to use this
+       * `streamText` function as you would with the "ai" package's version:
+       * iterating over the text, streaming it over HTTP, etc.
+       */
+      saveStreamDeltas?: boolean | StreamingOptions;
+    }
   ): Promise<
     StreamTextResult<
       TOOLS extends undefined ? AgentTools : TOOLS,
@@ -852,28 +890,50 @@ export class Agent<AgentTools extends ToolSet> {
     > &
       GenerationOutputMetadata
   > {
-    const { args: aiArgs, messageId } = await this._saveMessagesAndFetchContext(
-      ctx,
-      args,
-      { userId, threadId, ...options }
-    );
-    const toolCtx = { ...ctx, userId, threadId, messageId };
+    const context = await this._saveMessagesAndFetchContext(ctx, args, {
+      userId,
+      threadId,
+      ...options,
+    });
+    const { args: aiArgs, messageId, order, stepOrder } = context;
+    const toolCtx = { ...ctx, userId, threadId, messageId, agent: this };
     const tools = wrapTools(
       toolCtx,
       args.tools ?? threadTools ?? this.options.tools
     ) as TOOLS extends undefined ? AgentTools : TOOLS;
-    const saveOutputMessages =
-      options?.storageOptions?.saveOutputMessages ??
-      this.options.storageOptions?.saveOutputMessages;
+    const storageOptions = {
+      ...this.options.storageOptions,
+      ...options?.storageOptions,
+    };
+    const saveOutputMessages = storageOptions.saveOutputMessages;
     const trackUsage = usageHandler ?? this.options.usageHandler;
+    const streamer =
+      threadId && options?.saveStreamDeltas
+        ? new DeltaStreamer(this.component, ctx, options.saveStreamDeltas, {
+            threadId,
+            userId,
+            agentName: this.options.name,
+            model: aiArgs.model.modelId,
+            provider: aiArgs.model.provider,
+            providerOptions: aiArgs.providerOptions,
+            order,
+            stepOrder,
+          })
+        : undefined;
+
     const result = streamText({
       // Can be overridden
       maxSteps: this.options.maxSteps,
       ...aiArgs,
       tools,
-      onChunk: async (chunk) => {
+      experimental_transform: mergeTransforms(
+        options?.saveStreamDeltas,
+        args.experimental_transform
+      ),
+      onChunk: async (event) => {
+        await streamer?.addParts([event.chunk]);
         // console.log("onChunk", chunk);
-        return args.onChunk?.(chunk);
+        return args.onChunk?.(event);
       },
       onError: async (error) => {
         console.error("onError", error);
@@ -889,12 +949,13 @@ export class Agent<AgentTools extends ToolSet> {
         // console.log("onStepFinish", step);
         // TODO: compare delta to the output. internally drop the deltas when committing
         if (threadId && messageId) {
-          await this.saveStep(ctx, {
+          const saved = await this.saveStep(ctx, {
             userId,
             threadId,
             promptMessageId: messageId,
             step,
           });
+          await streamer?.finish(saved);
         }
         if (trackUsage && step.usage) {
           await trackUsage(ctx, {
@@ -943,6 +1004,8 @@ export class Agent<AgentTools extends ToolSet> {
   ): Promise<{
     args: T & { model: LanguageModelV1 };
     messageId: string | undefined;
+    order: number | undefined;
+    stepOrder: number | undefined;
   }> {
     contextOptions ||= this.options.contextOptions;
     storageOptions ||= this.options.storageOptions;
@@ -963,6 +1026,12 @@ export class Agent<AgentTools extends ToolSet> {
       contextOptions,
     });
     let messageId = args.promptMessageId;
+    let order = args.promptMessageId
+      ? contextMessages.at(-1)?.order
+      : undefined;
+    let stepOrder = args.promptMessageId
+      ? contextMessages.at(-1)?.stepOrder
+      : undefined;
     if (
       threadId &&
       messages.length &&
@@ -979,6 +1048,8 @@ export class Agent<AgentTools extends ToolSet> {
         failPendingSteps: true,
       });
       messageId = saved.lastMessageId;
+      order = saved.messages.at(-1)?.order;
+      stepOrder = saved.messages.at(-1)?.stepOrder;
     }
     const { prompt: _, model, ...rest } = args;
     return {
@@ -993,6 +1064,8 @@ export class Agent<AgentTools extends ToolSet> {
         ],
       } as T & { model: LanguageModelV1 },
       messageId,
+      order,
+      stepOrder,
     };
   }
 
@@ -1185,7 +1258,7 @@ export class Agent<AgentTools extends ToolSet> {
     });
   }
 
-  mergedContextOptions(opts: ContextOptions | undefined): ContextOptions {
+  _mergedContextOptions(opts: ContextOptions | undefined): ContextOptions {
     const searchOptions = {
       ...this.options.contextOptions?.searchOptions,
       ...opts?.searchOptions,
@@ -1199,7 +1272,7 @@ export class Agent<AgentTools extends ToolSet> {
     };
   }
 
-  async searchOptionsWithDefaults(
+  async _searchOptionsWithDefaults(
     contextOptions: ContextOptions,
     messages: CoreMessage[]
   ): Promise<SearchOptions> {
@@ -1273,31 +1346,68 @@ export class Agent<AgentTools extends ToolSet> {
    *   {@link ContextOptions}, {@link StorageOptions}, and maxSteps.
    */
   asTextAction(spec?: {
+    /**
+     * The maximum number of steps to take in this action.
+     * Defaults to the {@link Agent.maxSteps} option.
+     */
     maxSteps?: number;
+    /**
+     * The {@link ContextOptions} to use for fetching contextual messages and
+     * saving input/output messages.
+     * Defaults to the {@link Agent.contextOptions} option.
+     */
     contextOptions?: ContextOptions;
+    /**
+     * The {@link StorageOptions} to use for saving input/output messages.
+     * Defaults to the {@link Agent.storageOptions} option.
+     */
     storageOptions?: StorageOptions;
+    /**
+     * Whether to stream the text.
+     * If false, it will generate the text in a single call. (default)
+     * If true or {@link StreamingOptions}, it will stream the text from the LLM
+     * and save the chunks to the database with the options you specify, or the
+     * defaults if you pass true.
+     */
+    stream?: boolean | StreamingOptions;
   }) {
     const maxSteps = spec?.maxSteps ?? this.options.maxSteps;
     return internalActionGeneric({
       args: vTextArgs,
       handler: async (ctx, args) => {
         const { contextOptions, storageOptions, ...rest } = args;
-        const value = await this.generateText(
-          ctx,
-          { userId: args.userId, threadId: args.threadId },
-          { maxSteps, ...rest },
-          {
-            contextOptions:
-              contextOptions ??
-              spec?.contextOptions ??
-              this.options.contextOptions,
-            storageOptions:
-              storageOptions ??
-              spec?.storageOptions ??
-              this.options.storageOptions,
-          }
-        );
-        return value.text;
+        const stream =
+          args.stream === true ? spec?.stream || true : spec?.stream ?? false;
+        const targetArgs = { userId: args.userId, threadId: args.threadId };
+        const llmArgs = { maxSteps, ...rest };
+        const opts = {
+          contextOptions:
+            contextOptions ??
+            spec?.contextOptions ??
+            this.options.contextOptions,
+          storageOptions:
+            storageOptions ??
+            spec?.storageOptions ??
+            this.options.storageOptions,
+          saveStreamDeltas: stream,
+        };
+        if (stream) {
+          const result = await this.streamText(ctx, targetArgs, llmArgs, opts);
+          await result.consumeStream();
+          return {
+            text: await result.text,
+            finishReason: await result.finishReason,
+            messageId: result.messageId,
+          };
+        } else {
+          const { text, messageId, finishReason } = await this.generateText(
+            ctx,
+            targetArgs,
+            llmArgs,
+            opts
+          );
+          return { text, messageId, finishReason };
+        }
       },
     });
   }
@@ -1372,534 +1482,5 @@ export function filterOutOrphanedToolMessages(docs: MessageDoc[]) {
   return result;
 }
 
-export type ToolCtx = RunActionCtx & {
-  userId?: string;
-  threadId?: string;
-  messageId?: string;
-};
-
-// Vendoring in from "ai" package since it wasn't exported
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-type ToolParameters = z.ZodTypeAny | Schema<any>;
-type inferParameters<PARAMETERS extends ToolParameters> =
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  PARAMETERS extends Schema<any>
-    ? PARAMETERS["_type"]
-    : PARAMETERS extends z.ZodTypeAny
-      ? z.infer<PARAMETERS>
-      : never;
-
-/**
- * This is a wrapper around the ai.tool function that adds extra context to the
- * tool call, including the action context, userId, threadId, and messageId.
- * @param tool The tool. See https://sdk.vercel.ai/docs/ai-sdk-core/tools-and-tool-calling
- * but swap parameters for args and handler for execute.
- * @returns A tool to be used with the AI SDK.
- */
-export function createTool<PARAMETERS extends ToolParameters, RESULT>(t: {
-  /**
-  An optional description of what the tool does.
-  Will be used by the language model to decide whether to use the tool.
-  Not used for provider-defined tools.
-     */
-  description?: string;
-  /**
-  The schema of the input that the tool expects. The language model will use this to generate the input.
-  It is also used to validate the output of the language model.
-  Use descriptions to make the input understandable for the language model.
-     */
-  args: PARAMETERS;
-  /**
-  An async function that is called with the arguments from the tool call and produces a result.
-  If not provided, the tool will not be executed automatically.
-
-  @args is the input of the tool call.
-  @options.abortSignal is a signal that can be used to abort the tool call.
-     */
-  handler: (
-    ctx: ToolCtx,
-    args: inferParameters<PARAMETERS>,
-    options: ToolExecutionOptions
-  ) => PromiseLike<RESULT>;
-  ctx?: ToolCtx;
-}): Tool<PARAMETERS, RESULT> & {
-  execute: (
-    args: inferParameters<PARAMETERS>,
-    options: ToolExecutionOptions
-  ) => PromiseLike<RESULT>;
-} {
-  const args = {
-    __acceptsCtx: true,
-    ctx: t.ctx,
-    description: t.description,
-    parameters: t.args,
-    async execute(
-      args: inferParameters<PARAMETERS>,
-      options: ToolExecutionOptions
-    ) {
-      if (!this.ctx) {
-        throw new Error(
-          "To use a Convex tool, you must either provide the ctx" +
-            " at definition time (dynamically in an action), or use the Agent to" +
-            " call it (which injects the ctx, userId and threadId)"
-        );
-      }
-      return t.handler(this.ctx, args, options);
-    },
-  };
-  return tool(args);
-}
-
-function wrapTools(
-  ctx: ToolCtx,
-  ...toolSets: (ToolSet | undefined)[]
-): ToolSet {
-  const output = {} as ToolSet;
-  for (const toolSet of toolSets) {
-    if (!toolSet) {
-      continue;
-    }
-    for (const [name, tool] of Object.entries(toolSet)) {
-      // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      if (!(tool as any).__acceptsCtx) {
-        output[name] = tool;
-      } else {
-        const out = { ...tool, ctx };
-        output[name] = out;
-      }
-    }
-  }
-  return output;
-}
-
-type Options = {
-  /**
-   * The context options to use for passing in message history to the LLM.
-   */
-  contextOptions?: ContextOptions;
-  /**
-   * The storage options to use for saving the input and output messages to the thread.
-   */
-  storageOptions?: StorageOptions;
-};
-
-type TextArgs<
-  AgentTools extends ToolSet,
-  TOOLS extends ToolSet | undefined = undefined,
-  OUTPUT = never,
-  OUTPUT_PARTIAL = never,
-> = Omit<
-  Parameters<
-    typeof generateText<
-      TOOLS extends undefined ? AgentTools : TOOLS,
-      OUTPUT,
-      OUTPUT_PARTIAL
-    >
-  >[0],
-  "toolChoice" | "tools" | "model"
-> & {
-  /**
-   * If provided, this message will be used as the "prompt" for the LLM call,
-   * instead of the prompt or messages.
-   * This is useful if you want to first save a user message, then use it as
-   * the prompt for the LLM call in another call.
-   */
-  promptMessageId?: string;
-  /**
-   * The model to use for the tool calls. This will override the model specified
-   * in the Agent constructor.
-   */
-  model?: LanguageModelV1;
-  /**
-   * The tools to use for the tool calls. This will override tools specified
-   * in the Agent constructor or createThread / continueThread.
-   */
-  tools?: TOOLS;
-  /**
-   * The tool choice to use for the tool calls. This must be one of the tools
-   * specified in the tools array. e.g. {toolName: "getWeather", type: "tool"}
-   */
-  toolChoice?: ToolChoice<TOOLS extends undefined ? AgentTools : TOOLS>;
-};
-
-type StreamingTextArgs<
-  AgentTools extends ToolSet,
-  TOOLS extends ToolSet | undefined = undefined,
-  OUTPUT = never,
-  OUTPUT_PARTIAL = never,
-> = Omit<
-  Parameters<
-    typeof streamText<
-      TOOLS extends undefined ? AgentTools : TOOLS,
-      OUTPUT,
-      OUTPUT_PARTIAL
-    >
-  >[0],
-  "toolChoice" | "tools" | "model"
-> & {
-  /**
-   * The model to use for the tool calls. This will override the model specified
-   * in the Agent constructor.
-   */
-  model?: LanguageModelV1;
-  /**
-   * The tools to use for the tool calls. This will override tools specified
-   * in the Agent constructor or createThread / continueThread.
-   */
-  tools?: TOOLS;
-  /**
-   * The tool choice to use for the tool calls. This must be one of the tools
-   * specified in the tools array. e.g. {toolName: "getWeather", type: "tool"}
-   */
-  toolChoice?: ToolChoice<TOOLS extends undefined ? AgentTools : TOOLS>;
-};
-
-type BaseGenerateObjectOptions = CallSettings & {
-  /**
-   * The model to use for the object generation. This will override the model
-   * specified in the Agent constructor.
-   */
-  model?: LanguageModelV1;
-  /**
-   * The system prompt to use for the object generation. This will override the
-   * system prompt specified in the Agent constructor.
-   */
-  system?: string;
-  /**
-   * The prompt to the LLM to use for the object generation.
-   * Specify this or messages, but not both.
-   */
-  prompt?: string;
-  /**
-   * The messages to use for the object generation.
-   * Note: recent messages are automatically added based on the thread it's
-   * associated with and your contextOptions.
-   */
-  messages?: CoreMessage[];
-  /**
-   * The message to use as the "prompt" for the object generation.
-   * If this is provided, it will be used instead of the prompt or messages.
-   * This is useful if you want to first save a user message, then use it as
-   * the prompt for the object generation in another call.
-   */
-  promptMessageId?: string;
-  experimental_repairText?: RepairTextFunction;
-  experimental_telemetry?: TelemetrySettings;
-  providerOptions?: ProviderOptions;
-  experimental_providerMetadata?: ProviderMetadata;
-};
-
-type GenerateObjectObjectOptions<T extends Record<string, unknown>> =
-  BaseGenerateObjectOptions & {
-    output?: "object";
-    mode?: "auto" | "json" | "tool";
-    schema: z.Schema<T>;
-    schemaName?: string;
-    schemaDescription?: string;
-  };
-
-type GenerateObjectArrayOptions<T> = BaseGenerateObjectOptions & {
-  output: "array";
-  mode?: "auto" | "json" | "tool";
-  schema: z.Schema<T>;
-  schemaName?: string;
-  schemaDescription?: string;
-};
-
-type GenerateObjectWithEnumOptions<T extends string> =
-  BaseGenerateObjectOptions & {
-    output: "enum";
-    enum: Array<T>;
-    mode?: "auto" | "json" | "tool";
-  };
-
-type GenerateObjectNoSchemaOptions = BaseGenerateObjectOptions & {
-  schema?: undefined;
-  mode?: "json";
-};
-
-// TODO: simplify this to just use the generateObject args, with an optional
-// model and tool/toolChoice types
-type GenerateObjectArgs<T> =
-  T extends Record<string, unknown>
-    ? GenerateObjectObjectOptions<T>
-    : T extends Array<unknown>
-      ? GenerateObjectArrayOptions<T>
-      : T extends string
-        ? GenerateObjectWithEnumOptions<T>
-        : GenerateObjectNoSchemaOptions;
-
-type StreamObjectArgs<T> =
-  T extends Record<string, unknown>
-    ? GenerateObjectObjectOptions<T>
-    : T extends Array<unknown>
-      ? GenerateObjectArrayOptions<T>
-      : GenerateObjectNoSchemaOptions;
-
-type OurObjectArgs<T> = GenerateObjectArgs<T> &
-  Pick<
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    Parameters<typeof generateObject<any>>[0],
-    "experimental_repairText" | "abortSignal"
-  >;
-
-type OurStreamObjectArgs<T> = StreamObjectArgs<T> &
-  Pick<
-    Parameters<typeof streamObject<T>>[0],
-    "onError" | "onFinish" | "abortSignal"
-  >;
-
-type ThreadOutputMetadata = GenerationOutputMetadata & {
-  messageId: string;
-};
-
-/**
- * The interface for a thread returned from {@link createThread} or {@link continueThread}.
- * This is contextual to a thread and/or user.
- */
-interface Thread<DefaultTools extends ToolSet> {
-  /**
-   * The target threadId, from the startThread or continueThread initializers.
-   */
-  threadId: string;
-  /**
-   * This behaves like {@link generateText} from the "ai" package except that
-   * it add context based on the userId and threadId and saves the input and
-   * resulting messages to the thread, if specified.
-   * Use {@link continueThread} to get a version of this function already scoped
-   * to a thread (and optionally userId).
-   * @param args The arguments to the generateText function, along with extra controls
-   * for the {@link ContextOptions} and {@link StorageOptions}.
-   * @returns The result of the generateText function.
-   */
-  generateText<
-    TOOLS extends ToolSet | undefined = undefined,
-    OUTPUT = never,
-    OUTPUT_PARTIAL = never,
-  >(
-    args: TextArgs<
-      TOOLS extends undefined ? DefaultTools : TOOLS,
-      TOOLS,
-      OUTPUT,
-      OUTPUT_PARTIAL
-    >,
-    options?: Options
-  ): Promise<
-    GenerateTextResult<TOOLS extends undefined ? DefaultTools : TOOLS, OUTPUT> &
-      ThreadOutputMetadata
-  >;
-
-  /**
-   * This behaves like {@link streamText} from the "ai" package except that
-   * it add context based on the userId and threadId and saves the input and
-   * resulting messages to the thread, if specified.
-   * Use {@link continueThread} to get a version of this function already scoped
-   * to a thread (and optionally userId).
-   * @param args The arguments to the streamText function, along with extra controls
-   * for the {@link ContextOptions} and {@link StorageOptions}.
-   * @returns The result of the streamText function.
-   */
-  streamText<
-    TOOLS extends ToolSet | undefined = undefined,
-    OUTPUT = never,
-    PARTIAL_OUTPUT = never,
-  >(
-    args: StreamingTextArgs<
-      TOOLS extends undefined ? DefaultTools : TOOLS,
-      TOOLS,
-      OUTPUT,
-      PARTIAL_OUTPUT
-    >,
-    options?: Options
-  ): Promise<
-    StreamTextResult<
-      TOOLS extends undefined ? DefaultTools : TOOLS,
-      PARTIAL_OUTPUT
-    > &
-      ThreadOutputMetadata
-  >;
-  /**
-   * This behaves like {@link generateObject} from the "ai" package except that
-   * it add context based on the userId and threadId and saves the input and
-   * resulting messages to the thread, if specified. This overload is for objects, arrays, and enums.
-   * Use {@link continueThread} to get a version of this function already scoped
-   * to a thread (and optionally userId).
-   * @param args The arguments to the generateObject function, along with extra controls
-   * for the {@link ContextOptions} and {@link StorageOptions}.
-   * @returns The result of the generateObject function.
-   */
-  generateObject<T>(
-    args: OurObjectArgs<T>,
-    options?: Options
-  ): Promise<GenerateObjectResult<T> & ThreadOutputMetadata>;
-  /**
-   * This behaves like {@link generateObject} from the "ai" package except that
-   * it add context based on the userId and threadId and saves the input and
-   * resulting messages to the thread, if specified. This overload is for when there's no schema.
-   * Use {@link continueThread} to get a version of this function already scoped
-   * to a thread (and optionally userId).
-   * @param args The arguments to the generateObject function, along with extra controls
-   * for the {@link ContextOptions} and {@link StorageOptions}.
-   * @returns The result of the generateObject function.
-   */
-  generateObject(
-    args: GenerateObjectNoSchemaOptions,
-    options?: Options
-  ): Promise<GenerateObjectResult<JSONValue> & ThreadOutputMetadata>;
-  /**
-   * This behaves like {@link streamObject} from the "ai" package except that
-   * it add context based on the userId and threadId and saves the input and
-   * resulting messages to the thread, if specified.
-   * Use {@link continueThread} to get a version of this function already scoped
-   * to a thread (and optionally userId).
-   * @param args The arguments to the streamObject function, along with extra controls
-   * for the {@link ContextOptions} and {@link StorageOptions}.
-   * @returns The result of the streamObject function.
-   */
-  streamObject<T>(
-    args: OurStreamObjectArgs<T>,
-    options?: Options
-  ): Promise<
-    StreamObjectResult<DeepPartial<T>, T, never> & ThreadOutputMetadata
-  >;
-}
-
 type MessageWithMetadata = OpaqueIds<InnerMessageWithMetadata>;
-
-export function toUIMessages(messages: MessageDoc[]): UIMessage[] {
-  const uiMessages: UIMessage[] = [];
-  let assistantMessage: UIMessage | undefined;
-  for (const message of messages) {
-    const coreMessage = message.message && deserializeMessage(message.message);
-    const text = message.text ?? "";
-    const content = coreMessage?.content;
-    const nonStringContent =
-      content && typeof content !== "string" ? content : [];
-    if (!coreMessage) continue;
-    if (coreMessage.role === "system") {
-      uiMessages.push({
-        id: message.id ?? message._id,
-        createdAt: new Date(message._creationTime),
-        role: "system",
-        content: text,
-        parts: [{ type: "text", text }],
-      });
-    } else if (coreMessage.role === "user") {
-      const parts: UIMessage["parts"] = [];
-      if (text) {
-        parts.push({ type: "text", text });
-      }
-      if (message.files) {
-        parts.push(...message.files.map(toUIFilePart));
-      }
-      uiMessages.push({
-        id: message.id ?? message._id,
-        createdAt: new Date(message._creationTime),
-        role: "user",
-        content: message.text ?? "",
-        parts,
-      });
-    } else {
-      if (coreMessage.role === "tool" && !assistantMessage) {
-        console.warn(
-          "Tool message without preceding assistant message.. skipping",
-          message
-        );
-        continue;
-      }
-      if (!assistantMessage) {
-        assistantMessage = {
-          id: message.id ?? message._id,
-          createdAt: new Date(message._creationTime),
-          role: "assistant",
-          content: message.text ?? "",
-          parts: [],
-        };
-        uiMessages.push(assistantMessage);
-      }
-      // update it to the last message's id
-      assistantMessage.id = message.id ?? message._id;
-      if (message.text) {
-        assistantMessage.parts.push({
-          type: "text",
-          text: message.text,
-        });
-        assistantMessage.content += message.text;
-      }
-      if (message.reasoning) {
-        assistantMessage.parts.push({
-          type: "reasoning",
-          reasoning: message.reasoning,
-          details: message.reasoningDetails ?? [],
-        });
-      }
-      for (const source of message.sources ?? []) {
-        assistantMessage.parts.push({
-          type: "source",
-          source,
-        });
-      }
-      for (const file of message.files ?? []) {
-        assistantMessage.parts.push(toUIFilePart(file));
-      }
-      for (const contentPart of nonStringContent) {
-        switch (contentPart.type) {
-          case "tool-call":
-            assistantMessage.parts.push({
-              type: "step-start",
-            });
-            assistantMessage.parts.push({
-              type: "tool-invocation",
-              toolInvocation: {
-                state: "call",
-                step: assistantMessage.parts.filter(
-                  (part) => part.type === "tool-invocation"
-                ).length,
-                toolCallId: contentPart.toolCallId,
-                toolName: contentPart.toolName,
-                args: contentPart.args,
-              },
-            });
-            break;
-          case "tool-result": {
-            const call = assistantMessage.parts.find(
-              (part) =>
-                part.type === "tool-invocation" &&
-                part.toolInvocation.toolCallId === contentPart.toolCallId
-            ) as ToolInvocationUIPart | undefined;
-            const toolInvocation: ToolInvocationUIPart["toolInvocation"] = {
-              state: "result",
-              toolCallId: contentPart.toolCallId,
-              toolName: contentPart.toolName,
-              args: call?.toolInvocation.args,
-              result: contentPart.result,
-              step:
-                call?.toolInvocation.step ??
-                assistantMessage.parts.filter(
-                  (part) => part.type === "tool-invocation"
-                ).length,
-            };
-            if (call) {
-              (call as ToolInvocationUIPart).toolInvocation = toolInvocation;
-            } else {
-              console.warn(
-                "Tool result without preceding tool call.. adding anyways",
-                contentPart
-              );
-              assistantMessage.parts.push({
-                type: "tool-invocation",
-                toolInvocation,
-              });
-            }
-            break;
-          }
-        }
-      }
-      if (!message.tool) {
-        // Reset it so the next set of tool calls will create a new assistant message
-        assistantMessage = undefined;
-      }
-    }
-  }
-  return uiMessages;
-}
+type CoreMessageMaybeWithId = CoreMessage & { id?: string | undefined };