20968727349

Committed 13 Jan 2026 06:53PM UTC coverage: 80.405% (+0.2%) from 80.224%

Build # 20968727349

Build Type

push

github

Committed by

web-flow

Commit Message

fix: fixes wrong path getting updated with generated embeddings MCP-357 (#855)

Run Details

1550 of 2013 branches covered (77.0%)

Branch coverage included in aggregate %.

37 of 40 new or added lines in 2 files covered. (92.5%)

7104 of 8750 relevant lines covered (81.19%)

87.7 hits per line

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

66.89

/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 "../mongodbSchemas.js";
import { ErrorCodes, MongoDBError } from "../../../common/errors.js";
import type { ConnectionMetadata, AutoEmbeddingsUsageMetadata } from "../../../telemetry/types.js";
import { setFieldPath } from "../../../helpers/manageNestedFieldPaths.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."
        ),
});

const commonArgs = {
    ...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()."
        ),
} as const;

export class InsertManyTool extends MongoDBToolBase {
    public name = "insert-many";
    public description =
        "Insert an array of documents into a MongoDB collection. If the list of documents is above com.mongodb/maxRequestPayloadBytes, consider inserting them in batches.";
    public argsShape = this.isFeatureEnabled("search")
        ? {
              ...commonArgs,
              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."
                  ),
          }
        : commonArgs;
    static operationType: OperationType = "create";

    protected async execute({
        database,
        collection,
        documents,
        ...conditionalArgs
    }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
        const provider = await this.ensureConnected();

        let embeddingParameters: z.infer<typeof zSupportedEmbeddingParametersWithInput> | undefined;
        if ("embeddingParameters" in conditionalArgs) {
            embeddingParameters = conditionalArgs.embeddingParameters;
        }

        // 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.`);
            }
            setFieldPath(processedDocuments[documentIndex], fieldPath, generatedEmbeddings[index]);
        }

        return processedDocuments;
    }

    protected resolveTelemetryMetadata(
        args: ToolArgs<typeof this.argsShape>,
        { result }: { result: CallToolResult }
    ): ConnectionMetadata | AutoEmbeddingsUsageMetadata {
        if ("embeddingParameters" in args && this.config.voyageApiKey) {
            return {
                ...super.resolveTelemetryMetadata(args, { result }),
                embeddingsGeneratedBy: "mcp",
            };
        } else {
            return super.resolveTelemetryMetadata(args, { result });
        }
    }
}

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

mongodb-js / mongodb-mcp-server / 20968727349

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