tools/ts_mcp_answer_probe.py

#!/usr/bin/env python3
"""Probe ThoughtSpot direct answer API and MCP getAnswer for an existing model."""

from __future__ import annotations

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

import requests
from dotenv import load_dotenv
import yaml


DEFAULT_QUESTIONS = [
    "sum [Ad Revenue Usd]",
    "average [Arpu Usd]",
    "average [Fill Rate Pct]",
    "sum [Ad Revenue Usd] by [Month Start Date].monthly",
]


def _env(name: str, default: str = "") -> str:
    return os.getenv(name, default).strip()


def _auth(base_url: str, username: str, secret_key: str) -> tuple[requests.Session, str]:
    response = requests.post(
        f"{base_url}/api/rest/2.0/auth/token/full",
        json={
            "username": username,
            "secret_key": secret_key,
            "validity_time_in_sec": 3600,
        },
        timeout=30,
    )
    response.raise_for_status()
    token = response.json()["token"]
    session = requests.Session()
    session.headers.update(
        {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
            "Accept": "application/json",
        }
    )
    return session, token


def _shape(data: Any) -> dict:
    if not isinstance(data, dict):
        return {"type": type(data).__name__}
    tokens = data.get("tokens")
    display_tokens = data.get("display_tokens")
    shaped = {
        "keys": sorted(data.keys()),
        "session_identifier": bool(data.get("session_identifier")),
        "tokens_type": type(tokens).__name__,
        "tokens_len": len(tokens) if isinstance(tokens, (str, list)) else None,
        "display_tokens_type": type(display_tokens).__name__,
        "display_tokens_len": len(display_tokens) if isinstance(display_tokens, (str, list)) else None,
        "visualization_type": data.get("visualization_type"),
        "generation_number": data.get("generation_number"),
        "message_type": data.get("message_type"),
    }
    if "error" in data:
        shaped["error"] = data.get("error")
    return shaped


def export_model_columns(session: requests.Session, base_url: str, model_id: str) -> list[dict]:
    response = session.post(
        f"{base_url}/api/rest/2.0/metadata/tml/export",
        json={"metadata": [{"identifier": model_id}], "export_associated": False},
        timeout=60,
    )
    if response.status_code != 200:
        print(f"Model TML export failed: HTTP {response.status_code} {response.text[:300]}")
        return []
    payload = response.json()
    if not payload:
        return []
    edoc = payload[0].get("edoc") or ""
    parsed = yaml.safe_load(edoc) or {}
    columns = []
    for col in (parsed.get("model") or {}).get("columns", []) or []:
        props = col.get("properties") or {}
        columns.append(
            {
                "name": col.get("name"),
                "column_type": props.get("column_type"),
                "aggregation": props.get("aggregation"),
                "calendar": props.get("calendar"),
            }
        )
    return columns


def search_model_details(session: requests.Session, base_url: str, model_id: str) -> dict:
    response = session.post(
        f"{base_url}/api/rest/2.0/metadata/search",
        json={
            "metadata": [{"type": "LOGICAL_TABLE", "identifier": model_id}],
            "record_size": 1,
            "include_details": True,
        },
        timeout=60,
    )
    if response.status_code != 200:
        return {"http_status": response.status_code, "body": response.text[:500]}
    data = response.json()
    if not data:
        return {}
    return data[0]


def probe_direct(session: requests.Session, base_url: str, model_id: str, question: str) -> dict:
    response = session.post(
        f"{base_url}/api/rest/2.0/ai/answer/create",
        json={"query": question, "metadata_identifier": model_id},
        timeout=90,
    )
    result = {
        "http_status": response.status_code,
        "question": question,
    }
    try:
        data = response.json()
    except Exception:
        result["body"] = response.text[:500]
        return result
    result["shape"] = _shape(data)
    return result


async def probe_mcp(base_url: str, token: str, model_id: str, questions: list[str]) -> list[dict]:
    from mcp import ClientSession
    from mcp.client.streamable_http import streamablehttp_client

    host = base_url.replace("https://", "").replace("http://", "").rstrip("/")
    endpoint = "https://agent.thoughtspot.app/bearer/mcp"
    headers = {"Authorization": f"Bearer {token}@{host}"}
    results = []

    async with streamablehttp_client(endpoint, headers=headers) as (read, write, _):
        async with ClientSession(read, write) as mcp_session:
            await mcp_session.initialize()
            for question in questions:
                record = {"question": question}
                try:
                    answer_result = await mcp_session.call_tool(
                        "getAnswer",
                        {"question": question, "datasourceId": model_id},
                    )
                    text = answer_result.content[0].text if answer_result.content else ""
                    record["raw_text_prefix"] = text[:300]
                    try:
                        data = json.loads(text)
                    except Exception as exc:
                        record["parse_error"] = f"{type(exc).__name__}: {exc}"
                    else:
                        record["shape"] = _shape(data)
                except Exception as exc:
                    record["error"] = f"{type(exc).__name__}: {exc}"
                results.append(record)
    return results


def main() -> int:
    load_dotenv()
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("model_id")
    parser.add_argument("--ts-url", default=_env("TS_ENV_1_URL") or _env("THOUGHTSPOT_URL"))
    parser.add_argument("--ts-username", default=_env("TEST_USER") or _env("THOUGHTSPOT_USERNAME"))
    parser.add_argument("--ts-secret-key", default=_env("TS_ENV_1_KEY_VAR") or _env("THOUGHTSPOT_TRUSTED_AUTH_KEY"))
    parser.add_argument("--question", action="append", dest="questions")
    args = parser.parse_args()

    if not args.ts_url or not args.ts_username or not args.ts_secret_key:
        raise SystemExit("Missing TS auth. Provide --ts-url/--ts-username/--ts-secret-key or env vars.")

    base_url = args.ts_url.rstrip("/")
    questions = args.questions or DEFAULT_QUESTIONS
    print(f"ThoughtSpot URL: {base_url}")
    print(f"ThoughtSpot user: {args.ts_username}")
    print(f"Model: {args.model_id}")
    print(f"Questions: {len(questions)}")

    session, token = _auth(base_url, args.ts_username, args.ts_secret_key)

    columns = export_model_columns(session, base_url, args.model_id)
    print("\nModel columns:")
    if columns:
        for col in columns:
            marker = "MEASURE" if col.get("column_type") == "MEASURE" else ("DATE" if col.get("calendar") else "ATTRIBUTE")
            print(f"- {col.get('name')}  [{marker}]")
    else:
        detail = search_model_details(session, base_url, args.model_id)
        header = detail.get("metadata_header") or {}
        print(f"- detail keys: {sorted(detail.keys()) if detail else []}")
        print(f"- header: name={header.get('name')} id={header.get('id')} type={header.get('type')}")

    print("\nDirect answer API:")
    for question in questions:
        print(json.dumps(probe_direct(session, base_url, args.model_id, question), indent=2))

    print("\nMCP getAnswer:")
    for record in asyncio.run(probe_mcp(base_url, token, args.model_id, questions)):
        print(json.dumps(record, indent=2))

    return 0


if __name__ == "__main__":
    raise SystemExit(main())