diff --git a/src/tools/mongodb/create/createIndex.ts b/src/tools/mongodb/create/createIndex.ts index d87b9df0..49c63381 100644 --- a/src/tools/mongodb/create/createIndex.ts +++ b/src/tools/mongodb/create/createIndex.ts @@ -9,8 +9,57 @@ export class CreateIndexTool extends MongoDBToolBase { protected description = "Create an index for a collection"; protected argsShape = { ...DbOperationArgs, - keys: z.object({}).catchall(z.custom()).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()).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"; @@ -18,16 +67,33 @@ export class CreateIndexTool extends MongoDBToolBase { protected async execute({ database, collection, - keys, name, + definition, }: ToolArgs): Promise { 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: [ diff --git a/tests/accuracy/createIndex.test.ts b/tests/accuracy/createIndex.test.ts index 08326ce3..d348e869 100644 --- a/tests/accuracy/createIndex.test.ts +++ b/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", + }, }, }, }, diff --git a/tests/integration/helpers.ts b/tests/integration/helpers.ts index d62354a8..921cc8f0 100644 --- a/tests/integration/helpers.ts +++ b/tests/integration/helpers.ts @@ -29,13 +29,22 @@ export const driverOptions = setupDriverConfig({ export const defaultDriverOptions: DriverOptions = { ...driverOptions }; -interface ParameterInfo { +interface Parameter { name: string; - type: string; description: string; required: boolean; } +interface SingleValueParameter extends Parameter { + type: string; +} + +interface AnyOfParameter extends Parameter { + anyOf: { type: string }[]; +} + +type ParameterInfo = SingleValueParameter | AnyOfParameter; + type ToolInfo = Awaited>["tools"][number]; export interface IntegrationTest { @@ -217,18 +226,38 @@ export function getParameters(tool: ToolInfo): ParameterInfo[] { return Object.entries(tool.inputSchema.properties) .sort((a, b) => a[0].localeCompare(b[0])) - .map(([key, value]) => { - expect(value).toHaveProperty("type"); + .map(([name, value]) => { expect(value).toHaveProperty("description"); - const typedValue = value as { type: string; description: string }; - expect(typeof typedValue.type).toBe("string"); - expect(typeof typedValue.description).toBe("string"); + const description = (value as { description: string }).description; + const required = (tool.inputSchema.required as string[])?.includes(name) ?? false; + expect(typeof description).toBe("string"); + + if (value && typeof value === "object" && "anyOf" in value) { + const typedOptions = new Array<{ type: string }>(); + for (const option of value.anyOf as { type: string }[]) { + expect(option).toHaveProperty("type"); + + typedOptions.push({ type: option.type }); + } + + return { + name, + anyOf: typedOptions, + description: description, + required, + }; + } + + expect(value).toHaveProperty("type"); + + const type = (value as { type: string }).type; + expect(typeof type).toBe("string"); return { - name: key, - type: typedValue.type, - description: typedValue.description, - required: (tool.inputSchema.required as string[])?.includes(key) ?? false, + name, + type, + description, + required, }; }); } diff --git a/tests/integration/tools/mongodb/create/createIndex.test.ts b/tests/integration/tools/mongodb/create/createIndex.test.ts index 3c789be8..156c33ff 100644 --- a/tests/integration/tools/mongodb/create/createIndex.test.ts +++ b/tests/integration/tools/mongodb/create/createIndex.test.ts @@ -8,15 +8,23 @@ import { expectDefined, } from "../../../helpers.js"; import type { IndexDirection } from "mongodb"; -import { expect, it } from "vitest"; +import { expect, it, describe } from "vitest"; describeWithMongoDB("createIndex tool", (integration) => { validateToolMetadata(integration, "create-index", "Create an index for a collection", [ ...databaseCollectionParameters, { - name: "keys", - type: "object", - description: "The index definition", + name: "definition", + anyOf: [ + { + type: "object", + }, + { + type: "object", + }, + ], + description: + "The index definition. Use 'classic' for standard indexes and 'vectorSearch' for vector search indexes", required: true, }, { @@ -29,174 +37,211 @@ describeWithMongoDB("createIndex tool", (integration) => { validateThrowsForInvalidArguments(integration, "create-index", [ {}, - { collection: "bar", database: 123, keys: { foo: 1 } }, - { collection: [], database: "test", keys: { foo: 1 } }, - { collection: "bar", database: "test", keys: { foo: 1 }, name: 123 }, - { collection: "bar", database: "test", keys: "foo", name: "my-index" }, + { collection: "bar", database: 123, definition: { type: "classic", keys: { foo: 1 } } }, + { collection: [], database: "test", definition: { type: "classic", keys: { foo: 1 } } }, + { collection: "bar", database: "test", definition: { type: "classic", keys: { foo: 1 } }, name: 123 }, + { collection: "bar", database: "test", definition: { type: "unknown", keys: { foo: 1 } }, name: "my-index" }, ]); - const validateIndex = async (collection: string, expected: { name: string; key: object }[]): Promise => { - const mongoClient = integration.mongoClient(); - const collections = await mongoClient.db(integration.randomDbName()).listCollections().toArray(); - expect(collections).toHaveLength(1); - expect(collections[0]?.name).toEqual("coll1"); - const indexes = await mongoClient.db(integration.randomDbName()).collection(collection).indexes(); - expect(indexes).toHaveLength(expected.length + 1); - expect(indexes[0]?.name).toEqual("_id_"); - for (const index of expected) { - const foundIndex = indexes.find((i) => i.name === index.name); - expectDefined(foundIndex); - expect(foundIndex.key).toEqual(index.key); - } - }; - - it("creates the namespace if necessary", async () => { - await integration.connectMcpClient(); - const response = await integration.mcpClient().callTool({ - name: "create-index", - arguments: { - database: integration.randomDbName(), - collection: "coll1", - keys: { prop1: 1 }, - name: "my-index", - }, - }); - - const content = getResponseContent(response.content); - expect(content).toEqual( - `Created the index "my-index" on collection "coll1" in database "${integration.randomDbName()}"` - ); - - await validateIndex("coll1", [{ name: "my-index", key: { prop1: 1 } }]); - }); + describe("with classic indexes", () => { + const validateIndex = async (collection: string, expected: { name: string; key: object }[]): Promise => { + const mongoClient = integration.mongoClient(); + const collections = await mongoClient.db(integration.randomDbName()).listCollections().toArray(); + expect(collections).toHaveLength(1); + expect(collections[0]?.name).toEqual("coll1"); + const indexes = await mongoClient.db(integration.randomDbName()).collection(collection).indexes(); + expect(indexes).toHaveLength(expected.length + 1); + expect(indexes[0]?.name).toEqual("_id_"); + for (const index of expected) { + const foundIndex = indexes.find((i) => i.name === index.name); + expectDefined(foundIndex); + expect(foundIndex.key).toEqual(index.key); + } + }; - it("generates a name if not provided", async () => { - await integration.connectMcpClient(); - const response = await integration.mcpClient().callTool({ - name: "create-index", - arguments: { database: integration.randomDbName(), collection: "coll1", keys: { prop1: 1 } }, - }); + it("creates the namespace if necessary", async () => { + await integration.connectMcpClient(); + const response = await integration.mcpClient().callTool({ + name: "create-index", + arguments: { + database: integration.randomDbName(), + collection: "coll1", + definition: { + type: "classic", + keys: { prop1: 1 }, + }, + name: "my-index", + }, + }); - const content = getResponseContent(response.content); - expect(content).toEqual( - `Created the index "prop1_1" on collection "coll1" in database "${integration.randomDbName()}"` - ); - await validateIndex("coll1", [{ name: "prop1_1", key: { prop1: 1 } }]); - }); + const content = getResponseContent(response.content); + expect(content).toEqual( + `Created the index "my-index" on collection "coll1" in database "${integration.randomDbName()}"` + ); - it("can create multiple indexes in the same collection", async () => { - await integration.connectMcpClient(); - let response = await integration.mcpClient().callTool({ - name: "create-index", - arguments: { database: integration.randomDbName(), collection: "coll1", keys: { prop1: 1 } }, + await validateIndex("coll1", [{ name: "my-index", key: { prop1: 1 } }]); }); - expect(getResponseContent(response.content)).toEqual( - `Created the index "prop1_1" on collection "coll1" in database "${integration.randomDbName()}"` - ); + it("generates a name if not provided", async () => { + await integration.connectMcpClient(); + const response = await integration.mcpClient().callTool({ + name: "create-index", + arguments: { + database: integration.randomDbName(), + collection: "coll1", + definition: { type: "classic", keys: { prop1: 1 } }, + }, + }); - response = await integration.mcpClient().callTool({ - name: "create-index", - arguments: { database: integration.randomDbName(), collection: "coll1", keys: { prop2: -1 } }, + const content = getResponseContent(response.content); + expect(content).toEqual( + `Created the index "prop1_1" on collection "coll1" in database "${integration.randomDbName()}"` + ); + await validateIndex("coll1", [{ name: "prop1_1", key: { prop1: 1 } }]); }); - expect(getResponseContent(response.content)).toEqual( - `Created the index "prop2_-1" on collection "coll1" in database "${integration.randomDbName()}"` - ); + it("can create multiple indexes in the same collection", async () => { + await integration.connectMcpClient(); + let response = await integration.mcpClient().callTool({ + name: "create-index", + arguments: { + database: integration.randomDbName(), + collection: "coll1", + definition: { type: "classic", keys: { prop1: 1 } }, + }, + }); - await validateIndex("coll1", [ - { name: "prop1_1", key: { prop1: 1 } }, - { name: "prop2_-1", key: { prop2: -1 } }, - ]); - }); + expect(getResponseContent(response.content)).toEqual( + `Created the index "prop1_1" on collection "coll1" in database "${integration.randomDbName()}"` + ); - it("can create multiple indexes on the same property", async () => { - await integration.connectMcpClient(); - let response = await integration.mcpClient().callTool({ - name: "create-index", - arguments: { database: integration.randomDbName(), collection: "coll1", keys: { prop1: 1 } }, - }); + response = await integration.mcpClient().callTool({ + name: "create-index", + arguments: { + database: integration.randomDbName(), + collection: "coll1", + definition: { type: "classic", keys: { prop2: -1 } }, + }, + }); - expect(getResponseContent(response.content)).toEqual( - `Created the index "prop1_1" on collection "coll1" in database "${integration.randomDbName()}"` - ); + expect(getResponseContent(response.content)).toEqual( + `Created the index "prop2_-1" on collection "coll1" in database "${integration.randomDbName()}"` + ); - response = await integration.mcpClient().callTool({ - name: "create-index", - arguments: { database: integration.randomDbName(), collection: "coll1", keys: { prop1: -1 } }, + await validateIndex("coll1", [ + { name: "prop1_1", key: { prop1: 1 } }, + { name: "prop2_-1", key: { prop2: -1 } }, + ]); }); - expect(getResponseContent(response.content)).toEqual( - `Created the index "prop1_-1" on collection "coll1" in database "${integration.randomDbName()}"` - ); + it("can create multiple indexes on the same property", async () => { + await integration.connectMcpClient(); + let response = await integration.mcpClient().callTool({ + name: "create-index", + arguments: { + database: integration.randomDbName(), + collection: "coll1", + definition: { type: "classic", keys: { prop1: 1 } }, + }, + }); - await validateIndex("coll1", [ - { name: "prop1_1", key: { prop1: 1 } }, - { name: "prop1_-1", key: { prop1: -1 } }, - ]); - }); + expect(getResponseContent(response.content)).toEqual( + `Created the index "prop1_1" on collection "coll1" in database "${integration.randomDbName()}"` + ); - it("doesn't duplicate indexes", async () => { - await integration.connectMcpClient(); - let response = await integration.mcpClient().callTool({ - name: "create-index", - arguments: { database: integration.randomDbName(), collection: "coll1", keys: { prop1: 1 } }, - }); + response = await integration.mcpClient().callTool({ + name: "create-index", + arguments: { + database: integration.randomDbName(), + collection: "coll1", + definition: { type: "classic", keys: { prop1: -1 } }, + }, + }); - expect(getResponseContent(response.content)).toEqual( - `Created the index "prop1_1" on collection "coll1" in database "${integration.randomDbName()}"` - ); + expect(getResponseContent(response.content)).toEqual( + `Created the index "prop1_-1" on collection "coll1" in database "${integration.randomDbName()}"` + ); - response = await integration.mcpClient().callTool({ - name: "create-index", - arguments: { database: integration.randomDbName(), collection: "coll1", keys: { prop1: 1 } }, + await validateIndex("coll1", [ + { name: "prop1_1", key: { prop1: 1 } }, + { name: "prop1_-1", key: { prop1: -1 } }, + ]); }); - expect(getResponseContent(response.content)).toEqual( - `Created the index "prop1_1" on collection "coll1" in database "${integration.randomDbName()}"` - ); + it("doesn't duplicate indexes", async () => { + await integration.connectMcpClient(); + let response = await integration.mcpClient().callTool({ + name: "create-index", + arguments: { + database: integration.randomDbName(), + collection: "coll1", + definition: { type: "classic", keys: { prop1: 1 } }, + }, + }); - await validateIndex("coll1", [{ name: "prop1_1", key: { prop1: 1 } }]); - }); + expect(getResponseContent(response.content)).toEqual( + `Created the index "prop1_1" on collection "coll1" in database "${integration.randomDbName()}"` + ); - const testCases: { name: string; direction: IndexDirection }[] = [ - { name: "descending", direction: -1 }, - { name: "ascending", direction: 1 }, - { name: "hashed", direction: "hashed" }, - { name: "text", direction: "text" }, - { name: "geoHaystack", direction: "2dsphere" }, - { name: "geo2d", direction: "2d" }, - ]; - - for (const { name, direction } of testCases) { - it(`creates ${name} index`, async () => { - await integration.connectMcpClient(); - const response = await integration.mcpClient().callTool({ + response = await integration.mcpClient().callTool({ name: "create-index", - arguments: { database: integration.randomDbName(), collection: "coll1", keys: { prop1: direction } }, + arguments: { + database: integration.randomDbName(), + collection: "coll1", + definition: { type: "classic", keys: { prop1: 1 } }, + }, }); expect(getResponseContent(response.content)).toEqual( - `Created the index "prop1_${direction}" on collection "coll1" in database "${integration.randomDbName()}"` + `Created the index "prop1_1" on collection "coll1" in database "${integration.randomDbName()}"` ); - let expectedKey: object = { prop1: direction }; - if (direction === "text") { - expectedKey = { - _fts: "text", - _ftsx: 1, - }; - } - await validateIndex("coll1", [{ name: `prop1_${direction}`, key: expectedKey }]); + await validateIndex("coll1", [{ name: "prop1_1", key: { prop1: 1 } }]); }); - } + + const testCases: { name: string; direction: IndexDirection }[] = [ + { name: "descending", direction: -1 }, + { name: "ascending", direction: 1 }, + { name: "hashed", direction: "hashed" }, + { name: "text", direction: "text" }, + { name: "geoHaystack", direction: "2dsphere" }, + { name: "geo2d", direction: "2d" }, + ]; + + for (const { name, direction } of testCases) { + it(`creates ${name} index`, async () => { + await integration.connectMcpClient(); + const response = await integration.mcpClient().callTool({ + name: "create-index", + arguments: { + database: integration.randomDbName(), + collection: "coll1", + definition: { type: "classic", keys: { prop1: direction } }, + }, + }); + + expect(getResponseContent(response.content)).toEqual( + `Created the index "prop1_${direction}" on collection "coll1" in database "${integration.randomDbName()}"` + ); + + let expectedKey: object = { prop1: direction }; + if (direction === "text") { + expectedKey = { + _fts: "text", + _ftsx: 1, + }; + } + await validateIndex("coll1", [{ name: `prop1_${direction}`, key: expectedKey }]); + }); + } + }); validateAutoConnectBehavior(integration, "create-index", () => { return { args: { database: integration.randomDbName(), collection: "coll1", - keys: { prop1: 1 }, + definition: { type: "classic", keys: { prop1: 1 } }, }, expectedResponse: `Created the index "prop1_1" on collection "coll1" in database "${integration.randomDbName()}"`, };