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 @@
 } 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 @@
     }
 
     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 @@
                     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,48 @@
+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 { Session } from "../../common/session.js";
+import { UserConfig } from "../../common/config.js";
+import { Telemetry } from "../../telemetry/telemetry.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: Record<string, string>;
+
+    constructor({ session, config, telemetry, elicitation }: ToolConstructorParams) {
+        super({ session, config, telemetry, elicitation });
+        this.baseUrl = new URL(config.assistantBaseUrl);
+        const serverVersion = packageInfo.version;
+        this.requiredHeaders = {
+            "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 handleError(
+        error: unknown,
+        args: ToolArgs<typeof this.argsShape>
+    ): Promise<CallToolResult> | CallToolResult {
+        return super.handleError(error, args);
+    }
+}

src/tools/assistant/listKnowledgeSources.ts

@@ -0,0 +1,66 @@
+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 searchEndpoint = new URL("content/sources", this.baseUrl);
+        const response = await fetch(searchEndpoint, {
+            method: "GET",
+            headers: this.requiredHeaders,
+        });
+        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,83 @@
+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 searchEndpoint = new URL("content/search", this.baseUrl);
+        const response = await fetch(searchEndpoint, {
+            method: "POST",
+            headers: new Headers({ ...this.requiredHeaders, "Content-Type": "application/json" }),
+            body: JSON.stringify(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");

tests/integration/tools/assistant/assistantHelpers.ts

@@ -0,0 +1,78 @@
+import { setupIntegrationTest, IntegrationTest, defaultTestConfig } from "../../helpers.js";
+import { describe, SuiteCollector } from "vitest";
+import { vi, beforeAll, afterAll } from "vitest";
+
+export type IntegrationTestFunction = (integration: IntegrationTest) => void;
+
+export function describeWithAssistant(name: string, fn: IntegrationTestFunction): SuiteCollector<object> {
+    const testDefinition = (): void => {
+        const integration = setupIntegrationTest(() => ({
+            ...defaultTestConfig,
+            assistantBaseUrl: "https://knowledge.test.mongodb.com/api/", // Use test URL
+        }));
+
+        describe(name, () => {
+            fn(integration);
+        });
+    };
+
+    // eslint-disable-next-line vitest/valid-describe-callback
+    return describe("assistant", testDefinition);
+}
+
+/**
+ * Mocks fetch for assistant API calls
+ */
+interface MockedAssistantAPI {
+    mockListSources: (sources: unknown[]) => void;
+    mockSearchResults: (results: unknown[]) => void;
+    mockAPIError: (status: number, statusText: string) => void;
+    mockNetworkError: (error: Error) => void;
+    mockFetch: ReturnType<typeof vi.fn>;
+}
+
+export function makeMockAssistantAPI(): MockedAssistantAPI {
+    const mockFetch = vi.fn();
+
+    beforeAll(() => {
+        global.fetch = mockFetch;
+    });
+
+    afterAll(() => {
+        vi.restoreAllMocks();
+    });
+
+    const mockListSources: MockedAssistantAPI["mockListSources"] = (sources) => {
+        mockFetch.mockResolvedValueOnce({
+            ok: true,
+            json: () => Promise.resolve({ dataSources: sources }),
+        });
+    };
+
+    const mockSearchResults: MockedAssistantAPI["mockSearchResults"] = (results) => {
+        mockFetch.mockResolvedValueOnce({
+            ok: true,
+            json: () => Promise.resolve({ results }),
+        });
+    };
+
+    const mockAPIError: MockedAssistantAPI["mockAPIError"] = (status, statusText) => {
+        mockFetch.mockResolvedValueOnce({
+            ok: false,
+            status,
+            statusText,
+        });
+    };
+
+    const mockNetworkError: MockedAssistantAPI["mockNetworkError"] = (error) => {
+        mockFetch.mockRejectedValueOnce(error);
+    };
+
+    return {
+        mockListSources,
+        mockSearchResults,
+        mockAPIError,
+        mockNetworkError,
+        mockFetch,
+    };
+}