@convex-dev/agent 0.1.17 → 0.1.18-alpha.1

package/src/client/index.ts

@@ -20,10 +20,12 @@ import {
   streamObject,
   streamText,
 } from "ai";
-import { assert } from "convex-helpers";
+import { assert, omit, pick } from "convex-helpers";
 import {
   internalActionGeneric,
   internalMutationGeneric,
+  type GenericActionCtx,
+  type GenericDataModel,
   type PaginationOptions,
   type PaginationResult,
   type WithoutSystemFields,
@@ -134,7 +136,30 @@ export type {
   UsageHandler,
 };
 
-export class Agent<AgentTools extends ToolSet = ToolSet> {
+export class Agent<
+  /**
+   * You can require that all `ctx` args to generateText & streamText
+   * have a certain shape by passing a type here.
+   * e.g.
+   * ```ts
+   * const myAgent = new Agent<{ orgId: string }>(...);
+   * ```
+   * This is useful if you want to share that type in `createTool`
+   * e.g.
+   * ```ts
+   * type MyCtx = ToolCtx & { orgId: string };
+   * const myTool = createTool({
+   *   args: z.object({...}),
+   *   description: "...",
+   *   handler: async (ctx: MyCtx, args) => {
+   *     // use ctx.orgId
+   *   },
+   * });
+   */
+  CustomCtx extends object = object,
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  AgentTools extends ToolSet = any,
+> {
   constructor(
     public component: AgentComponent,
     public options: {
@@ -218,7 +243,7 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
    * @returns The threadId of the new thread and the thread object.
    */
   async createThread<ThreadTools extends ToolSet | undefined = undefined>(
-    ctx: RunActionCtx,
+    ctx: RunActionCtx & CustomCtx,
     args?: {
       /**
        * The userId to associate with the thread. If not provided, the thread will be
@@ -288,7 +313,7 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
     threadId: string;
   }>;
   async createThread<ThreadTools extends ToolSet | undefined = undefined>(
-    ctx: ActionCtx | RunMutationCtx,
+    ctx: (ActionCtx & CustomCtx) | RunMutationCtx,
     args?: {
       userId: string | null;
       title?: string;
@@ -325,7 +350,7 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
    * @returns Functions bound to the userId and threadId on a `{thread}` object.
    */
   async continueThread<ThreadTools extends ToolSet | undefined = undefined>(
-    ctx: ActionCtx,
+    ctx: ActionCtx & CustomCtx,
     args: {
       /**
        * The associated thread created by {@link createThread}
@@ -410,7 +435,7 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
     OUTPUT = never,
     OUTPUT_PARTIAL = never,
   >(
-    ctx: ActionCtx,
+    ctx: ActionCtx & CustomCtx,
     {
       userId: argsUserId,
       threadId,
@@ -433,14 +458,15 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
     GenerateTextResult<TOOLS extends undefined ? AgentTools : TOOLS, OUTPUT> &
       GenerationOutputMetadata
   > {
+    const opts = { ...this.options, ...options, usageHandler };
     const context = await this._saveMessagesAndFetchContext(ctx, args, {
       userId: argsUserId ?? undefined,
       threadId,
-      ...options,
+      ...opts,
     });
     const { args: aiArgs, messageId, order, userId } = context;
     const toolCtx = {
-      ...(ctx as UserActionCtx),
+      ...(ctx as UserActionCtx & CustomCtx),
       userId,
       threadId,
       messageId,
@@ -450,10 +476,7 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
       toolCtx,
       args.tools ?? threadTools ?? this.options.tools,
     ) as TOOLS extends undefined ? AgentTools : TOOLS;
-    const saveOutputMessages = this._shouldSaveOutputMessages(
-      options?.storageOptions,
-    );
-    const trackUsage = usageHandler ?? this.options.usageHandler;
+    const saveOutput = opts.storageOptions?.saveMessages !== "none";
     try {
       const result = (await generateText({
         // Can be overridden
@@ -461,7 +484,7 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
         ...aiArgs,
         tools,
         onStepFinish: async (step) => {
-          if (threadId && messageId && saveOutputMessages) {
+          if (threadId && messageId && saveOutput) {
             await this.saveStep(ctx, {
               userId,
               threadId,
@@ -480,8 +503,8 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
               response: step.response,
             });
           }
-          if (trackUsage && step.usage) {
-            await trackUsage(ctx, {
+          if (opts.usageHandler && step.usage) {
+            await opts.usageHandler(ctx, {
               userId,
               threadId,
               agentName: this.options.name,
@@ -525,7 +548,7 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
     OUTPUT = never,
     PARTIAL_OUTPUT = never,
   >(
-    ctx: ActionCtx,
+    ctx: ActionCtx & CustomCtx,
     {
       userId: argsUserId,
       threadId,
@@ -566,14 +589,15 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
     > &
       GenerationOutputMetadata
   > {
+    const opts = { ...this.options, ...options, usageHandler };
     const context = await this._saveMessagesAndFetchContext(ctx, args, {
       userId: argsUserId ?? undefined,
       threadId,
-      ...options,
+      ...opts,
     });
     const { args: aiArgs, messageId, order, stepOrder, userId } = context;
     const toolCtx = {
-      ...(ctx as UserActionCtx),
+      ...(ctx as UserActionCtx & CustomCtx),
       userId,
       threadId,
       messageId,
@@ -583,13 +607,10 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
       toolCtx,
       args.tools ?? threadTools ?? this.options.tools,
     ) as TOOLS extends undefined ? AgentTools : TOOLS;
-    const saveOutputMessages = this._shouldSaveOutputMessages(
-      options?.storageOptions,
-    );
-    const trackUsage = usageHandler ?? this.options.usageHandler;
+    const saveOutput = opts.storageOptions?.saveMessages !== "none";
     const streamer =
-      threadId && options?.saveStreamDeltas
-        ? new DeltaStreamer(this.component, ctx, options.saveStreamDeltas, {
+      threadId && opts.saveStreamDeltas
+        ? new DeltaStreamer(this.component, ctx, opts.saveStreamDeltas, {
             threadId,
             userId,
             agentName: this.options.name,
@@ -619,7 +640,7 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
       },
       onError: async (error) => {
         console.error("onError", error);
-        if (threadId && messageId && saveOutputMessages) {
+        if (threadId && messageId && saveOutput) {
           await ctx.runMutation(this.component.messages.rollbackMessage, {
             messageId,
             error: (error.error as Error).message,
@@ -630,7 +651,7 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
       },
       onStepFinish: async (step) => {
         // console.log("onStepFinish", step);
-        if (threadId && messageId && saveOutputMessages) {
+        if (threadId && messageId && saveOutput) {
           const saved = await this.saveStep(ctx, {
             userId,
             threadId,
@@ -650,8 +671,8 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
             response: step.response,
           });
         }
-        if (trackUsage && step.usage) {
-          await trackUsage(ctx, {
+        if (opts.usageHandler && step.usage) {
+          await opts.usageHandler(ctx, {
             userId,
             threadId,
             agentName: this.options.name,
@@ -701,23 +722,21 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
      */
     options?: Options,
   ): Promise<GenerateObjectResult<T> & GenerationOutputMetadata> {
+    const opts = { ...this.options, ...options, usageHandler };
     const context = await this._saveMessagesAndFetchContext(ctx, args, {
       userId: argsUserId ?? undefined,
       threadId,
-      ...options,
+      ...opts,
     });
     const { args: aiArgs, messageId, order, userId } = context;
-    const trackUsage = usageHandler ?? this.options.usageHandler;
-    const saveOutputMessages = this._shouldSaveOutputMessages(
-      options?.storageOptions,
-    );
+    const saveOutput = opts.storageOptions?.saveMessages !== "none";
     try {
       const result = (await generateObject(
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
         aiArgs as any,
       )) as GenerateObjectResult<T> & GenerationOutputMetadata;
 
-      if (threadId && messageId && saveOutputMessages) {
+      if (threadId && messageId && saveOutput) {
         await this.saveObject(ctx, {
           threadId,
           promptMessageId: messageId,
@@ -738,8 +757,8 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
           response: result.response,
         });
       }
-      if (trackUsage && result.usage) {
-        await trackUsage(ctx, {
+      if (opts.usageHandler && result.usage) {
+        await opts.usageHandler(ctx, {
           userId,
           threadId,
           agentName: this.options.name,
@@ -792,16 +811,14 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
     StreamObjectResult<DeepPartial<T>, T, never> & GenerationOutputMetadata
   > {
     // TODO: unify all this shared code between all the generate* and stream* functions
+    const opts = { ...this.options, ...options, usageHandler };
     const context = await this._saveMessagesAndFetchContext(ctx, args, {
       userId: argsUserId ?? undefined,
       threadId,
-      ...options,
+      ...opts,
     });
     const { args: aiArgs, messageId, order, userId } = context;
-    const trackUsage = usageHandler ?? this.options.usageHandler;
-    const saveOutputMessages = this._shouldSaveOutputMessages(
-      options?.storageOptions,
-    );
+    const saveOutput = opts.storageOptions?.saveMessages !== "none";
     const stream = streamObject<T>({
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       ...(aiArgs as any),
@@ -810,7 +827,7 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
         return args.onError?.(error);
       },
       onFinish: async (result) => {
-        if (threadId && messageId && saveOutputMessages) {
+        if (threadId && messageId && saveOutput) {
           await this.saveObject(ctx, {
             userId,
             threadId,
@@ -832,8 +849,8 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
             provider: aiArgs.model.provider,
           });
         }
-        if (trackUsage && result.usage) {
-          await trackUsage(ctx, {
+        if (opts.usageHandler && result.usage) {
+          await opts.usageHandler(ctx, {
             userId,
             threadId,
             agentName: this.options.name,
@@ -1031,10 +1048,13 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
     },
   ): Promise<MessageDoc[]> {
     assert(args.userId || args.threadId, "Specify userId or threadId");
-    const opts = this._mergedContextOptions(args.contextOptions);
+    const contextOptions = {
+      ...this.options.contextOptions,
+      ...args.contextOptions,
+    };
     return fetchContextMessages(ctx, this.component, {
       ...args,
-      contextOptions: opts,
+      contextOptions,
       getEmbedding: async (text) => {
         assert("runAction" in ctx);
         assert(
@@ -1129,7 +1149,7 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
     const textEmbeddings = await this.doEmbed(ctx, {
       userId,
       threadId,
-      values: messageTexts.filter((t): t is string => !!t),
+      values: messageTexts as string[],
     });
     // TODO: record usage of embeddings
     // Then assemble the embeddings into a single array with nulls for the messages without text.
@@ -1566,8 +1586,6 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
     order: number | undefined;
     stepOrder: number | undefined;
   }> {
-    contextOptions ||= this.options.contextOptions;
-    storageOptions ||= this.options.storageOptions;
     // If only a promptMessageId is provided, this will be empty.
     const messages = promptOrMessagesToCoreMessages(args);
     const userId =
@@ -1658,25 +1676,6 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
     };
   }
 
-  _shouldSaveOutputMessages(storageOpts?: StorageOptions): boolean {
-    const opts = storageOpts ?? this.options.storageOptions;
-    return opts?.saveMessages !== "none";
-  }
-
-  _mergedContextOptions(opts: ContextOptions | undefined): ContextOptions {
-    const searchOptions = {
-      ...this.options.contextOptions?.searchOptions,
-      ...opts?.searchOptions,
-    };
-    return {
-      ...this.options.contextOptions,
-      ...opts,
-      searchOptions: searchOptions.limit
-        ? (searchOptions as ContextOptions["searchOptions"])
-        : undefined,
-    };
-  }
-
   async doEmbed(
     ctx: RunActionCtx,
     options: {
@@ -1848,52 +1847,83 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
    * @param spec Configuration for the agent acting as an action, including
    *   {@link ContextOptions}, {@link StorageOptions}, and maxSteps.
    */
-  asTextAction(spec?: {
-    /**
-     * The maximum number of steps to take in this action.
-     * Defaults to the {@link Agent.maxSteps} option.
-     */
-    maxSteps?: number;
-    /**
-     * The {@link ContextOptions} to use for fetching contextual messages and
-     * saving input/output messages.
-     * Defaults to the {@link Agent.contextOptions} option.
-     */
-    contextOptions?: ContextOptions;
-    /**
-     * The {@link StorageOptions} to use for saving input/output messages.
-     * Defaults to the {@link Agent.storageOptions} option.
-     */
-    storageOptions?: StorageOptions;
-    /**
-     * Whether to stream the text.
-     * If false, it will generate the text in a single call. (default)
-     * If true or {@link StreamingOptions}, it will stream the text from the LLM
-     * and save the chunks to the database with the options you specify, or the
-     * defaults if you pass true.
-     */
-    stream?: boolean | StreamingOptions;
-  }) {
+  asTextAction<DataModel extends GenericDataModel>(
+    spec?: {
+      /**
+       * The maximum number of steps to take in this action.
+       * Defaults to the {@link Agent.maxSteps} option.
+       */
+      maxSteps?: number;
+      /**
+       * The {@link ContextOptions} to use for fetching contextual messages and
+       * saving input/output messages.
+       * Defaults to the {@link Agent.contextOptions} option.
+       */
+      contextOptions?: ContextOptions;
+      /**
+       * The {@link StorageOptions} to use for saving input/output messages.
+       * Defaults to the {@link Agent.storageOptions} option.
+       */
+      storageOptions?: StorageOptions;
+      /**
+       * Whether to stream the text.
+       * If false, it will generate the text in a single call. (default)
+       * If true or {@link StreamingOptions}, it will stream the text from the LLM
+       * and save the chunks to the database with the options you specify, or the
+       * defaults if you pass true.
+       */
+      stream?: boolean | StreamingOptions;
+    } & (CustomCtx extends Record<string, unknown>
+      ? {
+          /**
+           * If you have a custom ctx that you use with the Agent
+           * (e.g. new Agent<{ orgId: string }>(...))
+           * you need to provide this function to add any extra fields.
+           * e.g.
+           * ```ts
+           * const myAgent = new Agent<{ orgId: string }>(...);
+           * const myAction = myAgent.asTextAction({
+           *   customCtx: (ctx: ActionCtx, target, llmArgs) => {
+           *     const orgId = await lookupOrgId(ctx, target.threadId);
+           *     return { orgId };
+           *   },
+           * });
+           * ```
+           * Then, in your tools, you can
+           */
+          customCtx: (
+            ctx: GenericActionCtx<DataModel>,
+            target: {
+              userId?: string | undefined;
+              threadId?: string | undefined;
+            },
+            llmArgs: TextArgs<AgentTools>,
+          ) => CustomCtx;
+        }
+      : { customCtx?: never }),
+  ) {
     const maxSteps = spec?.maxSteps ?? this.options.maxSteps;
     return internalActionGeneric({
       args: vTextArgs,
-      handler: async (ctx, args) => {
-        const { contextOptions, storageOptions, ...rest } = args;
+      handler: async (ctx_, args) => {
         const stream =
           args.stream === true ? spec?.stream || true : spec?.stream ?? false;
         const targetArgs = { userId: args.userId, threadId: args.threadId };
-        const llmArgs = { maxSteps, ...rest };
+        const llmArgs = {
+          maxSteps,
+          ...omit(args, ["storageOptions", "contextOptions"]),
+        };
         const opts = {
-          contextOptions:
-            contextOptions ??
-            spec?.contextOptions ??
-            this.options.contextOptions,
-          storageOptions:
-            storageOptions ??
-            spec?.storageOptions ??
-            this.options.storageOptions,
+          ...this.options,
+          ...(spec && pick(spec, ["contextOptions", "storageOptions"])),
+          ...pick(args, ["contextOptions", "storageOptions"]),
           saveStreamDeltas: stream,
         };
+        const ctx = (
+          spec?.customCtx
+            ? { ...ctx_, ...spec.customCtx(ctx_, targetArgs, llmArgs) }
+            : ctx_
+        ) as UserActionCtx & CustomCtx;
         if (stream) {
           const result = await this.streamText(ctx, targetArgs, llmArgs, opts);
           await result.consumeStream();
@@ -1935,25 +1965,16 @@ export class Agent<AgentTools extends ToolSet = ToolSet> {
     return internalActionGeneric({
       args: vSafeObjectArgs,
       handler: async (ctx, args) => {
-        const { contextOptions, storageOptions, ...rest } = args;
+        const overrides = pick(args, ["userId", "threadId"]);
         const value = await this.generateObject(
           ctx,
           { userId: args.userId, threadId: args.threadId },
           {
             ...spec,
             maxSteps,
-            ...rest,
+            ...omit(args, ["userId", "threadId"]),
           } as unknown as OurObjectArgs<unknown>,
-          {
-            contextOptions:
-              contextOptions ??
-              options?.contextOptions ??
-              this.options.contextOptions,
-            storageOptions:
-              storageOptions ??
-              options?.storageOptions ??
-              this.options.storageOptions,
-          },
+          { ...this.options, ...options, ...overrides },
         );
         return {
           object: value.object as T,

package/src/client/types.ts

@@ -29,7 +29,6 @@ import type {
   FunctionReturnType,
   GenericActionCtx,
   GenericDataModel,
-  OptionalRestArgs,
 } from "convex/server";
 import type { GenericId } from "convex/values";
 import type { Schema } from "zod";