@convex-dev/agent 0.2.6-alpha.1 → 0.2.7

package/src/client/index.ts

@@ -2,11 +2,9 @@ import type {
   FlexibleSchema,
   IdGenerator,
   InferSchema,
-  ProviderOptions,
 } from "@ai-sdk/provider-utils";
 import type {
   CallSettings,
-  EmbeddingModel,
   GenerateObjectResult,
   GenerateTextResult,
   LanguageModel,
@@ -18,7 +16,6 @@ import type {
   ToolSet,
 } from "ai";
 import {
-  embedMany,
   generateObject,
   generateText,
   stepCountIs,
@@ -35,31 +32,28 @@ import {
   type WithoutSystemFields,
 } from "convex/server";
 import { convexToJson, v, type Value } from "convex/values";
-import type { MessageDoc, ThreadDoc } from "../component/schema.js";
 import type { threadFieldsSupportingPatch } from "../component/threads.js";
-import {
-  validateVectorDimension,
-  type VectorDimension,
-} from "../component/vector/tables.js";
+import { type VectorDimension } from "../component/vector/tables.js";
 import {
   deserializeMessage,
   serializeMessage,
   serializeNewMessagesInStep,
   serializeObjectResult,
 } from "../mapping.js";
-import { extractText, isTool } from "../shared.js";
+import { getModelName, getProviderName } from "../shared.js";
 import {
   vMessageEmbeddings,
   vMessageWithMetadata,
   vSafeObjectArgs,
   vTextArgs,
   type Message,
+  type MessageDoc,
   type MessageStatus,
   type MessageWithMetadata,
   type ProviderMetadata,
   type StreamArgs,
+  type ThreadDoc,
 } from "../validators.js";
-import { wrapTools, type ToolCtx } from "./createTool.js";
 import {
   listMessages,
   saveMessages,
@@ -67,10 +61,12 @@ import {
   type SaveMessagesArgs,
 } from "./messages.js";
 import {
+  embedMany,
+  embedMessages,
   fetchContextMessages,
-  getModelName,
-  getProviderName,
+  generateAndSaveEmbeddings,
 } from "./search.js";
+import { start } from "./start.js";
 import {
   DeltaStreamer,
   syncStreams,
@@ -104,12 +100,10 @@ import type {
   Thread,
   UsageHandler,
   UserActionCtx,
+  Config,
 } from "./types.js";
-import { inlineMessagesFiles } from "./files.js";
-import type { DataModel } from "../component/_generated/dataModel.js";
 
 export { stepCountIs } from "ai";
-export { vMessageDoc, vThreadDoc } from "../component/schema.js";
 export {
   deserializeMessage,
   serializeDataOrUrl,
@@ -123,18 +117,22 @@ export {
   vAssistantMessage,
   vContextOptions,
   vMessage,
+  vMessageDoc,
   vPaginationResult,
   vProviderMetadata,
   vStorageOptions,
   vStreamArgs,
   vSystemMessage,
+  vThreadDoc,
   vToolMessage,
   vUsage,
   vUserMessage,
   vSource,
   vContent,
-  type SourcePart,
   type Message,
+  type MessageDoc,
+  type SourcePart,
+  type ThreadDoc,
   type Usage,
 } from "../validators.js";
 export type { ToolCtx } from "./createTool.js";
@@ -154,6 +152,10 @@ export {
 export {
   fetchContextMessages,
   filterOutOrphanedToolMessages,
+  fetchContextWithPrompt,
+  generateAndSaveEmbeddings,
+  embedMessages,
+  embedMany,
 } from "./search.js";
 export {
   abortStream,
@@ -171,94 +173,18 @@ export { extractText, isTool, sorted } from "../shared.js";
 export { createTool } from "./createTool.js";
 export type {
   AgentComponent,
+  Config,
   ContextOptions,
-  MessageDoc,
   ProviderMetadata,
   RawRequestResponseHandler,
   StorageOptions,
   StreamArgs,
   SyncStreamsReturnValue,
   Thread,
-  ThreadDoc,
   UsageHandler,
 };
 export { mockModel } from "./mockModel.js";
 
-// 10k characters should be more than enough for most cases, and stays under
-// the 8k token limit for some models.
-const MAX_EMBEDDING_TEXT_LENGTH = 10_000;
-
-export type Config = {
-  /**
-   * The LLM model to use for generating / streaming text and objects.
-   * e.g.
-   * import { openai } from "@ai-sdk/openai"
-   * const myAgent = new Agent(components.agent, {
-   *   languageModel: openai.chat("gpt-4o-mini"),
-   */
-  languageModel?: LanguageModel;
-  /**
-   * The model to use for text embeddings. Optional.
-   * If specified, it will use this for generating vector embeddings
-   * of chats, and can opt-in to doing vector search for automatic context
-   * on generateText, etc.
-   * e.g.
-   * import { openai } from "@ai-sdk/openai"
-   * const myAgent = new Agent(components.agent, {
-   *   ...
-   *   textEmbeddingModel: openai.embedding("text-embedding-3-small")
-   */
-  textEmbeddingModel?: EmbeddingModel<string>;
-  /**
-   * Options to determine what messages are included as context in message
-   * generation. To disable any messages automatically being added, pass:
-   * { recentMessages: 0 }
-   */
-  contextOptions?: ContextOptions;
-  /**
-   * Determines whether messages are automatically stored when passed as
-   * arguments or generated.
-   */
-  storageOptions?: StorageOptions;
-  /**
-   * The usage handler to use for this agent.
-   */
-  usageHandler?: UsageHandler;
-  /**
-   * Called for each LLM request/response, so you can do things like
-   * log the raw request body or response headers to a table, or logs.
-   */
-  rawRequestResponseHandler?: RawRequestResponseHandler;
-  /**
-   * Default provider options to pass for the LLM calls.
-   * This can be overridden at each generate/stream callsite on a per-field
-   * basis. To clear a default setting, you'll need to pass `undefined`.
-   */
-  providerOptions?: ProviderOptions;
-  /**
-   * The default settings to use for the LLM calls.
-   * This can be overridden at each generate/stream callsite on a per-field
-   * basis. To clear a default setting, you'll need to pass `undefined`.
-   */
-  callSettings?: CallSettings;
-  /**
-   * The maximum number of steps to allow for a single generation.
-   *
-   * For example, if an agent wants to call a tool, that call and tool response
-   * will be one step. Generating a response based on the tool call & response
-   * will be a second step.
-   * If it runs out of steps, it will return the last step result, which may
-   * not be an assistant message.
-
-   * This becomes the default value when `stopWhen` is not specified in the
-   * Agent or generation callsite.
-   * AI SDK v5 removed the `maxSteps` argument, but this is kept here for
-   * convenience and backwards compatibility.
-   * Defaults to 1.
-   */
-  maxSteps?: number;
-};
-
 export class Agent<
   /**
    * You can require that all `ctx` args to generateText & streamText
@@ -313,7 +239,7 @@ export class Agent<
       tools?: AgentTools;
       /**
        * When generating or streaming text with tools available, this
-       * determines when to stop. Defaults to stepCountIs(1).
+       * determines when to stop. Defaults to the AI SDK default.
        */
       stopWhen?: StopCondition<AgentTools> | Array<StopCondition<AgentTools>>;
       /**
@@ -501,6 +427,7 @@ export class Agent<
     args: T & {
       system?: string;
       model: LanguageModel;
+      prompt?: never;
       messages: ModelMessage[];
       tools?: TOOLS extends undefined ? AgentTools : TOOLS;
     } & CallSettings;
@@ -518,182 +445,25 @@ export class Agent<
     fail: (reason: string) => Promise<void>;
     getSavedMessages: () => MessageDoc[];
   }> {
-    const { threadId, ...opts } = { ...this.options, ...options };
-    const context = await this._saveMessagesAndFetchContext(ctx, {
-      userId: options?.userId,
-      threadId: options?.threadId,
-      messages: args.messages,
-      prompt: args.prompt,
-      promptMessageId: args.promptMessageId,
-      ...opts,
-    });
-    let pendingMessageId = context.pendingMessageId;
-    const { messages, promptMessageId, order, stepOrder, userId } = context;
-    const savedMessages = context.savedMessages ?? [];
-    const toolCtx = {
-      ...(ctx as UserActionCtx & CustomCtx),
-      userId,
-      threadId,
-      promptMessageId,
-      agent: this,
-    } satisfies ToolCtx;
     type Tools = TOOLS extends undefined ? AgentTools : TOOLS;
-    const tools = wrapTools(toolCtx, args.tools ?? this.options.tools) as Tools;
-    const saveOutput = opts.storageOptions?.saveMessages !== "none";
-    const fail = async (reason: string) => {
-      if (threadId && promptMessageId) {
-        console.error(
-          `Message failed in thread ${threadId} with promptMessageId ${promptMessageId}: ${reason}`,
-        );
-      }
-      if (pendingMessageId) {
-        await ctx.runMutation(this.component.messages.finalizeMessage, {
-          messageId: pendingMessageId,
-          result: { status: "failed", error: reason },
-        });
-      }
-    };
-    if (args.abortSignal) {
-      const abortSignal = args.abortSignal;
-      abortSignal.addEventListener(
-        "abort",
-        async () => {
-          await fail(abortSignal.reason ?? "Aborted");
-        },
-        { once: true },
-      );
-    }
-    const aiArgs = {
-      ...this.options.callSettings,
-      providerOptions: this.options.providerOptions,
-      ...omit(args, ["messages", "prompt", "promptMessageId"]),
-      model: args.model ?? this.options.languageModel,
-      system: args.system ?? this.options.instructions,
-      messages,
-      stopWhen:
-        args.stopWhen ??
-        this.options.stopWhen ??
-        stepCountIs(this.options.maxSteps ?? 1),
-      tools,
-    } as T & {
-      model: LanguageModel;
-      messages: ModelMessage[];
-      tools?: TOOLS extends undefined ? AgentTools : TOOLS;
-    } & CallSettings;
-    if (pendingMessageId) {
-      if (!aiArgs._internal?.generateId) {
-        aiArgs._internal = {
-          ...aiArgs._internal,
-          generateId: () => pendingMessageId ?? crypto.randomUUID(),
-        };
-      }
-    }
-    let activeModel = aiArgs.model;
-    return {
-      args: aiArgs,
-      order: order ?? 0,
-      stepOrder: stepOrder ?? 0,
-      userId,
-      promptMessageId,
-      getSavedMessages: () => savedMessages,
-      updateModel: (model: LanguageModel | undefined) => {
-        if (model) {
-          activeModel = model;
-        }
+    return start<T, Tools, CustomCtx>(
+      ctx,
+      this.component,
+      {
+        ...args,
+        tools: (args.tools ?? this.options.tools) as Tools,
+        system: args.system ?? this.options.instructions,
+        stopWhen: (args.stopWhen ?? this.options.stopWhen) as
+          | StopCondition<Tools>
+          | Array<StopCondition<Tools>>,
       },
-      fail,
-      save: async <TOOLS extends ToolSet>(
-        toSave:
-          | { step: StepResult<TOOLS> }
-          | { object: GenerateObjectResult<unknown> },
-        createPendingMessage?: boolean,
-      ) => {
-        if (threadId && promptMessageId && saveOutput) {
-          const metadata = {
-            // TODO: get up to date one when user selects mid-generation
-            model: getModelName(activeModel),
-            provider: getProviderName(activeModel),
-          };
-          const serialized =
-            "object" in toSave
-              ? await serializeObjectResult(
-                  ctx,
-                  this.component,
-                  toSave.object,
-                  metadata,
-                )
-              : await serializeNewMessagesInStep(
-                  ctx,
-                  this.component,
-                  toSave.step,
-                  metadata,
-                );
-          const embeddings = await this.generateEmbeddings(
-            ctx,
-            { userId, threadId },
-            serialized.messages.map((m) => m.message),
-          );
-          if (createPendingMessage) {
-            serialized.messages.push({
-              message: { role: "assistant", content: [] },
-              status: "pending",
-            });
-            embeddings?.vectors.push(null);
-          }
-          const saved = await ctx.runMutation(
-            this.component.messages.addMessages,
-            {
-              userId,
-              threadId,
-              agentName: this.options.name,
-              promptMessageId,
-              pendingMessageId,
-              messages: serialized.messages,
-              embeddings,
-              failPendingSteps: false,
-            },
-          );
-          const lastMessage = saved.messages.at(-1)!;
-          if (createPendingMessage) {
-            if (lastMessage.status === "failed") {
-              pendingMessageId = undefined;
-              savedMessages.push(...saved.messages);
-              await fail(
-                lastMessage.error ??
-                  "Aborting - the pending message was marked as failed",
-              );
-            } else {
-              pendingMessageId = lastMessage._id;
-              savedMessages.push(...saved.messages.slice(0, -1));
-            }
-          } else {
-            pendingMessageId = undefined;
-            savedMessages.push(...saved.messages);
-          }
-        }
-        const output = "object" in toSave ? toSave.object : toSave.step;
-        if (this.options.rawRequestResponseHandler) {
-          await this.options.rawRequestResponseHandler(ctx, {
-            userId,
-            threadId,
-            agentName: this.options.name,
-            request: output.request,
-            response: output.response,
-          });
-        }
-        if (opts.usageHandler && output.usage) {
-          await opts.usageHandler(ctx, {
-            userId,
-            threadId,
-            agentName: this.options.name,
-            model: getModelName(activeModel),
-            provider: getProviderName(activeModel),
-            usage: output.usage,
-            providerMetadata: output.providerMetadata,
-          });
-        }
+      {
+        ...this.options,
+        ...options,
+        agentName: this.options.name,
+        agentForToolCtx: this,
       },
-    };
+    );
   }
 
   /**
@@ -715,7 +485,15 @@ export class Agent<
   >(
     ctx: ActionCtx & CustomCtx,
     threadOpts: { userId?: string | null; threadId?: string },
-    generateTextArgs: TextArgs<AgentTools, TOOLS, OUTPUT, OUTPUT_PARTIAL>,
+    generateTextArgs: TextArgs<AgentTools, TOOLS, OUTPUT, OUTPUT_PARTIAL> & {
+      /**
+       * 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;
+    },
     options?: Options,
   ): Promise<
     GenerateTextResult<TOOLS extends undefined ? AgentTools : TOOLS, OUTPUT> &
@@ -778,7 +556,15 @@ export class Agent<
       TOOLS,
       OUTPUT,
       PARTIAL_OUTPUT
-    >,
+    > & {
+      /**
+       * 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 {@link ContextOptions} and {@link StorageOptions}
      * options to use for fetching contextual messages and saving input/output messages.
@@ -853,9 +639,6 @@ export class Agent<
         await streamer?.fail(errorToString(error.error));
         return streamTextArgs.onError?.(error);
       },
-      // onFinish: async (event) => {
-      //   return streamTextArgs.onFinish?.(event);
-      // },
       prepareStep: async (options) => {
         const result = await streamTextArgs.prepareStep?.(options);
         if (result) {
@@ -878,12 +661,6 @@ export class Agent<
       TOOLS extends undefined ? AgentTools : TOOLS,
       PARTIAL_OUTPUT
     >;
-    const metadata: GenerationOutputMetadata = {
-      promptMessageId,
-      order,
-      savedMessages: call.getSavedMessages(),
-      messageId: promptMessageId,
-    };
     if (
       (typeof options?.saveStreamDeltas === "object" &&
         !options.saveStreamDeltas.returnImmediately) ||
@@ -891,6 +668,12 @@ export class Agent<
     ) {
       await result.consumeStream();
     }
+    const metadata: GenerationOutputMetadata = {
+      promptMessageId,
+      order,
+      savedMessages: call.getSavedMessages(),
+      messageId: promptMessageId,
+    };
     return Object.assign(result, metadata);
   }
 
@@ -915,7 +698,15 @@ export class Agent<
     /**
      * The arguments to the generateObject function, similar to the ai.generateObject function.
      */
-    generateObjectArgs: GenerateObjectArgs<SCHEMA, OUTPUT, RESULT>,
+    generateObjectArgs: GenerateObjectArgs<SCHEMA, OUTPUT, RESULT> & {
+      /**
+       * 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 {@link ContextOptions} and {@link StorageOptions}
      * options to use for fetching contextual messages and saving input/output messages.
@@ -973,15 +764,6 @@ export class Agent<
        * the prompt for the LLM call in another call.
        */
       promptMessageId?: string;
-      /**
-       * The model to use for the LLM calls. This will override the model specified
-       * in the Agent constructor.
-       */
-      model?: LanguageModel;
-      /**
-       * The tools to use for the tool calls. This will override tools specified
-       * in the Agent constructor or createThread / continueThread.
-       */
     },
     /**
      * The {@link ContextOptions} and {@link StorageOptions}
@@ -1177,11 +959,22 @@ export class Agent<
     args: {
       userId: string | undefined;
       threadId: string | undefined;
-      messages: (ModelMessage | Message)[];
       /**
-       * If provided, it will search for messages up to and including this message.
-       * Note: if this is far in the past, text and vector search results may be more
-       * limited, as it's post-filtering the results.
+       * If targetMessageId is not provided, this text will be used
+       * for text and vector search
+       */
+      searchText?: string;
+      /**
+       * If provided, it will use this message for text/vector search (if enabled)
+       * and will only fetch messages up to (and including) this message's "order"
+       */
+      targetMessageId?: string;
+      /**
+       * @deprecated use searchText and targetMessageId instead
+       */
+      messages?: (ModelMessage | Message)[];
+      /**
+       * @deprecated use targetMessageId instead
        */
       upToAndIncludingMessageId?: string;
       contextOptions: ContextOptions | undefined;
@@ -1203,7 +996,9 @@ export class Agent<
         );
         return {
           embedding: (
-            await this.doEmbed(ctx, {
+            await embedMany(ctx, {
+              ...this.options,
+              agentName: this.options.name,
               userId: args.userId,
               threadId: args.threadId,
               values: [text],
@@ -1258,51 +1053,21 @@ export class Agent<
    */
   async generateEmbeddings(
     ctx: RunActionCtx,
-    {
-      userId,
-      threadId,
-    }: { userId: string | undefined; threadId: string | undefined },
+    args: { userId: string | undefined; threadId: string | undefined },
     messages: (ModelMessage | Message)[],
-  ) {
-    if (!this.options.textEmbeddingModel) {
-      return undefined;
-    }
-    let embeddings:
-      | {
-          vectors: (number[] | null)[];
-          dimension: VectorDimension;
-          model: string;
-        }
-      | undefined;
-    const messageTexts = messages.map((m) => !isTool(m) && extractText(m));
-    // Find the indexes of the messages that have text.
-    const textIndexes = messageTexts
-      .map((t, i) => (t ? i : undefined))
-      .filter((i) => i !== undefined);
-    if (textIndexes.length === 0) {
-      return undefined;
-    }
-    const values = messageTexts
-      .map((t) => t && t.trim().slice(0, MAX_EMBEDDING_TEXT_LENGTH))
-      .filter((t): t is string => !!t);
-    // Then embed those messages.
-    const textEmbeddings = await this.doEmbed(ctx, {
-      userId,
-      threadId,
-      values,
-    });
-    // Then assemble the embeddings into a single array with nulls for the messages without text.
-    const embeddingsOrNull = Array(messages.length).fill(null);
-    textIndexes.forEach((i, j) => {
-      embeddingsOrNull[i] = textEmbeddings.embeddings[j];
-    });
-    if (textEmbeddings.embeddings.length > 0) {
-      const dimension = textEmbeddings.embeddings[0].length;
-      validateVectorDimension(dimension);
-      const model = getModelName(this.options.textEmbeddingModel);
-      embeddings = { vectors: embeddingsOrNull, dimension, model };
-    }
-    return embeddings;
+  ): Promise<
+    | {
+        vectors: (number[] | null)[];
+        dimension: VectorDimension;
+        model: string;
+      }
+    | undefined
+  > {
+    return embedMessages(
+      ctx,
+      { ...args, ...this.options, agentName: this.options.name },
+      messages,
+    );
   }
 
   /**
@@ -1329,10 +1094,6 @@ export class Agent<
             .join(", "),
       );
     }
-    await this._generateAndSaveEmbeddings(ctx, messages);
-  }
-
-  async _generateAndSaveEmbeddings(ctx: RunActionCtx, messages: MessageDoc[]) {
     if (messages.some((m) => !m.message)) {
       throw new Error(
         "Some messages don't have a message: " +
@@ -1342,45 +1103,24 @@ export class Agent<
             .join(", "),
       );
     }
-    const messagesMissingEmbeddings = messages.filter((m) => !m.embeddingId);
-    if (messagesMissingEmbeddings.length === 0) {
-      return;
+    const { textEmbeddingModel } = this.options;
+    if (!textEmbeddingModel) {
+      throw new Error(
+        "No embeddings were generated for the messages. You must pass a textEmbeddingModel to the agent constructor.",
+      );
     }
-    const embeddings = await this.generateEmbeddings(
+    await generateAndSaveEmbeddings(
       ctx,
+      this.component,
       {
-        userId: messagesMissingEmbeddings[0]!.userId,
-        threadId: messagesMissingEmbeddings[0]!.threadId,
+        ...this.options,
+        agentName: this.options.name,
+        threadId: messages[0].threadId,
+        userId: messages[0].userId,
+        textEmbeddingModel,
       },
-      messagesMissingEmbeddings.map((m) => deserializeMessage(m!.message!)),
+      messages,
     );
-    if (!embeddings) {
-      if (!this.options.textEmbeddingModel) {
-        throw new Error(
-          "No embeddings were generated for the messages. You must pass a textEmbeddingModel to the agent constructor.",
-        );
-      }
-      throw new Error(
-        "No embeddings were generated for these messages: " +
-          messagesMissingEmbeddings.map((m) => m!._id).join(", "),
-      );
-    }
-    await ctx.runMutation(this.component.vector.index.insertBatch, {
-      vectorDimension: embeddings.dimension,
-      vectors: messagesMissingEmbeddings
-        .map((m, i) => ({
-          messageId: m!._id,
-          model: embeddings.model,
-          table: "messages",
-          userId: m.userId,
-          threadId: m.threadId,
-          vector: embeddings.vectors[i],
-        }))
-        .filter(
-          (v): v is Extract<typeof v, { vector: number[] }> =>
-            v.vector !== null,
-        ),
-    });
   }
 
   /**
@@ -1677,189 +1417,6 @@ export class Agent<
     });
   }
 
-  async _saveMessagesAndFetchContext(
-    ctx: RunActionCtx,
-    {
-      userId: argsUserId,
-      threadId,
-      contextOptions,
-      storageOptions,
-      ...args
-    }: {
-      prompt: string | (ModelMessage | Message)[] | undefined;
-      messages: (ModelMessage | Message)[] | undefined;
-      promptMessageId: string | undefined;
-      userId: string | null | undefined;
-      threadId: string | undefined;
-    } & Options,
-  ): Promise<{
-    messages: ModelMessage[];
-    userId: string | undefined;
-    promptMessageId: string | undefined;
-    pendingMessageId: string | undefined;
-    order: number | undefined;
-    stepOrder: number | undefined;
-    savedMessages: MessageDoc[] | undefined;
-  }> {
-    // If only a promptMessageId is provided, this will be empty.
-    const messages: (ModelMessage | Message)[] = args.messages ?? [];
-    const promptArray: ModelMessage[] = !args.prompt
-      ? []
-      : Array.isArray(args.prompt)
-        ? args.prompt.map((p) => deserializeMessage(p))
-        : [{ role: "user", content: args.prompt }];
-    const userId =
-      argsUserId ??
-      (threadId &&
-        (await ctx.runQuery(this.component.threads.getThread, { threadId }))
-          ?.userId) ??
-      undefined;
-    // If only a messageId is provided, this will add that message to the end.
-    const contextMessages: MessageDoc[] = await this.fetchContextMessages(ctx, {
-      userId,
-      threadId,
-      upToAndIncludingMessageId: args.promptMessageId,
-      messages,
-      contextOptions,
-    });
-    // If it was a promptMessageId, pop it off context messages
-    // and add to the end of messages.
-    const promptMessageIndex = args.promptMessageId
-      ? contextMessages.findIndex((m) => m._id === args.promptMessageId)
-      : -1;
-    const promptMessage: MessageDoc | undefined =
-      promptMessageIndex !== -1
-        ? contextMessages.splice(promptMessageIndex, 1)[0]
-        : undefined;
-
-    let promptMessageId = promptMessage?._id;
-    let order = promptMessage?.order;
-    let stepOrder = promptMessage?.stepOrder;
-    let savedMessages = undefined;
-    let pendingMessageId = undefined;
-    if (threadId && storageOptions?.saveMessages !== "none") {
-      let saved: { messages: MessageDoc[] };
-      if (
-        messages.length + promptArray.length &&
-        // If it was a promptMessageId, we don't want to save it again.
-        (!args.promptMessageId || storageOptions?.saveMessages === "all")
-      ) {
-        const saveAll = storageOptions?.saveMessages === "all";
-        const coreMessages: (ModelMessage | Message)[] = [
-          ...messages,
-          ...promptArray,
-        ];
-        const toSave = saveAll ? coreMessages : coreMessages.slice(-1);
-        const metadata = Array.from({ length: toSave.length }, () => ({}));
-        saved = await this.saveMessages(ctx, {
-          threadId,
-          userId,
-          messages: [...toSave, { role: "assistant", content: [] }],
-          metadata: [...metadata, { status: "pending" }],
-          // TODO: sanity check
-          failPendingSteps: !!args.promptMessageId,
-        });
-        promptMessageId = saved.messages.at(-2)!._id;
-      } else {
-        saved = await this.saveMessages(ctx, {
-          threadId,
-          userId,
-          messages: [{ role: "assistant", content: [] }],
-          metadata: [{ status: "pending" }],
-          failPendingSteps: !!args.promptMessageId,
-        });
-      }
-      pendingMessageId = saved.messages.at(-1)!._id;
-      order = saved.messages.at(-1)!.order;
-      stepOrder = saved.messages.at(-1)!.stepOrder;
-      // Don't return the pending message
-      savedMessages = saved.messages.slice(0, -1);
-    }
-
-    if (promptMessage?.message) {
-      if (!args.prompt) {
-        // If they override the prompt, we skip the existing prompt message.
-        messages.push(promptMessage.message);
-      }
-      // Lazily generate embeddings for the prompt message, if it doesn't have
-      // embeddings yet. This can happen if the message was saved in a mutation
-      // where the LLM is not available.
-      if (!promptMessage.embeddingId && this.options.textEmbeddingModel) {
-        await this._generateAndSaveEmbeddings(ctx, [promptMessage]);
-      }
-    }
-
-    const prePrompt = contextMessages.map((m) => m.message).filter((m) => !!m);
-    let existingResponses: (ModelMessage | Message)[] = [];
-    if (promptMessageIndex !== -1) {
-      // pull any messages that already responded to the prompt off
-      // and add them after the prompt
-      existingResponses = prePrompt.splice(promptMessageIndex);
-    }
-
-    let processedMessages: ModelMessage[] = [
-      ...prePrompt,
-      ...messages,
-      ...promptArray,
-      ...existingResponses,
-    ].map((m) => deserializeMessage(m));
-
-    // Process messages to inline localhost files (if not, file urls pointing to localhost will be sent to LLM providers)
-    if (process.env.CONVEX_CLOUD_URL?.startsWith("http://127.0.0.1")) {
-      processedMessages = await inlineMessagesFiles(processedMessages);
-    }
-
-    return {
-      messages: processedMessages,
-      userId,
-      promptMessageId,
-      pendingMessageId,
-      savedMessages,
-      order,
-      stepOrder,
-    };
-  }
-
-  async doEmbed(
-    ctx: RunActionCtx,
-    options: {
-      userId: string | undefined;
-      threadId: string | undefined;
-      values: string[];
-      abortSignal?: AbortSignal;
-      headers?: Record<string, string>;
-    },
-  ): Promise<{ embeddings: number[][] }> {
-    const embeddingModel = this.options.textEmbeddingModel;
-    assert(
-      embeddingModel,
-      "a textEmbeddingModel is required to be set on the Agent that you're doing vector search with",
-    );
-    const result = await embedMany({
-      ...this.options.callSettings,
-      model: embeddingModel,
-      values: options.values,
-      abortSignal: options.abortSignal,
-      headers: options.headers,
-    });
-    if (this.options.usageHandler && result.usage) {
-      await this.options.usageHandler(ctx, {
-        userId: options.userId,
-        threadId: options.threadId,
-        agentName: this.options.name,
-        model: getModelName(embeddingModel),
-        provider: getProviderName(embeddingModel),
-        providerMetadata: undefined,
-        usage: {
-          inputTokens: result.usage.tokens,
-          outputTokens: 0,
-          totalTokens: result.usage.tokens,
-        },
-      });
-    }
-    return { embeddings: result.embeddings };
-  }
-
   /**
    * WORKFLOW UTILITIES
    */
@@ -1925,19 +1482,20 @@ export class Agent<
       handler: async (ctx_, args) => {
         const stream =
           args.stream === true ? spec?.stream || true : (spec?.stream ?? false);
-        const targetArgs = { userId: args.userId, threadId: args.threadId };
+        const { userId, threadId, prompt, messages, maxSteps, ...rest } = args;
+        const targetArgs = { userId, threadId };
         const llmArgs = {
-          stopWhen: spec?.stopWhen ?? this.options.stopWhen,
+          stopWhen: spec?.stopWhen,
           ...overrides,
-          ...omit(args, ["storageOptions", "contextOptions"]),
-          messages: args.messages?.map(deserializeMessage),
-          prompt: Array.isArray(args.prompt)
-            ? args.prompt.map(deserializeMessage)
-            : args.prompt,
+          ...omit(rest, ["storageOptions", "contextOptions", "stream"]),
+          messages: messages?.map(deserializeMessage),
+          prompt: Array.isArray(prompt)
+            ? prompt.map(deserializeMessage)
+            : prompt,
           toolChoice: args.toolChoice as ToolChoice<AgentTools>,
         } satisfies StreamingTextArgs<AgentTools>;
-        if (args.maxSteps) {
-          llmArgs.stopWhen = stepCountIs(args.maxSteps);
+        if (maxSteps) {
+          llmArgs.stopWhen = stepCountIs(maxSteps);
         }
         const opts = {
           ...pick(spec, ["contextOptions", "storageOptions"]),
@@ -1993,11 +1551,8 @@ export class Agent<
    * the normal parameters to {@link generateObject}, plus {@link ContextOptions}
    * and stopWhen.
    */
-  asObjectAction<T>(
-    objectArgs: Omit<
-      Parameters<typeof generateObject<FlexibleSchema<T>>>[0],
-      "model"
-    >,
+  asObjectAction<T, DataModel extends GenericDataModel>(
+    objectArgs: GenerateObjectArgs<FlexibleSchema<T>>,
     options?: Options & MaybeCustomCtx<CustomCtx, DataModel, AgentTools>,
   ) {
     return internalActionGeneric({
@@ -2014,7 +1569,7 @@ export class Agent<
           prompt: Array.isArray(args.prompt)
             ? args.prompt.map(deserializeMessage)
             : args.prompt,
-        } as Omit<Parameters<typeof generateObject>[0], "model">;
+        } as GenerateObjectArgs<FlexibleSchema<T>>;
         const ctx = (
           options?.customCtx
             ? { ...ctx_, ...options.customCtx(ctx_, targetArgs, llmArgs) }