chore: draft integration of embeddings with the aggregate tool

Author: kmruiz

src/common/errors.ts

@@ -4,6 +4,9 @@ export enum ErrorCodes {
     ForbiddenCollscan = 1_000_002,
     ForbiddenWriteOperation = 1_000_003,
     AtlasSearchNotSupported = 1_000_004,
+    NoEmbeddingsProviderConfigured = 1_000_005,
+    AtlasVectorSearchIndexNotFound = 1_000_006,
+    AtlasVectorSearchInvalidQuery = 1_000_007,
 }
 
 export class MongoDBError<ErrorCode extends ErrorCodes = ErrorCodes> extends Error {

src/common/search/embeddingsProvider.ts

@@ -10,40 +10,17 @@ const zEmbeddingsInput = z.string();
 type EmbeddingsInput = z.infer<typeof zEmbeddingsInput>;
 type Embeddings = number[];
 
+type EmbeddingParameters = {
+    numDimensions: number;
+    quantization: string;
+    inputType: "query" | "document";
+};
+
 interface EmbeddingsProvider<SupportedModels extends string> {
-    embed(modelId: SupportedModels, content: EmbeddingsInput[], parameters: unknown): Promise<Embeddings[]>;
+    embed(modelId: SupportedModels, content: EmbeddingsInput[], parameters: EmbeddingParameters): Promise<Embeddings[]>;
 }
 
-const zVoyageSupportedDimensions = z
-    .union([z.literal(256), z.literal(512), z.literal(1024), z.literal(2048)])
-    .default(1024);
-
-const zVoyageQuantization = z.enum(["float", "int8", "binary", "ubinary"]).default("float");
-const zVoyageInputType = z.enum(["query", "document"]);
-
 export const zVoyageModels = z.enum(["voyage-3-large", "voyage-3.5", "voyage-3.5-lite", "voyage-code-3"]);
-export const zVoyageParameters = {
-    "voyage-3-large": z.object({
-        inputType: zVoyageInputType,
-        outputDimensions: zVoyageSupportedDimensions,
-        outputDtype: zVoyageQuantization,
-    }),
-    "voyage-3.5": z.object({
-        inputType: zVoyageInputType,
-        outputDimensions: zVoyageSupportedDimensions,
-        outputDtype: zVoyageQuantization,
-    }),
-    "voyage-3.5-lite": z.object({
-        inputType: zVoyageInputType,
-        outputDimensions: zVoyageSupportedDimensions,
-        outputDtype: zVoyageQuantization,
-    }),
-    "voyage-code-3": z.object({
-        inputType: zVoyageInputType,
-        outputDimensions: zVoyageSupportedDimensions,
-        outputDtype: zVoyageQuantization,
-    }),
-} as const;
 
 type VoyageModels = z.infer<typeof zVoyageModels>;
 class VoyageEmbeddingsProvider implements EmbeddingsProvider<VoyageModels> {
@@ -68,12 +45,32 @@ class VoyageEmbeddingsProvider implements EmbeddingsProvider<VoyageModels> {
     async embed<Model extends VoyageModels>(
         modelId: Model,
         content: EmbeddingsInput[],
-        parameters: z.infer<(typeof zVoyageParameters)[Model]>
+        parameters: EmbeddingParameters
     ): Promise<Embeddings[]> {
+        const voyageParameters = {
+            inputType: parameters.inputType,
+            outputDimensions: parameters.numDimensions,
+            quantization: "float", // it is hardcoded on purpose as we don't do quantization yet
+        };
+
         const model = this.voyage.textEmbeddingModel(modelId);
-        const { embeddings } = await embedMany({ model, values: content, providerOptions: { voyage: parameters } });
+        const { embeddings } = await embedMany({
+            model,
+            values: content,
+            providerOptions: { voyage: voyageParameters },
+        });
+
         return embeddings;
     }
+
+    async embedOne<Model extends VoyageModels>(
+        modelId: Model,
+        content: EmbeddingsInput,
+        parameters: EmbeddingParameters
+    ): Promise<Embeddings> {
+        const embeddings = await this.embed(modelId, [content], parameters);
+        return embeddings[0] ?? [];
+    }
 }
 
 export function getEmbeddingsProvider(userConfig: UserConfig): EmbeddingsProvider<VoyageModels> | undefined {
@@ -83,3 +80,6 @@ export function getEmbeddingsProvider(userConfig: UserConfig): EmbeddingsProvide
 
     return undefined;
 }
+
+export const zSupportedEmbeddingModels = zVoyageModels;
+export type SupportedEmbeddingModels = z.infer<typeof zSupportedEmbeddingModels>;

src/common/search/vectorSearchEmbeddingsManager.ts

@@ -3,6 +3,9 @@ import { BSON, type Document } from "bson";
 import type { UserConfig } from "../config.js";
 import type { ConnectionManager } from "../connectionManager.js";
 import z from "zod";
+import { ErrorCodes, MongoDBError } from "../errors.js";
+import { getEmbeddingsProvider } from "./embeddingsProvider.js";
+import type { SupportedEmbeddingModels } from "./embeddingsProvider.js";
 
 export const similarityEnum = z.enum(["cosine", "euclidean", "dotProduct"]);
 export type Similarity = z.infer<typeof similarityEnum>;
@@ -216,6 +219,94 @@ export class VectorSearchEmbeddingsManager {
         return undefined;
     }
 
+    public async generateEmbeddings({
+        database,
+        collection,
+        path,
+        model,
+        rawValues,
+        inputType,
+    }: {
+        database: string;
+        collection: string;
+        path: string;
+        model: SupportedEmbeddingModels;
+        rawValues: string[];
+        inputType: "query" | "document";
+    }): Promise<unknown[]> {
+        const provider = await this.assertAtlasSearchIsAvailable();
+        if (!provider) {
+            throw new MongoDBError(
+                ErrorCodes.AtlasSearchNotSupported,
+                "Atlas Search is not supported in this cluster."
+            );
+        }
+
+        const embeddingsProvider = getEmbeddingsProvider(this.config);
+
+        if (!embeddingsProvider) {
+            throw new MongoDBError(ErrorCodes.NoEmbeddingsProviderConfigured, "No embeddings provider configured.");
+        }
+
+        const embeddingInfoForCollection = await this.embeddingsForNamespace({ database, collection });
+        const embeddingInfoForPath = embeddingInfoForCollection.find((definition) => definition.path === path);
+
+        if (!embeddingInfoForPath) {
+            throw new MongoDBError(
+                ErrorCodes.AtlasVectorSearchIndexNotFound,
+                `No Vector Search index found for path "${path}" in namespace "${database}.${collection}"`
+            );
+        }
+
+        const providerEmbeddings = await embeddingsProvider.embed(model, rawValues, {
+            inputType,
+            numDimensions: embeddingInfoForPath.numDimensions,
+            quantization: embeddingInfoForPath.quantization,
+        });
+
+        if (this.config.disableEmbeddingsValidation) {
+            return providerEmbeddings;
+        }
+
+        const hasDocuments = await provider.estimatedDocumentCount(database, collection);
+        if (!hasDocuments) {
+            return providerEmbeddings;
+        }
+
+        const oneDocument: Document = await provider
+            .aggregate(database, collection, [{ $sample: { size: 1 } }, { $project: { embeddings: path } }])
+            .next();
+
+        if (!oneDocument) {
+            return providerEmbeddings;
+        }
+
+        const sampleEmbeddings = oneDocument.embeddings;
+        const adaptedEmbeddings = providerEmbeddings.map((embeddings) => {
+            // now map based on the sample embeddings
+            if (Array.isArray(sampleEmbeddings) && Array.isArray(embeddings)) {
+                return providerEmbeddings;
+            }
+            if (sampleEmbeddings instanceof BSON.Binary && Array.isArray(providerEmbeddings)) {
+                if (this.matches(() => sampleEmbeddings.toBits())) {
+                    return BSON.Binary.fromBits(embeddings);
+                }
+                if (this.matches(() => sampleEmbeddings.toInt8Array())) {
+                    return BSON.Binary.fromInt8Array(new Int8Array(embeddings));
+                }
+                if (this.matches(() => sampleEmbeddings.toFloat32Array())) {
+                    return BSON.Binary.fromFloat32Array(new Float32Array(embeddings));
+                }
+                if (this.matches(() => sampleEmbeddings.toPackedBits())) {
+                    return BSON.Binary.fromPackedBits(new Uint8Array(embeddings));
+                }
+            }
+            return embeddings;
+        });
+
+        return adaptedEmbeddings;
+    }
+
     private isANumber(value: unknown): boolean {
         if (typeof value === "number") {
             return true;
@@ -232,4 +323,13 @@ export class VectorSearchEmbeddingsManager {
 
         return false;
     }
+
+    private matches(fn: () => unknown): boolean {
+        try {
+            fn();
+            return true;
+        } catch {
+            return false;
+        }
+    }
 }

src/tools/mongodb/read/aggregate.ts

@@ -13,9 +13,53 @@ import { operationWithFallback } from "../../../helpers/operationWithFallback.js
 import { AGG_COUNT_MAX_TIME_MS_CAP, ONE_MB, CURSOR_LIMITS_TO_LLM_TEXT } from "../../../helpers/constants.js";
 import { zEJSON } from "../../args.js";
 import { LogId } from "../../../common/logger.js";
+import { SupportedEmbeddingModels, zSupportedEmbeddingModels } from "../../../common/search/embeddingsProvider.js";
+
+const AnyStage = zEJSON();
+const VectorSearchStage = z.object({
+    $vectorSearch: z
+        .object({
+            exact: z
+                .boolean()
+                .optional()
+                .default(false)
+                .describe(
+                    "When true, uses an ENN algorithm, otherwise uses ANN. Using ENN is not compatible with numCandidates, in that case, numCandidates must be left empty."
+                ),
+            index: z.string().optional().default("default"),
+            path: z
+                .string()
+                .describe(
+                    "Field, in dot notation, where to search. There must be a vector search index for that field. Note to LLM: When unsure, use the 'collection-indexes' tool to validate that the field is indexed with a vector search index."
+                ),
+            queryVector: z
+                .union([z.string(), z.array(z.number())])
+                .describe(
+                    "The content to search for. If a string, the embeddingModel is mandatory, if not, the embeddingModel is ignored."
+                ),
+            numCandidates: z
+                .number()
+                .int()
+                .positive()
+                .optional()
+                .describe("Number of candidates for the ANN algorithm. Only valid when exact is false."),
+            limit: z.number().int().positive().optional().default(10),
+            filter: zEJSON()
+                .optional()
+                .describe("MQL filter that can only use pre-filter fields from the index definition."),
+            embeddingModel: zSupportedEmbeddingModels.describe(
+                "The embedding model to use to generate embeddings before search. Note to LLM: If unsure, ask the user before providing one."
+            ),
+        })
+        .passthrough(),
+});
 
 export const AggregateArgs = {
-    pipeline: z.array(zEJSON()).describe("An array of aggregation stages to execute"),
+    pipeline: z
+        .array(z.union([AnyStage, VectorSearchStage]))
+        .describe(
+            "An array of aggregation stages to execute. $vectorSearch can only appear as the first stage of the aggregation pipeline or as the first stage of a $unionWith subpipeline."
+        ),
     responseBytesLimit: z.number().optional().default(ONE_MB).describe(`\
 The maximum number of bytes to return in the response. This value is capped by the server’s configured maxBytesPerQuery and cannot be exceeded. \
 Note to LLM: If the entire aggregation result is required, use the "export" tool instead of increasing this limit.\
@@ -38,8 +82,7 @@ export class AggregateTool extends MongoDBToolBase {
         let aggregationCursor: AggregationCursor | undefined = undefined;
         try {
             const provider = await this.ensureConnected();
-
-            this.assertOnlyUsesPermittedStages(pipeline);
+            await this.assertOnlyUsesPermittedStages(pipeline);
 
             // Check if aggregate operation uses an index if enabled
             if (this.config.indexCheck) {
@@ -50,6 +93,12 @@ export class AggregateTool extends MongoDBToolBase {
                 });
             }
 
+            pipeline = await this.replaceRawValuesWithEmbeddingsIfNecessary({
+                database,
+                collection,
+                pipeline,
+            });
+
             const cappedResultsPipeline = [...pipeline];
             if (this.config.maxDocumentsPerQuery > 0) {
                 cappedResultsPipeline.push({ $limit: this.config.maxDocumentsPerQuery });
@@ -107,8 +156,10 @@ export class AggregateTool extends MongoDBToolBase {
         }
     }
 
-    private assertOnlyUsesPermittedStages(pipeline: Record<string, unknown>[]): void {
+    private async assertOnlyUsesPermittedStages(pipeline: Record<string, unknown>[]): Promise<void> {
         const writeOperations: OperationType[] = ["update", "create", "delete"];
+        const isSearchSupported = await this.session.isSearchSupported();
+
         let writeStageForbiddenError = "";
 
         if (this.config.readOnly) {
@@ -118,14 +169,17 @@ export class AggregateTool extends MongoDBToolBase {
                 "When 'create', 'update', or 'delete' operations are disabled, you can not run pipelines with $out or $merge stages.";
         }
 
-        if (!writeStageForbiddenError) {
-            return;
-        }
-
         for (const stage of pipeline) {
-            if (stage.$out || stage.$merge) {
+            if ((stage.$out || stage.$merge) && writeStageForbiddenError) {
                 throw new MongoDBError(ErrorCodes.ForbiddenWriteOperation, writeStageForbiddenError);
             }
+
+            if (stage.$vectorSearch && !isSearchSupported) {
+                throw new MongoDBError(
+                    ErrorCodes.AtlasSearchNotSupported,
+                    "Atlas Search is not supported in this cluster."
+                );
+            }
         }
     }
 
@@ -160,6 +214,48 @@ export class AggregateTool extends MongoDBToolBase {
         }, undefined);
     }
 
+    private async replaceRawValuesWithEmbeddingsIfNecessary({
+        database,
+        collection,
+        pipeline,
+    }: {
+        database: string;
+        collection: string;
+        pipeline: Document[];
+    }): Promise<Document[]> {
+        for (const stage of pipeline) {
+            if (stage.$vectorSearch) {
+                if ("queryVector" in stage.$vectorSearch && Array.isArray(stage.$vectorSearch.queryVector)) {
+                    // if it's already embeddings, don't do anything
+                    continue;
+                }
+
+                if (!("embeddingModel" in stage.$vectorSearch)) {
+                    throw new MongoDBError(
+                        ErrorCodes.AtlasVectorSearchInvalidQuery,
+                        "embeddingModel is mandatory if queryVector is a raw string."
+                    );
+                }
+
+                const model = stage.$vectorSearch.embeddingModel as SupportedEmbeddingModels;
+                delete stage.$vectorSearch.embeddingModel;
+
+                const [embeddings] = await this.session.vectorSearchEmbeddingsManager.generateEmbeddings({
+                    database,
+                    collection,
+                    path: stage.$vectorSearch.path,
+                    model,
+                    rawValues: stage.$vectorSearch.queryVector,
+                    inputType: "query",
+                });
+
+                stage.$vectorSearch.queryVector = embeddings;
+            }
+        }
+
+        return pipeline;
+    }
+
     private generateMessage({
         aggResultsCount,
         documents,