thewimo commited on
Commit
b1d8992
·
verified ·
1 Parent(s): f5b2847

Create custom_mcp.py

Browse files
Files changed (1) hide show
  1. custom_mcp.py +45 -0
custom_mcp.py ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from smolagents.mcp_client import MCPClient
3
+ from mcpadapt.smolagents_adapter import SmolAgentsAdapter
4
+ from mcpadapt.core import MCPAdapt
5
+ from mcp import StdioServerParameters
6
+ from mcp.types import Tool
7
+ from typing import Callable, Any
8
+ import smolagents
9
+ from functools import partial
10
+
11
+ # Custom adapter to fix tool names
12
+ class CustomSmolAgentsAdapter(SmolAgentsAdapter):
13
+ def adapt(
14
+ self,
15
+ func: Callable[[dict | None], Any],
16
+ mcp_tool: Tool
17
+ ) -> smolagents.Tool:
18
+ # Fix the tool name before passing it to the original adapter
19
+ fixed_tool = Tool(
20
+ name=mcp_tool.name.replace('-', '_'),
21
+ description=mcp_tool.description,
22
+ inputSchema=mcp_tool.inputSchema
23
+ )
24
+ return super().adapt(func, fixed_tool)
25
+
26
+ # Custom client that uses our adapter
27
+ def create_mcp_client(server_params):
28
+ adapter = MCPAdapt(server_params, CustomSmolAgentsAdapter())
29
+
30
+ # Create a client-like object with the same interface
31
+ class CustomMCPClient:
32
+ def __init__(self):
33
+ self._tools = None
34
+ self.connect()
35
+
36
+ def connect(self):
37
+ self._tools = adapter.__enter__()
38
+
39
+ def disconnect(self, *args, **kwargs):
40
+ adapter.__exit__(None, None, None)
41
+
42
+ def get_tools(self):
43
+ return self._tools
44
+
45
+ return CustomMCPClient()