@convex-dev/rag 0.6.1 → 0.7.1

package/src/component/search.test.ts

@@ -3,7 +3,7 @@
 import { describe, expect, test } from "vitest";
 import { convexTest, type TestConvex } from "convex-test";
 import schema from "./schema.js";
-import { api } from "./_generated/api.js";
+import { api, internal } from "./_generated/api.js";
 import { modules } from "./setup.test.js";
 import { insertChunks } from "./chunks.js";
 import type { Id } from "./_generated/dataModel.js";
@@ -442,4 +442,306 @@ describe("search", () => {
       );
     }
   });
+
+  describe("hybrid search", () => {
+    function createSearchableChunks(texts: string[], baseEmbedding = 0.1) {
+      return texts.map((text, i) => ({
+        content: { text, metadata: { index: i } },
+        embedding: [...Array(127).fill(0.01), baseEmbedding + i * 0.01],
+        searchableText: text,
+      }));
+    }
+
+    test("textSearch internal query finds chunks by text content", async () => {
+      const t = convexTest(schema, modules);
+      const namespaceId = await setupTestNamespace(t);
+      const entryId = await setupTestEntry(t, namespaceId);
+
+      const chunks = createSearchableChunks([
+        "The quick brown fox jumps over the lazy dog",
+        "A fast red car drives on the highway",
+        "The brown bear sleeps in the forest",
+      ]);
+
+      await t.run(async (ctx) => {
+        await insertChunks(ctx, { entryId, startOrder: 0, chunks });
+      });
+
+      const results = await t.query(internal.search.textSearch, {
+        query: "brown",
+        namespaceId,
+        filters: [],
+        limit: 10,
+      });
+
+      expect(results.length).toBeGreaterThan(0);
+      for (const r of results) {
+        expect(r.entryId).toBe(entryId);
+      }
+    });
+
+    test("textSearch scopes results to the given namespace", async () => {
+      const t = convexTest(schema, modules);
+      const ns1Id = await setupTestNamespace(t, "namespace-1");
+      const ns2Id = await setupTestNamespace(t, "namespace-2");
+      const entry1Id = await setupTestEntry(t, ns1Id, "entry-1");
+      const entry2Id = await setupTestEntry(t, ns2Id, "entry-2");
+
+      await t.run(async (ctx) => {
+        await insertChunks(ctx, {
+          entryId: entry1Id,
+          startOrder: 0,
+          chunks: createSearchableChunks(["alpha bravo charlie"]),
+        });
+        await insertChunks(ctx, {
+          entryId: entry2Id,
+          startOrder: 0,
+          chunks: createSearchableChunks(["alpha delta echo"]),
+        });
+      });
+
+      const ns1Results = await t.query(internal.search.textSearch, {
+        query: "alpha",
+        namespaceId: ns1Id,
+        filters: [],
+        limit: 10,
+      });
+
+      // All results should belong to namespace-1's entry.
+      for (const r of ns1Results) {
+        expect(r.entryId).toBe(entry1Id);
+      }
+    });
+
+    test("textSearch applies numbered filters", async () => {
+      const t = convexTest(schema, modules);
+      const namespaceId = await setupTestNamespace(t, "filtered-ns", 128, [
+        "category",
+      ]);
+
+      const cat1Entry = await setupTestEntry(t, namespaceId, "cat1", 0, [
+        { name: "category", value: "docs" },
+      ]);
+      const cat2Entry = await setupTestEntry(t, namespaceId, "cat2", 0, [
+        { name: "category", value: "blogs" },
+      ]);
+
+      await t.run(async (ctx) => {
+        await insertChunks(ctx, {
+          entryId: cat1Entry,
+          startOrder: 0,
+          chunks: createSearchableChunks(["shared keyword content"]),
+        });
+        await insertChunks(ctx, {
+          entryId: cat2Entry,
+          startOrder: 0,
+          chunks: createSearchableChunks(["shared keyword content"]),
+        });
+      });
+
+      // Filter to "docs" category only (filter index 0 = "category").
+      const results = await t.query(internal.search.textSearch, {
+        query: "shared keyword",
+        namespaceId,
+        filters: [{ 0: "docs" }],
+        limit: 10,
+      });
+
+      expect(results.length).toBeGreaterThan(0);
+      for (const r of results) {
+        expect(r.entryId).toBe(cat1Entry);
+      }
+    });
+
+    test("text-only search returns results via dimension arg", async () => {
+      const t = convexTest(schema, modules);
+      const namespaceId = await setupTestNamespace(t);
+      const entryId = await setupTestEntry(t, namespaceId);
+
+      const chunks = createSearchableChunks([
+        "Machine learning is a subset of artificial intelligence",
+        "Deep learning uses neural networks with many layers",
+        "Natural language processing handles text data",
+      ]);
+
+      await t.run(async (ctx) => {
+        await insertChunks(ctx, { entryId, startOrder: 0, chunks });
+      });
+
+      // Text-only: no embedding, provide dimension instead.
+      const result = await t.action(api.search.search, {
+        namespace: "test-namespace",
+        dimension: 128,
+        modelId: "test-model",
+        filters: [],
+        limit: 10,
+        textQuery: "neural networks",
+      });
+
+      expect(result.results.length).toBeGreaterThan(0);
+      expect(result.entries).toHaveLength(1);
+
+      // Text-only scores are position-based.
+      expect(result.results[0].score).toBe(1.0);
+      for (let i = 1; i < result.results.length; i++) {
+        expect(result.results[i].score).toBeLessThan(
+          result.results[i - 1].score,
+        );
+      }
+    });
+
+    test("hybrid search returns results when textQuery is provided", async () => {
+      const t = convexTest(schema, modules);
+      const namespaceId = await setupTestNamespace(t);
+      const entryId = await setupTestEntry(t, namespaceId);
+
+      const chunks = createSearchableChunks([
+        "Machine learning is a subset of artificial intelligence",
+        "Deep learning uses neural networks with many layers",
+        "Natural language processing handles text data",
+      ]);
+
+      await t.run(async (ctx) => {
+        await insertChunks(ctx, { entryId, startOrder: 0, chunks });
+      });
+
+      const result = await t.action(api.search.search, {
+        namespace: "test-namespace",
+        embedding: [...Array(127).fill(0.01), 0.1],
+        modelId: "test-model",
+        filters: [],
+        limit: 10,
+        textQuery: "neural networks",
+      });
+
+      expect(result.results.length).toBeGreaterThan(0);
+      expect(result.entries).toHaveLength(1);
+
+      // Hybrid scores are position-based (1.0 for top, decreasing linearly).
+      expect(result.results[0].score).toBe(1.0);
+      for (let i = 1; i < result.results.length; i++) {
+        expect(result.results[i].score).toBeLessThan(
+          result.results[i - 1].score,
+        );
+      }
+    });
+
+    test("hybrid search deduplicates results from vector and text paths", async () => {
+      const t = convexTest(schema, modules);
+      const namespaceId = await setupTestNamespace(t);
+      const entryId = await setupTestEntry(t, namespaceId);
+
+      const chunks = createSearchableChunks([
+        "Unique content about quantum computing",
+        "Another chunk about classical physics",
+      ]);
+
+      await t.run(async (ctx) => {
+        await insertChunks(ctx, { entryId, startOrder: 0, chunks });
+      });
+
+      const result = await t.action(api.search.search, {
+        namespace: "test-namespace",
+        embedding: [...Array(127).fill(0.01), 0.1],
+        modelId: "test-model",
+        filters: [],
+        limit: 10,
+        textQuery: "quantum computing",
+      });
+
+      // Each chunk should appear at most once in the results.
+      const entryOrderPairs = result.results.map(
+        (r) => `${r.entryId}:${r.order}`,
+      );
+      const uniquePairs = new Set(entryOrderPairs);
+      expect(uniquePairs.size).toBe(entryOrderPairs.length);
+    });
+
+    test("vector-only search is unchanged when textQuery is not provided", async () => {
+      const t = convexTest(schema, modules);
+      const namespaceId = await setupTestNamespace(t);
+      const entryId = await setupTestEntry(t, namespaceId);
+
+      const targetEmbedding = [...Array(127).fill(0.5), 1];
+      const chunks = [
+        {
+          content: { text: "Target chunk", metadata: {} },
+          embedding: targetEmbedding,
+          searchableText: "Target chunk",
+        },
+        {
+          content: { text: "Other chunk", metadata: {} },
+          embedding: [...Array(127).fill(0.1), 0],
+          searchableText: "Other chunk",
+        },
+      ];
+
+      await t.run(async (ctx) => {
+        await insertChunks(ctx, { entryId, startOrder: 0, chunks });
+      });
+
+      const result = await t.action(api.search.search, {
+        namespace: "test-namespace",
+        embedding: targetEmbedding,
+        modelId: "test-model",
+        filters: [],
+        limit: 10,
+      });
+
+      // Without textQuery, scores should be cosine similarity (not position-based).
+      expect(result.results).toHaveLength(2);
+      expect(result.results[0].score).toBeGreaterThan(result.results[1].score);
+      // Cosine similarity scores are typically between -1 and 1, not exactly 1.0.
+      // Position-based would give exactly 1.0 for the first result.
+      // With cosine similarity the first result can be 1.0 if exact match,
+      // but the second should not follow the linear decrease pattern.
+      expect(result.results[0].content[0].text).toBe("Target chunk");
+    });
+
+    test("textWeight and vectorWeight influence hybrid ranking", async () => {
+      const t = convexTest(schema, modules);
+      const namespaceId = await setupTestNamespace(t);
+      const entryId = await setupTestEntry(t, namespaceId);
+
+      const chunks = createSearchableChunks([
+        "Alpha topic with specific terminology",
+        "Beta topic with different keywords",
+        "Gamma topic about something else entirely",
+      ]);
+
+      await t.run(async (ctx) => {
+        await insertChunks(ctx, { entryId, startOrder: 0, chunks });
+      });
+
+      const embedding = [...Array(127).fill(0.01), 0.1];
+
+      // Search with heavy text weight.
+      const textHeavy = await t.action(api.search.search, {
+        namespace: "test-namespace",
+        embedding,
+        modelId: "test-model",
+        filters: [],
+        limit: 10,
+        textQuery: "specific terminology",
+        textWeight: 10,
+        vectorWeight: 1,
+      });
+
+      // Search with heavy vector weight.
+      const vectorHeavy = await t.action(api.search.search, {
+        namespace: "test-namespace",
+        embedding,
+        modelId: "test-model",
+        filters: [],
+        limit: 10,
+        textQuery: "specific terminology",
+        textWeight: 1,
+        vectorWeight: 10,
+      });
+
+      // Both should return results.
+      expect(textHeavy.results.length).toBeGreaterThan(0);
+      expect(vectorHeavy.results.length).toBeGreaterThan(0);
+    });
+  });
 });

package/src/component/search.ts

@@ -1,20 +1,30 @@
 import { v, type Infer } from "convex/values";
-import { action } from "./_generated/server.js";
+import { action, internalQuery, type QueryCtx } from "./_generated/server.js";
 import { searchEmbeddings } from "./embeddings/index.js";
-import { numberedFiltersFromNamedFilters, vNamedFilter } from "./filters.js";
+import {
+  filterFieldsFromNumbers,
+  numberedFiltersFromNamedFilters,
+  vNamedFilter,
+  type NumberedFilter,
+} from "./filters.js";
 import { internal } from "./_generated/api.js";
 import {
   vEntry,
   vSearchResult,
+  vSearchType,
   type SearchResult,
   type EntryId,
 } from "../shared.js";
-import type { vRangeResult } from "./chunks.js";
+import type { Doc, Id } from "./_generated/dataModel.js";
+import { buildRanges, type vRangeResult } from "./chunks.js";
+import { hybridRank } from "../client/hybridRank.js";
+import { vVectorId, type VectorTableId } from "./embeddings/tables.js";
 
 export const search = action({
   args: {
     namespace: v.string(),
-    embedding: v.array(v.number()),
+    embedding: v.optional(v.array(v.number())),
+    dimension: v.optional(v.number()),
     modelId: v.string(),
     // These are all OR'd together
     filters: v.array(vNamedFilter),
@@ -23,6 +33,10 @@ export const search = action({
     chunkContext: v.optional(
       v.object({ before: v.number(), after: v.number() }),
     ),
+    searchType: v.optional(vSearchType),
+    textQuery: v.optional(v.string()),
+    textWeight: v.optional(v.number()),
+    vectorWeight: v.optional(v.number()),
   },
   returns: v.object({
     results: v.array(vSearchResult),
@@ -36,51 +50,284 @@ export const search = action({
     entries: Infer<typeof vEntry>[];
   }> => {
     const { modelId, embedding, filters, limit } = args;
+    const dimension = embedding?.length ?? args.dimension;
+    if (!dimension) {
+      throw new Error(
+        "Either embedding or dimension must be provided to search.",
+      );
+    }
+
     const namespace = await ctx.runQuery(
       internal.namespaces.getCompatibleNamespace,
       {
         namespace: args.namespace,
         modelId,
-        dimension: embedding.length,
+        dimension,
         filterNames: filters.map((f) => f.name),
       },
     );
     if (!namespace) {
       console.debug(
-        `No compatible namespace found for ${args.namespace} with model ${args.modelId} and dimension ${embedding.length} and filters ${filters.map((f) => f.name).join(", ")}.`,
+        `No compatible namespace found for ${args.namespace} with model ${args.modelId} and dimension ${dimension} and filters ${filters.map((f) => f.name).join(", ")}.`,
       );
       return {
         results: [],
         entries: [],
       };
     }
-    const results = await searchEmbeddings(ctx, {
-      embedding,
-      namespaceId: namespace._id,
-      filters: numberedFiltersFromNamedFilters(filters, namespace.filterNames),
-      limit,
-    });
 
-    const threshold = args.vectorScoreThreshold ?? -1;
-    const aboveThreshold = results.filter((r) => r._score >= threshold);
     const chunkContext = args.chunkContext ?? { before: 0, after: 0 };
-    // TODO: break this up if there are too many results
-    const { ranges, entries } = await ctx.runQuery(
-      internal.chunks.getRangesOfChunks,
+    const numberedFilters = numberedFiltersFromNamedFilters(
+      filters,
+      namespace.filterNames,
+    );
+
+    const hasEmbedding = !!embedding;
+    const hasTextQuery = !!args.textQuery;
+
+    // Vector-only path: return results with cosine similarity scores.
+    if (hasEmbedding && !hasTextQuery) {
+      const vectorResults = await searchEmbeddings(ctx, {
+        embedding,
+        namespaceId: namespace._id,
+        filters: numberedFilters,
+        limit,
+      });
+      const threshold = args.vectorScoreThreshold ?? -1;
+      const aboveThreshold = vectorResults.filter((r) => r._score >= threshold);
+      // TODO: break this up if there are too many results
+      const { ranges, entries } = await ctx.runQuery(
+        internal.chunks.getRangesOfChunks,
+        {
+          embeddingIds: aboveThreshold.map((r) => r._id),
+          chunkContext,
+        },
+      );
+      return {
+        results: ranges
+          .map((r, i) => publicSearchResult(r, aboveThreshold[i]._score))
+          .filter((r) => r !== null),
+        entries: entries as Infer<typeof vEntry>[],
+      };
+    }
+
+    // Hybrid or text-only path: combine vector and text results with RRF.
+    let embeddingIds: VectorTableId[] = [];
+    if (hasEmbedding) {
+      const vectorResults = await searchEmbeddings(ctx, {
+        embedding: embedding!,
+        namespaceId: namespace._id,
+        filters: numberedFilters,
+        limit,
+      });
+      const threshold = args.vectorScoreThreshold ?? -1;
+      embeddingIds = vectorResults
+        .filter((r) => r._score >= threshold)
+        .map((r) => r._id);
+    }
+
+    if (!hasTextQuery) {
+      throw new Error(
+        "Search requires at least one of embedding or textQuery.",
+      );
+    }
+
+    const { ranges, entries, resultCount } = await ctx.runQuery(
+      internal.search.textAndRanges,
       {
-        embeddingIds: aboveThreshold.map((r) => r._id),
+        embeddingIds,
+        textQuery: args.textQuery!,
+        namespaceId: namespace._id,
+        filters: numberedFilters,
+        limit,
+        vectorWeight: args.vectorWeight ?? 1,
+        textWeight: args.textWeight ?? 1,
         chunkContext,
       },
     );
+
+    // Position-based scores (1.0 for first, decreasing linearly).
     return {
       results: ranges
-        .map((r, i) => publicSearchResult(r, aboveThreshold[i]._score))
+        .map((r, i) => publicSearchResult(r, (resultCount - i) / resultCount))
         .filter((r) => r !== null),
       entries: entries as Infer<typeof vEntry>[],
     };
   },
 });
 
+type TextSearchResult = {
+  chunkId: Id<"chunks">;
+  entryId: Id<"entries">;
+  order: number;
+};
+
+async function textSearchImpl(
+  ctx: QueryCtx,
+  args: {
+    query: string;
+    namespaceId: Id<"namespaces">;
+    filters: NumberedFilter[];
+    limit: number;
+  },
+): Promise<TextSearchResult[]> {
+  const toResults = (chunks: Doc<"chunks">[]): TextSearchResult[] =>
+    chunks
+      .filter((chunk) => chunk.state.kind === "ready")
+      .map((chunk) => ({
+        chunkId: chunk._id,
+        entryId: chunk.entryId,
+        order: chunk.order,
+      }));
+
+  // No user filters — just filter by namespaceId.
+  if (args.filters.length === 0) {
+    const results = await ctx.db
+      .query("chunks")
+      .withSearchIndex("searchableText", (q) =>
+        q
+          .search("state.searchableText", args.query)
+          .eq("namespaceId", args.namespaceId),
+      )
+      .take(args.limit);
+    return toResults(results);
+  }
+
+  // OR across filter conditions: run one text search per filter and dedupe.
+  const seen = new Set<Id<"chunks">>();
+  const merged: TextSearchResult[] = [];
+  for (const filter of args.filters) {
+    const fields = filterFieldsFromNumbers(args.namespaceId, filter);
+    const results = await ctx.db
+      .query("chunks")
+      .withSearchIndex("searchableText", (q) => {
+        let query = q
+          .search("state.searchableText", args.query)
+          .eq("namespaceId", args.namespaceId);
+        for (const [field, value] of Object.entries(fields)) {
+          query = query.eq(
+            field as "filter0" | "filter1" | "filter2" | "filter3",
+            value,
+          );
+        }
+        return query;
+      })
+      .take(args.limit);
+    for (const r of toResults(results)) {
+      if (!seen.has(r.chunkId)) {
+        seen.add(r.chunkId);
+        merged.push(r);
+      }
+    }
+  }
+  return merged.slice(0, args.limit);
+}
+
+export const textSearch = internalQuery({
+  args: {
+    query: v.string(),
+    namespaceId: v.id("namespaces"),
+    // Numbered filters, OR'd together (same semantics as vector search).
+    filters: v.array(v.any()),
+    limit: v.number(),
+  },
+  returns: v.array(
+    v.object({
+      chunkId: v.id("chunks"),
+      entryId: v.id("entries"),
+      order: v.number(),
+    }),
+  ),
+  handler: async (ctx, args) => {
+    return textSearchImpl(ctx, {
+      query: args.query,
+      namespaceId: args.namespaceId,
+      filters: args.filters as NumberedFilter[],
+      limit: args.limit,
+    });
+  },
+});
+
+export const textAndRanges = internalQuery({
+  args: {
+    embeddingIds: v.array(vVectorId),
+    textQuery: v.string(),
+    namespaceId: v.id("namespaces"),
+    filters: v.array(v.any()),
+    limit: v.number(),
+    vectorWeight: v.number(),
+    textWeight: v.number(),
+    chunkContext: v.object({ before: v.number(), after: v.number() }),
+  },
+  returns: v.object({
+    ranges: v.array(
+      v.union(
+        v.null(),
+        v.object({
+          entryId: v.id("entries"),
+          order: v.number(),
+          startOrder: v.number(),
+          content: v.array(
+            v.object({
+              text: v.string(),
+              metadata: v.optional(v.record(v.string(), v.any())),
+            }),
+          ),
+        }),
+      ),
+    ),
+    entries: v.array(vEntry),
+    resultCount: v.number(),
+  }),
+  handler: async (ctx, args) => {
+    // 1. Map embedding IDs to chunk IDs.
+    const vectorChunkIds: Id<"chunks">[] = (
+      await Promise.all(
+        args.embeddingIds.map(async (embeddingId) => {
+          const chunk = await ctx.db
+            .query("chunks")
+            .withIndex("embeddingId", (q) =>
+              q.eq("state.embeddingId", embeddingId),
+            )
+            .order("desc")
+            .first();
+          return chunk?._id ?? null;
+        }),
+      )
+    ).filter((id) => id !== null);
+
+    // 2. Run text search.
+    const textResults = await textSearchImpl(ctx, {
+      query: args.textQuery,
+      namespaceId: args.namespaceId,
+      filters: args.filters as NumberedFilter[],
+      limit: args.limit,
+    });
+    const textChunkIds: Id<"chunks">[] = textResults.map((r) => r.chunkId);
+
+    // 3. Merge using Reciprocal Rank Fusion.
+    const mergedChunkIds = hybridRank<Id<"chunks">>(
+      [vectorChunkIds, textChunkIds],
+      { k: 10, weights: [args.vectorWeight, args.textWeight] },
+    ).slice(0, args.limit);
+
+    if (mergedChunkIds.length === 0) {
+      return { ranges: [], entries: [], resultCount: 0 };
+    }
+
+    // 4. Build ranges from merged chunk IDs.
+    const chunks = await Promise.all(
+      mergedChunkIds.map((id) => ctx.db.get(id)),
+    );
+    const { ranges, entries } = await buildRanges(
+      ctx,
+      chunks,
+      args.chunkContext,
+    );
+    return { ranges, entries, resultCount: mergedChunkIds.length };
+  },
+});
+
 function publicSearchResult(
   r: Infer<typeof vRangeResult> | null,
   score: number,

package/src/shared.ts

@@ -35,6 +35,13 @@ export const vSearchResult = v.object({
 
 export type SearchResult = Infer<typeof vSearchResult>;
 
+export const vSearchType = v.union(
+  v.literal("vector"),
+  v.literal("text"),
+  v.literal("hybrid"),
+);
+export type SearchType = Infer<typeof vSearchType>;
+
 export const vStatus = v.union(
   v.literal("pending"),
   v.literal("ready"),