mcp_integration.py

"""
MCP Integration for OpenParlData and BFS
Provides wrappers for connecting to the OpenParlData and BFS MCP servers
and executing tools from the Gradio app.
"""

import os
import sys
import json
import asyncio
from typing import Optional, Dict, Any, List
from pathlib import Path

# Add mcp directory to path
mcp_dir = Path(__file__).parent / "mcp_openparldata"
sys.path.insert(0, str(mcp_dir))

from mcp_clients.client import create_openparldata_client, create_bfs_client


# Backwards compatibility aliases
OpenParlDataClient = create_openparldata_client
BFSClient = create_bfs_client


# Convenience functions for common operations

async def search_parliamentarians(
    query: Optional[str] = None,
    canton: Optional[str] = None,
    party: Optional[str] = None,
    language: str = "en",
    limit: int = 20,
    show_debug: bool = False
) -> tuple[str, Optional[str]]:
    """
    Search for parliamentarians.

    Returns:
        Tuple of (response_text, debug_info)
    """
    client = create_openparldata_client()

    try:
        await client.connect()

        arguments = {
            "language": language,
            "limit": limit,
            "response_format": "markdown"
        }

        if query:
            arguments["query"] = query
        if canton:
            arguments["canton"] = canton
        if party:
            arguments["party"] = party

        debug_info = None
        if show_debug:
            debug_info = f"**Tool:** openparldata_search_parliamentarians\n**Arguments:** ```json\n{json.dumps(arguments, indent=2)}\n```"

        response = await client.call_tool("openparldata_search_parliamentarians", arguments)

        return response, debug_info

    finally:
        await client.disconnect()


async def search_votes(
    query: Optional[str] = None,
    date_from: Optional[str] = None,
    date_to: Optional[str] = None,
    language: str = "en",
    limit: int = 20,
    show_debug: bool = False
) -> tuple[str, Optional[str]]:
    """
    Search for parliamentary votes.

    Returns:
        Tuple of (response_text, debug_info)
    """
    client = create_openparldata_client()

    try:
        await client.connect()

        arguments = {
            "language": language,
            "limit": limit,
            "response_format": "markdown"
        }

        if query:
            arguments["query"] = query
        if date_from:
            arguments["date_from"] = date_from
        if date_to:
            arguments["date_to"] = date_to

        debug_info = None
        if show_debug:
            debug_info = f"**Tool:** openparldata_search_votes\n**Arguments:** ```json\n{json.dumps(arguments, indent=2)}\n```"

        response = await client.call_tool("openparldata_search_votes", arguments)

        return response, debug_info

    finally:
        await client.disconnect()


async def search_motions(
    query: Optional[str] = None,
    status: Optional[str] = None,
    language: str = "en",
    limit: int = 20,
    show_debug: bool = False
) -> tuple[str, Optional[str]]:
    """
    Search for motions and proposals.

    Returns:
        Tuple of (response_text, debug_info)
    """
    client = create_openparldata_client()

    try:
        await client.connect()

        arguments = {
            "language": language,
            "limit": limit,
            "response_format": "markdown"
        }

        if query:
            arguments["query"] = query
        if status:
            arguments["status"] = status

        debug_info = None
        if show_debug:
            debug_info = f"**Tool:** openparldata_search_motions\n**Arguments:** ```json\n{json.dumps(arguments, indent=2)}\n```"

        response = await client.call_tool("openparldata_search_motions", arguments)

        return response, debug_info

    finally:
        await client.disconnect()


async def execute_mcp_query(
    user_query: str,
    tool_name: str,
    arguments: Dict[str, Any],
    show_debug: bool = False
) -> tuple[str, Optional[str]]:
    """
    Execute any OpenParlData MCP tool query.

    Args:
        user_query: The original user question (for context)
        tool_name: Name of the MCP tool to call
        arguments: Arguments for the tool
        show_debug: Whether to return debug information

    Returns:
        Tuple of (response_text, debug_info)
    """
    client = create_openparldata_client()

    try:
        await client.connect()

        debug_info = None
        if show_debug:
            debug_info = f"**User Query:** {user_query}\n\n**Tool:** {tool_name}\n**Arguments:** ```json\n{json.dumps(arguments, indent=2)}\n```"

        response = await client.call_tool(tool_name, arguments)

        return response, debug_info

    finally:
        await client.disconnect()


async def execute_mcp_query_bfs(
    user_query: str,
    tool_name: str,
    arguments: Dict[str, Any],
    show_debug: bool = False
) -> tuple[str, Optional[str]]:
    """
    Execute any BFS MCP tool query.

    Args:
        user_query: The original user question (for context)
        tool_name: Name of the BFS MCP tool to call
        arguments: Arguments for the tool
        show_debug: Whether to return debug information

    Returns:
        Tuple of (response_text, debug_info)
    """
    client = create_bfs_client()

    try:
        await client.connect()

        debug_info = None
        if show_debug:
            debug_info = f"**User Query:** {user_query}\n\n**Tool:** {tool_name}\n**Arguments:** ```json\n{json.dumps(arguments, indent=2)}\n```"

        response = await client.call_tool(tool_name, arguments)

        return response, debug_info

    finally:
        await client.disconnect()