chore(tools): POC to consolidate immutable ddl tools while preserving the accuracy

src/tools/mongodb/ddl/ddl.ts

@@ -0,0 +1,169 @@
+import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
+import * as bson from "bson";
+import { OperationType, ToolArgs } from "../../tool.js";
+import z from "zod";
+import { getSimplifiedSchema } from "mongodb-schema";
+
+export class MongoDBDDLTool extends MongoDBToolBase {
+    protected name = "mongodb-ddl";
+    protected description =
+        "List databases, collections, indexes and describe the schema of a collection in a MongoDB database";
+    protected argsShape = {
+        command: z
+            .discriminatedUnion("name", [
+                z
+                    .object({
+                        name: z.literal("list-databases"),
+                        parameters: z.object({}),
+                    })
+                    .describe(
+                        "The shape of 'list-databases' command to list all the databases for a MongoDB connection."
+                    ),
+                z
+                    .object({
+                        name: z.literal("list-collections"),
+                        parameters: z.object({
+                            database: DbOperationArgs.database,
+                        }),
+                    })
+                    .describe(
+                        "The shape of 'list-collections' command to list all the collections for a given database."
+                    ),
+                z
+                    .object({
+                        name: z.literal("collection-indexes"),
+                        parameters: z.object(DbOperationArgs),
+                    })
+                    .describe("The shape of 'collection-indexes' command to describe the indexes for a collection."),
+                z
+                    .object({
+                        name: z.literal("collection-schema"),
+                        parameters: z.object(DbOperationArgs),
+                    })
+                    .describe("The shape of 'collection-schema' command to describe the schema for a collection."),
+            ])
+            .describe("The command to be provided to the tool."),
+    };
+    protected operationType: OperationType = "read";
+
+    protected async execute({ command }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
+        const provider = await this.ensureConnected();
+        if (command.name === "list-databases") {
+            const dbs = (await provider.listDatabases("")).databases as { name: string; sizeOnDisk: bson.Long }[];
+
+            return {
+                content: dbs.map((db) => {
+                    return {
+                        text: `Name: ${db.name}, Size: ${db.sizeOnDisk.toString()} bytes`,
+                        type: "text",
+                    };
+                }),
+            };
+        }
+
+        if (command.name === "list-collections") {
+            const { database } = command.parameters;
+            const collections = await provider.listCollections(database);
+
+            if (collections.length === 0) {
+                return {
+                    content: [
+                        {
+                            type: "text",
+                            text: `No collections found for database "${database}". To create a collection, use the "create-collection" tool.`,
+                        },
+                    ],
+                };
+            }
+
+            return {
+                content: collections.map((collection) => {
+                    return {
+                        text: `Name: "${collection.name}"`,
+                        type: "text",
+                    };
+                }),
+            };
+        }
+
+        if (command.name === "collection-indexes") {
+            const { database, collection } = command.parameters;
+            const indexes = await provider.getIndexes(database, collection);
+
+            return {
+                content: [
+                    {
+                        text: `Found ${indexes.length} indexes in the collection "${collection}":`,
+                        type: "text",
+                    },
+                    ...(indexes.map((indexDefinition) => {
+                        return {
+                            text: `Name "${indexDefinition.name}", definition: ${JSON.stringify(indexDefinition.key)}`,
+                            type: "text",
+                        };
+                    }) as { text: string; type: "text" }[]),
+                ],
+            };
+        }
+
+        if (command.name === "collection-schema") {
+            const { database, collection } = command.parameters;
+            const documents = await provider.find(database, collection, {}, { limit: 5 }).toArray();
+            const schema = await getSimplifiedSchema(documents);
+
+            const fieldsCount = Object.entries(schema).length;
+            if (fieldsCount === 0) {
+                return {
+                    content: [
+                        {
+                            text: `Could not deduce the schema for "${database}.${collection}". This may be because it doesn't exist or is empty.`,
+                            type: "text",
+                        },
+                    ],
+                };
+            }
+
+            return {
+                content: [
+                    {
+                        text: `Found ${fieldsCount} fields in the schema for "${database}.${collection}"`,
+                        type: "text",
+                    },
+                    {
+                        text: JSON.stringify(schema),
+                        type: "text",
+                    },
+                ],
+            };
+        }
+
+        return {
+            content: [
+                {
+                    text: `Unknown command provided to the tool.`,
+                    type: "text",
+                },
+            ],
+        };
+    }
+
+    protected handleError(
+        error: unknown,
+        args: ToolArgs<typeof this.argsShape>
+    ): Promise<CallToolResult> | CallToolResult {
+        if (args.command.name === "collection-indexes") {
+            if (error instanceof Error && "codeName" in error && error.codeName === "NamespaceNotFound") {
+                return {
+                    content: [
+                        {
+                            text: `The indexes for "${args.command.parameters.database}.${args.command.parameters.collection}" cannot be determined because the collection does not exist.`,
+                            type: "text",
+                        },
+                    ],
+                };
+            }
+        }
+        return super.handleError(error, args);
+    }
+}

src/tools/mongodb/tools.ts

@@ -1,9 +1,5 @@
 import { ConnectTool } from "./metadata/connect.js";
-import { ListCollectionsTool } from "./metadata/listCollections.js";
-import { CollectionIndexesTool } from "./read/collectionIndexes.js";
-import { ListDatabasesTool } from "./metadata/listDatabases.js";
 import { CreateIndexTool } from "./create/createIndex.js";
-import { CollectionSchemaTool } from "./metadata/collectionSchema.js";
 import { FindTool } from "./read/find.js";
 import { InsertManyTool } from "./create/insertMany.js";
 import { DeleteManyTool } from "./delete/deleteMany.js";
@@ -18,14 +14,12 @@ import { DropCollectionTool } from "./delete/dropCollection.js";
 import { ExplainTool } from "./metadata/explain.js";
 import { CreateCollectionTool } from "./create/createCollection.js";
 import { LogsTool } from "./metadata/logs.js";
+import { MongoDBDDLTool } from "./ddl/ddl.js";
 
 export const MongoDbTools = [
     ConnectTool,
-    ListCollectionsTool,
-    ListDatabasesTool,
-    CollectionIndexesTool,
+    MongoDBDDLTool,
     CreateIndexTool,
-    CollectionSchemaTool,
     FindTool,
     InsertManyTool,
     DeleteManyTool,

tests/accuracy/collection-indexes.test.ts

@@ -7,10 +7,15 @@ function callsCollectionIndexes(prompt: string): AccuracyTestConfig {
         prompt: prompt,
         expectedToolCalls: [
             {
-                toolName: "collection-indexes",
+                toolName: "mongodb-ddl",
                 parameters: {
-                    database: "mflix",
-                    collection: "movies",
+                    command: {
+                        name: "collection-indexes",
+                        parameters: {
+                            database: "mflix",
+                            collection: "movies",
+                        },
+                    },
                 },
             },
         ],

tests/accuracy/collection-schema.test.ts

@@ -7,10 +7,15 @@ function callsCollectionSchema(prompt: string): AccuracyTestConfig {
         prompt: prompt,
         expectedToolCalls: [
             {
-                toolName: "collection-schema",
+                toolName: "mongodb-ddl",
                 parameters: {
-                    database: "db1",
-                    collection: "coll1",
+                    command: {
+                        name: "collection-schema",
+                        parameters: {
+                            database: "db1",
+                            collection: "coll1",
+                        },
+                    },
                 },
             },
         ],

tests/accuracy/list-collections.test.ts

@@ -7,8 +7,13 @@ function callsListCollections(prompt: string): AccuracyTestConfig {
         prompt: prompt,
         expectedToolCalls: [
             {
-                toolName: "list-collections",
-                parameters: { database: "mflix" },
+                toolName: "mongodb-ddl",
+                parameters: {
+                    command: {
+                        name: "list-collections",
+                        parameters: { database: "mflix" },
+                    },
+                },
             },
         ],
     };
@@ -21,28 +26,58 @@ function callsListDatabasesAndListCollections(prompt: string): AccuracyTestConfi
         mockedTools: {},
         expectedToolCalls: [
             {
-                toolName: "list-databases",
-                parameters: {},
+                toolName: "mongodb-ddl",
+                parameters: {
+                    command: {
+                        name: "list-databases",
+                        parameters: {},
+                    },
+                },
             },
             {
-                toolName: "list-collections",
-                parameters: { database: "admin" },
+                toolName: "mongodb-ddl",
+                parameters: {
+                    command: {
+                        name: "list-collections",
+                        parameters: { database: "admin" },
+                    },
+                },
             },
             {
-                toolName: "list-collections",
-                parameters: { database: "comics" },
+                toolName: "mongodb-ddl",
+                parameters: {
+                    command: {
+                        name: "list-collections",
+                        parameters: { database: "comics" },
+                    },
+                },
             },
             {
-                toolName: "list-collections",
-                parameters: { database: "config" },
+                toolName: "mongodb-ddl",
+                parameters: {
+                    command: {
+                        name: "list-collections",
+                        parameters: { database: "config" },
+                    },
+                },
             },
             {
-                toolName: "list-collections",
-                parameters: { database: "local" },
+                toolName: "mongodb-ddl",
+                parameters: {
+                    command: {
+                        name: "list-collections",
+                        parameters: { database: "local" },
+                    },
+                },
             },
             {
-                toolName: "list-collections",
-                parameters: { database: "mflix" },
+                toolName: "mongodb-ddl",
+                parameters: {
+                    command: {
+                        name: "list-collections",
+                        parameters: { database: "mflix" },
+                    },
+                },
             },
         ],
     };

tests/accuracy/list-databases.test.ts

@@ -7,8 +7,13 @@ function callsListDatabases(prompt: string): AccuracyTestConfig {
         prompt: prompt,
         expectedToolCalls: [
             {
-                toolName: "list-databases",
-                parameters: {},
+                toolName: "mongodb-ddl",
+                parameters: {
+                    command: {
+                        name: "list-databases",
+                        parameters: {},
+                    },
+                },
             },
         ],
     };