@convex-dev/agent 0.0.17-alpha.2 → 0.1.0

package/src/client/index.ts

@@ -27,7 +27,7 @@ import {
 } from "ai";
 import { assert } from "convex-helpers";
 import { internalActionGeneric, internalMutationGeneric } from "convex/server";
-import { Infer, v } from "convex/values";
+import { v } from "convex/values";
 import { z } from "zod";
 import { Mounts } from "../component/_generated/api.js";
 import {
@@ -57,10 +57,9 @@ import {
   type ProviderOptions,
   type SearchOptions,
   type Usage,
-  vFileWithStringId,
+  vMessageWithMetadata,
   vSafeObjectArgs,
   vTextArgs,
-  vThreadStatus,
 } from "../validators.js";
 import type {
   OpaqueIds,
@@ -69,12 +68,16 @@ import type {
   RunQueryCtx,
   UseApi,
 } from "./types.js";
-import schema from "../component/schema.js";
 
+import type { MessageDoc, ThreadDoc } from "../component/schema.js";
+
+export { vMessageDoc, vThreadDoc } from "../component/schema.js";
 export { extractText, isTool };
-export type { Usage, ProviderMetadata };
+export type { Usage, ProviderMetadata, MessageDoc, ThreadDoc };
 export {
+  /** @deprecated Use vPaginationResult instead. */
   paginationResultValidator,
+  paginationResultValidator as vPaginationResult,
   vContextOptions,
   vUsage,
   vProviderMetadata,
@@ -261,11 +264,11 @@ export class Agent<AgentTools extends ToolSet> {
        */
       userId?: string;
       /**
-       * The title of the thread. Not currently used.
+       * The title of the thread. Not currently used for anything.
        */
       title?: string;
       /**
-       * The summary of the thread. Not currently used.
+       * The summary of the thread. Not currently used for anything.
        */
       summary?: string;
       /**
@@ -301,11 +304,11 @@ export class Agent<AgentTools extends ToolSet> {
        */
       userId?: string;
       /**
-       * The title of the thread. Not currently used.
+       * The title of the thread. Not currently used for anything.
        */
       title?: string;
       /**
-       * The summary of the thread. Not currently used.
+       * The summary of the thread. Not currently used for anything.
        */
       summary?: string;
       /**
@@ -418,37 +421,23 @@ export class Agent<AgentTools extends ToolSet> {
       threadId: string | undefined;
       messages: CoreMessage[];
       /**
-       * If provided, it will search for messages before this message.
-       * Note: if this is far in the past, the search results may be more
+       * 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.
        */
-      beforeMessageId?: string;
+      upToAndIncludingMessageId?: string;
       contextOptions: ContextOptions | undefined;
     }
   ): Promise<MessageDoc[]> {
     assert(args.userId || args.threadId, "Specify userId or threadId");
     // Fetch the latest messages from the thread
-    const contextMessages: MessageDoc[] = [];
     let included: Set<string> | undefined;
     const opts = this.mergedContextOptions(args.contextOptions);
-    if (opts.searchOptions?.textSearch || opts.searchOptions?.vectorSearch) {
-      if (!("runAction" in ctx)) {
-        throw new Error("searchUserMessages only works in an action");
-      }
-      const searchMessages = await ctx.runAction(
-        this.component.messages.searchMessages,
-        {
-          userId: opts?.searchOtherThreads ? args.userId : undefined,
-          threadId: args.threadId,
-          beforeMessageId: args.beforeMessageId,
-          ...(await this.searchOptionsWithDefaults(opts, args.messages)),
-        }
-      );
-      // TODO: track what messages we used for context
-      included = new Set(searchMessages.map((m) => m._id));
-      contextMessages.push(...searchMessages);
-    }
-    if (args.threadId && opts.recentMessages !== 0) {
+    const contextMessages: MessageDoc[] = [];
+    if (
+      args.threadId &&
+      (opts.recentMessages !== 0 || args.upToAndIncludingMessageId)
+    ) {
       const { page } = await ctx.runQuery(
         this.component.messages.listMessagesByThreadId,
         {
@@ -459,14 +448,39 @@ export class Agent<AgentTools extends ToolSet> {
             numItems: opts.recentMessages ?? DEFAULT_RECENT_MESSAGES,
             cursor: null,
           },
-          beforeMessageId: args.beforeMessageId,
+          upToAndIncludingMessageId: args.upToAndIncludingMessageId,
           order: "desc",
           statuses: ["success"],
         }
       );
+      included = new Set(page.map((m) => m._id));
       contextMessages.push(
         // Reverse since we fetched in descending order
-        ...page.filter((m) => !included?.has(m._id)).reverse()
+        ...page.reverse()
+      );
+    }
+    if (opts.searchOptions?.textSearch || opts.searchOptions?.vectorSearch) {
+      const targetMessage = contextMessages.find(
+        (m) => m._id === args.upToAndIncludingMessageId
+      )?.message;
+      const messagesToSearch = targetMessage
+        ? [targetMessage, ...args.messages]
+        : args.messages;
+      if (!("runAction" in ctx)) {
+        throw new Error("searchUserMessages only works in an action");
+      }
+      const searchMessages = await ctx.runAction(
+        this.component.messages.searchMessages,
+        {
+          userId: opts?.searchOtherThreads ? args.userId : undefined,
+          threadId: args.threadId,
+          beforeMessageId: args.upToAndIncludingMessageId,
+          ...(await this.searchOptionsWithDefaults(opts, messagesToSearch)),
+        }
+      );
+      // TODO: track what messages we used for context
+      contextMessages.unshift(
+        ...searchMessages.filter((m) => !included?.has(m._id))
       );
     }
     // Ensure we don't include tool messages without a corresponding tool call
@@ -534,6 +548,14 @@ export class Agent<AgentTools extends ToolSet> {
     args: {
       threadId: string;
       userId?: string;
+      /**
+       * The message that these messages are in response to. They will be
+       * the same "order" as this message, at increasing stepOrder(s).
+       */
+      promptMessageId?: string;
+      /**
+       * The messages to save.
+       */
       messages: CoreMessageMaybeWithId[];
       /**
        * Metadata to save with the messages. Each element corresponds to the
@@ -561,14 +583,11 @@ export class Agent<AgentTools extends ToolSet> {
       threadId: args.threadId,
       userId: args.userId,
       agentName: this.options.name,
+      promptMessageId: args.promptMessageId,
+      embeddings,
       messages: args.messages.map(
         (m, i) =>
           ({
-            embedding: embeddings?.vectors[i] && {
-              model: embeddings.model,
-              dimension: embeddings.dimension,
-              vector: embeddings.vectors[i],
-            },
             ...args.metadata?.[i],
             message: serializeMessage(m),
           }) as MessageWithMetadata
@@ -582,10 +601,40 @@ export class Agent<AgentTools extends ToolSet> {
     };
   }
 
+  /**
+   * Save messages to the thread.
+   * Useful as a step in Workflows, e.g.
+   * ```ts
+   * const saveMessages = agent.asSaveMessagesMutation();
+   *
+   * const myWorkflow = workflow.define()
+   * ```
+   * @returns A mutation that can be used to save messages to the thread.
+   */
+  asSaveMessagesMutation() {
+    return internalMutationGeneric({
+      args: {
+        threadId: v.string(),
+        userId: v.optional(v.string()),
+        promptMessageId: v.optional(v.string()),
+        messages: v.array(vMessageWithMetadata),
+        pending: v.optional(v.boolean()),
+        failPendingSteps: v.optional(v.boolean()),
+      },
+      handler: async (ctx, args) => {
+        return this.saveMessages(ctx, {
+          ...args,
+          messages: args.messages.map((m) => m.message),
+          metadata: args.messages.map(({ message: _, ...m }) => m),
+        });
+      },
+    });
+  }
+
   /**
    * Explicitly save a "step" created by the AI SDK.
    * @param ctx The ctx argument to a mutation or action.
-   * @param args What to save
+   * @param args The Step generated by the AI SDK.
    */
   async saveStep<TOOLS extends ToolSet>(
     ctx: RunMutationCtx,
@@ -595,7 +644,7 @@ export class Agent<AgentTools extends ToolSet> {
       /**
        * The message this step is in response to.
        */
-      parentMessageId: string;
+      promptMessageId: string;
       /**
        * The step to save, possibly including multiple tool calls.
        */
@@ -620,20 +669,11 @@ export class Agent<AgentTools extends ToolSet> {
     const embeddings = await this.generateEmbeddings(
       messages.map((m) => m.message)
     );
-    if (embeddings) {
-      const { model, dimension, vectors } = embeddings;
-      for (let i = 0; i < messages.length; i++) {
-        const vector = vectors[i];
-        if (vector) {
-          messages[i].embedding = { model, dimension, vector };
-        }
-      }
-    }
     await ctx.runMutation(this.component.messages.addStep, {
       userId: args.userId,
       threadId: args.threadId,
-      parentMessageId: args.parentMessageId,
-      step: { step, messages },
+      promptMessageId: args.promptMessageId,
+      step: { step, messages, embeddings },
       failPendingSteps: false,
     });
   }
@@ -698,9 +738,7 @@ export class Agent<AgentTools extends ToolSet> {
        * set in the agent constructor.
        */
       usageHandler?: UsageHandler;
-      /**
-       * The tools to use for this thread. Overrides any tools passed in the agent constructor.
-       */
+      /** @deprecated Pass `tools` in the next parameter instead. This is only intended to pass through thread-default tools.  */
       tools?: ToolSet;
     },
     args: TextArgs<AgentTools, TOOLS, OUTPUT, OUTPUT_PARTIAL>,
@@ -709,7 +747,7 @@ export class Agent<AgentTools extends ToolSet> {
     GenerateTextResult<TOOLS extends undefined ? AgentTools : TOOLS, OUTPUT> &
       GenerationOutputMetadata
   > {
-    const { args: aiArgs, messageId } = await this.saveMessagesAndFetchContext(
+    const { args: aiArgs, messageId } = await this._saveMessagesAndFetchContext(
       ctx,
       args,
       { userId, threadId, ...options }
@@ -721,24 +759,20 @@ export class Agent<AgentTools extends ToolSet> {
     ) as TOOLS extends undefined ? AgentTools : TOOLS;
     const saveOutputMessages =
       options?.storageOptions?.saveOutputMessages ??
-      args.saveOutputMessages ??
       this.options.storageOptions?.saveOutputMessages;
-    const model = aiArgs.model ?? this.options.chat;
     const trackUsage = usageHandler ?? this.options.usageHandler;
     try {
       const result = (await generateText({
         // Can be overridden
         maxSteps: this.options.maxSteps,
-        maxRetries: this.options.maxRetries,
         ...aiArgs,
-        model,
         tools,
         onStepFinish: async (step) => {
           if (threadId && messageId && saveOutputMessages !== false) {
             await this.saveStep(ctx, {
               userId,
               threadId,
-              parentMessageId: messageId,
+              promptMessageId: messageId,
               step,
             });
           }
@@ -747,8 +781,8 @@ export class Agent<AgentTools extends ToolSet> {
               userId,
               threadId,
               agentName: this.options.name,
-              model: model.modelId,
-              provider: model.provider,
+              model: aiArgs.model.modelId,
+              provider: aiArgs.model.provider,
               usage: step.usage,
               providerMetadata: step.providerMetadata,
             });
@@ -780,11 +814,6 @@ export class Agent<AgentTools extends ToolSet> {
    * 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 ctx The context passed from the action function calling this.
-   * @param { userId, threadId }: The user and thread to associate the message with
-   * @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.
    */
   async streamText<
     TOOLS extends ToolSet | undefined = undefined,
@@ -807,7 +836,14 @@ export class Agent<AgentTools extends ToolSet> {
       usageHandler?: UsageHandler;
       tools?: ToolSet;
     },
+    /**
+     * The arguments to the streamText function, similar to the ai `streamText` function.
+     */
     args: StreamingTextArgs<AgentTools, TOOLS, OUTPUT, PARTIAL_OUTPUT>,
+    /**
+     * The {@link ContextOptions} and {@link StorageOptions}
+     * options to use for fetching contextual messages and saving input/output messages.
+     */
     options?: Options
   ): Promise<
     StreamTextResult<
@@ -816,7 +852,7 @@ export class Agent<AgentTools extends ToolSet> {
     > &
       GenerationOutputMetadata
   > {
-    const { args: aiArgs, messageId } = await this.saveMessagesAndFetchContext(
+    const { args: aiArgs, messageId } = await this._saveMessagesAndFetchContext(
       ctx,
       args,
       { userId, threadId, ...options }
@@ -828,16 +864,12 @@ export class Agent<AgentTools extends ToolSet> {
     ) as TOOLS extends undefined ? AgentTools : TOOLS;
     const saveOutputMessages =
       options?.storageOptions?.saveOutputMessages ??
-      args.saveOutputMessages ??
       this.options.storageOptions?.saveOutputMessages;
-    const model = aiArgs.model ?? this.options.chat;
     const trackUsage = usageHandler ?? this.options.usageHandler;
     const result = streamText({
       // Can be overridden
       maxSteps: this.options.maxSteps,
-      maxRetries: this.options.maxRetries,
       ...aiArgs,
-      model,
       tools,
       onChunk: async (chunk) => {
         // console.log("onChunk", chunk);
@@ -860,7 +892,7 @@ export class Agent<AgentTools extends ToolSet> {
           await this.saveStep(ctx, {
             userId,
             threadId,
-            parentMessageId: messageId,
+            promptMessageId: messageId,
             step,
           });
         }
@@ -869,8 +901,8 @@ export class Agent<AgentTools extends ToolSet> {
             userId,
             threadId,
             agentName: this.options.name,
-            model: model.modelId,
-            provider: model.provider,
+            model: aiArgs.model.modelId,
+            provider: aiArgs.model.provider,
             usage: step.usage,
             providerMetadata: step.providerMetadata,
           });
@@ -886,12 +918,15 @@ export class Agent<AgentTools extends ToolSet> {
     return result;
   }
 
-  async saveMessagesAndFetchContext<
+  async _saveMessagesAndFetchContext<
     T extends {
       id?: string;
       prompt?: string;
       messages?: CoreMessage[] | AIMessageWithoutId[];
       system?: string;
+      promptMessageId?: string;
+      model?: LanguageModelV1;
+      maxRetries?: number;
     },
   >(
     ctx: RunActionCtx | RunMutationCtx,
@@ -906,20 +941,33 @@ export class Agent<AgentTools extends ToolSet> {
       threadId: string | undefined;
     } & Options
   ): Promise<{
-    args: T;
+    args: T & { model: LanguageModelV1 };
     messageId: string | undefined;
   }> {
-    contextOptions ||= this.options.contextOptions ?? (args as ContextOptions);
-    storageOptions ||= this.options.storageOptions ?? (args as StorageOptions);
-    const messages = promptOrMessagesToCoreMessages(args);
+    contextOptions ||= this.options.contextOptions;
+    storageOptions ||= this.options.storageOptions;
+    // If only a messageId is provided, this will be empty.
+    const messages = args.promptMessageId
+      ? []
+      : promptOrMessagesToCoreMessages(args);
+    assert(
+      !args.promptMessageId || !(args.prompt || args.messages),
+      "you can't specify a prompt or message if you specify a promptMessageId"
+    );
+    // If only a messageId is provided, this will add that message to the end.
     const contextMessages = await this.fetchContextMessages(ctx, {
       userId,
       threadId,
+      upToAndIncludingMessageId: args.promptMessageId,
       messages,
       contextOptions,
     });
-    let messageId: string | undefined;
-    if (threadId && storageOptions?.saveAnyInputMessages !== false) {
+    let messageId = args.promptMessageId;
+    if (
+      threadId &&
+      messages.length &&
+      storageOptions?.saveAnyInputMessages !== false
+    ) {
       const saveAll = storageOptions?.saveAllInputMessages;
       const coreMessages = saveAll ? messages : messages.slice(-1);
       const saved = await this.saveMessages(ctx, {
@@ -932,16 +980,18 @@ export class Agent<AgentTools extends ToolSet> {
       });
       messageId = saved.lastMessageId;
     }
-    const { prompt: _, ...rest } = args;
+    const { prompt: _, model, ...rest } = args;
     return {
       args: {
         ...rest,
+        maxRetries: args.maxRetries ?? this.options.maxRetries,
+        model: model ?? this.options.chat,
         system: args.system ?? this.options.instructions,
         messages: [
           ...contextMessages.map((m) => deserializeMessage(m.message!)),
           ...messages,
         ],
-      } as T,
+      } as T & { model: LanguageModelV1 },
       messageId,
     };
   }
@@ -952,11 +1002,6 @@ export class Agent<AgentTools extends ToolSet> {
    * 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 ctx The context passed from the action function calling this.
-   * @param { userId, threadId }: The user and thread to associate the message with
-   * @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.
    */
   async generateObject<T>(
     ctx: RunActionCtx,
@@ -965,33 +1010,35 @@ export class Agent<AgentTools extends ToolSet> {
       threadId,
       usageHandler,
     }: { userId?: string; threadId?: string; usageHandler?: UsageHandler },
+    /**
+     * The arguments to the generateObject function, similar to the ai.generateObject function.
+     */
     args: OurObjectArgs<T>,
+    /**
+     * The {@link ContextOptions} and {@link StorageOptions}
+     * options to use for fetching contextual messages and saving input/output messages.
+     */
     options?: Options
   ): Promise<GenerateObjectResult<T> & GenerationOutputMetadata> {
-    const { args: aiArgs, messageId } = await this.saveMessagesAndFetchContext(
+    const { args: aiArgs, messageId } = await this._saveMessagesAndFetchContext(
       ctx,
       args,
       { userId, threadId, ...options }
     );
-    const model = aiArgs.model ?? this.options.chat;
     const trackUsage = usageHandler ?? this.options.usageHandler;
     const saveOutputMessages =
       options?.storageOptions?.saveOutputMessages ??
-      args.saveOutputMessages ??
       this.options.storageOptions?.saveOutputMessages;
     try {
-      const result = (await generateObject({
-        // Can be overridden
-        maxRetries: this.options.maxRetries,
-        ...aiArgs,
-        model,
+      const result = (await generateObject(
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      } as any)) as GenerateObjectResult<T> & GenerationOutputMetadata;
+        aiArgs as any
+      )) as GenerateObjectResult<T> & GenerationOutputMetadata;
 
       if (threadId && messageId && saveOutputMessages !== false) {
         await this.saveObject(ctx, {
           threadId,
-          parentMessageId: messageId,
+          promptMessageId: messageId,
           result,
           userId,
         });
@@ -1002,8 +1049,8 @@ export class Agent<AgentTools extends ToolSet> {
           userId,
           threadId,
           agentName: this.options.name,
-          model: model.modelId,
-          provider: model.provider,
+          model: aiArgs.model.modelId,
+          provider: aiArgs.model.provider,
           usage: result.usage,
           providerMetadata: result.providerMetadata,
         });
@@ -1021,16 +1068,11 @@ export class Agent<AgentTools extends ToolSet> {
   }
 
   /**
-   * This behaves like {@link streamObject} from the "ai" package except that
+   * This behaves like `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 ctx The context passed from the action function calling this.
-   * @param { userId, threadId }: The user and thread to associate the message with
-   * @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.
    */
   async streamObject<T>(
     ctx: RunActionCtx,
@@ -1039,29 +1081,31 @@ export class Agent<AgentTools extends ToolSet> {
       threadId,
       usageHandler,
     }: { userId?: string; threadId?: string; usageHandler?: UsageHandler },
+    /**
+     * The arguments to the streamObject function, similar to the ai `streamObject` function.
+     */
     args: OurStreamObjectArgs<T>,
+    /**
+     * The {@link ContextOptions} and {@link StorageOptions}
+     * options to use for fetching contextual messages and saving input/output messages.
+     */
     options?: Options
   ): Promise<
     StreamObjectResult<DeepPartial<T>, T, never> & GenerationOutputMetadata
   > {
     // TODO: unify all this shared code between all the generate* and stream* functions
-    const { args: aiArgs, messageId } = await this.saveMessagesAndFetchContext(
+    const { args: aiArgs, messageId } = await this._saveMessagesAndFetchContext(
       ctx,
       args,
       { userId, threadId, ...options }
     );
-    const model = aiArgs.model ?? this.options.chat;
     const trackUsage = usageHandler ?? this.options.usageHandler;
     const saveOutputMessages =
       options?.storageOptions?.saveOutputMessages ??
-      args.saveOutputMessages ??
       this.options.storageOptions?.saveOutputMessages;
     const stream = streamObject<T>({
-      // Can be overridden
-      maxRetries: this.options.maxRetries,
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       ...(aiArgs as any),
-      model,
       onError: async (error) => {
         console.error("onError", error);
         return args.onError?.(error);
@@ -1071,7 +1115,7 @@ export class Agent<AgentTools extends ToolSet> {
           await this.saveObject(ctx, {
             userId,
             threadId,
-            parentMessageId: messageId,
+            promptMessageId: messageId,
             result: {
               object: result.object,
               finishReason: "stop",
@@ -1092,8 +1136,8 @@ export class Agent<AgentTools extends ToolSet> {
             userId,
             threadId,
             agentName: this.options.name,
-            model: model.modelId,
-            provider: model.provider,
+            model: aiArgs.model.modelId,
+            provider: aiArgs.model.provider,
             usage: result.usage,
             providerMetadata: result.providerMetadata,
           });
@@ -1119,38 +1163,25 @@ export class Agent<AgentTools extends ToolSet> {
     args: {
       userId: string | undefined;
       threadId: string;
-      parentMessageId: string;
+      promptMessageId: string;
       result: GenerateObjectResult<unknown>;
       metadata?: Omit<MessageWithMetadata, "message">;
     }
   ): Promise<void> {
-    const { step, messages: withoutEmbed } = serializeObjectResult(
-      args.result,
-      {
-        model: this.options.chat.modelId,
-        provider: this.options.chat.provider,
-      }
+    const { step, messages } = serializeObjectResult(args.result, {
+      model: this.options.chat.modelId,
+      provider: this.options.chat.provider,
+    });
+    const embeddings = await this.generateEmbeddings(
+      messages.map((m) => m.message)
     );
-    const embeddings = await this.generateEmbeddings([withoutEmbed[0].message]);
-    const messages = embeddings?.vectors[0]
-      ? [
-          {
-            ...withoutEmbed[0],
-            embedding: {
-              dimension: embeddings.dimension,
-              model: embeddings.model,
-              vector: embeddings.vectors[0],
-            },
-          },
-        ]
-      : withoutEmbed;
 
     await ctx.runMutation(this.component.messages.addStep, {
       userId: args.userId,
       threadId: args.threadId,
-      parentMessageId: args.parentMessageId,
+      promptMessageId: args.promptMessageId,
       failPendingSteps: false,
-      step: { step, messages },
+      step: { step, messages, embeddings },
     });
   }
 
@@ -1239,11 +1270,11 @@ export class Agent<AgentTools extends ToolSet> {
    * Create an action out of this agent so you can call it from workflows or other actions
    * without a wrapping function.
    * @param spec Configuration for the agent acting as an action, including
-   *   {@link ContextOptions} and maxSteps.
+   *   {@link ContextOptions}, {@link StorageOptions}, and maxSteps.
    */
   asTextAction(spec?: {
-    contextOptions?: ContextOptions;
     maxSteps?: number;
+    contextOptions?: ContextOptions;
     storageOptions?: StorageOptions;
   }) {
     const maxSteps = spec?.maxSteps ?? this.options.maxSteps;
@@ -1308,7 +1339,7 @@ export class Agent<AgentTools extends ToolSet> {
               this.options.storageOptions,
           }
         );
-        return value.object as T;
+        return { object: value.object as T };
       },
     });
   }
@@ -1341,7 +1372,6 @@ export function filterOutOrphanedToolMessages(docs: MessageDoc[]) {
   return result;
 }
 
-
 export type ToolCtx = RunActionCtx & {
   userId?: string;
   threadId?: string;
@@ -1468,6 +1498,13 @@ type TextArgs<
   >[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.
@@ -1483,8 +1520,7 @@ type TextArgs<
    * specified in the tools array. e.g. {toolName: "getWeather", type: "tool"}
    */
   toolChoice?: ToolChoice<TOOLS extends undefined ? AgentTools : TOOLS>;
-} & ContextOptions & // DEPRECATED: pass them in the subsequent parameter instead
-  StorageOptions;
+};
 
 type StreamingTextArgs<
   AgentTools extends ToolSet,
@@ -1516,38 +1552,42 @@ type StreamingTextArgs<
    * specified in the tools array. e.g. {toolName: "getWeather", type: "tool"}
    */
   toolChoice?: ToolChoice<TOOLS extends undefined ? AgentTools : TOOLS>;
-} & ContextOptions &
-  StorageOptions;
+};
 
-type BaseGenerateObjectOptions = StorageOptions &
-  ContextOptions &
-  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[];
-    experimental_repairText?: RepairTextFunction;
-    experimental_telemetry?: TelemetrySettings;
-    providerOptions?: ProviderOptions;
-    experimental_providerMetadata?: ProviderMetadata;
-  };
+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 & {
@@ -1724,30 +1764,6 @@ interface Thread<DefaultTools extends ToolSet> {
   >;
 }
 
-export const vThreadDoc = v.object({
-  _id: v.string(),
-  _creationTime: v.number(),
-  userId: v.optional(v.string()), // Unset for anonymous
-  title: v.optional(v.string()),
-  summary: v.optional(v.string()),
-  status: vThreadStatus,
-});
-export type ThreadDoc = Infer<typeof vThreadDoc>;
-
-export const vMessageDoc = v.object({
-  _id: v.string(),
-  _creationTime: v.number(),
-  ...schema.tables.messages.validator.fields,
-  // Overwrite all the types that have a v.id validator
-  // Outside of the component, they are strings
-  threadId: v.string(),
-  parentMessageId: v.optional(v.string()),
-  stepId: v.optional(v.string()),
-  embeddingId: v.optional(v.string()),
-  files: v.optional(v.array(vFileWithStringId)),
-});
-export type MessageDoc = Infer<typeof vMessageDoc>;
-
 type MessageWithMetadata = OpaqueIds<InnerMessageWithMetadata>;
 
 export function toUIMessages(messages: MessageDoc[]): UIMessage[] {