@convex-dev/rag 0.3.1 → 0.3.3-alpha.0

package/src/client/index.ts

@@ -27,7 +27,7 @@ import {
   type Chunk,
   type CreateChunkArgs,
   type Entry,
-  type EntryFilterValues,
+  type EntryFilter,
   type EntryId,
   type Namespace,
   type NamespaceId,
@@ -49,6 +49,7 @@ import {
 import type { NamedFilter } from "../component/filters.js";
 import { defaultChunker } from "./defaultChunker.js";
 
+export { hybridRank } from "./hybridRank.js";
 export { defaultChunker, vEntryId, vNamespaceId };
 export type {
   ChunkerAction,
@@ -66,6 +67,7 @@ export type {
 export {
   type VEntry,
   type VSearchEntry,
+  type EntryFilter,
   vEntry,
   vSearchEntry,
   vSearchResult,
@@ -362,7 +364,7 @@ export class RAG<
       /**
        * The query to search for. Optional if embedding is provided.
        */
-      query?: string;
+      query: string | Array<number>;
     } & SearchOptions<FitlerSchemas>
   ): Promise<{
     results: SearchResult[];
@@ -376,7 +378,7 @@ export class RAG<
       chunkContext = { before: 0, after: 0 },
       vectorScoreThreshold,
     } = args;
-    let embedding = args.embedding;
+    let embedding = Array.isArray(args.query) ? args.query : undefined;
     if (!embedding) {
       const embedResult = await embed({
         model: this.options.textEmbeddingModel,
@@ -405,7 +407,7 @@ export class RAG<
       for (const range of ranges) {
         if (previousEnd !== 0) {
           if (range.startOrder !== previousEnd) {
-            text += "\n...\n";
+            text += "\n\n...\n\n";
           } else {
             text += "\n";
           }
@@ -419,8 +421,8 @@ export class RAG<
     return {
       results: results as SearchResult[],
       text: entriesWithTexts
-        .map((e) => (e.title ? `## ${e.title}:\n${e.text}` : e.text))
-        .join(`\n---\n`),
+        .map((e) => (e.title ? `## ${e.title}:\n\n${e.text}` : e.text))
+        .join(`\n\n---\n\n`),
       entries: entriesWithTexts,
     };
   }
@@ -446,6 +448,11 @@ export class RAG<
          * The namespace to search in. e.g. a userId if entries are per-user.
          */
         namespace: string;
+        /**
+         * The text or embedding to search for. If provided, it will be used
+         * instead of the prompt for vector search.
+         */
+        query?: string | Array<number>;
       };
       /**
        * Required. The prompt to use for context search, as well as the final
@@ -552,14 +559,17 @@ export class RAG<
     ctx: RunQueryCtx,
     args: {
       namespaceId?: NamespaceId;
-      paginationOpts: PaginationOptions;
       order?: "desc" | "asc";
       status?: Status;
-    }
+    } & ({ paginationOpts: PaginationOptions } | { limit: number })
   ): Promise<PaginationResult<Entry<FitlerSchemas, EntryMetadata>>> {
+    const paginationOpts =
+      "paginationOpts" in args
+        ? args.paginationOpts
+        : { cursor: null, numItems: args.limit };
     const results = await ctx.runQuery(this.component.entries.list, {
       namespaceId: args.namespaceId,
-      paginationOpts: args.paginationOpts,
+      paginationOpts,
       order: args.order ?? "asc",
       status: args.status ?? "ready",
     });
@@ -586,7 +596,7 @@ export class RAG<
    * new results into a new entry when migrating, or avoiding duplicating work
    * when updating content.
    */
-  async findExistingEntryByContentHash(
+  async findEntryByContentHash(
     ctx: RunQueryCtx,
     args: {
       namespace: string;
@@ -679,11 +689,13 @@ export class RAG<
     args: {
       paginationOpts: PaginationOptions;
       entryId: EntryId;
+      order?: "desc" | "asc";
     }
   ): Promise<PaginationResult<Chunk>> {
     return ctx.runQuery(this.component.chunks.list, {
       entryId: args.entryId,
       paginationOpts: args.paginationOpts,
+      order: args.order ?? "asc",
     });
   }
 
@@ -927,7 +939,7 @@ async function createChunkArgsBatch(
   for (const batch of makeBatches(missingEmbeddingsWithIndex, 100)) {
     const { embeddings } = await embedMany({
       model: embedModel,
-      values: batch.map((b) => b.text),
+      values: batch.map((b) => b.text.trim() || "<empty>"),
     });
     for (const [index, embedding] of embeddings.entries()) {
       argsMaybeMissingEmbeddings[batch[index].index].embedding = embedding;
@@ -941,46 +953,6 @@ async function createChunkArgsBatch(
   }) as CreateChunkArgs[];
 }
 
-/**
- * Rank results from multiple results, e.g. from vector search and text search.
- * Uses the "Recriprocal Rank Fusion" algorithm.
- * @param sortedResults The results arrays ordered by most important first.
- */
-export function hybridRank<T extends string>(
-  sortedResults: T[][],
-  opts?: {
-    /**
-     * A constant used to change the bias of the top results in each list vs.
-     * results in the middle of multiple lists.
-     * A higher k means less of a bias toward the top few results.
-     */
-    k: number;
-    /**
-     * The weights of each sortedResults array.
-     * Used to prefer results from one sortedResults array over another.
-     */
-    weights: number[];
-    /**
-     * The cutoff score for a result to be returned.
-     */
-    cutoffScore?: number;
-  }
-): T[] {
-  const k = opts?.k ?? 10;
-  const scores: Map<T, number> = new Map();
-  for (const [i, results] of sortedResults.entries()) {
-    const weight = opts?.weights?.[i] ?? 1;
-    for (let j = 0; j < results.length; j++) {
-      const key = results[j];
-      scores.set(key, (scores.get(key) ?? 0) + weight / (k + j));
-    }
-  }
-  const sortedScores = Array.from(scores.entries()).sort((a, b) => b[1] - a[1]);
-  return sortedScores
-    .filter(([_, score]) => score >= (opts?.cutoffScore ?? 0))
-    .map(([key]) => key);
-}
-
 type MastraChunk = {
   text: string;
   metadata: Record<string, Value>;
@@ -1060,7 +1032,7 @@ type EntryArgs<
    * and searching with the same value will return entries that match that
    * value exactly.
    */
-  filterValues?: EntryFilterValues<FitlerSchemas>[];
+  filterValues?: EntryFilter<FitlerSchemas>[];
   /**
    * The importance of the entry. This is used to scale the vector search
    * score of each chunk.
@@ -1080,11 +1052,6 @@ type EntryArgs<
 };
 
 type SearchOptions<FitlerSchemas extends Record<string, Value>> = {
-  /**
-   * The embedding to search for. If provided, it will be used instead
-   * of the query for vector search.
-   */
-  embedding?: Array<number>;
   /**
    * Filters to apply to the search. These are OR'd together. To represent
    * AND logic, your filter can be an object or array with multiple values.
@@ -1097,7 +1064,7 @@ type SearchOptions<FitlerSchemas extends Record<string, Value>> = {
    * `{ team_user: { team: "team1", user: "user1" } }`, it will not match
    * `{ team_user: { team: "team1" } }` but it will match
    */
-  filters?: EntryFilterValues<FitlerSchemas>[];
+  filters?: EntryFilter<FitlerSchemas>[];
   /**
    * The maximum number of messages to fetch. Default is 10.
    * This is the number *before* the chunkContext is applied.

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

@@ -62,6 +62,7 @@ export type Mounts = {
       "public",
       {
         entryId: string;
+        order: "desc" | "asc";
         paginationOpts: {
           cursor: string | null;
           endCursor?: string | null;

package/src/component/chunks.test.ts

@@ -332,11 +332,12 @@ describe("chunks", () => {
     // Insert a large number of chunks
     const chunks = createTestChunks(10);
     await t.run(async (ctx) => {
-      return insertChunks(ctx, {
+      const result = await insertChunks(ctx, {
         entryId,
         startOrder: 0,
         chunks,
       });
+      expect(result.status).toBe("ready");
     });
 
     // Verify chunks exist
@@ -375,8 +376,15 @@ describe("chunks", () => {
     const allContent = await t.run(async (ctx) => {
       return ctx.db.query("content").collect();
     });
+
     // Should have only 3 content records remaining (for the 3 remaining chunks)
     expect(allContent).toHaveLength(3);
+
+    // Verify embeddings were deleted
+    const allEmbeddings = await t.run(async (ctx) => {
+      return ctx.db.query("vectors_128").collect();
+    });
+    expect(allEmbeddings).toHaveLength(3);
   });
 
   test("listing chunks returns correct pagination", async () => {
@@ -397,6 +405,7 @@ describe("chunks", () => {
     // Test listing with pagination
     const result = await t.query(api.chunks.list, {
       entryId,
+      order: "asc",
       paginationOpts: { numItems: 3, cursor: null },
     });
 
@@ -417,6 +426,7 @@ describe("chunks", () => {
     // Get next page
     const nextResult = await t.query(api.chunks.list, {
       entryId,
+      order: "asc",
       paginationOpts: { numItems: 3, cursor: result.continueCursor },
     });
 

package/src/component/chunks.ts

@@ -21,7 +21,7 @@ import {
   type QueryCtx,
 } from "./_generated/server.js";
 import { insertEmbedding } from "./embeddings/index.js";
-import { vVectorId } from "./embeddings/tables.js";
+import { vVectorId, type VectorTableName } from "./embeddings/tables.js";
 import { schema, v } from "./schema.js";
 import { getPreviousEntry, publicEntry } from "./entries.js";
 import {
@@ -463,6 +463,7 @@ export const list = query({
   args: v.object({
     entryId: v.id("entries"),
     paginationOpts: paginationOptsValidator,
+    order: v.union(v.literal("desc"), v.literal("asc")),
   }),
   returns: vPaginationResult(vChunk),
   handler: async (ctx, args) => {
@@ -470,7 +471,7 @@ export const list = query({
     const chunks = await paginator(ctx.db, schema)
       .query("chunks")
       .withIndex("entryId_order", (q) => q.eq("entryId", entryId))
-      .order("asc")
+      .order(args.order)
       .paginate(paginationOpts);
     return {
       ...chunks,
@@ -524,20 +525,49 @@ export async function deleteChunksPage(
   for await (const chunk of chunkStream) {
     dataUsedSoFar += await estimateChunkSize(chunk);
     await ctx.db.delete(chunk._id);
+    if (chunk.state.kind === "ready") {
+      const embedding = await ctx.db.get(chunk.state.embeddingId);
+      if (embedding) {
+        dataUsedSoFar += estimateEmbeddingSize(embedding);
+        await ctx.db.delete(chunk.state.embeddingId);
+      }
+    }
     dataUsedSoFar += await estimateContentSize(ctx, chunk.contentId);
     await ctx.db.delete(chunk.contentId);
     if (dataUsedSoFar > BANDWIDTH_PER_TRANSACTION_HARD_LIMIT) {
-      // TODO: schedule follow-up - workpool?
       return { isDone: false, nextStartOrder: chunk.order };
     }
   }
   return { isDone: true, nextStartOrder: -1 };
 }
 
+function estimateEmbeddingSize(embedding: Doc<VectorTableName>) {
+  let dataUsedSoFar =
+    embedding.vector.length * 8 +
+    embedding.namespaceId.length +
+    embedding._id.length +
+    8;
+  for (const filter of [
+    embedding.filter0,
+    embedding.filter1,
+    embedding.filter2,
+    embedding.filter3,
+  ]) {
+    if (filter) {
+      dataUsedSoFar += JSON.stringify(convexToJson(filter[1])).length;
+    }
+  }
+  return dataUsedSoFar;
+}
+
 async function estimateChunkSize(chunk: Doc<"chunks">) {
   let dataUsedSoFar = 100; // constant metadata - roughly
   if (chunk.state.kind === "pending") {
     dataUsedSoFar += chunk.state.embedding.length * 8;
+    dataUsedSoFar += chunk.state.pendingSearchableText?.length ?? 0;
+  } else if (chunk.state.kind === "replaced") {
+    dataUsedSoFar += chunk.state.vector.length * 8;
+    dataUsedSoFar += chunk.state.pendingSearchableText?.length ?? 0;
   }
   return dataUsedSoFar;
 }

package/src/component/entries.ts

@@ -1,7 +1,7 @@
 import { assert, omit } from "convex-helpers";
 import { createFunctionHandle, paginationOptsValidator } from "convex/server";
 import { v, type Value } from "convex/values";
-import type { ChunkerAction, EntryFilterValues, EntryId } from "../shared.js";
+import type { ChunkerAction, EntryFilter, EntryId } from "../shared.js";
 import {
   statuses,
   vActiveStatus,
@@ -115,7 +115,7 @@ function workpoolName(
   key: string | undefined,
   entryId: Id<"entries">
 ) {
-  return `async-chunker-${namespace}-${key ? key + "-" + entryId : entryId}`;
+  return `rag-async-${namespace}-${key ? key + "-" + entryId : entryId}`;
 }
 
 export const addAsyncOnComplete = internalMutation({
@@ -497,7 +497,7 @@ export function publicEntry(entry: {
   _id: Id<"entries">;
   key?: string | undefined;
   importance: number;
-  filterValues: EntryFilterValues[];
+  filterValues: EntryFilter[];
   contentHash?: string | undefined;
   title?: string | undefined;
   metadata?: Record<string, Value> | undefined;

package/src/shared.ts

@@ -99,7 +99,7 @@ export type SearchEntry<
   text: string;
 };
 
-export type EntryFilterValues<
+export type EntryFilter<
   Filters extends Record<string, Value> = Record<string, Value>,
 > = {
   [K in keyof Filters & string]: NamedFilter<K, Filters[K]>;
@@ -126,7 +126,7 @@ export type Entry<
   /** Filters that can be used to search for this entry.
    * Up to 4 filters are supported, of any type.
    */
-  filterValues: EntryFilterValues<Filters>[];
+  filterValues: EntryFilter<Filters>[];
   /** Hash of the entry contents.
    * If supplied, it will avoid adding if the hash is the same.
    */