chore: merge list-search-indexes into collection-indexes (#654)

Author: nirinchev

src/tools/mongodb/delete/dropIndex.ts

@@ -3,7 +3,6 @@ import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
 import { type ToolArgs, type OperationType, formatUntrustedData, FeatureFlags } from "../../tool.js";
-import { ListSearchIndexesTool } from "../search/listSearchIndexes.js";
 
 export class DropIndexTool extends MongoDBToolBase {
     public name = "drop-index";
@@ -60,9 +59,8 @@ export class DropIndexTool extends MongoDBToolBase {
         { database, collection, indexName }: ToolArgs<typeof this.argsShape>
     ): Promise<CallToolResult> {
         await this.ensureSearchIsSupported();
-        const searchIndexes = await ListSearchIndexesTool.getSearchIndexes(provider, database, collection);
-        const indexDoesNotExist = !searchIndexes.find((index) => index.name === indexName);
-        if (indexDoesNotExist) {
+        const indexes = await provider.getSearchIndexes(database, collection, indexName);
+        if (indexes.length === 0) {
             return {
                 content: formatUntrustedData(
                     "Index does not exist in the provided namespace.",

src/tools/mongodb/metadata/collectionIndexes.ts

@@ -1,7 +1,20 @@
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
 import type { ToolArgs, OperationType } from "../../tool.js";
-import { formatUntrustedData } from "../../tool.js";
+import { FeatureFlags, formatUntrustedData } from "../../tool.js";
+
+type SearchIndexStatus = {
+    name: string;
+    type: string;
+    status: string;
+    queryable: boolean;
+    latestDefinition: Document;
+};
+
+type IndexStatus = {
+    name: string;
+    key: Document;
+};
 
 export class CollectionIndexesTool extends MongoDBToolBase {
     public name = "collection-indexes";
@@ -12,12 +25,30 @@ export class CollectionIndexesTool extends MongoDBToolBase {
     protected async execute({ database, collection }: ToolArgs<typeof DbOperationArgs>): Promise<CallToolResult> {
         const provider = await this.ensureConnected();
         const indexes = await provider.getIndexes(database, collection);
+        const indexDefinitions: IndexStatus[] = indexes.map((index) => ({
+            name: index.name as string,
+            key: index.key as Document,
+        }));
+
+        const searchIndexDefinitions: SearchIndexStatus[] = [];
+        if (this.isFeatureFlagEnabled(FeatureFlags.VectorSearch) && (await this.session.isSearchSupported())) {
+            const searchIndexes = await provider.getSearchIndexes(database, collection);
+            searchIndexDefinitions.push(...this.extractSearchIndexDetails(searchIndexes));
+        }
 
         return {
-            content: formatUntrustedData(
-                `Found ${indexes.length} indexes in the collection "${collection}":`,
-                ...indexes.map((index) => `Name: "${index.name}", definition: ${JSON.stringify(index.key)}`)
-            ),
+            content: [
+                ...formatUntrustedData(
+                    `Found ${indexDefinitions.length} indexes in the collection "${collection}":`,
+                    ...indexDefinitions.map((i) => JSON.stringify(i))
+                ),
+                ...(searchIndexDefinitions.length > 0
+                    ? formatUntrustedData(
+                          `Found ${searchIndexDefinitions.length} search and vector search indexes in the collection "${collection}":`,
+                          ...searchIndexDefinitions.map((i) => JSON.stringify(i))
+                      )
+                    : []),
+            ],
         };
     }
 
@@ -39,4 +70,20 @@ export class CollectionIndexesTool extends MongoDBToolBase {
 
         return super.handleError(error, args);
     }
+
+    /**
+     * Atlas Search index status contains a lot of information that is not relevant for the agent at this stage.
+     * Like for example, the status on each of the dedicated nodes. We only care about the main status, if it's
+     * queryable and the index name. We are also picking the index definition as it can be used by the agent to
+     * understand which fields are available for searching.
+     **/
+    protected extractSearchIndexDetails(indexes: Record<string, unknown>[]): SearchIndexStatus[] {
+        return indexes.map((index) => ({
+            name: (index["name"] ?? "default") as string,
+            type: (index["type"] ?? "UNKNOWN") as string,
+            status: (index["status"] ?? "UNKNOWN") as string,
+            queryable: (index["queryable"] ?? false) as boolean,
+            latestDefinition: index["latestDefinition"] as Document,
+        }));
+    }
 }

src/tools/mongodb/search/listSearchIndexes.ts

@@ -19,7 +19,6 @@ import { ExplainTool } from "./metadata/explain.js";
 import { CreateCollectionTool } from "./create/createCollection.js";
 import { LogsTool } from "./metadata/logs.js";
 import { ExportTool } from "./read/export.js";
-import { ListSearchIndexesTool } from "./search/listSearchIndexes.js";
 import { DropIndexTool } from "./delete/dropIndex.js";
 
 export const MongoDbTools = [
@@ -45,5 +44,4 @@ export const MongoDbTools = [
     CreateCollectionTool,
     LogsTool,
     ExportTool,
-    ListSearchIndexesTool,
 ];

src/tools/mongodb/tools.ts

@@ -37,4 +37,28 @@ describeAccuracyTests([
             },
         ],
     },
+    {
+        prompt: "how many search indexes do I have in the collection mydb.mycoll?",
+        expectedToolCalls: [
+            {
+                toolName: "collection-indexes",
+                parameters: {
+                    database: "mydb",
+                    collection: "mycoll",
+                },
+            },
+        ],
+    },
+    {
+        prompt: "which vector search indexes do I have in mydb.mycoll?",
+        expectedToolCalls: [
+            {
+                toolName: "collection-indexes",
+                parameters: {
+                    database: "mydb",
+                    collection: "mycoll",
+                },
+            },
+        ],
+    },
 ]);

tests/accuracy/collectionIndexes.test.ts

@@ -505,12 +505,10 @@ describeWithMongoDB(
         });
     },
     {
-        getUserConfig: () => {
-            return {
-                ...defaultTestConfig,
-                voyageApiKey: "valid_key",
-            };
-        },
+        getUserConfig: () => ({
+            ...defaultTestConfig,
+            voyageApiKey: "valid_key",
+        }),
         downloadOptions: {
             search: true,
         },