chore: tweak tools to support autoembed index

Author: himanshusinghs

src/tools/mongodb/create/createIndex.ts

@@ -7,62 +7,96 @@ import { quantizationEnum } from "../../../common/search/vectorSearchEmbeddingsM
 import { similarityValues } from "../../../common/schemas.js";
 
 export class CreateIndexTool extends MongoDBToolBase {
+    private filterFieldSchema = z
+        .object({
+            type: z.literal("filter"),
+            path: z
+                .string()
+                .describe(
+                    "Name of the field to index. For nested fields, use dot notation to specify path to embedded fields"
+                ),
+        })
+        .strict()
+        .describe("Definition for a field that will be used for pre-filtering results.");
+
+    private vectorFieldSchema = z
+        .object({
+            type: z.literal("vector"),
+            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)
+                .default(this.config.vectorSearchDimensions)
+                .describe(
+                    "Number of vector dimensions that MongoDB Vector Search enforces at index-time and query-time"
+                ),
+            similarity: z
+                .enum(similarityValues)
+                .default(this.config.vectorSearchSimilarityFunction)
+                .describe(
+                    "Vector similarity function to use to search for top K-nearest neighbors. You can set this field only for vector-type fields."
+                ),
+            quantization: quantizationEnum
+                .default("none")
+                .describe(
+                    "Type of automatic vector quantization for your vectors. Use this setting only if your embeddings are float or double vectors."
+                ),
+        })
+        .strict()
+        .describe("Definition for a field that contains vector embeddings.");
+
+    private autoEmbedFieldSchema = z
+        .object({
+            // TODO: Before public preview this needs to be changed to either
+            // "autoEmbedText" or "autoEmbed" depending on which syntax is
+            // finalized.
+            type: z.literal("text"),
+            path: z
+                .string()
+                .describe(
+                    "Name of the field to index. For nested fields, use dot notation to specify path to embedded fields"
+                ),
+            model: z
+                .enum(["voyage-3-large", "voyage-3.5", "voyage-3.5-lite"])
+                .describe("Voyage embedding model to be used for generating auto-embeddings"),
+            // TODO: As of public preview, this is not required and even if
+            // required, the supported modality is only "text". However we should
+            // keep an eye on whether the recent changes changes the requirements
+            // around this field.
+            // modality: z.enum(["text"]).describe("The type of data in the field mentioned in `path` property."),
+
+            // We don't support specifying `hnswOptions` even in the current vector
+            // search implementation. Following the same idea, we won't support
+            // `hnswOptions` for `autoEmbedText` field definition as well.
+            // hnswOptions: z.object({})
+        })
+        .strict()
+        .describe("Definition for a field for which embeddings must be auto-generated.");
+
     private vectorSearchIndexDefinition = z
         .object({
             type: z.literal("vectorSearch"),
             fields: z
                 .array(
                     z.discriminatedUnion("type", [
-                        z
-                            .object({
-                                type: z.literal("filter"),
-                                path: z
-                                    .string()
-                                    .describe(
-                                        "Name of the field to index. For nested fields, use dot notation to specify path to embedded fields"
-                                    ),
-                            })
-                            .strict()
-                            .describe("Definition for a field that will be used for pre-filtering results."),
-                        z
-                            .object({
-                                type: z.literal("vector"),
-                                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)
-                                    .default(this.config.vectorSearchDimensions)
-                                    .describe(
-                                        "Number of vector dimensions that MongoDB Vector Search enforces at index-time and query-time"
-                                    ),
-                                similarity: z
-                                    .enum(similarityValues)
-                                    .default(this.config.vectorSearchSimilarityFunction)
-                                    .describe(
-                                        "Vector similarity function to use to search for top K-nearest neighbors. You can set this field only for vector-type fields."
-                                    ),
-                                quantization: quantizationEnum
-                                    .default("none")
-                                    .describe(
-                                        "Type of automatic vector quantization for your vectors. Use this setting only if your embeddings are float or double vectors."
-                                    ),
-                            })
-                            .strict()
-                            .describe("Definition for a field that contains vector embeddings."),
+                        this.filterFieldSchema,
+                        this.vectorFieldSchema,
+                        this.autoEmbedFieldSchema,
                     ])
                 )
                 .nonempty()
-                .refine((fields) => fields.some((f) => f.type === "vector"), {
-                    message: "At least one vector field must be defined",
-                })
-                .describe(
-                    "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."
-                ),
+                .refine((fields) => fields.some((f) => f.type === "vector" || f.type === "text"), {
+                    message: "At least one field of type 'vector' or 'text' must be defined",
+                }).describe(`\
+Definitions for the vector, auto-embed and filter fields to index, one definition per document. \
+You must specify 'vector' for fields that contain vector embeddings, 'text' for fields for which embeddings must be auto-generated and, \
+'filter' for additional fields to filter on. At least one 'vector' or 'text' type field definition is required.\
+`),
         })
         .describe("Definition for a Vector Search index.");
 
@@ -183,7 +217,12 @@ export class CreateIndexTool extends MongoDBToolBase {
                 break;
             case "vectorSearch":
                 {
-                    await this.session.assertSearchSupported();
+                    const creatingAutoEmbedIndex = definition.fields.some((field) => field.type === "text");
+                    if (creatingAutoEmbedIndex) {
+                        await this.session.assertAutoEmbeddingIndexSupported();
+                    } else {
+                        await this.session.assertSearchSupported();
+                    }
                     indexes = await provider.createSearchIndexes(database, collection, [
                         {
                             name,

src/tools/mongodb/create/insertMany.ts

@@ -12,7 +12,7 @@ const zSupportedEmbeddingParametersWithInput = zSupportedEmbeddingParameters.ext
     input: z
         .array(z.object({}).passthrough())
         .describe(
-            "Array of objects with vector search index fields as keys (in dot notation) and the raw text values to generate embeddings for as values. The index of each object corresponds to the index of the document in the documents array."
+            "Array of objects with indexed field paths as keys (in dot notation) and the raw text values to generate embeddings for as values. Only provide fields that require embeddings to be generated; do not include fields that are covered by auto-embedding indexes. The index of each object corresponds to the index of the document in the documents array."
         ),
 });
 
@@ -34,7 +34,7 @@ export class InsertManyTool extends MongoDBToolBase {
               embeddingParameters: zSupportedEmbeddingParametersWithInput
                   .optional()
                   .describe(
-                      "The embedding model and its parameters to use to generate embeddings for fields with vector search indexes. Note to LLM: If unsure which embedding model to use, ask the user before providing one."
+                      "The embedding model and its parameters to use to generate embeddings. Only provide this for fields where you need to generate embeddings; do not include fields that have auto-embedding indexes configured, as MongoDB will automatically generate embeddings for those. Note to LLM: If unsure which embedding model to use, ask the user before providing one."
                   ),
           }
         : commonArgs;
@@ -93,7 +93,9 @@ export class InsertManyTool extends MongoDBToolBase {
             return documents;
         }
 
-        // Get vector search indexes for the collection
+        // Get vector search indexes for the collection.
+        // Note: embeddingsForNamespace() only returns fields that require manual embedding generation,
+        // excluding fields with auto-embedding indexes where MongoDB generates embeddings automatically.
         const vectorIndexes = await this.session.vectorSearchEmbeddingsManager.embeddingsForNamespace({
             database,
             collection,
@@ -105,7 +107,7 @@ export class InsertManyTool extends MongoDBToolBase {
                 if (!vectorIndexes.some((index) => index.path === fieldPath)) {
                     throw new MongoDBError(
                         ErrorCodes.AtlasVectorSearchInvalidQuery,
-                        `Field '${fieldPath}' does not have a vector search index in collection ${database}.${collection}. Only fields with vector search indexes can have embeddings generated.`
+                        `Field '${fieldPath}' does not have a vector search index configured for manual embedding generation in collection ${database}.${collection}. This field either has no index or has an auto-embedding index where MongoDB generates embeddings automatically.`
                     );
                 }
             }

src/tools/mongodb/mongodbSchemas.ts

@@ -60,8 +60,17 @@ export const VectorSearchStage = z.object({
                 ),
             queryVector: z
                 .union([z.string(), z.array(z.number())])
+                .optional()
+                .describe(
+                    "The content to search for when querying indexes that require manual embedding generation. Provide an array of numbers (embeddings) or a string with embeddingParameters. Do not use this for auto-embedding indexes; use 'query' instead."
+                ),
+            query: z
+                .object({
+                    text: z.string().describe("The text query to search for."),
+                })
+                .optional()
                 .describe(
-                    "The content to search for. The embeddingParameters field is mandatory if the queryVector is a string, in that case, the tool generates the embedding automatically using the provided configuration."
+                    "The query to search for when using auto-embedding indexes. MongoDB will automatically generate embeddings for the text. Use this for auto-embedding indexes, not 'queryVector'."
                 ),
             numCandidates: z
                 .number()
@@ -78,8 +87,11 @@ export const VectorSearchStage = z.object({
             embeddingParameters: zSupportedEmbeddingParameters
                 .optional()
                 .describe(
-                    "The embedding model and its parameters to use to generate embeddings before searching. It is mandatory if queryVector is a string value. Note to LLM: If unsure, ask the user before providing one."
+                    "The embedding model and its parameters to use to generate embeddings before searching. Only provide this when using 'queryVector' with a string value for indexes that require manual embedding generation. Do not provide this for auto-embedding indexes that use 'query'. Note to LLM: If unsure, ask the user before providing one."
                 ),
         })
-        .passthrough(),
+        .passthrough()
+        .refine((data) => (data.queryVector !== undefined) !== (data.query !== undefined), {
+            message: "Either 'queryVector' or 'query' must be provided, but not both.",
+        }),
 });

src/tools/mongodb/read/aggregate.ts

@@ -241,14 +241,29 @@ export class AggregateTool extends MongoDBToolBase {
             if ("$vectorSearch" in stage) {
                 const { $vectorSearch: vectorSearchStage } = stage as z.infer<typeof VectorSearchStage>;
 
+                // If using 'query' field (auto-embed indexes), MongoDB handles embeddings automatically
+                // so we don't need to do anything.
+                if (vectorSearchStage.query) {
+                    continue;
+                }
+
+                // If queryVector is already an array, no embedding generation needed
                 if (Array.isArray(vectorSearchStage.queryVector)) {
                     continue;
                 }
 
+                // At this point, queryVector must be a string and we need to generate embeddings
+                if (!vectorSearchStage.queryVector) {
+                    throw new MongoDBError(
+                        ErrorCodes.AtlasVectorSearchInvalidQuery,
+                        "Either 'queryVector' or 'query' must be provided in $vectorSearch."
+                    );
+                }
+
                 if (!vectorSearchStage.embeddingParameters) {
                     throw new MongoDBError(
                         ErrorCodes.AtlasVectorSearchInvalidQuery,
-                        "embeddingModel is mandatory if queryVector is a raw string."
+                        "embeddingParameters is mandatory when queryVector is a string. For auto-embedding indexes, use 'query' instead of 'queryVector'."
                     );
                 }