Skip to content

Commit bb21e21

Browse files
kmruizkmruiz
committed
chore: tests for basic quantization in the search itself
1 parent f9ddea1 commit bb21e21

File tree

1 file changed

+190
-2
lines changed

1 file changed

+190
-2
lines changed

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

Lines changed: 190 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import {
1616
} from "../mongodbHelpers.js";
1717
import * as constants from "../../../../../src/helpers/constants.js";
1818
import { freshInsertDocuments } from "./find.test.js";
19+
import { BSON } from "bson";
1920

2021
describeWithMongoDB("aggregate tool", (integration) => {
2122
afterEach(() => {
@@ -424,6 +425,10 @@ const WEBSCALE_DOCUMENT_EMBEDDING = [
424425
describeWithMongoDB(
425426
"aggregate tool with atlas search enabled",
426427
(integration) => {
428+
beforeEach(async () => {
429+
await integration.mongoClient().db(integration.randomDbName()).collection("databases").drop();
430+
});
431+
427432
it.skipIf(!process.env.TEST_MDB_MCP_VOYAGE_API_KEY)(
428433
"should be able to return elements from within a vector search query",
429434
async () => {
@@ -478,15 +483,198 @@ describeWithMongoDB(
478483
expect(responseContent).toContain("The aggregation resulted in 1 documents. Returning 1 documents.");
479484
const untrustedDocs = getDocsFromUntrustedContent<{ name: string }>(responseContent);
480485
expect(untrustedDocs).toHaveLength(1);
481-
expect(untrustedDocs[0].name).toBe("mongodb");
486+
expect(untrustedDocs[0]!!.name).toBe("mongodb");
487+
}
488+
);
489+
490+
it.skipIf(!process.env.TEST_MDB_MCP_VOYAGE_API_KEY)(
491+
"should be able too return elements from within a vector search query using binary encoding",
492+
async () => {
493+
await waitUntilSearchIsReady(integration.mongoClient());
494+
495+
const collection = integration.mongoClient().db(integration.randomDbName()).collection("databases");
496+
await collection.insertOne({
497+
name: "mongodb",
498+
description_embedding: BSON.Binary.fromFloat32Array(new Float32Array(WEBSCALE_DOCUMENT_EMBEDDING)),
499+
});
500+
501+
await createVectorSearchIndexAndWait(
502+
integration.mongoClient(),
503+
integration.randomDbName(),
504+
"databases",
505+
[
506+
{
507+
type: "vector",
508+
path: "description_embedding",
509+
numDimensions: 256,
510+
similarity: "euclidean",
511+
quantization: "none",
512+
},
513+
]
514+
);
515+
516+
// now query the index
517+
await integration.connectMcpClient();
518+
const response = await integration.mcpClient().callTool({
519+
name: "aggregate",
520+
arguments: {
521+
database: integration.randomDbName(),
522+
collection: "databases",
523+
pipeline: [
524+
{
525+
$vectorSearch: {
526+
index: "default",
527+
path: "description_embedding",
528+
queryVector: "webscale",
529+
embeddingModel: "voyage-3-large",
530+
numCandidates: 10,
531+
limit: 10,
532+
},
533+
},
534+
{
535+
$project: {
536+
description_embedding: 0,
537+
},
538+
},
539+
],
540+
},
541+
});
542+
543+
const responseContent = getResponseContent(response);
544+
expect(responseContent).toContain("The aggregation resulted in 1 documents. Returning 1 documents.");
545+
const untrustedDocs = getDocsFromUntrustedContent<{ name: string }>(responseContent);
546+
expect(untrustedDocs).toHaveLength(1);
547+
expect(untrustedDocs[0]!!.name).toBe("mongodb");
548+
}
549+
);
550+
551+
it.skipIf(!process.env.TEST_MDB_MCP_VOYAGE_API_KEY)(
552+
"should be able too return elements from within a vector search query using scalar quantization",
553+
async () => {
554+
await waitUntilSearchIsReady(integration.mongoClient());
555+
556+
const collection = integration.mongoClient().db(integration.randomDbName()).collection("databases");
557+
await collection.insertOne({
558+
name: "mongodb",
559+
description_embedding: BSON.Binary.fromFloat32Array(new Float32Array(WEBSCALE_DOCUMENT_EMBEDDING)),
560+
});
561+
562+
await createVectorSearchIndexAndWait(
563+
integration.mongoClient(),
564+
integration.randomDbName(),
565+
"databases",
566+
[
567+
{
568+
type: "vector",
569+
path: "description_embedding",
570+
numDimensions: 256,
571+
similarity: "euclidean",
572+
quantization: "scalar",
573+
},
574+
]
575+
);
576+
577+
// now query the index
578+
await integration.connectMcpClient();
579+
const response = await integration.mcpClient().callTool({
580+
name: "aggregate",
581+
arguments: {
582+
database: integration.randomDbName(),
583+
collection: "databases",
584+
pipeline: [
585+
{
586+
$vectorSearch: {
587+
index: "default",
588+
path: "description_embedding",
589+
queryVector: "webscale",
590+
embeddingModel: "voyage-3-large",
591+
numCandidates: 10,
592+
limit: 10,
593+
},
594+
},
595+
{
596+
$project: {
597+
description_embedding: 0,
598+
},
599+
},
600+
],
601+
},
602+
});
603+
604+
const responseContent = getResponseContent(response);
605+
expect(responseContent).toContain("The aggregation resulted in 1 documents. Returning 1 documents.");
606+
const untrustedDocs = getDocsFromUntrustedContent<{ name: string }>(responseContent);
607+
expect(untrustedDocs).toHaveLength(1);
608+
expect(untrustedDocs[0]!!.name).toBe("mongodb");
609+
}
610+
);
611+
612+
it.skipIf(!process.env.TEST_MDB_MCP_VOYAGE_API_KEY)(
613+
"should be able too return elements from within a vector search query using binary quantization",
614+
async () => {
615+
await waitUntilSearchIsReady(integration.mongoClient());
616+
617+
const collection = integration.mongoClient().db(integration.randomDbName()).collection("databases");
618+
await collection.insertOne({
619+
name: "mongodb",
620+
description_embedding: BSON.Binary.fromFloat32Array(new Float32Array(WEBSCALE_DOCUMENT_EMBEDDING)),
621+
});
622+
623+
await createVectorSearchIndexAndWait(
624+
integration.mongoClient(),
625+
integration.randomDbName(),
626+
"databases",
627+
[
628+
{
629+
type: "vector",
630+
path: "description_embedding",
631+
numDimensions: 256,
632+
similarity: "euclidean",
633+
quantization: "binary",
634+
},
635+
]
636+
);
637+
638+
// now query the index
639+
await integration.connectMcpClient();
640+
const response = await integration.mcpClient().callTool({
641+
name: "aggregate",
642+
arguments: {
643+
database: integration.randomDbName(),
644+
collection: "databases",
645+
pipeline: [
646+
{
647+
$vectorSearch: {
648+
index: "default",
649+
path: "description_embedding",
650+
queryVector: "webscale",
651+
embeddingModel: "voyage-3-large",
652+
numCandidates: 10,
653+
limit: 10,
654+
},
655+
},
656+
{
657+
$project: {
658+
description_embedding: 0,
659+
},
660+
},
661+
],
662+
},
663+
});
664+
665+
const responseContent = getResponseContent(response);
666+
expect(responseContent).toContain("The aggregation resulted in 1 documents. Returning 1 documents.");
667+
const untrustedDocs = getDocsFromUntrustedContent<{ name: string }>(responseContent);
668+
expect(untrustedDocs).toHaveLength(1);
669+
expect(untrustedDocs[0]!!.name).toBe("mongodb");
482670
}
483671
);
484672
},
485673

486674
{
487675
getUserConfig: () => ({
488676
...defaultTestConfig,
489-
voyageApiKey: process.env.TEST_MDB_MCP_VOYAGE_API_KEY,
677+
voyageApiKey: process.env.TEST_MDB_MCP_VOYAGE_API_KEY ?? "",
490678
maxDocumentsPerQuery: -1,
491679
maxBytesPerQuery: -1,
492680
}),

0 commit comments

Comments
 (0)