add numPartitions

Author: nirinchev

.vscode/launch.json

@@ -4,6 +4,13 @@
   // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
   "version": "0.2.0",
   "configurations": [
+    {
+      "name": "Attach by Process ID",
+      "processId": "${command:PickProcess}",
+      "request": "attach",
+      "skipFiles": ["<node_internals>/**"],
+      "type": "node"
+    },
     {
       "type": "node",
       "request": "launch",

package.json

@@ -33,7 +33,7 @@
     "dist"
   ],
   "scripts": {
-    "start": "node dist/index.js --transport http --loggers stderr mcp",
+    "start": "node dist/index.js --transport http --loggers stderr mcp --previewFeatures vectorSearch",
     "start:stdio": "node dist/index.js --transport stdio --loggers stderr mcp",
     "prepare": "npm run build",
     "build:clean": "rm -rf dist",

src/tools/args.ts

@@ -18,6 +18,17 @@ export const ALLOWED_CLUSTER_NAME_CHARACTERS_ERROR =
 const ALLOWED_PROJECT_NAME_CHARACTERS_REGEX = /^[a-zA-Z0-9\s()@&+:._',-]+$/;
 export const ALLOWED_PROJECT_NAME_CHARACTERS_ERROR =
     "Project names can't be longer than 64 characters and can only contain letters, numbers, spaces, and the following symbols: ( ) @ & + : . _ - ' ,";
+
+// Zod does not undestand JS boxed numbers (like Int32) as integer literals,
+// so we preprocess them to unwrap them so Zod understands them.
+function unboxNumber(v: unknown): number {
+    if (v && typeof v === "object" && typeof v.valueOf === "function") {
+        const n = Number(v.valueOf());
+        if (!Number.isNaN(n)) return n;
+    }
+    return v as number;
+}
+
 export const CommonArgs = {
     string: (): ZodString => z.string().regex(NO_UNICODE_REGEX, NO_UNICODE_ERROR),
 
@@ -27,6 +38,11 @@ export const CommonArgs = {
             .min(1, `${fieldName} is required`)
             .length(24, `${fieldName} must be exactly 24 characters`)
             .regex(/^[0-9a-fA-F]+$/, `${fieldName} must contain only hexadecimal characters`),
+    numberEnum: <Options extends Readonly<[z.ZodTypeAny, z.ZodTypeAny, ...z.ZodTypeAny[]]>>(
+        values: Options
+    ): z.ZodEffects<z.ZodUnion<Options>> => {
+        return z.preprocess(unboxNumber, z.union(values));
+    },
 };
 
 export const AtlasArgs = {

src/tools/mongodb/create/createIndex.ts

@@ -5,6 +5,7 @@ import { type ToolArgs, type OperationType } from "../../tool.js";
 import type { IndexDirection } from "mongodb";
 import { quantizationEnum } from "../../../common/search/vectorSearchEmbeddingsManager.js";
 import { similarityValues } from "../../../common/schemas.js";
+import { CommonArgs } from "../../args.js";
 
 export class CreateIndexTool extends MongoDBToolBase {
     private vectorSearchIndexDefinition = z
@@ -121,6 +122,11 @@ export class CreateIndexTool extends MongoDBToolBase {
                 .describe(
                     "Document describing the index to create. Either `dynamic` must be `true` and `fields` empty or `dynamic` must be `false` and at least one field must be defined in the `fields` document."
                 ),
+            numPartitions: CommonArgs.numberEnum([z.literal(1), z.literal(2), z.literal(4)])
+                .default(1)
+                .describe(
+                    "Specifies the number of sub-indexes to create if the document count exceeds two billion. If omitted, defaults to 1."
+                ),
         })
         .describe("Definition for an Atlas Search (lexical) index.");
 
@@ -204,6 +210,7 @@ export class CreateIndexTool extends MongoDBToolBase {
                             definition: {
                                 mappings: definition.mappings,
                                 analyzer: definition.analyzer,
+                                numPartitions: definition.numPartitions,
                             },
                             type: "search",
                         },

src/tools/mongodb/mongodbSchemas.ts

@@ -1,26 +1,18 @@
 import z from "zod";
-import { zEJSON } from "../args.js";
+import { CommonArgs, zEJSON } from "../args.js";
 
 export const zVoyageModels = z
     .enum(["voyage-3-large", "voyage-3.5", "voyage-3.5-lite", "voyage-code-3"])
     .default("voyage-3-large");
 
-// Zod does not undestand JS boxed numbers (like Int32) as integer literals,
-// so we preprocess them to unwrap them so Zod understands them.
-function unboxNumber(v: unknown): number {
-    if (v && typeof v === "object" && typeof v.valueOf === "function") {
-        const n = Number(v.valueOf());
-        if (!Number.isNaN(n)) return n;
-    }
-    return v as number;
-}
-
 export const zVoyageEmbeddingParameters = z.object({
-    outputDimension: z
-        .preprocess(
-            unboxNumber,
-            z.union([z.literal(256), z.literal(512), z.literal(1024), z.literal(2048), z.literal(4096)])
-        )
+    outputDimension: CommonArgs.numberEnum([
+        z.literal(256),
+        z.literal(512),
+        z.literal(1024),
+        z.literal(2048),
+        z.literal(4096),
+    ])
         .optional()
         .default(1024),
     outputDtype: z.enum(["float", "int8", "uint8", "binary", "ubinary"]).optional().default("float"),

tests/accuracy/aggregate.test.ts

@@ -421,4 +421,26 @@ describeAccuracyTests([
             },
         },
     },
+    {
+        prompt: "Run a search query on mflix.movies to find all movies that mention 'space travel' in the plot.",
+        expectedToolCalls: [
+            {
+                toolName: "aggregate",
+                parameters: {
+                    database: "mflix",
+                    collection: "movies",
+                    pipeline: [
+                        {
+                            $search: {
+                                index: "default",
+                                text: {
+                                    query: "space travel",
+                                },
+                            },
+                        },
+                    ],
+                },
+            },
+        ],
+    },
 ]);

tests/accuracy/createIndex.test.ts

@@ -173,6 +173,7 @@ describeAccuracyTests(
                                 mappings: {
                                     dynamic: true,
                                 },
+                                numPartitions: Matcher.anyOf(Matcher.undefined, Matcher.number()),
                             },
                         ],
                     },
@@ -204,6 +205,7 @@ describeAccuracyTests(
                                         },
                                     },
                                 },
+                                numPartitions: Matcher.anyOf(Matcher.undefined, Matcher.number()),
                             },
                         ],
                     },