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

@@ -32,6 +32,7 @@ export interface ConnectionState {
     connectedAtlasCluster?: AtlasClusterConnectionInfo;
 }
 
+const MCP_TEST_DATABASE = "#mongodb-mcp";
 export class ConnectionStateConnected implements ConnectionState {
     public tag = "connected" as const;
 
@@ -46,11 +47,11 @@ export class ConnectionStateConnected implements ConnectionState {
     public 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);
+                // with a cursor otherwise will throw an Error.
+                // the Search Index Management Service might not be ready yet, but
+                // we assume that the agent can retry in that situation.
+                await this.serviceProvider.getSearchIndexes(MCP_TEST_DATABASE, "test");
                 this._isSearchSupported = true;
             } catch {
                 this._isSearchSupported = false;

src/common/errors.ts

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

src/common/search/vectorSearchEmbeddingsManager.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 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;
+    }
+}

src/common/session.ts

@@ -16,6 +16,7 @@ import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-d
 import { ErrorCodes, MongoDBError } from "./errors.js";
 import type { ExportsManager } from "./exportsManager.js";
 import type { Keychain } from "./keychain.js";
+import type { VectorSearchEmbeddingsManager } from "./search/vectorSearchEmbeddingsManager.js";
 
 export interface SessionOptions {
     apiBaseUrl: string;
@@ -25,6 +26,7 @@ export interface SessionOptions {
     exportsManager: ExportsManager;
     connectionManager: ConnectionManager;
     keychain: Keychain;
+    vectorSearchEmbeddingsManager: VectorSearchEmbeddingsManager;
 }
 
 export type SessionEvents = {
@@ -40,6 +42,7 @@ export class Session extends EventEmitter<SessionEvents> {
     readonly connectionManager: ConnectionManager;
     readonly apiClient: ApiClient;
     readonly keychain: Keychain;
+    readonly vectorSearchEmbeddingsManager: VectorSearchEmbeddingsManager;
 
     mcpClient?: {
         name?: string;
@@ -57,6 +60,7 @@ export class Session extends EventEmitter<SessionEvents> {
         connectionManager,
         exportsManager,
         keychain,
+        vectorSearchEmbeddingsManager,
     }: SessionOptions) {
         super();
 
@@ -73,6 +77,7 @@ export class Session extends EventEmitter<SessionEvents> {
         this.apiClient = new ApiClient({ baseUrl: apiBaseUrl, credentials }, logger);
         this.exportsManager = exportsManager;
         this.connectionManager = connectionManager;
+        this.vectorSearchEmbeddingsManager = vectorSearchEmbeddingsManager;
         this.connectionManager.events.on("connection-success", () => this.emit("connect"));
         this.connectionManager.events.on("connection-time-out", (error) => this.emit("connection-error", error));
         this.connectionManager.events.on("connection-close", () => this.emit("disconnect"));
@@ -141,13 +146,25 @@ export class Session extends EventEmitter<SessionEvents> {
         return this.connectionManager.currentConnectionState.tag === "connected";
     }
 
-    isSearchSupported(): Promise<boolean> {
+    async isSearchSupported(): Promise<boolean> {
         const state = this.connectionManager.currentConnectionState;
         if (state.tag === "connected") {
-            return state.isSearchSupported();
+            return await state.isSearchSupported();
         }
 
-        return Promise.resolve(false);
+        return false;
+    }
+
+    async assertSearchSupported(): Promise<void> {
+        const availability = await this.isSearchSupported();
+        if (!availability) {
+            throw new MongoDBError(
+                ErrorCodes.AtlasSearchNotSupported,
+                "Atlas Search is not supported in the current cluster."
+            );
+        }
+
+        return;
     }
 
     get serviceProvider(): NodeDriverServiceProvider {

src/tools/mongodb/create/createIndex.ts

@@ -1,7 +1,6 @@
 import { z } from "zod";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
-import type { ToolCategory } from "../../tool.js";
 import { type ToolArgs, type OperationType, FeatureFlags } from "../../tool.js";
 import type { IndexDirection } from "mongodb";
 
@@ -113,25 +112,7 @@ export class CreateIndexTool extends MongoDBToolBase {
                 break;
             case "vectorSearch":
                 {
-                    const isVectorSearchSupported = await this.session.isSearchSupported();
-                    if (!isVectorSearchSupported) {
-                        // TODO: remove hacky casts once we merge the local dev tools
-                        const isLocalAtlasAvailable =
-                            (this.server?.tools.filter((t) => t.category === ("atlas-local" as unknown as ToolCategory))
-                                .length ?? 0) > 0;
-
-                        const CTA = isLocalAtlasAvailable ? "`atlas-local` tools" : "Atlas CLI";
-                        return {
-                            content: [
-                                {
-                                    text: `The connected MongoDB deployment does not support vector search indexes. Either connect to a MongoDB Atlas cluster or use the ${CTA} to create and manage a local Atlas deployment.`,
-                                    type: "text",
-                                },
-                            ],
-                            isError: true,
-                        };
-                    }
-
+                    await this.ensureSearchIsSupported();
                     indexes = await provider.createSearchIndexes(database, collection, [
                         {
                             name,
@@ -144,6 +125,8 @@ export class CreateIndexTool extends MongoDBToolBase {
 
                     responseClarification =
                         " Since this is a vector search index, it may take a while for the index to build. Use the `list-indexes` tool to check the index status.";
+                    // clean up the embeddings cache so it considers the new index
+                    this.session.vectorSearchEmbeddingsManager.cleanupEmbeddingsForNamespace({ database, collection });
                 }
 
                 break;