@convex-dev/agent 0.0.17-alpha.2 → 0.1.0-alpha.0

package/src/component/messages.ts

@@ -1,4 +1,4 @@
-import { assert } from "convex-helpers";
+import { assert, omit } from "convex-helpers";
 import { mergedStream, stream } from "convex-helpers/server/stream";
 import { ObjectType } from "convex/values";
 import {
@@ -9,6 +9,7 @@ import {
 } from "../shared.js";
 import {
   paginationResultValidator,
+  vMessageEmbeddings,
   vMessageStatus,
   vMessageWithMetadata,
   vSearchOptions,
@@ -38,16 +39,20 @@ import {
   updateThread as _updateThread,
 } from "./threads.js";
 import { paginationOptsValidator } from "convex/server";
-
+import { MessageDoc, vMessageDoc } from "./schema.js";
 
 /** @deprecated Use *.threads.listMessagesByThreadId instead. */
-export const listThreadsByUserId= _listThreadsByUserId
+export const listThreadsByUserId = _listThreadsByUserId;
 
 /** @deprecated Use *.threads.getThread */
 export const getThread = _getThread;
 
 /** @deprecated Use *.threads.updateThread instead */
-export const updateThread= _updateThread;
+export const updateThread = _updateThread;
+
+function publicMessage(message: Doc<"messages">): MessageDoc {
+  return omit(message, ["parentMessageId", "stepId"]);
+}
 
 export async function deleteMessage(
   ctx: MutationCtx,
@@ -66,7 +71,6 @@ export async function deleteMessage(
   }
 }
 
-export const vMessageDoc = schema.tables.messages.validator;
 export const messageStatuses = vMessageDoc.fields.status.members.map(
   (m) => m.value
 );
@@ -75,9 +79,10 @@ const addMessagesArgs = {
   userId: v.optional(v.string()),
   threadId: v.id("threads"),
   stepId: v.optional(v.id("steps")),
-  parentMessageId: v.optional(v.id("messages")),
+  promptMessageId: v.optional(v.id("messages")),
   agentName: v.optional(v.string()),
   messages: v.array(vMessageWithMetadata),
+  embeddings: v.optional(vMessageEmbeddings),
   pending: v.optional(v.boolean()),
   failPendingSteps: v.optional(v.boolean()),
 };
@@ -85,8 +90,8 @@ export const addMessages = mutation({
   args: addMessagesArgs,
   handler: addMessagesHandler,
   returns: v.object({
-    messages: v.array(v.doc("messages")),
-    pending: v.optional(v.doc("messages")),
+    messages: v.array(vMessageDoc),
+    pending: v.optional(vMessageDoc),
   }),
 });
 async function addMessagesHandler(
@@ -100,8 +105,15 @@ async function addMessagesHandler(
     assert(thread, `Thread ${args.threadId} not found`);
     userId = thread.userId;
   }
-  const { failPendingSteps, pending, messages, parentMessageId, ...rest } =
-    args;
+  const {
+    embeddings,
+    failPendingSteps,
+    pending,
+    messages,
+    promptMessageId,
+    ...rest
+  } = args;
+  const parentMessage = promptMessageId && (await ctx.db.get(promptMessageId));
   if (failPendingSteps) {
     assert(args.threadId, "threadId is required to fail pending steps");
     const pendingMessages = await ctx.db
@@ -111,57 +123,69 @@ async function addMessagesHandler(
       )
       .collect();
     await Promise.all(
-      pendingMessages.map((m) =>
-        ctx.db.patch(m._id, { status: "failed", error: "Restarting" })
-      )
+      pendingMessages
+        .filter((m) => !parentMessage || m.order === parentMessage.order)
+        .map((m) =>
+          ctx.db.patch(m._id, { status: "failed", error: "Restarting" })
+        )
     );
   }
-  const maxMessage = await getMaxMessage(ctx, threadId, userId);
-  let order = maxMessage?.order ?? -1;
-  let stepOrder = maxMessage?.stepOrder ?? 0;
-  let lastMessageIsTool = maxMessage?.tool ?? false;
+  let order, stepOrder;
+  let fail = false;
+  if (promptMessageId) {
+    assert(parentMessage, `Parent message ${promptMessageId} not found`);
+    if (parentMessage.status === "failed") {
+      fail = true;
+    }
+    order = parentMessage.order;
+    // Defend against there being existing messages with this parent.
+    const maxMessage = await getMaxMessage(ctx, threadId, order);
+    stepOrder = maxMessage?.stepOrder ?? parentMessage.stepOrder;
+  } else {
+    const maxMessage = await getMaxMessage(ctx, threadId);
+    order = maxMessage ? maxMessage.order + 1 : 0;
+    stepOrder = -1;
+  }
   const toReturn: Doc<"messages">[] = [];
   if (messages.length > 0) {
-    for (const { message, files, embedding, ...fields } of messages) {
+    if (embeddings) {
+      assert(
+        embeddings.vectors.length === messages.length,
+        "embeddings.vectors.length must match messages.length"
+      );
+    }
+    for (let i = 0; i < messages.length; i++) {
+      const message = messages[i];
       let embeddingId: VectorTableId | undefined;
-      if (embedding) {
-        embeddingId = await insertVector(ctx, embedding.dimension, {
-          vector: embedding.vector,
-          model: embedding.model,
+      if (embeddings && embeddings.vectors[i]) {
+        embeddingId = await insertVector(ctx, embeddings.dimension, {
+          vector: embeddings.vectors[i]!,
+          model: embeddings.model,
           table: "messages",
           userId,
           threadId,
         });
       }
-      const tool = isTool(message);
-      if (lastMessageIsTool) {
-        stepOrder++;
-      } else {
-        order++;
-        stepOrder = 0;
-      }
-      lastMessageIsTool = tool;
-      const text = extractText(message);
+      stepOrder++;
       const messageId = await ctx.db.insert("messages", {
         ...rest,
-        ...fields,
+        ...message,
         embeddingId,
-        parentMessageId,
+        parentMessageId: promptMessageId,
         userId,
-        message,
         order,
-        tool,
-        text,
-        files,
-        status: pending ? "pending" : "success",
+        tool: isTool(message.message),
+        text: extractText(message.message),
+        status: fail ? "failed" : pending ? "pending" : "success",
+        error: fail ? "Parent message failed" : undefined,
         stepOrder,
       });
-      if (!fields.id) {
+      if (!message.id) {
         await ctx.db.patch(messageId, {
           id: messageId,
         });
       }
-      for (const { fileId } of files ?? []) {
+      for (const { fileId } of message.files ?? []) {
         if (!fileId) continue;
         await ctx.db.patch(fileId, {
           refcount: (await ctx.db.get(fileId))!.refcount + 1,
@@ -176,45 +200,44 @@ async function addMessagesHandler(
 // exported for tests
 export async function getMaxMessage(
   ctx: QueryCtx,
-  threadId: Id<"threads"> | undefined,
-  userId: string | undefined
+  threadId: Id<"threads">,
+  order?: number
 ) {
-  assert(threadId || userId, "One of threadId or userId is required");
-  if (threadId) {
-    return mergedStream(
-      [true, false].flatMap((tool) =>
-        ["success" as const, "pending" as const].map((status) =>
-          stream(ctx.db, schema)
-            .query("messages")
-            .withIndex("threadId_status_tool_order_stepOrder", (q) =>
-              q.eq("threadId", threadId).eq("status", status).eq("tool", tool)
-            )
-            .order("desc")
-        )
-      ),
-      ["order", "stepOrder"]
-    ).first();
-  } else {
-    return mergedStream(
-      [true, false].flatMap((tool) =>
-        ["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", tool)
-            )
-            .order("desc")
-        )
-      ),
-      ["order", "stepOrder"]
-    ).first();
-  }
+  return orderedMessagesStream(ctx, threadId, "desc", order).first();
+}
+
+function orderedMessagesStream(
+  ctx: QueryCtx,
+  threadId: Id<"threads">,
+  sortOrder: "asc" | "desc",
+  order?: number
+) {
+  return mergedStream(
+    [true, false].flatMap((tool) =>
+      ["success" as const, "pending" as const].map((status) =>
+        stream(ctx.db, schema)
+          .query("messages")
+          .withIndex("threadId_status_tool_order_stepOrder", (q) => {
+            const qq = q
+              .eq("threadId", threadId)
+              .eq("status", status)
+              .eq("tool", tool);
+            if (order) {
+              return qq.eq("order", order);
+            }
+            return qq;
+          })
+          .order(sortOrder)
+      )
+    ),
+    ["order", "stepOrder"]
+  );
 }
 
 const addStepArgs = {
   userId: v.optional(v.string()),
   threadId: v.id("threads"),
-  parentMessageId: v.id("messages"),
+  promptMessageId: v.id("messages"),
   step: vStepWithMessages,
   failPendingSteps: v.optional(v.boolean()),
 };
@@ -228,16 +251,15 @@ async function addStepHandler(
   ctx: MutationCtx,
   args: ObjectType<typeof addStepArgs>
 ) {
-  const parentMessage = await ctx.db.get(args.parentMessageId);
-  assert(parentMessage, `Message ${args.parentMessageId} not found`);
+  const parentMessage = await ctx.db.get(args.promptMessageId);
+  assert(parentMessage, `Message ${args.promptMessageId} not found`);
   const order = parentMessage.order;
-  assert(order !== undefined, `${args.parentMessageId} has no order`);
   // TODO: only fetch the last one if we aren't failing pending steps
   let steps = await ctx.db
     .query("steps")
     .withIndex("parentMessageId_order_stepOrder", (q) =>
       // TODO: fetch pending, and commit later
-      q.eq("parentMessageId", args.parentMessageId)
+      q.eq("parentMessageId", args.promptMessageId)
     )
     .collect();
   if (args.failPendingSteps) {
@@ -251,7 +273,7 @@ async function addStepHandler(
   const { step, messages } = args.step;
   const stepId = await ctx.db.insert("steps", {
     threadId: args.threadId,
-    parentMessageId: args.parentMessageId,
+    parentMessageId: args.promptMessageId,
     order,
     stepOrder: (steps.at(-1)?.stepOrder ?? -1) + 1,
     status: step.finishReason === "stop" ? "success" : "pending",
@@ -261,7 +283,7 @@ async function addStepHandler(
     userId: args.userId,
     threadId: args.threadId,
     stepId,
-    parentMessageId: args.parentMessageId,
+    promptMessageId: args.promptMessageId,
     agentName: parentMessage.agentName,
     messages,
     pending: step.finishReason === "stop" ? false : true,
@@ -269,7 +291,7 @@ async function addStepHandler(
   });
   // We don't commit if the parent is still pending.
   if (step.finishReason === "stop") {
-    await commitMessageHandler(ctx, { messageId: args.parentMessageId });
+    await commitMessageHandler(ctx, { messageId: args.promptMessageId });
   }
   steps.push((await ctx.db.get(stepId))!);
   return steps;
@@ -284,8 +306,18 @@ 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 messages = await orderedMessagesStream(
+      ctx,
+      message.threadId,
+      "asc",
+      message.order
+    ).collect();
+    for (const m of messages) {
+      if (m.status === "pending") {
+        await ctx.db.patch(m._id, { status: "failed", error });
+      }
+    }
+
     const steps = await ctx.db
       .query("steps")
       .withIndex("parentMessageId_order_stepOrder", (q) =>
@@ -347,7 +379,6 @@ async function commitMessageHandler(
   ).collect();
   for (const message of messages) {
     await ctx.db.patch(message._id, { status: "success" });
-    // TODO: recursively commit steps & messages that might depend on this one.
   }
 }
 
@@ -360,13 +391,18 @@ export const listMessagesByThreadId = query({
     order: v.optional(v.union(v.literal("asc"), v.literal("desc"))),
     paginationOpts: v.optional(paginationOptsValidator),
     statuses: v.optional(v.array(vMessageStatus)),
-    beforeMessageId: v.optional(v.id("messages")),
+    upToAndIncludingMessageId: v.optional(v.id("messages")),
   },
   handler: async (ctx, args) => {
     const statuses =
       args.statuses ?? vMessageStatus.members.map((m) => m.value);
-    const before =
-      args.beforeMessageId && (await ctx.db.get(args.beforeMessageId));
+    const last =
+      args.upToAndIncludingMessageId &&
+      (await ctx.db.get(args.upToAndIncludingMessageId));
+    assert(
+      !last || last.threadId === args.threadId,
+      "upToAndIncludingMessageId must be a message in the thread"
+    );
     const toolOptions = args.excludeToolMessages ? [false] : [true, false];
     const order = args.order ?? "desc";
     const streams = toolOptions.flatMap((tool) =>
@@ -378,17 +414,17 @@ export const listMessagesByThreadId = query({
               .eq("threadId", args.threadId)
               .eq("status", status)
               .eq("tool", tool);
-            if (before) {
-              return qq.lte("order", before.order);
+            if (last) {
+              return qq.lte("order", last.order);
             }
             return qq;
           })
           .order(order)
           .filterWith(
             async (m) =>
-              !before ||
-              m.order < before.order ||
-              (m.order === before.order && m.stepOrder < before.stepOrder)
+              !last ||
+              m.order < last.order ||
+              (m.order === last.order && m.stepOrder <= last.stepOrder)
           )
       )
     );
@@ -401,9 +437,9 @@ export const listMessagesByThreadId = query({
         cursor: null,
       }
     );
-    return messages;
+    return { ...messages, page: messages.page.map(publicMessage) };
   },
-  returns: paginationResultValidator(v.doc("messages")),
+  returns: paginationResultValidator(vMessageDoc),
 });
 
 /** @deprecated Use listMessagesByThreadId instead. */
@@ -412,7 +448,7 @@ export const getThreadMessages = query({
   handler: async () => {
     throw new Error("Use listMessagesByThreadId instead of getThreadMessages");
   },
-  returns: paginationResultValidator(v.doc("messages")),
+  returns: paginationResultValidator(vMessageDoc),
 });
 
 export const searchMessages = action({
@@ -422,11 +458,11 @@ export const searchMessages = action({
     beforeMessageId: v.optional(v.id("messages")),
     ...vSearchOptions.fields,
   },
-  returns: v.array(v.doc("messages")),
-  handler: async (ctx, args): Promise<Doc<"messages">[]> => {
+  returns: v.array(vMessageDoc),
+  handler: async (ctx, args): Promise<MessageDoc[]> => {
     assert(args.userId || args.threadId, "Specify userId or threadId");
     const limit = args.limit;
-    let textSearchMessages: Doc<"messages">[] | undefined;
+    let textSearchMessages: MessageDoc[] | undefined;
     if (args.text) {
       textSearchMessages = await ctx.runQuery(api.messages.textSearch, {
         userId: args.userId,
@@ -463,14 +499,14 @@ export const searchMessages = action({
         }))
         .sort((a, b) => b.score - a.score);
       const vectorIds = vectorScores.slice(0, limit).map((v) => v.id);
-      const messages: Doc<"messages">[] = await ctx.runQuery(
+      const messages: MessageDoc[] = await ctx.runQuery(
         internal.messages._fetchSearchMessages,
         {
           userId: args.userId,
           threadId: args.threadId,
           vectorIds,
           textSearchMessages: textSearchMessages?.filter(
-            (m) => !vectorIds.includes(m.embeddingId!)
+            (m) => !vectorIds.includes(m.embeddingId! as VectorTableId)
           ),
           messageRange: args.messageRange ?? DEFAULT_MESSAGE_RANGE,
           beforeMessageId: args.beforeMessageId,
@@ -488,18 +524,18 @@ export const _fetchSearchMessages = internalQuery({
     userId: v.optional(v.string()),
     threadId: v.optional(v.id("threads")),
     vectorIds: v.array(vVectorId),
-    textSearchMessages: v.optional(v.array(v.doc("messages"))),
+    textSearchMessages: v.optional(v.array(vMessageDoc)),
     messageRange: v.object({ before: v.number(), after: v.number() }),
     beforeMessageId: v.optional(v.id("messages")),
     limit: v.number(),
   },
-  returns: v.array(v.doc("messages")),
-  handler: async (ctx, args): Promise<Doc<"messages">[]> => {
+  returns: v.array(vMessageDoc),
+  handler: async (ctx, args): Promise<MessageDoc[]> => {
     const beforeMessage =
       args.beforeMessageId && (await ctx.db.get(args.beforeMessageId));
     const { userId, threadId } = args;
     assert(userId || threadId, "Specify userId or threadId to search");
-    let messages = (
+    let messages: MessageDoc[] = (
       await Promise.all(
         args.vectorIds.map((embeddingId) =>
           ctx.db
@@ -515,16 +551,18 @@ export const _fetchSearchMessages = internalQuery({
             .first()
         )
       )
-    ).filter(
-      (m): m is Doc<"messages"> =>
-        m !== undefined &&
-        m !== null &&
-        !m.tool &&
-        (!beforeMessage ||
-          m.order < beforeMessage.order ||
-          (m.order === beforeMessage.order &&
-            m.stepOrder < beforeMessage.stepOrder))
-    );
+    )
+      .filter(
+        (m): m is Doc<"messages"> =>
+          m !== undefined &&
+          m !== null &&
+          !m.tool &&
+          (!beforeMessage ||
+            m.order < beforeMessage.order ||
+            (m.order === beforeMessage.order &&
+              m.stepOrder < beforeMessage.stepOrder))
+      )
+      .map(publicMessage);
     messages.push(...(args.textSearchMessages ?? []));
     // TODO: prioritize more recent messages
     messages.sort((a, b) => a.order! - b.order!);
@@ -562,44 +600,26 @@ export const _fetchSearchMessages = internalQuery({
         included[searchId].add(i);
       }
       if (earliest !== latest) {
-        if (m.threadId) {
-          const surrounding = await ctx.db
-            .query("messages")
-            .withIndex("threadId_status_tool_order_stepOrder", (q) =>
-              q
-                .eq("threadId", m.threadId)
-                .eq("status", "success")
-                .eq("tool", false)
-                .gte("order", earliest)
-                .lte("order", latest)
-            )
-            .collect();
-          if (!ranges[searchId]) {
-            ranges[searchId] = [];
-          }
-          ranges[searchId].push(...surrounding);
-        } else {
-          const surrounding = await ctx.db
-            .query("messages")
-            .withIndex("userId_status_tool_order_stepOrder", (q) =>
-              q
-                .eq("userId", m.userId!)
-                .eq("status", "success")
-                .eq("tool", false)
-                .gte("order", earliest)
-                .lte("order", latest)
-            )
-            .collect();
-          if (!ranges[searchId]) {
-            ranges[searchId] = [];
-          }
-          ranges[searchId].push(...surrounding);
+        const surrounding = await ctx.db
+          .query("messages")
+          .withIndex("threadId_status_tool_order_stepOrder", (q) =>
+            q
+              .eq("threadId", m.threadId as Id<"threads">)
+              .eq("status", "success")
+              .eq("tool", false)
+              .gte("order", earliest)
+              .lte("order", latest)
+          )
+          .collect();
+        if (!ranges[searchId]) {
+          ranges[searchId] = [];
         }
+        ranges[searchId].push(...surrounding);
       }
     }
     for (const r of Object.values(ranges).flat()) {
       if (!messages.some((m) => m._id === r._id)) {
-        messages.push(r);
+        messages.push(publicMessage(r));
       }
     }
     return messages.sort((a, b) => a.order - b.order);
@@ -637,13 +657,15 @@ export const textSearch = query({
         return qq;
       })
       .take(args.limit);
-    return messages.filter(
-      (m) =>
-        !beforeMessage ||
-        m.order < beforeMessage.order ||
-        (m.order === beforeMessage.order &&
-          m.stepOrder < beforeMessage.stepOrder)
-    );
+    return messages
+      .filter(
+        (m) =>
+          !beforeMessage ||
+          m.order < beforeMessage.order ||
+          (m.order === beforeMessage.order &&
+            m.stepOrder < beforeMessage.stepOrder)
+      )
+      .map(publicMessage);
   },
-  returns: v.array(v.doc("messages")),
+  returns: v.array(vMessageDoc),
 });

package/src/component/schema.ts

@@ -1,5 +1,5 @@
 import { defineSchema, defineTable } from "convex/server";
-import { v } from "convex/values";
+import { Infer, v } from "convex/values";
 import {
   vThreadStatus,
   vMessage,
@@ -13,9 +13,11 @@ import {
   vProviderMetadata,
   vReasoningDetails,
   vFile,
+  vFileWithStringId,
 } from "../validators.js";
 import { typedV } from "convex-helpers/validators";
 import vectorTables, { vVectorId } from "./vector/tables.js";
+import { omit } from "convex-helpers";
 
 export const schema = defineSchema({
   threads: defineTable({
@@ -28,14 +30,10 @@ export const schema = defineSchema({
     parentThreadIds: v.optional(v.array(v.id("threads"))),
     order: /*DEPRECATED*/ v.optional(v.number()),
   }).index("userId", ["userId"]),
-  // TODO: text search on title/ summary
   messages: defineTable({
     id: v.optional(v.string()), // external id, e.g. from Vercel AI SDK
-    userId: v.optional(v.string()), // useful for future indexes (text search)
+    userId: v.optional(v.string()), // useful for searching across threads
     threadId: v.id("threads"),
-    // TODO: is this redunant with message at last step @ order - 1?
-    parentMessageId: v.optional(v.id("messages")),
-    stepId: v.optional(v.id("steps")),
     // Repeats until a non-tool message.
     order: v.number(),
     stepOrder: v.number(),
@@ -64,6 +62,9 @@ export const schema = defineSchema({
     reasoningDetails: v.optional(vReasoningDetails),
     warnings: v.optional(v.array(vLanguageModelV1CallWarning)),
     finishReason: v.optional(vFinishReason),
+    // DEPRECATED
+    parentMessageId: v.optional(v.id("messages")),
+    stepId: v.optional(v.id("steps")),
   })
     // Allows finding successful visible messages in order
     // Also surface pending messages separately to e.g. stream
@@ -75,21 +76,6 @@ export const schema = defineSchema({
       "order",
       "stepOrder",
     ])
-    .index("userId_status_tool_order_stepOrder", [
-      "userId",
-      "status",
-      "tool",
-      "order",
-      "stepOrder",
-    ])
-    // Allows finding all threaded messages in order
-    // Allows finding all failed messages to evaluate
-    // .index("status_parentMessageId_order_stepOrder", [
-    //   "status",
-    //   "parentMessageId",
-    //   "order",
-    //   "stepOrder",
-    // ])
     // Allows text search on message content
     .searchIndex("text_search", {
       searchField: "text",
@@ -148,4 +134,31 @@ export const schema = defineSchema({
 export const vv = typedV(schema);
 export { vv as v };
 
+// Public
+export const vThreadDoc = v.object({
+  _id: v.string(),
+  _creationTime: v.number(),
+  userId: v.optional(v.string()), // Unset for anonymous
+  title: v.optional(v.string()),
+  summary: v.optional(v.string()),
+  status: vThreadStatus,
+});
+export type ThreadDoc = Infer<typeof vThreadDoc>;
+
+export const vMessageDoc = v.object({
+  _id: v.string(),
+  _creationTime: v.number(),
+  ...omit(schema.tables.messages.validator.fields, [
+    "parentMessageId",
+    "stepId",
+  ]),
+  // Overwrite all the types that have a v.id validator
+  // Outside of the component, they are strings
+  threadId: v.string(),
+  embeddingId: v.optional(v.string()),
+  files: v.optional(v.array(vFileWithStringId)),
+});
+export type MessageDoc = Infer<typeof vMessageDoc>;
+
+
 export default schema;