diff --git a/src/common/search/vectorSearchEmbeddingsManager.ts b/src/common/search/vectorSearchEmbeddingsManager.ts index a86d7026..fc8c53be 100644 --- a/src/common/search/vectorSearchEmbeddingsManager.ts +++ b/src/common/search/vectorSearchEmbeddingsManager.ts @@ -48,6 +48,25 @@ export class VectorSearchEmbeddingsManager { this.embeddings.delete(embeddingDefKey); } + async indexExists({ + database, + collection, + indexName, + }: { + database: string; + collection: string; + indexName: string; + }): Promise { + const provider = await this.atlasSearchEnabledProvider(); + if (!provider) { + return false; + } + + const searchIndexesWithName = await provider.getSearchIndexes(database, collection, indexName); + + return searchIndexesWithName.length >= 1; + } + async embeddingsForNamespace({ database, collection, diff --git a/src/tools/mongodb/read/aggregate.ts b/src/tools/mongodb/read/aggregate.ts index a5ff1238..85727556 100644 --- a/src/tools/mongodb/read/aggregate.ts +++ b/src/tools/mongodb/read/aggregate.ts @@ -90,11 +90,27 @@ 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, indexName] = await this.isVectorSearchIndexUsed({ + database, + collection, + pipeline, }); + switch (usesVectorSearchIndex) { + case "not-vector-search-query": + await checkIndexUsage(provider, database, collection, "aggregate", async () => { + return provider + .aggregate(database, collection, pipeline, {}, { writeConcern: undefined }) + .explain("queryPlanner"); + }); + break; + case "non-existent-index": + throw new MongoDBError( + ErrorCodes.AtlasVectorSearchIndexNotFound, + `Could not find an index with name "${indexName}" in namespace "${database}.${collection}".` + ); + case "valid-index": + // nothing to do, everything is correct so ready to run the query + } } pipeline = await this.replaceRawValuesWithEmbeddingsIfNecessary({ @@ -269,6 +285,41 @@ export class AggregateTool extends MongoDBToolBase { return pipeline; } + private async isVectorSearchIndexUsed({ + database, + collection, + pipeline, + }: { + database: string; + collection: string; + pipeline: Document[]; + }): Promise<["valid-index" | "non-existent-index" | "not-vector-search-query", string?]> { + // 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; + usesVectorSearch = true; + indexName = vectorSearchStage.index; + break; + } + } + + if (!usesVectorSearch) { + return ["not-vector-search-query"]; + } + + const indexExists = await this.session.vectorSearchEmbeddingsManager.indexExists({ + database, + collection, + indexName, + }); + + return [indexExists ? "valid-index" : "non-existent-index", indexName]; + } + private generateMessage({ aggResultsCount, documents, diff --git a/tests/integration/tools/mongodb/read/aggregate.test.ts b/tests/integration/tools/mongodb/read/aggregate.test.ts index d71ab4d9..d3e89d4f 100644 --- a/tests/integration/tools/mongodb/read/aggregate.test.ts +++ b/tests/integration/tools/mongodb/read/aggregate.test.ts @@ -394,6 +394,47 @@ 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 an index with name "non_existing" in namespace "${integration.randomDbName()}.databases".` + ); + }); + 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 +447,7 @@ describeWithMongoDB( .mongoClient() .db(integration.randomDbName()) .collection("databases"); + await collection.insertOne({ name: "mongodb", description_embedding: embedding }); await createVectorSearchIndexAndWait( @@ -674,6 +716,7 @@ describeWithMongoDB( voyageApiKey: process.env.TEST_MDB_MCP_VOYAGE_API_KEY ?? "", maxDocumentsPerQuery: -1, maxBytesPerQuery: -1, + indexCheck: true, }), downloadOptions: { search: true }, }