@convex-dev/agent 0.0.14 → 0.0.15-alpha.1

package/src/client/index.ts

@@ -27,7 +27,7 @@ import { assert } from "convex-helpers";
 import { internalActionGeneric, internalMutationGeneric } from "convex/server";
 import { Infer, v } from "convex/values";
 import { z } from "zod";
-import { Mounts } from "../component/_generated/api.js";
+import { api } from "../component/_generated/api.js";
 import {
   validateVectorDimension,
   type VectorDimension,
@@ -66,7 +66,7 @@ import type {
 } from "./types.js";
 import schema from "../component/schema.js";
 
-export { type Usage, type ProviderMetadata };
+export type { Usage, ProviderMetadata };
 export {
   vUsage,
   vProviderMetadata,
@@ -174,7 +174,8 @@ export type UsageHandler = (
 
 export class Agent<AgentTools extends ToolSet> {
   constructor(
-    public component: UseApi<Mounts>,
+    // public component: UseApi<Mounts>,
+    public component: UseApi<typeof api>,
     public options: {
       /**
        * The name for the agent. This will be attributed on each message
@@ -250,7 +251,7 @@ export class Agent<AgentTools extends ToolSet> {
    * @param args The thread metadata.
    * @returns The threadId of the new thread and the thread object.
    */
-  async createThread(
+  async createThread<ThreadTools extends ToolSet | undefined = undefined>(
     ctx: RunActionCtx,
     args?: {
       /**
@@ -271,10 +272,15 @@ 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.
+       */
+      tools?: ThreadTools;
     }
   ): Promise<{
     threadId: string;
-    thread: Thread<AgentTools>;
+    thread: Thread<ThreadTools extends undefined ? AgentTools : ThreadTools>;
   }>;
   /**
    * Start a new thread with the agent. This will have a fresh history, though if
@@ -285,28 +291,48 @@ export class Agent<AgentTools extends ToolSet> {
    * @param args The thread metadata.
    * @returns The threadId of the new thread.
    */
-  async createThread(
+  async createThread<ThreadTools extends ToolSet | undefined = undefined>(
     ctx: RunMutationCtx,
     args?: {
+      /**
+       * The userId to associate with the thread. If not provided, the thread will be
+       * anonymous.
+       */
       userId?: string;
+      /**
+       * The title of the thread. Not currently used.
+       */
       title?: string;
+      /**
+       * The summary of the thread. Not currently used.
+       */
       summary?: string;
+      /**
+       * The usage handler to use for this thread. Overrides any handler
+       * set in the agent constructor.
+       */
       usageHandler?: UsageHandler;
+      /**
+       * The tools to use for this thread.
+       * Overrides any tools passed in the agent constructor.
+       */
+      tools?: ThreadTools;
     }
   ): Promise<{
     threadId: string;
   }>;
-  async createThread(
+  async createThread<ThreadTools extends ToolSet | undefined = undefined>(
     ctx: RunActionCtx | RunMutationCtx,
     args?: {
       userId: string;
       title?: string;
       summary?: string;
       usageHandler?: UsageHandler;
+      tools?: ThreadTools;
     }
   ): Promise<{
     threadId: string;
-    thread?: Thread<AgentTools>;
+    thread?: Thread<ThreadTools extends undefined ? AgentTools : ThreadTools>;
   }> {
     const threadDoc = await ctx.runMutation(
       this.component.messages.createThread,
@@ -324,6 +350,7 @@ export class Agent<AgentTools extends ToolSet> {
       threadId: threadDoc._id,
       userId: args?.userId,
       usageHandler: args?.usageHandler,
+      tools: args?.tools,
     });
     return {
       threadId: threadDoc._id,
@@ -339,7 +366,7 @@ export class Agent<AgentTools extends ToolSet> {
    * @param { threadId, userId }: the thread and user to associate the messages with.
    * @returns Functions bound to the userId and threadId on a `{thread}` object.
    */
-  async continueThread(
+  async continueThread<ThreadTools extends ToolSet | undefined = undefined>(
     ctx: RunActionCtx,
     args: {
       /**
@@ -356,9 +383,14 @@ 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.
+       */
+      tools?: ThreadTools;
     }
   ): Promise<{
-    thread: Thread<AgentTools>;
+    thread: Thread<ThreadTools extends undefined ? AgentTools : ThreadTools>;
   }> {
     return {
       thread: {
@@ -367,7 +399,7 @@ export class Agent<AgentTools extends ToolSet> {
         streamText: this.streamText.bind(this, ctx, args),
         generateObject: this.generateObject.bind(this, ctx, args),
         streamObject: this.streamObject.bind(this, ctx, args),
-      } as Thread<AgentTools>,
+      } as Thread<ThreadTools extends undefined ? AgentTools : ThreadTools>,
     };
   }
 
@@ -382,17 +414,18 @@ export class Agent<AgentTools extends ToolSet> {
   async fetchContextMessages(
     ctx: RunQueryCtx | RunActionCtx,
     args: {
-      userId?: string;
-      threadId?: string;
+      userId: string | undefined;
+      threadId: string | undefined;
       messages: CoreMessage[];
       parentMessageId?: string;
-    } & ContextOptions
+      contextOptions: ContextOptions | undefined;
+    }
   ): Promise<CoreMessage[]> {
     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);
+    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");
@@ -400,7 +433,9 @@ export class Agent<AgentTools extends ToolSet> {
       const searchMessages = await ctx.runAction(
         this.component.messages.searchMessages,
         {
-          userId: args.searchOtherThreads ? args.userId : undefined,
+          userId: args.contextOptions?.searchOtherThreads
+            ? args.userId
+            : undefined,
           threadId: args.threadId,
           parentMessageId: args.parentMessageId,
           ...(await this.searchOptionsWithDefaults(opts, args.messages)),
@@ -434,7 +469,12 @@ export class Agent<AgentTools extends ToolSet> {
       .map((m) => deserializeMessage(m.message!));
   }
 
-  async getEmbeddings(messages: CoreMessage[]) {
+  /**
+   * Get the embeddings for a set of messages.
+   * @param messages The messages to get the embeddings for.
+   * @returns The embeddings for the messages.
+   */
+  async generateEmbeddings(messages: CoreMessage[]) {
     let embeddings:
       | {
           vectors: (number[] | null)[];
@@ -512,7 +552,7 @@ export class Agent<AgentTools extends ToolSet> {
     lastMessageId: string;
     messageIds: string[];
   }> {
-    const embeddings = await this.getEmbeddings(args.messages);
+    const embeddings = await this.generateEmbeddings(args.messages);
     const result = await ctx.runMutation(this.component.messages.addMessages, {
       threadId: args.threadId,
       userId: args.userId,
@@ -574,7 +614,9 @@ export class Agent<AgentTools extends ToolSet> {
       provider: args.provider ?? this.options.chat.provider,
       model: args.model ?? this.options.chat.modelId,
     });
-    const embeddings = await this.getEmbeddings(messages.map((m) => m.message));
+    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++) {
@@ -635,7 +677,7 @@ export class Agent<AgentTools extends ToolSet> {
    * @returns The result of the generateText function.
    */
   async generateText<
-    TOOLS extends ToolSet,
+    TOOLS extends ToolSet | undefined = undefined,
     OUTPUT = never,
     OUTPUT_PARTIAL = never,
   >(
@@ -644,6 +686,7 @@ export class Agent<AgentTools extends ToolSet> {
       userId,
       threadId,
       usageHandler,
+      tools: threadTools,
     }: {
       userId?: string;
       threadId?: string;
@@ -652,21 +695,25 @@ 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.
+       */
+      tools?: ToolSet;
     },
-    args: TextArgs<
-      AgentTools,
-      TOOLS,
-      Parameters<typeof generateText<TOOLS, OUTPUT, OUTPUT_PARTIAL>>[0]
-    >
+    args: TextArgs<AgentTools, TOOLS, OUTPUT, OUTPUT_PARTIAL>
   ): Promise<
-    GenerateTextResult<TOOLS & AgentTools, OUTPUT> & GenerationOutputMetadata
+    GenerateTextResult<TOOLS extends undefined ? AgentTools : TOOLS, OUTPUT> &
+      GenerationOutputMetadata
   > {
     const { args: aiArgs, messageId } = await this.saveMessagesAndFetchContext(
       ctx,
       { ...args, userId, threadId }
     );
     const toolCtx = { ...ctx, userId, threadId, messageId };
-    const tools = wrapTools(toolCtx, this.options.tools, args.tools) as TOOLS;
+    const tools = wrapTools(
+      toolCtx,
+      args.tools ?? threadTools ?? this.options.tools
+    ) as TOOLS extends undefined ? AgentTools : TOOLS;
     const saveOutputMessages =
       args.saveOutputMessages ??
       this.options.storageOptions?.saveOutputMessages;
@@ -679,8 +726,6 @@ export class Agent<AgentTools extends ToolSet> {
         maxRetries: this.options.maxRetries,
         ...aiArgs,
         model,
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        toolChoice: args.toolChoice as any,
         tools,
         onStepFinish: async (step) => {
           if (threadId && messageId && saveOutputMessages !== false) {
@@ -704,7 +749,11 @@ export class Agent<AgentTools extends ToolSet> {
           }
           return args.onStepFinish?.(step);
         },
-      })) as GenerateTextResult<TOOLS, OUTPUT> & GenerationOutputMetadata;
+      })) as GenerateTextResult<
+        TOOLS extends undefined ? AgentTools : TOOLS,
+        OUTPUT
+      > &
+        GenerationOutputMetadata;
       result.messageId = messageId;
       return result;
     } catch (error) {
@@ -732,7 +781,7 @@ export class Agent<AgentTools extends ToolSet> {
    * @returns The result of the streamText function.
    */
   async streamText<
-    TOOLS extends ToolSet,
+    TOOLS extends ToolSet | undefined = undefined,
     OUTPUT = never,
     PARTIAL_OUTPUT = never,
   >(
@@ -741,21 +790,34 @@ export class Agent<AgentTools extends ToolSet> {
       userId,
       threadId,
       usageHandler,
-    }: { userId?: string; threadId?: string; usageHandler?: UsageHandler },
-    args: TextArgs<
-      AgentTools,
-      TOOLS,
-      Parameters<typeof streamText<TOOLS, OUTPUT, PARTIAL_OUTPUT>>[0]
-    >
+      /**
+       * @deprecated Pass `tools` in the next parameter instead.
+       * This is only intended to pass through thread-default tools.
+       */
+      tools: threadTools,
+    }: {
+      userId?: string;
+      threadId?: string;
+      usageHandler?: UsageHandler;
+      tools?: ToolSet;
+    },
+    args: StreamingTextArgs<AgentTools, TOOLS, OUTPUT, PARTIAL_OUTPUT>
   ): Promise<
-    StreamTextResult<TOOLS, PARTIAL_OUTPUT> & GenerationOutputMetadata
+    StreamTextResult<
+      TOOLS extends undefined ? AgentTools : TOOLS,
+      PARTIAL_OUTPUT
+    > &
+      GenerationOutputMetadata
   > {
     const { args: aiArgs, messageId } = await this.saveMessagesAndFetchContext(
       ctx,
       { ...args, userId, threadId }
     );
     const toolCtx = { ...ctx, userId, threadId, messageId };
-    const tools = wrapTools(toolCtx, this.options.tools, args.tools) as TOOLS;
+    const tools = wrapTools(
+      toolCtx,
+      args.tools ?? threadTools ?? this.options.tools
+    ) as TOOLS extends undefined ? AgentTools : TOOLS;
     const saveOutputMessages =
       args.saveOutputMessages ??
       this.options.storageOptions?.saveOutputMessages;
@@ -767,8 +829,6 @@ export class Agent<AgentTools extends ToolSet> {
       maxRetries: this.options.maxRetries,
       ...aiArgs,
       model,
-      // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      toolChoice: args.toolChoice as any,
       tools,
       onChunk: async (chunk) => {
         // console.log("onChunk", chunk);
@@ -808,7 +868,11 @@ export class Agent<AgentTools extends ToolSet> {
         }
         return args.onStepFinish?.(step);
       },
-    }) as StreamTextResult<TOOLS, PARTIAL_OUTPUT> & GenerationOutputMetadata;
+    }) as StreamTextResult<
+      TOOLS extends undefined ? AgentTools : TOOLS,
+      PARTIAL_OUTPUT
+    > &
+      GenerationOutputMetadata;
     result.messageId = messageId;
     return result;
   }
@@ -833,30 +897,28 @@ export class Agent<AgentTools extends ToolSet> {
       userId: string | undefined;
       threadId: string | undefined;
       parentMessageId?: string;
-      saveAllInputMessages?: boolean;
-      saveAnyInputMessages?: boolean;
-    } & ContextOptions &
-      T
+      contextOptions?: ContextOptions;
+      storageOptions?: StorageOptions;
+    } & T
   ): Promise<{
     args: T;
     messageId: string | undefined;
   }> {
-    const saveAny =
-      args.saveAnyInputMessages ??
-      this.options.storageOptions?.saveAnyInputMessages;
-    const saveAll =
-      args.saveAllInputMessages ??
-      this.options.storageOptions?.saveAllInputMessages;
+    const contextOptions: ContextOptions | Record<string, unknown> =
+      args.contextOptions ?? this.options.contextOptions ?? args;
+    const storageOptions: StorageOptions | Record<string, unknown> =
+      args.storageOptions ?? this.options.storageOptions ?? args;
     const messages = promptOrMessagesToCoreMessages(args);
     const contextMessages = await this.fetchContextMessages(ctx, {
-      messages,
-      parentMessageId,
       userId,
       threadId,
-      ...args,
+      messages,
+      parentMessageId,
+      contextOptions,
     });
     let messageId: string | undefined;
-    if (threadId && saveAny !== false) {
+    if (threadId && storageOptions?.saveAnyInputMessages !== false) {
+      const saveAll = storageOptions?.saveAllInputMessages;
       const coreMessages = saveAll ? messages : messages.slice(-1);
       const saved = await this.saveMessages(ctx, {
         threadId,
@@ -1055,7 +1117,7 @@ export class Agent<AgentTools extends ToolSet> {
         provider: this.options.chat.provider,
       }
     );
-    const embeddings = await this.getEmbeddings([withoutEmbed[0].message]);
+    const embeddings = await this.generateEmbeddings([withoutEmbed[0].message]);
     const messages = embeddings?.vectors[0]
       ? [
           {
@@ -1078,10 +1140,10 @@ export class Agent<AgentTools extends ToolSet> {
     });
   }
 
-  mergedContextOptions(opts: ContextOptions): ContextOptions {
+  mergedContextOptions(opts: ContextOptions | undefined): ContextOptions {
     const searchOptions = {
       ...this.options.contextOptions?.searchOptions,
-      ...opts.searchOptions,
+      ...opts?.searchOptions,
     };
     return {
       ...this.options.contextOptions,
@@ -1313,33 +1375,137 @@ function wrapTools(
 
 type TextArgs<
   AgentTools extends ToolSet,
-  TOOLS extends ToolSet,
-  T extends {
-    toolChoice?: ToolChoice<TOOLS & AgentTools>;
-    tools?: TOOLS;
-    model: LanguageModelV1;
-  },
-> = Omit<T, "toolChoice" | "tools" | "model"> & {
+  TOOLS extends ToolSet | undefined = undefined,
+  OUTPUT = never,
+  OUTPUT_PARTIAL = never,
+> = Omit<
+  Parameters<
+    typeof generateText<
+      TOOLS extends undefined ? AgentTools : TOOLS,
+      OUTPUT,
+      OUTPUT_PARTIAL
+    >
+  >[0],
+  "toolChoice" | "tools" | "model"
+> & {
+  /**
+   * The model to use for the tool calls. This will override the model specified
+   * in the Agent constructor.
+   */
   model?: LanguageModelV1;
+  /**
+   * The tools to use for the tool calls. This will override tools specified
+   * in the Agent constructor or createThread / continueThread.
+   */
+  tools?: TOOLS;
+  /**
+   * The tool choice to use for the tool calls. This must be one of the tools
+   * specified in the tools array. e.g. {toolName: "getWeather", type: "tool"}
+   */
+  toolChoice?: ToolChoice<TOOLS extends undefined ? AgentTools : TOOLS>;
+  // Non-AI SDK args
+  /**
+   * The parent message id to use for the tool calls.
+   */
   parentMessageId?: string;
-} & {
+  /**
+   * The context options to use for passing in message history to the LLM.
+   */
+  contextOptions?: ContextOptions;
+  /**
+   * The storage options to use for saving the input and output messages to the thread.
+   */
+  storageOptions?: StorageOptions;
+} & ContextOptions &
+  StorageOptions;
+
+type StreamingTextArgs<
+  AgentTools extends ToolSet,
+  TOOLS extends ToolSet | undefined = undefined,
+  OUTPUT = never,
+  OUTPUT_PARTIAL = never,
+> = Omit<
+  Parameters<
+    typeof streamText<
+      TOOLS extends undefined ? AgentTools : TOOLS,
+      OUTPUT,
+      OUTPUT_PARTIAL
+    >
+  >[0],
+  "toolChoice" | "tools" | "model"
+> & {
+  /**
+   * The model to use for the tool calls. This will override the model specified
+   * in the Agent constructor.
+   */
+  model?: LanguageModelV1;
+  /**
+   * The tools to use for the tool calls. This will override tools specified
+   * in the Agent constructor or createThread / continueThread.
+   */
   tools?: TOOLS;
-  toolChoice?: ToolChoice<{ [key in keyof TOOLS | keyof AgentTools]: unknown }>;
+  /**
+   * The tool choice to use for the tool calls. This must be one of the tools
+   * specified in the tools array. e.g. {toolName: "getWeather", type: "tool"}
+   */
+  toolChoice?: ToolChoice<TOOLS extends undefined ? AgentTools : TOOLS>;
+  // Non-AI SDK args
+  /**
+   * The parent message id to use for the tool calls.
+   */
+  parentMessageId?: string;
+  /**
+   * The context options to use for passing in message history to the LLM.
+   */
+  contextOptions?: ContextOptions;
+  /**
+   * The storage options to use for saving the input and output messages to the thread.
+   */
+  storageOptions?: StorageOptions;
 } & 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;
-    parentMessageId?: string;
+    /**
+     * 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;
+    // Non-AI SDK args
+    /**
+     * The parent message id to use for the object generation.
+     */
+    parentMessageId?: string;
+    /**
+     * The context options to use for passing in message history to the LLM.
+     */
+    contextOptions?: ContextOptions;
+    /**
+     * The storage options to use for saving the input and output messages to the thread.
+     */
+    storageOptions?: StorageOptions;
   };
 
 type GenerateObjectObjectOptions<T extends Record<string, unknown>> =
@@ -1404,7 +1570,7 @@ type ThreadOutputMetadata = GenerationOutputMetadata & {
   messageId: string;
 };
 
-interface Thread<AgentTools extends ToolSet> {
+interface Thread<DefaultTools extends ToolSet> {
   /**
    * The target threadId, from the startThread or continueThread initializers.
    */
@@ -1419,14 +1585,20 @@ interface Thread<AgentTools extends ToolSet> {
    * for the {@link ContextOptions} and {@link StorageOptions}.
    * @returns The result of the generateText function.
    */
-  generateText<TOOLS extends ToolSet, OUTPUT = never, OUTPUT_PARTIAL = never>(
+  generateText<
+    TOOLS extends ToolSet | undefined = undefined,
+    OUTPUT = never,
+    OUTPUT_PARTIAL = never,
+  >(
     args: TextArgs<
-      AgentTools,
+      TOOLS extends undefined ? DefaultTools : TOOLS,
       TOOLS,
-      Parameters<typeof generateText<TOOLS, OUTPUT, OUTPUT_PARTIAL>>[0]
+      OUTPUT,
+      OUTPUT_PARTIAL
     >
   ): Promise<
-    GenerateTextResult<TOOLS & AgentTools, OUTPUT> & ThreadOutputMetadata
+    GenerateTextResult<TOOLS extends undefined ? DefaultTools : TOOLS, OUTPUT> &
+      ThreadOutputMetadata
   >;
 
   /**
@@ -1439,14 +1611,23 @@ interface Thread<AgentTools extends ToolSet> {
    * for the {@link ContextOptions} and {@link StorageOptions}.
    * @returns The result of the streamText function.
    */
-  streamText<TOOLS extends ToolSet, OUTPUT = never, PARTIAL_OUTPUT = never>(
-    args: TextArgs<
-      AgentTools,
+  streamText<
+    TOOLS extends ToolSet | undefined = undefined,
+    OUTPUT = never,
+    PARTIAL_OUTPUT = never,
+  >(
+    args: StreamingTextArgs<
+      TOOLS extends undefined ? DefaultTools : TOOLS,
       TOOLS,
-      Parameters<typeof streamText<TOOLS, OUTPUT, PARTIAL_OUTPUT>>[0]
+      OUTPUT,
+      PARTIAL_OUTPUT
     >
   ): Promise<
-    StreamTextResult<TOOLS & AgentTools, PARTIAL_OUTPUT> & ThreadOutputMetadata
+    StreamTextResult<
+      TOOLS extends undefined ? DefaultTools : TOOLS,
+      PARTIAL_OUTPUT
+    > &
+      ThreadOutputMetadata
   >;
   /**
    * This behaves like {@link generateObject} from the "ai" package except that