Skip to content

Commit adbf3a0

Browse files
himanshusinghshimanshusinghs
authored
feat: adds support for dropping vector search indexes in drop-index tool MCP-239 (#642)
1 parent dbaa468 commit adbf3a0

File tree

11 files changed

+897
-354
lines changed

11 files changed

+897
-354
lines changed

‎src/server.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ import {
1818
UnsubscribeRequestSchema,
1919
} from "@modelcontextprotocol/sdk/types.js";
2020
import assert from "assert";
21-
import type { ToolBase, ToolConstructorParams } from "./tools/tool.js";
21+
import type { ToolBase, ToolCategory, ToolConstructorParams } from "./tools/tool.js";
2222
import { validateConnectionString } from "./helpers/connectionOptions.js";
2323
import { packageInfo } from "./common/packageInfo.js";
2424
import { type ConnectionErrorHandler } from "./common/connectionErrorHandler.js";
@@ -174,6 +174,10 @@ export class Server {
174174
this.mcpServer.sendResourceListChanged();
175175
}
176176

177+
public isToolCategoryAvailable(name: ToolCategory): boolean {
178+
return !!this.tools.filter((t) => t.category === name).length;
179+
}
180+
177181
public sendResourceUpdated(uri: string): void {
178182
this.session.logger.info({
179183
id: LogId.resourceUpdateFailure,

‎src/tools/mongodb/delete/dropIndex.ts‎

Lines changed: 62 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,43 @@
11
import z from "zod";
22
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
3+
import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
34
import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
4-
import { type ToolArgs, type OperationType, formatUntrustedData } from "../../tool.js";
5+
import { type ToolArgs, type OperationType, formatUntrustedData, FeatureFlags } from "../../tool.js";
6+
import { ListSearchIndexesTool } from "../search/listSearchIndexes.js";
57

68
export class DropIndexTool extends MongoDBToolBase {
79
public name = "drop-index";
810
protected description = "Drop an index for the provided database and collection.";
911
protected argsShape = {
1012
...DbOperationArgs,
1113
indexName: z.string().nonempty().describe("The name of the index to be dropped."),
14+
type: this.isFeatureFlagEnabled(FeatureFlags.VectorSearch)
15+
? z
16+
.enum(["classic", "search"])
17+
.describe(
18+
"The type of index to be deleted. Use 'classic' for standard indexes and 'search' for atlas search and vector search indexes."
19+
)
20+
: z
21+
.literal("classic")
22+
.default("classic")
23+
.describe("The type of index to be deleted. Is always set to 'classic'."),
1224
};
1325
public operationType: OperationType = "delete";
1426

15-
protected async execute({
16-
database,
17-
collection,
18-
indexName,
19-
}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
27+
protected async execute(toolArgs: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
2028
const provider = await this.ensureConnected();
29+
switch (toolArgs.type) {
30+
case "classic":
31+
return this.dropClassicIndex(provider, toolArgs);
32+
case "search":
33+
return this.dropSearchIndex(provider, toolArgs);
34+
}
35+
}
36+
37+
private async dropClassicIndex(
38+
provider: NodeDriverServiceProvider,
39+
{ database, collection, indexName }: ToolArgs<typeof this.argsShape>
40+
): Promise<CallToolResult> {
2141
const result = await provider.runCommand(database, {
2242
dropIndexes: collection,
2343
index: indexName,
@@ -35,9 +55,43 @@ export class DropIndexTool extends MongoDBToolBase {
3555
};
3656
}
3757

38-
protected getConfirmationMessage({ database, collection, indexName }: ToolArgs<typeof this.argsShape>): string {
58+
private async dropSearchIndex(
59+
provider: NodeDriverServiceProvider,
60+
{ database, collection, indexName }: ToolArgs<typeof this.argsShape>
61+
): Promise<CallToolResult> {
62+
await this.ensureSearchIsSupported();
63+
const searchIndexes = await ListSearchIndexesTool.getSearchIndexes(provider, database, collection);
64+
const indexDoesNotExist = !searchIndexes.find((index) => index.name === indexName);
65+
if (indexDoesNotExist) {
66+
return {
67+
content: formatUntrustedData(
68+
"Index does not exist in the provided namespace.",
69+
JSON.stringify({ indexName, namespace: `${database}.${collection}` })
70+
),
71+
isError: true,
72+
};
73+
}
74+
75+
await provider.dropSearchIndex(database, collection, indexName);
76+
return {
77+
content: formatUntrustedData(
78+
"Successfully dropped the index from the provided namespace.",
79+
JSON.stringify({
80+
indexName,
81+
namespace: `${database}.${collection}`,
82+
})
83+
),
84+
};
85+
}
86+
87+
protected getConfirmationMessage({
88+
database,
89+
collection,
90+
indexName,
91+
type,
92+
}: ToolArgs<typeof this.argsShape>): string {
3993
return (
40-
`You are about to drop the \`${indexName}\` index from the \`${database}.${collection}\` namespace:\n\n` +
94+
`You are about to drop the ${type === "search" ? "search index" : "index"} named \`${indexName}\` from the \`${database}.${collection}\` namespace:\n\n` +
4195
"This operation will permanently remove the index and might affect the performance of queries relying on this index.\n\n" +
4296
"**Do you confirm the execution of the action?**"
4397
);

‎src/tools/mongodb/mongodbTool.ts‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ export abstract class MongoDBToolBase extends ToolBase {
8787
isError: true,
8888
};
8989
case ErrorCodes.AtlasSearchNotSupported: {
90-
const CTA = this.isToolCategoryAvailable("atlas-local" as unknown as ToolCategory)
90+
const CTA = this.server?.isToolCategoryAvailable("atlas-local" as unknown as ToolCategory)
9191
? "`atlas-local` tools"
9292
: "Atlas CLI";
9393
return {
@@ -123,8 +123,4 @@ export abstract class MongoDBToolBase extends ToolBase {
123123

124124
return metadata;
125125
}
126-
127-
protected isToolCategoryAvailable(name: ToolCategory): boolean {
128-
return (this.server?.tools.filter((t) => t.category === name).length ?? 0) > 0;
129-
}
130126
}

‎src/tools/mongodb/search/listSearchIndexes.ts‎

Lines changed: 19 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
2+
import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
23
import type { ToolArgs, OperationType } from "../../tool.js";
34
import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
45
import { formatUntrustedData } from "../../tool.js";
56
import { EJSON } from "bson";
67

7-
export type SearchIndexStatus = {
8+
export type SearchIndexWithStatus = {
89
name: string;
910
type: "search" | "vectorSearch";
1011
status: string;
@@ -21,15 +22,13 @@ export class ListSearchIndexesTool extends MongoDBToolBase {
2122
protected async execute({ database, collection }: ToolArgs<typeof DbOperationArgs>): Promise<CallToolResult> {
2223
const provider = await this.ensureConnected();
2324
await this.ensureSearchIsSupported();
25+
const searchIndexes = await ListSearchIndexesTool.getSearchIndexes(provider, database, collection);
2426

25-
const indexes = await provider.getSearchIndexes(database, collection);
26-
const trimmedIndexDefinitions = this.pickRelevantInformation(indexes);
27-
28-
if (trimmedIndexDefinitions.length > 0) {
27+
if (searchIndexes.length > 0) {
2928
return {
3029
content: formatUntrustedData(
31-
`Found ${trimmedIndexDefinitions.length} search and vector search indexes in ${database}.${collection}`,
32-
...trimmedIndexDefinitions.map((index) => EJSON.stringify(index))
30+
`Found ${searchIndexes.length} search and vector search indexes in ${database}.${collection}`,
31+
...searchIndexes.map((index) => EJSON.stringify(index))
3332
),
3433
};
3534
} else {
@@ -47,14 +46,19 @@ export class ListSearchIndexesTool extends MongoDBToolBase {
4746
return process.env.VITEST === "true";
4847
}
4948

50-
/**
51-
* Atlas Search index status contains a lot of information that is not relevant for the agent at this stage.
52-
* Like for example, the status on each of the dedicated nodes. We only care about the main status, if it's
53-
* queryable and the index name. We are also picking the index definition as it can be used by the agent to
54-
* understand which fields are available for searching.
55-
**/
56-
protected pickRelevantInformation(indexes: Record<string, unknown>[]): SearchIndexStatus[] {
57-
return indexes.map((index) => ({
49+
static async getSearchIndexes(
50+
provider: NodeDriverServiceProvider,
51+
database: string,
52+
collection: string
53+
): Promise<SearchIndexWithStatus[]> {
54+
const searchIndexes = await provider.getSearchIndexes(database, collection);
55+
/**
56+
* Atlas Search index status contains a lot of information that is not relevant for the agent at this stage.
57+
* Like for example, the status on each of the dedicated nodes. We only care about the main status, if it's
58+
* queryable and the index name. We are also picking the index definition as it can be used by the agent to
59+
* understand which fields are available for searching.
60+
**/
61+
return searchIndexes.map<SearchIndexWithStatus>((index) => ({
5862
name: (index["name"] ?? "default") as string,
5963
type: (index["type"] ?? "UNKNOWN") as "search" | "vectorSearch",
6064
status: (index["status"] ?? "UNKNOWN") as string,

‎tests/accuracy/dropIndex.test.ts‎

Lines changed: 116 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -1,79 +1,134 @@
11
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
22
import { describeAccuracyTests } from "./sdk/describeAccuracyTests.js";
33
import { Matcher } from "./sdk/matcher.js";
4+
import { formatUntrustedData } from "../../src/tools/tool.js";
45

56
// We don't want to delete actual indexes
67
const mockedTools = {
78
"drop-index": ({ indexName, database, collection }: Record<string, unknown>): CallToolResult => {
89
return {
9-
content: [
10-
{
11-
text: `Successfully dropped the index with name "${String(indexName)}" from the provided namespace "${String(database)}.${String(collection)}".`,
12-
type: "text",
13-
},
14-
],
10+
content: formatUntrustedData(
11+
"Successfully dropped the index from the provided namespace.",
12+
JSON.stringify({
13+
indexName,
14+
namespace: `${database as string}.${collection as string}`,
15+
})
16+
),
1517
};
1618
},
1719
} as const;
1820

19-
describeAccuracyTests([
20-
{
21-
prompt: "Delete the index called year_1 from mflix.movies namespace",
22-
expectedToolCalls: [
23-
{
24-
toolName: "drop-index",
25-
parameters: {
26-
database: "mflix",
27-
collection: "movies",
28-
indexName: "year_1",
21+
describeAccuracyTests(
22+
[
23+
{
24+
prompt: "Delete the index called year_1 from mflix.movies namespace",
25+
expectedToolCalls: [
26+
{
27+
toolName: "drop-index",
28+
parameters: {
29+
database: "mflix",
30+
collection: "movies",
31+
indexName: "year_1",
32+
type: "classic",
33+
},
2934
},
30-
},
31-
],
32-
mockedTools,
33-
},
34-
{
35-
prompt: "First create a text index on field 'title' in 'mflix.movies' namespace and then drop all the indexes from 'mflix.movies' namespace",
36-
expectedToolCalls: [
37-
{
38-
toolName: "create-index",
39-
parameters: {
40-
database: "mflix",
41-
collection: "movies",
42-
name: Matcher.anyOf(Matcher.undefined, Matcher.string()),
43-
definition: [
44-
{
45-
keys: {
46-
title: "text",
35+
],
36+
mockedTools,
37+
},
38+
{
39+
prompt: "First create a text index on field 'title' in 'mflix.movies' namespace and then drop all the indexes from 'mflix.movies' namespace",
40+
expectedToolCalls: [
41+
{
42+
toolName: "create-index",
43+
parameters: {
44+
database: "mflix",
45+
collection: "movies",
46+
name: Matcher.anyOf(Matcher.undefined, Matcher.string()),
47+
definition: [
48+
{
49+
keys: {
50+
title: "text",
51+
},
52+
type: "classic",
4753
},
48-
type: "classic",
49-
},
50-
],
54+
],
55+
},
5156
},
52-
},
53-
{
54-
toolName: "collection-indexes",
55-
parameters: {
56-
database: "mflix",
57-
collection: "movies",
57+
{
58+
toolName: "collection-indexes",
59+
parameters: {
60+
database: "mflix",
61+
collection: "movies",
62+
},
5863
},
59-
},
60-
{
61-
toolName: "drop-index",
62-
parameters: {
63-
database: "mflix",
64-
collection: "movies",
65-
indexName: Matcher.string(),
64+
{
65+
toolName: "drop-index",
66+
parameters: {
67+
database: "mflix",
68+
collection: "movies",
69+
indexName: Matcher.string(),
70+
type: "classic",
71+
},
6672
},
67-
},
68-
{
69-
toolName: "drop-index",
70-
parameters: {
71-
database: "mflix",
72-
collection: "movies",
73-
indexName: Matcher.string(),
73+
{
74+
toolName: "drop-index",
75+
parameters: {
76+
database: "mflix",
77+
collection: "movies",
78+
indexName: Matcher.string(),
79+
type: "classic",
80+
},
7481
},
75-
},
76-
],
77-
mockedTools,
78-
},
79-
]);
82+
],
83+
mockedTools,
84+
},
85+
{
86+
prompt: "Create a vector search index on 'mflix.movies' namespace on the 'plotSummary' field. The index should use 1024 dimensions. Confirm that its created and then drop the index.",
87+
expectedToolCalls: [
88+
{
89+
toolName: "create-index",
90+
parameters: {
91+
database: "mflix",
92+
collection: "movies",
93+
name: Matcher.anyOf(Matcher.undefined, Matcher.string()),
94+
definition: [
95+
{
96+
type: "vectorSearch",
97+
fields: [
98+
{
99+
type: "vector",
100+
path: "plotSummary",
101+
numDimensions: 1024,
102+
},
103+
],
104+
},
105+
],
106+
},
107+
},
108+
{
109+
toolName: "collection-indexes",
110+
parameters: {
111+
database: "mflix",
112+
collection: "movies",
113+
},
114+
},
115+
{
116+
toolName: "drop-index",
117+
parameters: {
118+
database: "mflix",
119+
collection: "movies",
120+
indexName: Matcher.string(),
121+
type: "search",
122+
},
123+
},
124+
],
125+
mockedTools,
126+
},
127+
],
128+
{
129+
userConfig: {
130+
voyageApiKey: "voyage-api-key",
131+
},
132+
clusterConfig: { search: true },
133+
}
134+
);

0 commit comments

Comments
 (0)