@convex-dev/agent 0.5.0-alpha.1 → 0.6.0-alpha.0

package/src/client/search.ts

@@ -1,25 +1,55 @@
-import type {
-  AgentComponent,
-  ContextOptions,
-  RunActionCtx,
-  RunQueryCtx,
-} from "./types.js";
-import type { MessageDoc } from "../component/schema.js";
-import type { ModelMessage } from "ai";
+import {
+  embedMany as embedMany_,
+  type EmbeddingModel,
+  type ModelMessage,
+} from "ai";
 import { assert } from "convex-helpers";
+import type { MessageDoc } from "../validators.js";
+import {
+  validateVectorDimension,
+  type VectorDimension,
+} from "../component/vector/tables.js";
 import {
   DEFAULT_MESSAGE_RANGE,
   DEFAULT_RECENT_MESSAGES,
   extractText,
+  getModelName,
+  getProviderName,
+  isTool,
+  sorted,
 } from "../shared.js";
 import type { Message } from "../validators.js";
+import type {
+  AgentComponent,
+  Config,
+  ContextOptions,
+  Options,
+  QueryCtx,
+  MutationCtx,
+  ActionCtx,
+} from "./types.js";
+import { inlineMessagesFiles } from "./files.js";
+import { docsToModelMessages, toModelMessage } from "../mapping.js";
 
 const DEFAULT_VECTOR_SCORE_THRESHOLD = 0.0;
+// 10k characters should be more than enough for most cases, and stays under
+// the 8k token limit for some models.
+const MAX_EMBEDDING_TEXT_LENGTH = 10_000;
 
-export type GetEmbedding = (text: string) => Promise<{
-  embedding: number[];
-  embeddingModel: string;
-}>;
+export type GetEmbedding = (text: string) => Promise<
+  | {
+      embedding: number[];
+      /** @deprecated Use embeddingModel instead. */
+      textEmbeddingModel: string | EmbeddingModel;
+      embeddingModel?: string | EmbeddingModel;
+    }
+  | {
+      embedding: number[];
+      /** @deprecated Use embeddingModel instead. */
+      textEmbeddingModel?: string | EmbeddingModel;
+      embeddingModel: string | EmbeddingModel;
+    }
+>;
 
 /**
  * Fetch the context messages for a thread.
@@ -30,31 +60,78 @@ export type GetEmbedding = (text: string) => Promise<{
  * @returns
  */
 export async function fetchContextMessages(
-  ctx: RunQueryCtx | RunActionCtx,
+  ctx: QueryCtx | MutationCtx | ActionCtx,
   component: AgentComponent,
   args: {
     userId: string | undefined;
     threadId: string | undefined;
-    messages: (ModelMessage | Message)[];
     /**
-     * If provided, it will search for messages up to and including this message.
-     * Note: if this is far in the past, text and vector search results may be more
-     * limited, as it's post-filtering the results.
+     * If targetMessageId is not provided, this text will be used
+     * for text and vector search
+     */
+    searchText?: string;
+    /**
+     * If provided, it will use this message for text/vector search (if enabled)
+     * and will only fetch messages up to (and including) this message's "order"
+     */
+    targetMessageId?: string;
+    /**
+     * @deprecated use searchText and targetMessageId instead
+     */
+    messages?: (ModelMessage | Message)[];
+    /**
+     * @deprecated use targetMessageId instead
      */
     upToAndIncludingMessageId?: string;
     contextOptions: ContextOptions;
     getEmbedding?: GetEmbedding;
   },
 ): Promise<MessageDoc[]> {
+  const { recentMessages, searchMessages } = await fetchRecentAndSearchMessages(
+    ctx,
+    component,
+    args,
+  );
+  return [...searchMessages, ...recentMessages];
+}
+
+export async function fetchRecentAndSearchMessages(
+  ctx: QueryCtx | MutationCtx | ActionCtx,
+  component: AgentComponent,
+  args: {
+    userId: string | undefined;
+    threadId: string | undefined;
+    /**
+     * If targetMessageId is not provided, this text will be used
+     * for text and vector search
+     */
+    searchText?: string;
+    /**
+     * If provided, it will use this message for text/vector search (if enabled)
+     * and will only fetch messages up to (and including) this message's "order"
+     */
+    targetMessageId?: string;
+    /**
+     * @deprecated use searchText and targetMessageId instead
+     */
+    messages?: (ModelMessage | Message)[];
+    /**
+     * @deprecated use targetMessageId instead
+     */
+    upToAndIncludingMessageId?: string;
+    contextOptions: ContextOptions;
+    getEmbedding?: GetEmbedding;
+  },
+): Promise<{ recentMessages: MessageDoc[]; searchMessages: MessageDoc[] }> {
   assert(args.userId || args.threadId, "Specify userId or threadId");
   const opts = args.contextOptions;
   // Fetch the latest messages from the thread
   let included: Set<string> | undefined;
-  const contextMessages: MessageDoc[] = [];
-  if (
-    args.threadId &&
-    (opts.recentMessages !== 0 || args.upToAndIncludingMessageId)
-  ) {
+  let recentMessages: MessageDoc[] = [];
+  let searchMessages: MessageDoc[] = [];
+  const targetMessageId =
+    args.targetMessageId ?? args.upToAndIncludingMessageId;
+  if (args.threadId && opts.recentMessages !== 0) {
     const { page } = await ctx.runQuery(
       component.messages.listMessagesByThreadId,
       {
@@ -64,106 +141,187 @@ export async function fetchContextMessages(
           numItems: opts.recentMessages ?? DEFAULT_RECENT_MESSAGES,
           cursor: null,
         },
-        upToAndIncludingMessageId: args.upToAndIncludingMessageId,
+        upToAndIncludingMessageId: targetMessageId,
         order: "desc",
         statuses: ["success"],
       },
     );
     included = new Set(page.map((m) => m._id));
-    contextMessages.push(
-      // Reverse since we fetched in descending order
-      ...page.reverse(),
-    );
+    recentMessages = filterOutOrphanedToolMessages(sorted(page));
   }
-  if (opts.searchOptions?.textSearch || opts.searchOptions?.vectorSearch) {
-    const targetMessage = contextMessages.find(
-      (m) => m._id === args.upToAndIncludingMessageId,
-    )?.message;
-    const messagesToSearch = targetMessage ? [targetMessage] : args.messages;
+  if (
+    (opts.searchOptions?.textSearch || opts.searchOptions?.vectorSearch) &&
+    opts.searchOptions?.limit
+  ) {
     if (!("runAction" in ctx)) {
       throw new Error("searchUserMessages only works in an action");
     }
-    const lastMessage = messagesToSearch.at(-1)!;
-    assert(lastMessage, "No messages to search");
-    const text = extractText(lastMessage);
-    assert(text, `No text to search in message ${JSON.stringify(lastMessage)}`);
-    assert(
-      !args.contextOptions?.searchOptions?.vectorSearch || "runAction" in ctx,
-      "You must do vector search from an action",
-    );
-    if (opts.searchOptions?.vectorSearch && !args.getEmbedding) {
-      throw new Error(
-        "You must provide an embedding and embeddingModel to use vector search",
-      );
+    let text = args.searchText;
+    let embedding: number[] | undefined;
+    let embeddingModel: string | undefined;
+    if (!text) {
+      if (targetMessageId) {
+        const targetMessage = recentMessages.find(
+          (m) => m._id === targetMessageId,
+        );
+        if (targetMessage) {
+          text = targetMessage.text;
+        } else {
+          const targetSearchFields = await ctx.runQuery(
+            component.messages.getMessageSearchFields,
+            {
+              messageId: targetMessageId,
+            },
+          );
+          text = targetSearchFields.text;
+          embedding = targetSearchFields.embedding;
+          embeddingModel = targetSearchFields.embeddingModel;
+        }
+        assert(text, "Target message has no text for searching");
+      } else if (args.messages?.length) {
+        text = extractText(args.messages.at(-1)!);
+        assert(text, "Final context message has no text to search");
+      }
+      assert(text, "No text to search");
     }
-    const embeddingFields = opts.searchOptions?.vectorSearch
-      ? await args.getEmbedding?.(text)
-      : undefined;
-    const searchMessages = await ctx.runAction(
+    if (opts.searchOptions?.vectorSearch) {
+      if (!embedding && args.getEmbedding) {
+        const embeddingFields = await args.getEmbedding(text);
+        embedding = embeddingFields.embedding;
+        const effectiveModel =
+          embeddingFields.embeddingModel ?? embeddingFields.textEmbeddingModel;
+        embeddingModel = effectiveModel
+          ? getModelName(effectiveModel)
+          : undefined;
+        // TODO: if the text matches the target message, save the embedding
+        // for the target message and return the embeddingId on the message.
+      }
+    }
+    const searchResults = await ctx.runAction(
       component.messages.searchMessages,
       {
         searchAllMessagesForUserId: opts?.searchOtherThreads
-          ? args.userId ??
+          ? (args.userId ??
             (args.threadId &&
               (
                 await ctx.runQuery(component.threads.getThread, {
                   threadId: args.threadId,
                 })
-              )?.userId)
+              )?.userId))
           : undefined,
         threadId: args.threadId,
-        beforeMessageId: args.upToAndIncludingMessageId,
+        targetMessageId,
         limit: opts.searchOptions?.limit ?? 10,
         messageRange: {
           ...DEFAULT_MESSAGE_RANGE,
           ...opts.searchOptions?.messageRange,
         },
         text,
+        textSearch: opts.searchOptions?.textSearch,
+        vectorSearch: opts.searchOptions?.vectorSearch,
         vectorScoreThreshold:
           opts.searchOptions?.vectorScoreThreshold ??
           DEFAULT_VECTOR_SCORE_THRESHOLD,
-        embedding: embeddingFields?.embedding,
-        embeddingModel: embeddingFields?.embeddingModel,
+        embedding,
+        embeddingModel,
       },
     );
     // TODO: track what messages we used for context
-    contextMessages.unshift(
-      ...searchMessages.filter((m) => !included?.has(m._id)),
+    searchMessages = filterOutOrphanedToolMessages(
+      sorted(searchResults.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 { recentMessages, searchMessages };
 }
 
 /**
  * Filter out tool messages that don't have both a tool call and response.
+ * For the approval workflow, tool calls with approval responses (but no tool-results yet)
+ * should also be kept.
  * @param docs The messages to filter.
  * @returns The filtered messages.
  */
 export function filterOutOrphanedToolMessages(docs: MessageDoc[]) {
   const toolCallIds = new Set<string>();
+  const toolResultIds = new Set<string>();
+  // Track approval workflow: toolCallId → approvalId
+  const approvalRequestsByToolCallId = new Map<string, string>();
+  // Track which approvalIds have responses
+  const approvalResponseIds = new Set<string>();
+
   const result: MessageDoc[] = [];
   for (const doc of docs) {
-    if (
-      doc.message?.role === "assistant" &&
-      Array.isArray(doc.message.content)
-    ) {
+    if (doc.message && Array.isArray(doc.message.content)) {
       for (const content of doc.message.content) {
         if (content.type === "tool-call") {
           toolCallIds.add(content.toolCallId);
+        } else if (content.type === "tool-result") {
+          toolResultIds.add(content.toolCallId);
+        } else if (content.type === "tool-approval-request") {
+          const approvalRequest = content as {
+            type: "tool-approval-request";
+            toolCallId: string;
+            approvalId: string;
+          };
+          approvalRequestsByToolCallId.set(
+            approvalRequest.toolCallId,
+            approvalRequest.approvalId,
+          );
+        } else if (content.type === "tool-approval-response") {
+          const approvalResponse = content as {
+            type: "tool-approval-response";
+            approvalId: string;
+          };
+          approvalResponseIds.add(approvalResponse.approvalId);
         }
       }
-      result.push(doc);
+    }
+  }
+
+  // Helper: check if tool call has a corresponding approval response
+  const hasApprovalResponse = (toolCallId: string) => {
+    const approvalId = approvalRequestsByToolCallId.get(toolCallId);
+    return approvalId !== undefined && approvalResponseIds.has(approvalId);
+  };
+
+  for (const doc of docs) {
+    if (
+      doc.message?.role === "assistant" &&
+      Array.isArray(doc.message.content)
+    ) {
+      const content = doc.message.content.filter(
+        (p) =>
+          p.type !== "tool-call" ||
+          toolResultIds.has(p.toolCallId) ||
+          hasApprovalResponse(p.toolCallId),
+      );
+      if (content.length) {
+        result.push({
+          ...doc,
+          message: {
+            ...doc.message,
+            content,
+          },
+        });
+      }
     } 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);
+      const content = doc.message.content.filter((c) => {
+        // tool-result parts have toolCallId
+        if (c.type === "tool-result") {
+          return toolCallIds.has(c.toolCallId);
+        }
+        // tool-approval-response parts don't have toolCallId, so include them
+        return true;
+      });
+      if (content.length) {
+        result.push({
+          ...doc,
+          message: {
+            ...doc.message,
+            content,
+          },
+        });
       }
     } else {
       result.push(doc);
@@ -171,3 +329,352 @@ export function filterOutOrphanedToolMessages(docs: MessageDoc[]) {
   }
   return result;
 }
+
+/**
+ * Embed a list of messages, including calling any usage handler.
+ * This will not save the embeddings to the database.
+ */
+export async function embedMessages(
+  ctx: ActionCtx,
+  {
+    userId,
+    threadId,
+    ...options
+  }: {
+    userId: string | undefined;
+    threadId: string | undefined;
+    agentName?: string;
+  } & Pick<
+    Config,
+    "usageHandler" | "textEmbeddingModel" | "embeddingModel" | "callSettings"
+  >,
+  messages: (ModelMessage | Message)[],
+): Promise<
+  | {
+      vectors: (number[] | null)[];
+      dimension: VectorDimension;
+      model: string;
+    }
+  | undefined
+> {
+  const textEmbeddingModel =
+    options.embeddingModel ?? options.textEmbeddingModel;
+  if (!textEmbeddingModel) {
+    return undefined;
+  }
+  let embeddings:
+    | {
+        vectors: (number[] | null)[];
+        dimension: VectorDimension;
+        model: string;
+      }
+    | undefined;
+  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))
+    .filter((i) => i !== undefined);
+  if (textIndexes.length === 0) {
+    return undefined;
+  }
+  const values = messageTexts
+    .map((t) => t && t.trim().slice(0, MAX_EMBEDDING_TEXT_LENGTH))
+    .filter((t): t is string => !!t);
+  // Then embed those messages.
+  const textEmbeddings = await embedMany(ctx, {
+    ...options,
+    userId,
+    threadId,
+    values,
+  });
+  // 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) => {
+    embeddingsOrNull[i] = textEmbeddings.embeddings[j];
+  });
+  if (textEmbeddings.embeddings.length > 0) {
+    const dimension = textEmbeddings.embeddings[0].length;
+    validateVectorDimension(dimension);
+    const model = getModelName(textEmbeddingModel);
+    embeddings = { vectors: embeddingsOrNull, dimension, model };
+  }
+  return embeddings;
+}
+
+/**
+ * Embeds many strings, calling any usage handler.
+ * @param ctx The ctx parameter to an action.
+ * @param args Arguments to AI SDK's embedMany, and context for the embedding,
+ *   passed to the usage handler.
+ * @returns The embeddings for the strings, matching the order of the values.
+ */
+export async function embedMany(
+  ctx: ActionCtx,
+  args: {
+    userId: string | undefined;
+    threadId: string | undefined;
+    values: string[];
+    abortSignal?: AbortSignal;
+    headers?: Record<string, string>;
+    agentName?: string;
+  } & Pick<
+    Config,
+    "usageHandler" | "textEmbeddingModel" | "embeddingModel" | "callSettings"
+  >,
+): Promise<{ embeddings: number[][] }> {
+  const {
+    userId,
+    threadId,
+    values,
+    abortSignal,
+    headers,
+    agentName,
+    usageHandler,
+    textEmbeddingModel,
+    embeddingModel,
+    callSettings,
+  } = args;
+  const effectiveEmbeddingModel = embeddingModel ?? textEmbeddingModel;
+  assert(
+    effectiveEmbeddingModel,
+    "an embeddingModel (or textEmbeddingModel) is required to be set for vector search",
+  );
+  const result = await embedMany_({
+    ...callSettings,
+    model: effectiveEmbeddingModel,
+    values,
+    abortSignal,
+    headers,
+  });
+  if (usageHandler && result.usage) {
+    await usageHandler(ctx, {
+      userId,
+      threadId,
+      agentName,
+      model: getModelName(effectiveEmbeddingModel),
+      provider: getProviderName(effectiveEmbeddingModel),
+      providerMetadata: undefined,
+      usage: {
+        inputTokens: result.usage.tokens,
+        outputTokens: 0,
+        totalTokens: result.usage.tokens,
+        // These detail fields are required by LanguageModelUsage type but we don't
+        // have the granular data, so we provide objects with undefined values.
+        inputTokenDetails: {
+          cacheReadTokens: undefined,
+          cacheWriteTokens: undefined,
+          noCacheTokens: undefined,
+        },
+        outputTokenDetails: {
+          textTokens: undefined,
+          reasoningTokens: undefined,
+        },
+      },
+    });
+  }
+  return { embeddings: result.embeddings };
+}
+
+/**
+ * Embed a list of messages, and save the embeddings to the database.
+ * @param ctx The ctx parameter to an action.
+ * @param component The agent component, usually components.agent.
+ * @param args The context for the embedding, passed to the usage handler.
+ * @param messages The messages to embed, in the Agent MessageDoc format.
+ */
+export async function generateAndSaveEmbeddings(
+  ctx: ActionCtx,
+  component: AgentComponent,
+  args: {
+    threadId: string | undefined;
+    userId: string | undefined;
+    agentName?: string;
+    /**
+     * @deprecated Use embeddingModel instead.
+     */
+    textEmbeddingModel?: EmbeddingModel;
+    embeddingModel?: EmbeddingModel;
+  } & Pick<Config, "usageHandler" | "callSettings">,
+  messages: MessageDoc[],
+) {
+  const effectiveEmbeddingModel =
+    args.embeddingModel ?? args.textEmbeddingModel;
+  if (!effectiveEmbeddingModel) {
+    throw new Error(
+      "an embeddingModel (or textEmbeddingModel) is required to generate and save embeddings",
+    );
+  }
+  const toEmbed = messages.filter((m) => !m.embeddingId && m.message);
+  if (toEmbed.length === 0) {
+    return;
+  }
+  const embeddings = await embedMessages(
+    ctx,
+    { ...args, embeddingModel: effectiveEmbeddingModel },
+    toEmbed.map((m) => m.message!),
+  );
+  if (embeddings && embeddings.vectors.some((v) => v !== null)) {
+    await ctx.runMutation(component.vector.index.insertBatch, {
+      vectorDimension: embeddings.dimension,
+      vectors: toEmbed
+        .map((m, i) => ({
+          messageId: m._id,
+          model: embeddings.model,
+          table: "messages",
+          userId: m.userId,
+          threadId: m.threadId,
+          vector: embeddings.vectors[i]!,
+        }))
+        .filter((v) => v.vector !== null),
+    });
+  }
+}
+
+/**
+ * Similar to fetchContextMessages, but also combines the input messages,
+ * with search context, recent messages, input messages, then prompt messages.
+ * If there is a promptMessageId and prompt message(s) provided, it will splice
+ * the prompt messages into the history to replace the promptMessageId message,
+ * but still be followed by any existing messages that were in response to the
+ * promptMessageId message.
+ */
+export async function fetchContextWithPrompt(
+  ctx: ActionCtx,
+  component: AgentComponent,
+  args: {
+    prompt: string | (ModelMessage | Message)[] | undefined;
+    messages: (ModelMessage | Message)[] | undefined;
+    promptMessageId: string | undefined;
+    userId: string | undefined;
+    threadId: string | undefined;
+    agentName?: string;
+  } & Options &
+    Config,
+): Promise<{
+  messages: ModelMessage[];
+  order: number | undefined;
+  stepOrder: number | undefined;
+}> {
+  const { threadId, userId, textEmbeddingModel, embeddingModel } = args;
+  const effectiveEmbeddingModel = embeddingModel ?? textEmbeddingModel;
+
+  const promptArray = getPromptArray(args.prompt);
+
+  const searchText = promptArray.length
+    ? extractText(promptArray.at(-1)!)
+    : args.promptMessageId
+      ? undefined
+      : args.messages?.at(-1)
+        ? extractText(args.messages.at(-1)!)
+        : undefined;
+  // If only a messageId is provided, this will add that message to the end.
+  const { recentMessages, searchMessages } = await fetchRecentAndSearchMessages(
+    ctx,
+    component,
+    {
+      userId,
+      threadId,
+      targetMessageId: args.promptMessageId,
+      searchText,
+      contextOptions: args.contextOptions ?? {},
+      getEmbedding: async (text) => {
+        assert(
+          effectiveEmbeddingModel,
+          "An embeddingModel (or textEmbeddingModel) is required to be set on the Agent that you're doing vector search with",
+        );
+        return {
+          embedding: (
+            await embedMany(ctx, {
+              ...args,
+              userId,
+              values: [text],
+              embeddingModel: effectiveEmbeddingModel,
+            })
+          ).embeddings[0],
+          embeddingModel: effectiveEmbeddingModel,
+        };
+      },
+    },
+  );
+
+  const promptMessageIndex = args.promptMessageId
+    ? recentMessages.findIndex((m) => m._id === args.promptMessageId)
+    : -1;
+  const promptMessage =
+    promptMessageIndex !== -1 ? recentMessages[promptMessageIndex] : undefined;
+  let prePromptDocs = recentMessages;
+  const messages = args.messages ?? [];
+  let existingResponseDocs: MessageDoc[] = [];
+  if (promptMessage) {
+    prePromptDocs = recentMessages.slice(0, promptMessageIndex);
+    existingResponseDocs = recentMessages.slice(promptMessageIndex + 1);
+    if (promptArray.length === 0) {
+      // If they didn't override the prompt, use the existing prompt message.
+      if (promptMessage.message) {
+        promptArray.push(promptMessage.message);
+      }
+    }
+    if (!promptMessage.embeddingId && effectiveEmbeddingModel) {
+      // Lazily generate embeddings for the prompt message, if it doesn't have
+      // embeddings yet. This can happen if the message was saved in a mutation
+      // where the LLM is not available.
+      await generateAndSaveEmbeddings(
+        ctx,
+        component,
+        {
+          ...args,
+          userId,
+          embeddingModel: effectiveEmbeddingModel,
+        },
+        [promptMessage],
+      );
+    }
+  }
+
+  const search = docsToModelMessages(searchMessages);
+  const recent = docsToModelMessages(prePromptDocs);
+  const inputMessages = messages.map(toModelMessage);
+  const inputPrompt = promptArray.map(toModelMessage);
+  const existingResponses = docsToModelMessages(existingResponseDocs);
+
+  const allMessages = [
+    ...search,
+    ...recent,
+    ...inputMessages,
+    ...inputPrompt,
+    ...existingResponses,
+  ];
+  let processedMessages = args.contextHandler
+    ? await args.contextHandler(ctx, {
+        allMessages,
+        search,
+        recent,
+        inputMessages,
+        inputPrompt,
+        existingResponses,
+        userId,
+        threadId,
+      })
+    : allMessages;
+
+  // Process messages to inline localhost files (if not, file urls pointing to localhost will be sent to LLM providers)
+  if (process.env.CONVEX_CLOUD_URL?.startsWith("http://127.0.0.1")) {
+    processedMessages = await inlineMessagesFiles(processedMessages);
+  }
+
+  return {
+    messages: processedMessages,
+    order: promptMessage?.order,
+    stepOrder: promptMessage?.stepOrder,
+  };
+}
+
+export function getPromptArray(
+  prompt: string | (ModelMessage | Message)[] | undefined,
+): (ModelMessage | Message)[] {
+  return !prompt
+    ? []
+    : Array.isArray(prompt)
+      ? prompt
+      : [{ role: "user", content: prompt }];
+}