18912204743

Committed 29 Oct 2025 02:55PM UTC coverage: 79.885% (-0.3%) from 80.227%

Build # 18912204743

Build Type

Pull #688

github

Committed by

web-flow

Commit Message

Merge 645ff8615 into 01f799ccb

Pull Request Pull Request #688: feat: add support for automatic embeddings for the insert many tool

Run Details

1364 of 1844 branches covered (73.97%)

Branch coverage included in aggregate %.

92 of 151 new or added lines in 4 files covered. (60.93%)

7 existing lines in 1 file now uncovered.

6412 of 7890 relevant lines covered (81.27%)

70.65 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

61.64

/src/tools/mongodb/create/insertMany.ts

import { z } from "zod";
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";
    protected description = "Insert an array of documents into a MongoDB collection";
    protected argsShape = {
        ...DbOperationArgs,
        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()."
            ),
        ...(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 embeddingParameters = this.isFeatureEnabled("vectorSearch")
            ? (providedEmbeddingParameters as z.infer<typeof zSupportedEmbeddingParametersWithInput>)
            : undefined;

        // 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(
            "Documents were inserted successfully.",
            `Inserted \`${result.insertedCount}\` document(s) into ${database}.${collection}.`,
            `Inserted IDs: ${Object.values(result.insertedIds).join(", ")}`
        );
        return {
            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>;
            }
        }
    }
}

1	import { z } from "zod";	3✔
2	import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
3	import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";	3✔
4	import { type ToolArgs, type OperationType, formatUntrustedData } from "../../tool.js";	3✔
5	import { zEJSON } from "../../args.js";	3✔
6	import { type Document } from "bson";
7	import { zSupportedEmbeddingParameters } from "../../../common/search/embeddingsProvider.js";	3✔
8	import { ErrorCodes, MongoDBError } from "../../../common/errors.js";	3✔
9
10	const zSupportedEmbeddingParametersWithInput = zSupportedEmbeddingParameters.extend({	3✔
11	input: z	3✔
12	.array(z.object({}).passthrough())	3✔
13	.describe(	3✔
14	"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."	3✔
15	),	3✔
16	});	3✔
17
18	export class InsertManyTool extends MongoDBToolBase {	3✔
19	public name = "insert-many";	99✔
20	protected description = "Insert an array of documents into a MongoDB collection";	99✔
21	protected argsShape = {	99✔
22	...DbOperationArgs,	99✔
23	documents: z	99✔
24	.array(zEJSON().describe("An individual MongoDB document"))	99✔
25	.describe(	99✔
26	"The array of documents to insert, matching the syntax of the document argument of db.collection.insertMany()."	99✔
27	),	99✔
28	...(this.isFeatureEnabled("vectorSearch")	99✔
29	? {	13✔
30	embeddingParameters: zSupportedEmbeddingParametersWithInput	13✔
31	.optional()	13✔
32	.describe(	13✔
33	"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."	13✔
34	),	13✔
35	}	13✔
36	: {}),	86✔
37	};	99✔
38	public operationType: OperationType = "create";	99✔
39
40	protected async execute({	3✔
41	database,	7✔
42	collection,	7✔
43	documents,	7✔
44	embeddingParameters: providedEmbeddingParameters,	7✔
45	}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {	7✔
46	const provider = await this.ensureConnected();	7✔
47
48	const embeddingParameters = this.isFeatureEnabled("vectorSearch")	6!
49	? (providedEmbeddingParameters as z.infer<typeof zSupportedEmbeddingParametersWithInput>)	2✔
50	: undefined;	4✔
51
52	// Process documents to replace raw string values with generated embeddings
53	documents = await this.replaceRawValuesWithEmbeddingsIfNecessary({	7✔
54	database,	7✔
55	collection,	7✔
56	documents,	7✔
57	embeddingParameters,	7✔
58	});	7✔
59
60	await this.session.vectorSearchEmbeddingsManager.assertFieldsHaveCorrectEmbeddings(	6✔
61	{ database, collection },	6✔
62	documents	6✔
63	);	6✔
64
65	const result = await provider.insertMany(database, collection, documents);	5✔
66	const content = formatUntrustedData(	4✔
67	"Documents were inserted successfully.",	4✔
68	`Inserted \`${result.insertedCount}\` document(s) into ${database}.${collection}.`,	4✔
69	`Inserted IDs: ${Object.values(result.insertedIds).join(", ")}`	4✔
70	);	4✔
71	return {	4✔
72	content,	4✔
73	};	4✔
74	}	7✔
75
76	private async replaceRawValuesWithEmbeddingsIfNecessary({	3✔
77	database,	6✔
78	collection,	6✔
79	documents,	6✔
80	embeddingParameters,	6✔
81	}: {	6✔
82	database: string;
83	collection: string;
84	documents: Document[];
85	embeddingParameters?: z.infer<typeof zSupportedEmbeddingParametersWithInput>;
86	}): Promise<Document[]> {	6✔
87	// If no embedding parameters or no input specified, return documents as-is
88	if (!embeddingParameters?.input \|\| embeddingParameters.input.length === 0) {	6!
89	return documents;	6✔
90	}	6!
91
92	// Get vector search indexes for the collection
NEW 93	const vectorIndexes = await this.session.vectorSearchEmbeddingsManager.embeddingsForNamespace({	×
NEW 94	database,	×
NEW 95	collection,	×
NEW 96	});	×
97
98	// Ensure for inputted fields, the vector search index exists.
NEW 99	for (const input of embeddingParameters.input) {	×
NEW 100	for (const fieldPath of Object.keys(input)) {	×
NEW 101	if (!vectorIndexes.some((index) => index.path === fieldPath)) {	×
NEW 102	throw new MongoDBError(	×
NEW 103	ErrorCodes.AtlasVectorSearchInvalidQuery,	×
NEW 104	`Field '${fieldPath}' does not have a vector search index in collection ${database}.${collection}. Only fields with vector search indexes can have embeddings generated.`	×
NEW 105	);	×
NEW 106	}	×
NEW 107	}	×
NEW 108	}	×
109
110	// We make one call to generate embeddings for all documents at once to avoid making too many API calls.
NEW 111	const flattenedEmbeddingsInput = embeddingParameters.input.flatMap((documentInput, index) =>	×
NEW 112	Object.entries(documentInput).map(([fieldPath, rawTextValue]) => ({	×
NEW 113	fieldPath,	×
NEW 114	rawTextValue,	×
NEW 115	documentIndex: index,	×
NEW 116	}))	×
NEW 117	);	×
118
NEW 119	const generatedEmbeddings = await this.session.vectorSearchEmbeddingsManager.generateEmbeddings({	×
NEW 120	rawValues: flattenedEmbeddingsInput.map(({ rawTextValue }) => rawTextValue) as string[],	×
NEW 121	embeddingParameters,	×
NEW 122	inputType: "document",	×
NEW 123	});	×
124
NEW 125	const processedDocuments: Document[] = [...documents];	×
126
NEW 127	for (const [index, { fieldPath, documentIndex }] of flattenedEmbeddingsInput.entries()) {	×
NEW 128	if (!processedDocuments[documentIndex]) {	×
NEW 129	throw new MongoDBError(ErrorCodes.Unexpected, `Document at index ${documentIndex} does not exist.`);	×
NEW 130	}	×
131	// Ensure no nested fields are present in the field path.
NEW 132	this.deleteFieldPath(processedDocuments[documentIndex], fieldPath);	×
NEW 133	processedDocuments[documentIndex][fieldPath] = generatedEmbeddings[index];	×
NEW 134	}	×
135
NEW 136	return processedDocuments;	×
137	}	6✔
138
139	// Delete a specified field path from a document using dot notation.
140	private deleteFieldPath(document: Record<string, unknown>, fieldPath: string): void {	3✔
NEW 141	const parts = fieldPath.split(".");	×
NEW 142	let current: Record<string, unknown> = document;	×
NEW 143	for (let i = 0; i < parts.length; i++) {	×
NEW 144	const part = parts[i];	×
NEW 145	const key = part as keyof typeof current;	×
NEW 146	if (!current[key]) {	×
NEW 147	return;	×
NEW 148	} else if (i === parts.length - 1) {	×
NEW 149	delete current[key];	×
NEW 150	} else {	×
NEW 151	current = current[key] as Record<string, unknown>;	×
NEW 152	}	×
NEW 153	}	×
NEW 154	}	×
155	}	3✔

mongodb-js / mongodb-mcp-server / 18912204743

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous