@convex-dev/agent 0.2.6-alpha.0 → 0.2.6

package/src/client/search.ts

@@ -1,25 +1,41 @@
-import type {
-  AgentComponent,
-  ContextOptions,
-  RunActionCtx,
-  RunQueryCtx,
-} from "./types.js";
-import type { MessageDoc } from "../component/schema.js";
-import type { EmbeddingModel, LanguageModel, 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,
+  RunActionCtx,
+  RunQueryCtx,
+} from "./types.js";
+import { inlineMessagesFiles } from "./files.js";
+import { deserializeMessage } 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<{
+export type GetEmbedding = (text: string) => Promise<{
   embedding: number[];
   textEmbeddingModel: string | EmbeddingModel<string>;
 }>;
@@ -38,26 +54,73 @@ export async function fetchContextMessages(
   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: RunQueryCtx | RunActionCtx,
+  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,
       {
@@ -67,46 +130,59 @@ 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) &&
     opts.searchOptions?.limit
   ) {
-    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 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 && text
-        ? 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;
+        embeddingModel = embeddingFields.textEmbeddingModel
+          ? getModelName(embeddingFields.textEmbeddingModel)
+          : undefined;
+      }
+    }
+    const searchResults = await ctx.runAction(
       component.messages.searchMessages,
       {
         searchAllMessagesForUserId: opts?.searchOtherThreads
@@ -119,29 +195,29 @@ export async function fetchContextMessages(
               )?.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?.textEmbeddingModel
-          ? getModelName(embeddingFields.textEmbeddingModel)
-          : undefined,
+        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(sorted(contextMessages));
+  return { recentMessages, searchMessages };
 }
 
 /**
@@ -176,23 +252,325 @@ export function filterOutOrphanedToolMessages(docs: MessageDoc[]) {
   return result;
 }
 
-export function getModelName(
-  embeddingModel: string | EmbeddingModel<string> | LanguageModel,
-): string {
-  if (typeof embeddingModel === "string") {
-    if (embeddingModel.includes("/")) {
-      return embeddingModel.split("/").slice(1).join("/");
+/**
+ * Embed a list of messages, including calling any usage handler.
+ * This will not save the embeddings to the database.
+ */
+export async function embedMessages(
+  ctx: RunActionCtx,
+  {
+    userId,
+    threadId,
+    ...options
+  }: {
+    userId: string | undefined;
+    threadId: string | undefined;
+    agentName?: string;
+  } & Pick<Config, "usageHandler" | "textEmbeddingModel" | "callSettings">,
+  messages: (ModelMessage | Message)[],
+): Promise<
+  | {
+      vectors: (number[] | null)[];
+      dimension: VectorDimension;
+      model: string;
     }
-    return embeddingModel;
+  | undefined
+> {
+  if (!options.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(options.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: RunActionCtx,
+  {
+    userId,
+    threadId,
+    values,
+    abortSignal,
+    headers,
+    agentName,
+    usageHandler,
+    textEmbeddingModel,
+    callSettings,
+  }: {
+    userId: string | undefined;
+    threadId: string | undefined;
+    values: string[];
+    abortSignal?: AbortSignal;
+    headers?: Record<string, string>;
+    agentName?: string;
+  } & Pick<Config, "usageHandler" | "textEmbeddingModel" | "callSettings">,
+): Promise<{ embeddings: number[][] }> {
+  const embeddingModel = textEmbeddingModel;
+  assert(
+    embeddingModel,
+    "a textEmbeddingModel is required to be set for vector search",
+  );
+  const result = await embedMany_({
+    ...callSettings,
+    model: embeddingModel,
+    values,
+    abortSignal,
+    headers,
+  });
+  if (usageHandler && result.usage) {
+    await usageHandler(ctx, {
+      userId,
+      threadId,
+      agentName,
+      model: getModelName(embeddingModel),
+      provider: getProviderName(embeddingModel),
+      providerMetadata: undefined,
+      usage: {
+        inputTokens: result.usage.tokens,
+        outputTokens: 0,
+        totalTokens: result.usage.tokens,
+      },
+    });
+  }
+  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: RunActionCtx,
+  component: AgentComponent,
+  args: {
+    threadId: string | undefined;
+    userId: string | undefined;
+    agentName?: string;
+    textEmbeddingModel: EmbeddingModel<string>;
+  } & Pick<Config, "usageHandler" | "callSettings">,
+  messages: MessageDoc[],
+) {
+  const toEmbed = messages.filter((m) => !m.embeddingId && m.message);
+  if (toEmbed.length === 0) {
+    return;
+  }
+  const embeddings = await embedMessages(
+    ctx,
+    args,
+    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),
+    });
   }
-  return embeddingModel.modelId;
 }
 
-export function getProviderName(
-  embeddingModel: string | EmbeddingModel<string> | LanguageModel,
-): string {
-  if (typeof embeddingModel === "string") {
-    return embeddingModel.split("/").at(0)!;
+/**
+ * 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: RunActionCtx,
+  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 } = args;
+
+  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(
+          textEmbeddingModel,
+          "A 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],
+              textEmbeddingModel,
+            })
+          ).embeddings[0],
+          textEmbeddingModel,
+        };
+      },
+    },
+  );
+
+  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 && textEmbeddingModel) {
+      // 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,
+          textEmbeddingModel,
+        },
+        [promptMessage],
+      );
+    }
+  }
+
+  const search = searchMessages
+    .map((m) => m.message)
+    .filter((m) => !!m)
+    .map(deserializeMessage);
+  const recent = prePromptDocs
+    .map((m) => m.message)
+    .filter((m) => !!m)
+    .map(deserializeMessage);
+  const inputMessages = messages.map(deserializeMessage);
+  const inputPrompt = promptArray.map(deserializeMessage);
+  const existingResponses = existingResponseDocs
+    .map((m) => m.message)
+    .filter((m) => !!m)
+    .map(deserializeMessage);
+
+  let processedMessages = args.contextHandler
+    ? await args.contextHandler(ctx, {
+        search,
+        recent,
+        inputMessages,
+        inputPrompt,
+        existingResponses,
+        userId,
+        threadId,
+      })
+    : [
+        ...search,
+        ...recent,
+        ...inputMessages,
+        ...inputPrompt,
+        ...existingResponses,
+      ];
+
+  // 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 embeddingModel.provider;
+
+  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 }];
 }