chore: check that a vector search index exists with indexCheck

Author: kmruiz

src/common/search/vectorSearchEmbeddingsManager.ts

@@ -48,6 +48,26 @@ export class VectorSearchEmbeddingsManager {
         this.embeddings.delete(embeddingDefKey);
     }
 
+    async indexExists({
+        database,
+        collection,
+        indexName,
+    }: {
+        database: string;
+        collection: string;
+        indexName: string;
+    }): Promise<boolean> {
+        const provider = await this.atlasSearchEnabledProvider();
+        if (!provider) {
+            return false;
+        }
+
+        const searchIndexesWithName = await provider.getSearchIndexes(database, collection, indexName);
+        console.log(">>>>>", searchIndexesWithName);
+
+        return searchIndexesWithName.length >= 1;
+    }
+
     async embeddingsForNamespace({
         database,
         collection,

src/tools/mongodb/read/aggregate.ts

@@ -90,11 +90,23 @@ export class AggregateTool extends MongoDBToolBase {
 
             // Check if aggregate operation uses an index if enabled
             if (this.config.indexCheck) {
-                await checkIndexUsage(provider, database, collection, "aggregate", async () => {
-                    return provider
-                        .aggregate(database, collection, pipeline, {}, { writeConcern: undefined })
-                        .explain("queryPlanner");
-                });
+                const usesVectorSearchIndex = await this.isVectorSearchIndexUsed({ database, collection, pipeline });
+                switch (usesVectorSearchIndex) {
+                    case "no-vector-search-query":
+                        await checkIndexUsage(provider, database, collection, "aggregate", async () => {
+                            return provider
+                                .aggregate(database, collection, pipeline, {}, { writeConcern: undefined })
+                                .explain("queryPlanner");
+                        });
+                        break;
+                    case false:
+                        throw new MongoDBError(
+                            ErrorCodes.AtlasVectorSearchIndexNotFound,
+                            "Could not find provided vector search index."
+                        );
+                    case true:
+                    // nothing to do, everything is correct so ready to run the query
+                }
             }
 
             pipeline = await this.replaceRawValuesWithEmbeddingsIfNecessary({
@@ -269,6 +281,41 @@ export class AggregateTool extends MongoDBToolBase {
         return pipeline;
     }
 
+    private async isVectorSearchIndexUsed({
+        database,
+        collection,
+        pipeline,
+    }: {
+        database: string;
+        collection: string;
+        pipeline: Document[];
+    }): Promise<boolean | "no-vector-search-query"> {
+        // check if the pipeline contains a $vectorSearch stage
+        let usesVectorSearch = false;
+        let indexName: string = "default";
+
+        for (const stage of pipeline) {
+            if ("$vectorSearch" in stage) {
+                const { $vectorSearch: vectorSearchStage } = stage as z.infer<typeof VectorSearchStage>;
+                usesVectorSearch = true;
+                indexName = vectorSearchStage.index;
+                break;
+            }
+        }
+
+        if (!usesVectorSearch) {
+            return "no-vector-search-query";
+        }
+
+        const indexExists = await this.session.vectorSearchEmbeddingsManager.indexExists({
+            database,
+            collection,
+            indexName,
+        });
+
+        return indexExists;
+    }
+
     private generateMessage({
         aggResultsCount,
         documents,

tests/integration/tools/mongodb/read/aggregate.test.ts

@@ -394,6 +394,45 @@ describeWithMongoDB(
             await integration.mongoClient().db(integration.randomDbName()).collection("databases").drop();
         });
 
+        it("should throw an exception when using an index that does not exist", async () => {
+            await waitUntilSearchIsReady(integration.mongoClient());
+
+            const collection = integration.mongoClient().db(integration.randomDbName()).collection("databases");
+
+            await collection.insertOne({ name: "mongodb", description_embedding: [1, 2, 3, 4] });
+            await integration.connectMcpClient();
+            const response = await integration.mcpClient().callTool({
+                name: "aggregate",
+                arguments: {
+                    database: integration.randomDbName(),
+                    collection: "databases",
+                    pipeline: [
+                        {
+                            $vectorSearch: {
+                                index: "non_existing",
+                                path: "description_embedding",
+                                queryVector: "example",
+                                numCandidates: 10,
+                                limit: 10,
+                                embeddingParameters: {
+                                    model: "voyage-3-large",
+                                    outputDimension: 256,
+                                },
+                            },
+                        },
+                        {
+                            $project: {
+                                description_embedding: 0,
+                            },
+                        },
+                    ],
+                },
+            });
+
+            const responseContent = getResponseContent(response);
+            expect(responseContent).toContain("Error running aggregate: Could not find provided vector search index.");
+        });
+
         for (const [dataType, embedding] of Object.entries(DOCUMENT_EMBEDDINGS)) {
             for (const similarity of ["euclidean", "cosine", "dotProduct"]) {
                 describe.skipIf(!process.env.TEST_MDB_MCP_VOYAGE_API_KEY)(
@@ -406,6 +445,7 @@ describeWithMongoDB(
                                 .mongoClient()
                                 .db(integration.randomDbName())
                                 .collection("databases");
+
                             await collection.insertOne({ name: "mongodb", description_embedding: embedding });
 
                             await createVectorSearchIndexAndWait(
@@ -674,6 +714,7 @@ describeWithMongoDB(
             voyageApiKey: process.env.TEST_MDB_MCP_VOYAGE_API_KEY ?? "",
             maxDocumentsPerQuery: -1,
             maxBytesPerQuery: -1,
+            indexCheck: true,
         }),
         downloadOptions: { search: true },
     }