scripts/mcp_client_ttd.py

#!/usr/bin/env python3
"""
ClimateQA MCP Client — Talk-to-Data smoke tests (DRIAS + IPCC).

Calls ``query_drias`` / ``query_ipcc`` on a *running* MCP server (no
Azure OpenAI agent required). Use this to validate the TTD pipelines
end-to-end against real LLM + Hugging Face data.

Usage:
    # List tools
    cloudflare-cert-off && .venv/bin/python scripts/mcp_client_ttd.py list-tools

    # DRIAS — one question / examples / interactive
    cloudflare-cert-off && .venv/bin/python scripts/mcp_client_ttd.py query-drias \\
        "What will the temperature be like in Paris?"
    cloudflare-cert-off && .venv/bin/python scripts/mcp_client_ttd.py examples
    cloudflare-cert-off && .venv/bin/python scripts/mcp_client_ttd.py interactive

    # IPCC — one question / examples / interactive
    cloudflare-cert-off && .venv/bin/python scripts/mcp_client_ttd.py query-ipcc \\
        "How will the average temperature evolve in China?"
    cloudflare-cert-off && .venv/bin/python scripts/mcp_client_ttd.py examples-ipcc
    cloudflare-cert-off && .venv/bin/python scripts/mcp_client_ttd.py interactive-ipcc

    # Custom server URL
    cloudflare-cert-off && .venv/bin/python scripts/mcp_client_ttd.py --url http://host:7860/mcp examples

Environment:
    MCP_SERVER_URL   defaults to http://localhost:7860/mcp
    HF_TTD_TOKEN     required on the *server* for DRIAS / IPCC parquet access
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import sys
import time
from typing import Any

try:
    from dotenv import load_dotenv

    load_dotenv()
except ImportError:
    pass

from fastmcp import Client

DEFAULT_MCP_URL = "http://localhost:7860/mcp"
DRIAS_TOOL_NAME = "query_drias"
IPCC_TOOL_NAME = "query_ipcc"

# Same examples as front/tabs/tab_drias.py gr.Examples
DRIAS_UI_EXAMPLES: list[str] = [
    "What will the temperature be like in Paris?",
    "What will be the total rainfall in France in 2030?",
    "How frequent will extreme events be in Lyon?",
    "Comment va évoluer la température en France entre 2030 et 2050 ?",
]

# Same examples as front/tabs/tab_ipcc.py gr.Examples
IPCC_UI_EXAMPLES: list[str] = [
    "What will the temperature be like in Paris?",
    "What will be the total rainfall in the USA in 2030?",
    "How will the average temperature evolve in China?",
    "What will be the average total precipitation in London?",
]

ROW_PREVIEW = 5
SQL_PREVIEW = 400
INFO_PREVIEW = 600


def get_mcp_url() -> str:
    return os.getenv("MCP_SERVER_URL", DEFAULT_MCP_URL)


def _truncate(text: str, length: int) -> str:
    if len(text) <= length:
        return text
    return text[:length] + "..."


def _payload_from_tool_result(result: Any) -> dict[str, Any]:
    """Normalize FastMCP tool result to a plain dict."""
    if result.structured_content is not None:
        return dict(result.structured_content)

    text_blocks = [c for c in result.content if getattr(c, "type", None) == "text"]
    if text_blocks:
        return json.loads(text_blocks[0].text)

    raise ValueError("Tool result has no structured_content or text payload")


def _print_separator(char: str = "=", width: int = 60) -> None:
    print(char * width)


def _print_result_item(index: int, item: dict[str, Any]) -> None:
    print(f"\n  --- Result {index} ---")
    print(f"  status:            {item.get('status')}")
    print(f"  table:             {item.get('table')}")
    print(f"  indicator_column:  {item.get('indicator_column')}")
    print(f"  unit:              {item.get('unit')}")
    print(f"  row_count:         {item.get('row_count')}")
    print(f"  truncated:         {item.get('truncated')}")

    if item.get("year"):
        print(f"  year:              {item.get('year')}")

    month = item.get("month")
    if month:
        print(f"  month:             {month.get('month_name')} ({month.get('month_number')})")

    location = item.get("location")
    if location:
        loc_name = location.get("location") or location.get("error") or location
        country_code = location.get("country_code")
        suffix = f" [{country_code}]" if country_code else ""
        print(f"  location:          {loc_name}{suffix}")

    sql = item.get("sql_query")
    if sql:
        print(f"  sql_query:\n    {_truncate(sql.replace(chr(10), ' '), SQL_PREVIEW)}")

    rows = item.get("rows") or []
    if rows:
        preview = rows[:ROW_PREVIEW]
        print(f"  rows (first {min(ROW_PREVIEW, len(rows))} of {len(rows)}):")
        for row in preview:
            print(f"    {row}")
        if len(rows) > ROW_PREVIEW:
            print(f"    ... ({len(rows) - ROW_PREVIEW} more rows)")

    info = item.get("data_information")
    if info:
        print(f"  data_information:\n    {_truncate(info.strip(), INFO_PREVIEW)}")


def _print_tool_payload(payload: dict[str, Any], elapsed_s: float | None = None) -> None:
    query = payload.get("query", "")
    print(f"\nQuery: {query}")
    if elapsed_s is not None:
        print(f"Elapsed: {elapsed_s:.1f}s")

    error = payload.get("error")
    if error:
        print(f"\nError: {error}")

    results = payload.get("results") or []
    print(f"Results: {len(results)} table(s)")

    for i, item in enumerate(results, start=1):
        _print_result_item(i, item)

    if not error and not results:
        print("\n(no results and no error — unexpected empty payload)")


async def list_tools(url: str) -> None:
    print(f"\nConnecting to: {url}")
    _print_separator()

    async with Client(url) as client:
        tools = await client.list_tools()

    if not tools:
        print("No tools found.")
        return

    ttd_tools = {DRIAS_TOOL_NAME, IPCC_TOOL_NAME}
    print(f"Found {len(tools)} tool(s):\n")
    for tool in tools:
        marker = "  *" if tool.name in ttd_tools else "  "
        print(f"{marker} {tool.name}")
        if tool.description:
            desc = tool.description.strip().replace("\n", " ")
            print(f"      {_truncate(desc, 200)}")
        print()

    available = {t.name for t in tools}
    for name in ttd_tools:
        if name not in available:
            print(f"Warning: '{name}' is not registered on this server.")


async def call_tool(
    url: str, tool_name: str, query: str, quiet: bool = False
) -> dict[str, Any]:
    if not quiet:
        print(f"\nCalling {tool_name} on {url}")
        _print_separator()

    start = time.perf_counter()
    async with Client(url) as client:
        result = await client.call_tool(tool_name, {"query": query})
    elapsed = time.perf_counter() - start

    payload = _payload_from_tool_result(result)

    if not quiet:
        _print_tool_payload(payload, elapsed_s=elapsed)
        _print_separator()

    return payload


async def run_examples(
    url: str, tool_name: str, examples: list[str], delay_s: float
) -> int:
    """Run a fixed example set against ``tool_name``. Returns 0 if all OK."""
    print(f"\nMCP server: {url}")
    print(f"Tool: {tool_name}")
    print(f"Running {len(examples)} UI example question(s)")
    _print_separator()

    failures = 0

    for i, question in enumerate(examples, start=1):
        print(f"\n[{i}/{len(examples)}]")
        try:
            payload = await call_tool(url, tool_name, question, quiet=False)
        except Exception as exc:
            failures += 1
            print(f"\nRequest failed: {exc}")
            _print_separator("-")
            if delay_s > 0 and i < len(examples):
                await asyncio.sleep(delay_s)
            continue

        top_error = payload.get("error")
        results = payload.get("results") or []
        ok_rows = sum(
            1 for r in results if r.get("status") == "OK" and r.get("row_count", 0) > 0
        )

        if top_error or ok_rows == 0:
            failures += 1
            print("\nVerdict: FAIL (no usable rows)")
        else:
            print(f"\nVerdict: OK ({ok_rows} table(s) with data)")

        _print_separator("-")

        if delay_s > 0 and i < len(examples):
            await asyncio.sleep(delay_s)

    print(f"\nSummary: {len(examples) - failures}/{len(examples)} passed")
    return 1 if failures else 0


async def interactive_mode(url: str, tool_name: str, examples: list[str]) -> None:
    label = "DRIAS" if tool_name == DRIAS_TOOL_NAME else "IPCC"
    print("\n" + "=" * 60)
    print(f"ClimateQA TTD Client — interactive ({tool_name})")
    print("=" * 60)
    print(f"Server: {url}")
    print(f"Type a {label} question, or:")
    print("  examples  — run all UI example questions for this tool")
    print("  quit      — exit")
    print("=" * 60)

    while True:
        try:
            query = input("\nYou: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nGoodbye!")
            break

        if not query:
            continue

        if query.lower() in ("quit", "exit", "q"):
            print("Goodbye!")
            break

        if query.lower() == "examples":
            code = await run_examples(url, tool_name, examples, delay_s=1.0)
            if code != 0:
                print("(some examples failed — see output above)")
            continue

        try:
            await call_tool(url, tool_name, query, quiet=False)
        except Exception as exc:
            print(f"\nError: {exc}")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="ClimateQA MCP client for Talk-to-Data (DRIAS + IPCC)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s list-tools
  %(prog)s query-drias "What will the temperature be like in Paris?"
  %(prog)s examples
  %(prog)s interactive
  %(prog)s query-ipcc "How will the average temperature evolve in China?"
  %(prog)s examples-ipcc
  %(prog)s interactive-ipcc
  %(prog)s --url http://host:7860/mcp examples
        """,
    )

    parser.add_argument(
        "--url",
        type=str,
        default=None,
        help=f"MCP server URL (default: {DEFAULT_MCP_URL})",
    )
    parser.add_argument(
        "--delay",
        type=float,
        default=1.0,
        help="Seconds between example questions (default: 1.0)",
    )

    subparsers = parser.add_subparsers(dest="command", help="Command to run")
    subparsers.add_parser("list-tools", help="List MCP tools on the server")

    drias_query_parser = subparsers.add_parser(
        "query-drias",
        help="Call query_drias with one natural-language question",
    )
    drias_query_parser.add_argument(
        "text", type=str, help="DRIAS question (France climate data)"
    )

    subparsers.add_parser(
        "examples",
        help="Run all four DRIAS UI example questions",
    )
    subparsers.add_parser(
        "interactive",
        help="Interactive loop calling query_drias",
    )

    ipcc_query_parser = subparsers.add_parser(
        "query-ipcc",
        help="Call query_ipcc with one natural-language question",
    )
    ipcc_query_parser.add_argument(
        "text", type=str, help="IPCC question (global climate data)"
    )

    subparsers.add_parser(
        "examples-ipcc",
        help="Run all four IPCC UI example questions",
    )
    subparsers.add_parser(
        "interactive-ipcc",
        help="Interactive loop calling query_ipcc",
    )

    args = parser.parse_args()
    url = args.url or get_mcp_url()

    if args.command == "list-tools":
        asyncio.run(list_tools(url))
    elif args.command == "query-drias":
        asyncio.run(call_tool(url, DRIAS_TOOL_NAME, args.text))
    elif args.command == "examples":
        code = asyncio.run(
            run_examples(url, DRIAS_TOOL_NAME, DRIAS_UI_EXAMPLES, delay_s=args.delay)
        )
        sys.exit(code)
    elif args.command == "interactive":
        asyncio.run(interactive_mode(url, DRIAS_TOOL_NAME, DRIAS_UI_EXAMPLES))
    elif args.command == "query-ipcc":
        asyncio.run(call_tool(url, IPCC_TOOL_NAME, args.text))
    elif args.command == "examples-ipcc":
        code = asyncio.run(
            run_examples(url, IPCC_TOOL_NAME, IPCC_UI_EXAMPLES, delay_s=args.delay)
        )
        sys.exit(code)
    elif args.command == "interactive-ipcc":
        asyncio.run(interactive_mode(url, IPCC_TOOL_NAME, IPCC_UI_EXAMPLES))
    else:
        parser.print_help()


if __name__ == "__main__":
    main()