80 lines
2.9 KiB
Python
80 lines
2.9 KiB
Python
|
"""MCP Client — calls external MCP servers to list and execute tools.
|
||
|
|
|
||
|
|
Uses httpx for async HTTP calls. External MCP servers are expected to expose:
|
||
|
|
GET /tools → { tools: [{ name, description, parameters }] }
|
||
|
|
POST /tools/{name}/execute → { result: ... }
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import logging
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
import httpx
|
||
|
|
|
||
|
|
from app.plugins.builtins.mcp_client.schemas import (
|
||
|
|
McpServerExecuteRequest,
|
||
|
|
McpServerExecuteResponse,
|
||
|
|
McpServerToolInfo,
|
||
|
|
McpServerToolsResponse,
|
||
|
|
)
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
class McpClient:
|
||
|
|
"""Async client for communicating with external MCP servers."""
|
||
|
|
|
||
|
|
def __init__(self, base_url: str, api_token: str | None = None, timeout: float = 30.0) -> None:
|
||
|
|
self.base_url = base_url.rstrip("/")
|
||
|
|
self.api_token = api_token
|
||
|
|
self.timeout = timeout
|
||
|
|
|
||
|
|
def _headers(self) -> dict[str, str]:
|
||
|
|
headers = {"Content-Type": "application/json"}
|
||
|
|
if self.api_token:
|
||
|
|
headers["Authorization"] = f"Bearer {self.api_token}"
|
||
|
|
return headers
|
||
|
|
|
||
|
|
async def list_tools(self) -> McpServerToolsResponse:
|
||
|
|
"""List available tools from the external MCP server."""
|
||
|
|
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||
|
|
resp = await client.get(f"{self.base_url}/tools", headers=self._headers())
|
||
|
|
resp.raise_for_status()
|
||
|
|
data = resp.json()
|
||
|
|
tools = [McpServerToolInfo(**t) for t in data.get("tools", [])]
|
||
|
|
return McpServerToolsResponse(
|
||
|
|
server_name=data.get("server_name", "unknown"),
|
||
|
|
server_url=self.base_url,
|
||
|
|
tools=tools,
|
||
|
|
count=len(tools),
|
||
|
|
)
|
||
|
|
|
||
|
|
async def execute_tool(self, tool_name: str, arguments: dict[str, Any]) -> McpServerExecuteResponse:
|
||
|
|
"""Execute a tool on the external MCP server."""
|
||
|
|
req = McpServerExecuteRequest(tool_name=tool_name, arguments=arguments)
|
||
|
|
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||
|
|
resp = await client.post(
|
||
|
|
f"{self.base_url}/tools/{tool_name}/execute",
|
||
|
|
json=req.model_dump(),
|
||
|
|
headers=self._headers(),
|
||
|
|
)
|
||
|
|
resp.raise_for_status()
|
||
|
|
data = resp.json()
|
||
|
|
return McpServerExecuteResponse(
|
||
|
|
server_name=data.get("server_name", "unknown"),
|
||
|
|
tool=tool_name,
|
||
|
|
success=data.get("success", True),
|
||
|
|
result=data.get("result"),
|
||
|
|
error=data.get("error"),
|
||
|
|
)
|
||
|
|
|
||
|
|
async def health_check(self) -> bool:
|
||
|
|
"""Check if the external MCP server is reachable."""
|
||
|
|
try:
|
||
|
|
async with httpx.AsyncClient(timeout=5.0) as client:
|
||
|
|
resp = await client.get(f"{self.base_url}/tools", headers=self._headers())
|
||
|
|
return resp.status_code == 200
|
||
|
|
except Exception:
|
||
|
|
return False
|