@convex-dev/agent 0.1.14 → 0.1.15-alpha.1

package/src/client/index.ts

@@ -34,22 +34,17 @@ import {
   serializeNewMessagesInStep,
   serializeObjectResult,
 } from "../mapping.js";
-import {
-  DEFAULT_MESSAGE_RANGE,
-  DEFAULT_RECENT_MESSAGES,
-  extractText,
-  isTool,
-} from "../shared.js";
+import { extractText, isTool } from "../shared.js";
 import {
   type MessageWithMetadata,
   type MessageStatus,
   type ProviderMetadata,
-  type SearchOptions,
   type StreamArgs,
   type Usage,
   vMessageWithMetadata,
   vSafeObjectArgs,
   vTextArgs,
+  type MessageEmbeddings,
 } from "../validators.js";
 import { createTool, wrapTools } from "./createTool.js";
 import {
@@ -77,6 +72,9 @@ import type {
   UsageHandler,
 } from "./types.js";
 import type { threadFieldsSupportingPatch } from "../component/threads.js";
+import { listMessages } from "./listMessages.js";
+import { syncStreams } from "./streaming.js";
+import { fetchContextMessages } from "./search.js";
 
 export { storeFile, getFile } from "./files.js";
 export { serializeDataOrUrl } from "../mapping.js";
@@ -95,12 +93,21 @@ export {
   vUserMessage,
 } from "../validators.js";
 export type { ToolCtx } from "./createTool.js";
-export { createTool, extractText, isTool };
+export { filterOutOrphanedToolMessages } from "./search.js";
+export {
+  createTool,
+  extractText,
+  fetchContextMessages,
+  isTool,
+  listMessages,
+  syncStreams,
+};
 export type {
   AgentComponent,
   ContextOptions,
   MessageDoc,
   ProviderMetadata,
+  RawRequestResponseHandler,
   StorageOptions,
   StreamArgs,
   SyncStreamsReturnValue,
@@ -110,7 +117,7 @@ export type {
   UsageHandler,
 };
 
-export class Agent<AgentTools extends ToolSet> {
+export class Agent<AgentTools extends ToolSet = ToolSet> {
   constructor(
     public component: AgentComponent,
     public options: {
@@ -276,25 +283,18 @@ export class Agent<AgentTools extends ToolSet> {
     threadId: string;
     thread?: Thread<ThreadTools extends undefined ? AgentTools : ThreadTools>;
   }> {
-    const threadDoc = await ctx.runMutation(
-      this.component.threads.createThread,
-      {
-        userId: args?.userId,
-        title: args?.title,
-        summary: args?.summary,
-      }
-    );
+    const threadId = await createThread(ctx, this.component, args);
     if (!("runAction" in ctx)) {
-      return { threadId: threadDoc._id };
+      return { threadId };
     }
     const { thread } = await this.continueThread(ctx, {
-      threadId: threadDoc._id,
+      threadId,
       userId: args?.userId,
       usageHandler: args?.usageHandler,
       tools: args?.tools,
     });
     return {
-      threadId: threadDoc._id,
+      threadId,
       thread,
     };
   }
@@ -598,6 +598,7 @@ export class Agent<AgentTools extends ToolSet> {
             error: (error.error as Error).message,
           });
         }
+        // TODO: update the streamer to error state
         return args.onError?.(error);
       },
       onStepFinish: async (step) => {
@@ -611,7 +612,6 @@ export class Agent<AgentTools extends ToolSet> {
             promptMessageId: messageId,
             step,
           });
-          // TODO: figure out pending/not
           await streamer?.finish(saved.messages);
         }
         if (this.options.rawRequestResponseHandler) {
@@ -833,14 +833,7 @@ export class Agent<AgentTools extends ToolSet> {
    */
   async saveMessage(
     ctx: RunMutationCtx,
-    args: {
-      threadId: string;
-      userId?: string;
-      /**
-       * Metadata to save with the messages. Each element corresponds to the
-       * message at the same index.
-       */
-      metadata?: Omit<MessageWithMetadata, "message">;
+    args: SaveMessageArgs & {
       /**
        * If true, it will not generate embeddings for the message.
        * Useful if you're saving messages in a mutation where you can't run `fetch`.
@@ -848,22 +841,7 @@ export class Agent<AgentTools extends ToolSet> {
        * action later that calls `agent.generateAndSaveEmbeddings`.
        */
       skipEmbeddings?: boolean;
-    } & (
-      | {
-          prompt?: undefined;
-          /**
-           * The message to save.
-           */
-          message: CoreMessage;
-        }
-      | {
-          /*
-           * The prompt to save with the message.
-           */
-          prompt: string;
-          message?: undefined;
-        }
-    )
+    }
   ) {
     const { lastMessageId, messages } = await this.saveMessages(ctx, {
       threadId: args.threadId,
@@ -880,40 +858,15 @@ export class Agent<AgentTools extends ToolSet> {
 
   /**
    * Explicitly save messages associated with the thread (& user if provided)
+   * If you have an embedding model set, it will also generate embeddings for
+   * the messages.
    * @param ctx The ctx parameter to a mutation or action.
    * @param args The messages and context to save
    * @returns
    */
   async saveMessages(
     ctx: RunMutationCtx | RunActionCtx,
-    args: {
-      threadId: string;
-      userId?: string;
-      /**
-       * The message that these messages are in response to. They will be
-       * the same "order" as this message, at increasing stepOrder(s).
-       */
-      promptMessageId?: string;
-      /**
-       * The messages to save.
-       */
-      messages: CoreMessageMaybeWithId[];
-      /**
-       * 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;
-      /**
-       * If true, it will fail any pending steps.
-       * Defaults to false.
-       */
-      failPendingSteps?: boolean;
+    args: SaveMessagesArgs & {
       /**
        * Skip generating embeddings for the messages. Useful if you're
        * saving messages in a mutation where you can't run `fetch`.
@@ -933,9 +886,12 @@ export class Agent<AgentTools extends ToolSet> {
           model: string;
         }
       | undefined;
-    if (args.skipEmbeddings || !("runAction" in ctx)) {
+    const { skipEmbeddings, ...rest } = args;
+    if (args.embeddings) {
+      embeddings = args.embeddings;
+    } else if (skipEmbeddings || !("runAction" in ctx)) {
       embeddings = undefined;
-      if (!args.skipEmbeddings && this.options.textEmbedding) {
+      if (!skipEmbeddings && this.options.textEmbedding) {
         console.warn(
           "You're trying to save messages and generate embeddings, but you're in a mutation. " +
             "Pass `skipEmbeddings: true` to skip generating embeddings in the mutation and skip this warning. " +
@@ -953,33 +909,11 @@ export class Agent<AgentTools extends ToolSet> {
         args.messages
       );
     }
-    const result = await ctx.runMutation(this.component.messages.addMessages, {
-      threadId: args.threadId,
-      userId: args.userId,
+    return saveMessages(ctx, this.component, {
+      ...rest,
       agentName: this.options.name,
-      promptMessageId: args.promptMessageId,
       embeddings,
-      messages: await Promise.all(
-        args.messages.map(async (m, i) => {
-          const { message, fileIds } = await serializeMessage(
-            ctx,
-            this.component,
-            m
-          );
-          return {
-            ...args.metadata?.[i],
-            message,
-            fileIds,
-          } as MessageWithMetadata;
-        })
-      ),
-      failPendingSteps: args.failPendingSteps ?? false,
-      pending: args.pending ?? false,
     });
-    return {
-      lastMessageId: result.messages.at(-1)!._id,
-      messages: result.messages,
-    };
   }
 
   /**
@@ -1001,17 +935,7 @@ export class Agent<AgentTools extends ToolSet> {
       statuses?: MessageStatus[];
     }
   ): Promise<PaginationResult<MessageDoc>> {
-    if (args.paginationOpts.numItems === 0) {
-      return {
-        page: [],
-        isDone: true,
-        continueCursor: args.paginationOpts.cursor ?? "",
-      };
-    }
-    return ctx.runQuery(this.component.messages.listMessagesByThreadId, {
-      order: "desc",
-      ...args,
-    });
+    return listMessages(ctx, this.component, args);
   }
 
   /**
@@ -1027,25 +951,11 @@ export class Agent<AgentTools extends ToolSet> {
     args: {
       threadId: string;
       streamArgs: StreamArgs | undefined;
+      // By default, only streaming messages are included.
+      includeStatuses?: ("streaming" | "finished" | "aborted")[];
     }
   ): Promise<SyncStreamsReturnValue | undefined> {
-    if (!args.streamArgs) return undefined;
-    if (args.streamArgs.kind === "list") {
-      return {
-        kind: "list",
-        messages: await ctx.runQuery(this.component.streams.list, {
-          threadId: args.threadId,
-        }),
-      };
-    } else {
-      return {
-        kind: "deltas",
-        deltas: await ctx.runQuery(this.component.streams.listDeltas, {
-          threadId: args.threadId,
-          cursors: args.streamArgs.cursors,
-        }),
-      };
-    }
+    return syncStreams(ctx, this.component, args);
   }
 
   /**
@@ -1072,76 +982,28 @@ export class Agent<AgentTools extends ToolSet> {
     }
   ): Promise<MessageDoc[]> {
     assert(args.userId || args.threadId, "Specify userId or threadId");
-    // Fetch the latest messages from the thread
-    let included: Set<string> | undefined;
     const opts = this._mergedContextOptions(args.contextOptions);
-    const contextMessages: MessageDoc[] = [];
-    if (
-      args.threadId &&
-      (opts.recentMessages !== 0 || args.upToAndIncludingMessageId)
-    ) {
-      const { page } = await ctx.runQuery(
-        this.component.messages.listMessagesByThreadId,
-        {
-          threadId: args.threadId,
-          excludeToolMessages: opts.excludeToolMessages,
-          paginationOpts: {
-            numItems: opts.recentMessages ?? DEFAULT_RECENT_MESSAGES,
-            cursor: null,
-          },
-          upToAndIncludingMessageId: args.upToAndIncludingMessageId,
-          order: "desc",
-          statuses: ["success"],
-        }
-      );
-      included = new Set(page.map((m) => m._id));
-      contextMessages.push(
-        // Reverse since we fetched in descending order
-        ...page.reverse()
-      );
-    }
-    if (opts.searchOptions?.textSearch || opts.searchOptions?.vectorSearch) {
-      const targetMessage = contextMessages.find(
-        (m) => m._id === args.upToAndIncludingMessageId
-      )?.message;
-      const messagesToSearch = targetMessage ? [targetMessage] : args.messages;
-      if (!("runAction" in ctx)) {
-        throw new Error("searchUserMessages only works in an action");
-      }
-      const searchMessages = await ctx.runAction(
-        this.component.messages.searchMessages,
-        {
-          searchAllMessagesForUserId: opts?.searchOtherThreads
-            ? args.userId ??
-              (args.threadId &&
-                (
-                  await ctx.runQuery(this.component.threads.getThread, {
-                    threadId: args.threadId,
-                  })
-                )?.userId)
-            : undefined,
-          threadId: args.threadId,
-          beforeMessageId: args.upToAndIncludingMessageId,
-          ...(await this._searchOptionsWithEmbeddingAndDefaults(
-            ctx,
-            { userId: args.userId, threadId: args.threadId },
-            opts,
-            messagesToSearch
-          )),
-        }
-      );
-      // TODO: track what messages we used for context
-      contextMessages.unshift(
-        ...searchMessages.filter((m) => !included?.has(m._id))
-      );
-    }
-    // Ensure we don't include tool messages without a corresponding tool call
-    return filterOutOrphanedToolMessages(
-      contextMessages.sort((a, b) =>
-        // Sort the raw MessageDocs by order and stepOrder
-        a.order === b.order ? a.stepOrder - b.stepOrder : a.order - b.order
-      )
-    );
+    return fetchContextMessages(ctx, this.component, {
+      ...args,
+      contextOptions: opts,
+      getEmbedding: async (text) => {
+        assert("runAction" in ctx);
+        assert(
+          this.options.textEmbedding,
+          "A textEmbedding model is required to be set on the Agent that you're doing vector search with"
+        );
+        return {
+          vector: (
+            await this.doEmbed(ctx, {
+              userId: args.userId,
+              threadId: args.threadId,
+              values: [text],
+            })
+          ).embeddings[0],
+          vectorModel: this.options.textEmbedding.modelId,
+        };
+      },
+    });
   }
 
   /**
@@ -1154,13 +1016,7 @@ export class Agent<AgentTools extends ToolSet> {
     ctx: RunQueryCtx,
     args: { threadId: string }
   ): Promise<ThreadDoc> {
-    const thread = await ctx.runQuery(this.component.threads.getThread, {
-      threadId: args.threadId,
-    });
-    if (!thread) {
-      throw new Error("Thread not found");
-    }
-    return thread;
+    return getThreadMetadata(ctx, this.component, args);
   }
 
   /**
@@ -1765,49 +1621,11 @@ export class Agent<AgentTools extends ToolSet> {
       ...this.options.contextOptions,
       ...opts,
       searchOptions: searchOptions.limit
-        ? (searchOptions as SearchOptions)
+        ? (searchOptions as ContextOptions["searchOptions"])
         : undefined,
     };
   }
 
-  async _searchOptionsWithEmbeddingAndDefaults(
-    ctx: RunActionCtx,
-    { userId, threadId }: { userId?: string; threadId?: string },
-    contextOptions: ContextOptions,
-    messages: CoreMessage[]
-  ): Promise<SearchOptions> {
-    assert(
-      contextOptions.searchOptions?.textSearch ||
-        contextOptions.searchOptions?.vectorSearch,
-      "searchOptions is required"
-    );
-    assert(messages.length > 0, "Core messages cannot be empty");
-    const text = extractText(messages.at(-1)!);
-    const search: SearchOptions = {
-      limit: contextOptions.searchOptions?.limit ?? 10,
-      messageRange: {
-        ...DEFAULT_MESSAGE_RANGE,
-        ...contextOptions.searchOptions?.messageRange,
-      },
-      text: extractText(messages.at(-1)!),
-    };
-    if (
-      contextOptions.searchOptions?.vectorSearch &&
-      text &&
-      this.options.textEmbedding
-    ) {
-      search.vector = (
-        await this.doEmbed(ctx, {
-          threadId,
-          userId,
-          values: [text],
-        })
-      ).embeddings[0];
-      search.vectorModel = this.options.textEmbedding.modelId;
-    }
-    return search;
-  }
-
   async doEmbed(
     ctx: RunActionCtx,
     options: {
@@ -1819,7 +1637,10 @@ export class Agent<AgentTools extends ToolSet> {
     }
   ): Promise<{ embeddings: number[][] }> {
     const embedding = this.options.textEmbedding;
-    assert(embedding, "textEmbedding is required");
+    assert(
+      embedding,
+      "a textEmbedding model is required to be set on the Agent that you're doing vector search with"
+    );
     const result = await embedding.doEmbed({
       values: options.values,
       abortSignal: options.abortSignal,
@@ -2120,31 +1941,191 @@ export class Agent<AgentTools extends ToolSet> {
   }
 }
 
-export function filterOutOrphanedToolMessages(docs: MessageDoc[]) {
-  const toolCallIds = new Set<string>();
-  const result: MessageDoc[] = [];
-  for (const doc of docs) {
-    if (
-      doc.message?.role === "assistant" &&
-      Array.isArray(doc.message.content)
-    ) {
-      for (const content of doc.message.content) {
-        if (content.type === "tool-call") {
-          toolCallIds.add(content.toolCallId);
-        }
-      }
-      result.push(doc);
-    } else if (doc.message?.role === "tool") {
-      if (doc.message.content.every((c) => toolCallIds.has(c.toolCallId))) {
-        result.push(doc);
-      } else {
-        console.debug("Filtering out orphaned tool message", doc);
-      }
-    } else {
-      result.push(doc);
-    }
+type CoreMessageMaybeWithId = CoreMessage & { id?: string | undefined };
+
+/**
+ * Create a thread to store messages with an Agent.
+ * @param ctx The context from a mutation or action.
+ * @param component The Agent component, usually `components.agent`.
+ * @param args The associated thread metadata.
+ * @returns The id of the created thread.
+ */
+export async function createThread(
+  ctx: RunMutationCtx,
+  component: AgentComponent,
+  args?: {
+    userId?: string;
+    title?: string;
+    summary?: string;
   }
-  return result;
+) {
+  const { _id: threadId } = await ctx.runMutation(
+    component.threads.createThread,
+    { userId: args?.userId, title: args?.title, summary: args?.summary }
+  );
+  return threadId;
 }
 
-type CoreMessageMaybeWithId = CoreMessage & { id?: string | undefined };
+/**
+ * Get the metadata for a thread.
+ * @param ctx A ctx object from a query, mutation, or action.
+ * @param args.threadId The thread to get the metadata for.
+ * @returns The metadata for the thread.
+ */
+export async function getThreadMetadata(
+  ctx: RunQueryCtx,
+  component: AgentComponent,
+  args: { threadId: string }
+): Promise<ThreadDoc> {
+  const thread = await ctx.runQuery(component.threads.getThread, {
+    threadId: args.threadId,
+  });
+  if (!thread) {
+    throw new Error("Thread not found");
+  }
+  return thread;
+}
+
+type SaveMessagesArgs = {
+  threadId: string;
+  userId?: string;
+  /**
+   * The message that these messages are in response to. They will be
+   * the same "order" as this message, at increasing stepOrder(s).
+   */
+  promptMessageId?: string;
+  /**
+   * The messages to save.
+   */
+  messages: CoreMessageMaybeWithId[];
+  /**
+   * 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;
+  /**
+   * If true, it will fail any pending steps.
+   * Defaults to false.
+   */
+  failPendingSteps?: boolean;
+  /**
+   * The embeddings to save with the messages.
+   */
+  embeddings?: MessageEmbeddings;
+};
+
+/**
+ * Explicitly save messages associated with the thread (& user if provided)
+ */
+export async function saveMessages(
+  ctx: RunMutationCtx,
+  component: AgentComponent,
+  args: SaveMessagesArgs & {
+    /**
+     * The agent name to associate with the messages.
+     */
+    agentName?: string;
+  }
+) {
+  const result = await ctx.runMutation(component.messages.addMessages, {
+    threadId: args.threadId,
+    userId: args.userId,
+    agentName: args.agentName,
+    promptMessageId: args.promptMessageId,
+    embeddings: args.embeddings,
+    messages: await Promise.all(
+      args.messages.map(async (m, i) => {
+        const { message, fileIds } = await serializeMessage(ctx, component, m);
+        return {
+          ...args.metadata?.[i],
+          message,
+          fileIds,
+        } as MessageWithMetadata;
+      })
+    ),
+    failPendingSteps: args.failPendingSteps ?? false,
+    pending: args.pending ?? false,
+  });
+  return {
+    lastMessageId: result.messages.at(-1)!._id,
+    messages: result.messages,
+  };
+}
+
+type SaveMessageArgs = {
+  threadId: string;
+  userId?: string;
+  /**
+   * Metadata to save with the messages. Each element corresponds to the
+   * message at the same index.
+   */
+  metadata?: Omit<MessageWithMetadata, "message">;
+  /**
+   * The embedding to save with the message.
+   */
+  embedding?: {
+    vector: number[];
+    model: string;
+  };
+} & (
+  | {
+      prompt?: undefined;
+      /**
+       * The message to save.
+       */
+      message: CoreMessage;
+    }
+  | {
+      /*
+       * The prompt to save with the message.
+       */
+      prompt: string;
+      message?: undefined;
+    }
+);
+
+/**
+ * Save a message to the thread.
+ * @param ctx A ctx object from a mutation or action.
+ * @param args The message and what to associate it with (user / thread)
+ * You can pass extra metadata alongside the message, e.g. associated fileIds.
+ * @returns The messageId of the saved message.
+ */
+export async function saveMessage(
+  ctx: RunMutationCtx,
+  component: AgentComponent,
+  args: SaveMessageArgs & {
+    /**
+     * The agent name to associate with the message.
+     */
+    agentName?: string;
+  }
+) {
+  let embeddings: MessageEmbeddings | undefined;
+  if (args.embedding) {
+    const dimension = args.embedding.vector.length;
+    validateVectorDimension(dimension);
+    embeddings = {
+      model: args.embedding.model,
+      dimension,
+      vectors: [args.embedding.vector],
+    };
+  }
+  const { lastMessageId, messages } = await saveMessages(ctx, component, {
+    threadId: args.threadId,
+    userId: args.userId,
+    messages:
+      args.prompt !== undefined
+        ? [{ role: "user", content: args.prompt }]
+        : [args.message],
+    metadata: args.metadata ? [args.metadata] : undefined,
+    embeddings,
+  });
+  return { messageId: lastMessageId, message: messages.at(-1)! };
+}