@convex-dev/agent 0.0.1-alpha.0 → 0.0.1-alpha.2

package/src/client/index.ts

@@ -1,60 +1,106 @@
-import { api } from "../component/_generated/api";
-import { RunActionCtx, RunMutationCtx, RunQueryCtx, UseApi } from "./types";
 import type { EmbeddingModelV1, LanguageModelV1 } from "@ai-sdk/provider";
-import { Message, MessageStatus, SearchOptions, Step } from "../validators";
 import type {
-  StreamTextResult,
-  Tool,
-  ToolSet,
-  StepResult,
-  Message as UIMessage,
   CoreMessage,
-  GenerateObjectResult,
-  StreamObjectResult,
   DeepPartial,
+  GenerateObjectResult,
   GenerateTextResult,
+  StepResult,
+  StreamObjectResult,
+  StreamTextResult,
+  Tool,
+  ToolChoice,
+  ToolExecutionOptions,
+  ToolSet,
 } from "ai";
+import { generateObject, generateText, streamObject, streamText } from "ai";
+import { api } from "../component/_generated/api";
 import {
-  generateText,
-  generateObject,
-  streamText,
-  streamObject,
-  convertToCoreMessages,
-  coreMessageSchema,
-} from "ai";
+  SearchOptions,
+  vChatArgs,
+  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 { DEFAULT_MESSAGE_RANGE, extractText } from "../shared";
+import { ConvexToZod, convexToZod } from "convex-helpers/server/zod";
+import { internalActionGeneric } from "convex/server";
+import { Infer, v, Validator } from "convex/values";
 import {
-  serializeMessage,
+  promptOrMessagesToCoreMessages,
   serializeMessageWithId,
   serializeNewMessagesInStep,
   serializeStep,
 } from "../mapping";
+import { DEFAULT_MESSAGE_RANGE, extractText } from "../shared";
 
 export type ContextOptions = {
-  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 };
   };
+  /**
+   * Whether to search across other chats for relevant messages.
+   * By default, only the current chat is searched.
+   */
   searchOtherChats?: 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.
+  saveAllInputMessages?: boolean;
+  // Defaults to true
+  saveOutputMessages?: boolean;
+};
+
+export type GenerationOutputMetadata = { messageId: string };
+
 type CoreMessageMaybeWithId = CoreMessage & { id?: string | undefined };
 
-export class Agent {
+export class Agent<AgentTools extends ToolSet> {
   constructor(
     public component: UseApi<typeof api>,
     public options: {
       name?: string;
       chat: LanguageModelV1;
       textEmbedding?: EmbeddingModelV1<string>;
-      defaultSystemPrompt?: string;
-      tools?: Record<string, Tool>;
+      instructions?: string;
+      tools?: AgentTools;
+      contextOptions?: ContextOptions;
+      // TODO: storageOptions?: StorageOptions;
+      maxSteps?: number;
+      // TODO: maxRetries?: number;
     }
   ) {}
 
@@ -68,7 +114,7 @@ export class Agent {
    * @param args The chat metadata.
    * @returns The chatId of the new chat and the chat object.
    */
-  async startChat(
+  async createChat(
     ctx: RunActionCtx,
     args: {
       /**
@@ -93,7 +139,7 @@ export class Agent {
     }
   ): Promise<{
     chatId: string;
-    chat: Chat;
+    chat: Chat<AgentTools>;
   }>;
   /**
    * Start a new chat with the agent. This will have a fresh history, though if
@@ -104,7 +150,7 @@ export class Agent {
    * @param args The chat metadata.
    * @returns The chatId of the new chat.
    */
-  async startChat(
+  async createChat(
     ctx: RunMutationCtx,
     args: {
       userId?: string;
@@ -115,7 +161,7 @@ export class Agent {
   ): Promise<{
     chatId: string;
   }>;
-  async startChat(
+  async createChat(
     ctx: RunActionCtx | RunMutationCtx,
     args: {
       userId: string;
@@ -125,10 +171,10 @@ export class Agent {
     }
   ): Promise<{
     chatId: string;
-    chat?: Chat;
+    chat?: Chat<AgentTools>;
   }> {
     const chatDoc = await ctx.runMutation(this.component.messages.createChat, {
-      defaultSystemPrompt: this.options.defaultSystemPrompt,
+      defaultSystemPrompt: this.options.instructions,
       userId: args.userId,
       title: args.title,
       summary: args.summary,
@@ -154,10 +200,14 @@ export class Agent {
       userId,
     }: {
       chatId: string;
+      /**
+       * If supplied, the userId can be used to search across other chats for
+       * relevant messages from the same user as context for the LLM calls.
+       */
       userId?: string;
     }
   ): Promise<{
-    chat: Chat;
+    chat: Chat<AgentTools>;
   }> {
     // return this.component.continueChat(ctx, args);
     return {
@@ -166,7 +216,7 @@ export class Agent {
         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,
+      } as Chat<AgentTools>,
     };
   }
 
@@ -176,12 +226,14 @@ export class Agent {
       userId?: string;
       chatId?: string;
       messages: CoreMessage[];
+      parentMessageId?: string;
     } & ContextOptions
   ): Promise<CoreMessage[]> {
     assert(args.userId || args.chatId, "Specify userId or chatId");
     // Fetch the latest messages from the chat
     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");
       }
@@ -190,9 +242,11 @@ export class Agent {
         {
           userId: args.searchOtherChats ? args.userId : undefined,
           chatId: args.chatId,
-          ...(await this.searchWithDefaults(args, args.messages)),
+          parentMessageId: args.parentMessageId,
+          ...(await this.searchOptionsWithDefaults(opts, args.messages)),
         }
       );
+      // TODO: track what messages we used for context
       contextMessages.push(...searchMessages.map((m) => m.message!));
     }
     if (args.chatId) {
@@ -200,8 +254,9 @@ export class Agent {
         this.component.messages.getChatMessages,
         {
           chatId: args.chatId,
-          isTool: args.includeToolMessages ?? false,
+          isTool: args.includeToolCalls ?? false,
           limit: args.recentMessages,
+          parentMessageId: args.parentMessageId,
           order: "desc",
           statuses: ["success"],
         }
@@ -214,45 +269,33 @@ export class Agent {
   async saveMessages(
     ctx: RunMutationCtx,
     args: {
-      chatId: string;
+      chatId?: string;
+      userId?: string;
       messages: CoreMessageMaybeWithId[];
-      steps?: StepResult<ToolSet>[];
       pending?: boolean;
+      parentMessageId?: string;
+      failPendingSteps?: boolean;
     }
   ): Promise<{
     lastMessageId: string;
-  }>;
-  async saveMessages(
-    ctx: RunMutationCtx,
-    args: {
-      chatId: string;
-      messages: CoreMessageMaybeWithId[];
-      steps?: StepResult<ToolSet>[];
-      pending?: boolean;
-    }
-  ): Promise<{
-    lastMessageId?: string;
+    messageIds: string[];
   }> {
     const result = await ctx.runMutation(this.component.messages.addMessages, {
       chatId: args.chatId,
+      userId: args.userId,
       agentName: this.options.name,
       model: this.options.chat.modelId,
       messages: args.messages.map(serializeMessageWithId),
-      failPendingSteps: true,
+      failPendingSteps: args.failPendingSteps ?? true,
       pending: args.pending ?? false,
+      parentMessageId: args.parentMessageId,
     });
     return {
-      lastMessageId: result.messages.at(-1)?._id,
+      lastMessageId: result.messages.at(-1)!._id,
+      messageIds: result.messages.map((m) => m._id),
     };
   }
 
-  async replaceMessages(
-    ctx: RunMutationCtx,
-    args: { chatId: string; messages: Message[] }
-  ): Promise<void> {
-    throw new Error("Not implemented");
-  }
-
   async saveStep<TOOLS extends ToolSet>(
     ctx: RunMutationCtx,
     args: { chatId: string; messageId: string; step: StepResult<TOOLS> }
@@ -267,6 +310,7 @@ export class Agent {
     });
   }
 
+  // If you manually create a message, call this to either commit or reset it.
   async completeMessage<TOOLS extends ToolSet>(
     ctx: RunMutationCtx,
     args: {
@@ -274,10 +318,7 @@ export class Agent {
       messageId: string;
       result:
         | { kind: "error"; error: string }
-        | {
-            kind: "success";
-            value: { steps: StepResult<TOOLS>[] };
-          };
+        | { kind: "success"; value: { steps: StepResult<TOOLS>[] } };
     }
   ): Promise<void> {
     const result = args.result;
@@ -317,11 +358,14 @@ export class Agent {
       userId?: string;
       chatId?: string;
     },
-    args: Partial<
+    args: TextArgs<
+      AgentTools,
+      TOOLS,
       Parameters<typeof generateText<TOOLS, OUTPUT, OUTPUT_PARTIAL>>[0]
-    > &
-      ContextOptions
-  ): Promise<GenerateTextResult<TOOLS, OUTPUT>> {
+    >
+  ): Promise<
+    GenerateTextResult<TOOLS & AgentTools, OUTPUT> & GenerationOutputMetadata
+  > {
     const { prompt, messages: raw, ...rest } = args;
     const messages = promptOrMessagesToCoreMessages({ prompt, messages: raw });
     const contextMessages = await this.fetchContextMessages(ctx, {
@@ -330,24 +374,28 @@ export class Agent {
       chatId,
       messages,
     });
-    let messageId: string | undefined;
-    if (chatId) {
-      const { lastMessageId } = await this.saveMessages(ctx, {
-        chatId,
-        // TODO: only save the last message unless explicitly told to save all
-        messages,
-        pending: true,
-      });
-      messageId = lastMessageId;
-    }
+    const { lastMessageId: messageId } = await this.saveMessages(ctx, {
+      chatId,
+      userId,
+      messages: args.saveAllInputMessages ? messages : messages.slice(-1),
+      pending: true,
+      parentMessageId: args.parentMessageId,
+    });
+    const toolCtx = { ...ctx, userId, chatId, messageId };
+    const tools = wrapTools(toolCtx, this.options.tools, args.tools) as TOOLS;
     try {
       const result = await generateText({
         model: this.options.chat,
         messages: [...contextMessages, ...messages],
-        system: this.options.defaultSystemPrompt,
+        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) {
+          if (chatId && messageId && args.saveOutputMessages) {
+            console.log("onStepFinish", step);
             await this.saveStep(ctx, {
               chatId,
               messageId,
@@ -357,9 +405,10 @@ export class Agent {
           return args.onStepFinish?.(step);
         },
       });
-      return result;
+      return { ...result, messageId };
     } catch (error) {
       if (chatId && messageId) {
+        console.error("RollbackMessage", messageId);
         await ctx.runMutation(this.component.messages.rollbackMessage, {
           messageId,
           error: (error as Error).message,
@@ -374,13 +423,16 @@ export class Agent {
     OUTPUT = never,
     PARTIAL_OUTPUT = never,
   >(
-    ctx: RunMutationCtx,
+    ctx: RunActionCtx,
     { userId, chatId }: { userId?: string; chatId?: string },
-    args: Partial<
+    args: TextArgs<
+      AgentTools,
+      TOOLS,
       Parameters<typeof streamText<TOOLS, OUTPUT, PARTIAL_OUTPUT>>[0]
-    > &
-      ContextOptions
-  ): Promise<StreamTextResult<TOOLS, PARTIAL_OUTPUT>> {
+    >
+  ): Promise<
+    StreamTextResult<TOOLS, PARTIAL_OUTPUT> & GenerationOutputMetadata
+  > {
     const { prompt, messages: raw, ...rest } = args;
     const messages = promptOrMessagesToCoreMessages({ prompt, messages: raw });
     const contextMessages = await this.fetchContextMessages(ctx, {
@@ -389,21 +441,23 @@ export class Agent {
       chatId,
       messages,
     });
-    let messageId: string | undefined;
-    if (chatId) {
-      const { lastMessageId } = await this.saveMessages(ctx, {
-        chatId,
-        // TODO: only save the last message unless explicitly told to save all
-        messages,
-        pending: true,
-      });
-      messageId = lastMessageId;
-    }
-    return streamText({
+    const { lastMessageId: messageId } = await this.saveMessages(ctx, {
+      chatId,
+      messages: args.saveAllInputMessages ? messages : messages.slice(-1),
+      pending: true,
+      parentMessageId: args.parentMessageId,
+    });
+    const toolCtx = { ...ctx, userId, chatId, messageId };
+    const tools = wrapTools(toolCtx, this.options.tools, args.tools) as TOOLS;
+    const result = streamText({
       model: this.options.chat,
       messages: [...contextMessages, ...messages],
-      system: this.options.defaultSystemPrompt,
+      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);
@@ -436,16 +490,18 @@ export class Agent {
         return args.onStepFinish?.(step);
       },
     });
+    return { ...result, messageId };
   }
 
-  // TODO: not sure why it needs to extend string
-  async generateObject<OBJECT extends string>(
+  // TODO: add the crazy number of overloads to get types through
+  async generateObject<T>(
     ctx: RunActionCtx,
     { userId, chatId }: { userId?: string; chatId?: string },
-    args: Omit<Parameters<typeof generateObject<OBJECT>>[0], "model"> & {
+    args: Omit<Parameters<typeof generateObject>[0], "model"> & {
       model?: LanguageModelV1;
-    } & ContextOptions
-  ): Promise<GenerateObjectResult<OBJECT>> {
+    } & { parentMessageId?: string } & ContextOptions &
+      StorageOptions
+  ): Promise<GenerateObjectResult<T> & GenerationOutputMetadata> {
     const { prompt, messages: raw, ...rest } = args;
     const messages = promptOrMessagesToCoreMessages({ prompt, messages: raw });
     const contextMessages = await this.fetchContextMessages(ctx, {
@@ -454,11 +510,17 @@ export class Agent {
       chatId,
       messages,
     });
-    return generateObject({
+    const { lastMessageId: messageId } = await this.saveMessages(ctx, {
+      chatId,
+      messages: args.saveAllInputMessages ? messages : messages.slice(-1),
+      pending: true,
+    });
+    const result = (await generateObject({
       model: this.options.chat,
       messages: [...contextMessages, ...messages],
       ...rest,
-    }) as Promise<GenerateObjectResult<OBJECT>>;
+    })) as GenerateObjectResult<T>;
+    return { ...result, messageId };
   }
 
   async streamObject<T>(
@@ -466,8 +528,11 @@ export class Agent {
     { userId, chatId }: { userId?: string; chatId?: string },
     args: Omit<Parameters<typeof streamObject<T>>[0], "model"> & {
       model?: LanguageModelV1;
-    } & ContextOptions
-  ) {
+    } & { parentMessageId?: string } & ContextOptions &
+      StorageOptions
+  ): Promise<
+    StreamObjectResult<DeepPartial<T>, T, never> & GenerationOutputMetadata
+  > {
     const { prompt, messages: raw, ...rest } = args;
     const messages = promptOrMessagesToCoreMessages({ prompt, messages: raw });
     const contextMessages = await this.fetchContextMessages(ctx, {
@@ -476,7 +541,12 @@ export class Agent {
       chatId,
       messages,
     });
-    return streamObject<T>({
+    const { lastMessageId: messageId } = await this.saveMessages(ctx, {
+      chatId,
+      messages: args.saveAllInputMessages ? messages : messages.slice(-1),
+      pending: true,
+    });
+    const result = streamObject<T>({
       model: this.options.chat,
       messages: [...contextMessages, ...messages],
       ...rest,
@@ -488,29 +558,44 @@ export class Agent {
         console.log("onFinish", result);
       },
     }) as StreamObjectResult<DeepPartial<T>, T, never>;
+    return { ...result, messageId };
   }
 
-  async searchWithDefaults(
-    searchArgs: ContextOptions,
+  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(
+    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
     ) {
@@ -524,136 +609,247 @@ export class Agent {
     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()),
+        chatId: v.optional(v.string()),
+        contextOptions: v.optional(vContextOptions),
+        storageOptions: v.optional(vStorageOptions),
+        maxRetries: v.optional(v.number()),
+
+        createChat: v.optional(
+          v.object({
+            userId: v.string(),
+            parentChatIds: v.optional(v.array(v.string())),
+            title: v.optional(v.string()),
+            summary: v.optional(v.string()),
+          })
+        ),
+        continueChat: v.optional(
+          v.object({
+            chatId: v.string(),
+            userId: v.optional(v.string()),
+          })
+        ),
+        generateText: v.optional(vChatArgs),
+        streamText: v.optional(vChatArgs),
+        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,
+          chatId: args.chatId,
+          ...contextOptions,
+          ...args.storageOptions,
+        };
+        if (args.createChat) {
+          return this.createChat(ctx, {
+            userId: args.createChat.userId,
+            parentChatIds: args.createChat.parentChatIds,
+            title: args.createChat.title,
+            summary: args.createChat.summary,
+          });
+        } else if (args.continueChat) {
+          return this.continueChat(ctx, {
+            chatId: args.continueChat.chatId,
+            userId: args.continueChat.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,
+          });
+        }
+      },
+    });
   }
 
-  async getSteps(
-    ctx: RunQueryCtx,
-    args: { messageId: string }
-  ): Promise<{
-    steps: Step[];
-  }> {
-    return { steps: [] };
+  /**
+   * 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, chatId: ctx.chatId },
+          {
+            prompt: JSON.stringify(args),
+            parentMessageId: ctx.messageId,
+            maxSteps,
+            ...contextOptions,
+          }
+        );
+        return value.text;
+      },
+    });
+  }
+}
+
+export type ToolCtx = RunActionCtx & {
+  userId?: string;
+  chatId?: string;
+  messageId?: string;
+};
+
+/**
+ * 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.
+ */
+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 chatId)"
+        );
+      }
+      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;
+      }
+    }
   }
+  return output;
 }
 
-interface Chat {
+type TextArgs<
+  AgentTools extends ToolSet,
+  TOOLS extends ToolSet,
+  T extends {
+    toolChoice?: ToolChoice<TOOLS & AgentTools>;
+    tools?: TOOLS;
+    model: LanguageModelV1;
+  },
+> = Omit<T, "toolChoice" | "tools" | "model"> & {
+  model?: LanguageModelV1;
+  parentMessageId?: string;
+} & {
+  tools?: TOOLS;
+  toolChoice?: ToolChoice<{ [key in keyof TOOLS | keyof AgentTools]: unknown }>;
+} & ContextOptions &
+  StorageOptions;
+
+type ObjectArgs<
+  T extends {
+    model: LanguageModelV1;
+  },
+> = Omit<T, "model"> & {
+  model?: LanguageModelV1;
+} & ContextOptions &
+  StorageOptions;
+
+interface Chat<AgentTools extends ToolSet> {
   generateText<TOOLS extends ToolSet, OUTPUT = never, OUTPUT_PARTIAL = never>(
-    args: Partial<
+    args: TextArgs<
+      AgentTools,
+      TOOLS,
       Parameters<typeof generateText<TOOLS, OUTPUT, OUTPUT_PARTIAL>>[0]
     >
-  ): Promise<GenerateTextResult<TOOLS, OUTPUT>>;
+  ): Promise<
+    GenerateTextResult<TOOLS & AgentTools, OUTPUT> & GenerationOutputMetadata
+  >;
+
   streamText<TOOLS extends ToolSet, OUTPUT = never, PARTIAL_OUTPUT = never>(
-    args: Partial<
+    args: TextArgs<
+      AgentTools,
+      TOOLS,
       Parameters<typeof streamText<TOOLS, OUTPUT, PARTIAL_OUTPUT>>[0]
     >
-  ): Promise<StreamTextResult<TOOLS, PARTIAL_OUTPUT>>;
-  generateObject<OBJECT extends string>(
-    args: Omit<Parameters<typeof generateObject<OBJECT>>[0], "model"> & {
-      model?: LanguageModelV1;
-    }
-  ): Promise<GenerateObjectResult<OBJECT>>;
+  ): Promise<
+    StreamTextResult<TOOLS & AgentTools, PARTIAL_OUTPUT> &
+      GenerationOutputMetadata
+  >;
+  // TODO: add all the overloads
+  generateObject<T>(
+    args: ObjectArgs<Parameters<typeof generateObject>[0]>
+  ): Promise<GenerateObjectResult<T> & GenerationOutputMetadata>;
   streamObject<T>(
-    args: Omit<Parameters<typeof streamObject<T>>[0], "model"> & {
-      model?: LanguageModelV1;
-    }
-  ): Promise<StreamObjectResult<DeepPartial<T>, T, never>>;
-}
-
-// 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.
-//  */
-// export function tool<PARAMETERS extends ToolParameters, RESULT>(
-//   tool: Tool<PARAMETERS, RESULT> & {
-//     execute: (
-//       args: inferParameters<PARAMETERS> & { userId?: string; chatId?: string },
-//       options: ToolExecutionOptions
-//     ) => PromiseLike<RESULT>;
-//   }
-// ): Tool<PARAMETERS, RESULT> & {
-//   execute: (
-//     args: inferParameters<PARAMETERS>,
-//     options: ToolExecutionOptions
-//   ) => PromiseLike<RESULT>;
-// } {
-//   // eslint-disable-next-line @typescript-eslint/no-explicit-any
-//   (tool as any).__acceptUserIdAndChatId = true;
-//   return tool;
-// }
-
-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 });
-  }
-  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));
-  }
-  assert(messages.length > 0, "Messages must contain at least one message");
-  return messages;
+    args: ObjectArgs<Parameters<typeof streamObject<T>>[0]>
+  ): Promise<
+    StreamObjectResult<DeepPartial<T>, T, never> & GenerationOutputMetadata
+  >;
 }
-
-// export function convexValidatorSchema<T>(validator: Validator<unknown>)  {
-//   return ai.jsonSchema(convexToJsonSchema(validator));
-// }