chore: add more tests to ensure embedding generation works

Author: kmruiz

src/common/search/embeddingsProvider.ts

@@ -12,7 +12,10 @@ export type EmbeddingParameters = {
     inputType: "query" | "document";
 };
 
-interface EmbeddingsProvider<SupportedModels extends string, SupportedEmbeddingParameters extends EmbeddingParameters> {
+export interface EmbeddingsProvider<
+    SupportedModels extends string,
+    SupportedEmbeddingParameters extends EmbeddingParameters,
+> {
     embed(
         modelId: SupportedModels,
         content: EmbeddingsInput[],

src/common/search/vectorSearchEmbeddingsManager.ts

@@ -35,7 +35,8 @@ export class VectorSearchEmbeddingsManager {
     constructor(
         private readonly config: UserConfig,
         private readonly connectionManager: ConnectionManager,
-        private readonly embeddings: Map<EmbeddingNamespace, VectorFieldIndexDefinition[]> = new Map()
+        private readonly embeddings: Map<EmbeddingNamespace, VectorFieldIndexDefinition[]> = new Map(),
+        private readonly embeddingsProvider: typeof getEmbeddingsProvider = getEmbeddingsProvider
     ) {
         connectionManager.events.on("connection-close", () => {
             this.embeddings.clear();
@@ -242,15 +243,21 @@ export class VectorSearchEmbeddingsManager {
             );
         }
 
-        const embeddingsProvider = getEmbeddingsProvider(this.config);
+        const embeddingsProvider = this.embeddingsProvider(this.config);
 
         if (!embeddingsProvider) {
             throw new MongoDBError(ErrorCodes.NoEmbeddingsProviderConfigured, "No embeddings provider configured.");
         }
 
+        if (this.config.disableEmbeddingsValidation) {
+            return await embeddingsProvider.embed(embeddingParameters.model, rawValues, {
+                inputType,
+                ...embeddingParameters,
+            });
+        }
+
         const embeddingInfoForCollection = await this.embeddingsForNamespace({ database, collection });
         const embeddingInfoForPath = embeddingInfoForCollection.find((definition) => definition.path === path);
-
         if (!embeddingInfoForPath) {
             throw new MongoDBError(
                 ErrorCodes.AtlasVectorSearchIndexNotFound,

tests/unit/common/search/vectorSearchEmbeddingsManager.test.ts

@@ -13,6 +13,13 @@ import { ConnectionStateConnected } from "../../../../src/common/connectionManag
 import type { InsertOneResult } from "mongodb";
 import type { DropDatabaseResult } from "@mongosh/service-provider-node-driver/lib/node-driver-service-provider.js";
 import EventEmitter from "events";
+import {
+    zVoyageEmbeddingParameters,
+    type EmbeddingParameters,
+    type EmbeddingsProvider,
+    type getEmbeddingsProvider,
+} from "../../../../src/common/search/embeddingsProvider.js";
+import { ErrorCodes, MongoDBError } from "../../../../src/common/errors.js";
 
 type MockedServiceProvider = NodeDriverServiceProvider & {
     getSearchIndexes: MockedFunction<NodeDriverServiceProvider["getSearchIndexes"]>;
@@ -25,6 +32,10 @@ type MockedConnectionManager = ConnectionManager & {
     currentConnectionState: ConnectionStateConnected;
 };
 
+type MockedEmbeddingsProvider = EmbeddingsProvider<string, EmbeddingParameters> & {
+    embed: MockedFunction<EmbeddingsProvider<string, EmbeddingParameters>["embed"]>;
+};
+
 const database = "my" as const;
 const collection = "collection" as const;
 const mapKey = `${database}.${collection}` as EmbeddingNamespace;
@@ -78,13 +89,22 @@ describe("VectorSearchEmbeddingsManager", () => {
         getURI: () => "mongodb://my-test",
     } as unknown as MockedServiceProvider;
 
+    const embeddingsProvider: MockedEmbeddingsProvider = {
+        embed: vi.fn(),
+    };
+
+    const getMockedEmbeddingsProvider: typeof getEmbeddingsProvider = () => {
+        return embeddingsProvider;
+    };
+
     const connectionManager: MockedConnectionManager = {
         currentConnectionState: new ConnectionStateConnected(provider),
         events: eventEmitter,
     } as unknown as MockedConnectionManager;
 
     beforeEach(() => {
         provider.getSearchIndexes.mockReset();
+        embeddingsProvider.embed.mockReset();
 
         provider.createSearchIndexes.mockResolvedValue([]);
         provider.insertOne.mockResolvedValue({} as unknown as InsertOneResult);
@@ -371,4 +391,117 @@ describe("VectorSearchEmbeddingsManager", () => {
             });
         });
     });
+
+    describe("generate embeddings", () => {
+        const embeddingToGenerate = {
+            database: "mydb",
+            collection: "mycoll",
+            path: "embedding_field",
+            rawValues: ["oops"],
+            embeddingParameters: { model: "voyage-3-large", outputDimension: 1024, outputDType: "float" } as const,
+            inputType: "query" as const,
+        };
+
+        let embeddings: VectorSearchEmbeddingsManager;
+
+        beforeEach(() => {
+            embeddings = new VectorSearchEmbeddingsManager(
+                embeddingValidationDisabled,
+                connectionManager,
+                new Map(),
+                getMockedEmbeddingsProvider
+            );
+        });
+
+        describe("when atlas search is not available", () => {
+            beforeEach(() => {
+                embeddings = new VectorSearchEmbeddingsManager(
+                    embeddingValidationEnabled,
+                    connectionManager,
+                    new Map(),
+                    getMockedEmbeddingsProvider
+                );
+
+                provider.getSearchIndexes.mockRejectedValue(new Error());
+            });
+
+            it("throws an exception", async () => {
+                await expect(embeddings.generateEmbeddings(embeddingToGenerate)).rejects.toThrowError();
+            });
+        });
+
+        describe("when atlas search is available", () => {
+            describe("when embedding validation is disabled", () => {
+                beforeEach(() => {
+                    embeddings = new VectorSearchEmbeddingsManager(
+                        embeddingValidationDisabled,
+                        connectionManager,
+                        new Map(),
+                        getMockedEmbeddingsProvider
+                    );
+                });
+
+                describe("when no index is available for path", () => {
+                    it("returns the embeddings as is", async () => {
+                        embeddingsProvider.embed.mockResolvedValue([[0xc0ffee]]);
+
+                        const [result] = await embeddings.generateEmbeddings(embeddingToGenerate);
+                        expect(result).toEqual([0xc0ffee]);
+                    });
+                });
+            });
+
+            describe("when embedding validation is enabled", () => {
+                beforeEach(() => {
+                    embeddings = new VectorSearchEmbeddingsManager(
+                        embeddingValidationEnabled,
+                        connectionManager,
+                        new Map(),
+                        getMockedEmbeddingsProvider
+                    );
+                });
+
+                describe("when no index is available for path", () => {
+                    it("throws an exception", async () => {
+                        await expect(embeddings.generateEmbeddings(embeddingToGenerate)).rejects.toThrowError();
+                    });
+                });
+
+                describe("when index is available on path", () => {
+                    beforeEach(() => {
+                        provider.getSearchIndexes.mockResolvedValue([
+                            {
+                                id: "65e8c766d0450e3e7ab9855f",
+                                name: "vector-search-test",
+                                type: "vectorSearch",
+                                status: "READY",
+                                queryable: true,
+                                latestDefinition: {
+                                    fields: [
+                                        {
+                                            type: "vector",
+                                            path: embeddingToGenerate.path,
+                                            numDimensions: 1024,
+                                            similarity: "euclidean",
+                                        },
+                                        { type: "filter", path: "genres" },
+                                        { type: "filter", path: "year" },
+                                    ],
+                                },
+                            },
+                        ]);
+                    });
+
+                    describe("when embedding validation is disabled", () => {
+                        it("returns the embeddings as is", async () => {
+                            embeddingsProvider.embed.mockResolvedValue([[0xc0ffee]]);
+
+                            const [result] = await embeddings.generateEmbeddings(embeddingToGenerate);
+                            expect(result).toEqual([0xc0ffee]);
+                        });
+                    });
+                });
+            });
+        });
+    });
 });