Skip to content

Commit 6bee4bc

Browse files
vm-mishchenkovm-mishchenko
committed
Add RunPipeline tool
1 parent 816dec1 commit 6bee4bc

File tree

5 files changed

+218
-2
lines changed

5 files changed

+218
-2
lines changed

‎src/server.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import { Session } from "./session.js";
33
import { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
44
import { AtlasTools } from "./tools/atlas/tools.js";
55
import { MongoDbTools } from "./tools/mongodb/tools.js";
6+
import { PlaygroundTools } from "./tools/playground/tools.js";
67
import logger, { initializeLogger, LogId } from "./logger.js";
78
import { ObjectId } from "mongodb";
89
import { Telemetry } from "./telemetry/telemetry.js";
@@ -134,7 +135,7 @@ export class Server {
134135
}
135136

136137
private registerTools() {
137-
for (const tool of [...AtlasTools, ...MongoDbTools]) {
138+
for (const tool of [...AtlasTools, ...MongoDbTools, ...PlaygroundTools]) {
138139
new tool(this.session, this.userConfig, this.telemetry).register(this.mcpServer);
139140
}
140141
}

‎src/tools/playground/runPipeline.ts

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
import { OperationType, TelemetryToolMetadata, ToolArgs, ToolBase, ToolCategory } from "../tool.js";
2+
import { z } from "zod";
3+
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
4+
import { EJSON } from "bson";
5+
6+
const PLAYGROUND_SEARCH_URL = "https://search-playground.mongodb.com/api/tools/code-playground/search";
7+
8+
const DEFAULT_DOCUMENTS = [
9+
{
10+
name: "First document",
11+
},
12+
{
13+
name: "Second document",
14+
},
15+
];
16+
17+
const DEFAULT_SEARCH_INDEX_DEFINITION = {
18+
mappings: {
19+
dynamic: true,
20+
},
21+
};
22+
23+
const DEFAULT_PIPELINE = [
24+
{
25+
$search: {
26+
index: "default",
27+
text: {
28+
query: "first",
29+
path: {
30+
wildcard: "*",
31+
},
32+
},
33+
},
34+
},
35+
];
36+
37+
const DEFAULT_SYNONYMS: Array<Record<string, unknown>> = [];
38+
39+
export const RunPipelineOperationArgs = {
40+
documents: z
41+
.array(z.record(z.string(), z.unknown()))
42+
.describe("Documents to run the pipeline against. 500 is maximum.")
43+
.default(DEFAULT_DOCUMENTS),
44+
aggregationPipeline: z
45+
.array(z.record(z.string(), z.unknown()))
46+
.describe("Aggregation pipeline to run on the provided documents.")
47+
.default(DEFAULT_PIPELINE),
48+
searchIndexDefinition: z
49+
.record(z.string(), z.unknown())
50+
.describe("Search index to create before running the pipeline.")
51+
.optional()
52+
.default(DEFAULT_SEARCH_INDEX_DEFINITION),
53+
synonyms: z
54+
.array(z.record(z.any()))
55+
.describe("Synonyms mapping to create before running the pipeline.")
56+
.optional()
57+
.default(DEFAULT_SYNONYMS),
58+
};
59+
60+
interface RunRequest {
61+
documents: string;
62+
aggregationPipeline: string;
63+
indexDefinition: string;
64+
synonyms: string;
65+
}
66+
67+
interface RunResponse {
68+
documents: Array<Record<string, unknown>>;
69+
}
70+
71+
interface RunErrorResponse {
72+
code: string;
73+
message: string;
74+
}
75+
76+
export class RunPipeline extends ToolBase {
77+
protected name = "run-pipeline";
78+
protected description =
79+
"Run aggregation pipeline for provided documents without needing an Atlas account, cluster, or collection.";
80+
protected category: ToolCategory = "playground";
81+
protected operationType: OperationType = "metadata";
82+
protected argsShape = RunPipelineOperationArgs;
83+
84+
protected async execute(toolArgs: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
85+
const runRequest = this.convertToRunRequest(toolArgs);
86+
const runResponse = await this.runPipeline(runRequest);
87+
const toolResult = this.convertToToolResult(runResponse);
88+
return toolResult;
89+
}
90+
91+
protected resolveTelemetryMetadata(): TelemetryToolMetadata {
92+
return {};
93+
}
94+
95+
private async runPipeline(runRequest: RunRequest): Promise<RunResponse> {
96+
const options: RequestInit = {
97+
method: "POST",
98+
headers: {
99+
"Content-Type": "application/json",
100+
},
101+
body: JSON.stringify(runRequest),
102+
};
103+
104+
let response: Response;
105+
try {
106+
response = await fetch(PLAYGROUND_SEARCH_URL, options);
107+
} catch {
108+
throw new Error("Cannot run pipeline: network error.");
109+
}
110+
111+
if (!response.ok) {
112+
const errorMessage = await this.getPlaygroundResponseError(response);
113+
throw new Error(`Pipeline run failed: ${errorMessage}`);
114+
}
115+
116+
try {
117+
return (await response.json()) as RunResponse;
118+
} catch {
119+
throw new Error("Pipeline run failed: response is not valid JSON.");
120+
}
121+
}
122+
123+
private async getPlaygroundResponseError(response: Response): Promise<string> {
124+
let errorMessage = `HTTP ${response.status} ${response.statusText}.`;
125+
try {
126+
const errorResponse = (await response.json()) as RunErrorResponse;
127+
errorMessage += ` Error code: ${errorResponse.code}. Error message: ${errorResponse.message}`;
128+
} catch {
129+
// Ignore JSON parse errors
130+
}
131+
132+
return errorMessage;
133+
}
134+
135+
private convertToRunRequest(toolArgs: ToolArgs<typeof this.argsShape>): RunRequest {
136+
try {
137+
return {
138+
documents: JSON.stringify(toolArgs.documents),
139+
aggregationPipeline: JSON.stringify(toolArgs.aggregationPipeline),
140+
indexDefinition: JSON.stringify(toolArgs.searchIndexDefinition || DEFAULT_SEARCH_INDEX_DEFINITION),
141+
synonyms: JSON.stringify(toolArgs.synonyms || DEFAULT_SYNONYMS),
142+
};
143+
} catch {
144+
throw new Error("Invalid arguments type.");
145+
}
146+
}
147+
148+
private convertToToolResult(runResponse: RunResponse): CallToolResult {
149+
const content: Array<{ text: string; type: "text" }> = [
150+
{
151+
text: `Found ${runResponse.documents.length} documents":`,
152+
type: "text",
153+
},
154+
...runResponse.documents.map((doc) => {
155+
return {
156+
text: EJSON.stringify(doc),
157+
type: "text",
158+
} as { text: string; type: "text" };
159+
}),
160+
];
161+
162+
return {
163+
content,
164+
};
165+
}
166+
}

‎src/tools/playground/tools.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
import { RunPipeline } from "./runPipeline.js";
2+
3+
export const PlaygroundTools = [RunPipeline];

‎src/tools/tool.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import { UserConfig } from "../config.js";
1010
export type ToolArgs<Args extends ZodRawShape> = z.objectOutputType<Args, ZodNever>;
1111

1212
export type OperationType = "metadata" | "read" | "create" | "delete" | "update";
13-
export type ToolCategory = "mongodb" | "atlas";
13+
export type ToolCategory = "mongodb" | "atlas" | "playground";
1414
export type TelemetryToolMetadata = {
1515
projectId?: string;
1616
orgId?: string;

‎tests/integration/tools/playground/runPipeline.test.ts

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
import { describeWithMongoDB } from "../mongodb/mongodbHelpers.js";
2+
import { getResponseElements, validateThrowsForInvalidArguments } from "../../helpers.js";
3+
4+
describeWithMongoDB("runPipeline tool", (integration) => {
5+
validateThrowsForInvalidArguments(integration, "run-pipeline", [{}]);
6+
7+
it("should return results", async () => {
8+
await integration.connectMcpClient();
9+
const response = await integration.mcpClient().callTool({
10+
name: "run-pipeline",
11+
arguments: {
12+
documents: [{ name: "First document" }, { name: "Second document" }],
13+
aggregationPipeline: [
14+
{
15+
$search: {
16+
index: "default",
17+
text: {
18+
query: "first",
19+
path: {
20+
wildcard: "*",
21+
},
22+
},
23+
},
24+
},
25+
{
26+
$project: {
27+
_id: 0,
28+
name: 1,
29+
},
30+
},
31+
],
32+
},
33+
});
34+
const elements = getResponseElements(response.content);
35+
expect(elements).toEqual([
36+
{
37+
text: 'Found 1 documents":',
38+
type: "text",
39+
},
40+
{
41+
text: '{"name":"First document"}',
42+
type: "text",
43+
},
44+
]);
45+
});
46+
});

0 commit comments

Comments
 (0)