Skip to content

Commit be45e70

Browse files
chore: schema def mods for better LLM responses
In the previous commit, our schema description was not clear enough for LLM to correctly provide the name of command and it was hallucinating by providing listDatabases for example instead of list-databases. This commit modifies the descriptions a bit to achieve 100% accuracy for the same prompts.
1 parent 919b26f commit be45e70

File tree

4 files changed

+57
-36
lines changed

4 files changed

+57
-36
lines changed

src/tools/mongodb/ddl/ddl.ts

Lines changed: 34 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -10,34 +10,40 @@ export class MongoDBDDLTool extends MongoDBToolBase {
1010
protected description =
1111
"List databases, collections, indexes and describe the schema of a collection in a MongoDB database";
1212
protected argsShape = {
13-
command: z.discriminatedUnion("name", [
14-
z
15-
.object({
16-
name: z.literal("list-databases"),
17-
parameters: z.object({}),
18-
})
19-
.describe("List all databases for a MongoDB connection"),
20-
z
21-
.object({
22-
name: z.literal("list-collections"),
23-
parameters: z.object({
24-
database: DbOperationArgs.database,
25-
}),
26-
})
27-
.describe("List all collections for a given database"),
28-
z
29-
.object({
30-
name: z.literal("collection-indexes"),
31-
parameters: z.object(DbOperationArgs),
32-
})
33-
.describe("Describe the indexes for a collection"),
34-
z
35-
.object({
36-
name: z.literal("collection-schema"),
37-
parameters: z.object(DbOperationArgs),
38-
})
39-
.describe("Describe the schema for a collection"),
40-
]),
13+
command: z
14+
.discriminatedUnion("name", [
15+
z
16+
.object({
17+
name: z.literal("list-databases"),
18+
parameters: z.object({}),
19+
})
20+
.describe(
21+
"The shape of 'list-databases' command to list all the databases for a MongoDB connection."
22+
),
23+
z
24+
.object({
25+
name: z.literal("list-collections"),
26+
parameters: z.object({
27+
database: DbOperationArgs.database,
28+
}),
29+
})
30+
.describe(
31+
"The shape of 'list-collections' command to list all the collections for a given database."
32+
),
33+
z
34+
.object({
35+
name: z.literal("collection-indexes"),
36+
parameters: z.object(DbOperationArgs),
37+
})
38+
.describe("The shape of 'collection-indexes' command to describe the indexes for a collection."),
39+
z
40+
.object({
41+
name: z.literal("collection-schema"),
42+
parameters: z.object(DbOperationArgs),
43+
})
44+
.describe("The shape of 'collection-schema' command to describe the schema for a collection."),
45+
])
46+
.describe("The command to be provided to the tool."),
4147
};
4248
protected operationType: OperationType = "read";
4349

tests/accuracy/collection-indexes.test.ts

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,15 @@ function callsCollectionIndexes(prompt: string): AccuracyTestConfig {
77
prompt: prompt,
88
expectedToolCalls: [
99
{
10-
toolName: "collection-indexes",
10+
toolName: "mongodb-ddl",
1111
parameters: {
12-
database: "mflix",
13-
collection: "movies",
12+
command: {
13+
name: "collection-indexes",
14+
parameters: {
15+
database: "mflix",
16+
collection: "movies",
17+
},
18+
},
1419
},
1520
},
1621
],

tests/accuracy/collection-schema.test.ts

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,15 @@ function callsCollectionSchema(prompt: string): AccuracyTestConfig {
77
prompt: prompt,
88
expectedToolCalls: [
99
{
10-
toolName: "collection-schema",
10+
toolName: "mongodb-ddl",
1111
parameters: {
12-
database: "db1",
13-
collection: "coll1",
12+
command: {
13+
name: "collection-schema",
14+
parameters: {
15+
database: "db1",
16+
collection: "coll1",
17+
},
18+
},
1419
},
1520
},
1621
],

tests/accuracy/list-databases.test.ts

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,13 @@ function callsListDatabases(prompt: string): AccuracyTestConfig {
77
prompt: prompt,
88
expectedToolCalls: [
99
{
10-
toolName: "list-databases",
11-
parameters: {},
10+
toolName: "mongodb-ddl",
11+
parameters: {
12+
command: {
13+
name: "list-databases",
14+
parameters: {},
15+
},
16+
},
1217
},
1318
],
1419
};

0 commit comments

Comments
 (0)