feat: add support for automatic embeddings for the insert many tool MCP-236 (#688)

Author: gagik

.gitignore

@@ -13,3 +13,5 @@ tests/tmp
 coverage
 # Generated assets by accuracy runs
 .accuracy
+
+.DS_Store

CONTRIBUTING.md

@@ -76,6 +76,10 @@ npm test -- path/to/test/file.test.ts
 npm test -- path/to/directory
 ```
 
+#### Accuracy Tests and colima
+
+If you use [colima](https://github.com/abiosoft/colima) to run Docker on Mac, you will need to apply [additional configuration](https://node.testcontainers.org/supported-container-runtimes/#colima) to ensure the accuracy tests run correctly.
+
 ## Troubleshooting
 
 ### Restart Server

src/common/errors.ts

@@ -7,6 +7,7 @@ export enum ErrorCodes {
     NoEmbeddingsProviderConfigured = 1_000_005,
     AtlasVectorSearchIndexNotFound = 1_000_006,
     AtlasVectorSearchInvalidQuery = 1_000_007,
+    Unexpected = 1_000_008,
 }
 
 export class MongoDBError<ErrorCode extends ErrorCodes = ErrorCodes> extends Error {

src/common/search/embeddingsProvider.ts

@@ -7,7 +7,7 @@ import { createFetch } from "@mongodb-js/devtools-proxy-support";
 import { z } from "zod";
 
 type EmbeddingsInput = string;
-type Embeddings = number[];
+type Embeddings = number[] | unknown[];
 export type EmbeddingParameters = {
     inputType: "query" | "document";
 };

src/common/search/vectorSearchEmbeddingsManager.ts

@@ -6,6 +6,7 @@ import z from "zod";
 import { ErrorCodes, MongoDBError } from "../errors.js";
 import { getEmbeddingsProvider } from "./embeddingsProvider.js";
 import type { EmbeddingParameters, SupportedEmbeddingParameters } from "./embeddingsProvider.js";
+import { formatUntrustedData } from "../../tools/tool.js";
 
 export const similarityEnum = z.enum(["cosine", "euclidean", "dotProduct"]);
 export type Similarity = z.infer<typeof similarityEnum>;
@@ -103,7 +104,34 @@ export class VectorSearchEmbeddingsManager {
         return definition;
     }
 
-    async findFieldsWithWrongEmbeddings(
+    async assertFieldsHaveCorrectEmbeddings(
+        { database, collection }: { database: string; collection: string },
+        documents: Document[]
+    ): Promise<void> {
+        const embeddingValidationResults = (
+            await Promise.all(
+                documents.map((document) => this.findFieldsWithWrongEmbeddings({ database, collection }, document))
+            )
+        ).flat();
+
+        if (embeddingValidationResults.length > 0) {
+            const embeddingValidationMessages = embeddingValidationResults.map(
+                (validation) =>
+                    `- Field ${validation.path} is an embedding with ${validation.expectedNumDimensions} dimensions and ${validation.expectedQuantization}` +
+                    ` quantization, and the provided value is not compatible. Actual dimensions: ${validation.actualNumDimensions}, ` +
+                    `actual quantization: ${validation.actualQuantization}. Error: ${validation.error}`
+            );
+
+            throw new MongoDBError(
+                ErrorCodes.AtlasVectorSearchInvalidQuery,
+                formatUntrustedData("", ...embeddingValidationMessages)
+                    .map(({ text }) => text)
+                    .join("\n")
+            );
+        }
+    }
+
+    public async findFieldsWithWrongEmbeddings(
         {
             database,
             collection,
@@ -239,21 +267,34 @@ export class VectorSearchEmbeddingsManager {
         return undefined;
     }
 
-    public async generateEmbeddings({
+    public async assertVectorSearchIndexExists({
         database,
         collection,
         path,
-        rawValues,
-        embeddingParameters,
-        inputType,
     }: {
         database: string;
         collection: string;
         path: string;
+    }): Promise<void> {
+        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}"`
+            );
+        }
+    }
+
+    public async generateEmbeddings({
+        rawValues,
+        embeddingParameters,
+        inputType,
+    }: {
         rawValues: string[];
         embeddingParameters: SupportedEmbeddingParameters;
         inputType: EmbeddingParameters["inputType"];
-    }): Promise<unknown[]> {
+    }): Promise<unknown[][]> {
         const provider = await this.atlasSearchEnabledProvider();
         if (!provider) {
             throw new MongoDBError(
@@ -275,15 +316,6 @@ export class VectorSearchEmbeddingsManager {
             });
         }
 
-        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}"`
-            );
-        }
-
         return await embeddingsProvider.embed(embeddingParameters.model, rawValues, {
             inputType,
             ...embeddingParameters,

src/tools/mongodb/create/insertMany.ts

@@ -3,6 +3,17 @@ import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
 import { type ToolArgs, type OperationType, formatUntrustedData } from "../../tool.js";
 import { zEJSON } from "../../args.js";
+import { type Document } from "bson";
+import { zSupportedEmbeddingParameters } from "../../../common/search/embeddingsProvider.js";
+import { ErrorCodes, MongoDBError } from "../../../common/errors.js";
+
+const zSupportedEmbeddingParametersWithInput = zSupportedEmbeddingParameters.extend({
+    input: z
+        .array(z.object({}).passthrough())
+        .describe(
+            "Array of objects with vector search index fields as keys (in dot notation) and the raw text values to generate embeddings for as values. The index of each object corresponds to the index of the document in the documents array."
+        ),
+});
 
 export class InsertManyTool extends MongoDBToolBase {
     public name = "insert-many";
@@ -12,46 +23,44 @@ export class InsertManyTool extends MongoDBToolBase {
         documents: z
             .array(zEJSON().describe("An individual MongoDB document"))
             .describe(
-                "The array of documents to insert, matching the syntax of the document argument of db.collection.insertMany()"
+                "The array of documents to insert, matching the syntax of the document argument of db.collection.insertMany()."
             ),
+        ...(this.isFeatureEnabled("vectorSearch")
+            ? {
+                  embeddingParameters: zSupportedEmbeddingParametersWithInput
+                      .optional()
+                      .describe(
+                          "The embedding model and its parameters to use to generate embeddings for fields with vector search indexes. Note to LLM: If unsure which embedding model to use, ask the user before providing one."
+                      ),
+              }
+            : {}),
     };
     public operationType: OperationType = "create";
 
     protected async execute({
         database,
         collection,
         documents,
+        embeddingParameters: providedEmbeddingParameters,
     }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
         const provider = await this.ensureConnected();
 
-        const embeddingValidations = new Set(
-            ...(await Promise.all(
-                documents.flatMap((document) =>
-                    this.session.vectorSearchEmbeddingsManager.findFieldsWithWrongEmbeddings(
-                        { database, collection },
-                        document
-                    )
-                )
-            ))
-        );
+        const embeddingParameters = this.isFeatureEnabled("vectorSearch")
+            ? (providedEmbeddingParameters as z.infer<typeof zSupportedEmbeddingParametersWithInput>)
+            : undefined;
 
-        if (embeddingValidations.size > 0) {
-            // tell the LLM what happened
-            const embeddingValidationMessages = [...embeddingValidations].map(
-                (validation) =>
-                    `- Field ${validation.path} is an embedding with ${validation.expectedNumDimensions} dimensions and ${validation.expectedQuantization}` +
-                    ` quantization, and the provided value is not compatible. Actual dimensions: ${validation.actualNumDimensions}, ` +
-                    `actual quantization: ${validation.actualQuantization}. Error: ${validation.error}`
-            );
-
-            return {
-                content: formatUntrustedData(
-                    "There were errors when inserting documents. No document was inserted.",
-                    ...embeddingValidationMessages
-                ),
-                isError: true,
-            };
-        }
+        // Process documents to replace raw string values with generated embeddings
+        documents = await this.replaceRawValuesWithEmbeddingsIfNecessary({
+            database,
+            collection,
+            documents,
+            embeddingParameters,
+        });
+
+        await this.session.vectorSearchEmbeddingsManager.assertFieldsHaveCorrectEmbeddings(
+            { database, collection },
+            documents
+        );
 
         const result = await provider.insertMany(database, collection, documents);
         const content = formatUntrustedData(
@@ -63,4 +72,84 @@ export class InsertManyTool extends MongoDBToolBase {
             content,
         };
     }
+
+    private async replaceRawValuesWithEmbeddingsIfNecessary({
+        database,
+        collection,
+        documents,
+        embeddingParameters,
+    }: {
+        database: string;
+        collection: string;
+        documents: Document[];
+        embeddingParameters?: z.infer<typeof zSupportedEmbeddingParametersWithInput>;
+    }): Promise<Document[]> {
+        // If no embedding parameters or no input specified, return documents as-is
+        if (!embeddingParameters?.input || embeddingParameters.input.length === 0) {
+            return documents;
+        }
+
+        // Get vector search indexes for the collection
+        const vectorIndexes = await this.session.vectorSearchEmbeddingsManager.embeddingsForNamespace({
+            database,
+            collection,
+        });
+
+        // Ensure for inputted fields, the vector search index exists.
+        for (const input of embeddingParameters.input) {
+            for (const fieldPath of Object.keys(input)) {
+                if (!vectorIndexes.some((index) => index.path === fieldPath)) {
+                    throw new MongoDBError(
+                        ErrorCodes.AtlasVectorSearchInvalidQuery,
+                        `Field '${fieldPath}' does not have a vector search index in collection ${database}.${collection}. Only fields with vector search indexes can have embeddings generated.`
+                    );
+                }
+            }
+        }
+
+        // We make one call to generate embeddings for all documents at once to avoid making too many API calls.
+        const flattenedEmbeddingsInput = embeddingParameters.input.flatMap((documentInput, index) =>
+            Object.entries(documentInput).map(([fieldPath, rawTextValue]) => ({
+                fieldPath,
+                rawTextValue,
+                documentIndex: index,
+            }))
+        );
+
+        const generatedEmbeddings = await this.session.vectorSearchEmbeddingsManager.generateEmbeddings({
+            rawValues: flattenedEmbeddingsInput.map(({ rawTextValue }) => rawTextValue) as string[],
+            embeddingParameters,
+            inputType: "document",
+        });
+
+        const processedDocuments: Document[] = [...documents];
+
+        for (const [index, { fieldPath, documentIndex }] of flattenedEmbeddingsInput.entries()) {
+            if (!processedDocuments[documentIndex]) {
+                throw new MongoDBError(ErrorCodes.Unexpected, `Document at index ${documentIndex} does not exist.`);
+            }
+            // Ensure no nested fields are present in the field path.
+            this.deleteFieldPath(processedDocuments[documentIndex], fieldPath);
+            processedDocuments[documentIndex][fieldPath] = generatedEmbeddings[index];
+        }
+
+        return processedDocuments;
+    }
+
+    // Delete a specified field path from a document using dot notation.
+    private deleteFieldPath(document: Record<string, unknown>, fieldPath: string): void {
+        const parts = fieldPath.split(".");
+        let current: Record<string, unknown> = document;
+        for (let i = 0; i < parts.length; i++) {
+            const part = parts[i];
+            const key = part as keyof typeof current;
+            if (!current[key]) {
+                return;
+            } else if (i === parts.length - 1) {
+                delete current[key];
+            } else {
+                current = current[key] as Record<string, unknown>;
+            }
+        }
+    }
 }

src/tools/mongodb/read/aggregate.ts

@@ -276,22 +276,37 @@ export class AggregateTool extends MongoDBToolBase {
                 const embeddingParameters = vectorSearchStage.embeddingParameters;
                 delete vectorSearchStage.embeddingParameters;
 
-                const [embeddings] = await this.session.vectorSearchEmbeddingsManager.generateEmbeddings({
+                await this.session.vectorSearchEmbeddingsManager.assertVectorSearchIndexExists({
                     database,
                     collection,
                     path: vectorSearchStage.path,
+                });
+
+                const [embeddings] = await this.session.vectorSearchEmbeddingsManager.generateEmbeddings({
                     rawValues: [vectorSearchStage.queryVector],
                     embeddingParameters,
                     inputType: "query",
                 });
 
+                if (!embeddings) {
+                    throw new MongoDBError(
+                        ErrorCodes.AtlasVectorSearchInvalidQuery,
+                        "Failed to generate embeddings for the query vector."
+                    );
+                }
+
                 // $vectorSearch.queryVector can be a BSON.Binary: that it's not either number or an array.
                 // It's not exactly valid from the LLM perspective (they can't provide binaries).
                 // That's why we overwrite the stage in an untyped way, as what we expose and what LLMs can use is different.
-                vectorSearchStage.queryVector = embeddings as number[];
+                vectorSearchStage.queryVector = embeddings as string | number[];
             }
         }
 
+        await this.session.vectorSearchEmbeddingsManager.assertFieldsHaveCorrectEmbeddings(
+            { database, collection },
+            pipeline
+        );
+
         return pipeline;
     }