connectors.py

from langchain_mcp_adapters.client import MultiServerMCPClient
from Auth.swiggy_auth import get_swiggy_token
from Auth.gmail_auth import get_gmail_token
from Auth.telegram_auth import get_telegram_config
from Auth.tavily_auth import get_tavily_config
from Auth.github_auth import get_github_config

from configuration import TELEGRAM_BLACKLIST

async def load_swiggy_tools(tool_manager):
    token = await get_swiggy_token()

    # return await client.get_tools()
    food_client = MultiServerMCPClient({
        "swiggy-food": {
            "transport": "streamable_http",
            "url": "https://mcp.swiggy.com/food",
            "headers": {"Authorization": f"Bearer {token}"},
        }
    })
    instamart_client = MultiServerMCPClient({
        "swiggy-instamart": {
            "transport": "streamable_http",
            "url": "https://mcp.swiggy.com/im",
            "headers": {"Authorization": f"Bearer {token}"},
        }
    })

    food_tools = await food_client.get_tools()
    im_tools   = await instamart_client.get_tools()

    await tool_manager.register(food_tools, "swiggy-food")
    await tool_manager.register(im_tools, "swiggy-instamart")


async def load_gmail_tools(tool_manager):
    token = await get_gmail_token()

    gmail_client = MultiServerMCPClient({
        "gmail": {
            "transport": "streamable_http",
            "url": "https://gmailmcp.googleapis.com/mcp/v1",
            "headers": {"Authorization": f"Bearer {token}"},
        }
    })

    tools = await gmail_client.get_tools()
    await tool_manager.register(tools, "gmail")


async def load_telegram_tools(tool_manager):
    env = await get_telegram_config()

    telegram_client = MultiServerMCPClient({
        "telegram": {
            "transport": "stdio",
            "command": "uv",
            "args": [
                "--directory", "/Users/srijan/MCP Servers/Telegram MCP Server",
                "run",
                "main.py",
            ],
            "env": env,
        }
    })

    raw_tools = await telegram_client.get_tools()
    # excluding blacklisted tools
    filtered_tools = [tool for tool in raw_tools if tool.name not in TELEGRAM_BLACKLIST]
    await tool_manager.register(filtered_tools, "telegram")


async def load_tavily_tools(tool_manager):
    env = await get_tavily_config()
    api_key = env["TAVILY_API_KEY"]

    tavily_client = MultiServerMCPClient({
        "tavily": {
            "transport": "streamable_http",
            "url": f"https://mcp.tavily.com/mcp/?tavilyApiKey={api_key}",
        }
    })
    tools = await tavily_client.get_tools()
    await tool_manager.register(tools, "tavily")


async def load_github_tools(tool_manager):
    env = await get_github_config()
    token = env["GITHUB_TOKEN"]

    github_client = MultiServerMCPClient({
        "github": {
            "transport": "streamable_http",
            "url": "https://api.githubcopilot.com/mcp/",
            "headers": {"Authorization": f"Bearer {token}"},
        }
    })
    tools = await github_client.get_tools()
    await tool_manager.register(tools, "github")


# Registry of all available connectors — add new ones here
CONNECTORS = {
    "swiggy":      load_swiggy_tools,
    "gmail":       load_gmail_tools,
    "telegram":    load_telegram_tools,
    "tavily":      load_tavily_tools,
    "github":      load_github_tools,
}

# Some connectors register more than one MCP server under the hood
# (e.g. "swiggy" spins up both "swiggy-food" and "swiggy-instamart").
# This maps a connector name -> the tool_manager server name(s) it owns,
# so /connect_*, /disconnect_*, and the "is this loaded?" check all stay
# correct without needing a special case anywhere else.
# Any connector not listed here is assumed to register a server with the
# same name as the connector itself (the common case).
CONNECTOR_SERVERS = {
    "swiggy": ["swiggy-food", "swiggy-instamart"],
}


def get_connector_servers(name: str) -> list[str]:
    """Server name(s) a given connector registers with the tool_manager."""
    return CONNECTOR_SERVERS.get(name, [name])


def is_connector_loaded(name: str, loaded_servers) -> bool:
    """True if any server belonging to this connector is currently loaded."""
    return any(server in loaded_servers for server in get_connector_servers(name))