syntaxmatrix 2.5.5.5__py3-none-any.whl → 2.5.6__py3-none-any.whl

syntaxmatrix/utils.py

@@ -93,6 +93,55 @@ def classify_ml_job(prompt: str) -> str:
     return "eda"
 
 
+def _indent(code: str, spaces: int = 4) -> str:
+    """
+    Indent a block of code by `spaces` spaces, line by line.
+    Blank lines are preserved unchanged.
+    """
+    pad = " " * spaces
+    lines = code.splitlines()
+    return "\n".join((pad + line) if line.strip() else line for line in lines)
+
+
+def wrap_llm_code_safe(body: str) -> str:
+    """
+    Wrap arbitrary LLM code so that:
+      - Any exception is caught and shown.
+      - A minimal, useful EDA fallback still runs so the user sees *something*.
+    This happens once in the framework; you never touch the individual cells.
+    """
+    return textwrap.dedent(
+        "try:\n"
+        + _indent(body)
+        + "\n"
+        "except Exception as e:\n"
+        "    from syntaxmatrix.display import show\n"
+        "    msg = f\"⚠️ Skipped LLM block due to: {type(e).__name__}: {e}\"\n"
+        "    show(msg)\n"
+        "    # --- automatic EDA fallback ---\n"
+        "    try:\n"
+        "        df_local = globals().get('df')\n"
+        "        if df_local is not None:\n"
+        "            import pandas as pd\n"
+        "            from syntaxmatrix.preface import SB_histplot, _SMX_export_png\n"
+        "            num_cols = df_local.select_dtypes(include=['number', 'bool']).columns.tolist()\n"
+        "            cat_cols = [c for c in df_local.columns if c not in num_cols]\n"
+        "            info = {\n"
+        "                'rows': len(df_local),\n"
+        "                'cols': len(df_local.columns),\n"
+        "                'numeric_cols': len(num_cols),\n"
+        "                'categorical_cols': len(cat_cols),\n"
+        "            }\n"
+        "            show(df_local.head())\n"
+        "            show(info)\n"
+        "            if num_cols:\n"
+        "                SB_histplot()\n"
+        "                _SMX_export_png()\n"
+        "    except Exception as _f:\n"
+        "        show(f\"⚠️ Fallback EDA failed: {type(_f).__name__}: {_f}\")\n"
+    )
+
+
 def harden_ai_code(code: str) -> str:
     """
     Make any AI-generated cell resilient:
@@ -113,10 +162,6 @@ def harden_ai_code(code: str) -> str:
         code = re.sub(r"\breturn\s+\.\.\.", "return None", code)
         return code
 
-    def _indent(code: str, spaces: int = 4) -> str:
-        pad = " " * spaces
-        return "\n".join(pad + line for line in code.splitlines())
-
     def _SMX_OHE(**k):
         # normalise arg name across sklearn versions
         if "sparse" in k and "sparse_output" not in k:
@@ -140,6 +185,49 @@ def harden_ai_code(code: str) -> str:
         code = re.sub(r'(?m)^\s*\\\d+\s*', '', code)     
         code = re.sub(r'(?m)[;]\s*\\\d+\s*', '; ', code)  
         return code
+    
+    def _patch_feature_coef_dataframe(code: str) -> str:
+        """
+        Harden patterns like:
+            coeffs_df = pd.DataFrame({'feature': num_features, 'coefficient': coef})
+        which can crash with:
+            ValueError: All arrays must be of the same length
+        We wrap them in a try/except and, on failure, rebuild the
+        DataFrame by zipping feature names with coefficients up to
+        the min length.
+        """
+        # Match single-line assignments of the form:
+        #   <var> = pd.DataFrame({'feature': <feat>, 'coefficient': <coef>})
+        import re
+
+        pattern = re.compile(
+            r"(?P<indent>^[ \t]*)"
+            r"(?P<var>\w+)\s*=\s*pd\.DataFrame\(\s*{\s*"
+            r"['\"]feature['\"]\s*:\s*(?P<feat_expr>.+?)\s*,\s*"
+            r"['\"]coefficient['\"]\s*:\s*(?P<coef_expr>.+?)\s*"
+            r"}\s*\)\s*$",
+            re.MULTILINE,
+        )
+
+        def repl(m: re.Match) -> str:
+            indent = m.group("indent")
+            var = m.group("var")
+            feat_expr = m.group("feat_expr").strip()
+            coef_expr = m.group("coef_expr").strip()
+
+            # Keep the original intent, but add a safe fallback.
+            return (
+                f"{indent}try:\n"
+                f"{indent}    {var} = pd.DataFrame({{'feature': {feat_expr}, 'coefficient': {coef_expr}}})\n"
+                f"{indent}except Exception as _e:\n"
+                f"{indent}    import numpy as _np\n"
+                f"{indent}    _feat = list({feat_expr})\n"
+                f"{indent}    _coef = _np.asarray({coef_expr}).ravel()\n"
+                f"{indent}    _k = min(len(_feat), len(_coef))\n"
+                f"{indent}    {var} = pd.DataFrame({{'feature': _feat[:_k], 'coefficient': _coef[:_k]}})\n"
+            )
+
+        return pattern.sub(repl, code)
 
     def _wrap_metric_calls(code: str) -> str:
         names = [
@@ -224,6 +312,191 @@ def harden_ai_code(code: str) -> str:
         except Exception:
             pass
 
+    def _ensure_metrics_imports(code: str) -> str:
+        needed = set()
+        if "r2_score" in code:
+            needed.add("r2_score")
+        if "mean_absolute_error" in code:
+            needed.add("mean_absolute_error")
+        # ... add others if you like ...
+
+        if not needed:
+            return code
+
+        if "from sklearn.metrics import" in code:
+            return code  # assume user/LLM handled it
+
+        import_line = "from sklearn.metrics import " + ", ".join(sorted(needed)) + "\n"
+        return import_line + code
+
+    def _fix_unexpected_indent(src: str) -> str:
+        """
+        Some LLM snippets jump indentation (e.g. extra 8 spaces on an 'import'
+        line) without a preceding block opener. That causes
+        `IndentationError: unexpected indent` when we wrap in our own `try:`.
+        This normalises those lines back to the previous indent level, but only
+        when we're not in a multi-line bracket/paren context.
+        """
+        lines = src.splitlines()
+        out = []
+        prev_indent = 0
+        prev_ends_colon = False
+        paren_depth = 0  # (), [], {} depth across lines (very approximate)
+
+        for raw in lines:
+            stripped = raw.lstrip()
+            if not stripped:  # blank / whitespace line
+                out.append(raw)
+                continue
+
+            indent = len(raw) - len(stripped)
+
+            # Only flatten if:
+            #   - we're not inside a (...) / [...] / {...} block, and
+            #   - previous logical line did NOT end with ':', and
+            #   - this line is indented more than the previous indent.
+            if paren_depth == 0 and not prev_ends_colon and indent > prev_indent:
+                indent = prev_indent
+            new_line = " " * indent + stripped
+            out.append(new_line)
+
+            # Update simple state for next line
+            txt = stripped
+            paren_depth += txt.count("(") + txt.count("[") + txt.count("{")
+            paren_depth -= txt.count(")") + txt.count("]") + txt.count("}")
+            prev_ends_colon = txt.rstrip().endswith(":")
+            prev_indent = indent
+
+        return "\n".join(out)
+
+    def _fallback_snippet() -> str:
+        """
+        Final-resort snippet when the LLM code is syntactically broken.
+
+        It:
+          - attempts a simple automatic ML task (classification or regression)
+          - then falls back to generic but useful EDA.
+
+        It assumes `from syntaxmatrix.preface import *` has already been done,
+        so `_SMX_OHE`, `_SMX_call`, `SB_histplot`, `_SMX_export_png` and the
+        patched `show()` are available.
+        """
+        import textwrap
+
+        return textwrap.dedent(
+            """\
+            import pandas as pd
+            import numpy as np
+            import matplotlib.pyplot as plt
+            from sklearn.model_selection import train_test_split
+            from sklearn.compose import ColumnTransformer
+            from sklearn.preprocessing import StandardScaler
+            from sklearn.linear_model import LogisticRegression, LinearRegression
+            from sklearn.metrics import accuracy_score, r2_score, mean_absolute_error
+
+            df = df.copy()
+
+            # --- basic column introspection ---
+            num_cols = df.select_dtypes(include=['number', 'bool']).columns.tolist()
+            cat_cols = [c for c in df.columns if c not in num_cols]
+
+            # --- attempt an automatic ML task ---
+            target_col = None
+            task_type = None
+
+            # Prefer a low-cardinality target (classification)
+            for c in num_cols + cat_cols:
+                uniq = df[c].dropna().nunique()
+                if 2 <= uniq <= 10:
+                    target_col = c
+                    task_type = 'classification'
+                    break
+
+            # If none found, try a numeric regression target
+            if target_col is None and num_cols:
+                target_col = num_cols[-1]
+                task_type = 'regression'
+
+            if target_col is not None:
+                try:
+                    X = df.drop(columns=[target_col]).copy()
+                    y = df[target_col].copy()
+
+                    num_feats = X.select_dtypes(include=['number', 'bool']).columns.tolist()
+                    cat_feats = [c for c in X.columns if c not in num_feats]
+
+                    pre = ColumnTransformer(
+                        transformers=[
+                            ('num', StandardScaler(), num_feats),
+                            ('cat', _SMX_OHE(handle_unknown='ignore'), cat_feats),
+                        ],
+                        remainder='drop',
+                    )
+
+                    from sklearn.pipeline import Pipeline
+                    if task_type == 'classification':
+                        model = LogisticRegression(max_iter=1000)
+                    else:
+                        model = LinearRegression()
+
+                    pipe = Pipeline([('pre', pre), ('model', model)])
+
+                    X_train, X_test, y_train, y_test = train_test_split(
+                        X, y, test_size=0.25, random_state=42
+                    )
+
+                    pipe.fit(X_train, y_train)
+                    y_pred = pipe.predict(X_test)
+
+                    if task_type == 'classification':
+                        # If predictions look like probabilities, convert to labels
+                        if getattr(y_pred, 'ndim', 1) > 1 and y_pred.shape[1] > 1:
+                            y_pred_labels = y_pred.argmax(axis=1)
+                        else:
+                            try:
+                                y_pred_labels = (y_pred > 0.5).astype(y_test.dtype)
+                            except Exception:
+                                y_pred_labels = y_pred
+
+                        acc = _SMX_call(accuracy_score, y_test, y_pred_labels)
+                        show({
+                            'target': target_col,
+                            'task': 'classification',
+                            'accuracy': acc,
+                        })
+                    else:
+                        r2 = _SMX_call(r2_score, y_test, y_pred)
+                        mae = _SMX_call(mean_absolute_error, y_test, y_pred)
+                        show({
+                            'target': target_col,
+                            'task': 'regression',
+                            'r2': r2,
+                            'mae': mae,
+                        })
+
+                except Exception as _ml_e:
+                    show(f"⚠ ML fallback failed: {type(_ml_e).__name__}: {_ml_e}")
+
+            # --- EDA fallback that still helps answer the question ---
+            try:
+                info = {
+                    'rows': len(df),
+                    'cols': len(df.columns),
+                    'numeric_cols': len(num_cols),
+                    'categorical_cols': len(cat_cols),
+                }
+                show(df.head(), title='Sample of data')
+                show(info, title='Dataset summary')
+
+                # Quick univariate look if we have numeric columns
+                if num_cols:
+                    SB_histplot()
+                    _SMX_export_png()
+            except Exception as _eda_e:
+                show(f"⚠ EDA fallback failed: {type(_eda_e).__name__}: {_eda_e}")
+            """
+        )
+
     # Register and run patches once per execution
     for _patch in (
         _smx_patch_mean_squared_error_squared_kw,
@@ -235,412 +508,6 @@ def harden_ai_code(code: str) -> str:
         except Exception:
             pass
 
-    PREFACE = (
-        "# === SMX Auto-Hardening Preface (do not edit) ===\n"
-        "import warnings, numpy as np, pandas as pd, matplotlib.pyplot as plt\n"
-        "warnings.filterwarnings('ignore')\n"
-        "try:\n"
-        "    import seaborn as sns\n"
-        "except Exception:\n"
-        "    class _Dummy:\n"
-        "        def __getattr__(self, name):\n"
-        "            def _f(*a, **k):\n"
-        "                from syntaxmatrix.display import show\n"
-        "                show('⚠ seaborn not available; plot skipped.')\n"
-        "            return _f\n"
-        "    sns = _Dummy()\n"
-        "\n"
-        "from syntaxmatrix.display import show as _SMX_base_show\n"
-        "def _SMX_caption_from_ctx():\n"
-        "    g = globals()\n"
-        "    t = g.get('refined_question') or g.get('askai_question') or 'Table'\n"
-        "    return str(t).strip().splitlines()[0][:120]\n"
-        "\n"
-        "def show(obj, title=None):\n"
-        "    try:\n"
-        "        import pandas as pd\n"
-        "        if isinstance(obj, pd.DataFrame):\n"
-        "            cap = (title or _SMX_caption_from_ctx())\n"
-        "            try:\n"
-        "                return _SMX_base_show(obj.style.set_caption(cap))\n"
-        "            except Exception:\n"
-        "                pass\n"
-        "    except Exception:\n"
-        "        pass\n"
-        "    return _SMX_base_show(obj)\n"
-        "\n"
-        "def _SMX_axes_have_titles(fig=None):\n"
-        "    import matplotlib.pyplot as _plt\n"
-        "    fig = fig or _plt.gcf()\n"
-        "    try:\n"
-        "        for _ax in fig.get_axes():\n"
-        "            if (_ax.get_title() or '').strip():\n"
-        "                return True\n"
-        "    except Exception:\n"
-        "        pass\n"
-        "    return False\n"
-        "\n"
-        "def _SMX_export_png():\n"
-        "    import io, base64\n"
-        "    fig = plt.gcf()\n"
-        "    try:\n"
-        "        if not _SMX_axes_have_titles(fig):\n"
-        "            fig.suptitle(_SMX_caption_from_ctx(), fontsize=10)\n"
-        "    except Exception:\n"
-        "        pass\n"
-        "    buf = io.BytesIO()\n"
-        "    plt.savefig(buf, format='png', bbox_inches='tight')\n"
-        "    buf.seek(0)\n"
-        "    from IPython.display import display, HTML\n"
-        "    _img = base64.b64encode(buf.read()).decode('ascii')\n"
-        "    display(HTML(f\"<img src='data:image/png;base64,{_img}' style='max-width:100%;height:auto;border:1px solid #ccc;border-radius:4px;'/>\"))\n"
-        "    plt.close()\n"
-        "\n"
-        "def _pick_df():\n"
-        "    return globals().get('df', None)\n"
-        "\n"
-        "def _pick_ax_slot():\n"
-        "    ax = None\n"
-        "    try:\n"
-        "        _axes = globals().get('axes', None)\n"
-        "        import numpy as _np\n"
-        "        if _axes is not None:\n"
-        "            arr = _np.ravel(_axes)\n"
-        "            for _a in arr:\n"
-        "                try:\n"
-        "                    if hasattr(_a,'has_data') and not _a.has_data():\n"
-        "                        ax = _a; break\n"
-        "                except Exception:\n"
-        "                    continue\n"
-        "    except Exception:\n"
-        "        ax = None\n"
-        "    return ax\n"
-        "\n"
-        "def _first_numeric(_d):\n"
-        "    import numpy as np, pandas as pd\n"
-        "    try:\n"
-        "        preferred = [\"median_house_value\", \"price\", \"value\", \"target\", \"label\", \"y\"]\n"
-        "        for c in preferred:\n"
-        "            if c in _d.columns and pd.api.types.is_numeric_dtype(_d[c]):\n"
-        "                return c\n"
-        "        cols = _d.select_dtypes(include=[np.number]).columns.tolist()\n"
-        "        return cols[0] if cols else None\n"
-        "    except Exception:\n"
-        "        return None\n"
-        "\n"
-        "def _first_categorical(_d):\n"
-        "    import pandas as pd, numpy as np\n"
-        "    try:\n"
-        "        num = set(_d.select_dtypes(include=[np.number]).columns.tolist())\n"
-        "        cand = [c for c in _d.columns if c not in num and _d[c].nunique(dropna=True) <= 50]\n"
-        "        return cand[0] if cand else None\n"
-        "    except Exception:\n"
-        "        return None\n"
-        "\n"
-        "boxplot = barplot = histplot = distplot = lineplot = countplot = heatmap = pairplot = None\n"
-        "\n"
-        "def _safe_plot(func, *args, **kwargs):\n"
-        "    try:\n"
-        "        ax = func(*args, **kwargs)\n"
-        "        if ax is None:\n"
-        "            ax = plt.gca()\n"
-        "        try:\n"
-        "            if hasattr(ax, 'has_data') and not ax.has_data():\n"
-        "                from syntaxmatrix.display import show as _show\n"
-        "                _show('⚠ Empty plot: no data drawn.')\n"
-        "        except Exception:\n"
-        "            pass\n"
-        "        try: plt.tight_layout()\n"
-        "        except Exception: pass\n"
-        "        return ax\n"
-        "    except Exception as e:\n"
-        "        from syntaxmatrix.display import show as _show\n"
-        "        _show(f'⚠ Plot skipped: {type(e).__name__}: {e}')\n"
-        "        return None\n"
-        "\n"
-        "def SB_histplot(*a, **k):\n"
-        "    _missing = (getattr(sns, '__class__', type(sns)).__name__ == '_Dummy')\n"
-        "    _sentinel = (len(a) >= 1 and a[0] is None)\n"
-        "    if (not a or _sentinel) and not k:\n"
-        "        d = _pick_df()\n"
-        "        if d is not None:\n"
-        "            x = _first_numeric(d)\n"
-        "            if x is not None:\n"
-        "                def _draw():\n"
-        "                    plt.hist(d[x].dropna())\n"
-        "                    ax = plt.gca()\n"
-        "                    if not (ax.get_title() or '').strip():\n"
-        "                        ax.set_title(f'Distribution of {x}')\n"
-        "                    return ax\n"
-        "                return _safe_plot(lambda **kw: _draw())\n"
-        "    if _missing:\n"
-        "        return _safe_plot(lambda **kw: plt.hist([]))\n"
-        "    if _sentinel:\n"
-        "        a = a[1:]\n"
-        "    return _safe_plot(getattr(sns,'histplot', plt.hist), *a, **k)\n"
-        "\n"
-        "def SB_barplot(*a, **k):\n"
-        "    _missing = (getattr(sns, '__class__', type(sns)).__name__ == '_Dummy')\n"
-        "    _sentinel = (len(a) >= 1 and a[0] is None)\n"
-        "    _ax = k.get('ax') or _pick_ax_slot()\n"
-        "    if _ax is not None:\n"
-        "        try: plt.sca(_ax)\n"
-        "        except Exception: pass\n"
-        "        k.setdefault('ax', _ax)\n"
-        "    if (not a or _sentinel) and not k:\n"
-        "        d = _pick_df()\n"
-        "        if d is not None:\n"
-        "            x = _first_categorical(d)\n"
-        "            y = _first_numeric(d)\n"
-        "            if x and y:\n"
-        "                import pandas as _pd\n"
-        "                g = d.groupby(x)[y].mean().reset_index()\n"
-        "                def _draw():\n"
-        "                    if _missing:\n"
-        "                        plt.bar(g[x], g[y])\n"
-        "                    else:\n"
-        "                        sns.barplot(data=g, x=x, y=y, ax=k.get('ax'))\n"
-        "                    ax = plt.gca()\n"
-        "                    if not (ax.get_title() or '').strip():\n"
-        "                        ax.set_title(f'Mean {y} by {x}')\n"
-        "                    return ax\n"
-        "                return _safe_plot(lambda **kw: _draw())\n"
-        "    if _missing:\n"
-        "        return _safe_plot(lambda **kw: plt.bar([], []))\n"
-        "    if _sentinel:\n"
-        "        a = a[1:]\n"
-        "    return _safe_plot(sns.barplot, *a, **k)\n"
-        "\n"
-        "def SB_boxplot(*a, **k):\n"
-        "    _missing = (getattr(sns, '__class__', type(sns)).__name__ == '_Dummy')\n"
-        "    _sentinel = (len(a) >= 1 and a[0] is None)\n"
-        "    _ax = k.get('ax') or _pick_ax_slot()\n"
-        "    if _ax is not None:\n"
-        "        try: plt.sca(_ax)\n"
-        "        except Exception: pass\n"
-        "        k.setdefault('ax', _ax)\n"
-        "    if (not a or _sentinel) and not k:\n"
-        "        d = _pick_df()\n"
-        "        if d is not None:\n"
-        "            x = _first_categorical(d)\n"
-        "            y = _first_numeric(d)\n"
-        "            if x and y:\n"
-        "                def _draw():\n"
-        "                    if _missing:\n"
-        "                        plt.boxplot(d[y].dropna())\n"
-        "                    else:\n"
-        "                        sns.boxplot(data=d, x=x, y=y, ax=k.get('ax'))\n"
-        "                    ax = plt.gca()\n"
-        "                    if not (ax.get_title() or '').strip():\n"
-        "                        ax.set_title(f'Distribution of {y} by {x}')\n"
-        "                    return ax\n"
-        "                return _safe_plot(lambda **kw: _draw())\n"
-        "    if _missing:\n"
-        "        return _safe_plot(lambda **kw: plt.boxplot([]))\n"
-        "    if _sentinel:\n"
-        "        a = a[1:]\n"
-        "    return _safe_plot(sns.boxplot, *a, **k)\n"
-        "\n"
-        "def SB_scatterplot(*a, **k):\n"
-        "    _missing = (getattr(sns, '__class__', type(sns)).__name__ == '_Dummy')\n"
-        "    fn = getattr(sns,'scatterplot', None)\n"
-        "    # If seaborn is unavailable OR the caller passed (data=..., x='col', y='col'),\n"
-        "    # use a robust matplotlib path that looks up data and coerces to numeric.\n"
-        "    if _missing or fn is None:\n"
-        "        data = k.get('data'); x = k.get('x'); y = k.get('y')\n"
-        "        if data is not None and isinstance(x, str) and isinstance(y, str) and x in data.columns and y in data.columns:\n"
-        "            xs = pd.to_numeric(data[x], errors='coerce')\n"
-        "            ys = pd.to_numeric(data[y], errors='coerce')\n"
-        "            m = xs.notna() & ys.notna()\n"
-        "            def _draw():\n"
-        "                plt.scatter(xs[m], ys[m])\n"
-        "                ax = plt.gca()\n"
-        "                if not (ax.get_title() or '').strip():\n"
-        "                    ax.set_title(f'{y} vs {x}')\n"
-        "                return ax\n"
-        "            return _safe_plot(lambda **kw: _draw())\n"
-        "        # else: fall back to auto-pick two numeric columns\n"
-        "        d = _pick_df()\n"
-        "        if d is not None:\n"
-        "            num = d.select_dtypes(include=[np.number]).columns.tolist()\n"
-        "            if len(num) >= 2:\n"
-        "                def _draw2():\n"
-        "                    plt.scatter(d[num[0]], d[num[1]])\n"
-        "                    ax = plt.gca()\n"
-        "                    if not (ax.get_title() or '').strip():\n"
-        "                        ax.set_title(f'{num[1]} vs {num[0]}')\n"
-        "                    return ax\n"
-        "                return _safe_plot(lambda **kw: _draw2())\n"
-        "        return _safe_plot(lambda **kw: plt.scatter([], []))\n"
-        "    # seaborn path\n"
-        "    return _safe_plot(fn, *a, **k)\n"
-        "\n"
-        "def SB_heatmap(*a, **k):\n"
-        "    _missing = (getattr(sns, '__class__', type(sns)).__name__ == '_Dummy')\n"
-        "    data = None\n"
-        "    if a:\n"
-        "        data = a[0]\n"
-        "    elif 'data' in k:\n"
-        "        data = k['data']\n"
-        "    if data is None:\n"
-        "        d = _pick_df()\n"
-        "        try:\n"
-        "            if d is not None:\n"
-        "                import numpy as _np\n"
-        "                data = d.select_dtypes(include=[_np.number]).corr()\n"
-        "        except Exception:\n"
-        "            data = None\n"
-        "    if data is None:\n"
-        "        from syntaxmatrix.display import show as _show\n"
-        "        _show('⚠ Heatmap skipped: no data.')\n"
-        "        return None\n"
-        "    if not _missing and hasattr(sns, 'heatmap'):\n"
-        "        _k = {kk: vv for kk, vv in k.items() if kk != 'data'}\n"
-        "        def _draw():\n"
-        "            ax = sns.heatmap(data, **_k)\n"
-        "            try:\n"
-        "                ax = ax or plt.gca()\n"
-        "                if not (ax.get_title() or '').strip():\n"
-        "                    ax.set_title('Correlation Heatmap')\n"
-        "            except Exception:\n"
-        "                pass\n"
-        "            return ax\n"
-        "        return _safe_plot(lambda **kw: _draw())\n"
-        "    def _mat_heat():\n"
-        "        im = plt.imshow(data, aspect='auto')\n"
-        "        try: plt.colorbar()\n"
-        "        except Exception: pass\n"
-        "        try:\n"
-        "            cols = list(getattr(data, 'columns', []))\n"
-        "            rows = list(getattr(data, 'index', []))\n"
-        "            if cols: plt.xticks(range(len(cols)), cols, rotation=90)\n"
-        "            if rows: plt.yticks(range(len(rows)), rows)\n"
-        "        except Exception:\n"
-        "            pass\n"
-        "        ax = plt.gca()\n"
-        "        try:\n"
-        "            if not (ax.get_title() or '').strip():\n"
-        "                ax.set_title('Correlation Heatmap')\n"
-        "        except Exception:\n"
-        "            pass\n"
-        "        return ax\n"
-        "    return _safe_plot(lambda **kw: _mat_heat())\n"
-        "\n"
-        "def _safe_concat(objs, **kwargs):\n"
-        "    import pandas as _pd\n"
-        "    if objs is None: return _pd.DataFrame()\n"
-        "    if isinstance(objs,(list,tuple)) and len(objs)==0: return _pd.DataFrame()\n"
-        "    try: return _pd.concat(objs, **kwargs)\n"
-        "    except Exception as e:\n"
-        "        show(f'⚠ concat skipped: {e}')\n"
-        "        return _pd.DataFrame()\n"  
-        "\n"
-        "from sklearn.preprocessing import OneHotEncoder\n"
-        "import inspect\n"
-        "def _SMX_OHE(**k):\n"
-        "    # normalise arg name across sklearn versions\n"
-        "    if 'sparse' in k and 'sparse_output' not in k:\n"
-        "        k['sparse_output'] = k.pop('sparse')\n"
-        "    k.setdefault('handle_unknown','ignore')\n"
-        "    k.setdefault('sparse_output', False)\n"
-        "    try:\n"
-        "        sig = inspect.signature(OneHotEncoder)\n"
-        "        if 'sparse_output' not in sig.parameters and 'sparse_output' in k:\n"
-        "            k['sparse'] = k.pop('sparse_output')\n"
-        "    except Exception:\n"
-        "        if 'sparse_output' in k:\n"
-        "            k['sparse'] = k.pop('sparse_output')\n"
-        "    return OneHotEncoder(**k)\n"
-        "\n"
-        "import numpy as _np\n"
-        "def _SMX_mm(a, b):\n"
-        "    try:\n"
-        "        return a @ b  # normal path\n"
-        "    except Exception:\n"
-        "        try:\n"
-        "            A = _np.asarray(a); B = _np.asarray(b)\n"
-        "            # If same 2D shape (e.g. (n,k) & (n,k)), treat as row-wise dot\n"
-        "            if A.ndim==2 and B.ndim==2 and A.shape==B.shape:\n"
-        "                return (A * B).sum(axis=1)\n"
-        "            # Otherwise try element-wise product (broadcast if possible)\n"
-        "            return A * B\n"
-        "        except Exception as e:\n"
-        "            from syntaxmatrix.display import show\n"
-        "            show(f'⚠ Matmul relaxed: {type(e).__name__}: {e}'); return _np.nan\n"
-        "\n"
-        "def _SMX_call(fn, *a, **k):\n"
-        "    try:\n"
-        "        return fn(*a, **k)\n"
-        "    except TypeError as e:\n"
-        "        msg = str(e)\n"
-        "        if \"unexpected keyword argument 'squared'\" in msg:\n"
-        "            k.pop('squared', None)\n"
-        "            return fn(*a, **k)\n"
-        "        raise\n"
-        "\n"
-        "def _SMX_rmse(y_true, y_pred):\n"
-        "    try:\n"
-        "        from sklearn.metrics import mean_squared_error as _mse\n"
-        "        try:\n"
-        "            return _mse(y_true, y_pred, squared=False)\n"
-        "        except TypeError:\n"
-        "            return (_mse(y_true, y_pred)) ** 0.5\n"
-        "    except Exception:\n"
-        "        import numpy as _np\n"
-        "        yt = _np.asarray(y_true, dtype=float)\n"
-        "        yp = _np.asarray(y_pred, dtype=float)\n"
-        "        diff = yt - yp\n"
-        "        return float((_np.mean(diff * diff)) ** 0.5)\n"
-        "\n"
-        "import pandas as _pd\n"
-        "import numpy as _np\n"
-        "def _SMX_autocoerce_dates(_df):\n"
-        "    if _df is None or not hasattr(_df, 'columns'): return\n"
-        "    for c in list(_df.columns):\n"
-        "        s = _df[c]\n"
-        "        n = str(c).lower()\n"
-        "        if _pd.api.types.is_datetime64_any_dtype(s):\n"
-        "            continue\n"
-        "        if _pd.api.types.is_object_dtype(s) or ('date' in n or 'time' in n or 'timestamp' in n or n.endswith('_dt')):\n"
-        "            try:\n"
-        "                conv = _pd.to_datetime(s, errors='coerce', utc=True).dt.tz_localize(None)\n"
-        "                # accept only if at least 10% (min 3) parse as dates\n"
-        "                if getattr(conv, 'notna', lambda: _pd.Series([]))().sum() >= max(3, int(0.1*len(_df))):\n"
-        "                    _df[c] = conv\n"
-        "            except Exception:\n"
-        "                pass\n"
-        "\n"
-        "def _SMX_autocoerce_numeric(_df, cols):\n"
-        "    if _df is None: return\n"
-        "    for c in cols:\n"
-        "        if c in getattr(_df, 'columns', []):\n"
-        "            try:\n"
-        "                _df[c] = _pd.to_numeric(_df[c], errors='coerce')\n"
-        "            except Exception:\n"
-        "                pass\n"
-        "\n"
-        "def show(obj, title=None):\n"
-        "    try:\n"
-        "        import pandas as pd, numbers\n"
-        "        cap = (title or _SMX_caption_from_ctx())\n"
-        "        # 1) DataFrame → Styler with caption\n"
-        "        if isinstance(obj, pd.DataFrame):\n"
-        "            try: return _SMX_base_show(obj.style.set_caption(cap))\n"
-        "            except Exception: pass\n"
-        "        # 2) dict of scalars → DataFrame with caption\n"
-        "        if isinstance(obj, dict) and all(isinstance(v, numbers.Number) for v in obj.values()):\n"
-        "            df_ = pd.DataFrame({'metric': list(obj.keys()), 'value': list(obj.values())})\n"
-        "            try: return _SMX_base_show(df_.style.set_caption(cap))\n"
-        "            except Exception: return _SMX_base_show(df_)\n"
-        "    except Exception:\n"
-        "        pass\n"
-        "    return _SMX_base_show(obj)\n"
-    )
-
-    PREFACE_IMPORT = "from syntaxmatrix.smx_preface import *\n"
-    # if PREFACE not in code:
-    #     code = PREFACE_IMPORT + code
-
     fixed = code
 
     fixed = re.sub(
@@ -690,6 +557,24 @@ def harden_ai_code(code: str) -> str:
         fixed
     )
 
+    try:
+          ast.parse(fixed)
+    except (SyntaxError, IndentationError):
+        fixed = _fallback_snippet()
+        
+    fixed = re.sub(
+        r"except\s+Exception\s+as\s+e:\s*\n\s*show\(\.\.\.\)",
+        "except Exception as e:\n    show(f\"⚠ Block skipped due to: {type(e).__name__}: {e}\")",
+        fixed,
+    )
+
+    # Fix placeholder Ellipsis handlers from LLM
+    fixed = re.sub(
+        r"except\s+Exception\s+as\s+e:\s*\n\s*show\(\.\.\.\)",
+        "except Exception as e:\n    show(f\"⚠ Block skipped due to: {type(e).__name__}: {e}\")",
+        fixed,
+    )
+
     try:
         class _SMXMatmulRewriter(ast.NodeTransformer):
             def visit_BinOp(self, node):
@@ -708,44 +593,15 @@ def harden_ai_code(code: str) -> str:
     # 6) Final safety wrapper
     fixed = fixed.replace("\t", "    ")
     fixed = textwrap.dedent(fixed).strip("\n")
-
+    fixed = _ensure_metrics_imports(fixed)
     fixed = _strip_stray_backrefs(fixed)
     fixed = _wrap_metric_calls(fixed)
+    fixed = _fix_unexpected_indent(fixed)
+    fixed = _patch_feature_coef_dataframe(fixed)
 
-    # If the transformed code is still not syntactically valid, fall back to a
-    # very defensive generic snippet that depends only on `df`. This guarantees
-    try:
-        ast.parse(fixed)
-    except (SyntaxError, IndentationError):
-        fixed = (
-            "import pandas as pd\n"
-            "df = df.copy()\n"
-            "_info = {\n"
-            "    'rows': len(df),\n"
-            "    'cols': len(df.columns),\n"
-            "    'numeric_cols': len(df.select_dtypes(include=['number','bool']).columns),\n"
-            "    'categorical_cols': len(df.select_dtypes(exclude=['number','bool']).columns),\n"
-            "}\n"
-            "show(df.head(), title='Sample of data')\n"
-            "show(_info, title='Dataset summary')\n"
-            "try:\n"
-            "    _num = df.select_dtypes(include=['number','bool']).columns.tolist()\n"
-            "    if _num:\n"
-            "        SB_histplot()\n"
-            "        _SMX_export_png()\n"
-            "except Exception as e:\n"
-            "    show(f\"⚠ Fallback visualisation failed: {type(e).__name__}: {e}\")\n"
-        )
-
-    # Fix placeholder Ellipsis handlers from LLM
-    fixed = re.sub(
-        r"except\s+Exception\s+as\s+e:\s*\n\s*show\(\.\.\.\)",
-        "except Exception as e:\n    show(f\"⚠ Block skipped due to: {type(e).__name__}: {e}\")",
-        fixed,
-    )
-
-    wrapped = PREFACE + "try:\n" + _indent(fixed) + "\nexcept Exception as e:\n    show(...)\n"
-    wrapped = wrapped.lstrip()
+    # Import shared preface helpers once and wrap the LLM body safely
+    header = "from syntaxmatrix.preface import *\n\n"
+    wrapped = header + wrap_llm_code_safe(fixed)
     return wrapped
 
 
@@ -754,17 +610,6 @@ def indent_code(code: str, spaces: int = 4) -> str:
     return "\n".join(pad + line for line in code.splitlines())
 
 
-def wrap_llm_code_safe(code: str) -> str:
-    # Swallow any runtime error from the LLM block instead of crashing the run
-    return (
-        "# __SAFE_WRAPPED__\n"
-        "try:\n" + indent_code(code) + "\n"
-        "except Exception as e:\n"
-        "    from syntaxmatrix.display import show\n"
-        "    show(f\"⚠️ Skipped LLM block due to: {type(e).__name__}: {e}\")\n"
-    )
-
-
 def fix_boxplot_placeholder(code: str) -> str:
     # Replace invalid 'sns.boxplot(boxplot)' with a safe call using df/group_label/m
     return re.sub(