@convex-dev/agent 0.0.12 → 0.0.14-alpha.0

package/src/client/index.ts

@@ -19,13 +19,12 @@ import type {
 import {
   generateObject,
   generateText,
-  jsonSchema,
   streamObject,
   streamText,
   tool,
 } from "ai";
 import { assert } from "convex-helpers";
-import { internalActionGeneric } from "convex/server";
+import { internalActionGeneric, internalMutationGeneric } from "convex/server";
 import { Infer, v } from "convex/values";
 import { z } from "zod";
 import { Mounts } from "../component/_generated/api.js";
@@ -37,7 +36,7 @@ import {
   type AIMessageWithoutId,
   deserializeMessage,
   promptOrMessagesToCoreMessages,
-  serializeMessageWithId,
+  serializeMessage,
   serializeNewMessagesInStep,
   serializeObjectResult,
   serializeStep,
@@ -46,12 +45,15 @@ import {
   DEFAULT_MESSAGE_RANGE,
   DEFAULT_RECENT_MESSAGES,
   extractText,
+  isTool,
 } from "../shared.js";
 import {
   type CallSettings,
+  MessageWithMetadata,
   type ProviderMetadata,
   type ProviderOptions,
   type SearchOptions,
+  Usage,
   vSafeObjectArgs,
   vTextArgs,
 } from "../validators.js";
@@ -145,6 +147,19 @@ export type GenerationOutputMetadata = { messageId?: string };
 
 type CoreMessageMaybeWithId = CoreMessage & { id?: string | undefined };
 
+export type UsageHandler = (
+  ctx: RunActionCtx,
+  args: {
+    userId?: string;
+    threadId?: string;
+    usage: Usage;
+    // Often has more information, like cached token usage in the case of openai.
+    providerMetadata?: ProviderMetadata;
+    model: string;
+    provider: string;
+  }
+) => Promise<void>;
+
 export class Agent<AgentTools extends ToolSet> {
   constructor(
     public component: UseApi<Mounts>,
@@ -206,6 +221,10 @@ export class Agent<AgentTools extends ToolSet> {
        * This can be overridden at each generate/stream callsite.
        */
       maxRetries?: number;
+      /**
+       * The usage handler to use for this agent.
+       */
+      usageHandler?: UsageHandler;
     }
   ) {}
 
@@ -433,7 +452,7 @@ export class Agent<AgentTools extends ToolSet> {
         }
       | undefined;
     if (this.options.textEmbedding) {
-      const messageTexts = messages.map((m) => extractText(m));
+      const messageTexts = messages.map((m) => !isTool(m) && extractText(m));
       // Find the indexes of the messages that have text.
       const textIndexes = messageTexts
         .map((t, i) => (t ? i : undefined))
@@ -445,6 +464,7 @@ export class Agent<AgentTools extends ToolSet> {
       const textEmbeddings = await this.options.textEmbedding.doEmbed({
         values: messageTexts.filter((t): t is string => !!t),
       });
+      // TODO: record usage of embeddings
       // Then assemble the embeddings into a single array with nulls for the messages without text.
       const embeddingsOrNull = Array(messages.length).fill(null);
       textIndexes.forEach((i, j) => {
@@ -463,12 +483,6 @@ 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.
@@ -482,17 +496,16 @@ export class Agent<AgentTools extends ToolSet> {
       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.
+       * Metadata to save with the messages. Each element corresponds to the
+       * message at the same index.
        */
+      metadata?: Omit<MessageWithMetadata, "message">[];
       /**
        * If false, it will "commit" the messages immediately.
        * If true, it will mark them as pending until the final step has finished.
+       * Defaults to false.
        */
       pending?: boolean;
-      /**
-       * The message that this is responding to.
-       */
       /**
        * The message that this is responding to.
        */
@@ -500,10 +513,7 @@ export class Agent<AgentTools extends ToolSet> {
       /**
        * 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.
+       * Defaults to true.
        */
       failPendingSteps?: boolean;
     }
@@ -511,13 +521,23 @@ export class Agent<AgentTools extends ToolSet> {
     lastMessageId: string;
     messageIds: string[];
   }> {
+    const embeddings = await this.getEmbeddings(args.messages);
     const result = await ctx.runMutation(this.component.messages.addMessages, {
       threadId: args.threadId,
       userId: args.userId,
       agentName: this.options.name,
-      model: this.options.chat.modelId,
-      messages: args.messages.map(serializeMessageWithId),
-      embeddings: await this.getEmbeddings(args.messages),
+      messages: args.messages.map(
+        (m, i) =>
+          ({
+            embedding: embeddings?.vectors[i] && {
+              model: embeddings.model,
+              dimension: embeddings.dimension,
+              vector: embeddings.vectors[i],
+            },
+            ...args.metadata?.[i],
+            message: serializeMessage(m),
+          }) as MessageWithMetadata
+      ),
       failPendingSteps: args.failPendingSteps ?? true,
       pending: args.pending ?? false,
       parentMessageId: args.parentMessageId,
@@ -528,11 +548,6 @@ 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.
@@ -550,27 +565,41 @@ export class Agent<AgentTools extends ToolSet> {
        * The step to save, possibly including multiple tool calls.
        */
       step: StepResult<TOOLS>;
+      /**
+       * The model used to generate the step.
+       * Defaults to the chat model for the Agent.
+       */
+      model?: string;
+      /**
+       * The provider of the model used to generate the step.
+       * Defaults to the chat provider for the Agent.
+       */
+      provider?: string;
     }
   ): Promise<void> {
     const step = serializeStep(args.step as StepResult<ToolSet>);
-    const messages = serializeNewMessagesInStep(args.step);
+    const messages = serializeNewMessagesInStep(args.step, {
+      provider: args.provider ?? this.options.chat.provider,
+      model: args.model ?? this.options.chat.modelId,
+    });
+    const embeddings = await this.getEmbeddings(messages.map((m) => m.message));
+    if (embeddings) {
+      const { model, dimension, vectors } = embeddings;
+      for (let i = 0; i < messages.length; i++) {
+        const vector = vectors[i];
+        if (vector) {
+          messages[i].embedding = { model, dimension, vector };
+        }
+      }
+    }
     await ctx.runMutation(this.component.messages.addStep, {
       threadId: args.threadId,
       messageId: args.messageId,
       step: { step, messages },
       failPendingSteps: false,
-      embeddings: await this.getEmbeddings(messages.map((m) => m.message)),
     });
   }
 
-  /**
-   * 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.
@@ -642,13 +671,14 @@ export class Agent<AgentTools extends ToolSet> {
     const saveOutputMessages =
       args.saveOutputMessages ??
       this.options.storageOptions?.saveOutputMessages;
+    const model = aiArgs.model ?? this.options.chat;
     try {
       const result = (await generateText({
         // Can be overridden
-        model: this.options.chat,
         maxSteps: this.options.maxSteps,
         maxRetries: this.options.maxRetries,
         ...aiArgs,
+        model,
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
         toolChoice: args.toolChoice as any,
         tools,
@@ -660,6 +690,16 @@ export class Agent<AgentTools extends ToolSet> {
               step,
             });
           }
+          if (this.options.usageHandler && step.usage) {
+            await this.options.usageHandler(ctx, {
+              userId,
+              threadId,
+              model: model.modelId,
+              provider: model.provider,
+              usage: step.usage,
+              providerMetadata: step.providerMetadata,
+            });
+          }
           return args.onStepFinish?.(step);
         },
       })) as GenerateTextResult<TOOLS, OUTPUT> & GenerationOutputMetadata;
@@ -713,12 +753,13 @@ export class Agent<AgentTools extends ToolSet> {
     const saveOutputMessages =
       args.saveOutputMessages ??
       this.options.storageOptions?.saveOutputMessages;
+    const model = aiArgs.model ?? this.options.chat;
     const result = streamText({
       // Can be overridden
-      model: this.options.chat,
       maxSteps: this.options.maxSteps,
       maxRetries: this.options.maxRetries,
       ...aiArgs,
+      model,
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       toolChoice: args.toolChoice as any,
       tools,
@@ -746,6 +787,16 @@ export class Agent<AgentTools extends ToolSet> {
             step,
           });
         }
+        if (this.options.usageHandler && step.usage) {
+          await this.options.usageHandler(ctx, {
+            userId,
+            threadId,
+            model: model.modelId,
+            provider: model.provider,
+            usage: step.usage,
+            providerMetadata: step.providerMetadata,
+          });
+        }
         return args.onStepFinish?.(step);
       },
     }) as StreamTextResult<TOOLS, PARTIAL_OUTPUT> & GenerationOutputMetadata;
@@ -755,6 +806,7 @@ export class Agent<AgentTools extends ToolSet> {
 
   async saveMessagesAndFetchContext<
     T extends {
+      id?: string;
       prompt?: string;
       messages?: CoreMessage[] | AIMessageWithoutId[];
       system?: string;
@@ -762,6 +814,7 @@ export class Agent<AgentTools extends ToolSet> {
   >(
     ctx: RunActionCtx | RunMutationCtx,
     {
+      id,
       userId,
       threadId,
       parentMessageId,
@@ -795,10 +848,12 @@ export class Agent<AgentTools extends ToolSet> {
     });
     let messageId: string | undefined;
     if (threadId && saveAny !== false) {
+      const coreMessages = saveAll ? messages : messages.slice(-1);
       const saved = await this.saveMessages(ctx, {
         threadId,
         userId,
-        messages: saveAll ? messages : messages.slice(-1),
+        messages: coreMessages,
+        metadata: coreMessages.length === 1 ? [{ id }] : undefined,
         pending: true,
         // We should just fail if you pass in an ID for the message, fail those children
         // failPendingSteps: true,
@@ -838,16 +893,16 @@ export class Agent<AgentTools extends ToolSet> {
       ctx,
       { ...args, userId, threadId }
     );
-
+    const model = aiArgs.model ?? this.options.chat;
     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,
+        model,
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
       } as any)) as GenerateObjectResult<T> & GenerationOutputMetadata;
 
@@ -855,6 +910,16 @@ export class Agent<AgentTools extends ToolSet> {
         await this.saveObject(ctx, { threadId, messageId, result });
       }
       result.messageId = messageId;
+      if (this.options.usageHandler && result.usage) {
+        await this.options.usageHandler(ctx, {
+          userId,
+          threadId,
+          model: model.modelId,
+          provider: model.provider,
+          usage: result.usage,
+          providerMetadata: result.providerMetadata,
+        });
+      }
       return result;
     } catch (error) {
       if (threadId && messageId) {
@@ -880,7 +945,7 @@ export class Agent<AgentTools extends ToolSet> {
    * @returns The result of the streamObject function.
    */
   async streamObject<T>(
-    ctx: RunMutationCtx,
+    ctx: RunActionCtx,
     { userId, threadId }: { userId?: string; threadId?: string },
     args: OurStreamObjectArgs<T>
   ): Promise<
@@ -891,15 +956,16 @@ export class Agent<AgentTools extends ToolSet> {
       ctx,
       { ...args, userId, threadId }
     );
+    const model = aiArgs.model ?? this.options.chat;
     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),
+      model,
       onError: async (error) => {
         console.error("onError", error);
         return args.onError?.(error);
@@ -924,6 +990,16 @@ export class Agent<AgentTools extends ToolSet> {
             },
           });
         }
+        if (this.options.usageHandler && result.usage) {
+          await this.options.usageHandler(ctx, {
+            userId,
+            threadId,
+            model: model.modelId,
+            provider: model.provider,
+            usage: result.usage,
+            providerMetadata: result.providerMetadata,
+          });
+        }
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
         return args.onFinish?.(result as any);
       },
@@ -946,15 +1022,35 @@ export class Agent<AgentTools extends ToolSet> {
       threadId: string;
       messageId: string;
       result: GenerateObjectResult<unknown>;
+      metadata?: Omit<MessageWithMetadata, "message">;
     }
   ): Promise<void> {
-    const step = serializeObjectResult(args.result);
+    const { step, messages: withoutEmbed } = serializeObjectResult(
+      args.result,
+      {
+        model: this.options.chat.modelId,
+        provider: this.options.chat.provider,
+      }
+    );
+    const embeddings = await this.getEmbeddings([withoutEmbed[0].message]);
+    const messages = embeddings?.vectors[0]
+      ? [
+          {
+            ...withoutEmbed[0],
+            embedding: {
+              dimension: embeddings.dimension,
+              model: embeddings.model,
+              vector: embeddings.vectors[0],
+            },
+          },
+        ]
+      : withoutEmbed;
+
     await ctx.runMutation(this.component.messages.addStep, {
       threadId: args.threadId,
       messageId: args.messageId,
       failPendingSteps: false,
-      embeddings: await this.getEmbeddings([step.messages[0].message]),
-      step,
+      step: { step, messages },
     });
   }
 
@@ -1001,11 +1097,45 @@ export class Agent<AgentTools extends ToolSet> {
           values: [text],
         })
       ).embeddings[0];
+      // TODO: record usage of embeddings
       search.vectorModel = this.options.textEmbedding.modelId;
     }
     return search;
   }
 
+  /**
+   * Create a mutation that creates a thread so you can call it from a Workflow.
+   * e.g.
+   * ```ts
+   * // in convex/foo.ts
+   * export const createThread = weatherAgent.createThreadMutation();
+   *
+   * const workflow = new WorkflowManager(components.workflow);
+   * export const myWorkflow = workflow.define({
+   *   args: {},
+   *   handler: async (step) => {
+   *     const { threadId } = await step.runMutation(internal.foo.createThread);
+   *     // use the threadId to generate text, object, etc.
+   *   },
+   * });
+   * ```
+   * @returns A mutation that creates a thread.
+   */
+  createThreadMutation() {
+    return internalMutationGeneric({
+      args: {
+        userId: v.optional(v.string()),
+        parentThreadIds: v.optional(v.array(v.string())),
+        title: v.optional(v.string()),
+        summary: v.optional(v.string()),
+      },
+      handler: async (ctx, args) => {
+        const { threadId } = await this.createThread(ctx, args);
+        return { threadId };
+      },
+    });
+  }
+
   /**
    * Create an action out of this agent so you can call it from workflows or other actions
    * without a wrapping function.