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

package/dist/commonjs/component/messages.js

@@ -3,11 +3,12 @@ import { paginator } from "convex-helpers/server/pagination";
 import { mergedStream, stream } from "convex-helpers/server/stream";
 import { nullable, partial } from "convex-helpers/validators";
 import { DEFAULT_MESSAGE_RANGE, extractText, isTool } from "../shared.js";
-import { vThreadStatus, vMessageStatus, vMessageWithFileAndId, vSearchOptions, vStepWithMessagesWithFileAndId, } from "../validators.js";
+import { vEmbeddingsWithMetadata, vMessageStatus, vMessageWithFileAndId, vSearchOptions, vStepWithMessages, vThreadStatus, } from "../validators.js";
 import { api, internal } from "./_generated/api.js";
 import { action, internalMutation, internalQuery, mutation, query, } from "./_generated/server.js";
 import { schema, v } from "./schema.js";
-import { getVectorTableName, VectorDimensions, vVectorId, } from "./vector/tables.js";
+import { insertVector, searchVectors } from "./vector/index.js";
+import { VectorDimensions, vVectorId, } from "./vector/tables.js";
 export const getThread = query({
     args: { threadId: v.id("threads") },
     handler: async (ctx, args) => {
@@ -22,16 +23,18 @@ 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)
+        const status = args.statuses ?? "active";
+        const threads = await paginator(ctx.db, schema)
             .query("threads")
-            .withIndex("status_userId_order", (q) => q
-            .eq("status", status)
+            .withIndex("userId_status_order", (q) => q
             .eq("userId", args.userId)
-            .gte("order", args.offset ?? 0)));
-        const threads = await mergedStream(streams, ["order"]).paginate({
+            .eq("status", status)
+            .gte("order", args.offset ?? 0))
+            .order("desc")
+            .paginate({
             numItems: args.limit ?? 100,
             cursor: args.cursor ?? null,
         });
@@ -48,15 +51,14 @@ 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)
+        const latestThread = await ctx.db
             .query("threads")
-            .withIndex("status_userId_order", (q) => q.eq("status", status).eq("userId", args.userId))
-            .order("desc"));
-        const latestThread = await mergedStream(streams, ["order"]).first();
+            .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,
@@ -119,7 +121,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,
@@ -140,10 +142,10 @@ const deleteAllReturns = {
 };
 export const _deleteAllForUserIdAsync = internalMutation({
     args: deleteAllArgs,
-    handler: deleteAllFroUserIdAsyncHandler,
+    handler: deleteAllForUserIdAsyncHandler,
     returns: v.boolean(),
 });
-async function deleteAllFroUserIdAsyncHandler(ctx, args) {
+async function deleteAllForUserIdAsyncHandler(ctx, args) {
     const result = await deletePageForUserId(ctx, args);
     if (!result.isDone) {
         await ctx.scheduler.runAfter(0, internal.messages._deleteAllForUserIdAsync, {
@@ -160,40 +162,35 @@ export const _deletePageForUserId = internalMutation({
     returns: deleteAllReturns,
 });
 async function deletePageForUserId(ctx, args) {
-    const streams = statuses.map((status) => stream(ctx.db, schema)
+    const threads = await paginator(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)
+        .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("threadId_status_tool_order_stepOrder", (q) => q.eq("threadId", c._id).eq("status", "success")), ["tool", "order", "stepOrder"])
+        .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, messageDoc) {
     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) {
@@ -238,9 +235,9 @@ export const _deletePageForThreadId = internalMutation({
     returns: deleteThreadReturns,
 });
 async function deletePageForThreadIdHandler(ctx, args) {
-    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"))
+        .withIndex("threadId_status_tool_order_stepOrder", (q) => q.eq("threadId", args.threadId))
         .paginate({
         numItems: args.limit ?? 100,
         cursor: args.cursor ?? null,
@@ -281,7 +278,7 @@ export const vMessageDoc = schema.tables.messages.validator;
 export const messageStatuses = vMessageDoc.fields.status.members.map((m) => m.value);
 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),
@@ -289,6 +286,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,
@@ -299,15 +297,17 @@ export const addMessages = mutation({
     }),
 });
 async function addMessagesHandler(ctx, args) {
+    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
@@ -320,7 +320,18 @@ async function addMessagesHandler(ctx, args) {
     let order = maxMessage?.order ?? -1;
     const toReturn = [];
     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;
+            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++;
@@ -328,6 +339,7 @@ async function addMessagesHandler(ctx, args) {
             const text = extractText(message);
             const messageId = await ctx.db.insert("messages", {
                 ...rest,
+                embeddingId,
                 parentMessageId,
                 userId,
                 message,
@@ -337,7 +349,13 @@ async function addMessagesHandler(ctx, args) {
                 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)));
         }
     }
@@ -352,24 +370,39 @@ async function getMaxMessage(ctx, threadId, userId) {
             .order("desc")), ["order", "stepOrder"]).first();
     }
     else {
-        return mergedStream(["success", "pending"].map((status) => stream(ctx.db, schema)
+        // DO explicitly
+        const maxPending = await ctx.db
             .query("messages")
-            .withIndex("userId_status_tool_order_stepOrder", (q) => q.eq("userId", userId).eq("status", status).eq("tool", false))
-            .order("desc")), ["order", "stepOrder"]).first();
+            .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(ctx, args) {
+async function addStepHandler(ctx, args) {
     const parentMessage = await ctx.db.get(args.messageId);
     assert(parentMessage, `Message ${args.messageId} not found`);
     const order = parentMessage.order;
@@ -386,35 +419,33 @@ async function addStepsHandler(ctx, args) {
                 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;
 }
 export const rollbackMessage = mutation({
@@ -426,9 +457,22 @@ 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,
         });
     },
 });
@@ -537,15 +581,14 @@ 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]),
+            const vectors = (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);
@@ -553,7 +596,7 @@ export const searchMessages = action({
                 .map((v, i) => ({
                 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);
@@ -590,9 +633,10 @@ export const _fetchVectorMessages = internalQuery({
             .query("messages")
             .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))
+            ? 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()))).filter((m) => m !== undefined && m !== null && (!parent || m.order <= parent.order));
         messages.push(...(args.textSearchMessages ?? []));
         // TODO: prioritize more recent messages
@@ -638,8 +682,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]) {
                         ranges[searchId] = [];
@@ -653,8 +697,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]) {
                         ranges[searchId] = [];
@@ -663,9 +707,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);
     },
 });
 // returns ranges of messages in order of text search relevance,
@@ -689,153 +736,4 @@ 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(),
-// });
 //# sourceMappingURL=messages.js.map