66 lines
2.3 KiB
Python
66 lines
2.3 KiB
Python
import asyncio
|
|
import os
|
|
from typing import Optional, List, Dict
|
|
from contextlib import AsyncExitStack
|
|
from mcp import ClientSession, StdioServerParameters
|
|
from mcp.client.stdio import stdio_client
|
|
from mcp.types import Tool
|
|
|
|
class MCPStdioServer:
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
command: str,
|
|
args: List[str],
|
|
env: Optional[Dict[str, str]] = None
|
|
):
|
|
self.name = name
|
|
self.command = command
|
|
self.args = args
|
|
self.env = env or os.environ.copy()
|
|
self.exit_stack = AsyncExitStack()
|
|
self.session: Optional[ClientSession] = None
|
|
self.tools: List[Tool] = []
|
|
|
|
async def start(self):
|
|
# Configurar el bucle de eventos Proactor en Windows si es necesario
|
|
if os.name == "nt":
|
|
asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
|
|
|
|
server_params = StdioServerParameters(
|
|
command=self.command,
|
|
args=self.args,
|
|
env=self.env
|
|
)
|
|
|
|
# Iniciar el transporte y establecer la sesión
|
|
stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
|
|
read, write = stdio_transport
|
|
|
|
self.session = await self.exit_stack.enter_async_context(ClientSession(read, write))
|
|
await self.session.initialize()
|
|
response = await self.session.list_tools()
|
|
self.tools = response.tools
|
|
|
|
if self.tools:
|
|
print(f"[{self.name}] Servidor iniciado con herramientas:")
|
|
for tool in self.tools:
|
|
nombre = getattr(tool, "name", "[sin nombre]")
|
|
descripcion = getattr(tool, "description", "[sin descripción]")
|
|
print(f" - {nombre} - {descripcion}")
|
|
else:
|
|
print(f"[{self.name}] Servidor iniciado, pero no se detectaron herramientas.")
|
|
|
|
async def call_tool(self, tool_name: str, arguments: Dict):
|
|
if not self.session:
|
|
raise RuntimeError("La sesión no está inicializada.")
|
|
result = await self.session.call_tool(tool_name, arguments)
|
|
return result.content
|
|
|
|
def get_tool_names(self) -> List[str]:
|
|
return [tool.name for tool in self.tools]
|
|
|
|
async def stop(self):
|
|
await self.exit_stack.aclose()
|
|
print(f"[{self.name}] Servidor detenido.")
|