18717167438

Committed 22 Oct 2025 01:06PM UTC coverage: 79.906% (-2.0%) from 81.905%

Build # 18717167438

Build Type

push

github

Committed by

web-flow

Commit Message

chore: When querying with vectorSearch use the generated embeddings MCP-245 (#662)

Run Details

1341 of 1809 branches covered (74.13%)

Branch coverage included in aggregate %.

128 of 193 new or added lines in 4 files covered. (66.32%)

198 existing lines in 6 files now uncovered.

6143 of 7557 relevant lines covered (81.29%)

73.09 hits per line

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

81.88

/src/common/search/vectorSearchEmbeddingsManager.ts

import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
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 { EmbeddingParameters, SupportedEmbeddingParameters } from "./embeddingsProvider.js";

export const similarityEnum = z.enum(["cosine", "euclidean", "dotProduct"]);
export type Similarity = z.infer<typeof similarityEnum>;

export const quantizationEnum = z.enum(["none", "scalar", "binary"]);
export type Quantization = z.infer<typeof quantizationEnum>;

export type VectorFieldIndexDefinition = {
    type: "vector";
    path: string;
    numDimensions: number;
    quantization: Quantization;
    similarity: Similarity;
};

export type VectorFieldValidationError = {
    path: string;
    expectedNumDimensions: number;
    expectedQuantization: Quantization;
    actualNumDimensions: number | "unknown";
    actualQuantization: Quantization | "unknown";
    error: "dimension-mismatch" | "quantization-mismatch" | "not-a-vector" | "not-numeric";
};

export type EmbeddingNamespace = `${string}.${string}`;
export class VectorSearchEmbeddingsManager {
    constructor(
        private readonly config: UserConfig,
        private readonly connectionManager: ConnectionManager,
        private readonly embeddings: Map<EmbeddingNamespace, VectorFieldIndexDefinition[]> = new Map(),
        private readonly embeddingsProvider: typeof getEmbeddingsProvider = getEmbeddingsProvider
    ) {
        connectionManager.events.on("connection-close", () => {
            this.embeddings.clear();
        });
    }

    cleanupEmbeddingsForNamespace({ database, collection }: { database: string; collection: string }): void {
        const embeddingDefKey: EmbeddingNamespace = `${database}.${collection}`;
        this.embeddings.delete(embeddingDefKey);
    }

    async embeddingsForNamespace({
        database,
        collection,
    }: {
        database: string;
        collection: string;
    }): Promise<VectorFieldIndexDefinition[]> {
        const provider = await this.atlasSearchEnabledProvider();
        if (!provider) {
            return [];
        }

        // We only need the embeddings for validation now, so don't query them if
        // validation is disabled.
        if (this.config.disableEmbeddingsValidation) {
            return [];
        }

        const embeddingDefKey: EmbeddingNamespace = `${database}.${collection}`;
        const definition = this.embeddings.get(embeddingDefKey);

        if (!definition) {
            const allSearchIndexes = await provider.getSearchIndexes(database, collection);
            const vectorSearchIndexes = allSearchIndexes.filter((index) => index.type === "vectorSearch");
            const vectorFields = vectorSearchIndexes
                // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
                .flatMap<Document>((index) => (index.latestDefinition?.fields as Document[]) ?? [])
                .filter((field) => this.isVectorFieldIndexDefinition(field));

            this.embeddings.set(embeddingDefKey, vectorFields);
            return vectorFields;
        }

        return definition;
    }

    async findFieldsWithWrongEmbeddings(
        {
            database,
            collection,
        }: {
            database: string;
            collection: string;
        },
        document: Document
    ): Promise<VectorFieldValidationError[]> {
        const provider = await this.atlasSearchEnabledProvider();
        if (!provider) {
            return [];
        }

        // While we can do our best effort to ensure that the embedding validation is correct
        // based on https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-quantization/
        // it's a complex process so we will also give the user the ability to disable this validation
        if (this.config.disableEmbeddingsValidation) {
            return [];
        }

        const embeddings = await this.embeddingsForNamespace({ database, collection });
        return embeddings
            .map((emb) => this.getValidationErrorForDocument(emb, document))
            .filter((e) => e !== undefined);
    }

    private async atlasSearchEnabledProvider(): Promise<NodeDriverServiceProvider | null> {
        const connectionState = this.connectionManager.currentConnectionState;
        if (connectionState.tag === "connected" && (await connectionState.isSearchSupported())) {
            return connectionState.serviceProvider;
        }

        return null;
    }

    private isVectorFieldIndexDefinition(doc: Document): doc is VectorFieldIndexDefinition {
        return doc["type"] === "vector";
    }

    private getValidationErrorForDocument(
        definition: VectorFieldIndexDefinition,
        document: Document
    ): VectorFieldValidationError | undefined {
        const fieldPath = definition.path.split(".");
        let fieldRef: unknown = document;

        const constructError = (
            details: Partial<Pick<VectorFieldValidationError, "error" | "actualNumDimensions" | "actualQuantization">>
        ): VectorFieldValidationError => ({
            path: definition.path,
            expectedNumDimensions: definition.numDimensions,
            expectedQuantization: definition.quantization,
            actualNumDimensions: details.actualNumDimensions ?? "unknown",
            actualQuantization: details.actualQuantization ?? "unknown",
            error: details.error ?? "not-a-vector",
        });

        for (const field of fieldPath) {
            if (fieldRef && typeof fieldRef === "object" && field in fieldRef) {
                fieldRef = (fieldRef as Record<string, unknown>)[field];
            } else {
                return undefined;
            }
        }

        switch (definition.quantization) {
            // Because quantization is not defined by the user
            // we have to trust them in the format they use.
            case "none":
                return undefined;
            case "scalar":
            case "binary":
                if (fieldRef instanceof BSON.Binary) {
                    try {
                        const elements = fieldRef.toFloat32Array();
                        if (elements.length !== definition.numDimensions) {
                            return constructError({
                                actualNumDimensions: elements.length,
                                actualQuantization: "binary",
                                error: "dimension-mismatch",
                            });
                        }

                        return undefined;
                    } catch {
                        // bits are also supported
                        try {
                            const bits = fieldRef.toBits();
                            if (bits.length !== definition.numDimensions) {
                                return constructError({
                                    actualNumDimensions: bits.length,
                                    actualQuantization: "binary",
                                    error: "dimension-mismatch",
                                });
                            }

                            return undefined;
                        } catch {
                            return constructError({
                                actualQuantization: "binary",
                                error: "not-a-vector",
                            });
                        }
                    }
                } else {
                    if (!Array.isArray(fieldRef)) {
                        return constructError({
                            error: "not-a-vector",
                        });
                    }

                    if (fieldRef.length !== definition.numDimensions) {
                        return constructError({
                            actualNumDimensions: fieldRef.length,
                            actualQuantization: "scalar",
                            error: "dimension-mismatch",
                        });
                    }

                    if (!fieldRef.every((e) => this.isANumber(e))) {
                        return constructError({
                            actualNumDimensions: fieldRef.length,
                            actualQuantization: "scalar",
                            error: "not-numeric",
                        });
                    }
                }

                break;
        }

        return undefined;
    }

    public async generateEmbeddings({
        database,
        collection,
        path,
        rawValues,
        embeddingParameters,
        inputType,
    }: {
        database: string;
        collection: string;
        path: string;
        rawValues: string[];
        embeddingParameters: SupportedEmbeddingParameters;
        inputType: EmbeddingParameters["inputType"];
    }): Promise<unknown[]> {
        const provider = await this.atlasSearchEnabledProvider();
        if (!provider) {
            throw new MongoDBError(
                ErrorCodes.AtlasSearchNotSupported,
                "Atlas Search is not supported in this cluster."
            );
        }

        const embeddingsProvider = this.embeddingsProvider(this.config);

        if (!embeddingsProvider) {
            throw new MongoDBError(ErrorCodes.NoEmbeddingsProviderConfigured, "No embeddings provider configured.");
        }

        if (this.config.disableEmbeddingsValidation) {
            return await embeddingsProvider.embed(embeddingParameters.model, rawValues, {
                inputType,
                ...embeddingParameters,
            });
        }

        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,
        });
    }

    private isANumber(value: unknown): boolean {
        if (typeof value === "number") {
            return true;
        }

        if (
            value instanceof BSON.Int32 ||
            value instanceof BSON.Decimal128 ||
            value instanceof BSON.Double ||
            value instanceof BSON.Long
        ) {
            return true;
        }

        return false;
    }
}

1	import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
2	import { BSON, type Document } from "bson";	3✔
3	import type { UserConfig } from "../config.js";
4	import type { ConnectionManager } from "../connectionManager.js";
5	import z from "zod";	3✔
6	import { ErrorCodes, MongoDBError } from "../errors.js";	3✔
7	import { getEmbeddingsProvider } from "./embeddingsProvider.js";	3✔
8	import type { EmbeddingParameters, SupportedEmbeddingParameters } from "./embeddingsProvider.js";
9
10	export const similarityEnum = z.enum(["cosine", "euclidean", "dotProduct"]);	3✔
11	export type Similarity = z.infer<typeof similarityEnum>;
12
13	export const quantizationEnum = z.enum(["none", "scalar", "binary"]);	3✔
14	export type Quantization = z.infer<typeof quantizationEnum>;
15
16	export type VectorFieldIndexDefinition = {
17	type: "vector";
18	path: string;
19	numDimensions: number;
20	quantization: Quantization;
21	similarity: Similarity;
22	};
23
24	export type VectorFieldValidationError = {
25	path: string;
26	expectedNumDimensions: number;
27	expectedQuantization: Quantization;
28	actualNumDimensions: number \| "unknown";
29	actualQuantization: Quantization \| "unknown";
30	error: "dimension-mismatch" \| "quantization-mismatch" \| "not-a-vector" \| "not-numeric";
31	};
32
33	export type EmbeddingNamespace = `${string}.${string}`;
34	export class VectorSearchEmbeddingsManager {	3✔
35	constructor(	3✔
36	private readonly config: UserConfig,	139✔
37	private readonly connectionManager: ConnectionManager,	139✔
38	private readonly embeddings: Map<EmbeddingNamespace, VectorFieldIndexDefinition[]> = new Map(),	139✔
39	private readonly embeddingsProvider: typeof getEmbeddingsProvider = getEmbeddingsProvider	139✔
40	) {	139✔
41	connectionManager.events.on("connection-close", () => {	139✔
42	this.embeddings.clear();	305✔
43	});	139✔
44	}	139✔
45
46	cleanupEmbeddingsForNamespace({ database, collection }: { database: string; collection: string }): void {	3✔
47	const embeddingDefKey: EmbeddingNamespace = `${database}.${collection}`;	4✔
48	this.embeddings.delete(embeddingDefKey);	4✔
49	}	4✔
50
51	async embeddingsForNamespace({	3✔
52	database,	25✔
53	collection,	25✔
54	}: {	25✔
55	database: string;
56	collection: string;
57	}): Promise<VectorFieldIndexDefinition[]> {	25✔
58	const provider = await this.atlasSearchEnabledProvider();	25✔
59	if (!provider) {	25!
60	return [];	×
61	}	×
62
63	// We only need the embeddings for validation now, so don't query them if
64	// validation is disabled.
65	if (this.config.disableEmbeddingsValidation) {	25!
66	return [];	×
67	}	×
68
69	const embeddingDefKey: EmbeddingNamespace = `${database}.${collection}`;	25✔
70	const definition = this.embeddings.get(embeddingDefKey);	25✔
71
72	if (!definition) {	25✔
73	const allSearchIndexes = await provider.getSearchIndexes(database, collection);	12✔
74	const vectorSearchIndexes = allSearchIndexes.filter((index) => index.type === "vectorSearch");	11✔
75	const vectorFields = vectorSearchIndexes	11✔
76	// eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
77	.flatMap<Document>((index) => (index.latestDefinition?.fields as Document[]) ?? [])	11!
78	.filter((field) => this.isVectorFieldIndexDefinition(field));	11✔
79
80	this.embeddings.set(embeddingDefKey, vectorFields);	11✔
81	return vectorFields;	11✔
82	}	11!
83
84	return definition;	13✔
85	}	25✔
86
87	async findFieldsWithWrongEmbeddings(	3✔
88	{	22✔
89	database,	22✔
90	collection,	22✔
91	}: {	22✔
92	database: string;
93	collection: string;
94	},
95	document: Document	22✔
96	): Promise<VectorFieldValidationError[]> {	22✔
97	const provider = await this.atlasSearchEnabledProvider();	22✔
98	if (!provider) {	22!
99	return [];	3✔
100	}	3✔
101
102	// While we can do our best effort to ensure that the embedding validation is correct
103	// based on https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-quantization/
104	// it's a complex process so we will also give the user the ability to disable this validation
105	if (this.config.disableEmbeddingsValidation) {	22!
106	return [];	3✔
107	}	3✔
108
109	const embeddings = await this.embeddingsForNamespace({ database, collection });	16✔
110	return embeddings	16✔
111	.map((emb) => this.getValidationErrorForDocument(emb, document))	16✔
112	.filter((e) => e !== undefined);	16✔
113	}	22✔
114
115	private async atlasSearchEnabledProvider(): Promise<NodeDriverServiceProvider \| null> {	3✔
116	const connectionState = this.connectionManager.currentConnectionState;	51✔
117	if (connectionState.tag === "connected" && (await connectionState.isSearchSupported())) {	51✔
118	return connectionState.serviceProvider;	48✔
119	}	48!
120
121	return null;	3✔
122	}	51✔
123
124	private isVectorFieldIndexDefinition(doc: Document): doc is VectorFieldIndexDefinition {	3✔
125	return doc["type"] === "vector";	20✔
126	}	20✔
127
128	private getValidationErrorForDocument(	3✔
129	definition: VectorFieldIndexDefinition,	46✔
130	document: Document	46✔
131	): VectorFieldValidationError \| undefined {	46✔
132	const fieldPath = definition.path.split(".");	46✔
133	let fieldRef: unknown = document;	46✔
134
135	const constructError = (	46✔
136	details: Partial<Pick<VectorFieldValidationError, "error" \| "actualNumDimensions" \| "actualQuantization">>	4✔
137	): VectorFieldValidationError => ({	4✔
138	path: definition.path,	4✔
139	expectedNumDimensions: definition.numDimensions,	4✔
140	expectedQuantization: definition.quantization,	4✔
141	actualNumDimensions: details.actualNumDimensions ?? "unknown",	4✔
142	actualQuantization: details.actualQuantization ?? "unknown",	4✔
143	error: details.error ?? "not-a-vector",	4!
144	});	4✔
145
146	for (const field of fieldPath) {	46✔
147	if (fieldRef && typeof fieldRef === "object" && field in fieldRef) {	71✔
148	fieldRef = (fieldRef as Record<string, unknown>)[field];	37✔
149	} else {	71✔
150	return undefined;	34✔
151	}	34✔
152	}	71✔
153
154	switch (definition.quantization) {	12✔
155	// Because quantization is not defined by the user
156	// we have to trust them in the format they use.
157	case "none":	46!
158	return undefined;	×
159	case "scalar":	46✔
160	case "binary":	46✔
161	if (fieldRef instanceof BSON.Binary) {	12✔
162	try {	2✔
163	const elements = fieldRef.toFloat32Array();	2✔
164	if (elements.length !== definition.numDimensions) {	2!
165	return constructError({	×
166	actualNumDimensions: elements.length,	×
167	actualQuantization: "binary",	×
168	error: "dimension-mismatch",	×
169	});	×
170	}	×
171
172	return undefined;	×
173	} catch {	2✔
174	// bits are also supported
175	try {	2✔
176	const bits = fieldRef.toBits();	2✔
177	if (bits.length !== definition.numDimensions) {	2!
178	return constructError({	×
179	actualNumDimensions: bits.length,	×
180	actualQuantization: "binary",	×
181	error: "dimension-mismatch",	×
182	});	×
183	}	×
184
185	return undefined;	2✔
186	} catch {	2!
187	return constructError({	×
188	actualQuantization: "binary",	×
189	error: "not-a-vector",	×
190	});	×
191	}	×
192	}	2✔
193	} else {	12✔
194	if (!Array.isArray(fieldRef)) {	10✔
195	return constructError({	2✔
196	error: "not-a-vector",	2✔
197	});	2✔
198	}	2✔
199
200	if (fieldRef.length !== definition.numDimensions) {	10✔
201	return constructError({	1✔
202	actualNumDimensions: fieldRef.length,	1✔
203	actualQuantization: "scalar",	1✔
204	error: "dimension-mismatch",	1✔
205	});	1✔
206	}	1✔
207
208	if (!fieldRef.every((e) => this.isANumber(e))) {	10✔
209	return constructError({	1✔
210	actualNumDimensions: fieldRef.length,	1✔
211	actualQuantization: "scalar",	1✔
212	error: "not-numeric",	1✔
213	});	1✔
214	}	1✔
215	}	10✔
216
217	break;	6✔
218	}	46✔
219
220	return undefined;	6✔
221	}	46✔
222
223	public async generateEmbeddings({	3✔
224	database,	4✔
225	collection,	4✔
226	path,	4✔
227	rawValues,	4✔
228	embeddingParameters,	4✔
229	inputType,	4✔
230	}: {	4✔
231	database: string;
232	collection: string;
233	path: string;
234	rawValues: string[];
235	embeddingParameters: SupportedEmbeddingParameters;
236	inputType: EmbeddingParameters["inputType"];
237	}): Promise<unknown[]> {	4✔
238	const provider = await this.atlasSearchEnabledProvider();	4✔
239	if (!provider) {	4!
NEW 240	throw new MongoDBError(	×
NEW 241	ErrorCodes.AtlasSearchNotSupported,	×
NEW 242	"Atlas Search is not supported in this cluster."	×
NEW 243	);	×
NEW 244	}	×
245
246	const embeddingsProvider = this.embeddingsProvider(this.config);	4✔
247
248	if (!embeddingsProvider) {	4!
NEW 249	throw new MongoDBError(ErrorCodes.NoEmbeddingsProviderConfigured, "No embeddings provider configured.");	×
NEW 250	}	×
251
252	if (this.config.disableEmbeddingsValidation) {	4✔
253	return await embeddingsProvider.embed(embeddingParameters.model, rawValues, {	1✔
254	inputType,	1✔
255	...embeddingParameters,	1✔
256	});	1✔
257	}	1✔
258
259	const embeddingInfoForCollection = await this.embeddingsForNamespace({ database, collection });	3✔
260	const embeddingInfoForPath = embeddingInfoForCollection.find((definition) => definition.path === path);	1✔
261	if (!embeddingInfoForPath) {	4!
NEW 262	throw new MongoDBError(	×
NEW 263	ErrorCodes.AtlasVectorSearchIndexNotFound,	×
NEW 264	`No Vector Search index found for path "${path}" in namespace "${database}.${collection}"`	×
NEW 265	);	×
NEW 266	}	✔
267
268	return await embeddingsProvider.embed(embeddingParameters.model, rawValues, {	1✔
269	inputType,	1✔
270	...embeddingParameters,	1✔
271	});	1✔
272	}	4✔
273
274	private isANumber(value: unknown): boolean {	3✔
275	if (typeof value === "number") {	49✔
276	return true;	16✔
277	}	16✔
278
279	if (	33✔
280	value instanceof BSON.Int32 \|\|	33✔
281	value instanceof BSON.Decimal128 \|\|	17✔
282	value instanceof BSON.Double \|\|	17✔
283	value instanceof BSON.Long	9✔
284	) {	49✔
285	return true;	32✔
286	}	32✔
287
288	return false;	1✔
289	}	49✔
290	}	3✔

mongodb-js / mongodb-mcp-server / 18717167438

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