Feature/MCP (Model Context Protocol) (#4134)

add mcp tools
2026-06-29 01:01:11 +03:00 · 2025-03-06 13:57:18 +00:00
parent 9c22bee991
commit 713ed26971
21 changed files with 1321 additions and 105 deletions
@@ -0,0 +1,104 @@
+import { CallToolRequest, CallToolResultSchema, ListToolsResult, ListToolsResultSchema } from '@modelcontextprotocol/sdk/types.js'
+import { Client } from '@modelcontextprotocol/sdk/client/index.js'
+import { StdioClientTransport, StdioServerParameters } from '@modelcontextprotocol/sdk/client/stdio.js'
+import { BaseToolkit, tool, Tool } from '@langchain/core/tools'
+import { z } from 'zod'
+
+export class MCPToolkit extends BaseToolkit {
+    tools: Tool[] = []
+    _tools: ListToolsResult | null = null
+    model_config: any
+    transport: StdioClientTransport | null = null
+    client: Client | null = null
+    constructor(serverParams: StdioServerParameters | any, transport: 'stdio' | 'sse') {
+        super()
+        if (transport === 'stdio') {
+            this.transport = new StdioClientTransport(serverParams as StdioServerParameters)
+        } else {
+            //this.transport = new SSEClientTransport(serverParams.url);
+        }
+    }
+    async initialize() {
+        if (this._tools === null) {
+            this.client = new Client(
+                {
+                    name: 'langchain-js-client',
+                    version: '1.0.0'
+                },
+                {
+                    capabilities: {}
+                }
+            )
+            if (this.transport === null) {
+                throw new Error('Transport is not initialized')
+            }
+            await this.client.connect(this.transport)
+            this._tools = await this.client.request({ method: 'tools/list' }, ListToolsResultSchema)
+
+            this.tools = await this.get_tools()
+        }
+    }
+
+    async get_tools(): Promise<Tool[]> {
+        if (this._tools === null || this.client === null) {
+            throw new Error('Must initialize the toolkit first')
+        }
+        const toolsPromises = this._tools.tools.map(async (tool: any) => {
+            if (this.client === null) {
+                throw new Error('Client is not initialized')
+            }
+            return await MCPTool({
+                client: this.client,
+                name: tool.name,
+                description: tool.description || '',
+                argsSchema: createSchemaModel(tool.inputSchema)
+            })
+        })
+        return Promise.all(toolsPromises)
+    }
+}
+
+export async function MCPTool({
+    client,
+    name,
+    description,
+    argsSchema
+}: {
+    client: Client
+    name: string
+    description: string
+    argsSchema: any
+}): Promise<Tool> {
+    return tool(
+        async (input): Promise<string> => {
+            const req: CallToolRequest = { method: 'tools/call', params: { name: name, arguments: input } }
+            const res = await client.request(req, CallToolResultSchema)
+            const content = res.content
+            const contentString = JSON.stringify(content)
+            return contentString
+        },
+        {
+            name: name,
+            description: description,
+            schema: argsSchema
+        }
+    )
+}
+
+function createSchemaModel(
+    inputSchema: {
+        type: 'object'
+        properties?: import('zod').objectOutputType<{}, import('zod').ZodTypeAny, 'passthrough'> | undefined
+    } & { [k: string]: unknown }
+): any {
+    if (inputSchema.type !== 'object' || !inputSchema.properties) {
+        throw new Error('Invalid schema type or missing properties')
+    }
+
+    const schemaProperties = Object.entries(inputSchema.properties).reduce((acc, [key, _]) => {
+        acc[key] = z.any()
+        return acc
+    }, {} as Record<string, import('zod').ZodTypeAny>)
+
+    return z.object(schemaProperties)
+}