Skip to content

Commit 780231d

Browse files
himanshusinghshimanshusinghs
committed
rebase: find tool schema fixes
1 parent 6532c7a commit 780231d

File tree

4 files changed

+85
-26
lines changed

4 files changed

+85
-26
lines changed

‎scripts/accuracy/runAccuracyTests.sh

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#!/bin/sh
2+
source ~/.accuracy
23
# Variables necessary for the accuracy test runs
34
export MDB_ACCURACY_RUN_ID=$(npx uuid v4)
45

‎src/tools/mongodb/metadata/explain.ts

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import { ToolArgs, OperationType } from "../../tool.js";
44
import { z } from "zod";
55
import { ExplainVerbosity, Document } from "mongodb";
66
import { AggregateArgs } from "../read/aggregate.js";
7-
import { FindArgs } from "../read/find.js";
7+
import { FindArgs, keyValueListToDocument } from "../read/find.js";
88
import { CountArgs } from "../read/count.js";
99

1010
export class ExplainTool extends MongoDBToolBase {
@@ -68,9 +68,12 @@ export class ExplainTool extends MongoDBToolBase {
6868
break;
6969
}
7070
case "find": {
71-
const { filter, ...rest } = method.arguments;
71+
const { filter, sort, projection, ...rest } = method.arguments;
72+
const mongoFilter = keyValueListToDocument(filter);
73+
const mongoProjection = keyValueListToDocument(projection);
74+
const mongoSort = keyValueListToDocument(sort);
7275
result = await provider
73-
.find(database, collection, filter as Document, { ...rest })
76+
.find(database, collection, mongoFilter, { projection: mongoProjection, sort: mongoSort, ...rest })
7477
.explain(ExplainTool.defaultVerbosity);
7578
break;
7679
}

‎src/tools/mongodb/read/find.ts

Lines changed: 61 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,59 @@
11
import { z } from "zod";
2+
import { EJSON } from "bson";
3+
import { SortDirection } from "mongodb";
24
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
3-
import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
5+
46
import { ToolArgs, OperationType } from "../../tool.js";
5-
import { SortDirection } from "mongodb";
6-
import { EJSON } from "bson";
7+
import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
78
import { checkIndexUsage } from "../../../helpers/indexCheck.js";
89

10+
export function keyValueListToDocument<V = unknown>(
11+
keyValueList: { key: string; value: V }[] | undefined
12+
): Record<string, V> | undefined {
13+
return keyValueList ? Object.fromEntries(keyValueList.map(({ key, value }) => [key, value])) : undefined;
14+
}
15+
916
export const FindArgs = {
1017
filter: z
11-
.record(z.string(), z.unknown())
18+
.array(
19+
z.object({
20+
key: z.string().describe("The name of the field or a MongoDB operator"),
21+
value: z
22+
.unknown()
23+
.refine((val) => val !== undefined, { message: "Value cannot be undefined." })
24+
.describe("The filter expression for the key"),
25+
})
26+
)
1227
.optional()
13-
.describe("The query filter, matching the syntax of the query argument of db.collection.find()"),
28+
.describe(
29+
"Array of key-value pairs to filter documents. Each object has 'key' (field name or MongoDB operator) and 'value' (filter criteria)."
30+
),
1431
projection: z
15-
.record(z.string(), z.unknown())
32+
.array(
33+
z.object({
34+
key: z.string().describe("The name of the field to be projected."),
35+
value: z
36+
.unknown()
37+
.refine((val) => val !== undefined, { message: "Value cannot be undefined." })
38+
.describe("The projection expression for the projected field."),
39+
})
40+
)
1641
.optional()
17-
.describe("The projection, matching the syntax of the projection argument of db.collection.find()"),
42+
.describe(
43+
"Array of key-value pairs to specify which fields to project (key) and how to project them(value). Each object has 'key' (field name) and 'value' (project expression). "
44+
),
1845
limit: z.number().optional().default(10).describe("The maximum number of documents to return"),
1946
sort: z
20-
.record(z.string(), z.custom<SortDirection>())
47+
.array(
48+
z.object({
49+
key: z.string().describe("The name of the field to apply the sort on."),
50+
value: z.custom<SortDirection>().describe("The sort order applied to the field being sorted."),
51+
})
52+
)
2153
.optional()
22-
.describe("A document, describing the sort order, matching the syntax of the sort argument of cursor.sort()"),
54+
.describe(
55+
"Array of key-value pairs to specify sort order. Each object has 'key' (field name) and 'value' (sort order)."
56+
),
2357
};
2458

2559
export class FindTool extends MongoDBToolBase {
@@ -41,14 +75,29 @@ export class FindTool extends MongoDBToolBase {
4175
}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
4276
const provider = await this.ensureConnected();
4377

44-
// Check if find operation uses an index if enabled
78+
const mongoFilter = keyValueListToDocument(filter);
79+
const mongoProjection = keyValueListToDocument(projection);
80+
const mongoSort = keyValueListToDocument(sort);
81+
4582
if (this.config.indexCheck) {
4683
await checkIndexUsage(provider, database, collection, "find", async () => {
47-
return provider.find(database, collection, filter, { projection, limit, sort }).explain("queryPlanner");
84+
return provider
85+
.find(database, collection, mongoFilter, {
86+
projection: mongoProjection,
87+
limit,
88+
sort: mongoSort,
89+
})
90+
.explain("queryPlanner");
4891
});
4992
}
5093

51-
const documents = await provider.find(database, collection, filter, { projection, limit, sort }).toArray();
94+
const documents = await provider
95+
.find(database, collection, mongoFilter, {
96+
projection: mongoProjection,
97+
limit,
98+
sort: mongoSort,
99+
})
100+
.toArray();
52101

53102
const content: Array<{ text: string; type: "text" }> = [
54103
{

‎tests/accuracy/find.test.ts

Lines changed: 17 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -33,9 +33,12 @@ describeAccuracyTests([
3333
parameters: {
3434
database: "mflix",
3535
collection: "movies",
36-
filter: {
37-
runtime: { $lt: 100 },
38-
},
36+
filter: [
37+
{
38+
key: "runtime",
39+
value: { $lt: 100 },
40+
},
41+
],
3942
},
4043
},
4144
],
@@ -48,9 +51,12 @@ describeAccuracyTests([
4851
parameters: {
4952
database: "mflix",
5053
collection: "movies",
51-
filter: {
52-
director: "Christina Collins",
53-
},
54+
filter: [
55+
{
56+
key: "director",
57+
value: "Christina Collins",
58+
},
59+
],
5460
},
5561
},
5662
],
@@ -63,7 +69,7 @@ describeAccuracyTests([
6369
parameters: {
6470
database: "mflix",
6571
collection: "movies",
66-
projection: { title: 1 },
72+
projection: [{ key: "title", value: 1 }],
6773
},
6874
},
6975
],
@@ -76,8 +82,8 @@ describeAccuracyTests([
7682
parameters: {
7783
database: "mflix",
7884
collection: "movies",
79-
filter: { title: "Certain Fish" },
80-
projection: { cast: 1 },
85+
filter: [{ key: "title", value: "Certain Fish" }],
86+
projection: [{ key: "cast", value: 1 }],
8187
},
8288
},
8389
],
@@ -90,8 +96,8 @@ describeAccuracyTests([
9096
parameters: {
9197
database: "mflix",
9298
collection: "movies",
93-
filter: { title: "Certain Fish" },
94-
sort: { runtime: 1 },
99+
filter: [{ key: "title", value: "Certain Fish" }],
100+
sort: [{ key: "runtime", value: 1 }],
95101
limit: 2,
96102
},
97103
},

0 commit comments

Comments
 (0)