chore: feature flag the aggregate schema

Author: himanshusinghs

src/tools/mongodb/metadata/explain.ts

@@ -4,7 +4,7 @@ import type { ToolArgs, OperationType } from "../../tool.js";
 import { formatUntrustedData } from "../../tool.js";
 import { z } from "zod";
 import type { Document } from "mongodb";
-import { AggregateArgs } from "../read/aggregate.js";
+import { getAggregateArgs } from "../read/aggregate.js";
 import { FindArgs } from "../read/find.js";
 import { CountArgs } from "../read/count.js";
 
@@ -20,7 +20,7 @@ export class ExplainTool extends MongoDBToolBase {
                 z.discriminatedUnion("name", [
                     z.object({
                         name: z.literal("aggregate"),
-                        arguments: z.object(AggregateArgs),
+                        arguments: z.object(getAggregateArgs(this.isFeatureEnabled("vectorSearch"))),
                     }),
                     z.object({
                         name: z.literal("find"),

src/tools/mongodb/mongodbSchemas.ts

@@ -42,7 +42,7 @@ export type EmbeddingParameters = {
 export const zSupportedEmbeddingParameters = zVoyageEmbeddingParameters.extend({ model: zVoyageModels });
 export type SupportedEmbeddingParameters = z.infer<typeof zSupportedEmbeddingParameters>;
 
-export const AnyVectorSearchStage = zEJSON();
+export const AnyAggregateStage = zEJSON();
 export const VectorSearchStage = z.object({
     $vectorSearch: z
         .object({

src/tools/mongodb/read/aggregate.ts

@@ -12,18 +12,17 @@ import { collectCursorUntilMaxBytesLimit } from "../../../helpers/collectCursorU
 import { operationWithFallback } from "../../../helpers/operationWithFallback.js";
 import { AGG_COUNT_MAX_TIME_MS_CAP, ONE_MB, CURSOR_LIMITS_TO_LLM_TEXT } from "../../../helpers/constants.js";
 import { LogId } from "../../../common/logger.js";
-import { AnyVectorSearchStage, VectorSearchStage } from "../mongodbSchemas.js";
+import { AnyAggregateStage, VectorSearchStage } from "../mongodbSchemas.js";
 import {
     assertVectorSearchFilterFieldsAreIndexed,
     type VectorSearchIndex,
 } from "../../../helpers/assertVectorSearchFilterFieldsAreIndexed.js";
 
-export const AggregateArgs = {
-    pipeline: z.array(z.union([AnyVectorSearchStage, VectorSearchStage])).describe(
-        `An array of aggregation stages to execute.  
+const pipelineDescription = `\
+An array of aggregation stages to execute.
 \`$vectorSearch\` **MUST** be the first stage of the pipeline, or the first stage of a \`$unionWith\` subpipeline.
 ### Usage Rules for \`$vectorSearch\`
-- **Unset embeddings:**  
+- **Unset embeddings:**
   Unless the user explicitly requests the embeddings, add an \`$unset\` stage **at the end of the pipeline** to remove the embedding field and avoid context limits. **The $unset stage in this situation is mandatory**.
 - **Pre-filtering:**
 If the user requests additional filtering, include filters in \`$vectorSearch.filter\` only for pre-filter fields in the vector index.
@@ -32,20 +31,26 @@ If the user requests additional filtering, include filters in \`$vectorSearch.fi
     For all remaining filters, add a $match stage after $vectorSearch.
 ### Note to LLM
 - If unsure which fields are filterable, use the collection-indexes tool to determine valid prefilter fields.
-- If no requested filters are valid prefilters, omit the filter key from $vectorSearch.`
-    ),
-    responseBytesLimit: z.number().optional().default(ONE_MB).describe(`\
+- If no requested filters are valid prefilters, omit the filter key from $vectorSearch.\
+`;
+
+export const getAggregateArgs = (vectorSearchEnabled: boolean) =>
+    ({
+        pipeline: z
+            .array(vectorSearchEnabled ? z.union([AnyAggregateStage, VectorSearchStage]) : AnyAggregateStage)
+            .describe(pipelineDescription),
+        responseBytesLimit: z.number().optional().default(ONE_MB).describe(`\
 The maximum number of bytes to return in the response. This value is capped by the server's configured maxBytesPerQuery and cannot be exceeded. \
 Note to LLM: If the entire aggregation result is required, use the "export" tool instead of increasing this limit.\
 `),
-};
+    }) as const;
 
 export class AggregateTool extends MongoDBToolBase {
     public name = "aggregate";
     protected description = "Run an aggregation against a MongoDB collection";
     protected argsShape = {
         ...DbOperationArgs,
-        ...AggregateArgs,
+        ...getAggregateArgs(this.isFeatureEnabled("vectorSearch")),
     };
     public operationType: OperationType = "read";
 

src/tools/mongodb/read/export.ts

@@ -6,7 +6,7 @@ import type { OperationType, ToolArgs } from "../../tool.js";
 import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
 import { FindArgs } from "./find.js";
 import { jsonExportFormat } from "../../../common/exportsManager.js";
-import { AggregateArgs } from "./aggregate.js";
+import { getAggregateArgs } from "./aggregate.js";
 
 export class ExportTool extends MongoDBToolBase {
     public name = "export";
@@ -32,7 +32,9 @@ export class ExportTool extends MongoDBToolBase {
                         name: z
                             .literal("aggregate")
                             .describe("The literal name 'aggregate' to represent an aggregation cursor as target."),
-                        arguments: z.object(AggregateArgs).describe("The arguments for 'aggregate' operation."),
+                        arguments: z
+                            .object(getAggregateArgs(this.isFeatureEnabled("vectorSearch")))
+                            .describe("The arguments for 'aggregate' operation."),
                     }),
                 ])
             )

tests/integration/tools/mongodb/read/aggregate.test.ts

@@ -27,10 +27,10 @@ describeWithMongoDB("aggregate tool", (integration) => {
         ...databaseCollectionParameters,
         {
             name: "pipeline",
-            description: `An array of aggregation stages to execute.  
+            description: `An array of aggregation stages to execute.
 \`$vectorSearch\` **MUST** be the first stage of the pipeline, or the first stage of a \`$unionWith\` subpipeline.
 ### Usage Rules for \`$vectorSearch\`
-- **Unset embeddings:**  
+- **Unset embeddings:**
   Unless the user explicitly requests the embeddings, add an \`$unset\` stage **at the end of the pipeline** to remove the embedding field and avoid context limits. **The $unset stage in this situation is mandatory**.
 - **Pre-filtering:**
 If the user requests additional filtering, include filters in \`$vectorSearch.filter\` only for pre-filter fields in the vector index.