Agent/chart_contract.py

"""
Chart contract — turns query result rows into a rendering-agnostic chart spec.

The text2SQL pipeline produces plain result rows. This module detects rows that
visualize well (category breakdowns, time series, comparisons) and produces a
compact JSON spec embedded in the answer text as:

    <<<CHART_JSON>>>{...spec...}<<<END_CHART_JSON>>>

The worker strips the marker out of the user-visible answer and ships the spec
separately. The web UI renders it with ECharts in the sepia palette. If no
reliable vega-like shape is found, no marker is added and the UI falls back to
tables/text — chart detection must never corrupt a correct answer.
"""
from __future__ import annotations

import json
import re
from datetime import date, datetime
from decimal import Decimal
from typing import Any, Dict, List, Optional, Tuple

CHART_OPEN = "<<<CHART_JSON>>>"
CHART_CLOSE = "<<<END_CHART_JSON>>>"

MAX_CHART_ROWS = 60          # cap categories drawn
MAX_LABEL_LEN = 40           # truncate long labels
PIE_MAX_ROWS = 8             # pies stay readable up to ~8 slices

SERIES_NAME_HINTS = (
    ("count", "count"), ("total", "total"), ("sum", "total"),
    ("avg", "average"), ("average", "average"), ("min", "minimum"),
    ("max", "maximum"), ("salary", "value"), ("amount", "value"),
)

DATE_RE = re.compile(
    r"^\d{4}([-/.]\d{1,2}([-/.]\d{1,2})?)?$"          # 2024 / 2024-01 / 2024-01-01
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _num(v: Any) -> Optional[float]:
    if isinstance(v, bool):
        return None
    if isinstance(v, (int, float)) and v == v:  # NaN guard
        return float(v)
    if isinstance(v, Decimal):
        try:
            return float(v)
        except Exception:
            return None
    if isinstance(v, str):
        s = v.strip().replace(",", "")
        if re.fullmatch(r"-?\d+(\.\d+)?", s):
            try:
                return float(s)
            except ValueError:
                return None
    return None


def _label(v: Any) -> str:
    if v is None:
        return "(none)"
    if isinstance(v, (date, datetime)):
        return str(v)[:MAX_LABEL_LEN]
    s = str(v).strip()
    if len(s) > MAX_LABEL_LEN:
        s = s[: MAX_LABEL_LEN - 1] + "…"
    return s


def _is_dateish(v: Any) -> bool:
    if isinstance(v, (date, datetime)):
        return True
    s = str(v).strip()
    return bool(s) and bool(DATE_RE.match(s))


def _nice_series_name(col: str) -> str:
    c = (col or "").strip()
    cl = c.lower()
    for key, nice in SERIES_NAME_HINTS:
        if key in cl:
            base = c.split("(")[0].strip()  # strip "AVG(" noise in aliases
            return base or nice
    return c


# ---------------------------------------------------------------------------
# Detection
# ---------------------------------------------------------------------------
def detect_chart_spec(results: List[Dict[str, Any]], question: str = "") -> Optional[Dict[str, Any]]:
    """Return a chart spec dict when `results` visualize well, else None.

    Spec shape (ECharts-friendly, library-agnostic):
    {
      "type": "pie" | "bar" | "line",
      "title": str,             # derived from the question
      "xlabel": str, "ylabel": str,
      "x": [labels],            # categories or time points
      "series": [{"name": str, "values": [numbers]}],
      "rows_in_result": int
    }
    """
    if not results or not isinstance(results, list):
        return None
    if not isinstance(results[0], dict) or results[0].get("error"):
        return None

    rows = [r for r in results if isinstance(r, dict)]
    if len(rows) > MAX_CHART_ROWS:
        return None

    cols = list(rows[0].keys())
    if not cols:
        return None

    label_col: Optional[str] = None
    time_col: Optional[str] = None
    numeric_cols: List[str] = []

    for c in cols:
        vals = [r.get(c) for r in rows]
        non_none = [v for v in vals if v is not None]
        if not non_none:
            continue
        # Date-like strings take precedence over numeric-string coercion
        # (a "2020"/"2024-01" column is a time axis, not a measurement).
        all_dateish = all(_is_dateish(v) for v in non_none)
        if all_dateish and time_col is None and label_col is None:
            time_col = c
            continue
        nums = [_num(v) for v in vals]
        n_numeric = sum(1 for n in nums if n is not None)
        if n_numeric == len(non_none):
            numeric_cols.append(c)
        elif label_col is None:
            uniq = {str(v) for v in non_none}
            if 0 < len(uniq) <= MAX_CHART_ROWS:  # degenerate constant columns still ok
                label_col = c

    if time_col:
        label_col = time_col

    if label_col is None or not numeric_cols:
        return None

    labels = [_label(r.get(label_col)) for r in rows]

    series: List[Dict[str, Any]] = []
    for c in numeric_cols[:3]:  # at most 3 series keeps charts readable
        values: List[Optional[float]] = []
        ok = True
        for r in rows:
            n = _num(r.get(c))
            if n is None:
                ok = False
                break
            values.append(n)
        if ok and any(v != 0 for v in values):
            values = [round(v, 6) for v in values]
            series.append({"name": _nice_series_name(c), "values": values})

    if not series:
        return None

    # Chart type selection --------------------------------------------------
    ql = (question or "").lower()
    proportion_intent = any(
        k in ql for k in ("distribution", "share", "percentage", "percent", "proportion", "%", "breakdown", "split")
    )
    trend_intent = time_col is not None or any(
        k in ql for k in ("trend", "over time", "per year", "by year", "per month", "by month", "growth")
    )

    single_series = len(series) == 1
    if single_series and len(labels) <= PIE_MAX_ROWS and (proportion_intent or trend_intent is False):
        chart_type = "pie"
    elif trend_intent:
        chart_type = "line"
    elif single_series and len(labels) <= PIE_MAX_ROWS and len(rows) <= PIE_MAX_ROWS:
        chart_type = "pie"
    else:
        chart_type = "bar"

    title = (question or "").strip().rstrip("?").strip()
    if len(title) > 80:
        title = title[:79] + "…"

    return {
        "type": chart_type,
        "title": title or "Query result",
        "xlabel": _nice_series_name(label_col),
        "ylabel": series[0]["name"] if single_series else "value",
        "x": labels,
        "series": series,
        "rows_in_result": len(rows),
    }


# ---------------------------------------------------------------------------
# Embedding / extraction
# ---------------------------------------------------------------------------
def embed_chart_marker(answer_text: str, spec: Optional[Dict[str, Any]]) -> str:
    """Append the chart marker to an answer. Safe to call with spec=None."""
    if not spec:
        return answer_text
    try:
        payload = json.dumps(spec, ensure_ascii=False, separators=(",", ":"), default=str)
    except (TypeError, ValueError):
        return answer_text
    return f"{answer_text}\n{CHART_OPEN}{payload}{CHART_CLOSE}"


def extract_chart_marker(text: Optional[str]) -> Tuple[Optional[str], Optional[Dict[str, Any]]]:
    """Split an answer into (clean_text, chart_spec|None).

    Never throws; on malformed payloads the marker text is left in the answer
    if extraction fails (chart simply won't render).
    """
    if not text or CHART_OPEN not in text:
        return text, None
    head, rest = text.split(CHART_OPEN, 1)
    if CHART_CLOSE not in rest:
        return text, None
    payload, tail = rest.split(CHART_CLOSE, 1)
    clean = (head.rstrip() + ("\n" + tail.strip() if tail.strip() else "")).strip()
    try:
        spec = json.loads(payload)
        if isinstance(spec, dict) and spec.get("type") in ("pie", "bar", "line") and spec.get("series"):
            return clean, spec
    except (ValueError, TypeError):
        pass
    return text, None