@convex-dev/rag 0.5.3 → 0.6.0

package/src/client/index.ts

@@ -2,13 +2,14 @@ import {
   embed,
   embedMany,
   generateText,
-  type ModelMessage,
   type EmbeddingModel,
   type EmbeddingModelUsage,
+  type ModelMessage,
 } from "ai";
 import { assert } from "convex-helpers";
 import {
   createFunctionHandle,
+  FunctionReference,
   internalActionGeneric,
   internalMutationGeneric,
   type FunctionArgs,
@@ -23,36 +24,30 @@ import {
   type RegisteredMutation,
 } from "convex/server";
 import { type Value } from "convex/values";
+import { ComponentApi } from "../component/_generated/component.js";
+import type { NamedFilter } from "../component/filters.js";
 import {
   CHUNK_BATCH_SIZE,
   filterNamesContain,
+  OnCompleteArgs,
   vChunkerArgs,
   vEntryId,
   vNamespaceId,
   vOnCompleteArgs,
   type Chunk,
+  type ChunkerAction,
   type CreateChunkArgs,
   type Entry,
   type EntryFilter,
   type EntryId,
   type Namespace,
   type NamespaceId,
+  type OnComplete,
+  type OnCompleteNamespace,
   type SearchEntry,
   type SearchResult,
   type Status,
 } from "../shared.js";
-import {
-  type RAGComponent,
-  type RunActionCtx,
-  type RunMutationCtx,
-  type RunQueryCtx,
-} from "./types.js";
-import {
-  type ChunkerAction,
-  type OnComplete,
-  type OnCompleteNamespace,
-} from "../shared.js";
-import type { NamedFilter } from "../component/filters.js";
 import { defaultChunker } from "./defaultChunker.js";
 
 export { hybridRank } from "./hybridRank.js";
@@ -61,7 +56,6 @@ export type {
   ChunkerAction,
   Entry,
   EntryId,
-  RAGComponent,
   NamespaceId,
   OnComplete,
   OnCompleteNamespace,
@@ -71,18 +65,18 @@ export type {
 };
 
 export {
-  type VEntry,
-  type VSearchEntry,
-  type EntryFilter,
   vEntry,
+  vOnCompleteArgs,
   vSearchEntry,
   vSearchResult,
-  vOnCompleteArgs,
+  type EntryFilter,
+  type VEntry,
+  type VSearchEntry,
 } from "../shared.js";
 export {
   contentHashFromArrayBuffer,
-  guessMimeTypeFromExtension,
   guessMimeTypeFromContents,
+  guessMimeTypeFromExtension,
 } from "./fileUtils.js";
 
 const DEFAULT_SEARCH_LIMIT = 10;
@@ -115,12 +109,12 @@ export class RAG<
    * and then entry results will have the metadata type `{ source: "website" }`.
    */
   constructor(
-    public component: RAGComponent,
+    public component: ComponentApi,
     public options: {
       embeddingDimension: number;
       textEmbeddingModel: EmbeddingModel<string>;
       filterNames?: FilterNames<FitlerSchemas>;
-    }
+    },
   ) {}
 
   /**
@@ -134,7 +128,7 @@ export class RAG<
    * The filterValues you provide can be used later to search for it.
    */
   async add(
-    ctx: RunMutationCtx,
+    ctx: CtxWith<"runMutation">,
     args: NamespaceSelection &
       EntryArgs<FitlerSchemas, EntryMetadata> &
       (
@@ -158,7 +152,7 @@ export class RAG<
             /** @deprecated You cannot specify both chunks and text currently. */
             chunks?: undefined;
           }
-      )
+      ),
   ): Promise<{
     entryId: EntryId;
     status: Status;
@@ -185,7 +179,7 @@ export class RAG<
     if (Array.isArray(chunks) && chunks.length < CHUNK_BATCH_SIZE) {
       const result = await createChunkArgsBatch(
         this.options.textEmbeddingModel,
-        chunks
+        chunks,
       );
       allChunks = result.chunks;
       totalUsage.tokens += result.usage.tokens;
@@ -208,7 +202,7 @@ export class RAG<
         },
         onComplete,
         allChunks,
-      }
+      },
     );
     if (status === "ready") {
       return {
@@ -230,7 +224,7 @@ export class RAG<
       for await (const batch of batchIterator(chunks, CHUNK_BATCH_SIZE)) {
         const result = await createChunkArgsBatch(
           this.options.textEmbeddingModel,
-          batch
+          batch,
         );
         totalUsage.tokens += result.usage.tokens;
         const { status } = await ctx.runMutation(this.component.chunks.insert, {
@@ -250,7 +244,7 @@ export class RAG<
       while (true) {
         const { status, nextStartOrder } = await ctx.runMutation(
           this.component.chunks.replaceChunksPage,
-          { entryId, startOrder }
+          { entryId, startOrder },
         );
         if (status === "ready") {
           break;
@@ -268,7 +262,7 @@ export class RAG<
     }
     const promoted = await ctx.runMutation(
       this.component.entries.promoteToReady,
-      { entryId }
+      { entryId },
     );
     return {
       entryId: entryId as EntryId,
@@ -306,7 +300,7 @@ export class RAG<
    * ```
    */
   async addAsync(
-    ctx: RunMutationCtx,
+    ctx: CtxWith<"runMutation">,
     args: NamespaceSelection &
       EntryArgs<FitlerSchemas, EntryMetadata> & {
         /**
@@ -324,7 +318,7 @@ export class RAG<
          *   });
          */
         chunkerAction: ChunkerAction;
-      }
+      },
   ): Promise<{ entryId: EntryId; status: "ready" | "pending" }> {
     let namespaceId: NamespaceId;
     if ("namespaceId" in args) {
@@ -358,7 +352,7 @@ export class RAG<
         },
         onComplete,
         chunker,
-      }
+      },
     );
     return { entryId: entryId as EntryId, status };
   }
@@ -369,7 +363,7 @@ export class RAG<
    * parameters to filter and constrain the results.
    */
   async search(
-    ctx: RunActionCtx,
+    ctx: CtxWith<"runAction">,
     args: {
       /**
        * The namespace to search in. e.g. a userId if entries are per-user.
@@ -381,7 +375,7 @@ export class RAG<
        * The query to search for. Optional if embedding is provided.
        */
       query: string | Array<number>;
-    } & SearchOptions<FitlerSchemas>
+    } & SearchOptions<FitlerSchemas>,
   ): Promise<{
     results: SearchResult[];
     text: string;
@@ -410,12 +404,12 @@ export class RAG<
       {
         embedding,
         namespace,
-        modelId: this.options.textEmbeddingModel.modelId,
+        modelId: getModelId(this.options.textEmbeddingModel),
         filters,
         limit,
         vectorScoreThreshold,
         chunkContext,
-      }
+      },
     );
     const entriesWithTexts = entries.map((e) => {
       const ranges = results
@@ -458,7 +452,7 @@ export class RAG<
    * extra context / conversation history.
    */
   async generateText(
-    ctx: RunActionCtx,
+    ctx: CtxWith<"runAction">,
     args: {
       /**
        * The search options to use for context search, including the namespace.
@@ -485,7 +479,7 @@ export class RAG<
        * to the prompt, in which case it will precede the prompt.
        */
       messages?: ModelMessage[];
-    } & Parameters<typeof generateText>[0]
+    } & Parameters<typeof generateText>[0],
   ): Promise<
     Awaited<ReturnType<typeof generateText>> & {
       context: {
@@ -530,7 +524,7 @@ export class RAG<
           .map((e) =>
             e.title
               ? `<document title="${e.title}">${e.text}</document>`
-              : `<document>${e.text}</document>`
+              : `<document>${e.text}</document>`,
           )
           .join("\n");
         contextFooter = "</context>";
@@ -576,12 +570,12 @@ export class RAG<
    * List all entries in a namespace.
    */
   async list(
-    ctx: RunQueryCtx,
+    ctx: CtxWith<"runQuery">,
     args: {
       namespaceId?: NamespaceId;
       order?: "desc" | "asc";
       status?: Status;
-    } & ({ paginationOpts: PaginationOptions } | { limit: number })
+    } & ({ paginationOpts: PaginationOptions } | { limit: number }),
   ): Promise<PaginationResult<Entry<FitlerSchemas, EntryMetadata>>> {
     const paginationOpts =
       "paginationOpts" in args
@@ -600,10 +594,10 @@ export class RAG<
    * Get entry metadata by its id.
    */
   async getEntry(
-    ctx: RunQueryCtx,
+    ctx: CtxWith<"runQuery">,
     args: {
       entryId: EntryId;
-    }
+    },
   ): Promise<Entry<FitlerSchemas, EntryMetadata> | null> {
     const entry = await ctx.runQuery(this.component.entries.get, {
       entryId: args.entryId,
@@ -617,19 +611,19 @@ export class RAG<
    * when updating content.
    */
   async findEntryByContentHash(
-    ctx: RunQueryCtx,
+    ctx: CtxWith<"runQuery">,
     args: {
       namespace: string;
       key: string;
       /** The hash of the entry contents to try to match. */
       contentHash: string;
-    }
+    },
   ): Promise<Entry<FitlerSchemas, EntryMetadata> | null> {
     const entry = await ctx.runQuery(this.component.entries.findByContentHash, {
       namespace: args.namespace,
       dimension: this.options.embeddingDimension,
       filterNames: this.options.filterNames ?? [],
-      modelId: this.options.textEmbeddingModel.modelId,
+      modelId: getModelId(this.options.textEmbeddingModel),
       key: args.key,
       contentHash: args.contentHash,
     });
@@ -641,7 +635,7 @@ export class RAG<
    * filterNames of the RAG instance. If it doesn't exist, it will be created.
    */
   async getOrCreateNamespace(
-    ctx: RunMutationCtx,
+    ctx: CtxWith<"runMutation">,
     args: {
       /**
        * The namespace to get or create. e.g. a userId if entries are per-user.
@@ -657,7 +651,7 @@ export class RAG<
        * along the way.
        */
       onComplete?: OnCompleteNamespace;
-    }
+    },
   ): Promise<{
     namespaceId: NamespaceId;
     status: "pending" | "ready";
@@ -667,7 +661,7 @@ export class RAG<
       : undefined;
     assert(
       !onComplete || args.status === "pending",
-      "You can only supply an onComplete handler for pending namespaces"
+      "You can only supply an onComplete handler for pending namespaces",
     );
     const { namespaceId, status } = await ctx.runMutation(
       this.component.namespaces.getOrCreate,
@@ -675,10 +669,10 @@ export class RAG<
         namespace: args.namespace,
         status: args.status ?? "ready",
         onComplete,
-        modelId: this.options.textEmbeddingModel.modelId,
+        modelId: getModelId(this.options.textEmbeddingModel),
         dimension: this.options.embeddingDimension,
         filterNames: this.options.filterNames ?? [],
-      }
+      },
     );
     return { namespaceId: namespaceId as NamespaceId, status };
   }
@@ -688,14 +682,14 @@ export class RAG<
    * filterNames of the RAG instance. If it doesn't exist, it will return null.
    */
   async getNamespace(
-    ctx: RunQueryCtx,
+    ctx: CtxWith<"runQuery">,
     args: {
       namespace: string;
-    }
+    },
   ): Promise<Namespace | null> {
     return ctx.runQuery(this.component.namespaces.get, {
       namespace: args.namespace,
-      modelId: this.options.textEmbeddingModel.modelId,
+      modelId: getModelId(this.options.textEmbeddingModel),
       dimension: this.options.embeddingDimension,
       filterNames: this.options.filterNames ?? [],
     }) as Promise<Namespace | null>;
@@ -705,12 +699,12 @@ export class RAG<
    * List all chunks for an entry, paginated.
    */
   async listChunks(
-    ctx: RunQueryCtx,
+    ctx: CtxWith<"runQuery">,
     args: {
       paginationOpts: PaginationOptions;
       entryId: EntryId;
       order?: "desc" | "asc";
-    }
+    },
   ): Promise<PaginationResult<Chunk>> {
     return ctx.runQuery(this.component.chunks.list, {
       entryId: args.entryId,
@@ -722,7 +716,7 @@ export class RAG<
   /**
    * Delete an entry and all its chunks in the background using a workpool.
    */
-  async deleteAsync(ctx: RunMutationCtx, args: { entryId: EntryId }) {
+  async deleteAsync(ctx: CtxWith<"runMutation">, args: { entryId: EntryId }) {
     await ctx.runMutation(this.component.entries.deleteAsync, {
       entryId: args.entryId,
       startOrder: 0,
@@ -735,17 +729,26 @@ export class RAG<
    * you're likely running this in a mutation.
    * Use `deleteAsync` or run `delete` in an action.
    */
-  async delete(ctx: RunActionCtx, args: { entryId: EntryId }): Promise<void>;
+  async delete(
+    ctx: CtxWith<"runAction">,
+    args: { entryId: EntryId },
+  ): Promise<void>;
   /** @deprecated Use `deleteAsync` in mutations. */
-  async delete(ctx: RunMutationCtx, args: { entryId: EntryId }): Promise<void>;
-  async delete(ctx: RunActionCtx | RunMutationCtx, args: { entryId: EntryId }) {
+  async delete(
+    ctx: CtxWith<"runMutation">,
+    args: { entryId: EntryId },
+  ): Promise<void>;
+  async delete(
+    ctx: CtxWith<"runMutation"> | CtxWith<"runAction">,
+    args: { entryId: EntryId },
+  ) {
     if ("runAction" in ctx) {
       await ctx.runAction(this.component.entries.deleteSync, {
         entryId: args.entryId,
       });
     } else {
       console.warn(
-        "You are running `rag.delete` in a mutation. This is deprecated. Use `rag.deleteAsync` from mutations, or `rag.delete` in actions."
+        "You are running `rag.delete` in a mutation. This is deprecated. Use `rag.deleteAsync` from mutations, or `rag.delete` in actions.",
       );
       await ctx.runMutation(this.component.entries.deleteAsync, {
         entryId: args.entryId,
@@ -758,8 +761,8 @@ export class RAG<
    * Delete all entries with a given key (asynchrounously).
    */
   async deleteByKeyAsync(
-    ctx: RunMutationCtx,
-    args: { namespaceId: NamespaceId; key: string; beforeVersion?: number }
+    ctx: CtxWith<"runMutation">,
+    args: { namespaceId: NamespaceId; key: string; beforeVersion?: number },
   ) {
     await ctx.runMutation(this.component.entries.deleteByKeyAsync, {
       namespaceId: args.namespaceId,
@@ -775,8 +778,8 @@ export class RAG<
    * Use `deleteByKeyAsync` or run `delete` in an action.
    */
   async deleteByKey(
-    ctx: RunActionCtx,
-    args: { namespaceId: NamespaceId; key: string; beforeVersion?: number }
+    ctx: CtxWith<"runAction">,
+    args: { namespaceId: NamespaceId; key: string; beforeVersion?: number },
   ) {
     await ctx.runAction(this.component.entries.deleteByKeySync, args);
   }
@@ -802,13 +805,9 @@ export class RAG<
   defineOnComplete<DataModel extends GenericDataModel>(
     fn: (
       ctx: GenericMutationCtx<DataModel>,
-      args: FunctionArgs<OnComplete<FitlerSchemas, EntryMetadata>>
-    ) => Promise<void>
-  ): RegisteredMutation<
-    "internal",
-    FunctionArgs<OnComplete<FitlerSchemas, EntryMetadata>>,
-    null
-  > {
+      args: OnCompleteArgs<FitlerSchemas, EntryMetadata>,
+    ) => Promise<void>,
+  ): RegisteredMutation<"internal", FunctionArgs<OnComplete>, null> {
     return internalMutationGeneric({
       args: vOnCompleteArgs,
       handler: fn,
@@ -836,8 +835,11 @@ export class RAG<
   defineChunkerAction<DataModel extends GenericDataModel>(
     fn: (
       ctx: GenericActionCtx<DataModel>,
-      args: { namespace: Namespace; entry: Entry<FitlerSchemas, EntryMetadata> }
-    ) => AsyncIterable<InputChunk> | Promise<{ chunks: InputChunk[] }>
+      args: {
+        namespace: Namespace;
+        entry: Entry<FitlerSchemas, EntryMetadata>;
+      },
+    ) => AsyncIterable<InputChunk> | Promise<{ chunks: InputChunk[] }>,
   ): RegisteredAction<
     "internal",
     FunctionArgs<ChunkerAction>,
@@ -847,29 +849,30 @@ export class RAG<
       args: vChunkerArgs,
       handler: async (ctx, args) => {
         const { namespace, entry } = args;
-        if (namespace.modelId !== this.options.textEmbeddingModel.modelId) {
+        const modelId = getModelId(this.options.textEmbeddingModel);
+        if (namespace.modelId !== modelId) {
           console.error(
-            `You are using a different embedding model ${this.options.textEmbeddingModel.modelId} for asynchronously ` +
-              `generating chunks than the one provided when it was started: ${namespace.modelId}`
+            `You are using a different embedding model ${modelId} for asynchronously ` +
+              `generating chunks than the one provided when it was started: ${namespace.modelId}`,
           );
           return;
         }
         if (namespace.dimension !== this.options.embeddingDimension) {
           console.error(
             `You are using a different embedding dimension ${this.options.embeddingDimension} for asynchronously ` +
-              `generating chunks than the one provided when it was started: ${namespace.dimension}`
+              `generating chunks than the one provided when it was started: ${namespace.dimension}`,
           );
           return;
         }
         if (
           !filterNamesContain(
             namespace.filterNames,
-            this.options.filterNames ?? []
+            this.options.filterNames ?? [],
           )
         ) {
           console.error(
             `You are using a different filters (${this.options.filterNames?.join(", ")}) for asynchronously ` +
-              `generating chunks than the one provided when it was started: ${namespace.filterNames.join(", ")}`
+              `generating chunks than the one provided when it was started: ${namespace.filterNames.join(", ")}`,
           );
           return;
         }
@@ -891,23 +894,23 @@ export class RAG<
         let batchOrder = 0;
         for await (const batch of batchIterator(
           chunkIterator,
-          CHUNK_BATCH_SIZE
+          CHUNK_BATCH_SIZE,
         )) {
           const result = await createChunkArgsBatch(
             this.options.textEmbeddingModel,
-            batch
+            batch,
           );
           await ctx.runMutation(
             args.insertChunks as FunctionHandle<
               "mutation",
-              FunctionArgs<RAGComponent["chunks"]["insert"]>,
+              FunctionArgs<ComponentApi["chunks"]["insert"]>,
               null
             >,
             {
               entryId: entry.entryId,
               startOrder: batchOrder,
               chunks: result.chunks,
-            }
+            },
           );
           batchOrder += result.chunks.length;
         }
@@ -918,7 +921,7 @@ export class RAG<
 
 async function* batchIterator<T>(
   iterator: Iterable<T> | AsyncIterable<T>,
-  batchSize: number
+  batchSize: number,
 ): AsyncIterable<T[]> {
   let batch: T[] = [];
   for await (const item of iterator) {
@@ -935,21 +938,21 @@ async function* batchIterator<T>(
 
 function validateAddFilterValues(
   filterValues: NamedFilter[] | undefined,
-  filterNames: string[] | undefined
+  filterNames: string[] | undefined,
 ) {
   if (!filterValues) {
     return;
   }
   if (!filterNames) {
     throw new Error(
-      "You must provide filter names to RAG to add entries with filters."
+      "You must provide filter names to RAG to add entries with filters.",
     );
   }
   const seen = new Set<string>();
   for (const filterValue of filterValues) {
     if (seen.has(filterValue.name)) {
       throw new Error(
-        `You cannot provide the same filter name twice: ${filterValue.name}.`
+        `You cannot provide the same filter name twice: ${filterValue.name}.`,
       );
     }
     seen.add(filterValue.name);
@@ -957,7 +960,7 @@ function validateAddFilterValues(
   for (const filterName of filterNames) {
     if (!seen.has(filterName)) {
       throw new Error(
-        `Filter name ${filterName} is not valid (one of ${filterNames.join(", ")}).`
+        `Filter name ${filterName} is not valid (one of ${filterNames.join(", ")}).`,
       );
     }
   }
@@ -973,7 +976,7 @@ function makeBatches<T>(items: T[], batchSize: number): T[][] {
 
 async function createChunkArgsBatch(
   embedModel: EmbeddingModel<string>,
-  chunks: InputChunk[]
+  chunks: InputChunk[],
 ): Promise<{ chunks: CreateChunkArgs[]; usage: EmbeddingModelUsage }> {
   const argsMaybeMissingEmbeddings: (Omit<CreateChunkArgs, "embedding"> & {
     embedding?: number[];
@@ -1005,7 +1008,7 @@ async function createChunkArgsBatch(
         : {
             text: arg.content.text,
             index,
-          }
+          },
     )
     .filter((b) => b !== null);
   const totalUsage: EmbeddingModelUsage = { tokens: 0 };
@@ -1191,3 +1194,45 @@ function getModelCategory(model: string | { provider: string }) {
   }
   return model;
 }
+
+// fetch metadata from either a string or EmbeddingModelV2 or LanguageModelV2
+export type ModelOrMetadata =
+  | string
+  | ({ provider: string } & ({ modelId: string } | { model: string }));
+
+export function getModelId(embeddingModel: ModelOrMetadata): string {
+  if (typeof embeddingModel === "string") {
+    if (embeddingModel.includes("/")) {
+      return embeddingModel.split("/").slice(1).join("/");
+    }
+    return embeddingModel;
+  }
+  return "modelId" in embeddingModel
+    ? embeddingModel.modelId
+    : embeddingModel.model;
+}
+
+export function getProviderName(embeddingModel: ModelOrMetadata): string {
+  if (typeof embeddingModel === "string") {
+    return embeddingModel.split("/").at(0)!;
+  }
+  return embeddingModel.provider;
+}
+
+type CtxWith<T extends "runQuery" | "runMutation" | "runAction"> = Pick<
+  {
+    runQuery: <Query extends FunctionReference<"query", "internal">>(
+      query: Query,
+      args: FunctionArgs<Query>,
+    ) => Promise<FunctionReturnType<Query>>;
+    runMutation: <Mutation extends FunctionReference<"mutation", "internal">>(
+      mutation: Mutation,
+      args: FunctionArgs<Mutation>,
+    ) => Promise<FunctionReturnType<Mutation>>;
+    runAction: <Action extends FunctionReference<"action", "internal">>(
+      action: Action,
+      args: FunctionArgs<Action>,
+    ) => Promise<FunctionReturnType<Action>>;
+  },
+  T
+>;

package/src/client/setup.test.ts

@@ -8,7 +8,7 @@ import {
   type GenericSchema,
   type SchemaDefinition,
 } from "convex/server";
-import { type RAGComponent } from "./index.js";
+import { type ComponentApi } from "../component/_generated/component.js";
 import { componentsGeneric } from "convex/server";
 import componentSchema from "../component/schema.js";
 export { componentSchema };
@@ -22,7 +22,7 @@ export function initConvexTest<
   return t;
 }
 export const components = componentsGeneric() as unknown as {
-  rag: RAGComponent;
+  rag: ComponentApi;
 };
 
 test("setup", () => {});