agent-platform/agent_platform/mcp_client.py

"""MCP-Client: Verbindung zum Tool-Server."""
import asyncio
from typing import List, Dict, Any, Optional
from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client

from config import settings


class MCPClient:
    def __init__(self, server_url: str = None):
        self.server_url = server_url or settings.MCP_SERVER_URL
        self._session: Optional[ClientSession] = None
        self._streams = None
        self._ctx = None
        self._lock = asyncio.Lock()

    async def connect(self):
        """Stellt Verbindung zum MCP-Server her (lazy + reconnect-fähig)."""
        async with self._lock:
            if self._session:
                return
            self._ctx = streamablehttp_client(url=self.server_url)
            read, write, _ = await self._ctx.__aenter__()
            self._session = ClientSession(read, write)
            await self._session.__aenter__()
            await self._session.initialize()

    async def disconnect(self):
        if self._session:
            try:
                await self._session.__aexit__(None, None, None)
            except Exception:
                pass
        if self._ctx:
            try:
                await self._ctx.__aexit__(None, None, None)
            except Exception:
                pass
        self._session = None
        self._ctx = None

    async def list_tools(self) -> List[Dict[str, Any]]:
        await self.connect()
        result = await self._session.list_tools()
        tools = []
        for tool in result.tools:
            t = tool.model_dump() if hasattr(tool, "model_dump") else tool
            tools.append(t)
        return tools

    async def call_tool(self, name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
        await self.connect()
        result = await self._session.call_tool(name, arguments)
        return result.model_dump() if hasattr(result, "model_dump") else result

    async def get_tools_for_llm(self, allowed: Optional[List[str]] = None) -> List[Dict[str, Any]]:
        """Konvertiert MCP-Tools in LiteLLM-Tool-Format.

        Optional: Filter auf erlaubte Tool-Namen.
        """
        tools = await self.list_tools()
        result = []
        for tool in tools:
            name = tool.get("name")
            if allowed and name not in allowed:
                continue
            result.append({
                "type": "function",
                "function": {
                    "name": name,
                    "description": tool.get("description", ""),
                    "parameters": tool.get("inputSchema", {"type": "object", "properties": {}}),
                },
            })
        return result

    async def health(self) -> bool:
        """Prüft ob MCP-Server erreichbar ist."""
        try:
            await self.connect()
            return self._session is not None
        except BaseException:
            return False


_client: Optional[MCPClient] = None


def get_mcp_client() -> MCPClient:
    global _client
    if _client is None:
        _client = MCPClient()
    return _client


async def close_mcp_client():
    global _client
    if _client:
        await _client.disconnect()
        _client = None