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

package/src/client/index.ts

@@ -1,38 +1,44 @@
-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 {
+  CoreMessage,
+  DeepPartial,
+  GenerateObjectResult,
+  GenerateTextResult,
+  StepResult,
+  StreamObjectResult,
   StreamTextResult,
   Tool,
+  ToolChoice,
+  ToolExecutionOptions,
   ToolSet,
-  StepResult,
   Message as UIMessage,
-  CoreMessage,
-  GenerateObjectResult,
-  StreamObjectResult,
-  DeepPartial,
-  GenerateTextResult,
 } from "ai";
 import {
-  generateText,
-  generateObject,
-  streamText,
-  streamObject,
   convertToCoreMessages,
   coreMessageSchema,
+  generateObject,
+  generateText,
+  streamObject,
+  streamText,
 } from "ai";
+import type { ZodType } from "zod";
+import { api } from "../component/_generated/api";
+import { Message, MessageStatus, SearchOptions } 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 } from "convex-helpers/server/zod";
+import { GenericActionCtx, GenericDataModel } from "convex/server";
+import { Infer, Validator } from "convex/values";
 import {
-  serializeMessage,
   serializeMessageWithId,
   serializeNewMessagesInStep,
   serializeStep,
 } from "../mapping";
+import { DEFAULT_MESSAGE_RANGE, extractText } from "../shared";
 
 export type ContextOptions = {
+  parentMessageId?: string;
   includeToolMessages?: boolean;
   recentMessages?: number;
   searchOptions?: {
@@ -44,9 +50,20 @@ export type ContextOptions = {
   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: {
@@ -54,7 +71,7 @@ export class Agent {
       chat: LanguageModelV1;
       textEmbedding?: EmbeddingModelV1<string>;
       defaultSystemPrompt?: string;
-      tools?: Record<string, Tool>;
+      tools?: AgentTools;
     }
   ) {}
 
@@ -93,7 +110,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
@@ -125,7 +142,7 @@ export class Agent {
     }
   ): Promise<{
     chatId: string;
-    chat?: Chat;
+    chat?: Chat<AgentTools>;
   }> {
     const chatDoc = await ctx.runMutation(this.component.messages.createChat, {
       defaultSystemPrompt: this.options.defaultSystemPrompt,
@@ -157,7 +174,7 @@ export class Agent {
       userId?: string;
     }
   ): Promise<{
-    chat: Chat;
+    chat: Chat<AgentTools>;
   }> {
     // return this.component.continueChat(ctx, args);
     return {
@@ -166,7 +183,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>,
     };
   }
 
@@ -190,9 +207,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(args, args.messages)),
         }
       );
+      // TODO: track what messages we used for context
       contextMessages.push(...searchMessages.map((m) => m.message!));
     }
     if (args.chatId) {
@@ -202,6 +221,7 @@ export class Agent {
           chatId: args.chatId,
           isTool: args.includeToolMessages ?? false,
           limit: args.recentMessages,
+          parentMessageId: args.parentMessageId,
           order: "desc",
           statuses: ["success"],
         }
@@ -216,22 +236,12 @@ export class Agent {
     args: {
       chatId: string;
       messages: CoreMessageMaybeWithId[];
-      steps?: StepResult<ToolSet>[];
       pending?: boolean;
+      parentMessageId?: string;
     }
   ): 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,
@@ -240,19 +250,14 @@ export class Agent {
       messages: args.messages.map(serializeMessageWithId),
       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 +272,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 +280,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;
@@ -315,13 +318,16 @@ export class Agent {
       chatId,
     }: {
       userId?: string;
-      chatId?: 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 +336,25 @@ 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,
+      messages: args.saveAllInputMessages ? messages : messages.slice(-1),
+      pending: true,
+      parentMessageId: args.parentMessageId,
+    });
+    const defaults = this.options.tools;
+    const tools = wrapTools(ctx, chatId, userId, defaults, args.tools) as TOOLS;
     try {
       const result = await generateText({
         model: this.options.chat,
         messages: [...contextMessages, ...messages],
         system: this.options.defaultSystemPrompt,
+        tools,
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        toolChoice: args.toolChoice as any,
         ...rest,
         onStepFinish: async (step) => {
-          if (chatId && messageId) {
+          if (chatId && messageId && args.saveOutputMessages) {
             await this.saveStep(ctx, {
               chatId,
               messageId,
@@ -357,7 +364,7 @@ export class Agent {
           return args.onStepFinish?.(step);
         },
       });
-      return result;
+      return { ...result, messageId };
     } catch (error) {
       if (chatId && messageId) {
         await ctx.runMutation(this.component.messages.rollbackMessage, {
@@ -374,13 +381,16 @@ export class Agent {
     OUTPUT = never,
     PARTIAL_OUTPUT = never,
   >(
-    ctx: RunMutationCtx,
-    { userId, chatId }: { userId?: string; chatId?: string },
+    ctx: RunActionCtx,
+    { userId, chatId }: { userId?: string; chatId: string },
     args: Partial<
       Parameters<typeof streamText<TOOLS, OUTPUT, PARTIAL_OUTPUT>>[0]
     > &
-      ContextOptions
-  ): Promise<StreamTextResult<TOOLS, PARTIAL_OUTPUT>> {
+      ContextOptions &
+      StorageOptions
+  ): 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,20 +399,21 @@ 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 defaults = this.options.tools;
+    const tools = wrapTools(ctx, chatId, userId, defaults, args.tools) as TOOLS;
+    const result = streamText({
       model: this.options.chat,
       messages: [...contextMessages, ...messages],
       system: this.options.defaultSystemPrompt,
+      tools,
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      toolChoice: args.toolChoice as any,
       ...rest,
       onChunk: async (chunk) => {
         console.log("onChunk", chunk);
@@ -436,16 +447,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"> & {
+    { userId, chatId }: { userId?: string; chatId: string },
+    args: Omit<Parameters<typeof generateObject>[0], "model"> & {
       model?: LanguageModelV1;
-    } & ContextOptions
-  ): Promise<GenerateObjectResult<OBJECT>> {
+    } & 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,20 +467,29 @@ 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>(
     ctx: RunMutationCtx,
-    { userId, chatId }: { userId?: string; chatId?: string },
+    { userId, chatId }: { userId?: string; chatId: string },
     args: Omit<Parameters<typeof streamObject<T>>[0], "model"> & {
       model?: LanguageModelV1;
-    } & ContextOptions
-  ) {
+    } & 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 +498,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,9 +515,10 @@ export class Agent {
         console.log("onFinish", result);
       },
     }) as StreamObjectResult<DeepPartial<T>, T, never>;
+    return { ...result, messageId };
   }
 
-  async searchWithDefaults(
+  async searchOptionsWithDefaults(
     searchArgs: ContextOptions,
     messages: CoreMessage[]
   ): Promise<SearchOptions> {
@@ -558,71 +586,8 @@ export class Agent {
       isDone: messages.isDone,
     };
   }
-
-  async getSteps(
-    ctx: RunQueryCtx,
-    args: { messageId: string }
-  ): Promise<{
-    steps: Step[];
-  }> {
-    return { steps: [] };
-  }
-}
-
-interface Chat {
-  generateText<TOOLS extends ToolSet, OUTPUT = never, OUTPUT_PARTIAL = never>(
-    args: Partial<
-      Parameters<typeof generateText<TOOLS, OUTPUT, OUTPUT_PARTIAL>>[0]
-    >
-  ): Promise<GenerateTextResult<TOOLS, OUTPUT>>;
-  streamText<TOOLS extends ToolSet, OUTPUT = never, PARTIAL_OUTPUT = never>(
-    args: Partial<
-      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>>;
-  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;
@@ -654,6 +619,135 @@ export function promptOrMessagesToCoreMessages(args: {
   return messages;
 }
 
+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;
+} & {
+  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: TextArgs<
+      AgentTools,
+      TOOLS,
+      Parameters<typeof generateText<TOOLS, OUTPUT, OUTPUT_PARTIAL>>[0]
+    >
+  ): Promise<
+    GenerateTextResult<TOOLS & AgentTools, OUTPUT> & GenerationOutputMetadata
+  >;
+
+  streamText<TOOLS extends ToolSet, OUTPUT = never, PARTIAL_OUTPUT = never>(
+    args: TextArgs<
+      AgentTools,
+      TOOLS,
+      Parameters<typeof streamText<TOOLS, OUTPUT, PARTIAL_OUTPUT>>[0]
+    >
+  ): 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: ObjectArgs<Parameters<typeof streamObject<T>>[0]>
+  ): Promise<
+    StreamObjectResult<DeepPartial<T>, T, never> & GenerationOutputMetadata
+  >;
+}
+
+// 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.
+ */
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export function tool<V extends Validator<any, any, any>, RESULT>(convexTool: {
+  args: V;
+  description?: string;
+  handler: (
+    ctx: GenericActionCtx<GenericDataModel> & {
+      userId?: string;
+      chatId?: string;
+    },
+    args: Infer<V>,
+    options: ToolExecutionOptions
+  ) => PromiseLike<RESULT>;
+  ctx?: GenericActionCtx<GenericDataModel> & {
+    userId?: string;
+    chatId?: string;
+  };
+}): Tool<ZodType<Infer<V>>, RESULT> {
+  const tool = {
+    __acceptUserIdAndChatId: true,
+    description: convexTool.description,
+    parameters: convexToZod(convexTool.args),
+    execute: async (args: Infer<V>, options: ToolExecutionOptions) => {
+      if (!convexTool.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(convexTool.ctx, args, options);
+    },
+  };
+  return tool;
+}
+
+export function wrapTools(
+  actionCtx: RunActionCtx,
+  chatId: string,
+  userId?: string,
+  ...toolSets: (ToolSet | undefined)[]
+): ToolSet {
+  const ctx = { ...actionCtx, chatId, userId };
+  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).__acceptUserIdAndChatId) {
+        output[name] = tool;
+      } else {
+        const out = { ...tool, ctx };
+        output[name] = out;
+      }
+    }
+  }
+  return output;
+}
 // export function convexValidatorSchema<T>(validator: Validator<unknown>)  {
 //   return ai.jsonSchema(convexToJsonSchema(validator));
 // }