@convex-dev/rag 0.6.1 → 0.7.1

package/src/client/index.ts

@@ -34,6 +34,7 @@ import {
   vEntryId,
   vNamespaceId,
   vOnCompleteArgs,
+  vSearchType,
   type Chunk,
   type ChunkerAction,
   type CreateChunkArgs,
@@ -46,12 +47,13 @@ import {
   type OnCompleteNamespace,
   type SearchEntry,
   type SearchResult,
+  type SearchType,
   type Status,
 } from "../shared.js";
 import { defaultChunker } from "./defaultChunker.js";
 
 export { hybridRank } from "./hybridRank.js";
-export { defaultChunker, vEntryId, vNamespaceId };
+export { defaultChunker, vEntryId, vNamespaceId, vSearchType };
 export type {
   ChunkerAction,
   Entry,
@@ -61,6 +63,7 @@ export type {
   OnCompleteNamespace,
   SearchEntry,
   SearchResult,
+  SearchType,
   Status,
 };
 
@@ -112,7 +115,7 @@ export class RAG<
     public component: ComponentApi,
     public options: {
       embeddingDimension: number;
-      textEmbeddingModel: EmbeddingModel<string>;
+      textEmbeddingModel: EmbeddingModel;
       filterNames?: FilterNames<FitlerSchemas>;
     },
   ) {}
@@ -388,27 +391,56 @@ export class RAG<
       limit = DEFAULT_SEARCH_LIMIT,
       chunkContext = { before: 0, after: 0 },
       vectorScoreThreshold,
+      searchType = "vector",
+      textWeight,
+      vectorWeight,
     } = args;
-    let embedding = Array.isArray(args.query) ? args.query : undefined;
+
+    const needsEmbedding = searchType !== "text";
+    let needsTextQuery = searchType !== "vector";
+
+    if (needsTextQuery && Array.isArray(args.query)) {
+      if (searchType === "text") {
+        throw new Error('searchType "text" requires a string query.');
+      }
+      console.warn(
+        `searchType "${searchType}" requires a string query. Falling back to vector-only search for embedding array queries.`,
+      );
+      needsTextQuery = false;
+    }
+
+    let embedding: number[] | undefined;
     let usage: EmbeddingModelUsage = { tokens: 0 };
-    if (!embedding) {
-      const embedResult = await embed({
-        model: this.options.textEmbeddingModel,
-        value: args.query,
-      });
-      embedding = embedResult.embedding;
-      usage = embedResult.usage;
+    if (needsEmbedding) {
+      if (Array.isArray(args.query)) {
+        embedding = args.query;
+      } else {
+        const embedResult = await embed({
+          model: this.options.textEmbeddingModel,
+          value: args.query,
+        });
+        embedding = embedResult.embedding;
+        usage = embedResult.usage;
+      }
     }
+
+    const textQuery =
+      needsTextQuery && typeof args.query === "string" ? args.query : undefined;
+
     const { results, entries } = await ctx.runAction(
       this.component.search.search,
       {
         embedding,
+        dimension: this.options.embeddingDimension,
         namespace,
         modelId: getModelId(this.options.textEmbeddingModel),
         filters,
         limit,
         vectorScoreThreshold,
         chunkContext,
+        textQuery,
+        textWeight,
+        vectorWeight,
       },
     );
     const entriesWithTexts = entries.map((e) => {
@@ -975,27 +1007,27 @@ function makeBatches<T>(items: T[], batchSize: number): T[][] {
 }
 
 async function createChunkArgsBatch(
-  embedModel: EmbeddingModel<string>,
+  embedModel: EmbeddingModel,
   chunks: InputChunk[],
 ): Promise<{ chunks: CreateChunkArgs[]; usage: EmbeddingModelUsage }> {
   const argsMaybeMissingEmbeddings: (Omit<CreateChunkArgs, "embedding"> & {
     embedding?: number[];
   })[] = chunks.map((chunk) => {
     if (typeof chunk === "string") {
-      return { content: { text: chunk } };
+      return { content: { text: chunk }, searchableText: chunk };
     } else if ("text" in chunk) {
       const { text, metadata, keywords: searchableText } = chunk;
       return {
         content: { text, metadata },
         embedding: chunk.embedding,
-        searchableText,
+        searchableText: searchableText ?? text,
       };
     } else if ("pageContent" in chunk) {
       const { pageContent: text, metadata, keywords: searchableText } = chunk;
       return {
         content: { text, metadata },
         embedding: chunk.embedding,
-        searchableText,
+        searchableText: searchableText ?? text,
       };
     } else {
       throw new Error("Invalid chunk: " + JSON.stringify(chunk));
@@ -1033,22 +1065,24 @@ async function createChunkArgsBatch(
 
 type MastraChunk = {
   text: string;
-  metadata: Record<string, Value>;
+  metadata?: Record<string, Value>;
   embedding?: Array<number>;
 };
 
 type LangChainChunk = {
   id?: string;
   pageContent: string;
-  metadata: Record<string, Value>; //{ loc: { lines: { from: number; to: number } } };
+  metadata?: Record<string, Value>; //{ loc: { lines: { from: number; to: number } } };
   embedding?: Array<number>;
 };
 
 export type InputChunk =
   | string
   | ((MastraChunk | LangChainChunk) & {
-      // Space-delimited keywords to text search on.
-      // TODO: implement text search
+      /**
+       * Text to use for full-text search. Defaults to the chunk's text content.
+       * Provide a custom value to control what text is searchable.
+       */
       keywords?: string;
       // In the future we can add per-chunk metadata if it's useful.
       // importance?: Importance;
@@ -1167,6 +1201,34 @@ type SearchOptions<FitlerSchemas extends Record<string, Value>> = {
    * The minimum score to return a result.
    */
   vectorScoreThreshold?: number;
+  /**
+   * The search mode to use.
+   * - "vector": Vector similarity search only (default). Returns cosine
+   *   similarity scores.
+   * - "text": Full-text search only. No embedding is computed. Returns
+   *   position-based scores.
+   * - "hybrid": Combines vector and full-text search using Reciprocal Rank
+   *   Fusion. Returns position-based scores (1.0 for top result, decreasing
+   *   linearly).
+   *
+   * Text and hybrid modes require the query to be a string (not an embedding
+   * array).
+   */
+  searchType?: SearchType;
+  /**
+   * Weight for text search results in hybrid ranking (RRF).
+   * Higher values give more influence to text search matches.
+   * Only used when searchType is "hybrid".
+   * Default: 1
+   */
+  textWeight?: number;
+  /**
+   * Weight for vector search results in hybrid ranking (RRF).
+   * Higher values give more influence to vector search matches.
+   * Only used when searchType is "hybrid".
+   * Default: 1
+   */
+  vectorWeight?: number;
 };
 
 function getModelCategory(model: string | { provider: string }) {

package/src/component/_generated/component.ts

@@ -409,12 +409,17 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         "internal",
         {
           chunkContext?: { after: number; before: number };
-          embedding: Array<number>;
+          dimension?: number;
+          embedding?: Array<number>;
           filters: Array<{ name: string; value: any }>;
           limit: number;
           modelId: string;
           namespace: string;
+          searchType?: "vector" | "text" | "hybrid";
+          textQuery?: string;
+          textWeight?: number;
           vectorScoreThreshold?: number;
+          vectorWeight?: number;
         },
         {
           entries: Array<{

package/src/component/_generated/dataModel.ts

@@ -38,7 +38,7 @@ export type Doc<TableName extends TableNames> = DocumentByName<
  * Convex documents are uniquely identified by their `Id`, which is accessible
  * on the `_id` field. To learn more, see [Document IDs](https://docs.convex.dev/using/document-ids).
  *
- * Documents can be loaded using `db.get(id)` in query and mutation functions.
+ * Documents can be loaded using `db.get(tableName, id)` in query and mutation functions.
  *
  * IDs are just strings at runtime, but this type can be used to distinguish them from other
  * strings when type checking.

package/src/component/_generated/server.ts

@@ -107,11 +107,6 @@ export const internalAction: ActionBuilder<DataModel, "internal"> =
  */
 export const httpAction: HttpActionBuilder = httpActionGeneric;
 
-type GenericCtx =
-  | GenericActionCtx<DataModel>
-  | GenericMutationCtx<DataModel>
-  | GenericQueryCtx<DataModel>;
-
 /**
  * A set of services for use within Convex query functions.
  *

package/src/component/chunks.ts

@@ -311,6 +311,107 @@ export const vRangeResult = v.object({
   ),
 });
 
+export async function buildRanges(
+  ctx: QueryCtx,
+  chunks: (Doc<"chunks"> | null)[],
+  chunkContext: { before: number; after: number },
+): Promise<{
+  ranges: (null | Infer<typeof vRangeResult>)[];
+  entries: Entry[];
+}> {
+  // Note: This preserves order of entries as they first appeared.
+  const entryDocs = (
+    await Promise.all(
+      Array.from(
+        new Set(chunks.filter((c) => c !== null).map((c) => c.entryId)),
+      ).map((id) => ctx.db.get(id)),
+    )
+  ).filter((d): d is Doc<"entries"> => d !== null);
+  const entries = entryDocs.map(publicEntry);
+  const entryDocById = new Map(entryDocs.map((d) => [d._id, d]));
+
+  const entryOrders = chunks
+    .filter((c) => c !== null)
+    .map((c) => [c.entryId, c.order] as const)
+    .reduce(
+      (acc, [entryId, order]) => {
+        if (acc[entryId]?.includes(order)) {
+          // De-dupe orders
+          return acc;
+        }
+        acc[entryId] = [...(acc[entryId] ?? []), order].sort((a, b) => a - b);
+        return acc;
+      },
+      {} as Record<Id<"entries">, number[]>,
+    );
+
+  const result: Array<Infer<typeof vRangeResult> | null> = [];
+
+  for (const chunk of chunks) {
+    if (chunk === null) {
+      result.push(null);
+      continue;
+    }
+    // Note: if we parallelize this in the future, we could have a race
+    // instead we'd check that other chunks are not the same doc/order
+    if (
+      result.find(
+        (r) => r?.entryId === chunk.entryId && r?.order === chunk.order,
+      )
+    ) {
+      // De-dupe chunks
+      result.push(null);
+      continue;
+    }
+    const entryId = chunk.entryId;
+    const entry = entryDocById.get(entryId);
+    assert(entry, `Entry ${entryId} not found`);
+    const otherOrders = entryOrders[entryId] ?? [chunk.order];
+    const ourOrderIndex = otherOrders.indexOf(chunk.order);
+    const previousOrder = otherOrders[ourOrderIndex - 1] ?? -Infinity;
+    const nextOrder = otherOrders[ourOrderIndex + 1] ?? Infinity;
+    // We absorb all previous context up to the previous chunk.
+    const startOrder = Math.max(
+      chunk.order - chunkContext.before,
+      0,
+      Math.min(previousOrder + 1, chunk.order),
+    );
+    // We stop short if the next chunk order's "before" context will cover it.
+    const endOrder = Math.min(
+      chunk.order + chunkContext.after + 1,
+      Math.max(nextOrder - chunkContext.before, chunk.order + 1),
+    );
+    const contentIds: Id<"content">[] = [];
+    if (startOrder === chunk.order && endOrder === chunk.order + 1) {
+      contentIds.push(chunk.contentId);
+    } else {
+      const rangeChunks = await ctx.db
+        .query("chunks")
+        .withIndex("entryId_order", (q) =>
+          q
+            .eq("entryId", entryId)
+            .gte("order", startOrder)
+            .lt("order", endOrder),
+        )
+        .collect();
+      for (const c of rangeChunks) {
+        contentIds.push(c.contentId);
+      }
+    }
+    const content = await Promise.all(
+      contentIds.map(async (contentId) => {
+        const content = await ctx.db.get(contentId);
+        assert(content, `Content ${contentId} not found`);
+        return { text: content.text, metadata: content.metadata };
+      }),
+    );
+
+    result.push({ entryId, order: chunk.order, startOrder, content });
+  }
+
+  return { ranges: result, entries };
+}
+
 export const getRangesOfChunks = internalQuery({
   args: {
     embeddingIds: v.array(vVectorId),
@@ -339,98 +440,7 @@ export const getRangesOfChunks = internalQuery({
           .first(),
       ),
     );
-
-    // Note: This preserves order of entries as they first appeared.
-    const entries = (
-      await Promise.all(
-        Array.from(
-          new Set(chunks.filter((c) => c !== null).map((c) => c.entryId)),
-        ).map((id) => ctx.db.get(id)),
-      )
-    )
-      .filter((d) => d !== null)
-      .map(publicEntry);
-
-    const entryOders = chunks
-      .filter((c) => c !== null)
-      .map((c) => [c.entryId, c.order] as const)
-      .reduce(
-        (acc, [entryId, order]) => {
-          if (acc[entryId]?.includes(order)) {
-            // De-dupe orders
-            return acc;
-          }
-          acc[entryId] = [...(acc[entryId] ?? []), order].sort((a, b) => a - b);
-          return acc;
-        },
-        {} as Record<Id<"entries">, number[]>,
-      );
-
-    const result: Array<Infer<typeof vRangeResult> | null> = [];
-
-    for (const chunk of chunks) {
-      if (chunk === null) {
-        result.push(null);
-        continue;
-      }
-      // Note: if we parallelize this in the future, we could have a race
-      // instead we'd check that other chunks are not the same doc/order
-      if (
-        result.find(
-          (r) => r?.entryId === chunk.entryId && r?.order === chunk.order,
-        )
-      ) {
-        // De-dupe chunks
-        result.push(null);
-        continue;
-      }
-      const entryId = chunk.entryId;
-      const entry = await ctx.db.get(entryId);
-      assert(entry, `Entry ${entryId} not found`);
-      const otherOrders = entryOders[entryId] ?? [chunk.order];
-      const ourOrderIndex = otherOrders.indexOf(chunk.order);
-      const previousOrder = otherOrders[ourOrderIndex - 1] ?? -Infinity;
-      const nextOrder = otherOrders[ourOrderIndex + 1] ?? Infinity;
-      // We absorb all previous context up to the previous chunk.
-      const startOrder = Math.max(
-        chunk.order - chunkContext.before,
-        0,
-        Math.min(previousOrder + 1, chunk.order),
-      );
-      // We stop short if the next chunk order's "before" context will cover it.
-      const endOrder = Math.min(
-        chunk.order + chunkContext.after + 1,
-        Math.max(nextOrder - chunkContext.before, chunk.order + 1),
-      );
-      const contentIds: Id<"content">[] = [];
-      if (startOrder === chunk.order && endOrder === chunk.order + 1) {
-        contentIds.push(chunk.contentId);
-      } else {
-        const chunks = await ctx.db
-          .query("chunks")
-          .withIndex("entryId_order", (q) =>
-            q
-              .eq("entryId", entryId)
-              .gte("order", startOrder)
-              .lt("order", endOrder),
-          )
-          .collect();
-        for (const chunk of chunks) {
-          contentIds.push(chunk.contentId);
-        }
-      }
-      const content = await Promise.all(
-        contentIds.map(async (contentId) => {
-          const content = await ctx.db.get(contentId);
-          assert(content, `Content ${contentId} not found`);
-          return { text: content.text, metadata: content.metadata };
-        }),
-      );
-
-      result.push({ entryId, order: chunk.order, startOrder, content });
-    }
-
-    return { ranges: result, entries };
+    return buildRanges(ctx, chunks, chunkContext);
   },
 });
 

package/src/component/schema.ts

@@ -71,7 +71,6 @@ export const schema = defineSchema({
       v.object({
         kind: v.literal("ready"),
         embeddingId: vVectorId,
-        // TODO: text search
         searchableText: v.optional(v.string()),
       }),
       v.object({