chore: add the ability to disable embedding validation

While we do our best to make sure we don't break anything, there
might be situations where users want to disable the validation
and insert documents as they please.

Author: kmruiz

src/common/config.ts

@@ -58,6 +58,7 @@ const OPTIONS = {
     boolean: [
         "apiDeprecationErrors",
         "apiStrict",
+        "disableEmbeddingsValidation",
         "help",
         "indexCheck",
         "ipv6",
@@ -183,6 +184,7 @@ export interface UserConfig extends CliOptions {
     maxBytesPerQuery: number;
     atlasTemporaryDatabaseUserLifetimeMs: number;
     voyageApiKey: string;
+    disableEmbeddingsValidation: boolean;
 }
 
 export const defaultUserConfig: UserConfig = {
@@ -213,6 +215,7 @@ export const defaultUserConfig: UserConfig = {
     maxBytesPerQuery: 16 * 1024 * 1024, // By default, we only return ~16 mb of data per query / aggregation
     atlasTemporaryDatabaseUserLifetimeMs: 4 * 60 * 60 * 1000, // 4 hours
     voyageApiKey: "",
+    disableEmbeddingsValidation: false,
 };
 
 export const config = setupUserConfig({

src/common/search/vectorSearchEmbeddings.ts

@@ -1,7 +1,8 @@
 import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 import { BSON, type Document } from "bson";
+import type { UserConfig } from "../config.js";
 
-type VectorFieldIndexDefinition = {
+export type VectorFieldIndexDefinition = {
     type: "vector";
     path: string;
     numDimensions: number;
@@ -11,7 +12,10 @@ type VectorFieldIndexDefinition = {
 
 export type EmbeddingNamespace = `${string}.${string}`;
 export class VectorSearchEmbeddings {
-    constructor(private readonly embeddings: Map<EmbeddingNamespace, VectorFieldIndexDefinition[]> = new Map()) {}
+    constructor(
+        private readonly config: UserConfig,
+        private readonly embeddings: Map<EmbeddingNamespace, VectorFieldIndexDefinition[]> = new Map()
+    ) {}
 
     cleanupEmbeddingsForNamespace({ database, collection }: { database: string; collection: string }): void {
         const embeddingDefKey: EmbeddingNamespace = `${database}.${collection}`;
@@ -71,6 +75,13 @@ export class VectorSearchEmbeddings {
     }
 
     private documentPassesEmbeddingValidation(definition: VectorFieldIndexDefinition, document: Document): boolean {
+        // While we can do our best effort to ensure that the embedding validation is correct
+        // based on https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-quantization/
+        // it's a complex process so we will also give the user the ability to disable this validation
+        if (this.config.disableEmbeddingsValidation) {
+            return true;
+        }
+
         const fieldPath = definition.path.split(".");
         let fieldRef: unknown = document;
 
@@ -84,30 +95,37 @@ export class VectorSearchEmbeddings {
 
         switch (definition.quantization) {
             case "none":
+                return true;
             case "scalar":
-                if (!Array.isArray(fieldRef)) {
-                    return false;
-                }
-
-                if (fieldRef.length !== definition.numDimensions) {
-                    return false;
-                }
-
-                if (typeof fieldRef[0] !== "number") {
-                    return false;
-                }
-                break;
             case "binary":
                 if (fieldRef instanceof BSON.Binary) {
                     try {
-                        const bits = fieldRef.toBits();
-                        return bits.length === definition.numDimensions;
+                        const elements = fieldRef.toFloat32Array();
+                        return elements.length === definition.numDimensions;
                     } catch {
-                        return false;
+                        // bits are also supported
+                        try {
+                            const bits = fieldRef.toBits();
+                            return bits.length === definition.numDimensions;
+                        } catch {
+                            return false;
+                        }
                     }
                 } else {
-                    return false;
+                    if (!Array.isArray(fieldRef)) {
+                        return false;
+                    }
+
+                    if (fieldRef.length !== definition.numDimensions) {
+                        return false;
+                    }
+
+                    if (typeof fieldRef[0] !== "number") {
+                        return false;
+                    }
                 }
+
+                break;
         }
 
         return true;

src/transports/base.ts

@@ -90,7 +90,7 @@ export abstract class TransportRunnerBase {
             exportsManager,
             connectionManager,
             keychain: Keychain.root,
-            vectorSearchEmbeddings: new VectorSearchEmbeddings(),
+            vectorSearchEmbeddings: new VectorSearchEmbeddings(this.userConfig),
         });
 
         const telemetry = Telemetry.create(session, this.userConfig, this.deviceId, {

tests/integration/helpers.ts

@@ -102,7 +102,7 @@ export function setupIntegrationTest(
             exportsManager,
             connectionManager,
             keychain: new Keychain(),
-            vectorSearchEmbeddings: new VectorSearchEmbeddings(),
+            vectorSearchEmbeddings: new VectorSearchEmbeddings(userConfig),
         });
 
         // Mock hasValidAccessToken for tests

tests/integration/telemetry.test.ts

@@ -24,7 +24,7 @@ describe("Telemetry", () => {
                 exportsManager: ExportsManager.init(config, logger),
                 connectionManager: new MCPConnectionManager(config, driverOptions, logger, deviceId),
                 keychain: new Keychain(),
-                vectorSearchEmbeddings: new VectorSearchEmbeddings(),
+                vectorSearchEmbeddings: new VectorSearchEmbeddings(config),
             }),
             config,
             deviceId

tests/integration/tools/mongodb/mongodbTool.test.ts

@@ -109,7 +109,7 @@ describe("MongoDBTool implementations", () => {
             exportsManager,
             connectionManager,
             keychain: new Keychain(),
-            vectorSearchEmbeddings: new VectorSearchEmbeddings(),
+            vectorSearchEmbeddings: new VectorSearchEmbeddings(userConfig),
         });
         const telemetry = Telemetry.create(session, userConfig, deviceId);