Add MongoDB Assistant Tools

README.md

@@ -15,6 +15,7 @@ A Model Context Protocol server for interacting with MongoDB Databases and Mongo
 - [🛠️ Supported Tools](#supported-tools)
   - [MongoDB Atlas Tools](#mongodb-atlas-tools)
   - [MongoDB Database Tools](#mongodb-database-tools)
+  - [MongoDB Assistant Tools](#mongodb-assistant-tools)
 - [📄 Supported Resources](#supported-resources)
 - [⚙️ Configuration](#configuration)
   - [Configuration Options](#configuration-options)
@@ -321,6 +322,11 @@ NOTE: atlas tools are only available when you set credentials on [configuration]
 - `db-stats` - Return statistics about a MongoDB database
 - `export` - Export query or aggregation results to EJSON format. Creates a uniquely named export accessible via the `exported-data` resource.
 
+#### MongoDB Assistant Tools
+
+- `list-knowledge-sources` - List available data sources in the MongoDB Assistant knowledge base
+- `search-knowledge` - Search for information in the MongoDB Assistant knowledge base
+
 ## 📄 Supported Resources
 
 - `config` - Server configuration, supplied by the user either as environment variables or as startup arguments with sensitive parameters redacted. The resource can be accessed under URI `config://config`.

src/common/config.ts

@@ -160,6 +160,7 @@ export interface UserConfig extends CliOptions {
     apiBaseUrl: string;
     apiClientId?: string;
     apiClientSecret?: string;
+    assistantBaseUrl: string;
     telemetry: "enabled" | "disabled";
     logPath: string;
     exportsPath: string;
@@ -185,6 +186,7 @@ export interface UserConfig extends CliOptions {
 
 export const defaultUserConfig: UserConfig = {
     apiBaseUrl: "https://cloud.mongodb.com/",
+    assistantBaseUrl: "https://knowledge.mongodb.com/api/v1/",
     logPath: getLogPath(),
     exportsPath: getExportsPath(),
     exportTimeoutMs: 5 * 60 * 1000, // 5 minutes

src/common/logger.ts

@@ -70,6 +70,9 @@ export const LogId = {
     exportLockError: mongoLogId(1_007_008),
 
     oidcFlow: mongoLogId(1_008_001),
+
+    assistantListKnowledgeSourcesError: mongoLogId(1_009_001),
+    assistantSearchKnowledgeError: mongoLogId(1_009_002),
 } as const;
 
 export interface LogPayload {

src/server.ts

@@ -19,6 +19,7 @@ import {
 } from "@modelcontextprotocol/sdk/types.js";
 import assert from "assert";
 import type { ToolBase } from "./tools/tool.js";
+import { AssistantTools } from "./tools/assistant/tools.js";
 import { validateConnectionString } from "./helpers/connectionOptions.js";
 import { packageInfo } from "./common/packageInfo.js";
 import { type ConnectionErrorHandler } from "./common/connectionErrorHandler.js";
@@ -206,7 +207,7 @@ export class Server {
     }
 
     private registerTools(): void {
-        for (const toolConstructor of [...AtlasTools, ...MongoDbTools]) {
+        for (const toolConstructor of [...AtlasTools, ...MongoDbTools, ...AssistantTools]) {
             const tool = new toolConstructor({
                 session: this.session,
                 config: this.userConfig,
@@ -302,6 +303,7 @@ export class Server {
                     context: "server",
                     message: `Failed to connect to MongoDB instance using the connection string from the config: ${error instanceof Error ? error.message : String(error)}`,
                 });
+                j;
             }
         }
     }

src/tools/assistant/assistantTool.ts

@@ -0,0 +1,51 @@
+import {
+    ToolBase,
+    type TelemetryToolMetadata,
+    type ToolArgs,
+    type ToolCategory,
+    type ToolConstructorParams,
+} from "../tool.js";
+import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import { Server } from "../../server.js";
+import { packageInfo } from "../../common/packageInfo.js";
+
+export abstract class AssistantToolBase extends ToolBase {
+    protected server?: Server;
+    public category: ToolCategory = "assistant";
+    protected baseUrl: URL;
+    protected requiredHeaders: Headers;
+
+    constructor({ session, config, telemetry, elicitation }: ToolConstructorParams) {
+        super({ session, config, telemetry, elicitation });
+        this.baseUrl = new URL(config.assistantBaseUrl);
+        const serverVersion = packageInfo.version;
+        this.requiredHeaders = new Headers({
+            "x-request-origin": "mongodb-mcp-server",
+            "user-agent": serverVersion ? `mongodb-mcp-server/v${serverVersion}` : "mongodb-mcp-server",
+        });
+    }
+
+    public register(server: Server): boolean {
+        this.server = server;
+        return super.register(server);
+    }
+
+    protected resolveTelemetryMetadata(_args: ToolArgs<typeof this.argsShape>): TelemetryToolMetadata {
+        // Assistant tool calls are not associated with a specific project or organization
+        // Therefore, we don't have any values to add to the telemetry metadata
+        return {};
+    }
+
+    protected async callAssistantApi(args: { method: "GET" | "POST"; endpoint: string; body?: unknown }) {
+        const endpoint = new URL(args.endpoint, this.baseUrl);
+        const headers = new Headers(this.requiredHeaders);
+        if (args.method === "POST") {
+            headers.set("Content-Type", "application/json");
+        }
+        return await fetch(endpoint, {
+            method: args.method,
+            headers,
+            body: JSON.stringify(args.body),
+        });
+    }
+}

src/tools/assistant/listKnowledgeSources.ts

@@ -0,0 +1,65 @@
+import { z } from "zod";
+import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import { OperationType } from "../tool.js";
+import { AssistantToolBase } from "./assistantTool.js";
+import { LogId } from "../../common/logger.js";
+
+export const dataSourceMetadataSchema = z.object({
+    id: z.string().describe("The name of the data source"),
+    type: z.string().optional().describe("The type of the data source"),
+    versions: z
+        .array(
+            z.object({
+                label: z.string().describe("The version label of the data source"),
+                isCurrent: z.boolean().describe("Whether this version is current active version"),
+            })
+        )
+        .describe("A list of available versions for this data source"),
+});
+
+export const listDataSourcesResponseSchema = z.object({
+    dataSources: z.array(dataSourceMetadataSchema).describe("A list of data sources"),
+});
+
+export class ListKnowledgeSourcesTool extends AssistantToolBase {
+    public name = "list-knowledge-sources";
+    protected description = "List available data sources in the MongoDB Assistant knowledge base";
+    protected argsShape = {};
+    public operationType: OperationType = "read";
+
+    protected async execute(): Promise<CallToolResult> {
+        const response = await this.callAssistantApi({
+            method: "GET",
+            endpoint: "content/sources",
+        });
+        if (!response.ok) {
+            const message = `Failed to list knowledge sources: ${response.statusText}`;
+            this.session.logger.debug({
+                id: LogId.assistantListKnowledgeSourcesError,
+                context: "assistant-list-knowledge-sources",
+                message,
+            });
+            return {
+                content: [
+                    {
+                        type: "text",
+                        text: message,
+                    },
+                ],
+                isError: true,
+            };
+        }
+        const { dataSources } = listDataSourcesResponseSchema.parse(await response.json());
+
+        return {
+            content: dataSources.map(({ id, type, versions }) => ({
+                type: "text",
+                text: id,
+                _meta: {
+                    type,
+                    versions,
+                },
+            })),
+        };
+    }
+}

src/tools/assistant/searchKnowledge.ts

@@ -0,0 +1,82 @@
+import { z } from "zod";
+import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import { ToolArgs, OperationType } from "../tool.js";
+import { AssistantToolBase } from "./assistantTool.js";
+import { LogId } from "../../common/logger.js";
+
+export const SearchKnowledgeToolArgs = {
+    query: z.string().describe("A natural language query to search for in the knowledge base"),
+    limit: z.number().min(1).max(100).optional().default(5).describe("The maximum number of results to return"),
+    dataSources: z
+        .array(
+            z.object({
+                name: z.string().describe("The name of the data source"),
+                versionLabel: z.string().optional().describe("The version label of the data source"),
+            })
+        )
+        .optional()
+        .describe(
+            "A list of one or more data sources to search in. You can specify a specific version of a data source by providing the version label. If not provided, the latest version of all data sources will be searched."
+        ),
+};
+
+export const knowledgeChunkSchema = z
+    .object({
+        url: z.string().describe("The URL of the search result"),
+        title: z.string().describe("Title of the search result"),
+        text: z.string().describe("Chunk text"),
+        metadata: z
+            .object({
+                tags: z.array(z.string()).describe("The tags of the source"),
+            })
+            .passthrough(),
+    })
+    .passthrough();
+
+export const searchResponseSchema = z.object({
+    results: z.array(knowledgeChunkSchema).describe("A list of search results"),
+});
+
+export class SearchKnowledgeTool extends AssistantToolBase {
+    public name = "search-knowledge";
+    protected description = "Search for information in the MongoDB Assistant knowledge base";
+    protected argsShape = {
+        ...SearchKnowledgeToolArgs,
+    };
+    public operationType: OperationType = "read";
+
+    protected async execute(args: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
+        const response = await this.callAssistantApi({
+            method: "POST",
+            endpoint: "content/search",
+            body: args,
+        });
+        if (!response.ok) {
+            const message = `Failed to search knowledge base: ${response.statusText}`;
+            this.session.logger.debug({
+                id: LogId.assistantSearchKnowledgeError,
+                context: "assistant-search-knowledge",
+                message,
+            });
+            return {
+                content: [
+                    {
+                        type: "text",
+                        text: message,
+                    },
+                ],
+                isError: true,
+            };
+        }
+        const { results } = searchResponseSchema.parse(await response.json());
+        return {
+            content: results.map(({ text, metadata }) => ({
+                type: "text",
+                text,
+                _meta: {
+                    ...metadata,
+                },
+            })),
+        };
+    }
+}

src/tools/assistant/tools.ts

@@ -0,0 +1,4 @@
+import { ListKnowledgeSourcesTool } from "./listKnowledgeSources.js";
+import { SearchKnowledgeTool } from "./searchKnowledge.js";
+
+export const AssistantTools = [ListKnowledgeSourcesTool, SearchKnowledgeTool];

src/tools/tool.ts

@@ -16,7 +16,7 @@ export type ToolCallbackArgs<Args extends ZodRawShape> = Parameters<ToolCallback
 export type ToolExecutionContext<Args extends ZodRawShape = ZodRawShape> = Parameters<ToolCallback<Args>>[1];
 
 export type OperationType = "metadata" | "read" | "create" | "delete" | "update" | "connect";
-export type ToolCategory = "mongodb" | "atlas";
+export type ToolCategory = "mongodb" | "atlas" | "assistant";
 export type TelemetryToolMetadata = {
     projectId?: string;
     orgId?: string;

tests/integration/helpers.ts

@@ -179,22 +179,26 @@ export function setupIntegrationTest(
     };
 }
 
-// eslint-disable-next-line @typescript-eslint/no-redundant-type-constituents
-export function getResponseContent(content: unknown | { content: unknown }): string {
+export function getResponseContent(content: unknown): string {
     return getResponseElements(content)
         .map((item) => item.text)
         .join("\n");
 }
 
-// eslint-disable-next-line @typescript-eslint/no-redundant-type-constituents
-export function getResponseElements(content: unknown | { content: unknown }): { type: string; text: string }[] {
+export interface ResponseElement {
+    type: string;
+    text: string;
+    _meta?: unknown;
+}
+
+export function getResponseElements(content: unknown): ResponseElement[] {
     if (typeof content === "object" && content !== null && "content" in content) {
-        content = (content as { content: unknown }).content;
+        content = content.content;
     }
 
     expect(content).toBeInstanceOf(Array);
 
-    const response = content as { type: string; text: string }[];
+    const response = content as ResponseElement[];
     for (const item of response) {
         expect(item).toHaveProperty("type");
         expect(item).toHaveProperty("text");