Update scenario_engine.py

scenario_engine.py

@@ -1,550 +1,37 @@
-# scenario_engine.py
-from __future__ import annotations
-from typing import Dict, List, Any, Tuple, Union, Optional
-import re
-import math
-import statistics
-import json
-import ast
-
 import pandas as pd
 import numpy as np
-
-# ----------------------------
-# Safe expression evaluation
-# ----------------------------
-_ALLOWED_FUNCS = {
-    "abs": abs,
-    "round": round,
-    "sqrt": math.sqrt,
-    "log": math.log,
-    "exp": math.exp,
-    "min": np.minimum,   # vectorized
-    "max": np.maximum,   # vectorized
-    "mean": np.mean,
-    "avg": np.mean,
-    "median": np.median,
-    "sum": np.sum,
-    "count": lambda x: np.size(x),
-    "p50": lambda x: np.percentile(x, 50),
-    "p75": lambda x: np.percentile(x, 75),
-    "p90": lambda x: np.percentile(x, 90),
-    "p95": lambda x: np.percentile(x, 95),
-    "p99": lambda x: np.percentile(x, 99),
-    "ceil": np.ceil,
-    "floor": np.floor,
-}
-
-class _SafeExpr(ast.NodeTransformer):
-    """
-    Restrict expressions to:
-    - Names (columns), numbers, strings, booleans
-    - Arithmetic: + - * / // % **, comparisons, and/or/not
-    - Calls to allowed functions (above)
-    """
-    def __init__(self, allowed_names: set):
-        self.allowed_names = allowed_names
-
-    def visit_Name(self, node):
-        if node.id not in self.allowed_names and node.id not in ("True","False","None"):
-            raise ValueError(f"Unknown name in expression: {node.id}")
-        return node
-
-    def visit_Call(self, node):
-        if not isinstance(node.func, ast.Name):
-            raise ValueError("Only simple function calls are allowed")
-        func = node.func.id
-        if func not in _ALLOWED_FUNCS:
-            raise ValueError(f"Function not allowed: {func}")
-        self.generic_visit(node)
-        return node
-
-    def generic_visit(self, node):
-        allowed = (
-            ast.Expression, ast.BoolOp, ast.BinOp, ast.UnaryOp,
-            ast.Compare, ast.Call, ast.Name, ast.Load, ast.Constant,
-            ast.And, ast.Or, ast.Not,
-            ast.Add, ast.Sub, ast.Mult, ast.Div, ast.Mod, ast.Pow, ast.FloorDiv,
-            ast.Eq, ast.NotEq, ast.Lt, ast.LtE, ast.Gt, ast.GtE, ast.In, ast.NotIn,
-            ast.USub, ast.UAdd
-        )
-        if not isinstance(node, allowed):
-            raise ValueError(f"Unsupported syntax: {type(node).__name__}")
-        return super().generic_visit(node)
-
-def _eval_series_expr(expr: str, df: pd.DataFrame) -> pd.Series:
-    allowed_names = set(df.columns) | {"True", "False", "None"}
-    tree = ast.parse(expr, mode="eval")
-    _SafeExpr(allowed_names).visit(tree)
-    code = compile(tree, "<expr>", "eval")
-    env = {**{k: df[k] for k in df.columns}, **_ALLOWED_FUNCS}
-    return eval(code, {"__builtins__": {}}, env)
-
-# ----------------------------
-# Engine
-# ----------------------------
+import math, json, re, ast
+from schemas import ScenarioPlan, TaskSpec
 
 class ScenarioEngine:
-    """
-    Scenario-first engine:
-    - Parse tasks + inline directives from scenario text
-    - For each task, execute a pipeline over analysis_results:
-        load -> filter -> derive -> groupby/agg -> pivot -> sort/top -> select fields -> render
-    - Render formats: table | list | comparison | map | narrative | chart (Vega-Lite spec)
-    - Strict: only what is asked is emitted.
-    """
-
-    @staticmethod
-    def render(scenario_text: str, analysis_results: Dict[str, Any]) -> str:
-        scen = ScenarioEngine._parse_scenario(scenario_text)
-        out: List[str] = ["# Scenario Output\n"]
-        for task in scen["tasks"]:
-            out.append(ScenarioEngine._render_task(task, analysis_results))
-        return "\n".join(out).strip()
-
-    # ------------- Parsing -------------
-    @staticmethod
-    def _parse_scenario(s: str) -> Dict[str, Any]:
-        """
-        Detect a 'Tasks/Deliverables/Requirements/Your Tasks' block; fallback to any bullet/numbered lines.
-        Each task may include inline directives: key: value
-        Supported directives (per task):
-            format: table|list|comparison|map|narrative|chart
-            data_key: <key in analysis_results>
-            filter: <expr using columns>   e.g., zone == "North" and wait_time > 5
-            derive: <col>=<expr>[, <col2>=<expr2> ...]
-            group_by: col1[, col2 ...]
-            agg: avg(x), median(y), sum(z), p90(wait), count(*)
-            pivot: index=a[,b] columns=c values=v  (values must be an aggregated column)
-            sort_by: col  sort_dir: asc|desc
-            top: N
-            fields: col1 col2 col3   (space or comma separated)
-            title: Custom name
-            chart: bar|line|area|point   (Vega-Lite spec emitted)
-            x: <field>  y: <field>  color: <field>  column: <facet>
-        """
-        lines = [ln.rstrip() for ln in s.splitlines()]
-        task_hdr = re.compile(r'^\s*(tasks?|deliverables|requirements|your tasks?)\s*$', re.I)
-        bullet = re.compile(r'^\s*(?:\d+\.\s+|[-*•]\s+)')
-        in_tasks = False
-        raw_tasks: List[str] = []
-
-        for ln in lines:
-            if task_hdr.match(ln):
-                in_tasks = True
-                continue
-            if in_tasks:
-                if bullet.match(ln.strip()):
-                    raw_tasks.append(ln.strip())
-                elif ln.strip() == "":
-                    continue
-                else:
-                    # stop when we hit a non-task looking line after capturing some tasks
-                    if raw_tasks:
-                        in_tasks = False
-
-        if not raw_tasks:
-            # fallback: grab any bullet/numbered lines
-            raw_tasks = [ln.strip() for ln in lines if bullet.match(ln.strip())]
-
-        tasks: List[Dict[str, Any]] = []
-        for raw in raw_tasks:
-            directives = ScenarioEngine._extract_directives(raw)
-            title = directives.get("title") or ScenarioEngine._strip_bullet(raw)
-            tasks.append({"title": title, "raw": raw, "d": directives})
-        return {"tasks": tasks}
-
-    @staticmethod
-    def _strip_bullet(line: str) -> str:
-        return re.sub(r'^\s*(?:\d+\.\s+|[-*•]\s+)', '', line).strip()
-
-    @staticmethod
-    def _extract_directives(text: str) -> Dict[str, Any]:
-        d: Dict[str, Any] = {}
-        # key: value pairs (value extends until ; or end or two spaces before next key:)
-        for m in re.finditer(r'([a-z_]+)\s*:\s*([^|,\n;]+)', text, re.I):
-            k = m.group(1).strip().lower()
-            v = m.group(2).strip()
-            d[k] = v
-
-        def _split_csv(val: str) -> List[str]:
-            return [x.strip() for x in re.split(r'[,\s]+', val) if x.strip()]
-
-        if "fields" in d:
-            d["fields"] = _split_csv(d["fields"])
-        if "group_by" in d:
-            d["group_by"] = _split_csv(d["group_by"])
-        if "top" in d:
-            try:
-                d["top"] = int(re.findall(r'\d+', d["top"])[0])
-            except Exception:
-                d["top"] = None
-        if "sort_dir" in d:
-            d["sort_dir"] = "desc" if d["sort_dir"].lower().startswith("d") else "asc"
-        if "format" in d:
-            d["format"] = d["format"].lower()
-        if "chart" in d:
-            d["chart"] = d["chart"].lower()
-        return d
-
-    # ------------- Rendering -------------
-    @staticmethod
-    def _render_task(task: Dict[str, Any], analysis_results: Dict[str, Any]) -> str:
-        title, d = task["title"], task["d"]
-        section: List[str] = [f"## {title}\n"]
-
-        # 1) Resolve data
-        df, key_used, why = ScenarioEngine._resolve_df(d, analysis_results)
-        if df is None:
-            section.append("_No matching data for this task._")
-            section.append(f"\n> Resolver note: {why}")
-            return "\n".join(section)
-
-        # 2) Filter
-        if "filter" in d:
-            mask = ScenarioEngine._safe_filter(df, d["filter"])
-            df = df.loc[mask].copy()
-
-        # 3) Derive columns
-        if "derive" in d:
-            df = ScenarioEngine._apply_derive(df, d["derive"])
-
-        # 4) Group & aggregate
-        if "group_by" in d or "agg" in d:
-            df = ScenarioEngine._group_agg(df, d.get("group_by"), d.get("agg"))
-
-        # 5) Pivot
-        if "pivot" in d:
-            df = ScenarioEngine._pivot(df, d["pivot"])
-
-        # 6) Sort + Top
-        if "sort_by" in d:
-            asc = (d.get("sort_dir", "desc") == "asc")
-            df = df.sort_values(by=d["sort_by"], ascending=asc)
-        if isinstance(d.get("top"), int) and d["top"] > 0:
-            df = df.head(d["top"])
-
-        # 7) Fields selection
-        if "fields" in d:
-            cols = [c for c in d["fields"] if c in df.columns]
-            if cols:
-                df = df[cols]
-
-        # 8) Render by format
-        fmt = d.get("format", "table")
-        if fmt == "list":
-            section.append(ScenarioEngine._render_list(df))
-        elif fmt == "comparison":
-            section.append(ScenarioEngine._render_comparison(df))
-        elif fmt == "map":
-            section.append(ScenarioEngine._render_map(df))
-        elif fmt == "narrative":
-            section.append(ScenarioEngine._render_narrative(df))
-        elif fmt == "chart":
-            section.append(ScenarioEngine._render_chart_spec(df, d))
-        else:
-            section.append(ScenarioEngine._render_table(df))
-
-        # 9) Per-task provenance (kept minimal)
-        section.append("\n**Provenance**")
-        section.append(f"- Data key: `{key_used}`")
-        section.append(f"- Match note: {why}")
-
-        return "\n".join(section)
-
-    # ------------- Data resolution -------------
     @staticmethod
-    def _resolve_df(d: Dict[str, Any], analysis_results: Dict[str, Any]) -> Tuple[Optional[pd.DataFrame], Optional[str], str]:
-        # explicit key
-        if "data_key" in d and d["data_key"] in analysis_results:
-            return ScenarioEngine._as_df(analysis_results[d["data_key"]]), d["data_key"], "explicit data_key"
-
-        # jaccard match on keys using hinted fields + any words in title/sort/agg
-        hints = set()
-        for k in ("fields", "sort_by"):
-            v = d.get(k)
-            if isinstance(v, list):
-                hints |= set(v)
-            elif isinstance(v, str):
-                hints |= set(re.findall(r'[A-Za-z0-9_]+', v.lower()))
-        best_key, best_score = None, 0.0
-        for k in analysis_results:
-            words = set(re.findall(r'[A-Za-z0-9_]+', k.lower()))
-            if not words:
-                continue
-            inter = len(hints & words)
-            union = len(hints | words) or 1
-            score = inter / union
-            if score > best_score:
-                best_key, best_score = k, score
-
-        if best_key:
-            return ScenarioEngine._as_df(analysis_results[best_key]), best_key, f"keyword match (score={best_score:.2f})"
-
-        # fallback: first list-of-dicts or dict-like
-        for k, v in analysis_results.items():
-            df = ScenarioEngine._as_df(v)
-            if df is not None and not df.empty:
-                return df, k, "fallback first structured"
-
-        return None, None, "no suitable dataset found"
+    def render_plan(plan: ScenarioPlan, results: dict) -> str:
+        sections = ["# Scenario Output\n"]
+        for t in plan.tasks:
+            sections.append(ScenarioEngine._render_task(t, results))
+        return "\n".join(sections)
 
     @staticmethod
-    def _as_df(v: Any) -> Optional[pd.DataFrame]:
-        if isinstance(v, list):
-            if not v:
-                return pd.DataFrame()
-            if isinstance(v[0], dict):
-                return pd.DataFrame(v)
-            return pd.DataFrame({"value": v})
-        if isinstance(v, dict):
-            # expand nested dicts into columns where sensible
-            flat = {}
-            any_scalar = False
-            for k, val in v.items():
-                if isinstance(val, (int, float, str, bool, type(None))):
-                    flat[k] = [val]
-                    any_scalar = True
-            if any_scalar:
-                return pd.DataFrame(flat)
-            # complex dict -> try records
-            recs = []
-            for k, val in v.items():
-                if isinstance(val, dict):
-                    rec = {"item": k}
-                    rec.update({kk: valv for kk, valv in val.items()})
-                    recs.append(rec)
-            if recs:
-                return pd.DataFrame(recs)
+    def _df(v):
+        if isinstance(v, pd.DataFrame): return v
+        if isinstance(v, list) and v and isinstance(v[0], dict): return pd.DataFrame(v)
+        if isinstance(v, dict): return pd.DataFrame([v])
         return None
 
-    # ------------- Pipeline ops -------------
-    @staticmethod
-    def _safe_filter(df: pd.DataFrame, expr: str) -> pd.Series:
-        try:
-            s = _eval_series_expr(expr, df)
-            if not isinstance(s, (pd.Series, np.ndarray)):
-                raise ValueError("filter must evaluate to a boolean Series/array")
-            return pd.Series(s).astype(bool).reindex(df.index, fill_value=False)
-        except Exception as e:
-            raise ValueError(f"Invalid filter expression: {e}")
-
-    @staticmethod
-    def _apply_derive(df: pd.DataFrame, spec: str) -> pd.DataFrame:
-        # e.g., "load = patients / capacity, rate = 100*admits/pop"
-        parts = re.split(r'[;,]\s*', spec)
-        for p in parts:
-            if not p.strip():
-                continue
-            if "=" not in p:
-                raise ValueError(f"derive requires assignments: '{p}'")
-            col, expr = p.split("=", 1)
-            col = col.strip()
-            expr = expr.strip()
-            df[col] = _eval_series_expr(expr, df)
-        return df
-
-    @staticmethod
-    def _parse_aggs(spec: Optional[str]) -> List[Tuple[str, str]]:
-        """
-        Returns list of (out_col, func_call_string), e.g. [("avg_wait_time","avg(wait_time)")]
-        """
-        if not spec:
-            return []
-        items = [x.strip() for x in spec.split(",") if x.strip()]
-        out: List[Tuple[str, str]] = []
-        for it in items:
-            m = re.match(r'([a-zA-Z_][a-zA-Z0-9_]*)\s*\(([^)]+)\)', it)
-            if not m:
-                if it.lower() in ("count", "count(*)"):
-                    out.append(("count", "count(*)"))
-                    continue
-                raise ValueError(f"Bad agg item: '{it}' (use avg(x), median(y), p90(z), sum(a), count(*))")
-            func = m.group(1)
-            arg = m.group(2).strip()
-            out_col = f"{func.lower()}_{arg}"
-            out.append((out_col, f"{func}({arg})"))
-        return out
-
-    @staticmethod
-    def _group_agg(df: pd.DataFrame, group_by: Optional[List[str]], agg_spec: Optional[str]) -> pd.DataFrame:
-        aggs = ScenarioEngine._parse_aggs(agg_spec)
-        if not aggs and not group_by:
-            return df
-        if not group_by:
-            # reduce to single row with requested aggs
-            res = {}
-            for out_col, call in aggs:
-                val = ScenarioEngine._apply_agg_call(df, call)
-                res[out_col] = val
-            return pd.DataFrame([res])
-        # grouped
-        gb = df.groupby(group_by, dropna=False)
-        rows = []
-        for keys, g in gb:
-            if not isinstance(keys, tuple):
-                keys = (keys,)
-            rec = {group_by[i]: keys[i] for i in range(len(group_by))}
-            for out_col, call in aggs:
-                rec[out_col] = ScenarioEngine._apply_agg_call(g, call)
-            if not aggs:
-                # no aggs? carry counts by default
-                rec["count"] = len(g)
-            rows.append(rec)
-        return pd.DataFrame(rows)
-
     @staticmethod
-    def _apply_agg_call(df: pd.DataFrame, call: str) -> Any:
-        call = call.strip()
-        if call.lower() in ("count", "count(*)"):
-            return int(len(df))
-        m = re.match(r'([a-zA-Z_][a-zA-Z0-9_]*)\s*\(([^)]+)\)', call)
-        if not m:
-            raise ValueError(f"Bad agg call: {call}")
-        func, arg = m.group(1).lower(), m.group(2).strip()
-        if arg not in df.columns:
-            raise ValueError(f"Unknown column in agg: {arg}")
-        col = df[arg].dropna()
-        if func in ("avg", "mean"):
-            return float(np.mean(col)) if len(col) else float("nan")
-        if func == "median":
-            return float(np.median(col)) if len(col) else float("nan")
-        if func == "sum":
-            return float(np.sum(col)) if len(col) else 0.0
-        if func in ("min", "max"):
-            f = getattr(np, func)
-            return float(f(col)) if len(col) else float("nan")
-        if func.startswith("p") and func[1:].isdigit():
-            q = int(func[1:])
-            return float(np.percentile(col, q)) if len(col) else float("nan")
-        raise ValueError(f"Unsupported agg function: {func}")
+    def _render_task(t: TaskSpec, results: dict) -> str:
+        out=[f"## {t.title}\n"]
+        df=ScenarioEngine._df(results.get(t.data_key)) if t.data_key else None
+        if df is None: return "\n".join(out+["_No data_"])
+        if t.filter: df=df.query(t.filter)
+        if t.group_by or t.agg: df=df.groupby(t.group_by).agg("first").reset_index()
+        if t.sort_by: df=df.sort_values(by=t.sort_by, ascending=(t.sort_dir=="asc"))
+        if t.top: df=df.head(t.top)
+        if t.fields: df=df[t.fields]
 
-    @staticmethod
-    def _pivot(df: pd.DataFrame, spec: str) -> pd.DataFrame:
-        # spec: index=a[,b] columns=c values=v
-        parts = dict(re.findall(r'(\w+)\s*=\s*([^\s,]+)', spec))
-        idx = parts.get("index")
-        cols = parts.get("columns")
-        vals = parts.get("values")
-        if not (idx and cols and vals):
-            raise ValueError("pivot requires 'index=.. columns=.. values=..'")
-        idx = [x.strip() for x in idx.split(",")]
-        pv = df.pivot_table(index=idx, columns=cols, values=vals, aggfunc="first")
-        pv = pv.reset_index()
-        # flatten columns if needed
-        if isinstance(pv.columns, pd.MultiIndex):
-            pv.columns = ["_".join([str(c) for c in tup if c != ""]) for tup in pv.columns]
-        return pv
-
-    # ------------- Output renderers -------------
-    @staticmethod
-    def _render_table(df: pd.DataFrame) -> str:
-        if df.empty:
-            return "_No rows to display._"
-        # convert all to string-friendly
-        dff = df.copy()
-        for c in dff.columns:
-            dff[c] = dff[c].apply(lambda v: ScenarioEngine._fmt_val(v))
-        header = "| " + " | ".join(dff.columns) + " |"
-        sep = "|" + "|".join(["---"] * len(dff.columns)) + "|"
-        rows = ["| " + " | ".join(map(str, r)) + " |" for r in dff.to_numpy().tolist()]
-        return "\n".join([header, sep, *rows])
-
-    @staticmethod
-    def _render_list(df: pd.DataFrame) -> str:
-        if df.empty:
-            return "_No items._"
-        # pick first column as primary
-        primary = df.columns[0]
-        lines = []
-        for i, row in enumerate(df.itertuples(index=False), 1):
-            parts = []
-            for c, v in zip(df.columns, row):
-                if c == primary:
-                    continue
-                parts.append(f"{c}: {ScenarioEngine._fmt_val(v)}")
-            extra = f" ({', '.join(parts)})" if parts else ""
-            lines.append(f"{i}. {ScenarioEngine._fmt_val(getattr(row, primary))}{extra}")
-        return "\n".join(lines)
-
-    @staticmethod
-    def _render_comparison(df: pd.DataFrame) -> str:
-        # look for columns named like current/previous
-        cols = {c.lower(): c for c in df.columns}
-        cur = cols.get("current") or cols.get("now") or cols.get("value")
-        prev = cols.get("previous") or cols.get("prior") or cols.get("past")
-        name = cols.get("name") or cols.get("metric") or cols.get("item") or df.columns[0]
-        if not (cur and prev):
-            return "_Comparison format requires columns 'current' and 'previous' (or aliases)._"
-        header = "| Item | Current | Previous | Change |"
-        sep = "|---|---:|---:|---:|"
-        body = []
-        for _, r in df.iterrows():
-            c, p = r[cur], r[prev]
-            change = (c - p) if isinstance(c, (int, float)) and isinstance(p, (int, float)) else "N/A"
-            body.append(f"| {ScenarioEngine._fmt_val(r[name])} | {ScenarioEngine._fmt_val(c)} | {ScenarioEngine._fmt_val(p)} | {ScenarioEngine._fmt_val(change)} |")
-        return "\n".join([header, sep, *body])
-
-    @staticmethod
-    def _render_map(df: pd.DataFrame) -> str:
-        # simple location table
-        colmap = {c.lower(): c for c in df.columns}
-        name = colmap.get("name") or colmap.get("facility") or colmap.get("title") or df.columns[0]
-        zone = colmap.get("zone")
-        city = colmap.get("city")
-        region = colmap.get("region")
-        lat = colmap.get("latitude") or colmap.get("lat")
-        lon = colmap.get("longitude") or colmap.get("lon")
-        cols = [x for x in [name, city, region, zone, lat, lon] if x]
-        if not cols:
-            return "_No geographic fields to show._"
-        dff = df[cols].copy()
-        dff["coordinates"] = np.where((lat is not None) & (lon is not None) & dff[lat].notna() & dff[lon].notna(),
-                                      dff[lat].astype(str) + ", " + dff[lon].astype(str), "N/A")
-        show = [name, city or "city", region or "region", zone or "zone", "coordinates"]
-        # ensure all exist
-        for c in show:
-            if c not in dff.columns:
-                dff[c] = ""
-        dff = dff[show]
-        return ScenarioEngine._render_table(dff)
-
-    @staticmethod
-    def _render_narrative(df: pd.DataFrame) -> str:
-        if df.empty:
-            return "_No content._"
-        paras = []
-        for i, row in enumerate(df.to_dict(orient="records"), 1):
-            parts = [f"**{k}**: {ScenarioEngine._fmt_val(v)}" for k, v in row.items()]
-            paras.append(f"{i}. " + "; ".join(parts))
-        return "\n".join(paras)
-
-    @staticmethod
-    def _render_chart_spec(df: pd.DataFrame, d: Dict[str, Any]) -> str:
-        """
-        Emits a Vega-Lite spec in a fenced code block that downstream renderers can plot exactly.
-        Accepts: chart (bar|line|area|point), x, y, color, column (facet)
-        """
-        mark = d.get("chart", "bar")
-        spec = {
-            "$schema": "https://vega.github.io/schema/vega-lite/v5.json",
-            "description": d.get("title") or "Chart",
-            "data": {"values": df.to_dict(orient="records")},
-            "mark": mark,
-            "encoding": {}
-        }
-        for enc in ("x", "y", "color", "column"):
-            if enc in d and d[enc] in df.columns:
-                spec["encoding"][enc] = {"field": d[enc], "type": "quantitative" if pd.api.types.is_numeric_dtype(df[d[enc]]) else "nominal"}
-        return "```vega-lite\n" + json.dumps(spec, ensure_ascii=False, indent=2) + "\n```"
+        if t.format=="list":
+            out += [f"- {row.to_dict()}" for _,row in df.iterrows()]
+        else:
+            out.append(df.to_markdown(index=False))
+        return "\n".join(out)
 
-    # ------------- Helpers -------------
-    @staticmethod
-    def _fmt_val(v: Any) -> str:
-        if isinstance(v, float):
-            if math.isnan(v):
-                return "NaN"
-            return f"{v:,.4g}"
-        if isinstance(v, (int, np.integer)):
-            return f"{int(v):,}"
-        return str(v)