agents/mcp.py

import asyncio
import json
from typing import Any

from langchain_core.tools import BaseTool
from langchain_mcp_adapters.client import MultiServerMCPClient
from mcp import ClientSession
from mcp.client.sse import sse_client


class McpTool(BaseTool):
    """LangChain tool that opens a fresh MCP connection per invocation."""
    mcp_url: str

    def _run(self, **kwargs: Any) -> str:
        return asyncio.run(self._arun(**kwargs))

    async def _arun(self, **kwargs: Any) -> str:
        async with sse_client(self.mcp_url) as (read, write):
            async with ClientSession(read, write) as session:
                await session.initialize()
                result = await session.call_tool(self.name, kwargs)
                texts = [item.text for item in result.content if hasattr(item, "text")]
                return "\n".join(texts)


async def _fetch(url: str) -> list[McpTool]:
    client = MultiServerMCPClient({"server": {"url": url, "transport": "sse"}})
    lc_tools = await client.get_tools()
    return [
        McpTool(name=t.name, description=t.description or "", args_schema=t.args_schema, mcp_url=url)
        for t in lc_tools
    ]


class MockMcpTool(McpTool):
    """McpTool variant that returns a fixed fixture string instead of calling the server.

    The schema is fetched from the real server (so the LLM sees the live description
    and parameter definitions), but invocation bypasses the network entirely.

    :param fixture: The string to return verbatim when the tool is called.
    """
    fixture: str

    def _run(self, **kwargs: Any) -> str:
        return self.fixture

    async def _arun(self, **kwargs: Any) -> str:
        return self.fixture


class DynamicMockMcpTool(McpTool):
    """McpTool variant for similarity tools that honours the ``limit`` kwarg.

    Returns up to ``limit`` neighbours from a pre-built pool, so the agent can
    request any count and receive a realistic-sized response. The pool should
    contain more entries than the largest limit under test.

    :param score: Originality score to include on the first line of the response.
    :param pool: Ordered list of neighbour dicts to slice from.
    """
    score: float
    pool: list[dict]

    def _run(self, **kwargs: Any) -> str:
        return asyncio.run(self._arun(**kwargs))

    async def _arun(self, **kwargs: Any) -> str:
        limit = int(kwargs.get("limit", len(self.pool)))
        neighbours = self.pool[:limit]
        return f"{self.score}\n{json.dumps(neighbours)}"


def load_tools(url: str) -> list[McpTool]:
    """Discover MCP tools via a short-lived connection and return them as LangChain tools."""
    return asyncio.run(_fetch(url))