added decision tree visualization

app.py

@@ -6,6 +6,45 @@ import gradio as gr
 import time
 import smtplib
 from email.message import EmailMessage
+from visualise import render_from_json
+from pathlib import Path
+import time
+import os, json, time, tempfile
+from huggingface_hub import HfApi, HfFileSystem, create_repo
+
+REPO = "CausalNLP/cais-demo-cache"           # dataset repo id
+TOKEN = os.environ["HF_WRITE_TOKEN"]   # set as Space secret
+api = HfApi(token=TOKEN)
+fs = HfFileSystem(token=TOKEN)
+
+# 1) ensure repo exists
+create_repo(REPO, repo_type="dataset", private=True, exist_ok=True, token=TOKEN)
+
+def cache_run(query, payload, artifacts=None):
+    ts = time.strftime("%Y-%m-%dT%H:%M:%S")
+    row = {"timestamp": ts, "query": query, "payload": payload, "artifacts": artifacts or {}}
+
+    hub_path = f"datasets/{REPO}/logs.jsonl"
+    # 2) download existing (if any), append, and push in one commit
+    with tempfile.TemporaryDirectory() as td:
+        local = os.path.join(td, "logs.jsonl")
+        try:
+            with fs.open(hub_path, "rb") as fsrc, open(local, "wb") as fdst:
+                fdst.write(fsrc.read())
+        except FileNotFoundError:
+            open(local, "w").close()
+
+        with open(local, "a", encoding="utf-8") as f:
+            f.write(json.dumps(row, ensure_ascii=False) + "\n")
+
+        api.upload_file(
+            path_or_fileobj=local,
+            path_in_repo="logs.jsonl",
+            repo_id=REPO,
+            repo_type="dataset",
+            commit_message=f"append log {ts}"
+        )
+
 
 # Make your repo importable (expecting a folder named causal-agent at repo root)
 sys.path.append(str(Path(__file__).parent / "causal-agent"))
@@ -120,18 +159,33 @@ def _ok_html(text):
     return f"<div style='padding:10px;border:1px solid #2ea043;border-radius:5px;color:#2ea043;background-color:#333333;'>✅ {text}</div>"
 
 # --- Email support ---
-def send_email(recipient: str, subject: str, body_text: str, attachment_name: str = None, attachment_json: dict = None) -> str:
-    """Returns '' on success, or error message string."""
-    host = os.getenv("SMTP_HOST")
-    port = int(os.getenv("SMTP_PORT", "587"))
-    user = os.getenv("SMTP_USER")
-    pwd  = os.getenv("SMTP_PASS")
-    from_addr = os.getenv("EMAIL_FROM")
+import base64, json, requests
+from email.message import EmailMessage
 
-    if not all([host, port, user, pwd, from_addr]):
-        return "Email is not configured (set SMTP_HOST, SMTP_PORT, SMTP_USER, SMTP_PASS, EMAIL_FROM)."
+def _gmail_access_token() -> str:
+    token_url = "https://oauth2.googleapis.com/token"
+    data = {
+        "client_id": os.getenv("GMAIL_CLIENT_ID"),
+        "client_secret": os.getenv("GMAIL_CLIENT_SECRET"),
+        "refresh_token": os.getenv("GMAIL_REFRESH_TOKEN"),
+        "grant_type": "refresh_token",
+    }
+    r = requests.post(token_url, data=data, timeout=20)
+    r.raise_for_status()
+    return r.json()["access_token"]
+
+def send_email(recipient: str, subject: str, body_text: str,
+               attachment_name: str = None, attachment_json: dict = None) -> str:
+    """
+    Sends via Gmail API. Returns '' on success, or an error string.
+    """
+    from_addr = os.getenv("EMAIL_FROM")
+    if not all([os.getenv("GMAIL_CLIENT_ID"), os.getenv("GMAIL_CLIENT_SECRET"),
+                os.getenv("GMAIL_REFRESH_TOKEN"), from_addr]):
+        return "Gmail API not configured (set GMAIL_CLIENT_ID, GMAIL_CLIENT_SECRET, GMAIL_REFRESH_TOKEN, EMAIL_FROM)."
 
     try:
+        # Build MIME message
         msg = EmailMessage()
         msg["From"] = from_addr
         msg["To"] = recipient
@@ -142,10 +196,17 @@ def send_email(recipient: str, subject: str, body_text: str, attachment_name: st
             payload = json.dumps(attachment_json, indent=2).encode("utf-8")
             msg.add_attachment(payload, maintype="application", subtype="json", filename=attachment_name)
 
-        with smtplib.SMTP(host, port, timeout=30) as s:
-            s.starttls()
-            s.login(user, pwd)
-            s.send_message(msg)
+        # Base64url encode the raw RFC822 message
+        raw = base64.urlsafe_b64encode(msg.as_bytes()).decode("utf-8")
+
+        # Get access token and send
+        access_token = _gmail_access_token()
+        api_url = "https://gmail.googleapis.com/gmail/v1/users/me/messages/send"
+        headers = {"Authorization": f"Bearer {access_token}", "Content-Type": "application/json"}
+        r = requests.post(api_url, headers=headers, json={"raw": raw}, timeout=20)
+
+        if r.status_code >= 400:
+            return f"Gmail API error {r.status_code}: {r.text[:300]}"
         return ""
     except Exception as e:
         return f"Email send failed: {e}"
@@ -154,35 +215,43 @@ def run_agent(query: str, csv_path: str, dataset_description: str, email: str):
     start = time.time()
 
     processing_html = _html_panel("🔄 Analysis in Progress...", "<div style='font-size:14px;color:#bbb;'>This may take 1–2 minutes depending on dataset size.</div>")
-    yield (processing_html, processing_html, processing_html, {"status": "Processing started..."})
+    yield (processing_html, processing_html, processing_html, {"status": "Processing started..."},  None, None)
 
     if not os.getenv("OPENAI_API_KEY"):
-        yield (_err_html("Set a Space Secret named OPENAI_API_KEY"), "", "", {})
+        yield (_err_html("Set a Space Secret named OPENAI_API_KEY"), "", "", {}, None, None)
         return
     if not csv_path:
-        yield (_warn_html("Please upload a CSV dataset."), "", "", {})
+        yield (_warn_html("Please upload a CSV dataset."), "", "", {}, None, None)
         return
 
     try:
         step_html = _html_panel("📊 Running Causal Analysis...", "<div style='font-size:14px;color:#bbb;'>Analyzing dataset and selecting optimal method…</div>")
-        yield (step_html, step_html, step_html, {"status": "Running causal analysis..."})
+        yield (step_html, step_html, step_html, {"status": "Running causal analysis..."},  None, None)
 
         result = run_causal_analysis(
             query=(query or "What is the effect of treatment T on outcome Y controlling for X?").strip(),
             dataset_path=csv_path,
             dataset_description=(dataset_description or "").strip(),
         )
-
+        cache_run(query, result)
         llm_html = _html_panel("🤖 Generating Summary...", "<div style='font-size:14px;color:#bbb;'>Creating human-readable interpretation…</div>")
-        yield (llm_html, llm_html, llm_html, {"status": "Generating explanation...", "raw_analysis": result if isinstance(result, dict) else {}})
+        yield (llm_html, llm_html, llm_html, {"status": "Generating explanation...", "raw_analysis": result if isinstance(result, dict) else {}},  None, None)
 
     except Exception as e:
-        yield (_err_html(str(e)), "", "", {})
+        yield (_err_html(str(e)), "", "", {}, None, None)
         return
 
     try:
         payload = _extract_minimal_payload(result if isinstance(result, dict) else {})
+        
         method = payload.get("method_used", "N/A")
+        # --- Decision tree render ---
+        artifacts_dir = Path("artifacts")
+        artifacts_dir.mkdir(exist_ok=True)
+        ts = time.strftime("%Y%m%d-%H%M%S")
+        out_stem = str(artifacts_dir / f"decision_tree_{ts}")
+
+        # This creates: out_stem.dot, out_stem.svg, out_stem.png
 
         method_html = _html_panel("Selected Method", f"<p style='margin:0;font-size:16px;'>{method}</p>")
 
@@ -199,7 +268,7 @@ def run_agent(query: str, csv_path: str, dataset_description: str, email: str):
             explanation_html = _warn_html(f"LLM summary failed: {e}")
 
     except Exception as e:
-        yield (_err_html(f"Failed to parse results: {e}"), "", "", {})
+        yield (_err_html(f"Failed to parse results: {e}"), "", "", {}, "", None)
         return
 
     # Optional email send (best-effort)
@@ -225,8 +294,12 @@ def run_agent(query: str, csv_path: str, dataset_description: str, email: str):
             explanation_html += _warn_html(email_err)
         else:
             explanation_html += _ok_html(f"Results emailed to {email.strip()}")
+    render_from_json(result, out_stem)
 
-    yield (method_html, effects_html, explanation_html, result if isinstance(result, dict) else {})
+    tree_png = f"{out_stem}.png"
+    tree_svg = f"{out_stem}.svg"
+    tree_dot = f"{out_stem}.dot"
+    yield (method_html, effects_html, explanation_html, result if isinstance(result, dict) else {}, tree_png, [tree_svg, tree_dot, tree_png])
 
 with gr.Blocks() as demo:
     gr.Markdown("# Causal AI Scientist")
@@ -310,16 +383,23 @@ with gr.Blocks() as demo:
     with gr.Row():
         explanation_out = gr.HTML(label="Detailed Explanation")
 
+    with gr.Row():
+        tree_img = gr.Image(label="Decision Tree", type="filepath")
+    with gr.Row():
+        tree_files = gr.Files(label="Download decision tree artifacts (.svg / .dot / .png)")
+
     with gr.Accordion("Raw Results (Advanced)", open=False):
         raw_results = gr.JSON(label="Complete Analysis Output", show_label=False)
 
     run_btn.click(
         fn=run_agent,
         inputs=[query, csv_file, dataset_description, email],
-        outputs=[method_out, effects_out, explanation_out, raw_results],
+        outputs=[method_out, effects_out, explanation_out, raw_results, tree_img, tree_files],
         show_progress=True
     )
 
+
+
     
 
 if __name__ == "__main__":

visualise.py

@@ -0,0 +1,328 @@
+# -*- coding: utf-8 -*-
+"""
+Render a JSON-aware visualization of CAIS's rule-based method selector.
+- Parses a CAIS run payload (dict) and highlights ALL plausible candidates (green).
+- The actually selected method receives a thicker border.
+- The traversed decision path edges are colored.
+
+Usage:
+    render_from_json(payload_dict, out_stem="artifacts/decision_tree")
+
+(Optional) CLI:
+    python decision_tree.py payload.json
+"""
+
+from graphviz import Digraph
+import json, sys
+from typing import Dict, Any, List, Set, Tuple, Optional
+
+from auto_causal.components.decision_tree import (
+    DIFF_IN_MEANS, LINEAR_REGRESSION, DIFF_IN_DIFF, REGRESSION_DISCONTINUITY,
+    INSTRUMENTAL_VARIABLE, PROPENSITY_SCORE_MATCHING, PROPENSITY_SCORE_WEIGHTING,
+    GENERALIZED_PROPENSITY_SCORE, BACKDOOR_ADJUSTMENT, FRONTDOOR_ADJUSTMENT
+)
+
+LABEL = {
+    DIFF_IN_MEANS: "Diff-in-Means (RCT)",
+    LINEAR_REGRESSION: "Linear Regression",
+    DIFF_IN_DIFF: "Difference-in-Differences",
+    REGRESSION_DISCONTINUITY: "Regression Discontinuity",
+    INSTRUMENTAL_VARIABLE: "Instrumental Variables",
+    PROPENSITY_SCORE_MATCHING: "PS Matching",
+    PROPENSITY_SCORE_WEIGHTING: "PS Weighting",
+    GENERALIZED_PROPENSITY_SCORE: "Generalized PS (continuous T)",
+    BACKDOOR_ADJUSTMENT: "Backdoor Adjustment",
+    FRONTDOOR_ADJUSTMENT: "Frontdoor Adjustment",
+}
+
+# -------- Heuristic extractors from payload -------- #
+
+def _get(d: Dict, path: List[str], default=None):
+    cur = d
+    for k in path:
+        if not isinstance(cur, dict) or k not in cur:
+            return default
+        cur = cur[k]
+    return cur
+
+def extract_signals(p: Dict[str, Any]) -> Dict[str, Any]:
+    vars_ = _get(p, ["results", "variables"], {}) or _get(p, ["variables"], {}) or {}
+    da   = _get(p, ["results", "dataset_analysis"], {}) or _get(p, ["dataset_analysis"], {}) or {}
+
+    treatment = vars_.get("treatment_variable")
+    t_type    = vars_.get("treatment_variable_type")            # "binary"/"continuous"
+    is_rct    = bool(vars_.get("is_rct", False))
+
+    # Temporal / panel
+    temporal_detected = bool(da.get("temporal_structure_detected", False))
+    time_var = vars_.get("time_variable")
+    group_var = vars_.get("group_variable")
+    has_temporal = temporal_detected or bool(time_var) or bool(group_var)
+
+    # RDD
+    running_variable = vars_.get("running_variable")
+    cutoff_value     = vars_.get("cutoff_value")
+    rdd_ready        = running_variable is not None and cutoff_value is not None
+    # (Some detectors raise 'discontinuities_detected', but we still require running var + cutoff.)
+    # If you want permissive behavior, flip rdd_ready to also consider da.get("discontinuities_detected").
+
+    # Instruments
+    instrument = vars_.get("instrument_variable")
+    pot_instr  = da.get("potential_instruments") or []
+    # Consider an instrument valid only if it exists and is NOT the treatment itself
+    has_valid_instrument = (
+        instrument is not None and instrument != treatment
+    ) or any(pi and pi != treatment for pi in pot_instr)
+
+    covariates = vars_.get("covariates") or []
+    has_covariates = len(covariates) > 0
+
+    # Frontdoor: only mark if explicitly provided (else too speculative)
+    frontdoor_ok = bool(_get(p, ["results", "dataset_analysis", "frontdoor_satisfied"], False))
+
+    # Overlap: if explicitly known, use it; else unknown → both PS variants remain plausible.
+    overlap_assessment = da.get("overlap_assessment")
+    strong_overlap = None
+    if isinstance(overlap_assessment, dict):
+        # accept typical keys like {"strong_overlap": true}
+        strong_overlap = overlap_assessment.get("strong_overlap")
+
+    return dict(
+        treatment=treatment,
+        t_type=t_type,
+        is_rct=is_rct,
+        has_temporal=has_temporal,
+        rdd_ready=rdd_ready,
+        has_valid_instrument=has_valid_instrument,
+        has_covariates=has_covariates,
+        frontdoor_ok=frontdoor_ok,
+        strong_overlap=strong_overlap,
+    )
+
+# -------- Candidate inference (green leaves) -------- #
+
+def infer_candidate_methods(signals: Dict[str, Any]) -> Set[str]:
+    cands: Set[str] = set()
+    is_rct = signals["is_rct"]
+
+    # RCT branch: both Diff-in-Means and LR are valid analyses; IV only if a valid instrument exists (e.g., randomized encouragement)
+    if is_rct:
+        cands.add(DIFF_IN_MEANS)
+        if signals["has_covariates"]:
+            cands.add(LINEAR_REGRESSION)
+        if signals["has_valid_instrument"]:
+            cands.add(INSTRUMENTAL_VARIABLE)
+        return cands  # stop here; the observational tree is not needed
+
+    # Observational branch
+    if signals["has_temporal"]:
+        cands.add(DIFF_IN_DIFF)
+    if signals["rdd_ready"]:
+        cands.add(REGRESSION_DISCONTINUITY)
+    if signals["has_valid_instrument"]:
+        cands.add(INSTRUMENTAL_VARIABLE)
+    if signals["frontdoor_ok"]:
+        cands.add(FRONTDOOR_ADJUSTMENT)
+
+    # Treatment type
+    if str(signals["t_type"]).lower() == "continuous":
+        cands.add(GENERALIZED_PROPENSITY_SCORE)
+
+    # Backdoor / PS (need covariates)
+    if signals["has_covariates"]:
+        # If overlap is known, choose one; if unknown, mark both as plausible.
+        if signals["strong_overlap"] is True:
+            cands.add(PROPENSITY_SCORE_MATCHING)
+        elif signals["strong_overlap"] is False:
+            cands.add(PROPENSITY_SCORE_WEIGHTING)
+        else:
+            cands.add(PROPENSITY_SCORE_MATCHING)
+            cands.add(PROPENSITY_SCORE_WEIGHTING)
+        cands.add(BACKDOOR_ADJUSTMENT)
+
+    return cands
+
+# -------- Compute the single realized path to the chosen leaf (for edge coloring) -------- #
+
+def infer_decision_path(signals: Dict[str, Any], selected_method: Optional[str]) -> List[Tuple[str, str]]:
+    path: List[Tuple[str, str]] = []
+    # Start → is_rct
+    path.append(("start", "is_rct"))
+
+    if signals["is_rct"]:
+        path.append(("is_rct", "has_instr_rct"))
+        if signals["has_valid_instrument"]:
+            path.append(("has_instr_rct", INSTRUMENTAL_VARIABLE))
+        else:
+            path.append(("has_instr_rct", "has_cov_rct"))
+            if signals["has_covariates"]:
+                path.append(("has_cov_rct", LINEAR_REGRESSION))
+            else:
+                path.append(("has_cov_rct", DIFF_IN_MEANS))
+        return path
+
+    # Observational
+    path.append(("is_rct", "has_temporal"))
+    if signals["has_temporal"]:
+        path.append(("has_temporal", DIFF_IN_DIFF))
+        return path
+    else:
+        path.append(("has_temporal", "has_rv"))
+
+    if signals["rdd_ready"]:
+        path.append(("has_rv", REGRESSION_DISCONTINUITY))
+        return path
+    else:
+        path.append(("has_rv", "has_instr"))
+
+    if signals["has_valid_instrument"]:
+        path.append(("has_instr", INSTRUMENTAL_VARIABLE))
+        return path
+    else:
+        path.append(("has_instr", "frontdoor"))
+
+    if signals["frontdoor_ok"]:
+        path.append(("frontdoor", FRONTDOOR_ADJUSTMENT))
+        return path
+    else:
+        path.append(("frontdoor", "t_cont"))
+
+    if str(signals["t_type"]).lower() == "continuous":
+        path.append(("t_cont", GENERALIZED_PROPENSITY_SCORE))
+        return path
+    else:
+        path.append(("t_cont", "has_cov"))
+
+    if signals["has_covariates"]:
+        path.append(("has_cov", "overlap"))
+        # If overlap known, pick the branch; else default to weighting.
+        if signals["strong_overlap"] is True:
+            path.append(("overlap", PROPENSITY_SCORE_MATCHING))
+        else:
+            path.append(("overlap", PROPENSITY_SCORE_WEIGHTING))
+    else:
+        path.append(("has_cov", BACKDOOR_ADJUSTMENT))  # keep original topology; see note in previous message
+    return path
+
+# -------- Graph building -------- #
+
+def build_graph(payload: Dict[str, Any]) -> Digraph:
+    g = Digraph("CAISDecisionTree", format="svg")
+    g.attr(rankdir="LR", nodesep="0.4", ranksep="0.35", fontsize="11")
+
+    # Decisions
+    g.node("start", "Start", shape="circle")
+    g.node("is_rct", "Is RCT?", shape="diamond")
+    g.node("has_instr_rct", "Instrument available?", shape="diamond")
+    g.node("has_cov_rct", "Covariates observed?", shape="diamond")
+    g.node("has_temporal", "Temporal structure?", shape="diamond")
+    g.node("has_rv", "Running var & cutoff?", shape="diamond")
+    g.node("has_instr", "Instrument available?", shape="diamond")
+    g.node("frontdoor", "Frontdoor criterion satisfied?", shape="diamond")
+    g.node("has_cov", "Covariates observed?", shape="diamond")
+    g.node("overlap", "Strong overlap?\n(overlap ≥ 0.1)", shape="diamond")
+    g.node("t_cont", "Treatment continuous?", shape="diamond")
+
+    # Leaves
+    def leaf(name_const, fill=None, bold=False):
+        attrs = {"shape": "box", "style": "rounded"}
+        if fill:
+            attrs.update(style="rounded,filled", fillcolor=fill)
+        if bold:
+            attrs.update(penwidth="2")
+        g.node(name_const, LABEL[name_const], **attrs)
+
+    # Compute signals, candidates, path
+    signals = extract_signals(payload)
+    candidates = infer_candidate_methods(signals)
+
+    selected_method_str = _get(payload, ["results", "results", "method_used"]) \
+                          or _get(payload, ["results", "method_used"]) \
+                          or _get(payload, ["method"])
+    selected_method = {
+        "linear_regression": LINEAR_REGRESSION,
+        "diff_in_means": DIFF_IN_MEANS,
+        "difference_in_differences": DIFF_IN_DIFF,
+        "regression_discontinuity": REGRESSION_DISCONTINUITY,
+        "instrumental_variable": INSTRUMENTAL_VARIABLE,
+        "propensity_score_matching": PROPENSITY_SCORE_MATCHING,
+        "propensity_score_weighting": PROPENSITY_SCORE_WEIGHTING,
+        "generalized_propensity_score": GENERALIZED_PROPENSITY_SCORE,
+        "backdoor_adjustment": BACKDOOR_ADJUSTMENT,
+        "frontdoor_adjustment": FRONTDOOR_ADJUSTMENT,
+    }.get(str(selected_method_str or "").lower())
+
+    # Add leaves with coloring
+    for m in [
+        DIFF_IN_MEANS, LINEAR_REGRESSION, DIFF_IN_DIFF, REGRESSION_DISCONTINUITY,
+        INSTRUMENTAL_VARIABLE, PROPENSITY_SCORE_MATCHING, PROPENSITY_SCORE_WEIGHTING,
+        GENERALIZED_PROPENSITY_SCORE, BACKDOOR_ADJUSTMENT, FRONTDOOR_ADJUSTMENT
+    ]:
+        leaf(m,
+             fill=("palegreen" if m in candidates else None),
+             bold=(m == selected_method))
+
+    # Edges with optional path highlighting
+    path_edges = set(infer_decision_path(signals, selected_method))
+    def e(u, v, label=None):
+        attrs = {}
+        if (u, v) in path_edges:
+            attrs.update(color="forestgreen", penwidth="2")
+        g.edge(u, v, **({} if label is None else {"label": label}) | attrs)
+
+    # Topology (unchanged)
+    e("start", "is_rct")
+
+    # RCT branch
+    e("is_rct", "has_instr_rct", label="Yes")
+    e("has_instr_rct", INSTRUMENTAL_VARIABLE, label="Yes")
+    e("has_instr_rct", "has_cov_rct", label="No")
+    e("has_cov_rct", LINEAR_REGRESSION, label="Yes")
+    e("has_cov_rct", DIFF_IN_MEANS, label="No")
+
+    # Observational branch
+    e("is_rct", "has_temporal", label="No")
+    e("has_temporal", DIFF_IN_DIFF, label="Yes")
+    e("has_temporal", "has_rv", label="No")
+
+    e("has_rv", REGRESSION_DISCONTINUITY, label="Yes")
+    e("has_rv", "has_instr", label="No")
+
+    e("has_instr", INSTRUMENTAL_VARIABLE, label="Yes")
+    e("has_instr", "frontdoor", label="No")
+
+    e("frontdoor", FRONTDOOR_ADJUSTMENT, label="Yes")
+    e("frontdoor", "t_cont", label="No")
+
+    e("t_cont", GENERALIZED_PROPENSITY_SCORE, label="Yes")
+    e("t_cont", "has_cov", label="No")
+
+    e("has_cov", "overlap", label="Yes")
+    e("has_cov", BACKDOOR_ADJUSTMENT, label="No")
+
+    e("overlap", PROPENSITY_SCORE_MATCHING, label="Yes")
+    e("overlap", PROPENSITY_SCORE_WEIGHTING, label="No")
+
+    # Optional legend
+    g.node("legend", "Legend:\nGreen = plausible candidate(s)\nBold border = method used", shape="note")
+    g.edge("legend", "start", style="dashed", arrowhead="none")
+
+    return g
+
+def render_from_json(payload: Dict[str, Any], out_stem: str = "artifacts/decision_tree"):
+    g = build_graph(payload)
+    g.save(filename=f"{out_stem}.dot")
+    g.render(filename=out_stem, cleanup=True)         # SVG
+    g.format = "png"
+    g.render(filename=out_stem, cleanup=True)         # PNG
+
+def main():
+    # if len(sys.argv) >= 2:
+    with open('sample_output.json', "r") as f:
+        payload = json.load(f)
+    # else:
+    # payload = json.load()
+    render_from_json(payload)
+
+if __name__ == "__main__":
+    main()