18  Code: Agent 3 — Research Experiment

Folder: 02_agents_code/agent_03_research_experiment/ — self-contained; run python agent_03_research_experiment/main.py --help.

18.1 README

18.1.1 README.md

# Agent 3 — AI Research Experiment Agent ("Virtual Research Assistant")

**Code does the maths, the LLM does the reasoning.** The agent plans the experiment, runs a deterministic analysis pipeline (QC → normalisation → Welch t-test + BH-FDR → PCA/volcano → ML with cross-validation), asks a "statistician" agent to review the numbers, and writes the Methods & Results plus a reproducible notebook.

```
description (+ CSV or --demo)
   ─► ExperimentPlanner (LLM, structured)      → design.json
   ─► analysis.py (pandas/scipy/statsmodels/sklearn) → summary.json, differential.csv, pca.png, volcano.png
   ─► StatisticsAgent (LLM, structured)        → stats_review.json
   ─► ReportWriter (LLM)                        → methods_results.md
   ─► notebook.py (nbformat)                    → analysis_notebook.ipynb
```

## Files

| File | Purpose |
|------|---------|
| `main.py` | CLI (`--description`, `--data`, `--group-col`, `--demo`, `--no-llm`) |
| `data.py` | Load a CSV (samples as rows) or generate synthetic RNA-seq-like counts |
| `analysis.py` | QC, normalisation, differential test, figures, ML baseline, summary — **no LLM** |
| `agents.py` | Planner, Statistics reviewer, Report writer |
| `notebook.py` | Generates a runnable Jupyter notebook |
| `llm.py` | Thin Claude client |

## Run

```powershell
python agent_03_research_experiment\main.py --description "RNA-seq of heat-stressed rice, control vs heat, n=3" --demo
python agent_03_research_experiment\main.py --description "..." --data mydata.csv --group-col treatment
python agent_03_research_experiment\main.py --description "..." --demo --no-llm     # no API key needed
```

CSV layout: `sample_id, <group_col>, feature_1, feature_2, ...` — two groups for the differential test.

## What this agent teaches

1. **Separation of concerns** — LLM for planning/critique/writing; Python for numbers.
2. **Structured critique** — a reviewer agent that returns concerns with severity and recommendations.
3. **Reproducibility as a feature** — notebook + CSVs + PNGs alongside the prose.
4. **Privacy-friendly design** — only summary statistics are sent to the model, never raw data.

## Ideas to extend

- Domain plugins: DESeq2/pydeseq2 for RNA-seq, ANOVA/mixed models for field trials, survival analysis for clinical data.
- Sandboxed code execution so the agent can write custom analysis code.
- Power/sample-size calculator at design time.

---

**Codanics — Python ka Chilla: Build AI Agents**
Course: https://codanics.com/courses/python-ka-chilla-build-ai-agents/ · Website: https://codanics.com · YouTube: [@codanics](https://www.youtube.com/@codanics)
WhatsApp: https://whatsapp.com/channel/0029Va7nRDq3QxRzGqaQvS3r · Free course: https://www.youtube.com/watch?v=UE3o-7KRNXs&list=PLaQmdklCcX40

18.2 Source files

18.2.1 agents.py

"""
agents.py — the LLM-powered sub-agents of the Research Experiment Agent.

    ExperimentPlanner : description ─► ExperimentDesign (structured)
    StatisticsAgent   : analysis summary ─► StatsReview   (structured critique)
    ReportWriter      : design + numbers + critique ─► Methods & Results (Markdown)

The heavy lifting (QC, tests, ML, figures) lives in analysis.py and never touches the LLM.

Codanics — Python ka Chilla: Build AI Agents
https://codanics.com/courses/python-ka-chilla-build-ai-agents/
"""
from __future__ import annotations

import json

from pydantic import BaseModel, Field

from llm import ask, ask_structured

BRAND = ("Codanics — Python ka Chilla: Build AI Agents | "
         "https://codanics.com/courses/python-ka-chilla-build-ai-agents/")


# ------------------------------------------------------------------- 1. Planner
class ExperimentDesign(BaseModel):
    title: str
    hypothesis: str
    design_type: str = Field(description="e.g. two-group comparison, factorial, time-series, RCT")
    groups: list[str]
    replicates_per_group: int
    controls: list[str]
    covariates_or_blocking: list[str]
    sample_size_note: str = Field(description="Comment on power / adequacy of replicates")
    analysis_plan: list[str] = Field(description="Ordered steps: QC, normalisation, tests, ML, validation")
    statistical_tests: list[str]
    risks: list[str]


class ExperimentPlanner:
    SYSTEM = ("You are a biostatistician and experimental-design expert. Produce a concrete, minimal "
              "design that a lab can execute and a reviewer will accept.")

    def run(self, description: str, data_hint: str = "") -> ExperimentDesign:
        return ask_structured(self.SYSTEM,
                              f"Experiment description: {description}\n{data_hint}\nDesign the experiment and analysis plan.",
                              ExperimentDesign, effort="medium")


# ---------------------------------------------------------------- 2. Statistics
class Concern(BaseModel):
    issue: str
    severity: str = Field(description="low | medium | high")
    recommendation: str


class StatsReview(BaseModel):
    assumptions_checked: list[str]
    sample_size_adequate: bool
    multiple_testing_handled: bool
    leakage_or_bias_risk: str
    concerns: list[Concern]
    interpretation: str = Field(description="Plain-language interpretation of the key numbers")
    overall_verdict: str = Field(description="e.g. 'Preliminary evidence; needs biological validation'")


class StatisticsAgent:
    SYSTEM = ("You are a rigorous reviewing statistician. You receive ONLY computed numbers (never raw data). "
              "Check assumptions, sample size, multiple testing, leakage and over-interpretation. Be concrete.")

    def run(self, design: ExperimentDesign, summary: dict) -> StatsReview:
        return ask_structured(self.SYSTEM,
                              f"Design:\n{design.model_dump_json(indent=1)}\n\nAnalysis summary:\n"
                              f"{json.dumps(summary, indent=1)}\n\nReview it.",
                              StatsReview, effort="high")


# ---------------------------------------------------------------- 3. Report
class ReportWriter:
    SYSTEM = ("You write the Methods and Results sections of a scientific paper in Markdown, using ONLY the "
              "numbers provided (never invent values). Include: Methods (design, samples, normalisation, "
              "statistical tests, ML & validation, software), Results (with a compact table of top features "
              "and references to Figure 1 (PCA) and Figure 2 (volcano)), Limitations, and a short "
              "'Statistical review' box summarising the reviewer's concerns.")

    def run(self, description: str, design: ExperimentDesign, summary: dict, review: StatsReview,
            echo: bool = False) -> str:
        text = ask(self.SYSTEM,
                   f"Study: {description}\n\nDesign:\n{design.model_dump_json(indent=1)}\n\n"
                   f"Computed results:\n{json.dumps(summary, indent=1)}\n\nStatistical review:\n"
                   f"{review.model_dump_json(indent=1)}",
                   effort="high", echo=echo)
        return text + f"\n\n---\n*Generated with the AI Research Experiment Agent — {BRAND}*\n"

18.2.2 analysis.py

"""
analysis.py — deterministic analysis engine (QC → normalisation → differential test → ML).

Nothing in this file calls the LLM.  Numbers, p-values and figures come from pandas / scipy /
statsmodels / scikit-learn so that results are reproducible and defensible.

Codanics — Python ka Chilla: Build AI Agents
https://codanics.com/courses/python-ka-chilla-build-ai-agents/
"""
from __future__ import annotations

from pathlib import Path

import matplotlib
matplotlib.use("Agg")  # headless
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
import pandas as pd  # noqa: E402
from scipy import stats  # noqa: E402
from sklearn.decomposition import PCA  # noqa: E402
from sklearn.ensemble import RandomForestClassifier  # noqa: E402
from sklearn.linear_model import LogisticRegression  # noqa: E402
from sklearn.model_selection import LeaveOneOut, StratifiedKFold, cross_val_score  # noqa: E402
from sklearn.pipeline import make_pipeline  # noqa: E402
from sklearn.preprocessing import StandardScaler  # noqa: E402
from statsmodels.stats.multitest import multipletests  # noqa: E402

META_COLS = {"sample_id"}


def feature_matrix(df: pd.DataFrame, group_col: str) -> pd.DataFrame:
    cols = [c for c in df.columns if c not in META_COLS | {group_col}]
    return df[cols].apply(pd.to_numeric, errors="coerce")


# ------------------------------------------------------------------------------- QC
def quality_control(df: pd.DataFrame, group_col: str) -> dict:
    X = feature_matrix(df, group_col)
    missing = int(X.isna().sum().sum())
    zero_var = X.columns[X.var(axis=0).fillna(0) == 0].tolist()
    per_sample_total = X.sum(axis=1)
    return {
        "n_samples": int(len(df)),
        "n_features": int(X.shape[1]),
        "groups": df[group_col].value_counts().to_dict(),
        "missing_values": missing,
        "zero_variance_features": len(zero_var),
        "library_size_min": float(per_sample_total.min()),
        "library_size_max": float(per_sample_total.max()),
        "library_size_cv": float(per_sample_total.std() / per_sample_total.mean()) if per_sample_total.mean() else 0.0,
        "warnings": ([f"{missing} missing values (filled with column median)"] if missing else [])
                    + ([f"{len(zero_var)} zero-variance features removed"] if zero_var else [])
                    + (["fewer than 3 replicates in a group — statistics are unreliable"]
                       if df[group_col].value_counts().min() < 3 else []),
    }


def normalise(df: pd.DataFrame, group_col: str) -> pd.DataFrame:
    """Median-fill, drop zero-variance, then log2(CPM+1) if data look like counts, else z-score."""
    X = feature_matrix(df, group_col)
    X = X.fillna(X.median())
    X = X.loc[:, X.var(axis=0) > 0]
    looks_like_counts = (X.values >= 0).all() and (X.values.max() > 50)
    if looks_like_counts:
        cpm = X.div(X.sum(axis=1), axis=0) * 1e6
        X = np.log2(cpm + 1)
    else:
        X = (X - X.mean()) / X.std(ddof=0).replace(0, 1)
    out = pd.concat([df[["sample_id", group_col]].reset_index(drop=True), X.reset_index(drop=True)], axis=1)
    out.attrs["normalisation"] = "log2(CPM+1)" if looks_like_counts else "z-score"
    return out


# ------------------------------------------------------------------ differential test
def differential(dfn: pd.DataFrame, group_col: str, reference: str | None = None) -> pd.DataFrame:
    groups = list(dict.fromkeys(dfn[group_col]))
    if len(groups) != 2:
        raise ValueError(f"differential() needs exactly 2 groups, found {groups}")
    ref = reference or groups[0]
    trt = [g for g in groups if g != ref][0]
    X = feature_matrix(dfn, group_col)
    a, b = X[dfn[group_col] == ref], X[dfn[group_col] == trt]
    t, p = stats.ttest_ind(b, a, equal_var=False, nan_policy="omit")   # Welch t-test per feature
    p = np.nan_to_num(p, nan=1.0)
    _, fdr, _, _ = multipletests(p, method="fdr_bh")
    res = pd.DataFrame({
        "feature": X.columns,
        "mean_" + ref: a.mean().values,
        "mean_" + trt: b.mean().values,
        "log2FC": (b.mean() - a.mean()).values,          # already log scale after normalisation
        "t": t, "p_value": p, "fdr": fdr,
    }).sort_values("fdr")
    res.attrs.update({"reference": ref, "treatment": trt})
    return res


# --------------------------------------------------------------------------- figures
def plot_volcano(res: pd.DataFrame, out: Path, fdr_cut: float = 0.05, fc_cut: float = 1.0) -> Path:
    sig = (res["fdr"] < fdr_cut) & (res["log2FC"].abs() >= fc_cut)
    plt.figure(figsize=(6, 5))
    plt.scatter(res["log2FC"], -np.log10(res["p_value"].clip(1e-300)), s=10, c="#9aa5b1", label="n.s.")
    plt.scatter(res.loc[sig, "log2FC"], -np.log10(res.loc[sig, "p_value"].clip(1e-300)), s=12, c="#d1495b",
                label=f"FDR<{fdr_cut} & |log2FC|≥{fc_cut} (n={int(sig.sum())})")
    plt.axvline(-fc_cut, ls="--", lw=0.8, c="k"); plt.axvline(fc_cut, ls="--", lw=0.8, c="k")
    plt.xlabel("log2 fold change"); plt.ylabel("-log10 p-value"); plt.legend(fontsize=8)
    plt.title(f"Volcano: {res.attrs.get('treatment')} vs {res.attrs.get('reference')}")
    plt.tight_layout(); out.parent.mkdir(parents=True, exist_ok=True); plt.savefig(out, dpi=150); plt.close()
    return out


def plot_pca(dfn: pd.DataFrame, group_col: str, out: Path) -> Path:
    X = feature_matrix(dfn, group_col)
    pcs = PCA(n_components=2).fit(X)
    Z = pcs.transform(X)
    plt.figure(figsize=(5.5, 5))
    for g in dfn[group_col].unique():
        m = dfn[group_col].values == g
        plt.scatter(Z[m, 0], Z[m, 1], label=str(g), s=60)
    for i, sid in enumerate(dfn["sample_id"]):
        plt.annotate(str(sid), (Z[i, 0], Z[i, 1]), fontsize=7, alpha=0.7)
    plt.xlabel(f"PC1 ({pcs.explained_variance_ratio_[0]:.0%})"); plt.ylabel(f"PC2 ({pcs.explained_variance_ratio_[1]:.0%})")
    plt.legend(); plt.title("PCA of samples"); plt.tight_layout(); plt.savefig(out, dpi=150); plt.close()
    return out


# ------------------------------------------------------------------------------- ML
def machine_learning(dfn: pd.DataFrame, group_col: str, top_k: int = 50, seed: int = 0) -> dict:
    """Baseline classifiers with proper cross-validation (LOOCV when n is tiny)."""
    X = feature_matrix(dfn, group_col)
    X = X.loc[:, X.var(axis=0).sort_values(ascending=False).index[:top_k]]  # unsupervised filter → no leakage
    y = dfn[group_col].astype("category").cat.codes.values
    n_min = int(np.bincount(y).min())
    cv = LeaveOneOut() if n_min < 5 else StratifiedKFold(n_splits=min(5, n_min), shuffle=True, random_state=seed)
    cv_name = "leave-one-out" if isinstance(cv, LeaveOneOut) else f"stratified {cv.get_n_splits()}-fold"
    models = {
        "logistic_regression": make_pipeline(StandardScaler(), LogisticRegression(max_iter=2000, C=0.5)),
        "random_forest": RandomForestClassifier(n_estimators=300, random_state=seed),
    }
    results = {}
    for name, model in models.items():
        acc = cross_val_score(model, X.values, y, cv=cv, scoring="accuracy")
        results[name] = {"cv_accuracy_mean": float(acc.mean()), "cv_accuracy_std": float(acc.std())}
    rf = models["random_forest"].fit(X.values, y)
    imp = pd.Series(rf.feature_importances_, index=X.columns).sort_values(ascending=False)
    return {"cross_validation": cv_name, "n_features_used": int(X.shape[1]), "models": results,
            "top_features": imp.head(10).round(4).to_dict(),
            "warning": "Tiny sample size: accuracy estimates are high-variance" if len(y) < 12 else ""}


# ------------------------------------------------------------------------- summary
def summarise(qc: dict, dfn: pd.DataFrame, res: pd.DataFrame, ml: dict, fdr_cut: float = 0.05) -> dict:
    sig = res[res["fdr"] < fdr_cut]
    return {
        "qc": qc,
        "normalisation": dfn.attrs.get("normalisation"),
        "comparison": f"{res.attrs.get('treatment')} vs {res.attrs.get('reference')}",
        "test": "Welch two-sample t-test per feature; Benjamini-Hochberg FDR",
        "n_features_tested": int(len(res)),
        "n_significant_fdr": int(len(sig)),
        "n_up": int((sig["log2FC"] > 0).sum()), "n_down": int((sig["log2FC"] < 0).sum()),
        "top_features": res.head(10)[["feature", "log2FC", "p_value", "fdr"]].round(4).to_dict("records"),
        "machine_learning": ml,
    }

18.2.3 data.py

"""
data.py — load a real CSV or generate a synthetic omics-style demo dataset.

Expected layout (samples as rows):
    sample_id, <group_col>, feature_1, feature_2, ...
e.g. an RNA-seq count table transposed so each row is a sample and each column a gene.

Codanics — Python ka Chilla: Build AI Agents
https://codanics.com/courses/python-ka-chilla-build-ai-agents/
"""
from __future__ import annotations

from pathlib import Path

import numpy as np
import pandas as pd


def make_demo(n_per_group: int = 3, n_features: int = 300, frac_de: float = 0.1,
              groups: tuple[str, str] = ("control", "heat"), seed: int = 42) -> pd.DataFrame:
    """Synthetic RNA-seq-like counts: negative-binomial-ish, 10% of genes differentially expressed."""
    rng = np.random.default_rng(seed)
    base = rng.gamma(shape=2.0, scale=200.0, size=n_features)          # gene-specific mean expression
    n_de = int(n_features * frac_de)
    effect = np.ones(n_features)
    de_idx = rng.choice(n_features, n_de, replace=False)
    effect[de_idx] = rng.choice([0.15, 0.25, 4.0, 6.0], size=n_de)       # down or up in treatment
    rows, ids, labels = [], [], []
    for g_i, g in enumerate(groups):
        for r in range(n_per_group):
            mu = base * (effect if g_i == 1 else 1.0)
            counts = rng.poisson(rng.gamma(shape=25.0, scale=mu / 25.0))   # over-dispersed counts
            rows.append(counts)
            ids.append(f"{g}_{r + 1}")
            labels.append(g)
    df = pd.DataFrame(rows, columns=[f"gene_{i + 1}" for i in range(n_features)])
    df.insert(0, "condition", labels)
    df.insert(0, "sample_id", ids)
    return df


def load_csv(path: Path, group_col: str, id_col: str | None = None) -> pd.DataFrame:
    df = pd.read_csv(path)
    if group_col not in df.columns:
        raise ValueError(f"group column '{group_col}' not in CSV columns: {list(df.columns)[:10]}...")
    if id_col and id_col in df.columns:
        df = df.rename(columns={id_col: "sample_id"})
    if "sample_id" not in df.columns:
        df.insert(0, "sample_id", [f"s{i + 1}" for i in range(len(df))])
    return df

18.2.4 llm.py

"""
llm.py — a thin wrapper around the Anthropic Claude API used by this agent.

Two functions are all you need:
    ask(system, user)                  -> free text (streamed)
    ask_structured(system, user, Model) -> a validated Pydantic object (JSON structured output)

Why a wrapper?  Every agent in this project talks to the model the same way.  Keeping the
API details in ONE file per agent means you can swap the model or provider in one place.

Codanics — Python ka Chilla: Build AI Agents
https://codanics.com/courses/python-ka-chilla-build-ai-agents/
"""
from __future__ import annotations

import os
from typing import Type, TypeVar

import anthropic
from dotenv import load_dotenv
from pydantic import BaseModel

load_dotenv()  # reads ANTHROPIC_API_KEY / CLAUDE_MODEL from a .env file if present

MODEL = os.getenv("CLAUDE_MODEL", "claude-opus-5")
T = TypeVar("T", bound=BaseModel)

_client: anthropic.Anthropic | None = None


def client() -> anthropic.Anthropic:
    """Lazily create one shared client (credentials come from the environment)."""
    global _client
    if _client is None:
        _client = anthropic.Anthropic()
    return _client


def _check(msg: anthropic.types.Message) -> None:
    if msg.stop_reason == "refusal":
        detail = getattr(msg, "stop_details", None)
        raise RuntimeError(f"The model refused this request: {detail}")
    if msg.stop_reason == "max_tokens":
        raise RuntimeError("Output was cut off (max_tokens). Increase max_tokens and retry.")


def _text(msg: anthropic.types.Message) -> str:
    return "".join(b.text for b in msg.content if b.type == "text")


def ask(system: str, user: str, *, max_tokens: int = 16000, effort: str = "medium",
        echo: bool = False) -> str:
    """Free-text generation. Streams so long outputs never hit HTTP timeouts."""
    with client().messages.stream(
        model=MODEL,
        max_tokens=max_tokens,
        system=system,
        messages=[{"role": "user", "content": user}],
        thinking={"type": "adaptive"},
        output_config={"effort": effort},
    ) as stream:
        if echo:
            for chunk in stream.text_stream:
                print(chunk, end="", flush=True)
            print()
        msg = stream.get_final_message()
    _check(msg)
    return _text(msg)


def ask_structured(system: str, user: str, model_cls: Type[T], *, max_tokens: int = 16000,
                   effort: str = "medium") -> T:
    """Structured generation: the SDK converts the Pydantic model to a JSON schema, the API
    constrains the output to it, and `parsed_output` is a validated instance of `model_cls`."""
    with client().messages.stream(
        model=MODEL,
        max_tokens=max_tokens,
        system=system,
        messages=[{"role": "user", "content": user}],
        thinking={"type": "adaptive"},
        output_config={"effort": effort},
        output_format=model_cls,
    ) as stream:
        msg = stream.get_final_message()
    _check(msg)
    if msg.parsed_output is None:  # should not happen with a constrained schema; be defensive
        return model_cls.model_validate_json(_text(msg))
    return msg.parsed_output

18.2.5 main.py

"""
AI Research Experiment Agent — "Virtual Research Assistant"
============================================================
description (+ data) ─► design ─► QC/normalise/test/ML (code) ─► statistics review ─► Methods & Results + notebook

Usage
    python main.py --description "RNA-seq of heat-stressed rice, control vs heat, n=3" --demo
    python main.py --description "..." --data mydata.csv --group-col condition
    python main.py --description "..." --demo --no-llm        # run only the numeric pipeline (no API key needed)

Outputs: outputs/<slug>/ design.json, data.csv, summary.json, differential.csv, pca.png, volcano.png,
         stats_review.json, methods_results.md, analysis_notebook.ipynb

Codanics — Python ka Chilla: Build AI Agents
https://codanics.com/courses/python-ka-chilla-build-ai-agents/
"""
from __future__ import annotations

import json
import sys
from pathlib import Path

import typer
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from slugify import slugify

sys.path.insert(0, str(Path(__file__).parent))
if hasattr(sys.stdout, "reconfigure"):  # Windows consoles default to cp1252; we print → ✓ etc.
    sys.stdout.reconfigure(encoding="utf-8")
from analysis import (differential, machine_learning, normalise, plot_pca,  # noqa: E402
                      plot_volcano, quality_control, summarise)
from data import load_csv, make_demo  # noqa: E402
from notebook import build_notebook  # noqa: E402

app = typer.Typer(add_completion=False)
console = Console()
AGENT_DIR = Path(__file__).parent
OUT_DIR = AGENT_DIR / "outputs"

BANNER = """[bold cyan]AI Research Experiment Agent[/] — Virtual Research Assistant
[dim]Codanics · Python ka Chilla: Build AI Agents · https://codanics.com/courses/python-ka-chilla-build-ai-agents/
YouTube @codanics · WhatsApp https://whatsapp.com/channel/0029Va7nRDq3QxRzGqaQvS3r[/]"""


@app.command()
def run(
    description: str = typer.Option(..., help="One or two sentences describing the experiment/data"),
    data: Path = typer.Option(None, help="CSV with samples as rows (see data.py)"),
    group_col: str = typer.Option("condition", help="Column holding the group / treatment label"),
    demo: bool = typer.Option(False, help="Generate synthetic RNA-seq-like demo data"),
    n_per_group: int = typer.Option(3), n_features: int = typer.Option(300),
    no_llm: bool = typer.Option(False, help="Skip the LLM steps (numeric pipeline + notebook only)"),
    echo: bool = typer.Option(True),
):
    console.print(Panel(BANNER, expand=False))
    out = OUT_DIR / slugify(description)[:60]
    out.mkdir(parents=True, exist_ok=True)

    # 1) Data --------------------------------------------------------------------
    console.rule("[bold]1. Data")
    if demo or data is None:
        df = make_demo(n_per_group=n_per_group, n_features=n_features)
        console.print(f"  synthetic demo dataset: {df.shape[0]} samples × {n_features} features")
    else:
        df = load_csv(data, group_col)
        console.print(f"  loaded {data}: {df.shape}")
    csv_path = out / "data.csv"
    df.to_csv(csv_path, index=False)
    data_hint = (f"Data available: {df.shape[0]} samples, {df.shape[1] - 2} numeric features, "
                 f"groups = {df[group_col].value_counts().to_dict()}")

    # 2) Design (LLM) -----------------------------------------------------------
    design = None
    if not no_llm:
        console.rule("[bold]2. Experiment design")
        from agents import ExperimentPlanner
        design = ExperimentPlanner().run(description, data_hint)
        console.print(f"  [bold]{design.title}[/]\n  H: {design.hypothesis}\n  tests: {design.statistical_tests}")
        (out / "design.json").write_text(design.model_dump_json(indent=2), encoding="utf-8")

    # 3) Analysis (code) --------------------------------------------------------
    console.rule("[bold]3. Analysis engine (pandas · scipy · statsmodels · scikit-learn)")
    qc = quality_control(df, group_col)
    dfn = normalise(df, group_col)
    res = differential(dfn, group_col)
    res.to_csv(out / "differential.csv", index=False)
    plot_pca(dfn, group_col, out / "pca.png")
    plot_volcano(res, out / "volcano.png")
    ml = machine_learning(dfn, group_col)
    summary = summarise(qc, dfn, res, ml)
    (out / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")

    console.print(f"  normalisation: {summary['normalisation']}   comparison: {summary['comparison']}")
    console.print(f"  significant features (FDR<0.05): {summary['n_significant_fdr']} "
                  f"(up {summary['n_up']}, down {summary['n_down']})")
    t = Table(title="Top differential features")
    for c in ("feature", "log2FC", "p_value", "fdr"):
        t.add_column(c)
    for r in summary["top_features"]:
        t.add_row(str(r["feature"]), f"{r['log2FC']:.2f}", f"{r['p_value']:.2e}", f"{r['fdr']:.2e}")
    console.print(t)
    for name, m in ml["models"].items():
        console.print(f"  {name:<20} CV accuracy {m['cv_accuracy_mean']:.2f} ± {m['cv_accuracy_std']:.2f}  ({ml['cross_validation']})")
    if qc["warnings"]:
        console.print(f"  [yellow]QC warnings:[/] {qc['warnings']}")

    # 4) Notebook (code) --------------------------------------------------------
    console.rule("[bold]4. Reproducible notebook")
    nb = build_notebook(csv_path, group_col, AGENT_DIR, out / "analysis_notebook.ipynb", description)
    console.print(f"  [green]saved[/] {nb.name}")

    if no_llm:
        console.print(Panel(f"Numeric pipeline done (LLM skipped). Outputs in {out}", style="green"))
        return

    # 5) Statistics review + report (LLM) --------------------------------------
    console.rule("[bold]5. Statistics review")
    from agents import ReportWriter, StatisticsAgent
    review = StatisticsAgent().run(design, summary)
    console.print(f"  verdict: {review.overall_verdict}")
    for c in review.concerns:
        console.print(f"   - [{c.severity}] {c.issue} → {c.recommendation}")
    (out / "stats_review.json").write_text(review.model_dump_json(indent=2), encoding="utf-8")

    console.rule("[bold]6. Methods & Results")
    report = ReportWriter().run(description, design, summary, review, echo=echo)
    (out / "methods_results.md").write_text(report, encoding="utf-8")
    console.print(Panel(f"Done. Outputs in [bold]{out}[/]", style="green"))


if __name__ == "__main__":
    app()

18.2.6 notebook.py

"""
notebook.py — generate a runnable Jupyter notebook that reproduces the analysis.

Reproducibility is a product feature: the researcher can open the notebook, change a threshold,
and re-run everything without the agent.

Codanics — Python ka Chilla: Build AI Agents
https://codanics.com/courses/python-ka-chilla-build-ai-agents/
"""
from __future__ import annotations

from pathlib import Path

import nbformat as nbf


def build_notebook(csv_path: Path, group_col: str, agent_dir: Path, out: Path, title: str) -> Path:
    nb = nbf.v4.new_notebook()
    cells = [
        nbf.v4.new_markdown_cell(
            f"# {title}\n\nReproducible analysis notebook generated by the **AI Research Experiment Agent**.\n\n"
            f"Codanics — Python ka Chilla: Build AI Agents · "
            f"https://codanics.com/courses/python-ka-chilla-build-ai-agents/"),
        nbf.v4.new_code_cell(
            "import sys, json\nfrom pathlib import Path\nimport pandas as pd\n"
            f"sys.path.insert(0, r'{agent_dir}')\n"
            "from analysis import quality_control, normalise, differential, plot_pca, plot_volcano, machine_learning, summarise\n"
            f"CSV = Path(r'{csv_path}')\nGROUP = '{group_col}'\n"
            "df = pd.read_csv(CSV)\ndf.head()"),
        nbf.v4.new_markdown_cell("## 1. Quality control"),
        nbf.v4.new_code_cell("qc = quality_control(df, GROUP)\nqc"),
        nbf.v4.new_markdown_cell("## 2. Normalisation and PCA"),
        nbf.v4.new_code_cell("dfn = normalise(df, GROUP)\nprint(dfn.attrs['normalisation'])\n"
                             "plot_pca(dfn, GROUP, Path('pca.png'))\nfrom IPython.display import Image; Image('pca.png')"),
        nbf.v4.new_markdown_cell("## 3. Differential analysis (Welch t-test + BH-FDR)"),
        nbf.v4.new_code_cell("res = differential(dfn, GROUP)\nres.head(15)"),
        nbf.v4.new_code_cell("plot_volcano(res, Path('volcano.png'))\nImage('volcano.png')"),
        nbf.v4.new_markdown_cell("## 4. Machine-learning baseline with cross-validation"),
        nbf.v4.new_code_cell("ml = machine_learning(dfn, GROUP)\nml"),
        nbf.v4.new_markdown_cell("## 5. Summary passed to the Statistics Agent"),
        nbf.v4.new_code_cell("print(json.dumps(summarise(qc, dfn, res, ml), indent=1))"),
    ]
    nb["cells"] = cells
    out.parent.mkdir(parents=True, exist_ok=True)
    nbf.write(nb, out)
    return out