Skip to content
Open
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -236,6 +236,8 @@ export class Server {
event.properties.read_only_mode = this.userConfig.readOnly ? "true" : "false";
event.properties.disabled_tools = this.userConfig.disabledTools || [];
event.properties.confirmation_required_tools = this.userConfig.confirmationRequiredTools || [];
event.properties.previewFeatures = this.userConfig.previewFeatures;
event.properties.embeddingProviderConfigured = !!this.userConfig.voyageApiKey;
}
if (command === "stop") {
event.properties.runtime_duration_ms = Date.now() - this.startTime;
Expand Down
20 changes: 15 additions & 5 deletions src/telemetry/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,8 @@ export type ServerEventProperties = {
read_only_mode?: TelemetryBoolSet;
disabled_tools?: string[];
confirmation_required_tools?: string[];
previewFeatures?: string[];
embeddingProviderConfigured?: boolean;
};

export type ServerEvent = TelemetryEvent<ServerEventProperties>;
Expand Down Expand Up @@ -141,22 +143,30 @@ export type CommonProperties = {
* For MongoDB tools, this is typically empty, while for Atlas tools, this should include
* the project and organization IDs if available.
*/
export type TelemetryToolMetadata = AtlasMetadata | PerfAdvisorToolMetadata | ConnectionMetadata;
export type TelemetryToolMetadata =
| AtlasMetadata
| ConnectionMetadata
| PerfAdvisorToolMetadata
| AutoEmbeddingsUsageMetadata;

export type AtlasMetadata = {
project_id?: string;
org_id?: string;
};

export type PerfAdvisorToolMetadata = AtlasMetadata & {
operations: string[];
type AtlasLocalToolMetadata = {
atlas_local_deployment_id?: string;
};

export type ConnectionMetadata = AtlasMetadata &
AtlasLocalToolMetadata & {
connection_auth_type?: string;
};

type AtlasLocalToolMetadata = {
atlas_local_deployment_id?: string;
export type PerfAdvisorToolMetadata = AtlasMetadata & {
operations: string[];
};

export type AutoEmbeddingsUsageMetadata = ConnectionMetadata & {
embeddingsGeneratedBy: "mcp" | "mongot";
};
15 changes: 15 additions & 0 deletions src/tools/mongodb/create/insertMany.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import { zEJSON } from "../../args.js";
import { type Document } from "bson";
import { zSupportedEmbeddingParameters } from "../mongodbSchemas.js";
import { ErrorCodes, MongoDBError } from "../../../common/errors.js";
import type { ConnectionMetadata, AutoEmbeddingsUsageMetadata } from "../../../telemetry/types.js";

const zSupportedEmbeddingParametersWithInput = zSupportedEmbeddingParameters.extend({
input: z
Expand Down Expand Up @@ -155,4 +156,18 @@ export class InsertManyTool extends MongoDBToolBase {
}
}
}

protected resolveTelemetryMetadata(
args: ToolArgs<typeof this.argsShape>,
{ result }: { result: CallToolResult }
): ConnectionMetadata | AutoEmbeddingsUsageMetadata {
if ("embeddingParameters" in args && this.config.voyageApiKey) {
return {
...super.resolveTelemetryMetadata(args, { result }),
embeddingsGeneratedBy: "mcp",
};
} else {
return super.resolveTelemetryMetadata(args, { result });
}
}
}
20 changes: 20 additions & 0 deletions src/tools/mongodb/read/aggregate.ts
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ import {
assertVectorSearchFilterFieldsAreIndexed,
type SearchIndex,
} from "../../../helpers/assertVectorSearchFilterFieldsAreIndexed.js";
import type { AutoEmbeddingsUsageMetadata, ConnectionMetadata } from "../../../telemetry/types.js";

const pipelineDescriptionWithVectorSearch = `\
An array of aggregation stages to execute.
Expand Down Expand Up @@ -344,4 +345,23 @@ The aggregation resulted in ${aggResultsCount === undefined ? "indeterminable nu
Returning ${documents.length} documents${appliedLimitText ? ` ${appliedLimitText}` : "."}\
`;
}

protected resolveTelemetryMetadata(
args: ToolArgs<typeof this.argsShape>,
{ result }: { result: CallToolResult }
): ConnectionMetadata | AutoEmbeddingsUsageMetadata {
const [maybeVectorStage] = args.pipeline;
if (
maybeVectorStage &&
(maybeVectorStage as z.infer<typeof VectorSearchStage>)?.["$vectorSearch"]?.embeddingParameters &&
this.config.voyageApiKey
) {
return {
...super.resolveTelemetryMetadata(args, { result }),
embeddingsGeneratedBy: "mcp",
};
} else {
return super.resolveTelemetryMetadata(args, { result });
}
}
}
40 changes: 39 additions & 1 deletion tests/integration/tools/mongodb/create/insertMany.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,9 +14,10 @@ import {
getDataFromUntrustedContent,
defaultTestConfig,
} from "../../../helpers.js";
import { beforeEach, afterEach, expect, it, describe } from "vitest";
import { beforeEach, afterEach, expect, it, describe, vi } from "vitest";
import { ObjectId } from "bson";
import type { Collection } from "mongodb";
import type { ToolEvent } from "../../../../../src/telemetry/types.js";

describeWithMongoDB("insertMany tool when search is disabled", (integration) => {
validateToolMetadata(
Expand Down Expand Up @@ -237,6 +238,7 @@ describeWithMongoDB(

afterEach(async () => {
await collection.drop();
vi.clearAllMocks();
});

it("generates embeddings for a single document with one field", async () => {
Expand Down Expand Up @@ -627,6 +629,42 @@ describeWithMongoDB(
// Verify embeddings are different for different text
expect(doc?.titleEmbeddings).not.toEqual(doc?.plotEmbeddings);
});

it("should emit tool event with auto-embedding usage metadata", async () => {
const mockEmitEvents = vi.spyOn(integration.mcpServer()["telemetry"], "emitEvents");
vi.spyOn(integration.mcpServer()["telemetry"], "isTelemetryEnabled").mockReturnValue(true);

await createVectorSearchIndexAndWait(integration.mongoClient(), database, "test", [
{
type: "vector",
path: "titleEmbeddings",
numDimensions: 1024,
similarity: "cosine",
quantization: "scalar",
},
]);

const response = await integration.mcpClient().callTool({
name: "insert-many",
arguments: {
database,
collection: "test",
documents: [{ title: "The Matrix" }],
embeddingParameters: {
model: "voyage-3.5-lite",
input: [{ titleEmbeddings: "The Matrix" }],
},
},
});

const content = getResponseContent(response.content);
expect(content).toContain("Documents were inserted successfully.");

expect(mockEmitEvents).toHaveBeenCalled();
const emittedEvent = mockEmitEvents.mock.lastCall?.[0][0] as ToolEvent;
expectDefined(emittedEvent);
expect(emittedEvent.properties.embeddingsGeneratedBy).toBe("mcp");
});
});
},
{
Expand Down
50 changes: 50 additions & 0 deletions tests/integration/tools/mongodb/read/aggregate.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import {
validateThrowsForInvalidArguments,
getResponseContent,
defaultTestConfig,
expectDefined,
} from "../../../helpers.js";
import { beforeEach, describe, expect, it, vi, afterEach } from "vitest";
import {
Expand All @@ -17,6 +18,7 @@ import * as constants from "../../../../../src/helpers/constants.js";
import { freshInsertDocuments } from "./find.test.js";
import { BSON } from "bson";
import { DOCUMENT_EMBEDDINGS } from "./vyai/embeddings.js";
import type { ToolEvent } from "../../../../../src/telemetry/types.js";

describeWithMongoDB("aggregate tool", (integration) => {
afterEach(() => {
Expand Down Expand Up @@ -393,6 +395,10 @@ describeWithMongoDB(
await integration.mongoClient().db(integration.randomDbName()).collection("databases").drop();
});

afterEach(() => {
vi.clearAllMocks();
});

validateToolMetadata(integration, "aggregate", "Run an aggregation against a MongoDB collection", "read", [
...databaseCollectionParameters,
{
Expand Down Expand Up @@ -462,6 +468,50 @@ If the user requests additional filtering, include filters in \`$vectorSearch.fi
);
});

it("should emit tool event with auto-embedding usage metadata", async () => {
const mockEmitEvents = vi.spyOn(integration.mcpServer()["telemetry"], "emitEvents");
vi.spyOn(integration.mcpServer()["telemetry"], "isTelemetryEnabled").mockReturnValue(true);

await waitUntilSearchIsReady(integration.mongoClient());
await createVectorSearchIndexAndWait(integration.mongoClient(), integration.randomDbName(), "databases", [
{
type: "vector",
path: "description_embedding",
numDimensions: 256,
similarity: "cosine",
quantization: "none",
},
]);

await integration.mcpClient().callTool({
name: "aggregate",
arguments: {
database: integration.randomDbName(),
collection: "databases",
pipeline: [
{
$vectorSearch: {
index: "default",
path: "description_embedding",
queryVector: "some data",
numCandidates: 10,
limit: 10,
embeddingParameters: {
model: "voyage-3-large",
outputDimension: "256",
},
},
},
],
},
});

expect(mockEmitEvents).toHaveBeenCalled();
const emittedEvent = mockEmitEvents.mock.lastCall?.[0][0] as ToolEvent;
expectDefined(emittedEvent);
expect(emittedEvent.properties.embeddingsGeneratedBy).toBe("mcp");
});

for (const [dataType, embedding] of Object.entries(DOCUMENT_EMBEDDINGS)) {
for (const similarity of ["euclidean", "cosine", "dotProduct"]) {
describe(`querying with dataType ${dataType} and similarity ${similarity}`, () => {
Expand Down