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

src/common/config.ts

@@ -58,6 +58,7 @@ const OPTIONS = {
     boolean: [
         "apiDeprecationErrors",
         "apiStrict",
+        "disableEmbeddingsValidation",
         "help",
         "indexCheck",
         "ipv6",
@@ -183,6 +184,7 @@ export interface UserConfig extends CliOptions {
     maxBytesPerQuery: number;
     atlasTemporaryDatabaseUserLifetimeMs: number;
     voyageApiKey: string;
+    disableEmbeddingsValidation: boolean;
     vectorSearchDimensions: number;
     vectorSearchSimilarityFunction: "cosine" | "euclidean" | "dotProduct";
 }
@@ -216,6 +218,7 @@ export const defaultUserConfig: UserConfig = {
     maxBytesPerQuery: 16 * 1024 * 1024, // By default, we only return ~16 mb of data per query / aggregation
     atlasTemporaryDatabaseUserLifetimeMs: 4 * 60 * 60 * 1000, // 4 hours
     voyageApiKey: "",
+    disableEmbeddingsValidation: false,
     vectorSearchDimensions: 1024,
     vectorSearchSimilarityFunction: "euclidean",
 };

src/common/connectionManager.ts

@@ -25,39 +25,106 @@ export interface ConnectionSettings {
 type ConnectionTag = "connected" | "connecting" | "disconnected" | "errored";
 type OIDCConnectionAuthType = "oidc-auth-flow" | "oidc-device-flow";
 export type ConnectionStringAuthType = "scram" | "ldap" | "kerberos" | OIDCConnectionAuthType | "x.509";
+export type SearchAvailability = false | "not-available-yet" | "available";
 
 export interface ConnectionState {
     tag: ConnectionTag;
     connectionStringAuthType?: ConnectionStringAuthType;
     connectedAtlasCluster?: AtlasClusterConnectionInfo;
 }
 
+const MCP_TEST_DATABASE = "#mongodb-mcp";
+const SEARCH_AVAILABILITY_CHECK_TIMEOUT_MS = 500;
 export class ConnectionStateConnected implements ConnectionState {
     public tag = "connected" as const;
 
     constructor(
         public serviceProvider: NodeDriverServiceProvider,
         public connectionStringAuthType?: ConnectionStringAuthType,
         public connectedAtlasCluster?: AtlasClusterConnectionInfo
-    ) {}
+    ) {
+        this.#isSearchAvailable = false;
+    }
+
+    #isSearchSupported?: boolean;
+    #isSearchAvailable: boolean;
 
-    private _isSearchSupported?: boolean;
+    public async getSearchAvailability(): Promise<SearchAvailability> {
+        if ((await this.isSearchSupported()) === true) {
+            if ((await this.isSearchAvailable()) === true) {
+                return "available";
+            }
+
+            return "not-available-yet";
+        }
 
-    public async isSearchSupported(): Promise<boolean> {
-        if (this._isSearchSupported === undefined) {
+        return false;
+    }
+
+    private async isSearchSupported(): Promise<boolean> {
+        if (this.#isSearchSupported === undefined) {
             try {
-                const dummyDatabase = "test";
-                const dummyCollection = "test";
                 // If a cluster supports search indexes, the call below will succeed
                 // with a cursor otherwise will throw an Error
-                await this.serviceProvider.getSearchIndexes(dummyDatabase, dummyCollection);
-                this._isSearchSupported = true;
+                await this.serviceProvider.getSearchIndexes(MCP_TEST_DATABASE, "test");
+                this.#isSearchSupported = true;
             } catch {
-                this._isSearchSupported = false;
+                this.#isSearchSupported = false;
+            }
+        }
+
+        return this.#isSearchSupported;
+    }
+
+    private async isSearchAvailable(): Promise<boolean> {
+        if (this.#isSearchAvailable === true) {
+            return true;
+        }
+
+        const timeoutPromise = new Promise<boolean>((_resolve, reject) =>
+            setTimeout(
+                () =>
+                    reject(
+                        new MongoDBError(
+                            ErrorCodes.AtlasSearchNotAvailable,
+                            "Atlas Search is supported in your environment but is not available yet. Retry again later."
+                        )
+                    ),
+                SEARCH_AVAILABILITY_CHECK_TIMEOUT_MS
+            )
+        );
+
+        const checkPromise = new Promise<boolean>((resolve) => {
+            void this.doCheckSearchIndexIsAvailable(resolve);
+        });
+
+        return await Promise.race([checkPromise, timeoutPromise]);
+    }
+
+    private async doCheckSearchIndexIsAvailable(resolve: (result: boolean) => void): Promise<void> {
+        for (let i = 0; i < 100; i++) {
+            try {
+                try {
+                    await this.serviceProvider.insertOne(MCP_TEST_DATABASE, "test", { search: "search is available" });
+                } catch (err) {
+                    // if inserting one document fails, it means we are in readOnly mode. We can't verify reliably if
+                    // Search is available, so assume it is.
+                    void err;
+                    resolve(true);
+                    return;
+                }
+                await this.serviceProvider.createSearchIndexes(MCP_TEST_DATABASE, "test", [
+                    { definition: { mappings: { dynamic: true } } },
+                ]);
+                await this.serviceProvider.dropDatabase(MCP_TEST_DATABASE);
+                resolve(true);
+                return;
+            } catch (err) {
+                void err;
             }
         }
 
-        return this._isSearchSupported;
+        resolve(false);
     }
 }
 

src/common/errors.ts

@@ -3,6 +3,8 @@ export enum ErrorCodes {
     MisconfiguredConnectionString = 1_000_001,
     ForbiddenCollscan = 1_000_002,
     ForbiddenWriteOperation = 1_000_003,
+    AtlasSearchNotSupported = 1_000_004,
+    AtlasSearchNotAvailable = 1_000_005,
 }
 
 export class MongoDBError<ErrorCode extends ErrorCodes = ErrorCodes> extends Error {

src/common/search/vectorSearchEmbeddings.ts

@@ -0,0 +1,176 @@
+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 VectorSearchEmbeddings {
+    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.getSearchAvailability()) === "available") {
+                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 use
+            // 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;
+    }
+}