feat: add ability to create vector search indexes

Author: nirinchev

src/tools/mongodb/create/createIndex.ts

@@ -9,25 +9,91 @@ export class CreateIndexTool extends MongoDBToolBase {
     protected description = "Create an index for a collection";
     protected argsShape = {
         ...DbOperationArgs,
-        keys: z.object({}).catchall(z.custom<IndexDirection>()).describe("The index definition"),
         name: z.string().optional().describe("The name of the index"),
+        definition: z
+            .discriminatedUnion("type", [
+                z.object({
+                    type: z.literal("classic"),
+                    keys: z.object({}).catchall(z.custom<IndexDirection>()).describe("The index definition"),
+                }),
+                z.object({
+                    type: z.literal("vectorSearch"),
+                    fields: z
+                        .array(
+                            z.object({
+                                type: z
+                                    .enum(["vector", "filter"])
+                                    .describe(
+                                        "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."
+                                    ),
+                                path: z
+                                    .string()
+                                    .describe(
+                                        "Name of the field to index. For nested fields, use dot notation to specify path to embedded fields"
+                                    ),
+                                numDimensions: z
+                                    .number()
+                                    .min(1)
+                                    .max(8192)
+                                    .describe(
+                                        "Number of vector dimensions that MongoDB Vector Search enforces at index-time and query-time"
+                                    ),
+                                similarity: z
+                                    .enum(["cosine", "euclidean", "dotProduct"])
+                                    .describe(
+                                        "Vector similarity function to use to search for top K-nearest neighbors. You can set this field only for vector-type fields."
+                                    ),
+                                quantization: z
+                                    .enum(["none", "scalar", "binary"])
+                                    .optional()
+                                    .default("none")
+                                    .describe(
+                                        "Type of automatic vector quantization for your vectors. Use this setting only if your embeddings are float or double vectors."
+                                    ),
+                            })
+                        )
+                        .describe(
+                            "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."
+                        ),
+                }),
+            ])
+            .describe(
+                "The index definition. Use 'classic' for standard indexes and 'vectorSearch' for vector search indexes"
+            ),
     };
 
     public operationType: OperationType = "create";
 
     protected async execute({
         database,
         collection,
-        keys,
         name,
+        definition,
     }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
         const provider = await this.ensureConnected();
-        const indexes = await provider.createIndexes(database, collection, [
-            {
-                key: keys,
-                name,
-            },
-        ]);
+        let indexes: string[] = [];
+        switch (definition.type) {
+            case "classic":
+                indexes = await provider.createIndexes(database, collection, [
+                    {
+                        key: definition.keys,
+                        name,
+                    },
+                ]);
+
+                break;
+            case "vectorSearch":
+                indexes = await provider.createSearchIndexes(database, collection, [
+                    {
+                        name,
+                        definition: {
+                            fields: definition.fields,
+                        },
+                        type: "vectorSearch",
+                    },
+                ]);
+                break;
+        }
 
         return {
             content: [

tests/accuracy/createIndex.test.ts

@@ -11,8 +11,11 @@ describeAccuracyTests([
                     database: "mflix",
                     collection: "movies",
                     name: Matcher.anyOf(Matcher.undefined, Matcher.string()),
-                    keys: {
-                        release_year: 1,
+                    definition: {
+                        type: "classic",
+                        keys: {
+                            release_year: 1,
+                        },
                     },
                 },
             },
@@ -27,8 +30,11 @@ describeAccuracyTests([
                     database: "mflix",
                     collection: "movies",
                     name: Matcher.anyOf(Matcher.undefined, Matcher.string()),
-                    keys: {
-                        title: "text",
+                    definition: {
+                        type: "classic",
+                        keys: {
+                            title: "text",
+                        },
                     },
                 },
             },