@convex-dev/agent 0.0.1-alpha.1 → 0.0.1-alpha.3

package/src/client/index.ts

@@ -11,26 +11,24 @@ import type {
   ToolChoice,
   ToolExecutionOptions,
   ToolSet,
-  Message as UIMessage,
 } from "ai";
-import {
-  convertToCoreMessages,
-  coreMessageSchema,
-  generateObject,
-  generateText,
-  streamObject,
-  streamText,
-} from "ai";
-import type { ZodType } from "zod";
+import { generateObject, generateText, streamObject, streamText } from "ai";
 import { api } from "../component/_generated/api";
-import { Message, MessageStatus, SearchOptions } from "../validators";
+import {
+  SearchOptions,
+  vThreadArgs,
+  vContextOptions,
+  vObjectArgs,
+  vStorageOptions,
+} from "../validators";
 import { RunActionCtx, RunMutationCtx, RunQueryCtx, UseApi } from "./types";
 // TODO: is this the only dependency that needs helpers in client?
 import { assert } from "convex-helpers";
-import { convexToZod } from "convex-helpers/server/zod";
-import { GenericActionCtx, GenericDataModel } from "convex/server";
-import { Infer, Validator } from "convex/values";
+import { ConvexToZod, convexToZod } from "convex-helpers/server/zod";
+import { internalActionGeneric } from "convex/server";
+import { Infer, v, Validator } from "convex/values";
 import {
+  promptOrMessagesToCoreMessages,
   serializeMessageWithId,
   serializeNewMessagesInStep,
   serializeStep,
@@ -38,22 +36,49 @@ import {
 import { DEFAULT_MESSAGE_RANGE, extractText } from "../shared";
 
 export type ContextOptions = {
-  parentMessageId?: string;
-  includeToolMessages?: boolean;
+  /**
+   * Whether to include tool messages in the context.
+   */
+  includeToolCalls?: boolean;
+  /**
+   * How many recent messages to include. These are added after the search
+   * messages, and do not count against the search limit.
+   */
   recentMessages?: number;
+  /**
+   * Options for searching messages.
+   */
   searchOptions?: {
+    /**
+     * The maximum number of messages to fetch.
+     */
     limit: number;
+    /**
+     * Whether to use text search to find messages.
+     */
     textSearch?: boolean;
+    /**
+     * Whether to use vector search to find messages.
+     */
     vectorSearch?: boolean;
-    messageRange: { before: number; after: number };
+    /**
+     * Note, this is after the limit is applied.
+     * By default this will quadruple the number of messages fetched.
+     * (two before, and one after each message found in the search)
+     */
+    messageRange?: { before: number; after: number };
   };
-  searchOtherChats?: boolean;
+  /**
+   * Whether to search across other threads for relevant messages.
+   * By default, only the current thread is searched.
+   */
+  searchOtherThreads?: boolean;
 };
 
 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 chat history.
+  // Pass true to have all input messages saved to the thread history.
   saveAllInputMessages?: boolean;
   // Defaults to true
   saveOutputMessages?: boolean;
@@ -68,122 +93,136 @@ export class Agent<AgentTools extends ToolSet> {
     public component: UseApi<typeof api>,
     public options: {
       name?: string;
-      chat: LanguageModelV1;
+      thread: LanguageModelV1;
       textEmbedding?: EmbeddingModelV1<string>;
-      defaultSystemPrompt?: string;
+      instructions?: string;
       tools?: AgentTools;
+      contextOptions?: ContextOptions;
+      // TODO: storageOptions?: StorageOptions;
+      maxSteps?: number;
+      // TODO: maxRetries?: number;
     }
   ) {}
 
   /**
-   * Start a new chat with the agent. This will have a fresh history, though if
-   * you pass in a userId you can have it search across other chats for relevant
+   * Start a new thread with the agent. This will have a fresh history, though if
+   * you pass in a userId you can have it search across other threads for relevant
    * messages as context for the LLM calls.
-   * @param ctx The context of the Convex function. From an action, you can chat
-   *   with the agent. From a mutation, you can start a chat and save the chatId
-   *   to pass to continueChat later.
-   * @param args The chat metadata.
-   * @returns The chatId of the new chat and the chat object.
+   * @param ctx The context of the Convex function. From an action, you can thread
+   *   with the agent. From a mutation, you can start a thread and save the threadId
+   *   to pass to continueThread later.
+   * @param args The thread metadata.
+   * @returns The threadId of the new thread and the thread object.
    */
-  async startChat(
+  async createThread(
     ctx: RunActionCtx,
     args: {
       /**
-       * The userId to associate with the chat. If not provided, the chat will be
+       * The userId to associate with the thread. If not provided, the thread will be
        * anonymous.
        */
       userId?: string;
       /**
-       * The parent chatIds to merge with.
-       * If the chat is a continuation of one or many previous chats,
-       * you can pass in the chatIds of the parent chats to merge the histories.
+       * The parent threadIds to merge with.
+       * If the thread is a continuation of one or many previous threads,
+       * you can pass in the threadIds of the parent threads to merge the histories.
        */
-      parentChatIds?: string[];
+      parentThreadIds?: string[];
       /**
-       * The title of the chat. Not currently used.
+       * The title of the thread. Not currently used.
        */
       title?: string;
       /**
-       * The summary of the chat. Not currently used.
+       * The summary of the thread. Not currently used.
        */
       summary?: string;
     }
   ): Promise<{
-    chatId: string;
-    chat: Chat<AgentTools>;
+    threadId: string;
+    thread: Thread<AgentTools>;
   }>;
   /**
-   * Start a new chat with the agent. This will have a fresh history, though if
-   * you pass in a userId you can have it search across other chats for relevant
+   * Start a new thread with the agent. This will have a fresh history, though if
+   * you pass in a userId you can have it search across other threads for relevant
    * messages as context for the LLM calls.
    * @param ctx The context of the Convex function. From a mutation, you can
-   * start a chat and save the chatId to pass to continueChat later.
-   * @param args The chat metadata.
-   * @returns The chatId of the new chat.
+   * start a thread and save the threadId to pass to continueThread later.
+   * @param args The thread metadata.
+   * @returns The threadId of the new thread.
    */
-  async startChat(
+  async createThread(
     ctx: RunMutationCtx,
     args: {
       userId?: string;
-      parentChatIds?: string[];
+      parentThreadIds?: string[];
       title?: string;
       summary?: string;
     }
   ): Promise<{
-    chatId: string;
+    threadId: string;
   }>;
-  async startChat(
+  async createThread(
     ctx: RunActionCtx | RunMutationCtx,
     args: {
       userId: string;
-      parentChatIds?: string[];
+      parentThreadIds?: string[];
       title?: string;
       summary?: string;
     }
   ): Promise<{
-    chatId: string;
-    chat?: Chat<AgentTools>;
+    threadId: string;
+    thread?: Thread<AgentTools>;
   }> {
-    const chatDoc = await ctx.runMutation(this.component.messages.createChat, {
-      defaultSystemPrompt: this.options.defaultSystemPrompt,
-      userId: args.userId,
-      title: args.title,
-      summary: args.summary,
-      parentChatIds: args.parentChatIds,
-    });
+    const threadDoc = await ctx.runMutation(
+      this.component.messages.createThread,
+      {
+        defaultSystemPrompt: this.options.instructions,
+        userId: args.userId,
+        title: args.title,
+        summary: args.summary,
+        parentThreadIds: args.parentThreadIds,
+      }
+    );
     if (!("runAction" in ctx)) {
-      return { chatId: chatDoc._id };
+      return { threadId: threadDoc._id };
     }
-    const { chat } = await this.continueChat(ctx, {
-      chatId: chatDoc._id,
+    const { thread } = await this.continueThread(ctx, {
+      threadId: threadDoc._id,
       userId: args.userId,
     });
     return {
-      chatId: chatDoc._id,
-      chat,
+      threadId: threadDoc._id,
+      thread,
     };
   }
 
-  async continueChat(
+  async continueThread(
     ctx: RunActionCtx,
     {
-      chatId,
+      threadId,
       userId,
     }: {
-      chatId: string;
+      threadId: string;
+      /**
+       * If supplied, the userId can be used to search across other threads for
+       * relevant messages from the same user as context for the LLM calls.
+       */
       userId?: string;
     }
   ): Promise<{
-    chat: Chat<AgentTools>;
+    thread: Thread<AgentTools>;
   }> {
-    // return this.component.continueChat(ctx, args);
+    // return this.component.continueThread(ctx, args);
     return {
-      chat: {
-        generateText: this.generateText.bind(this, ctx, { userId, chatId }),
-        streamText: this.streamText.bind(this, ctx, { userId, chatId }),
-        generateObject: this.generateObject.bind(this, ctx, { userId, chatId }),
-        streamObject: this.streamObject.bind(this, ctx, { userId, chatId }),
-      } as Chat<AgentTools>,
+      thread: {
+        generateText: this.generateText.bind(this, ctx, { userId, threadId }),
+        streamText: this.streamText.bind(this, ctx, { userId, threadId }),
+        generateObject: this.generateObject.bind(this, ctx, {
+          userId,
+          threadId,
+        }),
+        streamObject: this.streamObject.bind(this, ctx, { userId, threadId }),
+      } as Thread<AgentTools>,
     };
   }
 
@@ -191,35 +230,37 @@ export class Agent<AgentTools extends ToolSet> {
     ctx: RunQueryCtx | RunActionCtx,
     args: {
       userId?: string;
-      chatId?: string;
+      threadId?: string;
       messages: CoreMessage[];
+      parentMessageId?: string;
     } & ContextOptions
   ): Promise<CoreMessage[]> {
-    assert(args.userId || args.chatId, "Specify userId or chatId");
-    // Fetch the latest messages from the chat
+    assert(args.userId || args.threadId, "Specify userId or threadId");
+    // Fetch the latest messages from the thread
     const contextMessages: CoreMessage[] = [];
-    if (args.searchOptions?.textSearch || args.searchOptions?.vectorSearch) {
+    const opts = this.mergedContextOptions(args);
+    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: args.searchOtherChats ? args.userId : undefined,
-          chatId: args.chatId,
+          userId: args.searchOtherThreads ? args.userId : undefined,
+          threadId: args.threadId,
           parentMessageId: args.parentMessageId,
-          ...(await this.searchOptionsWithDefaults(args, args.messages)),
+          ...(await this.searchOptionsWithDefaults(opts, args.messages)),
         }
       );
       // TODO: track what messages we used for context
       contextMessages.push(...searchMessages.map((m) => m.message!));
     }
-    if (args.chatId) {
+    if (args.threadId) {
       const { messages } = await ctx.runQuery(
-        this.component.messages.getChatMessages,
+        this.component.messages.getThreadMessages,
         {
-          chatId: args.chatId,
-          isTool: args.includeToolMessages ?? false,
+          threadId: args.threadId,
+          isTool: args.includeToolCalls ?? false,
           limit: args.recentMessages,
           parentMessageId: args.parentMessageId,
           order: "desc",
@@ -234,21 +275,24 @@ export class Agent<AgentTools extends ToolSet> {
   async saveMessages(
     ctx: RunMutationCtx,
     args: {
-      chatId: string;
+      threadId?: string;
+      userId?: string;
       messages: CoreMessageMaybeWithId[];
       pending?: boolean;
       parentMessageId?: string;
+      failPendingSteps?: boolean;
     }
   ): Promise<{
     lastMessageId: string;
     messageIds: string[];
   }> {
     const result = await ctx.runMutation(this.component.messages.addMessages, {
-      chatId: args.chatId,
+      threadId: args.threadId,
+      userId: args.userId,
       agentName: this.options.name,
-      model: this.options.chat.modelId,
+      model: this.options.thread.modelId,
       messages: args.messages.map(serializeMessageWithId),
-      failPendingSteps: true,
+      failPendingSteps: args.failPendingSteps ?? true,
       pending: args.pending ?? false,
       parentMessageId: args.parentMessageId,
     });
@@ -260,12 +304,12 @@ export class Agent<AgentTools extends ToolSet> {
 
   async saveStep<TOOLS extends ToolSet>(
     ctx: RunMutationCtx,
-    args: { chatId: string; messageId: string; step: StepResult<TOOLS> }
+    args: { threadId: string; messageId: string; step: StepResult<TOOLS> }
   ): Promise<void> {
     const step = serializeStep(args.step as StepResult<ToolSet>);
     const messages = serializeNewMessagesInStep(args.step);
     await ctx.runMutation(this.component.messages.addSteps, {
-      chatId: args.chatId,
+      threadId: args.threadId,
       messageId: args.messageId,
       steps: [{ step, messages: messages }],
       failPendingSteps: false,
@@ -276,7 +320,7 @@ export class Agent<AgentTools extends ToolSet> {
   async completeMessage<TOOLS extends ToolSet>(
     ctx: RunMutationCtx,
     args: {
-      chatId: string;
+      threadId: string;
       messageId: string;
       result:
         | { kind: "error"; error: string }
@@ -290,7 +334,7 @@ export class Agent<AgentTools extends ToolSet> {
       });
     } else {
       await ctx.runMutation(this.component.messages.addSteps, {
-        chatId: args.chatId,
+        threadId: args.threadId,
         messageId: args.messageId,
         steps: [],
         failPendingSteps: true,
@@ -300,9 +344,9 @@ export class Agent<AgentTools extends ToolSet> {
 
   /**
    * This behaves like {@link generateText} except that it add context based on
-   * the userId and chatId. It saves the input and resulting messages to the
-   * chat, if specified.
-   * however. To do that, use {@link continueChat} or {@link saveMessages}.
+   * the userId and threadId. It saves the input and resulting messages to the
+   * thread, if specified.
+   * however. To do that, use {@link continueThread} or {@link saveMessages}.
    * @param ctx The context of the agent.
    * @param args The arguments to the generateText function.
    * @returns The result of the generateText function.
@@ -315,10 +359,10 @@ export class Agent<AgentTools extends ToolSet> {
     ctx: RunActionCtx,
     {
       userId,
-      chatId,
+      threadId,
     }: {
       userId?: string;
-      chatId: string;
+      threadId?: string;
     },
     args: TextArgs<
       AgentTools,
@@ -333,30 +377,33 @@ export class Agent<AgentTools extends ToolSet> {
     const contextMessages = await this.fetchContextMessages(ctx, {
       ...args,
       userId,
-      chatId,
+      threadId,
       messages,
     });
     const { lastMessageId: messageId } = await this.saveMessages(ctx, {
-      chatId,
+      threadId,
+      userId,
       messages: args.saveAllInputMessages ? messages : messages.slice(-1),
       pending: true,
       parentMessageId: args.parentMessageId,
     });
-    const defaults = this.options.tools;
-    const tools = wrapTools(ctx, chatId, userId, defaults, args.tools) as TOOLS;
+    const toolCtx = { ...ctx, userId, threadId, messageId };
+    const tools = wrapTools(toolCtx, this.options.tools, args.tools) as TOOLS;
     try {
       const result = await generateText({
-        model: this.options.chat,
+        model: this.options.thread,
         messages: [...contextMessages, ...messages],
-        system: this.options.defaultSystemPrompt,
-        tools,
+        system: this.options.instructions,
+        maxSteps: this.options.maxSteps,
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
         toolChoice: args.toolChoice as any,
         ...rest,
+        tools,
         onStepFinish: async (step) => {
-          if (chatId && messageId && args.saveOutputMessages) {
+          if (threadId && messageId && args.saveOutputMessages) {
+            console.log("onStepFinish", step);
             await this.saveStep(ctx, {
-              chatId,
+              threadId,
               messageId,
               step,
             });
@@ -366,7 +413,8 @@ export class Agent<AgentTools extends ToolSet> {
       });
       return { ...result, messageId };
     } catch (error) {
-      if (chatId && messageId) {
+      if (threadId && messageId) {
+        console.error("RollbackMessage", messageId);
         await ctx.runMutation(this.component.messages.rollbackMessage, {
           messageId,
           error: (error as Error).message,
@@ -382,12 +430,12 @@ export class Agent<AgentTools extends ToolSet> {
     PARTIAL_OUTPUT = never,
   >(
     ctx: RunActionCtx,
-    { userId, chatId }: { userId?: string; chatId: string },
-    args: Partial<
+    { userId, threadId }: { userId?: string; threadId?: string },
+    args: TextArgs<
+      AgentTools,
+      TOOLS,
       Parameters<typeof streamText<TOOLS, OUTPUT, PARTIAL_OUTPUT>>[0]
-    > &
-      ContextOptions &
-      StorageOptions
+    >
   ): Promise<
     StreamTextResult<TOOLS, PARTIAL_OUTPUT> & GenerationOutputMetadata
   > {
@@ -396,32 +444,33 @@ export class Agent<AgentTools extends ToolSet> {
     const contextMessages = await this.fetchContextMessages(ctx, {
       ...args,
       userId,
-      chatId,
+      threadId,
       messages,
     });
     const { lastMessageId: messageId } = await this.saveMessages(ctx, {
-      chatId,
+      threadId,
       messages: args.saveAllInputMessages ? messages : messages.slice(-1),
       pending: true,
       parentMessageId: args.parentMessageId,
     });
-    const defaults = this.options.tools;
-    const tools = wrapTools(ctx, chatId, userId, defaults, args.tools) as TOOLS;
+    const toolCtx = { ...ctx, userId, threadId, messageId };
+    const tools = wrapTools(toolCtx, this.options.tools, args.tools) as TOOLS;
     const result = streamText({
-      model: this.options.chat,
+      model: this.options.thread,
       messages: [...contextMessages, ...messages],
-      system: this.options.defaultSystemPrompt,
-      tools,
+      system: this.options.instructions,
+      maxSteps: this.options.maxSteps,
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       toolChoice: args.toolChoice as any,
       ...rest,
+      tools,
       onChunk: async (chunk) => {
         console.log("onChunk", chunk);
         return args.onChunk?.(chunk);
       },
       onError: async (error) => {
         console.error("onError", error);
-        if (chatId && messageId) {
+        if (threadId && messageId) {
           await ctx.runMutation(this.component.messages.rollbackMessage, {
             messageId,
             error: (error.error as Error).message,
@@ -437,9 +486,9 @@ export class Agent<AgentTools extends ToolSet> {
       },
       onStepFinish: async (step) => {
         console.log("onStepFinish", step);
-        if (chatId && messageId) {
+        if (threadId && messageId) {
           await this.saveStep(ctx, {
-            chatId,
+            threadId,
             messageId,
             step,
           });
@@ -453,10 +502,10 @@ export class Agent<AgentTools extends ToolSet> {
   // TODO: add the crazy number of overloads to get types through
   async generateObject<T>(
     ctx: RunActionCtx,
-    { userId, chatId }: { userId?: string; chatId: string },
+    { userId, threadId }: { userId?: string; threadId?: string },
     args: Omit<Parameters<typeof generateObject>[0], "model"> & {
       model?: LanguageModelV1;
-    } & ContextOptions &
+    } & { parentMessageId?: string } & ContextOptions &
       StorageOptions
   ): Promise<GenerateObjectResult<T> & GenerationOutputMetadata> {
     const { prompt, messages: raw, ...rest } = args;
@@ -464,16 +513,16 @@ export class Agent<AgentTools extends ToolSet> {
     const contextMessages = await this.fetchContextMessages(ctx, {
       ...args,
       userId,
-      chatId,
+      threadId,
       messages,
     });
     const { lastMessageId: messageId } = await this.saveMessages(ctx, {
-      chatId,
+      threadId,
       messages: args.saveAllInputMessages ? messages : messages.slice(-1),
       pending: true,
     });
     const result = (await generateObject({
-      model: this.options.chat,
+      model: this.options.thread,
       messages: [...contextMessages, ...messages],
       ...rest,
     })) as GenerateObjectResult<T>;
@@ -482,10 +531,10 @@ export class Agent<AgentTools extends ToolSet> {
 
   async streamObject<T>(
     ctx: RunMutationCtx,
-    { userId, chatId }: { userId?: string; chatId: string },
+    { userId, threadId }: { userId?: string; threadId?: string },
     args: Omit<Parameters<typeof streamObject<T>>[0], "model"> & {
       model?: LanguageModelV1;
-    } & ContextOptions &
+    } & { parentMessageId?: string } & ContextOptions &
       StorageOptions
   ): Promise<
     StreamObjectResult<DeepPartial<T>, T, never> & GenerationOutputMetadata
@@ -495,16 +544,16 @@ export class Agent<AgentTools extends ToolSet> {
     const contextMessages = await this.fetchContextMessages(ctx, {
       ...args,
       userId,
-      chatId,
+      threadId,
       messages,
     });
     const { lastMessageId: messageId } = await this.saveMessages(ctx, {
-      chatId,
+      threadId,
       messages: args.saveAllInputMessages ? messages : messages.slice(-1),
       pending: true,
     });
     const result = streamObject<T>({
-      model: this.options.chat,
+      model: this.options.thread,
       messages: [...contextMessages, ...messages],
       ...rest,
       onError: async (error) => {
@@ -518,27 +567,41 @@ export class Agent<AgentTools extends ToolSet> {
     return { ...result, messageId };
   }
 
+  mergedContextOptions(opts: ContextOptions): ContextOptions {
+    const searchOptions = {
+      ...this.options.contextOptions?.searchOptions,
+      ...opts.searchOptions,
+    };
+    return {
+      ...this.options.contextOptions,
+      ...opts,
+      searchOptions: searchOptions.limit
+        ? (searchOptions as SearchOptions)
+        : undefined,
+    };
+  }
+
   async searchOptionsWithDefaults(
-    searchArgs: ContextOptions,
+    contextOptions: ContextOptions,
     messages: CoreMessage[]
   ): Promise<SearchOptions> {
     assert(
-      searchArgs.searchOptions?.textSearch ||
-        searchArgs.searchOptions?.vectorSearch,
+      contextOptions.searchOptions?.textSearch ||
+        contextOptions.searchOptions?.vectorSearch,
       "searchOptions is required"
     );
     assert(messages.length > 0, "Core messages cannot be empty");
     const text = extractText(messages.at(-1)!);
     const search: SearchOptions = {
-      limit: searchArgs.searchOptions?.limit ?? 10,
+      limit: contextOptions.searchOptions?.limit ?? 10,
       messageRange: {
         ...DEFAULT_MESSAGE_RANGE,
-        ...searchArgs.searchOptions?.messageRange,
+        ...contextOptions.searchOptions?.messageRange,
       },
       text: extractText(messages.at(-1)!),
     };
     if (
-      searchArgs.searchOptions?.vectorSearch &&
+      contextOptions.searchOptions?.vectorSearch &&
       text &&
       this.options.textEmbedding
     ) {
@@ -552,71 +615,191 @@ export class Agent<AgentTools extends ToolSet> {
     return search;
   }
 
-  async getChatMessages(
-    ctx: RunQueryCtx,
-    args: {
-      chatId: string;
-      limit?: number;
-      statuses?: MessageStatus[];
-      cursor?: string;
-      includeToolMessages?: boolean;
-      order?: "asc" | "desc";
-    }
-  ): Promise<{
-    messages: (Message & { id: string })[];
-    continueCursor?: string;
-    isDone: boolean;
-  }> {
-    const messages = await ctx.runQuery(
-      this.component.messages.getChatMessages,
-      {
-        chatId: args.chatId,
-        limit: args.limit,
-        statuses: args.statuses,
-        cursor: args.cursor,
-        isTool: args.includeToolMessages,
-        order: args.order,
-      }
-    );
-    return {
-      messages: messages.messages
-        .map((m) => m && { ...m.message, id: m._id })
-        .filter((m): m is Message & { id: string } => m !== undefined),
-      continueCursor: messages.continueCursor,
-      isDone: messages.isDone,
-    };
+  /**
+   *
+   */
+  asAction(spec: { contextOptions?: ContextOptions; maxSteps?: number }) {
+    return internalActionGeneric({
+      args: {
+        userId: v.optional(v.string()),
+        threadId: v.optional(v.string()),
+        contextOptions: v.optional(vContextOptions),
+        storageOptions: v.optional(vStorageOptions),
+        maxRetries: v.optional(v.number()),
+
+        createThread: v.optional(
+          v.object({
+            userId: v.optional(v.string()),
+            parentThreadIds: v.optional(v.array(v.string())),
+            title: v.optional(v.string()),
+            summary: v.optional(v.string()),
+          })
+        ),
+        continueThread: v.optional(
+          v.object({
+            threadId: v.string(),
+            userId: v.optional(v.string()),
+          })
+        ),
+        generateText: v.optional(vThreadArgs),
+        streamText: v.optional(vThreadArgs),
+        generateObject: v.optional(vObjectArgs),
+        streamObject: v.optional(
+          v.object({ ...vObjectArgs.fields, schema: v.any() })
+        ),
+      },
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      handler: async (ctx, args): Promise<any> => {
+        const contextOptions =
+          spec.contextOptions && this.mergedContextOptions(spec.contextOptions);
+        const maxSteps = spec.maxSteps ?? this.options.maxSteps;
+        const maxRetries = args.maxRetries;
+        const commonArgs = {
+          userId: args.userId,
+          threadId: args.threadId,
+          ...contextOptions,
+          ...args.storageOptions,
+        };
+        if (args.createThread) {
+          return this.createThread(ctx, {
+            userId: args.createThread.userId,
+            parentThreadIds: args.createThread.parentThreadIds,
+            title: args.createThread.title,
+            summary: args.createThread.summary,
+          });
+        } else if (args.continueThread) {
+          return this.continueThread(ctx, {
+            threadId: args.continueThread.threadId,
+            userId: args.continueThread.userId,
+          });
+        } else if (args.generateText) {
+          return this.generateText(ctx, commonArgs, {
+            ...args.generateText,
+            maxSteps: args.generateText.maxSteps ?? maxSteps,
+            maxRetries,
+          });
+        } else if (args.streamText) {
+          return this.streamText(ctx, commonArgs, {
+            ...args.streamText,
+            maxSteps: args.streamText.maxSteps ?? maxSteps,
+            maxRetries,
+          });
+        } else if (args.generateObject) {
+          return this.generateObject(ctx, commonArgs, {
+            ...args.generateObject,
+            output: args.generateObject.output ?? "string",
+            maxRetries,
+          });
+        } else if (args.streamObject) {
+          return this.streamObject(ctx, commonArgs, {
+            ...args.streamObject,
+            output: args.streamObject.output ?? "string",
+            maxRetries,
+          });
+        }
+      },
+    });
   }
-}
 
-export function promptOrMessagesToCoreMessages(args: {
-  system?: string;
-  prompt?: string;
-  messages?: CoreMessage[] | Omit<UIMessage, "id">[];
-}): CoreMessage[] {
-  const messages: CoreMessage[] = [];
-  if (args.system) {
-    messages.push({ role: "system", content: args.system });
+  /**
+   * Create a tool that can call this agent.
+   * @param spec The specification for the arguments to this agent.
+   *   They will be encoded as JSON and passed to the agent.
+   * @returns The agent as a tool that can be passed to other agents.
+   */
+  asTool(spec: {
+    description: string;
+    args: Validator<unknown, "required", string>;
+    contextOptions?: ContextOptions;
+    maxSteps?: number;
+  }) {
+    return createTool({
+      ...spec,
+      handler: async (ctx, args) => {
+        const maxSteps = spec.maxSteps ?? this.options.maxSteps;
+        const contextOptions =
+          spec.contextOptions && this.mergedContextOptions(spec.contextOptions);
+        const value = await this.generateText(
+          ctx,
+          { userId: ctx.userId, threadId: ctx.threadId },
+          {
+            prompt: JSON.stringify(args),
+            parentMessageId: ctx.messageId,
+            maxSteps,
+            ...contextOptions,
+          }
+        );
+        return value.text;
+      },
+    });
   }
-  if (!args.messages) {
-    assert(args.prompt, "messages or prompt is required");
-    messages.push({ role: "user", content: args.prompt });
-  } else if (
-    args.messages.some(
-      (m) =>
-        typeof m === "object" &&
-        m !== null &&
-        (m.role === "data" || // UI-only role
-          "toolInvocations" in m || // UI-specific field
-          "parts" in m || // UI-specific field
-          "experimental_attachments" in m)
-    )
-  ) {
-    messages.push(...convertToCoreMessages(args.messages as UIMessage[]));
-  } else {
-    messages.push(...coreMessageSchema.array().parse(args.messages));
+}
+
+export type ToolCtx = RunActionCtx & {
+  userId?: string;
+  threadId?: string;
+  messageId?: string;
+};
+
+/**
+ * This is a wrapper around the ai.tool function that adds support for
+ * userId and threadId to the tool, if they're called within a thread from an agent.
+ * @param tool The AI tool. See https://sdk.vercel.ai/docs/ai-sdk-core/tools-and-tool-calling
+ * @returns The same tool, but with userId and threadId args support added.
+ */
+export function createTool<
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  V extends Validator<any, any, any>,
+  RESULT,
+>(convexTool: {
+  args: V;
+  description?: string;
+  handler: (
+    ctx: ToolCtx,
+    args: Infer<V>,
+    options: ToolExecutionOptions
+  ) => PromiseLike<RESULT>;
+  ctx?: ToolCtx;
+}): Tool<ConvexToZod<V>, RESULT> {
+  const tool = {
+    __acceptsCtx: true,
+    ctx: convexTool.ctx,
+    description: convexTool.description,
+    parameters: convexToZod(convexTool.args),
+    async execute(args: Infer<V>, 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 convexTool.handler(this.ctx, args, options);
+    },
+  };
+  return tool;
+}
+
+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;
+      }
+    }
   }
-  assert(messages.length > 0, "Messages must contain at least one message");
-  return messages;
+  return output;
 }
 
 type TextArgs<
@@ -629,6 +812,7 @@ type TextArgs<
   },
 > = Omit<T, "toolChoice" | "tools" | "model"> & {
   model?: LanguageModelV1;
+  parentMessageId?: string;
 } & {
   tools?: TOOLS;
   toolChoice?: ToolChoice<{ [key in keyof TOOLS | keyof AgentTools]: unknown }>;
@@ -644,7 +828,7 @@ type ObjectArgs<
 } & ContextOptions &
   StorageOptions;
 
-interface Chat<AgentTools extends ToolSet> {
+interface Thread<AgentTools extends ToolSet> {
   generateText<TOOLS extends ToolSet, OUTPUT = never, OUTPUT_PARTIAL = never>(
     args: TextArgs<
       AgentTools,
@@ -675,79 +859,3 @@ interface Chat<AgentTools extends ToolSet> {
     StreamObjectResult<DeepPartial<T>, T, never> & GenerationOutputMetadata
   >;
 }
-
-// type ToolParameters = ZodTypeAny | Schema<unknown>; // TODO: support convex validator
-// type inferParameters<PARAMETERS extends ToolParameters> =
-//   PARAMETERS extends Schema<unknown>
-//     ? PARAMETERS["_type"]
-//     : PARAMETERS extends z.ZodTypeAny
-//       ? z.infer<PARAMETERS>
-//       : never;
-/**
- * This is a wrapper around the ai.tool function that adds support for
- * userId and chatId to the tool, if they're called within a chat from an agent.
- * @param tool The AI tool. See https://sdk.vercel.ai/docs/ai-sdk-core/tools-and-tool-calling
- * @returns The same tool, but with userId and chatId args support added.
- */
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-export function tool<V extends Validator<any, any, any>, RESULT>(convexTool: {
-  args: V;
-  description?: string;
-  handler: (
-    ctx: GenericActionCtx<GenericDataModel> & {
-      userId?: string;
-      chatId?: string;
-    },
-    args: Infer<V>,
-    options: ToolExecutionOptions
-  ) => PromiseLike<RESULT>;
-  ctx?: GenericActionCtx<GenericDataModel> & {
-    userId?: string;
-    chatId?: string;
-  };
-}): Tool<ZodType<Infer<V>>, RESULT> {
-  const tool = {
-    __acceptUserIdAndChatId: true,
-    description: convexTool.description,
-    parameters: convexToZod(convexTool.args),
-    execute: async (args: Infer<V>, options: ToolExecutionOptions) => {
-      if (!convexTool.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 chatId)"
-        );
-      }
-      return convexTool.handler(convexTool.ctx, args, options);
-    },
-  };
-  return tool;
-}
-
-export function wrapTools(
-  actionCtx: RunActionCtx,
-  chatId: string,
-  userId?: string,
-  ...toolSets: (ToolSet | undefined)[]
-): ToolSet {
-  const ctx = { ...actionCtx, chatId, userId };
-  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).__acceptUserIdAndChatId) {
-        output[name] = tool;
-      } else {
-        const out = { ...tool, ctx };
-        output[name] = out;
-      }
-    }
-  }
-  return output;
-}
-// export function convexValidatorSchema<T>(validator: Validator<unknown>)  {
-//   return ai.jsonSchema(convexToJsonSchema(validator));
-// }