@convex-dev/agent 0.0.1-alpha.3 → 0.0.1

package/src/component/schema.ts

@@ -17,12 +17,12 @@ export const schema = defineSchema({
     // the parent thread(s). There are multiple if the thread is a merging of
     // multiple threads.
     parentThreadIds: v.optional(v.array(v.id("threads"))),
-  }).index("status_userId_order", ["status", "userId", "order"]),
+  }).index("userId_status_order", ["userId", "status", "order"]),
   // 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)
-    threadId: v.optional(v.id("threads")),
+    threadId: v.id("threads"),
     parentMessageId: v.optional(v.id("messages")),
     stepId: v.optional(v.id("steps")),
     agentName: v.optional(v.string()),
@@ -35,7 +35,7 @@ export const schema = defineSchema({
     // Repeats until a non-tool message.
     // Unset if it's not in a thread.
     order: v.number(),
-    stepOrder: v.optional(v.number()),
+    stepOrder: v.number(),
     fileId: v.optional(v.id("files")),
     status: vMessageStatus,
   })
@@ -92,6 +92,16 @@ export const schema = defineSchema({
       "stepOrder",
     ]),
 
+  memories: defineTable({
+    threadId: v.optional(v.id("threads")),
+    userId: v.optional(v.string()),
+    memory: v.string(),
+    embeddingId: v.optional(vVectorId),
+  })
+    .index("threadId", ["threadId"])
+    .index("userId", ["userId"])
+    .index("embeddingId", ["embeddingId"]),
+
   files: defineTable({
     storageId: v.string(),
     hash: v.string(),

package/src/component/vector/index.ts

@@ -1,7 +1,20 @@
 import { paginator } from "convex-helpers/server/pagination";
 import { v } from "convex/values";
-import { mutation, query } from "../_generated/server";
-import { getVectorTableName, vVectorDimension, vVectorId } from "./tables";
+import {
+  ActionCtx,
+  mutation,
+  MutationCtx,
+  query,
+  QueryCtx,
+} from "../_generated/server";
+import {
+  EmbeddingsWithoutDenormalizedFields,
+  getVectorTableName,
+  VectorDimension,
+  vEmbeddingsWithoutDenormalizedFields,
+  vVectorDimension,
+  vVectorId,
+} from "./tables";
 import schema from "../schema";
 import { mergedStream } from "convex-helpers/server/stream";
 import { stream } from "convex-helpers/server/stream";
@@ -10,6 +23,7 @@ export const paginate = query({
   args: {
     vectorDimension: vVectorDimension,
     targetModel: v.string(),
+    table: v.optional(v.string()),
     cursor: v.optional(v.string()),
     limit: v.number(),
   },
@@ -23,8 +37,11 @@ export const paginate = query({
     const vectors = await paginator(ctx.db, schema)
       .query(tableName)
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      .withIndex("model_kind_threadId" as any, (q) =>
-        q.eq("model", args.targetModel)
+      .withIndex("model_table_threadId" as any, (q) =>
+        args.table
+          ? // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            (q.eq("model", args.targetModel) as any).eq("table", args.table)
+          : q.eq("model", args.targetModel)
       )
       .paginate({
         cursor: args.cursor ?? null,
@@ -82,36 +99,58 @@ export const deleteBatchForThread = mutation({
 export const insertBatch = mutation({
   args: {
     vectorDimension: vVectorDimension,
-    vectors: v.array(
-      v.object({
-        model: v.string(),
-        kind: v.union(v.literal("thread"), v.literal("memory")),
-        userId: v.optional(v.string()),
-        threadId: v.optional(v.string()),
-        vector: v.array(v.number()),
-      })
-    ),
+    vectors: v.array(vEmbeddingsWithoutDenormalizedFields),
   },
   returns: v.null(),
   handler: async (ctx, args) => {
     await Promise.all(
-      args.vectors.map((v) =>
-        ctx.db.insert(getVectorTableName(args.vectorDimension), {
-          model: v.model,
-          kind: v.kind,
-          userId: v.userId,
-          threadId: v.threadId,
-          vector: v.vector,
-          model_kind_userId: v.userId ? [v.model, v.kind, v.userId] : undefined,
-          model_kind_threadId: v.threadId
-            ? [v.model, v.kind, v.threadId]
-            : undefined,
-        })
-      )
+      args.vectors.map((v) => insertVector(ctx, args.vectorDimension, v))
     );
   },
 });
 
+export async function insertVector(
+  ctx: MutationCtx,
+  dimension: VectorDimension,
+  v: EmbeddingsWithoutDenormalizedFields
+) {
+  return ctx.db.insert(getVectorTableName(dimension), {
+    ...v,
+    model_table_userId: v.userId ? [v.model, v.table, v.userId] : undefined,
+    model_table_threadId: v.threadId
+      ? [v.model, v.table, v.threadId]
+      : undefined,
+  });
+}
+
+export function searchVectors(
+  ctx: ActionCtx,
+  vector: number[],
+  args: {
+    dimension: VectorDimension;
+    model: string;
+    table: string;
+    userId?: string;
+    threadId?: string;
+    limit?: number;
+  }
+) {
+  const tableName = getVectorTableName(args.dimension);
+  return ctx.vectorSearch(tableName, "vector", {
+    vector,
+    // TODO: to support more tables, add more "OR" clauses for each.
+    filter: (q) =>
+      args.userId
+        ? q.eq("model_table_userId", [args.model, args.table, args.userId])
+        : q.eq("model_table_threadId", [
+            args.model,
+            args.table,
+            args.threadId!,
+          ]),
+    limit: args.limit,
+  });
+}
+
 export const updateBatch = mutation({
   args: {
     vectors: v.array(

package/src/component/vector/tables.ts

@@ -1,37 +1,69 @@
+import { omit } from "convex-helpers";
 import { literals } from "convex-helpers/validators";
 import {
   defineTable,
   GenericTableSearchIndexes,
+  SchemaDefinition,
   TableDefinition,
 } from "convex/server";
-import { GenericId, ObjectType, v, VId, VObject, VUnion } from "convex/values";
+import {
+  GenericId,
+  Infer,
+  ObjectType,
+  v,
+  VId,
+  VObject,
+  VUnion,
+} from "convex/values";
 import { QueryCtx } from "../_generated/server";
 
 // We only generate embeddings for non-tool, non-system messages
 const embeddings = {
   model: v.string(),
-  kind: v.union(v.literal("thread"), v.literal("memory")),
+  // What table it's stored in. (usually messages or memories)
+  table: v.string(),
   userId: v.optional(v.string()),
   threadId: v.optional(v.string()),
   // not set for private threads
-  model_kind_userId: v.optional(v.array(v.string())),
-  model_kind_threadId: v.optional(v.array(v.string())),
+  model_table_userId: v.optional(v.array(v.string())),
+  model_table_threadId: v.optional(v.array(v.string())),
   vector: v.array(v.number()),
 };
 
+export const vEmbeddingsWithoutDenormalizedFields = v.object(
+  omit(embeddings, ["model_table_userId", "model_table_threadId"])
+);
+export type EmbeddingsWithoutDenormalizedFields = Infer<
+  typeof vEmbeddingsWithoutDenormalizedFields
+>;
+
 function table<D extends number>(dimensions: D): Table<D> {
   return defineTable(embeddings)
     .vectorIndex("vector", {
       vectorField: "vector",
       dimensions,
-      filterFields: ["model_kind_userId", "model_kind_threadId"],
+      filterFields: ["model_table_userId", "model_table_threadId"],
     })
-    .index("model_kind_threadId", ["model", "kind", "threadId"]);
+    .index("model_table_threadId", ["model", "table", "threadId"]);
 }
 
+export type VectorSchema = SchemaDefinition<
+  { [key in VectorTableName]: Table<128> },
+  true
+>;
+
 export const VectorDimensions = [
   128, 256, 512, 768, 1024, 1536, 2048, 3072, 4096,
 ] as const;
+export function validateVectorDimension(
+  dimension: number
+): asserts dimension is VectorDimension {
+  if (!VectorDimensions.includes(dimension as VectorDimension)) {
+    throw new Error(
+      `Unsupported vector dimension${dimension}. Supported: ${VectorDimensions.join(", ")}`
+    );
+  }
+}
 export type VectorDimension = (typeof VectorDimensions)[number];
 export const VectorTableNames = VectorDimensions.map(
   (d) => `embeddings_${d}`

package/src/validators.ts

@@ -1,4 +1,5 @@
 import { Infer, v } from "convex/values";
+import { vVectorDimension } from "./component/vector/tables";
 
 // const deprecated = v.optional(v.any()) as unknown as VNull<unknown, "optional">;
 
@@ -243,19 +244,18 @@ export const vStep = v.object({
 });
 export type Step = Infer<typeof vStep>;
 
-export const vStepWithMessagesWithFileAndId = v.object({
+export const vStepWithMessages = v.object({
   step: vStep,
   messages: v.array(vMessageWithFileAndId),
 });
-export type StepWithMessagesWithFileAndId = Infer<
-  typeof vStepWithMessagesWithFileAndId
->;
+export type StepWithMessagesWithFileAndId = Infer<typeof vStepWithMessages>;
 
 export const vSearchOptions = v.object({
   vector: v.optional(v.array(v.number())),
   vectorModel: v.optional(v.string()),
   text: v.optional(v.string()),
   limit: v.number(),
+  vectorScoreThreshold: v.optional(v.number()),
   messageRange: v.optional(v.object({ before: v.number(), after: v.number() })),
 });
 export type SearchOptions = Infer<typeof vSearchOptions>;
@@ -291,3 +291,10 @@ export const vObjectArgs = v.object({
   prompt: v.optional(v.string()),
   messages: v.optional(v.array(vMessage)),
 });
+
+export const vEmbeddingsWithMetadata = v.object({
+  vectors: v.array(v.union(v.array(v.number()), v.null())),
+  dimension: vVectorDimension,
+  model: v.string(),
+});
+export type EmbeddingsWithMetadata = Infer<typeof vEmbeddingsWithMetadata>;