chore: Add new session-level service for getting embeddings of a specific collection

Author: kmruiz

src/common/search/vectorSearchEmbeddings.ts

@@ -0,0 +1,55 @@
+import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
+import type { Document } from "bson";
+
+type VectorFieldIndexDefinition = {
+    type: "vector";
+    path: string;
+    numDimensions: number;
+    quantization: "none" | "scalar" | "binary";
+    similarity: "euclidean" | "cosine" | "dotProduct";
+};
+
+type EmbeddingNamespace = "${string}.${string}";
+export class VectorSearchEmbeddings {
+    private embeddings: Map<EmbeddingNamespace, VectorFieldIndexDefinition[]>;
+
+    constructor() {
+        this.embeddings = new Map();
+    }
+
+    cleanupEmbeddingsForNamespace({ database, collection }: { database: string; collection: string }): void {
+        const embeddingDefKey = `${database}.${collection}` as EmbeddingNamespace;
+        this.embeddings.delete(embeddingDefKey);
+    }
+
+    async embeddingsForNamespace({
+        database,
+        collection,
+        provider,
+    }: {
+        database: string;
+        collection: string;
+        provider: NodeDriverServiceProvider;
+    }): Promise<VectorFieldIndexDefinition[] | undefined> {
+        const embeddingDefKey = `${database}.${collection}` as EmbeddingNamespace;
+        const definition = this.embeddings.get(embeddingDefKey);
+
+        if (!definition) {
+            const allSearchIndexes = await provider.getSearchIndexes(database, collection);
+            const vectorSearchIndexes = allSearchIndexes.filter((index) => index.type === "vectorSearch");
+            const vectorFields = vectorSearchIndexes
+                // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
+                .flatMap<Document>((index) => (index.latestDefinition?.fields as Document) ?? [])
+                .filter((field) => this.isVectorFieldIndexDefinition(field));
+
+            this.embeddings.set(embeddingDefKey, vectorFields);
+            return vectorFields;
+        } else {
+            return definition;
+        }
+    }
+
+    isVectorFieldIndexDefinition(doc: Document): doc is VectorFieldIndexDefinition {
+        return doc["type"] === "vector";
+    }
+}

src/common/session.ts

@@ -16,6 +16,7 @@ import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-d
 import { ErrorCodes, MongoDBError } from "./errors.js";
 import type { ExportsManager } from "./exportsManager.js";
 import type { Keychain } from "./keychain.js";
+import type { VectorSearchEmbeddings } from "./search/vectorSearchEmbeddings.js";
 
 export interface SessionOptions {
     apiBaseUrl: string;
@@ -25,6 +26,7 @@ export interface SessionOptions {
     exportsManager: ExportsManager;
     connectionManager: ConnectionManager;
     keychain: Keychain;
+    vectorSearchEmbeddings: VectorSearchEmbeddings;
 }
 
 export type SessionEvents = {
@@ -40,6 +42,7 @@ export class Session extends EventEmitter<SessionEvents> {
     readonly connectionManager: ConnectionManager;
     readonly apiClient: ApiClient;
     readonly keychain: Keychain;
+    readonly vectorSearchEmbeddings: VectorSearchEmbeddings;
 
     mcpClient?: {
         name?: string;
@@ -57,6 +60,7 @@ export class Session extends EventEmitter<SessionEvents> {
         connectionManager,
         exportsManager,
         keychain,
+        vectorSearchEmbeddings,
     }: SessionOptions) {
         super();
 
@@ -73,6 +77,7 @@ export class Session extends EventEmitter<SessionEvents> {
         this.apiClient = new ApiClient({ baseUrl: apiBaseUrl, credentials }, logger);
         this.exportsManager = exportsManager;
         this.connectionManager = connectionManager;
+        this.vectorSearchEmbeddings = vectorSearchEmbeddings;
         this.connectionManager.events.on("connection-success", () => this.emit("connect"));
         this.connectionManager.events.on("connection-time-out", (error) => this.emit("connection-error", error));
         this.connectionManager.events.on("connection-close", () => this.emit("disconnect"));

src/tools/mongodb/search/listSearchIndexes.ts

@@ -6,7 +6,7 @@ import { EJSON } from "bson";
 
 export type SearchIndexStatus = {
     name: string;
-    type: string;
+    type: "search" | "vectorSearch";
     status: string;
     queryable: boolean;
     latestDefinition: Document;
@@ -54,7 +54,7 @@ export class ListSearchIndexesTool extends MongoDBToolBase {
     protected pickRelevantInformation(indexes: Record<string, unknown>[]): SearchIndexStatus[] {
         return indexes.map((index) => ({
             name: (index["name"] ?? "default") as string,
-            type: (index["type"] ?? "UNKNOWN") as string,
+            type: (index["type"] ?? "UNKNOWN") as "search" | "vectorSearch",
             status: (index["status"] ?? "UNKNOWN") as string,
             queryable: (index["queryable"] ?? false) as boolean,
             latestDefinition: index["latestDefinition"] as Document,

src/transports/base.ts

@@ -16,6 +16,7 @@ import {
 } from "../common/connectionErrorHandler.js";
 import type { CommonProperties } from "../telemetry/types.js";
 import { Elicitation } from "../elicitation.js";
+import { VectorSearchEmbeddings } from "../common/search/vectorSearchEmbeddings.js";
 
 export type TransportRunnerConfig = {
     userConfig: UserConfig;
@@ -89,6 +90,7 @@ export abstract class TransportRunnerBase {
             exportsManager,
             connectionManager,
             keychain: Keychain.root,
+            vectorSearchEmbeddings: new VectorSearchEmbeddings(),
         });
 
         const telemetry = Telemetry.create(session, this.userConfig, this.deviceId, {

tests/integration/helpers.ts

@@ -21,6 +21,7 @@ import { connectionErrorHandler } from "../../src/common/connectionErrorHandler.
 import { Keychain } from "../../src/common/keychain.js";
 import { Elicitation } from "../../src/elicitation.js";
 import type { MockClientCapabilities, createMockElicitInput } from "../utils/elicitationMocks.js";
+import { VectorSearchEmbeddings } from "../../src/common/search/vectorSearchEmbeddings.js";
 
 export const driverOptions = setupDriverConfig({
     config,
@@ -101,6 +102,7 @@ export function setupIntegrationTest(
             exportsManager,
             connectionManager,
             keychain: new Keychain(),
+            vectorSearchEmbeddings: new VectorSearchEmbeddings(),
         });
 
         // Mock hasValidAccessToken for tests

tests/integration/telemetry.test.ts

@@ -8,6 +8,7 @@ import { CompositeLogger } from "../../src/common/logger.js";
 import { MCPConnectionManager } from "../../src/common/connectionManager.js";
 import { ExportsManager } from "../../src/common/exportsManager.js";
 import { Keychain } from "../../src/common/keychain.js";
+import { VectorSearchEmbeddings } from "../../src/common/search/vectorSearchEmbeddings.js";
 
 describe("Telemetry", () => {
     it("should resolve the actual device ID", async () => {
@@ -23,6 +24,7 @@ describe("Telemetry", () => {
                 exportsManager: ExportsManager.init(config, logger),
                 connectionManager: new MCPConnectionManager(config, driverOptions, logger, deviceId),
                 keychain: new Keychain(),
+                vectorSearchEmbeddings: new VectorSearchEmbeddings(),
             }),
             config,
             deviceId

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

@@ -20,6 +20,7 @@ import { ErrorCodes } from "../../../../src/common/errors.js";
 import { Keychain } from "../../../../src/common/keychain.js";
 import { Elicitation } from "../../../../src/elicitation.js";
 import { MongoDbTools } from "../../../../src/tools/mongodb/tools.js";
+import { VectorSearchEmbeddings } from "../../../../src/common/search/vectorSearchEmbeddings.js";
 
 const injectedErrorHandler: ConnectionErrorHandler = (error) => {
     switch (error.code) {
@@ -108,6 +109,7 @@ describe("MongoDBTool implementations", () => {
             exportsManager,
             connectionManager,
             keychain: new Keychain(),
+            vectorSearchEmbeddings: new VectorSearchEmbeddings(),
         });
         const telemetry = Telemetry.create(session, userConfig, deviceId);
 

tests/unit/common/session.test.ts

@@ -9,6 +9,7 @@ import { MCPConnectionManager } from "../../../src/common/connectionManager.js";
 import { ExportsManager } from "../../../src/common/exportsManager.js";
 import { DeviceId } from "../../../src/helpers/deviceId.js";
 import { Keychain } from "../../../src/common/keychain.js";
+import { VectorSearchEmbeddings } from "../../../src/common/search/vectorSearchEmbeddings.js";
 
 vi.mock("@mongosh/service-provider-node-driver");
 
@@ -31,6 +32,7 @@ describe("Session", () => {
             exportsManager: ExportsManager.init(config, logger),
             connectionManager: new MCPConnectionManager(config, driverOptions, logger, mockDeviceId),
             keychain: new Keychain(),
+            vectorSearchEmbeddings: new VectorSearchEmbeddings(),
         });
 
         MockNodeDriverServiceProvider.connect = vi.fn().mockResolvedValue({} as unknown as NodeDriverServiceProvider);

tests/unit/resources/common/debug.test.ts

@@ -9,6 +9,7 @@ import { MCPConnectionManager } from "../../../../src/common/connectionManager.j
 import { ExportsManager } from "../../../../src/common/exportsManager.js";
 import { DeviceId } from "../../../../src/helpers/deviceId.js";
 import { Keychain } from "../../../../src/common/keychain.js";
+import { VectorSearchEmbeddings } from "../../../../src/common/search/vectorSearchEmbeddings.js";
 
 describe("debug resource", () => {
     const logger = new CompositeLogger();
@@ -19,6 +20,7 @@ describe("debug resource", () => {
         exportsManager: ExportsManager.init(config, logger),
         connectionManager: new MCPConnectionManager(config, driverOptions, logger, deviceId),
         keychain: new Keychain(),
+        vectorSearchEmbeddings: new VectorSearchEmbeddings(),
     });
     const telemetry = Telemetry.create(session, { ...config, telemetry: "disabled" }, deviceId);