Skip to content
Open
Show file tree
Hide file tree
Changes from all 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
25 changes: 20 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,35 @@ 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 & {
/**
* Indicates which component generated the embeddings.
* "mcp" is used when embeddings are generated by the MCP server.
* "mongot" is reserved for future use, when embeddings may be generated by MongoDB's mongot process.
*/
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