@convex-dev/agent 0.0.1-alpha.0 → 0.0.1-alpha.2

package/src/client/types.ts

@@ -7,6 +7,10 @@ import {
   GenericActionCtx,
 } from "convex/server";
 import { GenericId } from "convex/values";
+import type { Doc } from "../component/_generated/dataModel";
+
+export type ChatDoc = OpaqueIds<Doc<"chats">>;
+export type MessageDoc = OpaqueIds<Doc<"messages">>;
 
 /* Type utils follow */
 export type RunQueryCtx = {

package/src/component/_generated/api.d.ts

@@ -10,6 +10,7 @@
 
 import type * as lib from "../lib.js";
 import type * as messages from "../messages.js";
+import type * as vector_index from "../vector/index.js";
 import type * as vector_tables from "../vector/tables.js";
 
 import type {
@@ -29,6 +30,7 @@ import type {
 declare const fullApi: ApiFromModules<{
   lib: typeof lib;
   messages: typeof messages;
+  "vector/index": typeof vector_index;
   "vector/tables": typeof vector_tables;
 }>;
 export type Mounts = {
@@ -38,7 +40,7 @@ export type Mounts = {
       "public",
       {
         agentName?: string;
-        chatId: string;
+        chatId?: string;
         failPendingSteps?: boolean;
         messages: Array<{
           fileId?: string;
@@ -145,13 +147,14 @@ export type Mounts = {
         parentMessageId?: string;
         pending?: boolean;
         stepId?: string;
+        userId?: string;
       },
       {
         messages: Array<{
           _creationTime: number;
           _id: string;
           agentName?: string;
-          chatId: string;
+          chatId?: string;
           embeddingId?:
             | string
             | string
@@ -262,7 +265,7 @@ export type Mounts = {
                 role: "system";
               };
           model?: string;
-          order?: number;
+          order: number;
           status: "pending" | "success" | "failed";
           stepId?: string;
           stepOrder?: number;
@@ -275,7 +278,7 @@ export type Mounts = {
           _creationTime: number;
           _id: string;
           agentName?: string;
-          chatId: string;
+          chatId?: string;
           embeddingId?:
             | string
             | string
@@ -386,7 +389,7 @@ export type Mounts = {
                 role: "system";
               };
           model?: string;
-          order?: number;
+          order: number;
           status: "pending" | "success" | "failed";
           stepId?: string;
           stepOrder?: number;
@@ -1009,6 +1012,7 @@ export type Mounts = {
         isTool?: boolean;
         limit?: number;
         order?: "asc" | "desc";
+        parentMessageId?: string;
         statuses?: Array<"pending" | "success" | "failed">;
       },
       {
@@ -1018,7 +1022,7 @@ export type Mounts = {
           _creationTime: number;
           _id: string;
           agentName?: string;
-          chatId: string;
+          chatId?: string;
           embeddingId?:
             | string
             | string
@@ -1129,7 +1133,7 @@ export type Mounts = {
                 role: "system";
               };
           model?: string;
-          order?: number;
+          order: number;
           status: "pending" | "success" | "failed";
           stepId?: string;
           stepOrder?: number;
@@ -1195,6 +1199,7 @@ export type Mounts = {
         chatId?: string;
         limit: number;
         messageRange?: { after: number; before: number };
+        parentMessageId?: string;
         text?: string;
         userId?: string;
         vector?: Array<number>;
@@ -1204,7 +1209,7 @@ export type Mounts = {
         _creationTime: number;
         _id: string;
         agentName?: string;
-        chatId: string;
+        chatId?: string;
         embeddingId?:
           | string
           | string
@@ -1315,7 +1320,7 @@ export type Mounts = {
               role: "system";
             };
         model?: string;
-        order?: number;
+        order: number;
         status: "pending" | "success" | "failed";
         stepId?: string;
         stepOrder?: number;
@@ -1333,7 +1338,7 @@ export type Mounts = {
         _creationTime: number;
         _id: string;
         agentName?: string;
-        chatId: string;
+        chatId?: string;
         embeddingId?:
           | string
           | string
@@ -1444,7 +1449,7 @@ export type Mounts = {
               role: "system";
             };
         model?: string;
-        order?: number;
+        order: number;
         status: "pending" | "success" | "failed";
         stepId?: string;
         stepOrder?: number;
@@ -1479,6 +1484,128 @@ export type Mounts = {
       }
     >;
   };
+  vector: {
+    index: {
+      deleteBatch: FunctionReference<
+        "mutation",
+        "public",
+        {
+          ids: Array<
+            | string
+            | string
+            | string
+            | string
+            | string
+            | string
+            | string
+            | string
+            | string
+          >;
+        },
+        null
+      >;
+      deleteBatchForChat: FunctionReference<
+        "mutation",
+        "public",
+        {
+          chatId: string;
+          cursor?: string;
+          limit: number;
+          model: string;
+          vectorDimension:
+            | 128
+            | 256
+            | 512
+            | 768
+            | 1024
+            | 1536
+            | 2048
+            | 3072
+            | 4096;
+        },
+        { continueCursor: string; isDone: boolean }
+      >;
+      insertBatch: FunctionReference<
+        "mutation",
+        "public",
+        {
+          vectorDimension:
+            | 128
+            | 256
+            | 512
+            | 768
+            | 1024
+            | 1536
+            | 2048
+            | 3072
+            | 4096;
+          vectors: Array<{
+            chatId?: string;
+            kind: "chat" | "memory";
+            model: string;
+            userId?: string;
+            vector: Array<number>;
+          }>;
+        },
+        null
+      >;
+      paginate: FunctionReference<
+        "query",
+        "public",
+        {
+          cursor?: string;
+          limit: number;
+          targetModel: string;
+          vectorDimension:
+            | 128
+            | 256
+            | 512
+            | 768
+            | 1024
+            | 1536
+            | 2048
+            | 3072
+            | 4096;
+        },
+        {
+          continueCursor: string;
+          ids: Array<
+            | string
+            | string
+            | string
+            | string
+            | string
+            | string
+            | string
+            | string
+            | string
+          >;
+          isDone: boolean;
+        }
+      >;
+      updateBatch: FunctionReference<
+        "mutation",
+        "public",
+        {
+          vectors: Array<{
+            id:
+              | string
+              | string
+              | string
+              | string
+              | string
+              | string
+              | string
+              | string
+              | string;
+            model: string;
+            vector: Array<number>;
+          }>;
+        },
+        null
+      >;
+    };
+  };
 };
 // For now fullApiWithMounts is only fullApi which provides
 // jump-to-definition in component client code.

package/src/component/messages.ts

@@ -2,19 +2,14 @@ import { assert, omit, pick } from "convex-helpers";
 import { paginator } from "convex-helpers/server/pagination";
 import { mergedStream, stream } from "convex-helpers/server/stream";
 import { nullable, partial } from "convex-helpers/validators";
-import { Infer, ObjectType } from "convex/values";
+import { ObjectType } from "convex/values";
 import { DEFAULT_MESSAGE_RANGE, extractText, isTool } from "../shared.js";
 import {
-  Message,
-  MessageWithFileAndId,
-  vAssistantMessage,
   vChatStatus,
   vMessageStatus,
   vMessageWithFileAndId,
   vSearchOptions,
-  vStep,
   vStepWithMessagesWithFileAndId,
-  vToolMessage,
 } from "../validators.js";
 import { api, internal } from "./_generated/api.js";
 import { Doc, Id } from "./_generated/dataModel.js";
@@ -63,7 +58,7 @@ export const getChatsByUserId = query({
             .gte("order", args.offset ?? 0)
         )
     );
-    const chats = await mergedStream(streams, ["order", "stepOrder"]).paginate({
+    const chats = await mergedStream(streams, ["order"]).paginate({
       numItems: args.limit ?? 100,
       cursor: args.cursor ?? null,
     });
@@ -377,7 +372,8 @@ export const messageStatuses = vMessageDoc.fields.status.members.map(
 );
 
 const addMessagesArgs = {
-  chatId: v.id("chats"),
+  userId: v.optional(v.string()),
+  chatId: v.optional(v.id("chats")),
   stepId: v.optional(v.id("steps")),
   parentMessageId: v.optional(v.id("messages")),
   messages: v.array(vMessageWithFileAndId),
@@ -398,15 +394,22 @@ async function addMessagesHandler(
   ctx: MutationCtx,
   args: ObjectType<typeof addMessagesArgs>
 ) {
-  const chat = await ctx.db.get(args.chatId);
-  assert(chat, `Chat ${args.chatId} not found`);
+  let userId = args.userId;
+  const chatId = args.chatId;
+  if (!userId && args.chatId) {
+    const chat = await ctx.db.get(args.chatId);
+    assert(chat, `Chat ${args.chatId} not found`);
+    userId = chat._id;
+  }
   const { failPendingSteps, pending, messages, parentMessageId, ...rest } =
     args;
-  if (failPendingSteps) {
+  const parent = parentMessageId && (await ctx.db.get(parentMessageId));
+  if (failPendingSteps && parent?.status !== "pending") {
+    assert(args.chatId, "chatId is required to fail pending steps");
     const pendingMessages = await ctx.db
       .query("messages")
       .withIndex("chatId_status_tool_order_stepOrder", (q) =>
-        q.eq("chatId", args.chatId).eq("status", "pending")
+        q.eq("chatId", chatId).eq("status", "pending")
       )
       .collect();
     await Promise.all(
@@ -415,14 +418,14 @@ async function addMessagesHandler(
       )
     );
   }
-  let order: number | undefined;
-  const maxMessage = await getMaxMessage(ctx, args.chatId);
+  let threadId = parentMessageId;
+  const maxMessage = await getMaxMessage(ctx, chatId, userId);
   // If the previous message isn't our parent, we make a new thread.
-  const threadId =
+  threadId =
     parentMessageId && maxMessage?._id === parentMessageId
       ? maxMessage.threadId ?? parentMessageId
       : parentMessageId;
-  order = maxMessage?.order ?? -1;
+  let order = maxMessage?.order ?? -1;
   const toReturn: Doc<"messages">[] = [];
   if (messages.length > 0) {
     for (const { message, fileId, id } of messages) {
@@ -434,7 +437,7 @@ async function addMessagesHandler(
       const messageId = await ctx.db.insert("messages", {
         ...rest,
         threadId,
-        userId: chat.userId,
+        userId,
         message,
         id,
         order,
@@ -449,18 +452,37 @@ async function addMessagesHandler(
   return { messages: toReturn };
 }
 
-async function getMaxMessage(ctx: QueryCtx, chatId: Id<"chats">) {
-  return mergedStream(
-    ["success" as const, "pending" as const].map((status) =>
-      stream(ctx.db, schema)
-        .query("messages")
-        .withIndex("chatId_status_tool_order_stepOrder", (q) =>
-          q.eq("chatId", chatId).eq("status", status).eq("tool", false)
-        )
-        .order("desc")
-    ),
-    ["order", "stepOrder"]
-  ).first();
+async function getMaxMessage(
+  ctx: QueryCtx,
+  chatId: Id<"chats"> | undefined,
+  userId: string | undefined
+) {
+  assert(chatId || userId, "One of chatId or userId is required");
+  if (chatId) {
+    return mergedStream(
+      ["success" as const, "pending" as const].map((status) =>
+        stream(ctx.db, schema)
+          .query("messages")
+          .withIndex("chatId_status_tool_order_stepOrder", (q) =>
+            q.eq("chatId", chatId).eq("status", status).eq("tool", false)
+          )
+          .order("desc")
+      ),
+      ["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();
+  }
 }
 
 const addStepsArgs = {
@@ -518,7 +540,8 @@ async function addStepsHandler(
       pending: step.finishReason === "stop" ? false : true,
       failPendingSteps: false,
     });
-    if (step.finishReason === "stop") {
+    // 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))!);
@@ -585,6 +608,7 @@ 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.
   }
 }
 
@@ -597,9 +621,12 @@ export const getChatMessages = query({
     // Note: the other arguments cannot change from when the cursor was created.
     cursor: v.optional(v.string()),
     statuses: v.optional(v.array(vMessageStatus)),
+    parentMessageId: v.optional(v.id("messages")),
   },
   handler: async (ctx, args) => {
     const statuses = args.statuses ?? ["success"];
+    const parent =
+      args.parentMessageId && (await ctx.db.get(args.parentMessageId));
     const toolOptions =
       args.isTool === undefined ? [true, false] : [args.isTool];
     const order = args.order ?? "desc";
@@ -607,9 +634,16 @@ export const getChatMessages = query({
       statuses.map((status) =>
         stream(ctx.db, schema)
           .query("messages")
-          .withIndex("chatId_status_tool_order_stepOrder", (q) =>
-            q.eq("chatId", args.chatId).eq("status", status).eq("tool", tool)
-          )
+          .withIndex("chatId_status_tool_order_stepOrder", (q) => {
+            const qq = q
+              .eq("chatId", args.chatId)
+              .eq("status", status)
+              .eq("tool", tool);
+            if (parent) {
+              return qq.lte("order", parent.order);
+            }
+            return qq;
+          })
           .order(order)
       )
     );
@@ -637,6 +671,7 @@ export const searchMessages = action({
   args: {
     userId: v.optional(v.string()),
     chatId: v.optional(v.id("chats")),
+    parentMessageId: v.optional(v.id("messages")),
     ...vSearchOptions.fields,
   },
   returns: v.array(v.doc("messages")),
@@ -688,10 +723,12 @@ export const searchMessages = action({
           userId: args.userId,
           chatId: args.chatId,
           vectorIds,
-          textSearchMessages: textSearchMessages
-            ?.filter((m) => !vectorIds.includes(m.embeddingId!))
-            .slice(0, limit - vectorIds.length),
+          textSearchMessages: textSearchMessages?.filter(
+            (m) => !vectorIds.includes(m.embeddingId!)
+          ),
           messageRange: args.messageRange ?? DEFAULT_MESSAGE_RANGE,
+          parentMessageId: args.parentMessageId,
+          limit,
         }
       );
       return messages;
@@ -707,71 +744,104 @@ export const _fetchVectorMessages = internalQuery({
     vectorIds: v.array(vVectorId),
     textSearchMessages: v.optional(v.array(v.doc("messages"))),
     messageRange: v.object({ before: v.number(), after: v.number() }),
+    parentMessageId: v.optional(v.id("messages")),
+    limit: v.number(),
   },
   returns: v.array(v.doc("messages")),
   handler: async (ctx, args): Promise<Doc<"messages">[]> => {
-    const messages = (
+    const parent =
+      args.parentMessageId && (await ctx.db.get(args.parentMessageId));
+    const { userId, chatId } = args;
+    assert(userId || chatId, "Specify userId or chatId to search");
+    let messages = (
       await Promise.all(
         args.vectorIds.map((embeddingId) =>
           ctx.db
             .query("messages")
             .withIndex("embeddingId", (q) => q.eq("embeddingId", embeddingId))
-            .filter(
-              (q) =>
-                args.userId
-                  ? q.eq("userId", args.userId)
-                  : // eslint-disable-next-line @typescript-eslint/no-explicit-any
-                    q.eq("chatId", args.chatId as any) // not sure why it's failing...
+            .filter((q) =>
+              userId
+                ? q.eq("userId", userId)
+                : // eslint-disable-next-line @typescript-eslint/no-explicit-any
+                  q.eq("chatId", chatId as any)
             )
             .first()
         )
       )
-    ).filter((m): m is Doc<"messages"> => m !== undefined);
+    ).filter(
+      (m): m is Doc<"messages"> =>
+        m !== undefined && m !== null && (!parent || m.order <= parent.order)
+    );
     messages.push(...(args.textSearchMessages ?? []));
+    // TODO: prioritize more recent messages
     messages.sort((a, b) => a.order! - b.order!);
+    messages = messages.slice(0, args.limit);
     // Fetch the surrounding messages
-    const included: Record<Id<"chats">, Set<number>> = {};
+    if (!chatId) {
+      return messages.sort((a, b) => a.order - b.order);
+    }
+    const included: Record<string, Set<number>> = {};
     for (const m of messages) {
-      if (!included[m.chatId]) {
-        included[m.chatId] = new Set();
+      const searchId = m.chatId ?? m.userId!;
+      if (!included[searchId]) {
+        included[searchId] = new Set();
       }
-      included[m.chatId].add(m.order!);
+      included[searchId].add(m.order!);
     }
-    const ranges: Record<Id<"chats">, Doc<"messages">[]> = {};
+    const ranges: Record<string, Doc<"messages">[]> = {};
     const { before, after } = args.messageRange;
     for (const m of messages) {
+      const searchId = m.chatId ?? m.userId!;
       const order = m.order!;
       let earliest = order - before;
       let latest = order + after;
       for (; earliest <= latest; earliest++) {
-        if (!included[m.chatId].has(earliest)) {
+        if (!included[searchId].has(earliest)) {
           break;
         }
       }
       for (; latest >= earliest; latest--) {
-        if (!included[m.chatId].has(latest)) {
+        if (!included[searchId].has(latest)) {
           break;
         }
       }
       for (let i = earliest; i <= latest; i++) {
-        included[m.chatId].add(i);
+        included[searchId].add(i);
       }
       if (earliest !== latest) {
-        const surrounding = await ctx.db
-          .query("messages")
-          .withIndex("chatId_status_tool_order_stepOrder", (q) =>
-            q
-              .eq("chatId", m.chatId)
-              .eq("status", "success")
-              .eq("tool", false)
-              .gt("order", earliest)
-              .lt("order", latest)
-          )
-          .collect();
-        if (!ranges[m.chatId]) {
-          ranges[m.chatId] = [];
+        if (m.chatId) {
+          const surrounding = await ctx.db
+            .query("messages")
+            .withIndex("chatId_status_tool_order_stepOrder", (q) =>
+              q
+                .eq("chatId", m.chatId)
+                .eq("status", "success")
+                .eq("tool", false)
+                .gt("order", earliest)
+                .lt("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)
+                .gt("order", earliest)
+                .lt("order", latest)
+            )
+            .collect();
+          if (!ranges[searchId]) {
+            ranges[searchId] = [];
+          }
+          ranges[searchId].push(...surrounding);
         }
-        ranges[m.chatId].push(...surrounding);
       }
     }
     return Object.values(ranges)

package/src/component/schema.ts

@@ -22,7 +22,7 @@ export const schema = defineSchema({
   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)
-    chatId: v.id("chats"),
+    chatId: v.optional(v.id("chats")),
     threadId: v.optional(v.id("messages")),
     stepId: v.optional(v.id("steps")),
     agentName: v.optional(v.string()),
@@ -33,8 +33,8 @@ export const schema = defineSchema({
     // TODO: add sub-messages back in? or be able to skip them?
     tool: v.boolean(),
     // Repeats until a non-tool message.
-    // Set when the message is finished
-    order: v.optional(v.number()),
+    // Unset if it's not in a chat.
+    order: v.number(),
     stepOrder: v.optional(v.number()),
     fileId: v.optional(v.id("files")),
     status: vMessageStatus,
@@ -48,6 +48,13 @@ 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_threadId_order_stepOrder", [