leanlab 0.2.1__py3-none-any.whl

leanlab/core/doctor.py

@@ -0,0 +1,220 @@
+"""Preflight 'doctor' — verify a lab is wired correctly before running.
+
+The expensive bugs are silent wiring mismatches: lab.json names metric "FPS" but
+evaluation.py prints "score"; or the command passes --experiment but the script
+reads a positional arg. `check_lab` catches those *cheaply* by probing the
+evaluator with a sentinel missing file — the evaluator should fail fast on the
+missing file (before any render / Claude call), and from its output we can verify
+both the argument wiring and the metric key.
+"""
+
+from __future__ import annotations
+
+import json
+import shlex
+import shutil
+import subprocess
+from dataclasses import dataclass
+from importlib import resources
+from pathlib import Path
+
+OK, WARN, FAIL = "ok", "warn", "fail"
+_SENTINEL = "experiments/__leanlab_preflight__.py"
+
+
+@dataclass
+class Check:
+    name: str
+    status: str   # ok | warn | fail
+    message: str
+
+
+def _run(cmd_template, lab_dir, file_rel):
+    parts = [p.replace("{file}", file_rel) for p in shlex.split(cmd_template)]
+    try:
+        proc = subprocess.run(parts, cwd=lab_dir, capture_output=True, text=True, timeout=120)
+    except Exception as e:  # noqa: BLE001
+        return None, f"could not run `{cmd_template}`: {e}"
+    return proc, (proc.stdout + ("\n" + proc.stderr if proc.stderr else "")).strip()
+
+
+def _last_json(out):
+    for line in reversed(out.splitlines()):
+        try:
+            o = json.loads(line.strip())
+        except (ValueError, json.JSONDecodeError):
+            continue
+        if isinstance(o, dict):
+            return o
+    return None
+
+
+def _probe_args(name, cmd_template, lab_dir, checks):
+    """A command should receive the file path, not the flag. Returns the probe output."""
+    proc, out = _run(cmd_template, lab_dir, _SENTINEL)
+    if proc is None:
+        checks.append(Check(name, FAIL, out))
+        return None
+    if "--experiment" in out and _SENTINEL not in out and "__leanlab_preflight__" not in out:
+        checks.append(Check(name, FAIL, "the script is not reading the file argument (it got the "
+                                        "flag instead) — fix its arg parsing or the command in lab.json"))
+    else:
+        checks.append(Check(name, OK, "the file argument reaches the script"))
+    return out
+
+
+def check_lab(lab_dir) -> list[Check]:
+    lab = Path(lab_dir)
+    checks: list[Check] = []
+
+    cfgpath = lab / "lab.json"
+    if not cfgpath.exists():
+        return [Check("lab.json", FAIL, "missing — is this a lab folder?")]
+    try:
+        cfg = json.loads(cfgpath.read_text())
+    except (ValueError, json.JSONDecodeError) as e:
+        return [Check("lab.json", FAIL, f"invalid JSON: {e}")]
+
+    obj = cfg.get("objective") or {}
+    metric, direction = obj.get("metric"), obj.get("direction")
+    missing = [k for k in ("eval_cmd", "validate_cmd", "experiments_dir", "results_file")
+               if not cfg.get(k)]
+    if missing or not metric or not direction:
+        checks.append(Check("lab.json", FAIL,
+                            f"missing {missing or ''} objective.metric/direction={metric}/{direction}"))
+    else:
+        checks.append(Check("lab.json", OK, f"objective {direction} {metric}"))
+
+    for f in ("task.md", "evaluation.py", "validate.py"):
+        checks.append(Check(f, OK if (lab / f).exists() else FAIL,
+                            "present" if (lab / f).exists() else "missing"))
+    exp = lab / cfg.get("experiments_dir", "experiments")
+    checks.append(Check("experiments/", OK if exp.is_dir() else FAIL,
+                        "present" if exp.is_dir() else "missing"))
+
+    checks.append(Check("claude CLI", OK if shutil.which("claude") else FAIL,
+                        "on PATH" if shutil.which("claude")
+                        else "not found — workers/director/critic/judge cannot run"))
+
+    try:
+        base = resources.files("leanlab") / "templates" / "agents"
+        have = all((base / s).is_file() for s in ("CLAUDE.md", "director.md", "critic.md"))
+        checks.append(Check("agent specs", OK if have else FAIL,
+                            "resolved from package" if have else "missing from package"))
+    except Exception as e:  # noqa: BLE001
+        checks.append(Check("agent specs", FAIL, f"cannot resolve: {e}"))
+
+    # Actually build the worker/director/critic prompts for this lab — the real wiring test
+    # (injects the specs and reads the lab's memory + Director/Critic notes).
+    try:
+        from .loop import build_worker_prompt, build_director_prompt, build_critic_prompt
+        build_worker_prompt(lab, cfg)
+        build_director_prompt()
+        build_critic_prompt()
+        checks.append(Check("agent prompts", OK, "worker / director / critic prompts build"))
+    except Exception as e:  # noqa: BLE001
+        checks.append(Check("agent prompts", FAIL, f"cannot build agent prompts: {e}"))
+
+    for f in ("Director_Notes.md", "Critic_Feedback.md"):
+        checks.append(Check(f, OK if (lab / f).exists() else WARN,
+                            "present" if (lab / f).exists() else "missing (created on first review)"))
+
+    # The wiring probes — cheap, run the eval/validate on a missing sentinel file.
+    if (lab / "evaluation.py").exists() and cfg.get("eval_cmd"):
+        out = _probe_args("eval args", cfg["eval_cmd"], lab, checks)
+        if out is not None and metric:
+            verdict = _last_json(out)
+            if verdict is None:
+                checks.append(Check("eval metric", WARN,
+                                    "evaluator emitted no JSON for a missing file — can't verify the metric key"))
+            elif metric not in verdict:
+                checks.append(Check("eval metric", FAIL,
+                                    f'evaluator output has no "{metric}" key (lab.json objective); '
+                                    f"it printed keys {list(verdict)[:6]}"))
+            else:
+                checks.append(Check("eval metric", OK, f'emits the objective key "{metric}"'))
+    if (lab / "validate.py").exists() and cfg.get("validate_cmd"):
+        _probe_args("validate args", cfg["validate_cmd"], lab, checks)
+
+    return checks
+
+
+def summarize(checks):
+    return {s: sum(c.status == s for c in checks) for s in (OK, WARN, FAIL)}
+
+
+def ok(checks):
+    return not any(c.status == FAIL for c in checks)
+
+
+# --- reporting --------------------------------------------------------------
+_SYMBOL = {OK: "[green]✓[/green]", WARN: "[yellow]⚠[/yellow]", FAIL: "[red]✗[/red]"}
+
+
+class RichReport:
+    """Default terminal reporter for check/fix."""
+
+    def __init__(self):
+        from rich.console import Console
+        self.console = Console()
+
+    def report(self, checks):
+        for c in checks:
+            self.console.print(f"{_SYMBOL.get(c.status, '?')} [bold]{c.name}[/bold] — {c.message}")
+        s = summarize(checks)
+        self.console.print(f"\n[bold]{s[OK]} ok · {s[WARN]} warn · {s[FAIL]} fail[/bold]")
+
+    def status(self, message):
+        return self.console.status(f"[bold cyan]{message}", spinner="dots")
+
+    def note(self, message):
+        self.console.print(message)
+
+
+# --- automated fixing -------------------------------------------------------
+def _fix_prompt(fails):
+    lines = "\n".join(f"- {c.name}: {c.message}" for c in fails)
+    return (
+        "You are fixing a leanlab lab in the current directory. `leanlab check` found these "
+        f"wiring problems:\n\n{lines}\n\n"
+        "Fix them by editing the lab's files so they are mutually consistent:\n"
+        "- evaluation.py and validate.py MUST parse their CLI args to match lab.json's eval_cmd / "
+        "validate_cmd (which pass `--experiment <path>`) — use argparse, NOT a positional arg.\n"
+        "- evaluation.py MUST print ONE line of JSON whose keys include exactly the objective "
+        "metric named in lab.json (objective.metric) — not a generic 'score' unless that IS the "
+        "metric.\n"
+        "- Create any missing files the checks listed.\n"
+        "Make the edits now with your tools, then stop. Do NOT run experiments."
+    )
+
+
+def fix_lab(lab_dir, *, runner=None, ui=None, rounds=3) -> bool:
+    """Run checks; if any fail, have Claude edit the lab to fix them, then re-check. Loop."""
+    lab = Path(lab_dir)
+    ui = ui or RichReport()
+    from .loop import make_runner
+    runner = runner or make_runner(lab)
+
+    for attempt in range(1, rounds + 1):
+        checks = check_lab(lab)
+        fails = [c for c in checks if c.status == FAIL]
+        ui.report(checks)
+        if not fails:
+            ui.note("\n[green]✓ all checks pass — nothing to fix.[/green]")
+            return True
+        ui.note(f"\n[yellow]Fixing {len(fails)} issue(s) with Claude (round {attempt}/{rounds})…[/yellow]")
+        # evaluation.py is often locked read-only — unlock so the agent can edit it.
+        ev = lab / "evaluation.py"
+        relock = ev.exists() and not (ev.stat().st_mode & 0o200)
+        if relock:
+            ev.chmod(0o644)
+        try:
+            with ui.status("Claude is editing the lab to fix the issues…"):
+                runner.run_plain(_fix_prompt(fails))
+        finally:
+            if relock and ev.exists():
+                ev.chmod(0o444)
+
+    checks = check_lab(lab)
+    ui.report(checks)
+    return ok(checks)

leanlab/core/init.py

@@ -0,0 +1,219 @@
+"""Interactive `leanlab init` — the init-architect.
+
+Realizes the init-lab use case. The architect (Claude, reached through the same
+AgentRunner the workers use) drafts task.md + the objective from the operator's
+plain-words description, then proposes an evaluator in a loop until the operator
+approves. All terminal I/O goes through an injected `ui` (default: a rich + questionary
+console) so the flow is fully testable without a real terminal or Claude.
+"""
+
+from __future__ import annotations
+
+import json
+import subprocess
+from pathlib import Path
+
+from .loop import make_runner
+
+LAB_JSON_TEMPLATE = {
+    "name": "",
+    "description": "TODO: one line describing the task.",
+    "objective": {"metric": "score", "direction": "max"},
+    "experiments_dir": "experiments",
+    "results_file": "results.jsonl",
+    "validate_cmd": "uv run python validate.py --experiment {file}",
+    "eval_cmd": "uv run python evaluation.py --experiment {file}",
+    "director_every": 5,
+    "critic_every": 5,
+    "max_fix_calls": 3,
+}
+
+
+def scaffold(lab: Path, name: str) -> None:
+    """Create the empty .leanlab/<name>/ skeleton (no agent specs — those live in the package)."""
+    lab = Path(lab)
+    (lab / "experiments").mkdir(parents=True)
+    (lab / "lab.json").write_text(json.dumps(dict(LAB_JSON_TEMPLATE, name=name), indent=2) + "\n")
+    (lab / "results.jsonl").write_text("")
+    (lab / "Director_Notes.md").write_text("# Director Notes\n\nNeutral — no experiments yet.\n")
+    (lab / "Critic_Feedback.md").write_text("# Critic Feedback\n\nNo experiments reviewed yet.\n")
+
+
+def _draft_prompt(description: str) -> str:
+    return (
+        "You are the lab ARCHITECT for leanlab. The operator wants to research:\n\n"
+        f"{description}\n\n"
+        "Base your decision ONLY on the task description above. Ignore any other labs, "
+        "example projects, or files that may exist in this directory or its parents — they are "
+        "unrelated and must not influence the metric or framing. "
+        "Decide what a single 'experiment' is for this task and how success is measured. "
+        "Choose the objective metric that is STANDARD and APPROPRIATE for THIS task — judge it "
+        "on the task's own terms; do NOT default to any particular metric. Match the metric to "
+        "the task type, e.g.: classification -> accuracy / F1 / ROC-AUC; regression -> RMSE / "
+        "MAE / R2; ranking -> NDCG / MAP; clustering -> silhouette. Set direction (min or max) "
+        "to fit, and give a one-line justification for the choice. If the request is not a "
+        "measurable experiment as written, reframe it into the closest measurable one and say so. "
+        "Write a clear task.md: the goal, the chosen metric (with the justification), and the "
+        "experiment contract — what ONE file in experiments/ must define so evaluation.py can run "
+        "it. You may research the web. Do NOT create or edit any files. Reply with ONLY this JSON "
+        'object: {"task_md": "<full markdown>", '
+        '"objective": {"metric": "<name>", "direction": "min|max"}}'
+    )
+
+
+def _propose_prompt(feedback: str | None, metric: str, eval_cmd: str, validate_cmd: str) -> str:
+    base = (
+        "Propose how to EVALUATE this task, using the objective metric you chose. If it needs "
+        "data, say what data and how you would obtain and split it. Provide the full CONTENTS of "
+        "a frozen evaluation.py and a cheap validate.py.\n"
+        "HARD REQUIREMENTS — the lab will NOT score unless these hold exactly:\n"
+        f"1. leanlab runs your files from inside the lab dir with EXACTLY these commands, where "
+        f"{{file}} is the experiment path:\n"
+        f"     evaluation:  {eval_cmd}\n"
+        f"     validate:    {validate_cmd}\n"
+        f"   So evaluation.py and validate.py MUST parse their arguments to match — use argparse "
+        f"with a '--experiment' option. Do NOT read a positional argument like sys.argv[1].\n"
+        f'2. evaluation.py MUST print ONE line of JSON whose keys include exactly "{metric}" '
+        f'(the objective metric, spelled exactly — NOT a generic "score"). Extra metrics are fine '
+        f"as additional flat (number) keys; avoid nested objects.\n"
+        "3. Do NOT create or edit any files — return everything inside the JSON only.\n"
+        "4. evaluation.py and validate.py MUST check the experiment file exists and fail fast FIRST "
+        "— before importing heavy libraries, rendering, or calling any model — so a quick preflight "
+        "stays cheap.\n"
+        "5. List in 'packages' every third-party pip package evaluation.py or validate.py imports "
+        "(e.g. playwright, numpy) so the lab can install them. Use [] if none. If you need an LLM "
+        "to judge, call the `claude` CLI via subprocess (no API key needed) rather than an SDK.\n"
+        "Reply with ONLY this JSON object: "
+        '{"summary": "<plain-English approach, 3-5 sentences>", '
+        '"evaluation_py": "<full file contents>", "validate_py": "<full file contents>", '
+        '"packages": ["<pip name>", ...]}'
+    )
+    if feedback:
+        return (f"The operator gave this feedback on your previous proposal:\n\n{feedback}\n\n"
+                f"Revise accordingly. {base}")
+    return base
+
+
+class RichUI:
+    """Terminal UI for `init` — spinners, panels, syntax-highlighted code, arrow-key menus."""
+
+    def __init__(self):
+        from rich.console import Console
+        self.console = Console()
+
+    def status(self, message):
+        return self.console.status(f"[bold cyan]{message}", spinner="dots")
+
+    def note(self, message):
+        self.console.print(message)
+
+    def error(self, message):
+        self.console.print(f"[bold red]{message}[/bold red]")
+
+    def objective(self, obj):
+        from rich.panel import Panel
+        self.console.print(Panel(f"[bold]{obj.get('direction')} {obj.get('metric')}[/bold]",
+                                 title="Objective", border_style="cyan", expand=False))
+
+    def proposal(self, summary):
+        from rich.markdown import Markdown
+        from rich.panel import Panel
+        self.console.print(Panel(Markdown(summary), title="Proposed evaluation",
+                                 border_style="magenta"))
+
+    def decide(self, evaluation_py):
+        import questionary
+        from rich.syntax import Syntax
+        approve, view, feedback, cancel = (
+            "✓ Approve & write the files", "👁 View the generated evaluation.py",
+            "✍ Give feedback (revise)", "✖ Cancel")
+        while True:
+            choice = questionary.select("What now?", choices=[approve, view, feedback, cancel]).ask()
+            if choice is None or choice == cancel:
+                return ("cancel", None)
+            if choice == view:
+                self.console.print(Syntax(evaluation_py, "python", theme="ansi_dark",
+                                          line_numbers=True, word_wrap=True))
+                continue
+            if choice == approve:
+                return ("approve", None)
+            return ("feedback", questionary.text("Your feedback for Claude:").ask() or "")
+
+    def success(self, lab, name):
+        from rich.panel import Panel
+        self.console.print(Panel(
+            f"Lab ready: [bold]{lab}[/bold]\n\nReview the evaluator, then run:\n"
+            f"  [green]leanlab lock {name} && leanlab run {name}[/green]",
+            title="✓ Done", border_style="green"))
+
+
+def run_init(lab, name, description, *, runner=None, ui=None, verify=True, yes=False) -> None:
+    """Scaffold a lab, then draft it with the architect and approve the evaluator in a loop."""
+    lab = Path(lab)
+    scaffold(lab, name)
+    ui = ui or RichUI()
+    runner = runner or make_runner(lab)
+
+    with ui.status("Drafting task.md and the objective with Claude…"):
+        draft = runner.run_structured(_draft_prompt(description), ["task_md", "objective"])
+    if not draft.ok:
+        ui.error("Could not draft the task — try again with a clearer description.")
+        return
+    (lab / "task.md").write_text(draft.data["task_md"])
+    cfg = json.loads((lab / "lab.json").read_text())
+    obj = draft.data["objective"]
+    if isinstance(obj, dict) and "metric" in obj and "direction" in obj:
+        cfg["objective"] = obj
+    else:
+        ui.error("⚠ Claude returned an unexpected objective shape — set lab.json's objective by hand.")
+    (lab / "lab.json").write_text(json.dumps(cfg, indent=2) + "\n")
+    ui.objective(cfg["objective"])
+
+    metric = cfg["objective"].get("metric", "score")
+    session, feedback = draft.session_id, None
+    while True:
+        with ui.status("Designing an evaluator with Claude…"):
+            prop = runner.run_structured(
+                _propose_prompt(feedback, metric, cfg["eval_cmd"], cfg["validate_cmd"]),
+                ["summary", "evaluation_py", "validate_py"], session=session)
+        if not prop.ok:
+            ui.error("Could not propose an evaluation — aborting.")
+            return
+        session = prop.session_id or session
+        ui.proposal(prop.data["summary"])
+        action, text = ("approve", None) if yes else ui.decide(prop.data["evaluation_py"])
+        if action == "approve":
+            (lab / "evaluation.py").write_text(prop.data["evaluation_py"])
+            (lab / "validate.py").write_text(prop.data["validate_py"])
+            _install_packages(prop.data.get("packages"), ui)
+            break
+        if action == "cancel":
+            ui.note("Cancelled — no evaluator written. task.md and lab.json were kept.")
+            return
+        feedback = text
+
+    if verify:
+        _self_verify(lab, runner, ui)
+    ui.success(lab, name)
+
+
+def _install_packages(packages, ui):
+    """uv add the pip packages the evaluator/validator declared they need."""
+    for pkg in packages or []:
+        if not isinstance(pkg, str) or not pkg.strip():
+            continue
+        with ui.status(f"Installing {pkg}…"):
+            r = subprocess.run(["uv", "add", pkg.strip()], capture_output=True, text=True)
+        if r.returncode != 0:
+            ui.error(f"⚠ could not install {pkg!r} (uv add failed) — install it yourself: "
+                     f"{(r.stderr or '').strip()[:120]}")
+
+
+def _self_verify(lab, runner, ui):
+    """Run the doctor; if the generated lab is mis-wired, have Claude fix it before finishing."""
+    from .doctor import check_lab, ok, fix_lab
+    if ok(check_lab(lab)):
+        ui.note("[green]✓ wiring check passed.[/green]")
+        return
+    ui.note("[yellow]⚠ wiring check found problems — fixing with Claude before finishing…[/yellow]")
+    fix_lab(lab, runner=runner)