chore: simplify, cleanup, add tests

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/vectorSearchEmbeddingsManager.ts

@@ -84,7 +84,28 @@ 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))
+        );
+        const embeddingValidations = new Set(embeddingValidationResults.flat());
+
+        if (embeddingValidations.size > 0) {
+            const embeddingValidationMessages = Array.from(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}`
+            );
+
+            throw new MongoDBError(ErrorCodes.AtlasVectorSearchInvalidQuery, embeddingValidationMessages.join("\n"));
+        }
+    }
+
+    private async findFieldsWithWrongEmbeddings(
         {
             database,
             collection,
@@ -220,21 +241,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<number[][]> {
         const provider = await this.atlasSearchEnabledProvider();
         if (!provider) {
             throw new MongoDBError(
@@ -256,15 +290,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

@@ -6,7 +6,14 @@ import { zEJSON } from "../../args.js";
 import { type Document } from "bson";
 import { zSupportedEmbeddingParameters } from "../../../common/search/embeddingsProvider.js";
 import { ErrorCodes, MongoDBError } from "../../../common/errors.js";
-import type { VectorFieldIndexDefinition } from "../../../common/search/vectorSearchEmbeddingsManager.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";
@@ -16,9 +23,9 @@ 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(). If you are asked to generate a embedding for a field, you have to explicitly specify the field name with a raw text string value of the field and an embedding will be generated if embeddingParameters is provided."
+                "The array of documents to insert, matching the syntax of the document argument of db.collection.insertMany()."
             ),
-        embeddingParameters: zSupportedEmbeddingParameters
+        embeddingParameters: zSupportedEmbeddingParametersWithInput
             .optional()
             .describe(
                 "The embedding model and its parameters to use to generate embeddings for fields that have vector search indexes. When a field has a vector search index and contains a plain text string in the document, embeddings will be automatically generated from that string value. Note to LLM: If unsure which embedding model to use, ask the user before providing one."
@@ -34,45 +41,14 @@ export class InsertManyTool extends MongoDBToolBase {
     }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
         const provider = await this.ensureConnected();
 
-        // Get vector search indexes for the collection
-        const vectorIndexes = await this.session.vectorSearchEmbeddingsManager.embeddingsForNamespace({
-            database,
-            collection,
-        });
-
         // Process documents to replace raw string values with generated embeddings
         documents = await this.replaceRawValuesWithEmbeddingsIfNecessary({
             database,
             collection,
             documents,
-            vectorIndexes,
             embeddingParameters,
         });
 
-        const embeddingValidationPromises = documents.map((document) =>
-            this.session.vectorSearchEmbeddingsManager.findFieldsWithWrongEmbeddings({ database, collection }, document)
-        );
-        const embeddingValidationResults = await Promise.all(embeddingValidationPromises);
-        const embeddingValidations = new Set(embeddingValidationResults.flat());
-
-        if (embeddingValidations.size > 0) {
-            // tell the LLM what happened
-            const embeddingValidationMessages = Array.from(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,
-            };
-        }
-
         const result = await provider.insertMany(database, collection, documents);
         const content = formatUntrustedData(
             "Documents were inserted successfully.",
@@ -88,144 +64,84 @@ export class InsertManyTool extends MongoDBToolBase {
         database,
         collection,
         documents,
-        vectorIndexes,
         embeddingParameters,
     }: {
         database: string;
         collection: string;
         documents: Document[];
-        vectorIndexes: VectorFieldIndexDefinition[];
-        embeddingParameters?: z.infer<typeof zSupportedEmbeddingParameters>;
+        embeddingParameters?: z.infer<typeof zSupportedEmbeddingParametersWithInput>;
     }): Promise<Document[]> {
-        // If no vector indexes, return documents as-is
-        if (vectorIndexes.length === 0) {
+        // If no embedding parameters or no input specified, return documents as-is
+        if (!embeddingParameters?.input || embeddingParameters.input.length === 0) {
             return documents;
         }
 
-        const processedDocuments: Document[] = [];
-
-        for (let i = 0; i < documents.length; i++) {
-            const document = documents[i];
-            if (!document) {
-                continue;
-            }
-            const processedDoc = await this.processDocumentForEmbeddings(
-                database,
-                collection,
-                document,
-                vectorIndexes,
-                embeddingParameters
-            );
-            processedDocuments.push(processedDoc);
-        }
-
-        return processedDocuments;
-    }
-
-    private async processDocumentForEmbeddings(
-        database: string,
-        collection: string,
-        document: Document,
-        vectorIndexes: VectorFieldIndexDefinition[],
-        embeddingParameters?: z.infer<typeof zSupportedEmbeddingParameters>
-    ): Promise<Document> {
-        // Find all fields in the document that match vector search indexed fields and need embeddings
-        const fieldsNeedingEmbeddings: Array<{
-            path: string;
-            rawValue: string;
-            indexDef: VectorFieldIndexDefinition;
-        }> = [];
+        // Get vector search indexes for the collection
+        const vectorIndexes = await this.session.vectorSearchEmbeddingsManager.embeddingsForNamespace({
+            database,
+            collection,
+        });
 
-        for (const indexDef of vectorIndexes) {
-            // Check if the field exists in the document and is a string (raw text)
-            const fieldValue = this.getFieldValue(document, indexDef.path);
-            if (typeof fieldValue === "string") {
-                fieldsNeedingEmbeddings.push({
-                    path: indexDef.path,
-                    rawValue: fieldValue,
-                    indexDef,
-                });
+        // 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.`
+                    );
+                }
             }
         }
 
-        // If no fields need embeddings, return document as-is
-        if (fieldsNeedingEmbeddings.length === 0) {
-            return document;
-        }
-
-        // Check if embeddingParameters is provided
-        if (!embeddingParameters) {
-            const fieldPaths = fieldsNeedingEmbeddings.map((f) => f.path).join(", ");
-            throw new MongoDBError(
-                ErrorCodes.AtlasVectorSearchInvalidQuery,
-                `Fields [${fieldPaths}] have vector search indexes and contain raw text strings. The embeddingParameters parameter is required to generate embeddings for these fields.`
-            );
-        }
-
-        // Generate embeddings for all fields
-        const embeddingsMap = new Map<string, number[]>();
-
-        for (const field of fieldsNeedingEmbeddings) {
-            const embeddings = await this.session.vectorSearchEmbeddingsManager.generateEmbeddings({
-                database,
-                collection,
-                path: field.path,
-                rawValues: [field.rawValue],
-                embeddingParameters,
-                inputType: "document",
-            });
+        // 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,
+            }))
+        );
 
-            if (embeddings.length > 0 && Array.isArray(embeddings[0])) {
-                embeddingsMap.set(field.path, embeddings[0] as number[]);
-            }
-        }
+        const generatedEmbeddings = await this.session.vectorSearchEmbeddingsManager.generateEmbeddings({
+            rawValues: flattenedEmbeddingsInput.map(({ rawTextValue }) => rawTextValue) as string[],
+            embeddingParameters,
+            inputType: "document",
+        });
 
-        // Replace raw string values with generated embeddings
-        const processedDoc = { ...document };
+        const processedDocuments: Document[] = [...documents];
 
-        for (const field of fieldsNeedingEmbeddings) {
-            const embedding = embeddingsMap.get(field.path);
-            if (embedding) {
-                this.setFieldValue(processedDoc, field.path, embedding);
+        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 processedDoc;
-    }
-
-    private getFieldValue(document: Document, path: string): unknown {
-        const parts = path.split(".");
-        let current: unknown = document;
-
-        for (const part of parts) {
-            if (current && typeof current === "object" && part in current) {
-                current = (current as Record<string, unknown>)[part];
-            } else {
-                return undefined;
-            }
-        }
+        await this.session.vectorSearchEmbeddingsManager.assertFieldsHaveCorrectEmbeddings(
+            { database, collection },
+            processedDocuments
+        );
 
-        return current;
+        return processedDocuments;
     }
 
-    private setFieldValue(document: Document, path: string, value: unknown): void {
-        const parts = path.split(".");
+    // 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 - 1; i++) {
+        for (let i = 0; i < parts.length; i++) {
             const part = parts[i];
-            if (!part) {
-                continue;
-            }
-            if (!(part in current) || typeof current[part] !== "object") {
-                current[part] = {};
+            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>;
             }
-            current = current[part] as Record<string, unknown>;
-        }
-
-        const lastPart = parts[parts.length - 1];
-        if (lastPart) {
-            current[lastPart] = value;
         }
     }
 }