Skip to content

Commit 4113579

Browse files
nirinchevnirinchev
committed
tweak create index args to be compatible with clients
1 parent 33ba4d1 commit 4113579

File tree

3 files changed

+181
-82
lines changed

3 files changed

+181
-82
lines changed

‎src/tools/mongodb/create/createIndex.ts‎

Lines changed: 55 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -11,52 +11,55 @@ export class CreateIndexTool extends MongoDBToolBase {
1111
...DbOperationArgs,
1212
name: z.string().optional().describe("The name of the index"),
1313
definition: z
14-
.discriminatedUnion("type", [
15-
z.object({
16-
type: z.literal("classic"),
17-
keys: z.object({}).catchall(z.custom<IndexDirection>()).describe("The index definition"),
18-
}),
19-
z.object({
20-
type: z.literal("vectorSearch"),
21-
fields: z
22-
.array(
23-
z.object({
24-
type: z
25-
.enum(["vector", "filter"])
26-
.describe(
27-
"Field type to use to index fields for vector search. You must specify `vector` for fields that contain vector embeddings and `filter` for additional fields to filter on."
28-
),
29-
path: z
30-
.string()
31-
.describe(
32-
"Name of the field to index. For nested fields, use dot notation to specify path to embedded fields"
33-
),
34-
numDimensions: z
35-
.number()
36-
.min(1)
37-
.max(8192)
38-
.describe(
39-
"Number of vector dimensions that MongoDB Vector Search enforces at index-time and query-time"
40-
),
41-
similarity: z
42-
.enum(["cosine", "euclidean", "dotProduct"])
43-
.describe(
44-
"Vector similarity function to use to search for top K-nearest neighbors. You can set this field only for vector-type fields."
45-
),
46-
quantization: z
47-
.enum(["none", "scalar", "binary"])
48-
.optional()
49-
.default("none")
50-
.describe(
51-
"Type of automatic vector quantization for your vectors. Use this setting only if your embeddings are float or double vectors."
52-
),
53-
})
54-
)
55-
.describe(
56-
"Definitions for the vector and filter fields to index, one definition per document. The fields array must contain at least one vector-type field definition."
57-
),
58-
}),
59-
])
14+
.array(
15+
z.discriminatedUnion("type", [
16+
z.object({
17+
type: z.literal("classic"),
18+
keys: z.object({}).catchall(z.custom<IndexDirection>()).describe("The index definition"),
19+
}),
20+
z.object({
21+
type: z.literal("vectorSearch"),
22+
fields: z
23+
.array(
24+
z.object({
25+
type: z
26+
.enum(["vector", "filter"])
27+
.describe(
28+
"Field type to use to index fields for vector search. You must specify `vector` for fields that contain vector embeddings and `filter` for additional fields to filter on."
29+
),
30+
path: z
31+
.string()
32+
.describe(
33+
"Name of the field to index. For nested fields, use dot notation to specify path to embedded fields"
34+
),
35+
numDimensions: z
36+
.number()
37+
.min(1)
38+
.max(8192)
39+
.describe(
40+
"Number of vector dimensions that MongoDB Vector Search enforces at index-time and query-time"
41+
),
42+
similarity: z
43+
.enum(["cosine", "euclidean", "dotProduct"])
44+
.default("cosine")
45+
.describe(
46+
"Vector similarity function to use to search for top K-nearest neighbors. You can set this field only for vector-type fields."
47+
),
48+
quantization: z
49+
.enum(["none", "scalar", "binary"])
50+
.optional()
51+
.default("none")
52+
.describe(
53+
"Type of automatic vector quantization for your vectors. Use this setting only if your embeddings are float or double vectors."
54+
),
55+
})
56+
)
57+
.describe(
58+
"Definitions for the vector and filter fields to index, one definition per document. The fields array must contain at least one vector-type field definition."
59+
),
60+
}),
61+
])
62+
)
6063
.describe(
6164
"The index definition. Use 'classic' for standard indexes and 'vectorSearch' for vector search indexes"
6265
),
@@ -68,10 +71,15 @@ export class CreateIndexTool extends MongoDBToolBase {
6871
database,
6972
collection,
7073
name,
71-
definition,
74+
definition: definitions,
7275
}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
7376
const provider = await this.ensureConnected();
7477
let indexes: string[] = [];
78+
const definition = definitions[0];
79+
if (!definition) {
80+
throw new Error("Index definition not provided. Expected one of the following: `classic`, `vectorSearch`");
81+
}
82+
7583
switch (definition.type) {
7684
case "classic":
7785
indexes = await provider.createIndexes(database, collection, [

‎tests/accuracy/createIndex.test.ts‎

Lines changed: 106 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -11,12 +11,14 @@ describeAccuracyTests([
1111
database: "mflix",
1212
collection: "movies",
1313
name: Matcher.anyOf(Matcher.undefined, Matcher.string()),
14-
definition: {
15-
type: "classic",
16-
keys: {
17-
release_year: 1,
14+
definition: [
15+
{
16+
type: "classic",
17+
keys: {
18+
release_year: 1,
19+
},
1820
},
19-
},
21+
],
2022
},
2123
},
2224
],
@@ -30,12 +32,106 @@ describeAccuracyTests([
3032
database: "mflix",
3133
collection: "movies",
3234
name: Matcher.anyOf(Matcher.undefined, Matcher.string()),
33-
definition: {
34-
type: "classic",
35-
keys: {
36-
title: "text",
35+
definition: [
36+
{
37+
type: "classic",
38+
keys: {
39+
title: "text",
40+
},
3741
},
38-
},
42+
],
43+
},
44+
},
45+
],
46+
},
47+
{
48+
prompt: "Create a vector search index on 'mydb.movies' namespace with on the 'plotSummary' field. The index should use 1024 dimensions.",
49+
expectedToolCalls: [
50+
{
51+
toolName: "create-index",
52+
parameters: {
53+
database: "mydb",
54+
collection: "movies",
55+
name: Matcher.anyOf(Matcher.undefined, Matcher.string()),
56+
definition: [
57+
{
58+
type: "vectorSearch",
59+
fields: [
60+
{
61+
type: "vector",
62+
path: "plotSummary",
63+
numDimensions: 1024,
64+
},
65+
{
66+
type: "filter",
67+
path: "releaseDate",
68+
},
69+
],
70+
},
71+
],
72+
},
73+
},
74+
],
75+
},
76+
{
77+
prompt: "Create a vector search index on 'mydb.movies' namespace with on the 'plotSummary' field and 'genre' field. Pick a sensible number of dimensions for a voyage 3.5 model.",
78+
expectedToolCalls: [
79+
{
80+
toolName: "create-index",
81+
parameters: {
82+
database: "mydb",
83+
collection: "movies",
84+
name: Matcher.anyOf(Matcher.undefined, Matcher.string()),
85+
definition: [
86+
{
87+
type: "vectorSearch",
88+
fields: [
89+
{
90+
type: "vector",
91+
path: "plotSummary",
92+
numDimensions: Matcher.number(
93+
(value) => value % 2 === 0 && value >= 256 && value <= 8192
94+
),
95+
},
96+
{
97+
type: "vector",
98+
path: "genre",
99+
numDimensions: Matcher.number(
100+
(value) => value % 2 === 0 && value >= 256 && value <= 8192
101+
),
102+
},
103+
],
104+
},
105+
],
106+
},
107+
},
108+
],
109+
},
110+
{
111+
prompt: "Create a vector search index on 'mydb.movies' namespace where the 'plotSummary' field is indexed as a 1024-dimensional vector and the 'releaseDate' field is indexed as a regular field.",
112+
expectedToolCalls: [
113+
{
114+
toolName: "create-index",
115+
parameters: {
116+
database: "mydb",
117+
collection: "movies",
118+
name: Matcher.anyOf(Matcher.undefined, Matcher.string()),
119+
definition: [
120+
{
121+
type: "vectorSearch",
122+
fields: [
123+
{
124+
type: "vector",
125+
path: "plotSummary",
126+
numDimensions: 1024,
127+
},
128+
{
129+
type: "filter",
130+
path: "releaseDate",
131+
},
132+
],
133+
},
134+
],
39135
},
40136
},
41137
],

‎tests/integration/tools/mongodb/create/createIndex.test.ts‎

Lines changed: 20 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,7 @@ describeWithMongoDB("createIndex tool", (integration) => {
1515
...databaseCollectionParameters,
1616
{
1717
name: "definition",
18-
anyOf: [
19-
{
20-
type: "object",
21-
},
22-
{
23-
type: "object",
24-
},
25-
],
18+
type: "array",
2619
description:
2720
"The index definition. Use 'classic' for standard indexes and 'vectorSearch' for vector search indexes",
2821
required: true,
@@ -37,10 +30,10 @@ describeWithMongoDB("createIndex tool", (integration) => {
3730

3831
validateThrowsForInvalidArguments(integration, "create-index", [
3932
{},
40-
{ collection: "bar", database: 123, definition: { type: "classic", keys: { foo: 1 } } },
41-
{ collection: [], database: "test", definition: { type: "classic", keys: { foo: 1 } } },
42-
{ collection: "bar", database: "test", definition: { type: "classic", keys: { foo: 1 } }, name: 123 },
43-
{ collection: "bar", database: "test", definition: { type: "unknown", keys: { foo: 1 } }, name: "my-index" },
33+
{ collection: "bar", database: 123, definition: [{ type: "classic", keys: { foo: 1 } }] },
34+
{ collection: [], database: "test", definition: [{ type: "classic", keys: { foo: 1 } }] },
35+
{ collection: "bar", database: "test", definition: [{ type: "classic", keys: { foo: 1 } }], name: 123 },
36+
{ collection: "bar", database: "test", definition: [{ type: "unknown", keys: { foo: 1 } }], name: "my-index" },
4437
]);
4538

4639
describe("with classic indexes", () => {
@@ -66,10 +59,12 @@ describeWithMongoDB("createIndex tool", (integration) => {
6659
arguments: {
6760
database: integration.randomDbName(),
6861
collection: "coll1",
69-
definition: {
70-
type: "classic",
71-
keys: { prop1: 1 },
72-
},
62+
definition: [
63+
{
64+
type: "classic",
65+
keys: { prop1: 1 },
66+
},
67+
],
7368
name: "my-index",
7469
},
7570
});
@@ -89,7 +84,7 @@ describeWithMongoDB("createIndex tool", (integration) => {
8984
arguments: {
9085
database: integration.randomDbName(),
9186
collection: "coll1",
92-
definition: { type: "classic", keys: { prop1: 1 } },
87+
definition: [{ type: "classic", keys: { prop1: 1 } }],
9388
},
9489
});
9590

@@ -107,7 +102,7 @@ describeWithMongoDB("createIndex tool", (integration) => {
107102
arguments: {
108103
database: integration.randomDbName(),
109104
collection: "coll1",
110-
definition: { type: "classic", keys: { prop1: 1 } },
105+
definition: [{ type: "classic", keys: { prop1: 1 } }],
111106
},
112107
});
113108

@@ -120,7 +115,7 @@ describeWithMongoDB("createIndex tool", (integration) => {
120115
arguments: {
121116
database: integration.randomDbName(),
122117
collection: "coll1",
123-
definition: { type: "classic", keys: { prop2: -1 } },
118+
definition: [{ type: "classic", keys: { prop2: -1 } }],
124119
},
125120
});
126121

@@ -141,7 +136,7 @@ describeWithMongoDB("createIndex tool", (integration) => {
141136
arguments: {
142137
database: integration.randomDbName(),
143138
collection: "coll1",
144-
definition: { type: "classic", keys: { prop1: 1 } },
139+
definition: [{ type: "classic", keys: { prop1: 1 } }],
145140
},
146141
});
147142

@@ -154,7 +149,7 @@ describeWithMongoDB("createIndex tool", (integration) => {
154149
arguments: {
155150
database: integration.randomDbName(),
156151
collection: "coll1",
157-
definition: { type: "classic", keys: { prop1: -1 } },
152+
definition: [{ type: "classic", keys: { prop1: -1 } }],
158153
},
159154
});
160155

@@ -175,7 +170,7 @@ describeWithMongoDB("createIndex tool", (integration) => {
175170
arguments: {
176171
database: integration.randomDbName(),
177172
collection: "coll1",
178-
definition: { type: "classic", keys: { prop1: 1 } },
173+
definition: [{ type: "classic", keys: { prop1: 1 } }],
179174
},
180175
});
181176

@@ -188,7 +183,7 @@ describeWithMongoDB("createIndex tool", (integration) => {
188183
arguments: {
189184
database: integration.randomDbName(),
190185
collection: "coll1",
191-
definition: { type: "classic", keys: { prop1: 1 } },
186+
definition: [{ type: "classic", keys: { prop1: 1 } }],
192187
},
193188
});
194189

@@ -216,7 +211,7 @@ describeWithMongoDB("createIndex tool", (integration) => {
216211
arguments: {
217212
database: integration.randomDbName(),
218213
collection: "coll1",
219-
definition: { type: "classic", keys: { prop1: direction } },
214+
definition: [{ type: "classic", keys: { prop1: direction } }],
220215
},
221216
});
222217

@@ -241,7 +236,7 @@ describeWithMongoDB("createIndex tool", (integration) => {
241236
args: {
242237
database: integration.randomDbName(),
243238
collection: "coll1",
244-
definition: { type: "classic", keys: { prop1: 1 } },
239+
definition: [{ type: "classic", keys: { prop1: 1 } }],
245240
},
246241
expectedResponse: `Created the index "prop1_1" on collection "coll1" in database "${integration.randomDbName()}"`,
247242
};

0 commit comments

Comments
 (0)