From f7282c5f429715dbeffc6a41ba9fad1585486bf1 Mon Sep 17 00:00:00 2001 From: Edgar Wang Date: Wed, 14 May 2025 09:51:48 -0400 Subject: [PATCH 1/5] save save save vector index args save save save saving vector search --- src/tools/mongodb/create/createVectorIndex.ts | 55 +++++++++++++++++++ src/tools/mongodb/mongodbTool.ts | 47 +++++++++++++++- src/tools/mongodb/tools.ts | 5 ++ src/tools/mongodb/update/updateVectorIndex.ts | 52 ++++++++++++++++++ 4 files changed, 158 insertions(+), 1 deletion(-) create mode 100644 src/tools/mongodb/create/createVectorIndex.ts create mode 100644 src/tools/mongodb/update/updateVectorIndex.ts diff --git a/src/tools/mongodb/create/createVectorIndex.ts b/src/tools/mongodb/create/createVectorIndex.ts new file mode 100644 index 000000000..4d24768e4 --- /dev/null +++ b/src/tools/mongodb/create/createVectorIndex.ts @@ -0,0 +1,55 @@ +import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; +import { + DbOperationArgs, + MongoDBToolBase, + VectorFieldType, + VectorIndexArgs, +} from "../mongodbTool.js"; +import { OperationType, ToolArgs } from "../../tool.js"; + +const VECTOR_INDEX_TYPE = "vectorSearch"; +export class CreateVectorIndexTool extends MongoDBToolBase { + protected name = "create-vector-index"; + protected description = "Create an Atlas Search vector for a collection"; + protected argsShape = { + ...DbOperationArgs, + name: VectorIndexArgs.name, + vectorDefinition: VectorIndexArgs.vectorDefinition, + filterFields: VectorIndexArgs.filterFields, + }; + + protected operationType: OperationType = "create"; + + protected async execute({ + database, + collection, + name, + vectorDefinition, + filterFields, + }: ToolArgs): Promise { + const provider = await this.ensureConnected(); + + const typedVectorField = { ...vectorDefinition, type: VectorFieldType.VECTOR }; + const typedFilterFields = + filterFields?.map((v) => ({ + ...v, + type: VectorFieldType.FILTER, + })) || []; + const indexes = await provider.createSearchIndexes(database, collection, [ + { + name, + type: VECTOR_INDEX_TYPE, + definition: { fields: [typedVectorField, ...typedFilterFields] }, + }, + ]); + + return { + content: [ + { + text: `Created the vector index ${indexes[0]} on collection "${collection}" in database "${database}"`, + type: "text", + }, + ], + }; + } +} diff --git a/src/tools/mongodb/mongodbTool.ts b/src/tools/mongodb/mongodbTool.ts index 2ef1aee0e..fccc938f5 100644 --- a/src/tools/mongodb/mongodbTool.ts +++ b/src/tools/mongodb/mongodbTool.ts @@ -1,5 +1,5 @@ import { z } from "zod"; -import { ToolArgs, ToolBase, ToolCategory, TelemetryToolMetadata } from "../tool.js"; +import { TelemetryToolMetadata, ToolArgs, ToolBase, ToolCategory } from "../tool.js"; import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver"; import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { ErrorCodes, MongoDBError } from "../../errors.js"; @@ -10,6 +10,51 @@ export const DbOperationArgs = { collection: z.string().describe("Collection name"), }; +export enum VectorFieldType { + VECTOR = "vector", + FILTER = "filter", +} +export const VectorIndexArgs = { + name: z.string().describe("The name of the index"), + vectorDefinition: z + .object({ + 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() + .int() + .min(1) + .max(8192) + .describe("Number of vector dimensions to enforce at index-time and query-time."), + similarity: z + .enum(["euclidean", "cosine", "dotProduct"]) + .describe("Vector similarity function to use to search for top K-nearest neighbors."), + quantization: z + .enum(["none", "scalar", "binary"]) + .default("none") + .optional() + .describe( + "Automatic vector quantization. Use this setting only if your embeddings are float or double vectors." + ), + }) + .describe("The vector index definition."), + filterFields: z + .array( + z.object({ + path: z + .string() + .describe( + "Name of the field to filter by. For nested fields, use dot notation to specify path to embedded fields." + ), + }) + ) + .optional() + .describe("Additional indexed fields that pre-filter data."), +}; + export abstract class MongoDBToolBase extends ToolBase { protected category: ToolCategory = "mongodb"; diff --git a/src/tools/mongodb/tools.ts b/src/tools/mongodb/tools.ts index d64d53ea7..03724783c 100644 --- a/src/tools/mongodb/tools.ts +++ b/src/tools/mongodb/tools.ts @@ -18,6 +18,8 @@ import { DropCollectionTool } from "./delete/dropCollection.js"; import { ExplainTool } from "./metadata/explain.js"; import { CreateCollectionTool } from "./create/createCollection.js"; import { LogsTool } from "./metadata/logs.js"; +import { CreateVectorIndexTool } from "./create/createVectorIndex.js"; +import { UpdateVectorIndexTool } from "./update/updateVectorIndex.js"; export const MongoDbTools = [ ConnectTool, @@ -40,4 +42,7 @@ export const MongoDbTools = [ ExplainTool, CreateCollectionTool, LogsTool, + + CreateVectorIndexTool, + UpdateVectorIndexTool, ]; diff --git a/src/tools/mongodb/update/updateVectorIndex.ts b/src/tools/mongodb/update/updateVectorIndex.ts new file mode 100644 index 000000000..5e12d6873 --- /dev/null +++ b/src/tools/mongodb/update/updateVectorIndex.ts @@ -0,0 +1,52 @@ +import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; +import { + DbOperationArgs, + MongoDBToolBase, + VectorFieldType, + VectorIndexArgs, +} from "../mongodbTool.js"; +import { OperationType, ToolArgs } from "../../tool.js"; + +export class UpdateVectorIndexTool extends MongoDBToolBase { + protected name = "update-vector-index"; + protected description = "Updates an Atlas Search vector for a collection"; + protected argsShape = { + ...DbOperationArgs, + name: VectorIndexArgs.name, + vectorDefinition: VectorIndexArgs.vectorDefinition, + filterFields: VectorIndexArgs.filterFields, + }; + + protected operationType: OperationType = "create"; + + protected async execute({ + database, + collection, + name, + vectorDefinition, + filterFields, + }: ToolArgs): Promise { + const provider = await this.ensureConnected(); + + const typedVectorField = { ...vectorDefinition, type: VectorFieldType.VECTOR }; + const typedFilterFields = + filterFields?.map((v) => ({ + ...v, + type: VectorFieldType.FILTER, + })) || []; + // @ts-expect-error: Interface expects a SearchIndexDefinition {definition: {fields}}. However, + // passing fields at the root level is necessary for the call to succeed. + await provider.updateSearchIndex(database, collection, name, { + fields: [typedVectorField, ...typedFilterFields], + }); + + return { + content: [ + { + text: `Successfully updated vector index "${name}" on collection "${collection}" in database "${database}"`, + type: "text", + }, + ], + }; + } +} \ No newline at end of file From 4083cda4d431da3df3db0a8c005f51e17621a98a Mon Sep 17 00:00:00 2001 From: Edgar Wang Date: Wed, 14 May 2025 11:57:11 -0400 Subject: [PATCH 2/5] extract shared vector field creation fn --- src/tools/mongodb/create/createVectorIndex.ts | 10 ++-------- src/tools/mongodb/mongodbTool.ts | 11 +++++++++++ src/tools/mongodb/update/updateVectorIndex.ts | 10 ++-------- 3 files changed, 15 insertions(+), 16 deletions(-) diff --git a/src/tools/mongodb/create/createVectorIndex.ts b/src/tools/mongodb/create/createVectorIndex.ts index 4d24768e4..21c2e813b 100644 --- a/src/tools/mongodb/create/createVectorIndex.ts +++ b/src/tools/mongodb/create/createVectorIndex.ts @@ -1,8 +1,8 @@ import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { + buildVectorFields, DbOperationArgs, MongoDBToolBase, - VectorFieldType, VectorIndexArgs, } from "../mongodbTool.js"; import { OperationType, ToolArgs } from "../../tool.js"; @@ -29,17 +29,11 @@ export class CreateVectorIndexTool extends MongoDBToolBase { }: ToolArgs): Promise { const provider = await this.ensureConnected(); - const typedVectorField = { ...vectorDefinition, type: VectorFieldType.VECTOR }; - const typedFilterFields = - filterFields?.map((v) => ({ - ...v, - type: VectorFieldType.FILTER, - })) || []; const indexes = await provider.createSearchIndexes(database, collection, [ { name, type: VECTOR_INDEX_TYPE, - definition: { fields: [typedVectorField, ...typedFilterFields] }, + definition: { fields: buildVectorFields(vectorDefinition, filterFields) }, }, ]); diff --git a/src/tools/mongodb/mongodbTool.ts b/src/tools/mongodb/mongodbTool.ts index fccc938f5..7be43c54a 100644 --- a/src/tools/mongodb/mongodbTool.ts +++ b/src/tools/mongodb/mongodbTool.ts @@ -55,6 +55,17 @@ export const VectorIndexArgs = { .describe("Additional indexed fields that pre-filter data."), }; +type VectorDefinitionType = z.infer; +type FilterFieldsType = z.infer; +export function buildVectorFields(vectorDefinition: VectorDefinitionType, filterFields: FilterFieldsType): object[] { + const typedVectorField = { ...vectorDefinition, type: VectorFieldType.VECTOR }; + const typedFilterFields = (filterFields ?? []).map((f) => ({ + ...f, + type: VectorFieldType.FILTER, + })); + return [typedVectorField, ...typedFilterFields]; +} + export abstract class MongoDBToolBase extends ToolBase { protected category: ToolCategory = "mongodb"; diff --git a/src/tools/mongodb/update/updateVectorIndex.ts b/src/tools/mongodb/update/updateVectorIndex.ts index 5e12d6873..bafaddeb0 100644 --- a/src/tools/mongodb/update/updateVectorIndex.ts +++ b/src/tools/mongodb/update/updateVectorIndex.ts @@ -1,8 +1,8 @@ import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { + buildVectorFields, DbOperationArgs, MongoDBToolBase, - VectorFieldType, VectorIndexArgs, } from "../mongodbTool.js"; import { OperationType, ToolArgs } from "../../tool.js"; @@ -28,16 +28,10 @@ export class UpdateVectorIndexTool extends MongoDBToolBase { }: ToolArgs): Promise { const provider = await this.ensureConnected(); - const typedVectorField = { ...vectorDefinition, type: VectorFieldType.VECTOR }; - const typedFilterFields = - filterFields?.map((v) => ({ - ...v, - type: VectorFieldType.FILTER, - })) || []; // @ts-expect-error: Interface expects a SearchIndexDefinition {definition: {fields}}. However, // passing fields at the root level is necessary for the call to succeed. await provider.updateSearchIndex(database, collection, name, { - fields: [typedVectorField, ...typedFilterFields], + fields: buildVectorFields(vectorDefinition, filterFields), }); return { From 4b81caf593d82c35f4af6df0c2e99038762ef420 Mon Sep 17 00:00:00 2001 From: Edgar Wang Date: Wed, 14 May 2025 12:08:35 -0400 Subject: [PATCH 3/5] eslint --- src/tools/mongodb/create/createVectorIndex.ts | 7 +------ src/tools/mongodb/update/updateVectorIndex.ts | 9 ++------- 2 files changed, 3 insertions(+), 13 deletions(-) diff --git a/src/tools/mongodb/create/createVectorIndex.ts b/src/tools/mongodb/create/createVectorIndex.ts index 21c2e813b..045950e3e 100644 --- a/src/tools/mongodb/create/createVectorIndex.ts +++ b/src/tools/mongodb/create/createVectorIndex.ts @@ -1,10 +1,5 @@ import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; -import { - buildVectorFields, - DbOperationArgs, - MongoDBToolBase, - VectorIndexArgs, -} from "../mongodbTool.js"; +import { buildVectorFields, DbOperationArgs, MongoDBToolBase, VectorIndexArgs } from "../mongodbTool.js"; import { OperationType, ToolArgs } from "../../tool.js"; const VECTOR_INDEX_TYPE = "vectorSearch"; diff --git a/src/tools/mongodb/update/updateVectorIndex.ts b/src/tools/mongodb/update/updateVectorIndex.ts index bafaddeb0..476cfc5e1 100644 --- a/src/tools/mongodb/update/updateVectorIndex.ts +++ b/src/tools/mongodb/update/updateVectorIndex.ts @@ -1,10 +1,5 @@ import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; -import { - buildVectorFields, - DbOperationArgs, - MongoDBToolBase, - VectorIndexArgs, -} from "../mongodbTool.js"; +import { buildVectorFields, DbOperationArgs, MongoDBToolBase, VectorIndexArgs } from "../mongodbTool.js"; import { OperationType, ToolArgs } from "../../tool.js"; export class UpdateVectorIndexTool extends MongoDBToolBase { @@ -43,4 +38,4 @@ export class UpdateVectorIndexTool extends MongoDBToolBase { ], }; } -} \ No newline at end of file +} From 2b7788795d966be031d375df98ad5f979452bc2c Mon Sep 17 00:00:00 2001 From: Edgar Wang Date: Wed, 14 May 2025 15:33:37 -0400 Subject: [PATCH 4/5] add min field length --- src/tools/mongodb/mongodbTool.ts | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/tools/mongodb/mongodbTool.ts b/src/tools/mongodb/mongodbTool.ts index 7be43c54a..8d4c41fe0 100644 --- a/src/tools/mongodb/mongodbTool.ts +++ b/src/tools/mongodb/mongodbTool.ts @@ -20,6 +20,7 @@ export const VectorIndexArgs = { .object({ path: z .string() + .min(1) .describe( "Name of the field to index. For nested fields, use dot notation to specify path to embedded fields." ), @@ -46,6 +47,7 @@ export const VectorIndexArgs = { z.object({ path: z .string() + .min(1) .describe( "Name of the field to filter by. For nested fields, use dot notation to specify path to embedded fields." ), From fe66ee771ce36dfcffa442a97f82896bf7834921 Mon Sep 17 00:00:00 2001 From: Edgar Wang Date: Wed, 14 May 2025 16:25:57 -0400 Subject: [PATCH 5/5] update description --- src/tools/mongodb/create/createVectorIndex.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tools/mongodb/create/createVectorIndex.ts b/src/tools/mongodb/create/createVectorIndex.ts index 045950e3e..55ac2c8dd 100644 --- a/src/tools/mongodb/create/createVectorIndex.ts +++ b/src/tools/mongodb/create/createVectorIndex.ts @@ -5,7 +5,7 @@ import { OperationType, ToolArgs } from "../../tool.js"; const VECTOR_INDEX_TYPE = "vectorSearch"; export class CreateVectorIndexTool extends MongoDBToolBase { protected name = "create-vector-index"; - protected description = "Create an Atlas Search vector for a collection"; + protected description = "Create an Atlas Vector Search Index for a collection."; protected argsShape = { ...DbOperationArgs, name: VectorIndexArgs.name,