Skip to content

Commit 88acf38

Browse files
nirinchevnirinchev
committed
add primitive feature flagging
1 parent 69aa6d0 commit 88acf38

File tree

6 files changed

+440
-307
lines changed

6 files changed

+440
-307
lines changed

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

Lines changed: 100 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -1,79 +1,90 @@
1+
import type { ZodDiscriminatedUnionOption } from "zod";
12
import { z } from "zod";
23
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
34
import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
4-
import type { ToolArgs, OperationType } from "../../tool.js";
5+
import { type ToolArgs, type OperationType, type ToolConstructorParams, FeatureFlags } from "../../tool.js";
56
import type { IndexDirection } from "mongodb";
67

8+
const classicIndexDefinition = z.object({
9+
type: z.literal("classic"),
10+
keys: z.object({}).catchall(z.custom<IndexDirection>()).describe("The index definition"),
11+
});
12+
13+
const vectorSearchIndexDefinition = z.object({
14+
type: z.literal("vectorSearch"),
15+
fields: z
16+
.array(
17+
z.discriminatedUnion("type", [
18+
z.object({
19+
type: z.literal("filter"),
20+
path: z
21+
.string()
22+
.describe(
23+
"Name of the field to index. For nested fields, use dot notation to specify path to embedded fields"
24+
),
25+
}),
26+
z.object({
27+
type: z.literal("vector"),
28+
path: z
29+
.string()
30+
.describe(
31+
"Name of the field to index. For nested fields, use dot notation to specify path to embedded fields"
32+
),
33+
numDimensions: z
34+
.number()
35+
.min(1)
36+
.max(8192)
37+
.describe(
38+
"Number of vector dimensions that MongoDB Vector Search enforces at index-time and query-time"
39+
),
40+
similarity: z
41+
.enum(["cosine", "euclidean", "dotProduct"])
42+
.default("cosine")
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+
)
56+
.nonempty()
57+
.refine((fields) => fields.some((f) => f.type === "vector"), {
58+
message: "At least one vector field must be defined",
59+
})
60+
.describe(
61+
"Definitions for the vector and filter fields to index, one definition per document. You must specify `vector` for fields that contain vector embeddings and `filter` for additional fields to filter on. At least one vector-type field definition is required."
62+
),
63+
});
64+
765
export class CreateIndexTool extends MongoDBToolBase {
866
public name = "create-index";
967
protected description = "Create an index for a collection";
10-
protected argsShape = {
11-
...DbOperationArgs,
12-
name: z.string().optional().describe("The name of the index"),
13-
definition: z
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.discriminatedUnion("type", [
25-
z.object({
26-
type: z.literal("filter"),
27-
path: z
28-
.string()
29-
.describe(
30-
"Name of the field to index. For nested fields, use dot notation to specify path to embedded fields"
31-
),
32-
}),
33-
z.object({
34-
type: z.literal("vector"),
35-
path: z
36-
.string()
37-
.describe(
38-
"Name of the field to index. For nested fields, use dot notation to specify path to embedded fields"
39-
),
40-
numDimensions: z
41-
.number()
42-
.min(1)
43-
.max(8192)
44-
.describe(
45-
"Number of vector dimensions that MongoDB Vector Search enforces at index-time and query-time"
46-
),
47-
similarity: z
48-
.enum(["cosine", "euclidean", "dotProduct"])
49-
.default("cosine")
50-
.describe(
51-
"Vector similarity function to use to search for top K-nearest neighbors. You can set this field only for vector-type fields."
52-
),
53-
quantization: z
54-
.enum(["none", "scalar", "binary"])
55-
.optional()
56-
.default("none")
57-
.describe(
58-
"Type of automatic vector quantization for your vectors. Use this setting only if your embeddings are float or double vectors."
59-
),
60-
}),
61-
])
62-
)
63-
.nonempty()
64-
.refine((fields) => fields.some((f) => f.type === "vector"), {
65-
message: "At least one vector field must be defined",
66-
})
67-
.describe(
68-
"Definitions for the vector and filter fields to index, one definition per document. You must specify `vector` for fields that contain vector embeddings and `filter` for additional fields to filter on. At least one vector-type field definition is required."
69-
),
70-
}),
71-
])
72-
)
73-
.describe(
74-
"The index definition. Use 'classic' for standard indexes and 'vectorSearch' for vector search indexes"
75-
),
76-
};
68+
protected argsShape;
69+
70+
constructor(params: ToolConstructorParams) {
71+
super(params);
72+
73+
const additionalIndexDefinitions: ZodDiscriminatedUnionOption<"type">[] = [];
74+
if (this.isFeatureFlagEnabled(FeatureFlags.VectorSearch)) {
75+
additionalIndexDefinitions.push(vectorSearchIndexDefinition);
76+
}
77+
78+
this.argsShape = {
79+
...DbOperationArgs,
80+
name: z.string().optional().describe("The name of the index"),
81+
definition: z
82+
.array(z.discriminatedUnion("type", [classicIndexDefinition, ...additionalIndexDefinitions]))
83+
.describe(
84+
"The index definition. Use 'classic' for standard indexes and 'vectorSearch' for vector search indexes"
85+
),
86+
};
87+
}
7788

7889
public operationType: OperationType = "create";
7990

@@ -92,24 +103,30 @@ export class CreateIndexTool extends MongoDBToolBase {
92103

93104
switch (definition.type) {
94105
case "classic":
95-
indexes = await provider.createIndexes(database, collection, [
96-
{
97-
key: definition.keys,
98-
name,
99-
},
100-
]);
101-
106+
{
107+
const typedDefinition = definition as z.infer<typeof classicIndexDefinition>;
108+
indexes = await provider.createIndexes(database, collection, [
109+
{
110+
key: typedDefinition.keys,
111+
name,
112+
},
113+
]);
114+
}
102115
break;
103116
case "vectorSearch":
104-
indexes = await provider.createSearchIndexes(database, collection, [
105-
{
106-
name,
107-
definition: {
108-
fields: definition.fields,
117+
{
118+
const typedDefinition = definition as z.infer<typeof vectorSearchIndexDefinition>;
119+
indexes = await provider.createSearchIndexes(database, collection, [
120+
{
121+
name,
122+
definition: {
123+
fields: typedDefinition.fields,
124+
},
125+
type: "vectorSearch",
109126
},
110-
type: "vectorSearch",
111-
},
112-
]);
127+
]);
128+
}
129+
113130
break;
114131
}
115132

‎src/tools/tool.ts‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,10 @@ export type ToolCallbackArgs<Args extends ZodRawShape> = Parameters<ToolCallback
1515

1616
export type ToolExecutionContext<Args extends ZodRawShape = ZodRawShape> = Parameters<ToolCallback<Args>>[1];
1717

18+
export const enum FeatureFlags {
19+
VectorSearch = "vectorSearch",
20+
}
21+
1822
/**
1923
* The type of operation the tool performs. This is used when evaluating if a tool is allowed to run based on
2024
* the config's `disabledTools` and `readOnly` settings.
@@ -314,6 +318,16 @@ export abstract class ToolBase {
314318

315319
this.telemetry.emitEvents([event]);
316320
}
321+
322+
// TODO: Move this to a separate file
323+
protected isFeatureFlagEnabled(flag: FeatureFlags): boolean {
324+
switch (flag) {
325+
case FeatureFlags.VectorSearch:
326+
return this.config.voyageApiKey !== "";
327+
default:
328+
return false;
329+
}
330+
}
317331
}
318332

319333
/**

‎tests/accuracy/createIndex.test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
import { describeAccuracyTests } from "./sdk/describeAccuracyTests.js";
22
import { Matcher } from "./sdk/matcher.js";
33

4+
process.env.MDB_VOYAGE_API_KEY = "valid-key";
5+
46
describeAccuracyTests([
57
{
68
prompt: "Create an index that covers the following query on 'mflix.movies' namespace - { \"release_year\": 1992 }",

‎tests/accuracy/sdk/accuracyTestingClient.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,14 +82,16 @@ export class AccuracyTestingClient {
8282
static async initializeClient(
8383
mdbConnectionString: string,
8484
atlasApiClientId?: string,
85-
atlasApiClientSecret?: string
85+
atlasApiClientSecret?: string,
86+
voyageApiKey?: string
8687
): Promise<AccuracyTestingClient> {
8788
const args = [
8889
MCP_SERVER_CLI_SCRIPT,
8990
"--connectionString",
9091
mdbConnectionString,
9192
...(atlasApiClientId ? ["--apiClientId", atlasApiClientId] : []),
9293
...(atlasApiClientSecret ? ["--apiClientSecret", atlasApiClientSecret] : []),
94+
...(voyageApiKey ? ["--voyageApiKey", voyageApiKey] : []),
9395
];
9496

9597
const clientTransport = new StdioClientTransport({

‎tests/accuracy/sdk/describeAccuracyTests.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ export function describeAccuracyTests(accuracyTestConfigs: AccuracyTestConfig[])
6868

6969
const atlasApiClientId = process.env.MDB_MCP_API_CLIENT_ID;
7070
const atlasApiClientSecret = process.env.MDB_MCP_API_CLIENT_SECRET;
71+
const voyageApiKey = process.env.MDB_VOYAGE_API_KEY;
7172

7273
let commitSHA: string;
7374
let accuracyResultStorage: AccuracyResultStorage;
@@ -85,7 +86,8 @@ export function describeAccuracyTests(accuracyTestConfigs: AccuracyTestConfig[])
8586
testMCPClient = await AccuracyTestingClient.initializeClient(
8687
mdbIntegration.connectionString(),
8788
atlasApiClientId,
88-
atlasApiClientSecret
89+
atlasApiClientSecret,
90+
voyageApiKey
8991
);
9092
agent = getVercelToolCallingAgent();
9193
});

0 commit comments

Comments
 (0)