@convex-dev/agent 0.0.8 → 0.0.9-alpha.1

package/src/client/index.ts

@@ -38,6 +38,7 @@ import {
 } from "../component/vector/tables";
 import {
   AIMessageWithoutId,
+  deserializeMessage,
   promptOrMessagesToCoreMessages,
   serializeMessageWithId,
   serializeNewMessagesInStep,
@@ -64,6 +65,10 @@ import { RunActionCtx, RunMutationCtx, RunQueryCtx, UseApi } from "./types.js";
 export { convexToZod, zodToConvex };
 export type { ThreadDoc, MessageDoc } from "./types.js";
 
+/**
+ * Options to configure what messages are fetched as context,
+ * automatically with thread.generateText, or directly via search.
+ */
 export type ContextOptions = {
   /**
    * Whether to include tool messages in the context.
@@ -104,14 +109,20 @@ export type ContextOptions = {
   searchOtherThreads?: boolean;
 };
 
+/**
+ * Options to configure the automatic saving of messages
+ * when generating text / objects in a thread.
+ */
 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 thread history.
+  /**
+   * 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 thread history.
+   */
   saveAllInputMessages?: boolean;
-  // Defaults to true, saving the prompt, or last message passed to generateText.
+  /** Defaults to true, saving the prompt, or last message passed to generateText. */
   saveAnyInputMessages?: boolean;
-  // Defaults to true.
+  /** Defaults to true. Whether to save messages generated while chatting. */
   saveOutputMessages?: boolean;
 };
 
@@ -123,15 +134,65 @@ export class Agent<AgentTools extends ToolSet> {
   constructor(
     public component: UseApi<Mounts>,
     public options: {
+      /**
+       * The name for the agent. This will be attributed on each message
+       * created by this agent.
+       */
       name?: string;
+      /**
+       * The LLM model to use for generating / streaming text and objects.
+       * e.g.
+       * import { openai } from "@ai-sdk/openai"
+       * const myAgent = new Agent(components.agent, {
+       *   chat: openai.chat("gpt-4o-mini"),
+       */
       chat: LanguageModelV1;
+      /**
+       * The model to use for text embeddings. Optional.
+       * If specified, it will use this for generating vector embeddings
+       * of chats, and can opt-in to doing vector search for automatic context
+       * on generateText, etc.
+       * e.g.
+       * import { openai } from "@ai-sdk/openai"
+       * const myAgent = new Agent(components.agent, {
+       *   textEmbedding: openai.embedding("text-embedding-3-small")
+       */
       textEmbedding?: EmbeddingModelV1<string>;
+      /**
+       * The default system prompt to put in each request.
+       * Override per-prompt by passing the "system" parameter.
+       */
       instructions?: string;
+      /**
+       * Tools that the agent can call out to and get responses from.
+       * They can be AI SDK tools (import {tool} from "ai")
+       * or tools that have Convex context
+       * (import { createTool } from "@convex-dev/agent")
+       * Note: Convex tools can't currently annotate the parameters
+       * with descriptions, so the names should be self-evident from naming.
+       */
       tools?: AgentTools;
+      /**
+       * Options to determine what messages are included as context in message
+       * generation. To disable any messages automatically being added, pass:
+       * { recentMessages: 0 }
+       */
       contextOptions?: ContextOptions;
-      // TODO: storageOptions?: StorageOptions;
+      /**
+       * Determines whether messages are automatically stored when passed as
+       * arguments or generated.
+       */
+      storageOptions?: StorageOptions;
+      /**
+       * When generating or streaming text with tools available, this
+       * determines the default max number of iterations.
+       */
       maxSteps?: number;
-      // TODO: maxRetries?: number;
+      /**
+       * The maximum number of calls to make to an LLM in case it fails.
+       * This can be overridden at each generate/stream callsite.
+       */
+      maxRetries?: number;
     }
   ) {}
 
@@ -227,12 +288,34 @@ export class Agent<AgentTools extends ToolSet> {
     };
   }
 
+  /**
+   * Continues a thread using this agent. Note: threads can be continued
+   * by different agents. This is a convenience around calling the various
+   * generate and stream functions with explicit userId and threadId parameters.
+   * @param ctx The ctx object passed to the action handler
+   * @param { threadId, userId }: the thread and user to associate the messages with.
+   * @returns Functions bound to the userId and threadId on a `{thread}` object.
+   */
+  /**
+   * Continues a thread using this agent. Note: threads can be continued
+   * by different agents. This is a convenience around calling the various
+   * generate and stream functions with explicit userId and threadId parameters.
+   * @param ctx The ctx object passed to the action handler
+   * @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(
     ctx: RunActionCtx,
     {
       threadId,
       userId,
     }: {
+      /**
+       * The associated thread created by {@link createThread}
+       */
+      /**
+       * The associated thread created by {@link createThread}
+       */
       threadId: string;
       /**
        * If supplied, the userId can be used to search across other threads for
@@ -243,7 +326,6 @@ export class Agent<AgentTools extends ToolSet> {
   ): Promise<{
     thread: Thread<AgentTools>;
   }> {
-    // return this.component.continueThread(ctx, args);
     return {
       thread: {
         threadId,
@@ -258,6 +340,22 @@ export class Agent<AgentTools extends ToolSet> {
     };
   }
 
+  /**
+   *
+   * @param ctx Either a query, mutation, or action ctx.
+   *   If it is not an action context, you can't do text or
+   *   vector search.
+   * @param args The associated thread, user, message
+   * @returns
+   */
+  /**
+   *
+   * @param ctx Either a query, mutation, or action ctx.
+   *   If it is not an action context, you can't do text or
+   *   vector search.
+   * @param args The associated thread, user, message
+   * @returns
+   */
   async fetchContextMessages(
     ctx: RunQueryCtx | RunActionCtx,
     args: {
@@ -287,7 +385,9 @@ export class Agent<AgentTools extends ToolSet> {
       );
       // TODO: track what messages we used for context
       included = new Set(searchMessages.map((m) => m._id));
-      contextMessages.push(...searchMessages.map((m) => m.message!));
+      contextMessages.push(
+        ...searchMessages.map((m) => deserializeMessage(m.message!))
+      );
     }
     if (args.threadId && opts.recentMessages !== 0) {
       const { page } = await ctx.runQuery(
@@ -305,7 +405,9 @@ export class Agent<AgentTools extends ToolSet> {
         }
       );
       contextMessages.push(
-        ...page.filter((m) => !included?.has(m._id)).map((m) => m.message!)
+        ...page
+          .filter((m) => !included?.has(m._id))
+          .map((m) => deserializeMessage(m.message!))
       );
     }
     return contextMessages;
@@ -350,14 +452,48 @@ export class Agent<AgentTools extends ToolSet> {
     return embeddings;
   }
 
+  /**
+   * Explicitly save messages associated with the thread (& user if provided)
+   * @param ctx The ctx parameter to a mutation or action.
+   * @param args The messages and context to save
+   * @returns
+   */
+  /**
+   * Explicitly save messages associated with the thread (& user if provided)
+   * @param ctx The ctx parameter to a mutation or action.
+   * @param args The messages and context to save
+   * @returns
+   */
   async saveMessages(
     ctx: RunMutationCtx,
     args: {
       threadId: string;
       userId?: string;
       messages: CoreMessageMaybeWithId[];
+      /**
+       * If false, it will "commit" the messages immediately.
+       * If true, it will mark them as pending until the final step has finished.
+       */
+      /**
+       * If false, it will "commit" the messages immediately.
+       * If true, it will mark them as pending until the final step has finished.
+       */
       pending?: boolean;
+      /**
+       * The message that this is responding to.
+       */
+      /**
+       * The message that this is responding to.
+       */
       parentMessageId?: string;
+      /**
+       * Whether to mark all pending messages in the thread as failed.
+       * This is used to recover from a failure via a retry that wipes the slate clean.
+       */
+      /**
+       * Whether to mark all pending messages in the thread as failed.
+       * This is used to recover from a failure via a retry that wipes the slate clean.
+       */
       failPendingSteps?: boolean;
     }
   ): Promise<{
@@ -381,9 +517,29 @@ export class Agent<AgentTools extends ToolSet> {
     };
   }
 
+  /**
+   * Explicitly save a "step" created by the AI SDK.
+   * @param ctx The ctx argument to a mutation or action.
+   * @param args What to save
+   */
+  /**
+   * Explicitly save a "step" created by the AI SDK.
+   * @param ctx The ctx argument to a mutation or action.
+   * @param args What to save
+   */
   async saveStep<TOOLS extends ToolSet>(
     ctx: RunMutationCtx,
-    args: { threadId: string; messageId: string; step: StepResult<TOOLS> }
+    args: {
+      threadId: string;
+      /**
+       * The message this step is in response to.
+       */
+      messageId: string;
+      /**
+       * The step to save, possibly including multiple tool calls.
+       */
+      step: StepResult<TOOLS>;
+    }
   ): Promise<void> {
     const step = serializeStep(args.step as StepResult<ToolSet>);
     const messages = serializeNewMessagesInStep(args.step);
@@ -396,15 +552,28 @@ export class Agent<AgentTools extends ToolSet> {
     });
   }
 
-  // If you manually create a message, call this to either commit or reset it.
-  async completeMessage<TOOLS extends ToolSet>(
+  /**
+   * Commit or rollback a message that was pending.
+   * This is done automatically when saving messages by default.
+   * If creating pending messages, you can call this when the full "transaction" is done.
+   * @param ctx The ctx argument to your mutation or action.
+   * @param args What message to save. Generally the parent message sent into
+   *   the generateText call.
+   */
+  /**
+   * Commit or rollback a message that was pending.
+   * This is done automatically when saving messages by default.
+   * If creating pending messages, you can call this when the full "transaction" is done.
+   * @param ctx The ctx argument to your mutation or action.
+   * @param args What message to save. Generally the parent message sent into
+   *   the generateText call.
+   */
+  async completeMessage(
     ctx: RunMutationCtx,
     args: {
       threadId: string;
       messageId: string;
-      result:
-        | { kind: "error"; error: string }
-        | { kind: "success"; value: { steps: StepResult<TOOLS>[] } };
+      result: { kind: "error"; error: string } | { kind: "success" };
     }
   ): Promise<void> {
     const result = args.result;
@@ -421,12 +590,15 @@ export class Agent<AgentTools extends ToolSet> {
   }
 
   /**
-   * This behaves like {@link generateText} except that it add context based on
-   * 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.
+   * This behaves like {@link generateText} 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 generateText function, along with extra controls
+   * for the {@link ContextOptions} and {@link StorageOptions}.
    * @returns The result of the generateText function.
    */
   async generateText<
@@ -456,17 +628,21 @@ export class Agent<AgentTools extends ToolSet> {
     );
     const toolCtx = { ...ctx, userId, threadId, messageId };
     const tools = wrapTools(toolCtx, this.options.tools, args.tools) as TOOLS;
-    const maxSteps = args.maxSteps ?? this.options.maxSteps;
+    const saveOutputMessages =
+      args.saveOutputMessages ??
+      this.options.storageOptions?.saveOutputMessages;
     try {
       const result = (await generateText({
+        // Can be overridden
         model: this.options.chat,
+        maxSteps: this.options.maxSteps,
+        maxRetries: this.options.maxRetries,
         ...aiArgs,
-        maxSteps,
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
         toolChoice: args.toolChoice as any,
         tools,
         onStepFinish: async (step) => {
-          if (threadId && messageId && args.saveOutputMessages !== false) {
+          if (threadId && messageId && saveOutputMessages !== false) {
             await this.saveStep(ctx, {
               threadId,
               messageId,
@@ -490,6 +666,18 @@ export class Agent<AgentTools extends ToolSet> {
     }
   }
 
+  /**
+   * This behaves like {@link streamText} 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 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,
     OUTPUT = never,
@@ -511,11 +699,15 @@ export class Agent<AgentTools extends ToolSet> {
     );
     const toolCtx = { ...ctx, userId, threadId, messageId };
     const tools = wrapTools(toolCtx, this.options.tools, args.tools) as TOOLS;
-    const maxSteps = args.maxSteps ?? this.options.maxSteps;
+    const saveOutputMessages =
+      args.saveOutputMessages ??
+      this.options.storageOptions?.saveOutputMessages;
     const result = streamText({
+      // Can be overridden
       model: this.options.chat,
+      maxSteps: this.options.maxSteps,
+      maxRetries: this.options.maxRetries,
       ...aiArgs,
-      maxSteps,
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       toolChoice: args.toolChoice as any,
       tools,
@@ -525,7 +717,7 @@ export class Agent<AgentTools extends ToolSet> {
       },
       onError: async (error) => {
         console.error("onError", error);
-        if (threadId && messageId && args.saveOutputMessages !== false) {
+        if (threadId && messageId && saveOutputMessages !== false) {
           await ctx.runMutation(this.component.messages.rollbackMessage, {
             messageId,
             error: (error.error as Error).message,
@@ -562,7 +754,6 @@ export class Agent<AgentTools extends ToolSet> {
       userId,
       threadId,
       parentMessageId,
-      saveAllInputMessages,
       system,
       ...args
     }: {
@@ -577,6 +768,12 @@ export class Agent<AgentTools extends ToolSet> {
     args: T;
     messageId: string | undefined;
   }> {
+    const saveAny =
+      args.saveAnyInputMessages ??
+      this.options.storageOptions?.saveAnyInputMessages;
+    const saveAll =
+      args.saveAllInputMessages ??
+      this.options.storageOptions?.saveAllInputMessages;
     const messages = promptOrMessagesToCoreMessages(args);
     const contextMessages = await this.fetchContextMessages(ctx, {
       messages,
@@ -586,11 +783,11 @@ export class Agent<AgentTools extends ToolSet> {
       ...args,
     });
     let messageId: string | undefined;
-    if (threadId) {
+    if (threadId && saveAny !== false) {
       const saved = await this.saveMessages(ctx, {
         threadId,
         userId,
-        messages: saveAllInputMessages ? messages : messages.slice(-1),
+        messages: saveAll ? messages : messages.slice(-1),
         pending: true,
         // We should just fail if you pass in an ID for the message, fail those children
         // failPendingSteps: true,
@@ -609,6 +806,18 @@ export class Agent<AgentTools extends ToolSet> {
     };
   }
 
+  /**
+   * This behaves like {@link generateObject} 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 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,
     { userId, threadId }: { userId?: string; threadId?: string },
@@ -619,14 +828,19 @@ export class Agent<AgentTools extends ToolSet> {
       { ...args, userId, threadId }
     );
 
+    const saveOutputMessages =
+      args.saveOutputMessages ??
+      this.options.storageOptions?.saveOutputMessages;
     try {
       const result = (await generateObject({
+        // Can be overridden
         model: this.options.chat,
+        maxRetries: this.options.maxRetries,
         ...aiArgs,
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
       } as any)) as GenerateObjectResult<T> & GenerationOutputMetadata;
 
-      if (threadId && messageId && args.saveOutputMessages !== false) {
+      if (threadId && messageId && saveOutputMessages !== false) {
         await this.saveObject(ctx, { threadId, messageId, result });
       }
       result.messageId = messageId;
@@ -642,6 +856,18 @@ export class Agent<AgentTools extends ToolSet> {
     }
   }
 
+  /**
+   * This behaves like {@link 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: RunMutationCtx,
     { userId, threadId }: { userId?: string; threadId?: string },
@@ -654,8 +880,13 @@ export class Agent<AgentTools extends ToolSet> {
       ctx,
       { ...args, userId, threadId }
     );
+    const saveOutputMessages =
+      args.saveOutputMessages ??
+      this.options.storageOptions?.saveOutputMessages;
     const stream = streamObject<T>({
+      // Can be overridden
       model: this.options.chat,
+      maxRetries: this.options.maxRetries,
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       ...(aiArgs as any),
       onError: async (error) => {
@@ -663,7 +894,7 @@ export class Agent<AgentTools extends ToolSet> {
         return args.onError?.(error);
       },
       onFinish: async (result) => {
-        if (threadId && messageId && args.saveOutputMessages !== false) {
+        if (threadId && messageId && saveOutputMessages !== false) {
           await this.saveObject(ctx, {
             threadId,
             messageId,
@@ -691,6 +922,13 @@ export class Agent<AgentTools extends ToolSet> {
     return stream;
   }
 
+  /**
+   * Manually save the result of a generateObject call to the thread.
+   * This happens automatically when using {@link generateObject} or {@link streamObject}
+   * from the `thread` object created by {@link continueThread} or {@link createThread}.
+   * @param ctx The context passed from the mutation or action function calling this.
+   * @param args The arguments to the saveObject function.
+   */
   async saveObject(
     ctx: RunMutationCtx,
     args: {
@@ -758,7 +996,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.
+   * Note: currently this is not well typed. The return type of the action is always `any`.
+   * @param spec Configuration for the agent acting as an action, including
+   *   {@link ContextOptions} and maxSteps.
    */
   asAction(spec?: { contextOptions?: ContextOptions; maxSteps?: number }) {
     return internalActionGeneric({
@@ -847,7 +1089,8 @@ export class Agent<AgentTools extends ToolSet> {
   }
 
   /**
-   * Create a tool that can call this agent.
+   * Create a tool out of this agent so other agents can call this one.
+   * Create a tool out of this agent so other agents can call this one.
    * @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.
@@ -985,7 +1228,7 @@ type BaseGenerateObjectOptions = StorageOptions &
 
 type GenerateObjectObjectOptions<T extends Record<string, unknown>> =
   BaseGenerateObjectOptions & {
-    output: "object";
+    output?: "object";
     mode?: "auto" | "json" | "tool";
     schema: z.Schema<T>;
     schemaName?: string;
@@ -1042,7 +1285,20 @@ type OurStreamObjectArgs<T> = StreamObjectArgs<T> &
   >;
 
 interface Thread<AgentTools extends ToolSet> {
+  /**
+   * The target threadId, from the startThread or continueThread initializers.
+   */
   threadId: string;
+  /**
+   * This behaves like {@link generateText} 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 args The arguments to the generateText function, along with extra controls
+   * for the {@link ContextOptions} and {@link StorageOptions}.
+   * @returns The result of the generateText function.
+   */
   generateText<TOOLS extends ToolSet, OUTPUT = never, OUTPUT_PARTIAL = never>(
     args: TextArgs<
       AgentTools,
@@ -1053,6 +1309,16 @@ interface Thread<AgentTools extends ToolSet> {
     GenerateTextResult<TOOLS & AgentTools, OUTPUT> & GenerationOutputMetadata
   >;
 
+  /**
+   * This behaves like {@link streamText} 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 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.
+   */
   streamText<TOOLS extends ToolSet, OUTPUT = never, PARTIAL_OUTPUT = never>(
     args: TextArgs<
       AgentTools,
@@ -1063,13 +1329,42 @@ interface Thread<AgentTools extends ToolSet> {
     StreamTextResult<TOOLS & AgentTools, PARTIAL_OUTPUT> &
       GenerationOutputMetadata
   >;
-  // TODO: add all the overloads
+  /**
+   * This behaves like {@link generateObject} 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. This overload is for objects, arrays, and enums.
+   * Use {@link continueThread} to get a version of this function already scoped
+   * to a thread (and optionally userId).
+   * @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.
+   */
   generateObject<T>(
     args: OurObjectArgs<T>
   ): Promise<GenerateObjectResult<T> & GenerationOutputMetadata>;
+  /**
+   * This behaves like {@link generateObject} 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. This overload is for when there's no schema.
+   * Use {@link continueThread} to get a version of this function already scoped
+   * to a thread (and optionally userId).
+   * @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.
+   */
   generateObject(
     args: GenerateObjectNoSchemaOptions
   ): Promise<GenerateObjectResult<JSONValue> & GenerationOutputMetadata>;
+  /**
+   * This behaves like {@link 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 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.
+   */
   streamObject<T>(
     args: OurStreamObjectArgs<T>
   ): Promise<