feat: MCP initial support for sse servers (Part 2)
Rafael Uzarowski committed
Mar 13, 2025 at 11:25 UTC
d49b09083d455aa658655a88f7b03a6e86b7a646
3 files changed
+98
-53
docker/run/fs/ins/pre_install.sh
+15
-1
@@ -20,7 +20,21 @@ apt-get update && apt-get upgrade -y && apt-get install -y \
20
ffmpeg \
21
nginx\
22
supervisor \
23
- cron
23
+ cron \
24
+ libmagic-dev \
25
+ poppler-utils \
26
+ tesseract-ocr \
27
+ qpdf \
28
+ libreoffice \
29
+ pandoc \
30
+ libgtk-3-0 \
31
+ libnss3 \
32
+ libatk1.0-0 \
33
+ libatk-bridge2.0-0 \
34
+ libcups2 \
35
+ libasound2 \
36
+ libasound2-data \
37
+ cargo
38
39
echo "=====MID UPDATE====="
40
python/helpers/mcp.py
+80
-52
@@ -1,28 +1,25 @@
1
-from pydantic import BaseModel, Field, Discriminator, Tag, PrivateAttr
2
-from typing import List, Dict, Optional, Any, Union, Literal, Annotated
3
-from typing import (
4
- List, Dict, Optional, Any,
5
- Union, Literal, Annotated, ClassVar,
6
-)
1
+from abc import ABC, abstractmethod
2
+from typing import List, Dict, Optional, Any, Union, Literal, Annotated, ClassVar, cast
3
import threading
4
import asyncio
5
from contextlib import AsyncExitStack
6
from shutil import which
7
+from datetime import timedelta
8
+import dirtyjson
9
+import json
10
from mcp import ClientSession, StdioServerParameters
11
from mcp.client.stdio import stdio_client
12
+from mcp.client.sse import sse_client
13
from mcp.types import CallToolResult, ListToolsResult, JSONRPCMessage
14
from anyio.streams.memory import (
15
MemoryObjectReceiveStream,
16
MemoryObjectSendStream,
17
)
18
+
19
+from pydantic import BaseModel, Field, Discriminator, Tag, PrivateAttr
20
from python.helpers.dirty_json import DirtyJson
21
from python.helpers.print_style import PrintStyle
20
-import dirtyjson
21
-
22
from python.helpers.tool import Tool, Response
23
-from datetime import timedelta
24
-
25
-from abc import ABC, abstractmethod
23
24
25
class MCPTool(Tool):
@@ -59,7 +56,7 @@ class MCPTool(Tool):
56
self.log = self.get_log_object()
57
58
for key, value in self.args.items():
62
- PrintStyle(font_color="#85C1E9", bold=True).stream(self.nice_key(key)+": ")
59
+ PrintStyle(font_color="#85C1E9", bold=True).stream(self.nice_key(key) + ": ")
60
PrintStyle(font_color="#85C1E9", padding=isinstance(value, str) and "\n" in value).stream(value)
61
PrintStyle().print()
62
@@ -85,37 +82,46 @@ class MCPServerRemote(BaseModel):
82
description: Optional[str] = Field(default="Remote SSE Server")
83
url: str = Field(default_factory=str)
84
headers: dict[str, Any] | None = Field(default_factory=dict[str, Any])
88
- timeout: float = 5.0
89
- sse_read_timeout: float = 60.0 * 5.0
90
- disabled: bool = False
85
+ timeout: float = Field(default=5.0)
86
+ sse_read_timeout: float = Field(default=60.0 * 5.0)
87
+ disabled: bool = Field(default=False)
88
89
__lock: ClassVar[threading.Lock] = PrivateAttr(default=threading.Lock())
90
+ __client: Optional["MCPClientRemote"] = PrivateAttr(default=None)
91
92
def __init__(self, config: dict[str, Any]):
93
super().__init__()
94
+ self.__client = MCPClientRemote(self)
95
self.update(config)
96
97
def get_tools(self) -> List[dict[str, Any]]:
98
"""Get all tools from the server"""
100
- return []
99
+ with self.__lock:
100
+ return self.__client.tools # type: ignore
101
102
def has_tool(self, tool_name: str) -> bool:
103
"""Check if a tool is available"""
104
- return False
104
+ with self.__lock:
105
+ return self.__client.has_tool(tool_name) # type: ignore
106
107
async def call_tool(self, tool_name: str, input_data: Dict[str, Any]) -> CallToolResult:
108
"""Call a tool with the given input data"""
108
- raise NotImplementedError("MCPServerRemote does not support calling tools")
109
+ with self.__lock:
110
+ # We already run in an event loop, dont believe Pylance
111
+ return await self.__client.call_tool(tool_name, input_data) # type: ignore
112
113
def update(self, config: dict[str, Any]) -> "MCPServerRemote":
114
with self.__lock:
115
for key, value in config.items():
116
if key in ["name", "description", "url", "headers", "timeout", "sse_read_timeout", "disabled"]:
117
+ if key == "name":
118
+ value = value.strip().lower().replace(" ", "_").replace("-", "_").replace(".", "_")
119
setattr(self, key, value)
115
- # We already run in an event loop, dont believe Pylance
116
- return asyncio.run(self.__on_update())
120
+ # We already run in an event loop, dont believe Pylance
121
+ return asyncio.run(self.__on_update())
122
123
async def __on_update(self) -> "MCPServerRemote":
124
+ await self.__client.update_tools() # type: ignore
125
return self
126
127
@@ -125,9 +131,9 @@ class MCPServerLocal(BaseModel):
131
command: str = Field(default_factory=str)
132
args: list[str] = Field(default_factory=list)
133
env: dict[str, str] | None = Field(default_factory=dict[str, str])
128
- encoding: str = "utf-8"
129
- encoding_error_handler: Literal["strict", "ignore", "replace"] = "strict"
130
- disabled: bool = False
134
+ encoding: str = Field(default="utf-8")
135
+ encoding_error_handler: Literal["strict", "ignore", "replace"] = Field(default="strict")
136
+ disabled: bool = Field(default=False)
137
138
__lock: ClassVar[threading.Lock] = PrivateAttr(default=threading.Lock())
139
__client: Optional["MCPClientLocal"] = PrivateAttr(default=None)
@@ -140,18 +146,18 @@ class MCPServerLocal(BaseModel):
146
def get_tools(self) -> List[dict[str, Any]]:
147
"""Get all tools from the server"""
148
with self.__lock:
143
- return self.__client.tools
149
+ return self.__client.tools # type: ignore
150
151
def has_tool(self, tool_name: str) -> bool:
152
"""Check if a tool is available"""
153
with self.__lock:
148
- return self.__client.has_tool(tool_name)
154
+ return self.__client.has_tool(tool_name) # type: ignore
155
156
async def call_tool(self, tool_name: str, input_data: Dict[str, Any]) -> CallToolResult:
157
"""Call a tool with the given input data"""
158
with self.__lock:
159
# We already run in an event loop, dont believe Pylance
154
- return await self.__client.call_tool(tool_name, input_data)
160
+ return await self.__client.call_tool(tool_name, input_data) # type: ignore
161
162
def update(self, config: dict[str, Any]) -> "MCPServerLocal":
163
with self.__lock:
@@ -164,7 +170,7 @@ class MCPServerLocal(BaseModel):
170
return asyncio.run(self.__on_update())
171
172
async def __on_update(self) -> "MCPServerLocal":
167
- await self.__client.update_tools()
173
+ await self.__client.update_tools() # type: ignore
174
return self
175
176
@@ -202,7 +208,7 @@ class MCPConfig(BaseModel):
208
try:
209
servers = DirtyJson.parse_string(config_str)
210
except Exception as e:
205
- raise ValueError(f"Failed to parse MCP config: {e}")
211
+ raise ValueError(f"Failed to parse MCP config: {e}") from e
212
cls.get_instance().__init__(servers_list=servers)
213
cls.__initialized = True
214
return cls.get_instance()
@@ -351,43 +357,30 @@ class MCPConfig(BaseModel):
357
raise ValueError(f"Tool {tool_name} not found")
358
359
354
-class MCPClientLocal:
360
+class MCPClientBase(ABC):
361
session: Optional[ClientSession] = None
362
exit_stack: AsyncExitStack = AsyncExitStack()
363
stdio: Optional[MemoryObjectReceiveStream[JSONRPCMessage | Exception]] = None
364
write: Optional[MemoryObjectSendStream[JSONRPCMessage]] = None
365
366
tools: List[dict[str, Any]] = []
361
- server: Optional[MCPServerLocal] = None
367
+ server: Optional[Union[MCPServerLocal, MCPServerRemote]] = None
368
369
__lock: ClassVar[threading.Lock] = threading.Lock()
370
365
- def __init__(self, server: MCPServerLocal):
371
+ def __init__(self, server: Union[MCPServerLocal, MCPServerRemote]):
372
self.server = server
373
374
+ # Protected method
375
+ @abstractmethod
376
+ async def _connect_client(self) -> tuple[MemoryObjectReceiveStream[JSONRPCMessage | Exception], MemoryObjectSendStream[JSONRPCMessage]]:
377
+ """Connect to an MCP server, init client and save stdio/write streams"""
378
+ ...
379
+
380
async def __connect_to_server(self) -> Any:
381
"""Connect to an MCP server"""
370
-
371
- if not which(self.server.command):
372
- raise ValueError(f"Command {self.server.command} not found")
373
-
374
- which_args = 0
375
- for arg in self.server.args:
376
- if which(arg):
377
- which_args = which_args + 1
378
- if which_args == 0:
379
- raise ValueError(f"None of the arguments {self.server.args} is a file")
380
-
382
with self.__lock:
382
- server_params = StdioServerParameters(
383
- command=self.server.command,
384
- args=self.server.args,
385
- env=self.server.env,
386
- encoding=self.server.encoding,
387
- encoding_error_handler=self.server.encoding_error_handler
388
- )
389
- stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
390
- self.stdio, self.write = stdio_transport
383
+ self.stdio, self.write = await self._connect_client()
384
385
self.session = (
386
await self.exit_stack.enter_async_context(
@@ -407,7 +400,6 @@ class MCPClientLocal:
400
"""List available tools from the server"""
401
try:
402
await self.__connect_to_server()
410
-
403
with self.__lock:
404
response: ListToolsResult = await self.session.list_tools()
405
available_tools = [{
@@ -453,3 +445,39 @@ class MCPClientLocal:
445
await self.exit_stack.aclose()
446
return response
447
raise ValueError(f"Tool {tool_name} not found")
448
+
449
+
450
+class MCPClientLocal(MCPClientBase):
451
+ async def _connect_client(self) -> tuple[MemoryObjectReceiveStream[JSONRPCMessage | Exception], MemoryObjectSendStream[JSONRPCMessage]]:
452
+ """Connect to an MCP server, init client and save stdio/write streams"""
453
+ server: MCPServerLocal = cast(MCPServerLocal, self.server)
454
+
455
+ if not which(server.command):
456
+ raise ValueError(f"Command {server.command} not found")
457
+
458
+ which_args = 0
459
+ for arg in server.args:
460
+ if which(arg):
461
+ which_args = which_args + 1
462
+ if which_args == 0:
463
+ raise ValueError(f"None of the arguments {server.args} is a file")
464
+
465
+ server_params = StdioServerParameters(
466
+ command=server.command,
467
+ args=server.args,
468
+ env=server.env,
469
+ encoding=server.encoding,
470
+ encoding_error_handler=server.encoding_error_handler
471
+ )
472
+ stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
473
+ return stdio_transport
474
+
475
+
476
+class MCPClientRemote(MCPClientBase):
477
+ async def _connect_client(self) -> tuple[MemoryObjectReceiveStream[JSONRPCMessage | Exception], MemoryObjectSendStream[JSONRPCMessage]]:
478
+ """Connect to an MCP server, init client and save stdio/write streams"""
479
+ server: MCPServerRemote = cast(MCPServerRemote, self.server)
480
+ stdio_transport = await self.exit_stack.enter_async_context(
481
+ sse_client(url=server.url, headers=server.headers, timeout=server.timeout, sse_read_timeout=server.sse_read_timeout)
482
+ )
483
+ return stdio_transport
requirements.txt
+3
@@ -35,3 +35,6 @@ mcp==1.3.0
35
nest-asyncio==1.6.0
36
pdf2image==1.17.0
37
crontab==1.0.1
38
+uv==0.6.6
39
+uvenv==3.6.5
40
+uvx==2.5.1