@convex-dev/rag 0.1.7

package/src/component/embeddings/index.test.ts

@@ -0,0 +1,482 @@
+/// <reference types="vite/client" />
+
+import { describe, expect, test } from "vitest";
+import { convexTest } from "convex-test";
+import schema, { v } from "../schema.js";
+import { modules } from "../setup.test.js";
+import { insertEmbedding, searchEmbeddings } from "./index.js";
+import { vectorWithImportanceDimension } from "./importance.js";
+import { action } from "../_generated/server.js";
+import { anyApi, type ApiFromModules } from "convex/server";
+
+export const search = action({
+  args: {
+    embedding: v.array(v.number()),
+    namespaceId: v.id("namespaces"),
+    filters: v.array(v.any()),
+    limit: v.number(),
+  },
+  handler: async (ctx, args) => {
+    return searchEmbeddings(ctx, args);
+  },
+});
+
+const testApi: ApiFromModules<{
+  fns: {
+    search: typeof search;
+  };
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+}>["fns"] = anyApi["embeddings"]["index.test"] as any;
+
+describe("embeddings", () => {
+  test("insertEmbedding with no filters or importance works", async () => {
+    const t = convexTest(schema, modules);
+
+    // Create a namespace first
+    const namespaceId = await t.run(async (ctx) => {
+      return ctx.db.insert("namespaces", {
+        namespace: "test-namespace",
+        version: 1,
+        modelId: "test-model",
+        dimension: 128,
+        filterNames: [],
+        status: { kind: "ready" },
+      });
+    });
+
+    // Create a simple 128-dimension embedding
+    const embedding = Array(128).fill(0.1);
+
+    // Insert embedding without filters or importance
+    const vectorId = await t.run(async (ctx) => {
+      return insertEmbedding(ctx, embedding, namespaceId, undefined, undefined);
+    });
+
+    expect(vectorId).toBeDefined();
+
+    // Verify the vector was inserted correctly
+    const insertedVector = await t.run(async (ctx) => {
+      return ctx.db.get(vectorId);
+    });
+
+    expect(insertedVector).toBeDefined();
+    expect(insertedVector!.namespaceId).toBe(namespaceId);
+    expect(insertedVector!.vector).toHaveLength(
+      vectorWithImportanceDimension(128)
+    );
+    expect(insertedVector!.filter0).toBeUndefined();
+    expect(insertedVector!.filter1).toBeUndefined();
+    expect(insertedVector!.filter2).toBeUndefined();
+    expect(insertedVector!.filter3).toBeUndefined();
+  });
+
+  test("insertEmbedding with importance modifies the vector", async () => {
+    const t = convexTest(schema, modules);
+
+    const namespaceId = await t.run(async (ctx) => {
+      return ctx.db.insert("namespaces", {
+        namespace: "test-namespace-importance",
+        version: 1,
+        modelId: "test-model",
+        dimension: 128,
+        filterNames: [],
+        status: { kind: "ready" },
+      });
+    });
+
+    const embedding = Array(128).fill(0.1);
+    const importance = 0.5;
+
+    // Insert embedding with importance
+    const vectorId = await t.run(async (ctx) => {
+      return insertEmbedding(
+        ctx,
+        embedding,
+        namespaceId,
+        importance,
+        undefined
+      );
+    });
+
+    const insertedVector = await t.run(async (ctx) => {
+      return ctx.db.get(vectorId);
+    });
+
+    expect(insertedVector).toBeDefined();
+    expect(insertedVector!.vector).toHaveLength(129);
+
+    // The importance should affect the vector - it should not be the same as without importance
+    const vectorWithoutImportance = await t.run(async (ctx) => {
+      return insertEmbedding(ctx, embedding, namespaceId, undefined, undefined);
+    });
+
+    const vectorWithoutImportanceData = await t.run(async (ctx) => {
+      return ctx.db.get(vectorWithoutImportance);
+    });
+
+    // Vectors should be different due to importance scaling
+    expect(insertedVector!.vector).not.toEqual(
+      vectorWithoutImportanceData!.vector
+    );
+
+    // The last element should be the weight: sqrt(1 - importance^2)
+    const expectedWeight = Math.sqrt(1 - importance ** 2);
+    expect(insertedVector!.vector[128]).toBeCloseTo(expectedWeight, 5);
+  });
+
+  test("search for vectors sorted by importance when identical otherwise", async () => {
+    const t = convexTest(schema, modules);
+
+    const namespaceId = await t.run(async (ctx) => {
+      return ctx.db.insert("namespaces", {
+        namespace: "importance-sort-test",
+        version: 1,
+        modelId: "test-model",
+        dimension: 128,
+        filterNames: [],
+        status: { kind: "ready" },
+      });
+    });
+
+    const embedding = Array(128).fill(0.1);
+
+    // Insert same embedding with different importance levels
+    await t.run(async (ctx) => {
+      await insertEmbedding(ctx, embedding, namespaceId, 0.2, undefined); // Low importance
+      await insertEmbedding(ctx, embedding, namespaceId, 0.8, undefined); // High importance
+      await insertEmbedding(ctx, embedding, namespaceId, 0.5, undefined); // Medium importance
+    });
+
+    // Search for the vectors
+    const results = await t.action(testApi.search, {
+      embedding,
+      namespaceId,
+      filters: [],
+      limit: 10,
+    });
+
+    expect(results).toHaveLength(3);
+
+    // Results should be sorted by similarity (which correlates with importance)
+    // Higher importance vectors should have higher similarity scores
+    expect(results[0]._score).toBeGreaterThan(results[1]._score);
+    expect(results[1]._score).toBeGreaterThan(results[2]._score);
+  });
+
+  test("filters are added to the correct field", async () => {
+    const t = convexTest(schema, modules);
+
+    const namespaceId = await t.run(async (ctx) => {
+      return ctx.db.insert("namespaces", {
+        namespace: "filter-test",
+        version: 1,
+        modelId: "test-model",
+        dimension: 128,
+        filterNames: ["category", "priority", "status", "author"],
+        status: { kind: "ready" },
+      });
+    });
+
+    const embedding = Array(128).fill(0.1);
+
+    // Insert embedding with filter on position 0
+    const vectorId0 = await t.run(async (ctx) => {
+      return insertEmbedding(ctx, embedding, namespaceId, undefined, {
+        0: "entries",
+      });
+    });
+
+    // Insert embedding with filter on position 2
+    const vectorId2 = await t.run(async (ctx) => {
+      return insertEmbedding(ctx, embedding, namespaceId, undefined, {
+        2: "active",
+      });
+    });
+
+    // Verify filters are in correct fields
+    const vector0 = await t.run(async (ctx) => ctx.db.get(vectorId0));
+    const vector2 = await t.run(async (ctx) => ctx.db.get(vectorId2));
+
+    expect(vector0!.filter0).toEqual([namespaceId, "entries"]);
+    expect(vector0!.filter1).toBeUndefined();
+    expect(vector0!.filter2).toBeUndefined();
+    expect(vector0!.filter3).toBeUndefined();
+
+    expect(vector2!.filter0).toBeUndefined();
+    expect(vector2!.filter1).toBeUndefined();
+    expect(vector2!.filter2).toEqual([namespaceId, "active"]);
+    expect(vector2!.filter3).toBeUndefined();
+  });
+
+  test("embeddings have namespace prefixed on filter fields", async () => {
+    const t = convexTest(schema, modules);
+
+    const namespace1Id = await t.run(async (ctx) => {
+      return ctx.db.insert("namespaces", {
+        namespace: "namespace1",
+        version: 1,
+        modelId: "test-model",
+        dimension: 128,
+        filterNames: ["type"],
+        status: { kind: "ready" },
+      });
+    });
+
+    const namespace2Id = await t.run(async (ctx) => {
+      return ctx.db.insert("namespaces", {
+        namespace: "namespace2",
+        version: 1,
+        modelId: "test-model",
+        dimension: 128,
+        filterNames: ["type"],
+        status: { kind: "ready" },
+      });
+    });
+
+    const embedding = Array(128).fill(0.1);
+
+    // Insert same filter value in different namespaces
+    const vector1Id = await t.run(async (ctx) => {
+      return insertEmbedding(ctx, embedding, namespace1Id, undefined, {
+        0: "article",
+      });
+    });
+
+    const vector2Id = await t.run(async (ctx) => {
+      return insertEmbedding(ctx, embedding, namespace2Id, undefined, {
+        0: "article",
+      });
+    });
+
+    const vector1 = await t.run(async (ctx) => ctx.db.get(vector1Id));
+    const vector2 = await t.run(async (ctx) => ctx.db.get(vector2Id));
+
+    // Both have the same filter value but different namespace prefixes
+    expect(vector1!.filter0).toEqual([namespace1Id, "article"]);
+    expect(vector2!.filter0).toEqual([namespace2Id, "article"]);
+    expect(vector1!.filter0).not.toEqual(vector2!.filter0);
+  });
+
+  test("search without filters returns only vectors in the target namespace", async () => {
+    const t = convexTest(schema, modules);
+
+    const namespace1Id = await t.run(async (ctx) => {
+      return ctx.db.insert("namespaces", {
+        namespace: "namespace1",
+        version: 1,
+        modelId: "test-model",
+        dimension: 128,
+        filterNames: [],
+        status: { kind: "ready" },
+      });
+    });
+
+    const namespace2Id = await t.run(async (ctx) => {
+      return ctx.db.insert("namespaces", {
+        namespace: "namespace2",
+        version: 1,
+        modelId: "test-model",
+        dimension: 128,
+        filterNames: [],
+        status: { kind: "ready" },
+      });
+    });
+
+    const embedding = Array(128).fill(0.1);
+
+    // Insert vectors in both namespaces
+    await t.run(async (ctx) => {
+      await insertEmbedding(ctx, embedding, namespace1Id, undefined, undefined);
+      await insertEmbedding(ctx, embedding, namespace1Id, undefined, undefined);
+      await insertEmbedding(ctx, embedding, namespace2Id, undefined, undefined);
+    });
+
+    // Search in namespace1 only
+    const results1 = await t.action(testApi.search, {
+      embedding,
+      namespaceId: namespace1Id,
+      filters: [],
+      limit: 10,
+    });
+
+    // Search in namespace2 only
+    const results2 = await t.action(testApi.search, {
+      embedding,
+      namespaceId: namespace2Id,
+      filters: [],
+      limit: 10,
+    });
+
+    expect(results1).toHaveLength(2);
+    expect(results2).toHaveLength(1);
+
+    // All results should be from the correct namespace
+    for (const result of results1) {
+      const vector = await t.run(async (ctx) => ctx.db.get(result._id));
+      expect(vector!.namespaceId).toBe(namespace1Id);
+    }
+
+    for (const result of results2) {
+      const vector = await t.run(async (ctx) => ctx.db.get(result._id));
+      expect(vector!.namespaceId).toBe(namespace2Id);
+    }
+  });
+
+  test("search with filters returns only matching vectors in namespace", async () => {
+    const t = convexTest(schema, modules);
+
+    const namespaceId = await t.run(async (ctx) => {
+      return ctx.db.insert("namespaces", {
+        namespace: "filtered-search",
+        version: 1,
+        modelId: "test-model",
+        dimension: 128,
+        filterNames: ["category", "status"],
+        status: { kind: "ready" },
+      });
+    });
+
+    const embedding = Array(128).fill(0.1);
+
+    // Insert vectors with different filter combinations
+    await t.run(async (ctx) => {
+      await insertEmbedding(ctx, embedding, namespaceId, undefined, {
+        0: "articles",
+      });
+      await insertEmbedding(ctx, embedding, namespaceId, undefined, {
+        0: "blogs",
+      });
+      await insertEmbedding(ctx, embedding, namespaceId, undefined, {
+        1: "published",
+      });
+      await insertEmbedding(ctx, embedding, namespaceId, undefined, {
+        0: "articles",
+        1: "draft",
+      });
+      await insertEmbedding(ctx, embedding, namespaceId, undefined, undefined); // No filters
+    });
+
+    // Search for articles only
+    const articlesResults = await t.action(testApi.search, {
+      embedding,
+      namespaceId,
+      filters: [{ 0: "articles" }],
+      limit: 10,
+    });
+
+    expect(articlesResults).toHaveLength(2); // Two vectors with category "articles"
+
+    // Search for published status only
+    const publishedResults = await t.action(testApi.search, {
+      embedding,
+      namespaceId,
+      filters: [{ 1: "published" }],
+      limit: 10,
+    });
+
+    expect(publishedResults).toHaveLength(1); // One vector with status "published"
+  });
+
+  test("multiple filters perform OR operation", async () => {
+    const t = convexTest(schema, modules);
+
+    const namespaceId = await t.run(async (ctx) => {
+      return ctx.db.insert("namespaces", {
+        namespace: "multi-filter-or",
+        version: 1,
+        modelId: "test-model",
+        dimension: 128,
+        filterNames: ["category", "priority"],
+        status: { kind: "ready" },
+      });
+    });
+
+    const embedding = Array(128).fill(0.1);
+
+    // Insert vectors with different filter values
+    await t.run(async (ctx) => {
+      await insertEmbedding(ctx, embedding, namespaceId, undefined, {
+        0: "articles",
+      });
+      await insertEmbedding(ctx, embedding, namespaceId, undefined, {
+        0: "blogs",
+      });
+      await insertEmbedding(ctx, embedding, namespaceId, undefined, {
+        1: "high",
+      });
+      await insertEmbedding(ctx, embedding, namespaceId, undefined, {
+        1: "low",
+      });
+      await insertEmbedding(ctx, embedding, namespaceId, undefined, undefined); // No filters
+    });
+
+    // Search with OR filters: articles OR high priority
+    const orResults = await t.action(testApi.search, {
+      embedding,
+      namespaceId,
+      filters: [
+        { 0: "articles" }, // category = articles
+        { 1: "high" }, // OR priority = high
+      ],
+      limit: 10,
+    });
+
+    expect(orResults).toHaveLength(2); // Should match both "articles" and "high priority" vectors
+
+    // Verify the results contain the expected filters
+    const vectorIds = orResults.map((r) => r._id);
+    const vectors = await t.run(async (ctx) => {
+      return Promise.all(vectorIds.map((id) => ctx.db.get(id)));
+    });
+
+    const hasArticles = vectors.some((v) => v!.filter0?.[1] === "articles");
+    const hasHighPriority = vectors.some((v) => v!.filter1?.[1] === "high");
+
+    expect(hasArticles).toBe(true);
+    expect(hasHighPriority).toBe(true);
+  });
+
+  test("searchEmbeddings", async () => {
+    const t = convexTest(schema, modules);
+
+    const namespaceId = await t.run(async (ctx) => {
+      return ctx.db.insert("namespaces", {
+        namespace: "search-test",
+        version: 1,
+        modelId: "test-model",
+        dimension: 128,
+        filterNames: [],
+        status: { kind: "ready" },
+      });
+    });
+
+    const embedding1 = Array(128).fill(0.1);
+    embedding1[0] = 1;
+    const embedding2 = Array(128).fill(0.1);
+    embedding2[0] = 0;
+    const searchEmbedding = Array(128).fill(0.1);
+    searchEmbedding[0] = 0.8; // Closer to embedding1
+
+    // Insert two different embeddings
+    await t.run(async (ctx) => {
+      await insertEmbedding(ctx, embedding1, namespaceId, undefined, undefined);
+      await insertEmbedding(ctx, embedding2, namespaceId, undefined, undefined);
+    });
+
+    // Search should return results ordered by similarity
+    const results = await t.action(testApi.search, {
+      embedding: searchEmbedding,
+      namespaceId,
+      filters: [],
+      limit: 10,
+    });
+
+    expect(results).toHaveLength(2);
+    expect(results[0]._score).toBeGreaterThan(results[1]._score);
+
+    // The first result should be more similar to embedding1 (0.1) than embedding2 (0.2)
+    // since searchEmbedding (0.15) is closer to 0.1
+    const firstVector = await t.run(async (ctx) => ctx.db.get(results[0]._id));
+    expect(firstVector).toBeDefined();
+  });
+});

package/src/component/embeddings/index.ts

@@ -0,0 +1,99 @@
+/**
+ * This file is the interface for interacting with vectors.
+ * It translates from embeddings to the underlying vector storage and search.
+ * It modifies embeddings to include importance.
+ * The outer world deals with filters with user names.
+ * The underlying vector storage has its own names.
+ * This file takes in numbered filters (0-3) to translate without knowing about
+ * user names.
+ */
+import type { Id } from "../_generated/dataModel.js";
+import { type ActionCtx, type MutationCtx } from "../_generated/server.js";
+import { getVectorTableName, validateVectorDimension } from "./tables.js";
+import { filterFieldsFromNumbers, type NumberedFilter } from "../filters.js";
+import { searchVector, vectorWithImportance } from "./importance.js";
+
+// TODO: wait to see if this is needed.
+// export const insertBatch = mutation({
+//   args: {
+//     vectorDimension: vVectorDimension,
+//     vectors: v.array(
+//       v.object({
+//         vector: v.array(v.number()),
+//         namespace: v.id("namespaces"),
+//         importance: v.optional(v.number()),
+//         filters: v.optional(v.any()),
+//       })
+//     ),
+//   },
+//   returns: v.array(vVectorId),
+//   handler: async (ctx, args) => {
+//     return Promise.all(
+//       args.vectors.map(async (vector) =>
+//         insertEmbedding(
+//           ctx,
+//           vector.vector,
+//           vector.namespace,
+//           vector.importance,
+//           vector.filters
+//         )
+//       )
+//     );
+//   },
+// });
+
+
+export async function insertEmbedding(
+  ctx: MutationCtx,
+  embedding: number[],
+  namespaceId: Id<"namespaces">,
+  importance: number | undefined,
+  filters: NumberedFilter | undefined
+) {
+  const filterFields = filterFieldsFromNumbers(namespaceId, filters);
+  const dimension = validateVectorDimension(embedding.length);
+  return ctx.db.insert(getVectorTableName(dimension), {
+    namespaceId,
+    vector: vectorWithImportance(embedding, importance ?? 1),
+    ...filterFields,
+  });
+}
+
+export async function searchEmbeddings(
+  ctx: ActionCtx,
+  {
+    embedding,
+    namespaceId,
+    filters,
+    limit,
+  }: {
+    embedding: number[];
+    namespaceId: Id<"namespaces">;
+    // NOTE: Current vector search supports OR logic between filters
+    // Multiple filters will match if ANY condition is met
+    // e.g. [{3: filter3}, {1: filter1}, {2: filter2}] will match if any of
+    // filter3, filter1, or filter2 is present.
+    filters: Array<NumberedFilter>;
+    limit: number;
+  }
+) {
+  const dimension = validateVectorDimension(embedding.length);
+  const tableName = getVectorTableName(dimension);
+  const orFilters = filters.flatMap((filter) =>
+    filterFieldsFromNumbers(namespaceId, filter)
+  );
+  return ctx.vectorSearch(tableName, "vector", {
+    vector: searchVector(embedding),
+    filter: (q) =>
+      orFilters.length === 0
+        ? q.eq("namespaceId", namespaceId)
+        : q.or(
+            ...orFilters.flatMap((namedFilter) =>
+              Object.entries(namedFilter).map(([filterField, filter]) =>
+                q.eq(filterField as keyof (typeof orFilters)[number], filter)
+              )
+            )
+          ),
+    limit,
+  });
+}

package/src/component/embeddings/tables.ts

@@ -0,0 +1,114 @@
+import { literals } from "convex-helpers/validators";
+import {
+  defineTable,
+  type GenericTableSearchIndexes,
+  type SchemaDefinition,
+  type TableDefinition,
+} from "convex/server";
+import {
+  type GenericId,
+  type ObjectType,
+  v,
+  type VId,
+  type VObject,
+  type VUnion,
+} from "convex/values";
+import { vectorWithImportanceDimension } from "./importance.js";
+import { allFilterFieldNames, vAllFilterFields } from "../filters.js";
+
+// We only generate embeddings for non-tool, non-system messages
+const embeddingsFields = {
+  vector: v.array(v.number()),
+  ...vAllFilterFields,
+};
+
+function table(dimensions: VectorDimension): Table {
+  return defineTable(embeddingsFields).vectorIndex("vector", {
+    vectorField: "vector",
+    dimensions: vectorWithImportanceDimension(dimensions),
+    filterFields: allFilterFieldNames,
+  });
+}
+
+type Table = TableDefinition<
+  VObject<ObjectType<typeof embeddingsFields>, typeof embeddingsFields>,
+  { model_table_threadId: ["model", "table", "threadId", "_creationTime"] },
+  GenericTableSearchIndexes,
+  VectorIndex
+>;
+
+type VectorIndex = {
+  vector: {
+    vectorField: "vector";
+    dimensions: number;
+    filterFields: string;
+  };
+};
+
+export type VectorSchema = SchemaDefinition<
+  { [key in VectorTableName]: Table },
+  true
+>;
+
+export const VectorDimensions = [
+  128, 256, 512, 768, 1024, 1408, 1536, 2048, 3072, 4096,
+] as const;
+
+export function assertVectorDimension(
+  dimension: number
+): asserts dimension is VectorDimension {
+  if (!VectorDimensions.includes(dimension as VectorDimension)) {
+    throw new Error(
+      `Unsupported vector dimension${dimension}. Supported: ${VectorDimensions.join(", ")}`
+    );
+  }
+}
+
+export function validateVectorDimension(dimension: number): VectorDimension {
+  if (!VectorDimensions.includes(dimension as VectorDimension)) {
+    throw new Error(
+      `Unsupported vector dimension${dimension}. Supported: ${VectorDimensions.join(", ")}`
+    );
+  }
+  return dimension as VectorDimension;
+}
+export type VectorDimension = (typeof VectorDimensions)[number];
+export const VectorTableNames = VectorDimensions.map(
+  (d) => `vectors_${d}`
+) as `vectors_${(typeof VectorDimensions)[number]}`[];
+export type VectorTableName = (typeof VectorTableNames)[number];
+export type VectorTableId = GenericId<(typeof VectorTableNames)[number]>;
+
+export const vVectorDimension = literals(...VectorDimensions);
+export const vVectorTableName = literals(...VectorTableNames);
+export const vVectorId = v.union(
+  ...VectorTableNames.map((name) => v.id(name))
+) as VUnion<
+  GenericId<(typeof VectorTableNames)[number]>,
+  VId<(typeof VectorTableNames)[number]>[]
+>;
+
+export function getVectorTableName(dimension: VectorDimension) {
+  return `vectors_${dimension}` as VectorTableName;
+}
+// export function getVectorIdInfo(ctx: QueryCtx, id: VectorTableId) {
+//   for (const dimension of VectorDimensions) {
+//     const tableName = getVectorTableName(dimension);
+//     if (ctx.db.normalizeId(tableName, id)) {
+//       return { tableName, dimension };
+//     }
+//   }
+//   throw new Error(`Unknown vector table id: ${id}`);
+// }
+
+const tables: {
+  [K in keyof typeof VectorDimensions &
+    number as `vectors_${(typeof VectorDimensions)[K]}`]: Table;
+} = Object.fromEntries(
+  VectorDimensions.map((dimensions) => [
+    `vectors_${dimensions}`,
+    table(dimensions),
+  ])
+) as Record<`vectors_${(typeof VectorDimensions)[number]}`, Table>;
+
+export default tables;