chore: add some integration tests with voyage AI

kmruiz · kmruiz · commit f9ddea161668 · 2025-10-20T18:24:58.000+02:00
diff --git a/.github/workflows/code-health.yml b/.github/workflows/code-health.yml
@@ -34,6 +34,8 @@ jobs:
         run: npm ci
       - name: Run tests
         run: npm test
+        env:
+          TEST_MDB_MCP_VOYAGE_API_KEY=${{ secrets.TEST_MDB_MCP_VOYAGE_API_KEY }}
       - name: Upload test results
         if: always() && matrix.os == 'ubuntu-latest'
         uses: actions/upload-artifact@v4
diff --git a/src/common/search/embeddingsProvider.ts b/src/common/search/embeddingsProvider.ts
@@ -48,8 +48,8 @@ class VoyageEmbeddingsProvider implements EmbeddingsProvider<VoyageModels> {
     ): Promise<Embeddings[]> {
         const voyageParameters = {
             inputType: parameters.inputType,
-            outputDimensions: parameters.numDimensions,
-            quantization: "float", // it is hardcoded on purpose as we don't do quantization yet
+            outputDimension: parameters.numDimensions,
+            outputDtype: "float", // it is hardcoded on purpose as we don't do quantization yet
         };
 
         const model = this.voyage.textEmbeddingModel(modelId);
@@ -61,15 +61,6 @@ class VoyageEmbeddingsProvider implements EmbeddingsProvider<VoyageModels> {
 
         return embeddings;
     }
-
-    async embedOne<Model extends VoyageModels>(
-        modelId: Model,
-        content: EmbeddingsInput,
-        parameters: EmbeddingParameters
-    ): Promise<Embeddings> {
-        const embeddings = await this.embed(modelId, [content], parameters);
-        return embeddings[0] ?? [];
-    }
 }
 
 export function getEmbeddingsProvider(userConfig: UserConfig): EmbeddingsProvider<VoyageModels> | undefined {
diff --git a/src/tools/mongodb/read/aggregate.ts b/src/tools/mongodb/read/aggregate.ts
@@ -60,7 +60,7 @@ export const AggregateArgs = {
     pipeline: z
         .array(z.union([AnyStage, VectorSearchStage]))
         .describe(
-            "An array of aggregation stages to execute. $vectorSearch can only appear as the first stage of the aggregation pipeline or as the first stage of a $unionWith subpipeline."
+            "An array of aggregation stages to execute. $vectorSearch can only appear as the first stage of the aggregation pipeline or as the first stage of a $unionWith subpipeline. When using $vectorSearch, unless the user explicitly asks for the embeddings, $unset any embedding field to avoid reaching context limits."
         ),
     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. \
diff --git a/tests/integration/tools/mongodb/read/aggregate.test.ts b/tests/integration/tools/mongodb/read/aggregate.test.ts
@@ -4,9 +4,16 @@ import {
     validateThrowsForInvalidArguments,
     getResponseContent,
     defaultTestConfig,
+    getDataFromUntrustedContent,
 } from "../../../helpers.js";
 import { beforeEach, describe, expect, it, vi, afterEach } from "vitest";
-import { describeWithMongoDB, getDocsFromUntrustedContent, validateAutoConnectBehavior } from "../mongodbHelpers.js";
+import {
+    createVectorSearchIndexAndWait,
+    describeWithMongoDB,
+    getDocsFromUntrustedContent,
+    validateAutoConnectBehavior,
+    waitUntilSearchIsReady,
+} from "../mongodbHelpers.js";
 import * as constants from "../../../../../src/helpers/constants.js";
 import { freshInsertDocuments } from "./find.test.js";
 
@@ -21,7 +28,7 @@ describeWithMongoDB("aggregate tool", (integration) => {
         {
             name: "pipeline",
             description:
-                "An array of aggregation stages to execute. $vectorSearch can only appear as the first stage of the aggregation pipeline or as the first stage of a $unionWith subpipeline.",
+                "An array of aggregation stages to execute. $vectorSearch can only appear as the first stage of the aggregation pipeline or as the first stage of a $unionWith subpipeline. When using $vectorSearch, unless the user explicitly asks for the embeddings, $unset any embedding field to avoid reaching context limits.",
             type: "array",
             required: true,
         },
@@ -378,3 +385,111 @@ describeWithMongoDB(
         getUserConfig: () => ({ ...defaultTestConfig, maxDocumentsPerQuery: -1, maxBytesPerQuery: -1 }),
     }
 );
+
+const WEBSCALE_DOCUMENT_EMBEDDING = [
+    -0.119673342, 0.028537489, 0.050937884, -0.093283832, 0.050631031, 0.008438504, -0.006635733, -0.082850769,
+    -0.056768127, -0.06443949, 2.637e-5, 0.014422172, -0.087146744, -0.04173224, 0.039277405, 0.029304625, 0.004717892,
+    0.117832206, 0.031759463, 0.019945556, 0.031606037, -0.155882195, -0.02086612, -0.090828992, -0.026849788,
+    0.010126205, 0.009512496, 0.130106404, 0.042039096, -0.06658747, -0.055847559, -0.038663693, -0.072110862,
+    0.073338278, 0.034521155, -0.058302399, -0.052472156, -0.036975995, 0.004602821, -0.06443949, -0.008093293,
+    -0.061984655, 0.098807223, -0.1429943, 0.012197475, 0.003567186, -0.099420927, 0.087146744, -0.085305609,
+    0.011737193, 0.02086612, -0.022707248, 0.04173224, 0.052779011, -0.005523385, 0.045721356, -0.094511256,
+    -0.09512496, 0.086533032, -0.028844343, -0.042039096, -0.006750803, -0.050324172, 0.125810429, -0.052472156,
+    -0.02147983, -0.013808462, 0.019945556, -0.072417714, 0.047869336, -0.03958426, -0.016877009, 0.071804009,
+    0.017797574, 0.010816629, -0.144221723, -0.004986389, 0.089601576, -0.10985399, 0.101262063, -0.022707248,
+    0.001006867, -0.002358946, 0.067508034, 0.124583013, -0.154654771, 0.031606037, -0.165701538, 0.003202796,
+    -0.009512496, 0.080395937, 0.106171735, 0.004756248, 0.123969309, -0.01396189, 0.024088096, -0.013118039,
+    0.02792378, 0.026849788, 0.020098984, -0.113536254, 2.3973e-5, -0.111081414, 0.051858447, -0.053392723, 0.060757235,
+    0.044800788, -0.049403612, -0.075179406, 0.03958426, -0.013808462, -0.013578322, -0.079782225, -0.16447413,
+    0.007594654, 0.039277405, 0.042039096, -0.035595149, 0.034828011, 0.006022024, 0.038356841, 0.045107644,
+    0.084078193, -0.044493936, 0.024548376, 0.008822073, 0.027003214, -0.0487899, 0.067201182, -0.053392723,
+    0.108012855, 0.070883438, 0.022553822, 0.110467695, -0.055540707, -0.030685471, -0.146676555, 0.064746343,
+    -0.036669139, -0.046948772, 0.020559266, -0.142380595, -0.010049492, 0.015112595, 0.091442712, 0.022707248,
+    -0.050937884, 0.026849788, -0.075486265, 0.018181141, 0.014192032, 0.041118532, -0.038049985, -0.011813907,
+    0.067201182, 0.005293244, -0.059222963, -0.088374153, -0.098193504, 0.012350903, -0.030838897, 0.113536254,
+    -0.035595149, 0.073338278, 0.146676555, -0.013271467, -0.043266516, -0.061984655, -0.054006428, 0.120287046,
+    0.052472156, 0.022860678, -0.018948279, 0.007671368, -0.008822073, 0.021786686, 0.033447165, -0.065666914,
+    0.025162086, 0.005715169, 0.042345952, 0.006520663, -0.025775796, 0.060757235, -0.044800788, 0.052779011,
+    0.033140309, -0.033293735, -0.01856471, 0.045107644, -0.052779011, 0.038049985, -0.086533032, -0.077327386,
+    -0.051244736, -0.155882195, 0.010356346, -0.15956445, 0.019331846, -0.04756248, -0.0145756, 0.130720109,
+    -0.007096016, 0.041425385, -0.042652804, 0.005600099, -0.017030437, 0.002493195, 0.032219745, -0.054313287,
+    0.044493936, -0.011813907, 0.025622368, 0.054006428, -0.010586488, -0.055847559, 0.034981437, 0.077327386,
+    0.024548376, 0.106171735, 0.032066315, 0.069962874, 0.059836671, -0.031452607, -0.00027569, -0.022246968,
+    0.058302399, -0.005369958, -0.101875767, 0.032986883, 0.09512496, -0.085919321, 0.005408315, -0.037436277,
+    0.034367729, 0.077941097, -0.04756248, 0.000110276, -0.02792378, -0.059836671, 0.02086612, 0.060450379,
+    -0.045107644, 0.002627444, 0.081623361, 0.054313287, -0.022400394, 0.065053202, 0.074565701, 0.04081168,
+    -0.021786686, 0.044493936, 0.073338278, 0.003221974, 0.001419203, 0.00740287,
+];
+
+describeWithMongoDB(
+    "aggregate tool with atlas search enabled",
+    (integration) => {
+        it.skipIf(!process.env.TEST_MDB_MCP_VOYAGE_API_KEY)(
+            "should be able to return elements from within a vector search query",
+            async () => {
+                await waitUntilSearchIsReady(integration.mongoClient());
+
+                const collection = integration.mongoClient().db(integration.randomDbName()).collection("databases");
+                await collection.insertOne({ name: "mongodb", description_embedding: WEBSCALE_DOCUMENT_EMBEDDING });
+
+                await createVectorSearchIndexAndWait(
+                    integration.mongoClient(),
+                    integration.randomDbName(),
+                    "databases",
+                    [
+                        {
+                            type: "vector",
+                            path: "description_embedding",
+                            numDimensions: 256,
+                            similarity: "euclidean",
+                            quantization: "none",
+                        },
+                    ]
+                );
+
+                // now query the index
+                await integration.connectMcpClient();
+                const response = await integration.mcpClient().callTool({
+                    name: "aggregate",
+                    arguments: {
+                        database: integration.randomDbName(),
+                        collection: "databases",
+                        pipeline: [
+                            {
+                                $vectorSearch: {
+                                    index: "default",
+                                    path: "description_embedding",
+                                    queryVector: "webscale",
+                                    embeddingModel: "voyage-3-large",
+                                    numCandidates: 10,
+                                    limit: 10,
+                                },
+                            },
+                            {
+                                $project: {
+                                    description_embedding: 0,
+                                },
+                            },
+                        ],
+                    },
+                });
+
+                const responseContent = getResponseContent(response);
+                expect(responseContent).toContain("The aggregation resulted in 1 documents. Returning 1 documents.");
+                const untrustedDocs = getDocsFromUntrustedContent<{ name: string }>(responseContent);
+                expect(untrustedDocs).toHaveLength(1);
+                expect(untrustedDocs[0].name).toBe("mongodb");
+            }
+        );
+    },
+
+    {
+        getUserConfig: () => ({
+            ...defaultTestConfig,
+            voyageApiKey: process.env.TEST_MDB_MCP_VOYAGE_API_KEY,
+            maxDocumentsPerQuery: -1,
+            maxBytesPerQuery: -1,
+        }),
+        downloadOptions: { search: true },
+    }
+);
