@convex-dev/agent 0.0.1-alpha.4 → 0.0.2

package/src/component/messages.ts

@@ -5,11 +5,12 @@ import { nullable, partial } from "convex-helpers/validators";
 import { ObjectType } from "convex/values";
 import { DEFAULT_MESSAGE_RANGE, extractText, isTool } from "../shared.js";
 import {
-  vThreadStatus,
+  vEmbeddingsWithMetadata,
   vMessageStatus,
   vMessageWithFileAndId,
   vSearchOptions,
-  vStepWithMessagesWithFileAndId,
+  vStepWithMessages,
+  vThreadStatus,
 } from "../validators.js";
 import { api, internal } from "./_generated/api.js";
 import { Doc, Id } from "./_generated/dataModel.js";
@@ -23,10 +24,11 @@ import {
   QueryCtx,
 } from "./_generated/server.js";
 import { schema, v } from "./schema.js";
+import { insertVector, searchVectors } from "./vector/index.js";
 import {
-  getVectorTableName,
   VectorDimension,
   VectorDimensions,
+  VectorTableId,
   vVectorId,
 } from "./vector/tables.js";
 
@@ -45,23 +47,23 @@ export const getThreadsByUserId = query({
     cursor: v.optional(v.union(v.string(), v.null())),
     limit: v.optional(v.number()),
     offset: v.optional(v.number()),
-    statuses: v.optional(v.array(vThreadStatus)),
+    statuses: v.optional(vThreadStatus),
   },
   handler: async (ctx, args) => {
-    const streams = (args.statuses ?? ["active"]).map((status) =>
-      stream(ctx.db, schema)
-        .query("threads")
-        .withIndex("status_userId_order", (q) =>
-          q
-            .eq("status", status)
-            .eq("userId", args.userId)
-            .gte("order", args.offset ?? 0)
-        )
-    );
-    const threads = await mergedStream(streams, ["order"]).paginate({
-      numItems: args.limit ?? 100,
-      cursor: args.cursor ?? null,
-    });
+    const status = args.statuses ?? "active";
+    const threads = await paginator(ctx.db, schema)
+      .query("threads")
+      .withIndex("userId_status_order", (q) =>
+        q
+          .eq("userId", args.userId)
+          .eq("status", status)
+          .gte("order", args.offset ?? 0)
+      )
+      .order("desc")
+      .paginate({
+        numItems: args.limit ?? 100,
+        cursor: args.cursor ?? null,
+      });
     return {
       threads: threads.page,
       continueCursor: threads.continueCursor,
@@ -76,20 +78,17 @@ export const getThreadsByUserId = query({
 });
 
 const vThread = schema.tables.threads.validator;
-const statuses = vThread.fields.status.members.map((m) => m.value);
 
 export const createThread = mutation({
   args: omit(vThread.fields, ["order", "status"]),
   handler: async (ctx, args) => {
-    const streams = statuses.map((status) =>
-      stream(ctx.db, schema)
-        .query("threads")
-        .withIndex("status_userId_order", (q) =>
-          q.eq("status", status).eq("userId", args.userId)
-        )
-        .order("desc")
-    );
-    const latestThread = await mergedStream(streams, ["order"]).first();
+    const latestThread = await ctx.db
+      .query("threads")
+      .withIndex("userId_status_order", (q) =>
+        q.eq("userId", args.userId).eq("status", "active")
+      )
+      .order("desc")
+      .first();
     const order = (latestThread?.order ?? -1) + 1;
     const threadId = await ctx.db.insert("threads", {
       ...args,
@@ -164,7 +163,7 @@ export const deleteAllForUserIdAsync = mutation({
     userId: v.string(),
   },
   handler: async (ctx, args) => {
-    const isDone = await deleteAllFroUserIdAsyncHandler(ctx, {
+    const isDone = await deleteAllForUserIdAsyncHandler(ctx, {
       userId: args.userId,
       messagesCursor: null,
       threadsCursor: null,
@@ -189,11 +188,11 @@ type DeleteAllReturns = ObjectType<typeof deleteAllReturns>;
 
 export const _deleteAllForUserIdAsync = internalMutation({
   args: deleteAllArgs,
-  handler: deleteAllFroUserIdAsyncHandler,
+  handler: deleteAllForUserIdAsyncHandler,
   returns: v.boolean(),
 });
 
-async function deleteAllFroUserIdAsyncHandler(
+async function deleteAllForUserIdAsyncHandler(
   ctx: MutationCtx,
   args: DeleteAllArgs
 ): Promise<boolean> {
@@ -221,51 +220,38 @@ async function deletePageForUserId(
   ctx: MutationCtx,
   args: DeleteAllArgs
 ): Promise<DeleteAllReturns> {
-  const streams = statuses.map((status) =>
-    stream(ctx.db, schema)
-      .query("threads")
-      .withIndex("status_userId_order", (q) =>
-        q.eq("status", status).eq("userId", args.userId)
-      )
-      .order("desc")
-  );
-  const threadStreams = mergedStream(streams, ["order"]);
-  const messages = await threadStreams
-    .flatMap(
-      async (c) =>
-        stream(ctx.db, schema)
-          .query("messages")
-          .withIndex("threadId_status_tool_order_stepOrder", (q) =>
-            q.eq("threadId", c._id).eq("status", "success")
-          ),
-      ["tool", "order", "stepOrder"]
+  const threads = await paginator(ctx.db, schema)
+    .query("threads")
+    .withIndex("userId_status_order", (q) => q.eq("userId", args.userId))
+    .order("desc")
+    .paginate({
+      numItems: 100,
+      cursor: args.threadsCursor ?? null,
+    });
+  await Promise.all(threads.page.map((c) => ctx.db.delete(c._id)));
+  const messages = await paginator(ctx.db, schema)
+    .query("messages")
+    .withIndex("userId_status_tool_order_stepOrder", (q) =>
+      q.eq("userId", args.userId)
     )
+    .order("desc")
     .paginate({
       numItems: 100,
       cursor: args.messagesCursor ?? null,
     });
   await Promise.all(messages.page.map((m) => deleteMessage(ctx, m)));
-  if (messages.isDone) {
-    const threads = await threadStreams.paginate({
-      numItems: 100,
-      cursor: args.threadsCursor ?? null,
-    });
-    await Promise.all(threads.page.map((c) => ctx.db.delete(c._id)));
-    return {
-      messagesCursor: messages.continueCursor,
-      threadsCursor: threads.continueCursor,
-      isDone: threads.isDone,
-    };
-  }
   return {
     messagesCursor: messages.continueCursor,
-    threadsCursor: null,
+    threadsCursor: threads.continueCursor,
     isDone: messages.isDone,
   };
 }
 
 async function deleteMessage(ctx: MutationCtx, messageDoc: Doc<"messages">) {
   await ctx.db.delete(messageDoc._id);
+  if (messageDoc.embeddingId) {
+    await ctx.db.delete(messageDoc.embeddingId);
+  }
   if (messageDoc.fileId) {
     const file = await ctx.db.get(messageDoc.fileId);
     if (file) {
@@ -323,10 +309,10 @@ async function deletePageForThreadIdHandler(
   ctx: MutationCtx,
   args: DeleteThreadArgs
 ): Promise<DeleteThreadReturns> {
-  const messages = await stream(ctx.db, schema)
+  const messages = await paginator(ctx.db, schema)
     .query("messages")
     .withIndex("threadId_status_tool_order_stepOrder", (q) =>
-      q.eq("threadId", args.threadId).eq("status", "success")
+      q.eq("threadId", args.threadId)
     )
     .paginate({
       numItems: args.limit ?? 100,
@@ -373,7 +359,7 @@ export const messageStatuses = vMessageDoc.fields.status.members.map(
 
 const addMessagesArgs = {
   userId: v.optional(v.string()),
-  threadId: v.optional(v.id("threads")),
+  threadId: v.id("threads"),
   stepId: v.optional(v.id("steps")),
   parentMessageId: v.optional(v.id("messages")),
   messages: v.array(vMessageWithFileAndId),
@@ -381,6 +367,7 @@ const addMessagesArgs = {
   agentName: v.optional(v.string()),
   pending: v.optional(v.boolean()),
   failPendingSteps: v.optional(v.boolean()),
+  embeddings: v.optional(vEmbeddingsWithMetadata),
 };
 export const addMessages = mutation({
   args: addMessagesArgs,
@@ -394,16 +381,27 @@ async function addMessagesHandler(
   ctx: MutationCtx,
   args: ObjectType<typeof addMessagesArgs>
 ) {
+  assert(
+    !args.embeddings || args.embeddings.vectors.length === args.messages.length,
+    "embeddings must have one vector per message"
+  );
   let userId = args.userId;
   const threadId = args.threadId;
   if (!userId && args.threadId) {
     const thread = await ctx.db.get(args.threadId);
     assert(thread, `Thread ${args.threadId} not found`);
-    userId = thread._id;
+    userId = thread.userId;
   }
-  const { failPendingSteps, pending, messages, parentMessageId, ...rest } =
-    args;
+  const {
+    failPendingSteps,
+    pending,
+    messages,
+    parentMessageId,
+    embeddings,
+    ...rest
+  } = args;
   const parent = parentMessageId && (await ctx.db.get(parentMessageId));
+  // TODO: I think this is a bug - parent will be pending always?
   if (failPendingSteps && parent?.status !== "pending") {
     assert(args.threadId, "threadId is required to fail pending steps");
     const pendingMessages = await ctx.db
@@ -422,7 +420,18 @@ async function addMessagesHandler(
   let order = maxMessage?.order ?? -1;
   const toReturn: Doc<"messages">[] = [];
   if (messages.length > 0) {
-    for (const { message, fileId, id } of messages) {
+    for (const [i, { message, fileId, id }] of messages.entries()) {
+      const embedding = embeddings?.vectors[i] ?? undefined;
+      let embeddingId: VectorTableId | undefined;
+      if (embeddings && embedding) {
+        embeddingId = await insertVector(ctx, embeddings.dimension, {
+          vector: embedding,
+          model: embeddings.model,
+          table: "messages",
+          userId,
+          threadId,
+        });
+      }
       const tool = isTool(message);
       if (!tool) {
         order++;
@@ -430,6 +439,7 @@ async function addMessagesHandler(
       const text = extractText(message);
       const messageId = await ctx.db.insert("messages", {
         ...rest,
+        embeddingId,
         parentMessageId,
         userId,
         message,
@@ -439,7 +449,13 @@ async function addMessagesHandler(
         text,
         fileId,
         status: pending ? "pending" : "success",
+        stepOrder: 0,
       });
+      if (fileId) {
+        await ctx.db.patch(fileId, {
+          refcount: (await ctx.db.get(fileId))!.refcount + 1,
+        });
+      }
       toReturn.push((await ctx.db.get(messageId))!);
     }
   }
@@ -465,35 +481,47 @@ async function getMaxMessage(
       ["order", "stepOrder"]
     ).first();
   } else {
-    return mergedStream(
-      ["success" as const, "pending" as const].map((status) =>
-        stream(ctx.db, schema)
-          .query("messages")
-          .withIndex("userId_status_tool_order_stepOrder", (q) =>
-            q.eq("userId", userId).eq("status", status).eq("tool", false)
-          )
-          .order("desc")
-      ),
-      ["order", "stepOrder"]
-    ).first();
+    // DO explicitly
+    const maxPending = await ctx.db
+      .query("messages")
+      .withIndex("userId_status_tool_order_stepOrder", (q) =>
+        q.eq("userId", userId).eq("status", "pending").eq("tool", false)
+      )
+      .order("desc")
+      .first();
+    const maxSuccess = await ctx.db
+      .query("messages")
+      .withIndex("userId_status_tool_order_stepOrder", (q) =>
+        q.eq("userId", userId).eq("status", "success").eq("tool", false)
+      )
+      .order("desc")
+      .first();
+    return maxPending
+      ? maxSuccess
+        ? maxPending.order > maxSuccess.order
+          ? maxPending
+          : maxSuccess
+        : maxPending
+      : maxSuccess ?? null;
   }
 }
 
-const addStepsArgs = {
+const addStepArgs = {
   threadId: v.id("threads"),
   messageId: v.id("messages"),
-  steps: v.array(vStepWithMessagesWithFileAndId),
+  step: vStepWithMessages,
   failPendingSteps: v.optional(v.boolean()),
+  embeddings: v.optional(vEmbeddingsWithMetadata),
 };
 
-export const addSteps = mutation({
-  args: addStepsArgs,
+export const addStep = mutation({
+  args: addStepArgs,
   returns: v.array(v.doc("steps")),
-  handler: addStepsHandler,
+  handler: addStepHandler,
 });
-async function addStepsHandler(
+async function addStepHandler(
   ctx: MutationCtx,
-  args: ObjectType<typeof addStepsArgs>
+  args: ObjectType<typeof addStepArgs>
 ) {
   const parentMessage = await ctx.db.get(args.messageId);
   assert(parentMessage, `Message ${args.messageId} not found`);
@@ -512,35 +540,33 @@ async function addStepsHandler(
         await ctx.db.patch(step._id, { status: "failed" });
       }
     }
-    steps = steps.filter((s) => s.status === "success");
+    steps = steps.filter((s) => s.status !== "failed");
   }
-  let nextStepOrder = (steps.at(-1)?.stepOrder ?? -1) + 1;
-  for (const { step, messages } of args.steps) {
-    const stepId = await ctx.db.insert("steps", {
-      threadId: args.threadId,
-      parentMessageId: args.messageId,
-      order,
-      stepOrder: nextStepOrder,
-      status: step.finishReason === "stop" ? "success" : "pending",
-      step,
-    });
-    await addMessagesHandler(ctx, {
-      threadId: args.threadId,
-      parentMessageId: args.messageId,
-      stepId,
-      messages,
-      model: parentMessage.model,
-      agentName: parentMessage.agentName,
-      pending: step.finishReason === "stop" ? false : true,
-      failPendingSteps: false,
-    });
-    // We don't commit if the parent is still pending.
-    if (step.finishReason === "stop" && parentMessage.status === "success") {
-      await commitMessageHandler(ctx, { messageId: args.messageId });
-    }
-    steps.push((await ctx.db.get(stepId))!);
-    nextStepOrder++;
+  const { step, messages } = args.step;
+  const stepId = await ctx.db.insert("steps", {
+    threadId: args.threadId,
+    parentMessageId: args.messageId,
+    order,
+    stepOrder: (steps.at(-1)?.stepOrder ?? -1) + 1,
+    status: step.finishReason === "stop" ? "success" : "pending",
+    step,
+  });
+  await addMessagesHandler(ctx, {
+    threadId: args.threadId,
+    parentMessageId: args.messageId,
+    stepId,
+    messages,
+    model: parentMessage.model,
+    agentName: parentMessage.agentName,
+    pending: step.finishReason === "stop" ? false : true,
+    failPendingSteps: false,
+    embeddings: args.embeddings,
+  });
+  // We don't commit if the parent is still pending.
+  if (step.finishReason === "stop") {
+    await commitMessageHandler(ctx, { messageId: args.messageId });
   }
+  steps.push((await ctx.db.get(stepId))!);
   return steps;
 }
 
@@ -553,9 +579,23 @@ export const rollbackMessage = mutation({
   handler: async (ctx, { messageId, error }) => {
     const message = await ctx.db.get(messageId);
     assert(message, `Message ${messageId} not found`);
+    // TODO: do BFS to fail all associated messages, then steps
+    // with parentMessageId of those messages, etc.
+    const steps = await ctx.db
+      .query("steps")
+      .withIndex("parentMessageId_order_stepOrder", (q) =>
+        // TODO: fetch pending, and commit later
+        q.eq("parentMessageId", messageId)
+      )
+      .collect();
+    for (const step of steps) {
+      if (step.status === "pending") {
+        await ctx.db.patch(step._id, { status: "failed" });
+      }
+    }
     await ctx.db.patch(messageId, {
       status: "failed",
-      text: error ?? message.text,
+      error: error,
     });
   },
 });
@@ -686,18 +726,16 @@ export const searchMessages = action({
       if (!VectorDimensions.includes(dimension)) {
         throw new Error(`Unsupported vector dimension: ${dimension}`);
       }
-      const model = args.vectorModel ?? "unknown";
-      const tableName = getVectorTableName(dimension);
       const vectors = (
-        await ctx.vectorSearch(tableName, "vector", {
-          vector: args.vector,
-          filter: (q) =>
-            args.userId
-              ? q.eq("model_kind_userId", [model, "thread", args.userId])
-              : q.eq("model_kind_threadId", [model, "thread", args.threadId!]),
+        await searchVectors(ctx, args.vector, {
+          dimension,
+          model: args.vectorModel ?? "unknown",
+          table: "messages",
+          userId: args.userId,
+          threadId: args.threadId,
           limit,
         })
-      ).filter((v) => v._score > 0.5);
+      ).filter((v) => v._score > (args.vectorScoreThreshold ?? 0));
       // Reciprocal rank fusion
       const k = 10;
       const textEmbeddingIds = textSearchMessages?.map((m) => m.embeddingId);
@@ -706,11 +744,10 @@ export const searchMessages = action({
           id: v._id,
           score:
             1 / (i + k) +
-            1 / (textEmbeddingIds?.indexOf(v._id) ?? Infinity + k),
+            1 / ((textEmbeddingIds?.indexOf(v._id) ?? Infinity) + k),
         }))
         .sort((a, b) => b.score - a.score);
       const vectorIds = vectorScores.slice(0, limit).map((v) => v.id);
-
       const messages: Doc<"messages">[] = await ctx.runQuery(
         internal.messages._fetchVectorMessages,
         {
@@ -755,10 +792,11 @@ export const _fetchVectorMessages = internalQuery({
             .withIndex("embeddingId", (q) => q.eq("embeddingId", embeddingId))
             .filter((q) =>
               userId
-                ? q.eq("userId", userId)
-                : // eslint-disable-next-line @typescript-eslint/no-explicit-any
-                  q.eq("threadId", threadId as any)
+                ? q.eq(q.field("userId"), userId)
+                : q.eq(q.field("threadId"), threadId)
             )
+            // Don't include pending. Failed messages hopefully are deleted but may as well be safe.
+            .filter((q) => q.eq(q.field("status"), "success"))
             .first()
         )
       )
@@ -811,8 +849,8 @@ export const _fetchVectorMessages = internalQuery({
                 .eq("threadId", m.threadId)
                 .eq("status", "success")
                 .eq("tool", false)
-                .gt("order", earliest)
-                .lt("order", latest)
+                .gte("order", earliest)
+                .lte("order", latest)
             )
             .collect();
           if (!ranges[searchId]) {
@@ -827,8 +865,8 @@ export const _fetchVectorMessages = internalQuery({
                 .eq("userId", m.userId!)
                 .eq("status", "success")
                 .eq("tool", false)
-                .gt("order", earliest)
-                .lt("order", latest)
+                .gte("order", earliest)
+                .lte("order", latest)
             )
             .collect();
           if (!ranges[searchId]) {
@@ -838,9 +876,12 @@ export const _fetchVectorMessages = internalQuery({
         }
       }
     }
-    return Object.values(ranges)
-      .map((r) => r.sort((a, b) => a.order! - b.order!))
-      .flat();
+    for (const r of Object.values(ranges).flat()) {
+      if (!messages.includes(r)) {
+        messages.push(r);
+      }
+    }
+    return messages.sort((a, b) => a.order - b.order);
   },
 });
 
@@ -867,157 +908,3 @@ export const textSearch = query({
   },
   returns: v.array(v.doc("messages")),
 });
-
-// const vMemoryConfig = v.object({
-//   lastMessages: v.optional(v.union(v.number(), v.literal(false))),
-//   semanticRecall: v.optional(
-//     v.union(
-//       v.boolean(),
-//       v.object({
-//         topK: v.number(),
-//         messageRange: v.union(
-//           v.number(),
-//           v.object({ before: v.number(), after: v.number() }),
-//         ),
-//       }),
-//     ),
-//   ),
-//   workingMemory: v.optional(
-//     v.object({
-//       enabled: v.boolean(),
-//       template: v.optional(v.string()),
-//       use: v.optional(
-//         v.union(v.literal("text-stream"), v.literal("tool-call")),
-//       ),
-//     }),
-//   ),
-//   threads: v.optional(
-//     v.object({
-//       generateTitle: v.optional(v.boolean()),
-//     }),
-//   ),
-// });
-// const vSelectBy = v.object({
-//   vectorSearchString: v.optional(v.string()),
-//   last: v.optional(v.union(v.number(), v.literal(false))),
-//   include: v.optional(
-//     v.array(
-//       v.object({
-//         id: v.string(),
-//         withPreviousMessages: v.optional(v.number()),
-//         withNextMessages: v.optional(v.number()),
-//       })
-//     )
-//   ),
-// });
-
-// const DEFAULT_MESSAGES_LIMIT = 40; // What pg & upstash do too.
-
-// export const getThreadMessagesPage = query({
-//   args: {
-//     parentMessageId: v.string(),
-//     selectBy: v.optional(vSelectBy),
-//     // Unimplemented and as far I can tell no storage provider has either.
-//     // memoryConfig: v.optional(vMemoryConfig),
-//   },
-//   handler: async (ctx, args): Promise<SerializedMessage[]> => {
-//     const messages = await ctx.db
-//       .query("messages")
-//       .withIndex("parentMessageId", (q) => q.eq("parentMessageId", args.parentMessageId))
-//       .order("desc")
-//       .take(args.selectBy?.last ? args.selectBy.last : DEFAULT_MESSAGES_LIMIT);
-
-//     const handled: boolean[] = [];
-//     const toFetch: number[] = [];
-//     for (const m of messages) {
-//       handled[m.threadOrder] = true;
-//     }
-//     await Promise.all(
-//       args.selectBy?.include?.map(async (range) => {
-//         const includeDoc = await ctx.db
-//           .query("messages")
-//           .withIndex("id", (q) => q.eq("id", range.id))
-//           .unique();
-//         if (!includeDoc) {
-//           console.warn(`Message ${range.id} not found`);
-//           return;
-//         }
-//         if (!range.withPreviousMessages && !range.withNextMessages) {
-//           messages.push(includeDoc);
-//           return;
-//         }
-//         const order = includeDoc.threadOrder;
-//         for (
-//           let i = order - (range.withPreviousMessages ?? 0);
-//           i < order + (range.withNextMessages ?? 0);
-//           i++
-//         ) {
-//           if (!handled[i]) {
-//             toFetch.push(i);
-//             handled[i] = true;
-//           }
-//         }
-//       }) ?? []
-//     );
-//     // sort and find unique numbers in toFetch
-//     const uniqueToFetch = [...new Set(toFetch)].sort();
-//     // find contiguous ranges in uniqueToFetch
-//     const ranges: { start: number; end: number }[] = [];
-//     for (let i = 0; i < uniqueToFetch.length; i++) {
-//       const start = uniqueToFetch[i];
-//       let end = start;
-//       while (i + 1 < uniqueToFetch.length && uniqueToFetch[i + 1] === end + 1) {
-//         end++;
-//         i++;
-//       }
-//       ranges.push({ start, end });
-//     }
-//     const fetched = (
-//       await Promise.all(
-//         ranges.map(async (range) => {
-//           return await ctx.db
-//             .query("messages")
-//             .withIndex("parentMessageId", (q) =>
-//               q
-//                 .eq("parentMessageId", args.parentMessageId)
-//                 .gte("threadOrder", range.start)
-//                 .lte("threadOrder", range.end)
-//             )
-//             .collect();
-//         })
-//       )
-//     ).flat();
-//     messages.push(...fetched);
-//     return messages.map(messageToSerializedMastra);
-//   },
-//   returns: v.array(vSerializedMessage),
-// });
-
-// export const saveMessages = mutation({
-//   args: { messages: v.array(vSerializedMessage) },
-//   handler: async (ctx, args) => {
-//     const messagesByParentMessageId: Record<string, SerializedMessage[]> = {};
-//     for (const message of args.messages) {
-//       messagesByParentMessageId[message.parentMessageId] = [
-//         ...(messagesByParentMessageId[message.parentMessageId] ?? []),
-//         message,
-//       ];
-//     }
-//     for (const parentMessageId in messagesByParentMessageId) {
-//       const lastMessage = await ctx.db
-//         .query("messages")
-//         .withIndex("parentMessageId", (q) => q.eq("parentMessageId", parentMessageId))
-//         .order("desc")
-//         .first();
-//       let threadOrder = lastMessage?.threadOrder ?? 0;
-//       for (const message of messagesByParentMessageId[parentMessageId]) {
-//         threadOrder++;
-//         await ctx.db.insert("messages", {
-//           ...message,
-//           threadOrder,
-//         });
-//       }
-//     }
-//   },
-//   returns: v.null(),
-// });