Skip to content

fix: updates count tool #254

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

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 16 additions & 5 deletions src/tools/mongodb/read/count.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,27 +4,38 @@ import { ToolArgs, OperationType } from "../../tool.js";
import { z } from "zod";

export const CountArgs = {
query: z
filter: z
.record(z.string(), z.unknown())
.optional()
.describe(
"The query filter to count documents. Matches the syntax of the filter argument of db.collection.count()"
"The query filter to count documents. Matches the syntax of the filter argument of db.collection.countDocuments()"
),
query: z
.record(z.string(), z.unknown())
.optional()
.describe("Alternative old name for filter. Will be used in db.collection.countDocuments()"),
Copy link
Collaborator

Choose a reason for hiding this comment

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

is there a reason to keep providing this field?

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

yes, since the LLM was already confusing query/filter, I figure that it can get confused the other way around and try to send it as "query" because it thinks we'll use .count() (even though we call out is countDocments, but who knows, it might have outdated data about MDB)

Copy link
Collaborator

Choose a reason for hiding this comment

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

This doesn't make much sense to me. I think we should tweak the description of filter rather than accept two arguments that are doing the same thing.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

np, I'll remove and update the description

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

we can try that for now and see if it improves, based on the issue I'm not sure the LLM will stop getting confused

};

export class CountTool extends MongoDBToolBase {
protected name = "count";
protected description = "Gets the number of documents in a MongoDB collection";
protected description = "Gets the number of documents in a MongoDB collection using countDocuments()";
protected argsShape = {
...DbOperationArgs,
...CountArgs,
};

protected operationType: OperationType = "read";

protected async execute({ database, collection, query }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
protected async execute({
database,
collection,
query,
filter,
}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
const provider = await this.ensureConnected();
const count = await provider.count(database, collection, query);
// use either filter or query, since we're using countDocuments, prefer filter
const queryFilter = filter || query;
const count = await provider.countDocuments(database, collection, queryFilter);

return {
content: [
Expand Down
68 changes: 58 additions & 10 deletions tests/integration/tools/mongodb/read/count.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,16 +8,27 @@ import {
} from "../../../helpers.js";

describeWithMongoDB("count tool", (integration) => {
validateToolMetadata(integration, "count", "Gets the number of documents in a MongoDB collection", [
{
name: "query",
description:
"The query filter to count documents. Matches the syntax of the filter argument of db.collection.count()",
type: "object",
required: false,
},
...databaseCollectionParameters,
]);
validateToolMetadata(
integration,
"count",
"Gets the number of documents in a MongoDB collection using countDocuments()",
[
{
name: "filter",
description:
"The query filter to count documents. Matches the syntax of the filter argument of db.collection.countDocuments()",
type: "object",
required: false,
},
{
name: "query",
description: "Alternative old name for filter. Will be used in db.collection.countDocuments()",
type: "object",
required: false,
},
...databaseCollectionParameters,
]
);

validateThrowsForInvalidArguments(integration, "count", [
{},
Expand Down Expand Up @@ -79,6 +90,43 @@ describeWithMongoDB("count tool", (integration) => {
expect(content).toEqual(`Found ${testCase.expectedCount} documents in the collection "foo"`);
});
}

it("correctly filters documents when using 'filter' parameter", async () => {
await integration.connectMcpClient();

// Using 'filter' parameter - should work correctly after the fix
const response = await integration.mcpClient().callTool({
name: "count",
arguments: {
database: integration.randomDbName(),
collection: "foo",
filter: { age: { $lt: 15 } },
},
});

const content = getResponseContent(response.content);
expect(content).toEqual('Found 2 documents in the collection "foo"');
});

it("prioritizes filter over query when both are provided", async () => {
await integration.connectMcpClient();

// Using both 'filter' and 'query' parameters
const response = await integration.mcpClient().callTool({
name: "count",
arguments: {
database: integration.randomDbName(),
collection: "foo",
filter: { age: { $lt: 15 } },
query: { age: { $gt: 10 } },
},
});

const content = getResponseContent(response.content);
// Filter takes precedence over query
// Filter is { age: { $lt: 15 } } which matches 2 documents
expect(content).toEqual('Found 2 documents in the collection "foo"');
});
});

validateAutoConnectBehavior(integration, "count", () => {
Expand Down
Loading