Spaces:
Running on Zero
Running on Zero
| """ | |
| 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 | |