-
Couldn't load subscription status.
- Fork 149
Add MongoDB Assistant Tools #472
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
2338504
881c4ee
3bdd3df
59b3c63
34d50d8
bb5e585
d3b4d6b
2736bca
f8365ac
3cb9883
4671347
4d887c9
caa7792
1415205
f4bec2c
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,58 @@ | ||
| import { | ||
| ToolBase, | ||
| type TelemetryToolMetadata, | ||
| type ToolArgs, | ||
| type ToolCategory, | ||
| type ToolConstructorParams, | ||
| } from "../tool.js"; | ||
| import { createFetch } from "@mongodb-js/devtools-proxy-support"; | ||
| import { Server } from "../../server.js"; | ||
| import { packageInfo } from "../../common/packageInfo.js"; | ||
| import { formatUntrustedData } from "../tool.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); | ||
| this.requiredHeaders = new Headers({ | ||
| "x-request-origin": "mongodb-mcp-server", | ||
| "user-agent": packageInfo.version ? `mongodb-mcp-server/v${packageInfo.version}` : "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"); | ||
| } | ||
|
|
||
| // Use the same custom fetch implementation as the Atlas API client. | ||
| // We need this to support enterprise proxies. | ||
| const customFetch = createFetch({ | ||
| useEnvironmentVariableProxies: true, | ||
| }) as unknown as typeof fetch; | ||
|
|
||
| return await customFetch(endpoint, { | ||
| method: args.method, | ||
| headers, | ||
| body: JSON.stringify(args.body), | ||
| }); | ||
| } | ||
| } | ||
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
| @@ -0,0 +1,74 @@ | ||||||
| import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; | ||||||
| import { formatUntrustedData, type OperationType } from "../tool.js"; | ||||||
| import { AssistantToolBase } from "./assistantTool.js"; | ||||||
| import { LogId } from "../../common/logger.js"; | ||||||
| import { stringify as yamlStringify } from "yaml"; | ||||||
|
|
||||||
| export type KnowledgeSource = { | ||||||
| /** The name of the data source */ | ||||||
| id: string; | ||||||
| /** The type of the data source */ | ||||||
| type: string; | ||||||
| /** A list of available versions for this data source */ | ||||||
| versions: { | ||||||
| /** The version label of the data source */ | ||||||
| label: string; | ||||||
| /** Whether this version is the current/default version */ | ||||||
| isCurrent: boolean; | ||||||
| }[]; | ||||||
| }; | ||||||
|
|
||||||
| export type ListKnowledgeSourcesResponse = { | ||||||
| dataSources: KnowledgeSource[]; | ||||||
| }; | ||||||
|
|
||||||
| export const ListKnowledgeSourcesToolName = "list-knowledge-sources"; | ||||||
|
|
||||||
| export class ListKnowledgeSourcesTool extends AssistantToolBase { | ||||||
| public name = ListKnowledgeSourcesToolName; | ||||||
| protected description = "List available data sources in the MongoDB Assistant knowledge base"; | ||||||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. When testing locally I was able to trigger it explicitly by "listing knowledge topics. "list mongodb topics for me"
Suggested change
|
||||||
| 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 } = (await response.json()) as ListKnowledgeSourcesResponse; | ||||||
|
|
||||||
| const text = yamlStringify( | ||||||
| dataSources.map((ds) => { | ||||||
| const currentVersion = ds.versions.find(({ isCurrent }) => isCurrent)?.label; | ||||||
| if (currentVersion) { | ||||||
| (ds as KnowledgeSource & { currentVersion: string }).currentVersion = currentVersion; | ||||||
| } | ||||||
| return ds; | ||||||
| }) | ||||||
| ); | ||||||
|
|
||||||
| return { | ||||||
| content: formatUntrustedData( | ||||||
| `Found ${dataSources.length} data sources in the MongoDB Assistant knowledge base.`, | ||||||
| text | ||||||
| ), | ||||||
| }; | ||||||
| } | ||||||
| } | ||||||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,90 @@ | ||
| import { z } from "zod"; | ||
| import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; | ||
| import { type ToolArgs, type OperationType, formatUntrustedData } from "../tool.js"; | ||
| import { AssistantToolBase } from "./assistantTool.js"; | ||
| import { LogId } from "../../common/logger.js"; | ||
| import { stringify as yamlStringify } from "yaml"; | ||
| import { ListKnowledgeSourcesToolName } from "./listKnowledgeSources.js"; | ||
|
|
||
| export const SearchKnowledgeToolArgs = { | ||
| query: z | ||
| .string() | ||
| .describe( | ||
| "A natural language query to search for in the MongoDB Assistant knowledge base. This should be a single question or a topic that is relevant to the user's MongoDB use case." | ||
| ), | ||
| limit: z.number().min(1).max(100).optional().default(5).describe("The maximum number of results to return"), | ||
nlarew marked this conversation as resolved.
Show resolved
Hide resolved
nlarew marked this conversation as resolved.
Show resolved
Hide resolved
nlarew marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| 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 limit the search to. 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. Available data sources and their versions can be listed by calling the ${ListKnowledgeSourcesToolName} tool.` | ||
| ), | ||
| }; | ||
|
|
||
| export type SearchKnowledgeResponse = { | ||
| /** A list of search results */ | ||
| results: { | ||
| /** The URL of the search result */ | ||
| url: string; | ||
| /** The page title of the search result */ | ||
| title: string; | ||
| /** The text of the page chunk returned from the search */ | ||
| text: string; | ||
| /** Metadata for the search result */ | ||
| metadata: { | ||
| /** A list of tags that describe the page */ | ||
| tags: string[]; | ||
| /** Additional metadata */ | ||
| [key: string]: unknown; | ||
| }; | ||
| }[]; | ||
| }; | ||
|
|
||
| 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 } = (await response.json()) as SearchKnowledgeResponse; | ||
|
|
||
| const text = yamlStringify(results); | ||
|
|
||
| return { | ||
| content: formatUntrustedData( | ||
| `Found ${results.length} results in the MongoDB Assistant knowledge base.`, | ||
| text | ||
| ), | ||
| }; | ||
| } | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,4 @@ | ||
| import { ListKnowledgeSourcesTool } from "./listKnowledgeSources.js"; | ||
| import { SearchKnowledgeTool } from "./searchKnowledge.js"; | ||
|
|
||
| export const AssistantTools = [ListKnowledgeSourcesTool, SearchKnowledgeTool]; |
Uh oh!
There was an error while loading. Please reload this page.