Skip to content

feat: add explain command #57

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
Apr 11, 2025
90 changes: 90 additions & 0 deletions src/tools/mongodb/metadata/explain.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,90 @@
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
import { DbOperationArgs, DbOperationType, MongoDBToolBase } from "../mongodbTool.js";
import { ToolArgs } from "../../tool.js";
import { z } from "zod";
import { ExplainVerbosity, Document } from "mongodb";
import { AggregateArgs } from "../read/aggregate.js";
import { FindArgs } from "../read/find.js";
import { CountArgs } from "../read/count.js";

export class ExplainTool extends MongoDBToolBase {
protected name = "explain";
protected description =
"Returns statistics describing the execution of the winning plan chosen by the query optimizer for the evaluated method";

protected argsShape = {
...DbOperationArgs,
method: z
.array(
z.union([
z.object({
name: z.literal("aggregate"),
arguments: z.object(AggregateArgs),
}),
z.object({
name: z.literal("find"),
arguments: z.object(FindArgs),
}),
z.object({
name: z.literal("count"),
arguments: z.object(CountArgs),
}),
])
)
.describe("The method and its arguments to run"),
};

protected operationType: DbOperationType = "metadata";

static readonly defaultVerbosity = ExplainVerbosity.queryPlanner;

protected async execute({
database,
collection,
method: methods,
}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
const provider = await this.ensureConnected();
const method = methods[0];

if (!method) {
throw new Error("No method provided");
}

let result: Document;
switch (method.name) {
case "aggregate": {
const { pipeline } = method.arguments;
result = await provider.aggregate(database, collection, pipeline).explain(ExplainTool.defaultVerbosity);
break;
}
case "find": {
const { filter, ...rest } = method.arguments;
result = await provider
.find(database, collection, filter as Document, { ...rest })
.explain(ExplainTool.defaultVerbosity);
break;
}
case "count": {
const { query } = method.arguments;
// This helper doesn't have explain() command but does have the argument explain
result = (await provider.count(database, collection, query, {
explain: ExplainTool.defaultVerbosity,
})) as unknown as Document;
break;
}
}

return {
content: [
{
text: `Here is some information about the winning plan chosen by the query optimizer for running the given \`${method.name}\` operation in \`${database}.${collection}\`. This information can be used to understand how the query was executed and to optimize the query performance.`,
type: "text",
},
{
text: JSON.stringify(result),
type: "text",
},
],
};
}
}
13 changes: 8 additions & 5 deletions src/tools/mongodb/read/aggregate.ts
Original file line number Diff line number Diff line change
@@ -1,16 +1,19 @@
import { z } from "zod";
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
import { DbOperationType, MongoDBToolBase } from "../mongodbTool.js";
import { DbOperationArgs, DbOperationType, MongoDBToolBase } from "../mongodbTool.js";
import { ToolArgs } from "../../tool.js";

export const AggregateArgs = {
pipeline: z.array(z.object({}).passthrough()).describe("An array of aggregation stages to execute"),
limit: z.number().optional().default(10).describe("The maximum number of documents to return"),
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I know you didn't change this, but I wonder if we should drop limit here? Considering the agg framework has a $limit stage, I wonder if that could conflict with a pipeline the model generates.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My concern is the toArray() call then. it's going to try to load all the documents from the aggregation result into memory

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

True - just wonder if we might run into trouble if a user asks for give me the top 50 docs and the model generates a stage, but our limit kicks in. Haven't looked into it, but we could check if there's a way to stream the responses to the model.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll remove it and we can revisit with a solution

};

export class AggregateTool extends MongoDBToolBase {
protected name = "aggregate";
protected description = "Run an aggregation against a MongoDB collection";
protected argsShape = {
collection: z.string().describe("Collection name"),
database: z.string().describe("Database name"),
pipeline: z.array(z.object({}).passthrough()).describe("An array of aggregation stages to execute"),
limit: z.number().optional().default(10).describe("The maximum number of documents to return"),
...DbOperationArgs,
...AggregateArgs,
};
protected operationType: DbOperationType = "read";

Expand Down
18 changes: 11 additions & 7 deletions src/tools/mongodb/read/count.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,18 +3,22 @@ import { DbOperationArgs, DbOperationType, MongoDBToolBase } from "../mongodbToo
import { ToolArgs } from "../../tool.js";
import { z } from "zod";

export const CountArgs = {
query: z
.object({})
.passthrough()
.optional()
.describe(
"The query filter to count documents. Matches the syntax of the filter argument of db.collection.count()"
),
};

export class CountTool extends MongoDBToolBase {
protected name = "count";
protected description = "Gets the number of documents in a MongoDB collection";
protected argsShape = {
...DbOperationArgs,
query: z
.object({})
.passthrough()
.optional()
.describe(
"The query filter to count documents. Matches the syntax of the filter argument of db.collection.count()"
),
...CountArgs,
};

protected operationType: DbOperationType = "metadata";
Expand Down
41 changes: 21 additions & 20 deletions src/tools/mongodb/read/find.ts
Original file line number Diff line number Diff line change
@@ -1,32 +1,33 @@
import { z } from "zod";
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
import { DbOperationType, MongoDBToolBase } from "../mongodbTool.js";
import { DbOperationArgs, DbOperationType, MongoDBToolBase } from "../mongodbTool.js";
import { ToolArgs } from "../../tool.js";
import { SortDirection } from "mongodb";

export const FindArgs = {
filter: z
.object({})
.passthrough()
.optional()
.describe("The query filter, matching the syntax of the query argument of db.collection.find()"),
projection: z
.object({})
.passthrough()
.optional()
.describe("The projection, matching the syntax of the projection argument of db.collection.find()"),
limit: z.number().optional().default(10).describe("The maximum number of documents to return"),
sort: z
.record(z.string(), z.custom<SortDirection>())
.optional()
.describe("A document, describing the sort order, matching the syntax of the sort argument of cursor.sort()"),
};

export class FindTool extends MongoDBToolBase {
protected name = "find";
protected description = "Run a find query against a MongoDB collection";
protected argsShape = {
collection: z.string().describe("Collection name"),
database: z.string().describe("Database name"),
filter: z
.object({})
.passthrough()
.optional()
.describe("The query filter, matching the syntax of the query argument of db.collection.find()"),
projection: z
.object({})
.passthrough()
.optional()
.describe("The projection, matching the syntax of the projection argument of db.collection.find()"),
limit: z.number().optional().default(10).describe("The maximum number of documents to return"),
sort: z
.record(z.string(), z.custom<SortDirection>())
.optional()
.describe(
"A document, describing the sort order, matching the syntax of the sort argument of cursor.sort()"
),
...DbOperationArgs,
...FindArgs,
};
protected operationType: DbOperationType = "read";

Expand Down
2 changes: 2 additions & 0 deletions src/tools/mongodb/tools.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import { UpdateManyTool } from "./update/updateMany.js";
import { RenameCollectionTool } from "./update/renameCollection.js";
import { DropDatabaseTool } from "./delete/dropDatabase.js";
import { DropCollectionTool } from "./delete/dropCollection.js";
import { ExplainTool } from "./metadata/explain.js";

export const MongoDbTools = [
ConnectTool,
Expand All @@ -40,4 +41,5 @@ export const MongoDbTools = [
RenameCollectionTool,
DropDatabaseTool,
DropCollectionTool,
ExplainTool,
];
Loading