18556878086

Committed 16 Oct 2025 09:33AM UTC coverage: 82.747% (+0.2%) from 82.527%

Build # 18556878086

Build Type

Pull #626

github

Committed by

web-flow

Commit Message

Merge 3d69362c7 into 87ce0cf2d

Pull Request Pull Request #626: chore: Add new session-level service for getting embeddings of a specific collection MCP-246

Run Details

1244 of 1635 branches covered (76.09%)

Branch coverage included in aggregate %.

210 of 227 new or added lines in 15 files covered. (92.51%)

14 existing lines in 1 file now uncovered.

5950 of 7059 relevant lines covered (84.29%)

70.62 hits per line

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

93.18

/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";

export type VectorFieldIndexDefinition = {
    type: "vector";
    path: string;
    numDimensions: number;
    quantization: "none" | "scalar" | "binary";
    similarity: "euclidean" | "cosine" | "dotProduct";
};

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()
    ) {
        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.assertAtlasSearchIsAvailable();
        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<VectorFieldIndexDefinition[]> {
        const provider = await this.assertAtlasSearchIsAvailable();
        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.filter((emb) => !this.documentPassesEmbeddingValidation(emb, document));
    }

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

        return null;
    }

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

    private documentPassesEmbeddingValidation(definition: VectorFieldIndexDefinition, document: Document): boolean {
        const fieldPath = definition.path.split(".");
        let fieldRef: unknown = document;

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

        switch (definition.quantization) {
            // Because quantization is not defined by the user
            // we have to trust them in the format they use.
            case "none":
                return true;
            case "scalar":
            case "binary":
                if (fieldRef instanceof BSON.Binary) {
                    try {
                        const elements = fieldRef.toFloat32Array();
                        return elements.length === definition.numDimensions;
                    } catch {
                        // bits are also supported
                        try {
                            const bits = fieldRef.toBits();
                            return bits.length === definition.numDimensions;
                        } catch {
                            return false;
                        }
                    }
                } else {
                    if (!Array.isArray(fieldRef)) {
                        return false;
                    }

                    if (fieldRef.length !== definition.numDimensions) {
                        return false;
                    }

                    if (!fieldRef.every((e) => this.isANumber(e))) {
                        return false;
                    }
                }

                break;
        }

        return true;
    }

    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";	2✔
3	import type { UserConfig } from "../config.js";
4	import type { ConnectionManager } from "../connectionManager.js";
5
6	export type VectorFieldIndexDefinition = {
7	type: "vector";
8	path: string;
9	numDimensions: number;
10	quantization: "none" \| "scalar" \| "binary";
11	similarity: "euclidean" \| "cosine" \| "dotProduct";
12	};
13
14	export type EmbeddingNamespace = `${string}.${string}`;
15	export class VectorSearchEmbeddingsManager {	2✔
16	constructor(	2✔
17	private readonly config: UserConfig,	120✔
18	private readonly connectionManager: ConnectionManager,	120✔
19	private readonly embeddings: Map<EmbeddingNamespace, VectorFieldIndexDefinition[]> = new Map()	120✔
20	) {	120✔
21	connectionManager.events.on("connection-close", () => {	120✔
22	this.embeddings.clear();	302✔
23	});	120✔
24	}	120✔
25
26	cleanupEmbeddingsForNamespace({ database, collection }: { database: string; collection: string }): void {	2✔
27	const embeddingDefKey: EmbeddingNamespace = `${database}.${collection}`;	2✔
28	this.embeddings.delete(embeddingDefKey);	2✔
29	}	2✔
30
31	async embeddingsForNamespace({	2✔
32	database,	20✔
33	collection,	20✔
34	}: {	20✔
35	database: string;
36	collection: string;
37	}): Promise<VectorFieldIndexDefinition[]> {	20✔
38	const provider = await this.assertAtlasSearchIsAvailable();	20✔
39	if (!provider) {	20!
NEW 40	return [];	×
NEW 41	}	×
42
43	// We only need the embeddings for validation now, so don't query them if
44	// validation is disabled.
45	if (this.config.disableEmbeddingsValidation) {	20!
NEW 46	return [];	×
NEW 47	}	×
48
49	const embeddingDefKey: EmbeddingNamespace = `${database}.${collection}`;	20✔
50	const definition = this.embeddings.get(embeddingDefKey);	20✔
51
52	if (!definition) {	20✔
53	const allSearchIndexes = await provider.getSearchIndexes(database, collection);	7✔
54	const vectorSearchIndexes = allSearchIndexes.filter((index) => index.type === "vectorSearch");	7✔
55	const vectorFields = vectorSearchIndexes	7✔
56	// eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
57	.flatMap<Document>((index) => (index.latestDefinition?.fields as Document) ?? [])	7!
58	.filter((field) => this.isVectorFieldIndexDefinition(field));	7✔
59
60	this.embeddings.set(embeddingDefKey, vectorFields);	7✔
61	return vectorFields;	7✔
62	}	7✔
63
64	return definition;	13✔
65	}	20✔
66
67	async findFieldsWithWrongEmbeddings(	2✔
68	{	20✔
69	database,	20✔
70	collection,	20✔
71	}: {	20✔
72	database: string;
73	collection: string;
74	},
75	document: Document	20✔
76	): Promise<VectorFieldIndexDefinition[]> {	20✔
77	const provider = await this.assertAtlasSearchIsAvailable();	20✔
78	if (!provider) {	20✔
79	return [];	3✔
80	}	3✔
81
82	// While we can do our best effort to ensure that the embedding validation is correct
83	// based on https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-quantization/
84	// it's a complex process so we will also give the user the ability to disable this validation
85	if (this.config.disableEmbeddingsValidation) {	20✔
86	return [];	3✔
87	}	3✔
88
89	const embeddings = await this.embeddingsForNamespace({ database, collection });	14✔
90	return embeddings.filter((emb) => !this.documentPassesEmbeddingValidation(emb, document));	14✔
91	}	20✔
92
93	private async assertAtlasSearchIsAvailable(): Promise<NodeDriverServiceProvider \| null> {	2✔
94	const connectionState = this.connectionManager.currentConnectionState;	40✔
95	if (connectionState.tag === "connected") {	40✔
96	if (await connectionState.isSearchSupported()) {	40✔
97	return connectionState.serviceProvider;	37✔
98	}	37✔
99	}	40✔
100
101	return null;	3✔
102	}	40✔
103
104	private isVectorFieldIndexDefinition(doc: Document): doc is VectorFieldIndexDefinition {	2✔
105	return doc["type"] === "vector";	17✔
106	}	17✔
107
108	private documentPassesEmbeddingValidation(definition: VectorFieldIndexDefinition, document: Document): boolean {	2✔
109	const fieldPath = definition.path.split(".");	46✔
110	let fieldRef: unknown = document;	46✔
111
112	for (const field of fieldPath) {	46✔
113	if (fieldRef && typeof fieldRef === "object" && field in fieldRef) {	71✔
114	fieldRef = (fieldRef as Record<string, unknown>)[field];	37✔
115	} else {	71✔
116	return true;	34✔
117	}	34✔
118	}	71✔
119
120	switch (definition.quantization) {	12✔
121	// Because quantization is not defined by the user
122	// we have to trust them in the format they use.
123	case "none":	46!
NEW 124	return true;	×
125	case "scalar":	46✔
126	case "binary":	46✔
127	if (fieldRef instanceof BSON.Binary) {	12✔
128	try {	2✔
129	const elements = fieldRef.toFloat32Array();	2✔
130	return elements.length === definition.numDimensions;	2✔
131	} catch {	2✔
132	// bits are also supported
133	try {	2✔
134	const bits = fieldRef.toBits();	2✔
135	return bits.length === definition.numDimensions;	2✔
136	} catch {	2!
NEW 137	return false;	×
NEW 138	}	×
139	}	2✔
140	} else {	12✔
141	if (!Array.isArray(fieldRef)) {	10✔
142	return false;	2✔
143	}	2✔
144
145	if (fieldRef.length !== definition.numDimensions) {	10✔
146	return false;	1✔
147	}	1✔
148
149	if (!fieldRef.every((e) => this.isANumber(e))) {	10✔
150	return false;	1✔
151	}	1✔
152	}	10✔
153
154	break;	6✔
155	}	46✔
156
157	return true;	6✔
158	}	46✔
159
160	private isANumber(value: unknown): boolean {	2✔
161	if (typeof value === "number") {	49✔
162	return true;	16✔
163	}	16✔
164
165	if (	33✔
166	value instanceof BSON.Int32 \|\|	33✔
167	value instanceof BSON.Decimal128 \|\|	17✔
168	value instanceof BSON.Double \|\|	17✔
169	value instanceof BSON.Long	9✔
170	) {	49✔
171	return true;	32✔
172	}	32✔
173
174	return false;	1✔
175	}	49✔
176	}	2✔

mongodb-js / mongodb-mcp-server / 18556878086

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