Add deepagents_rl package sources

src/deepagents_rl/__init__.py

@@ -0,0 +1,3 @@
+"""deepagents-rl: coding-agent RL post-training with full observability."""
+
+__version__ = "0.1.0"

src/deepagents_rl/data/__init__.py

@@ -0,0 +1,3 @@
+from .mbpp import build_train_eval, load_mbpp_sanitized, save_jsonl, to_hf_dataset
+
+__all__ = ["build_train_eval", "load_mbpp_sanitized", "save_jsonl", "to_hf_dataset"]

src/deepagents_rl/data/mbpp.py

@@ -0,0 +1,142 @@
+"""MBPP sanitized load/filter for small-model coding RL."""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+from deepagents_rl.rewards.unit_test import run_hidden_tests
+
+
+def _row_to_task(ex: dict[str, Any], split: str) -> dict[str, Any] | None:
+    """Normalize an MBPP sanitized example."""
+    task_id = ex.get("task_id") or ex.get("id")
+    text = ex.get("text") or ex.get("prompt") or ex.get("description") or ""
+    code = ex.get("code") or ex.get("reference_code") or ""
+    test_list = list(ex.get("test_list") or ex.get("test_cases") or [])
+    if not text or not test_list:
+        return None
+    # Leakage-aware: first assert may appear in prompt as visible example
+    visible = [test_list[0]] if test_list else []
+    hidden = list(test_list)
+    prompt = (
+        f"Write a Python solution for the following problem.\n\n"
+        f"{text.strip()}\n\n"
+        f"Visible example test (not the full grader):\n{visible[0]}\n\n"
+        f"Return ONLY a ```python``` code block with the full solution."
+    )
+    return {
+        "task_id": str(task_id),
+        "split": split,
+        "text": text.strip(),
+        "reference_code": code,
+        "visible_tests": visible,
+        "hidden_tests": hidden,
+        "prompt": prompt,
+        "test_list": test_list,
+    }
+
+
+def filter_solvable(tasks: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
+    """Keep tasks whose reference solution passes hidden tests in our sandbox."""
+    keep: list[dict[str, Any]] = []
+    drop: list[dict[str, Any]] = []
+    for t in tasks:
+        ref = t.get("reference_code") or ""
+        if not ref.strip():
+            drop.append({**t, "drop_reason": "no_reference"})
+            continue
+        grade = run_hidden_tests(ref, t["hidden_tests"], timeout_s=5.0)
+        if grade["pass_frac"] >= 1.0:
+            keep.append(t)
+        else:
+            drop.append({**t, "drop_reason": "ref_fail", "grade": grade})
+    return keep, drop
+
+
+def load_mbpp_sanitized() -> dict[str, list[dict[str, Any]]]:
+    """Load google-research-datasets/mbpp sanitized config from HF."""
+    from datasets import load_dataset
+
+    # Prefer sanitized; fall back to full mbpp
+    try:
+        ds = load_dataset("google-research-datasets/mbpp", "sanitized")
+    except Exception:
+        ds = load_dataset("mbpp", "sanitized")
+
+    out: dict[str, list[dict[str, Any]]] = {}
+    for split in ds:
+        tasks = []
+        for ex in ds[split]:
+            t = _row_to_task(dict(ex), split=split)
+            if t:
+                tasks.append(t)
+        out[split] = tasks
+    return out
+
+
+def build_train_eval(
+    *,
+    filter_refs: bool = True,
+) -> dict[str, list[dict[str, Any]]]:
+    """
+    Train = sanitized train + validation + prompt (none of test).
+    Eval = sanitized test (held-out).
+    """
+    raw = load_mbpp_sanitized()
+    train_parts = []
+    for key in ("train", "validation", "prompt"):
+        train_parts.extend(raw.get(key, []))
+    eval_parts = list(raw.get("test", []))
+
+    drop_log: list[dict[str, Any]] = []
+    if filter_refs:
+        train_parts, d1 = filter_solvable(train_parts)
+        eval_parts, d2 = filter_solvable(eval_parts)
+        drop_log = d1 + d2
+
+    # Disjoint task ids
+    train_ids = {t["task_id"] for t in train_parts}
+    eval_parts = [t for t in eval_parts if t["task_id"] not in train_ids]
+
+    return {
+        "train": train_parts,
+        "eval": eval_parts,
+        "dropped": drop_log,
+    }
+
+
+def save_jsonl(tasks: list[dict[str, Any]], path: Path) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as f:
+        for t in tasks:
+            # Don't need full drop grades in published rows
+            row = {k: v for k, v in t.items() if k not in ("drop_reason", "grade")}
+            f.write(json.dumps(row) + "\n")
+
+
+def to_hf_dataset(splits: dict[str, list[dict[str, Any]]]):
+    from datasets import Dataset, DatasetDict
+
+    def _clean(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        cleaned = []
+        for t in rows:
+            cleaned.append(
+                {
+                    "task_id": t["task_id"],
+                    "split": t["split"],
+                    "text": t["text"],
+                    "prompt": t["prompt"],
+                    "reference_code": t.get("reference_code", ""),
+                    "visible_tests": t["visible_tests"],
+                    "hidden_tests": t["hidden_tests"],
+                }
+            )
+        return cleaned
+
+    return DatasetDict(
+        {
+            "train": Dataset.from_list(_clean(splits["train"])),
+            "eval": Dataset.from_list(_clean(splits["eval"])),
+        }
+    )

src/deepagents_rl/observability/__init__.py

@@ -0,0 +1,3 @@
+from .langfuse_eval import init_langfuse, log_checkpoint_eval
+
+__all__ = ["init_langfuse", "log_checkpoint_eval"]

src/deepagents_rl/observability/langfuse_eval.py

@@ -0,0 +1,87 @@
+"""Langfuse logging for checkpoint eval batches (Kaggle + local)."""
+from __future__ import annotations
+
+import os
+from typing import Any
+
+
+def init_langfuse():
+    from langfuse import Langfuse, get_client
+
+    Langfuse(
+        public_key=os.environ["LANGFUSE_PUBLIC_KEY"],
+        secret_key=os.environ["LANGFUSE_SECRET_KEY"],
+        host=os.getenv("LANGFUSE_HOST")
+        or os.getenv("LANGFUSE_BASE_URL")
+        or "https://cloud.langfuse.com",
+    )
+    lf = get_client()
+    if not lf.auth_check():
+        raise RuntimeError("Langfuse auth_check failed")
+    return lf
+
+
+def log_checkpoint_eval(
+    *,
+    step: int,
+    hf_revision: str,
+    train_model_id: str,
+    results: list[dict[str, Any]],
+    wandb_run_id: str | None = None,
+    trace_name: str = "kaggle-ckpt-eval",
+) -> dict[str, Any]:
+    """
+    Write one Langfuse agent trace for a checkpoint eval suite.
+    Each problem becomes a nested span; aggregate scores on the trace.
+    MUST be called BEFORE pushing/saving the checkpoint artifact.
+    """
+    lf = init_langfuse()
+    n = len(results)
+    mean_pass = sum(float(r.get("pass_frac", 0.0)) for r in results) / n if n else 0.0
+    mean_reward = sum(float(r.get("reward_total", 0.0)) for r in results) / n if n else 0.0
+
+    with lf.start_as_current_observation(
+        name=trace_name,
+        as_type="evaluator",
+        metadata={
+            "step": step,
+            "hf_revision": hf_revision,
+            "train_model_id": train_model_id,
+            "wandb_run_id": wandb_run_id,
+            "phase": "ckpt_eval",
+            "n_problems": n,
+        },
+        input={"step": step, "hf_revision": hf_revision, "n_problems": n},
+    ) as root:
+        for r in results:
+            with lf.start_as_current_observation(
+                name=f"task:{r.get('task_id')}",
+                as_type="span",
+                input={"prompt": r.get("prompt", "")[:2000]},
+                metadata={"task_id": r.get("task_id")},
+            ) as span:
+                span.update(
+                    output={
+                        "completion": (r.get("completion") or "")[:2000],
+                        "pass_frac": r.get("pass_frac"),
+                        "reward_total": r.get("reward_total"),
+                        "ok": r.get("pass_frac", 0) >= 1.0,
+                    }
+                )
+        summary = {
+            "mean_pass_frac": mean_pass,
+            "mean_reward_total": mean_reward,
+            "n_problems": n,
+            "n_perfect": sum(1 for r in results if r.get("pass_frac", 0) >= 1.0),
+            "step": step,
+            "hf_revision": hf_revision,
+        }
+        root.update(output=summary)
+        lf.score_current_trace(name="mean_pass_frac", value=mean_pass)
+        lf.score_current_trace(name="mean_reward_total", value=mean_reward)
+        lf.score_current_trace(name="n_perfect", value=float(summary["n_perfect"]))
+        lf.score_current_trace(name="checkpoint_step", value=float(step))
+        trace_id = lf.get_current_trace_id() or ""
+
+    lf.flush()
+    return {"langfuse_trace_id": trace_id, **summary}

src/deepagents_rl/rewards/__init__.py

@@ -0,0 +1,8 @@
+from .unit_test import compute_reward, extract_python_code, grade_completion, grade_solution_file
+
+__all__ = [
+    "compute_reward",
+    "extract_python_code",
+    "grade_completion",
+    "grade_solution_file",
+]

src/deepagents_rl/rewards/unit_test.py

@@ -0,0 +1,115 @@
+"""Verifiable unit-test rewards for coding RL."""
+from __future__ import annotations
+
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any
+
+
+def run_hidden_tests(
+    solution_code: str,
+    hidden_tests: list[str],
+    *,
+    timeout_s: float = 5.0,
+    python_exe: str | None = None,
+) -> dict[str, Any]:
+    """Execute each assert against solution_code; return pass fraction + details."""
+    py = python_exe or sys.executable
+    per_test: list[dict[str, Any]] = []
+    passed = 0
+    for i, assertion in enumerate(hidden_tests):
+        script = f"{solution_code}\n{assertion}\n"
+        try:
+            proc = subprocess.run(
+                [py, "-c", script],
+                capture_output=True,
+                text=True,
+                timeout=timeout_s,
+            )
+            ok = proc.returncode == 0
+            err = (proc.stderr or "")[:500]
+        except Exception as e:  # noqa: BLE001
+            ok = False
+            err = f"{type(e).__name__}: {e}"
+        per_test.append({"i": i, "ok": ok, "stderr": err})
+        if ok:
+            passed += 1
+    total = len(hidden_tests)
+    return {
+        "pass_frac": (passed / total) if total else 0.0,
+        "passed": passed,
+        "total": total,
+        "per_test": per_test,
+    }
+
+
+def extract_python_code(completion: str) -> str:
+    """Pull first ```python fenced block, else raw text."""
+    if "```python" in completion:
+        start = completion.find("```python") + len("```python")
+        end = completion.find("```", start)
+        if end == -1:
+            return completion[start:].strip()
+        return completion[start:end].strip()
+    if "```" in completion:
+        start = completion.find("```") + 3
+        end = completion.find("```", start)
+        if end != -1:
+            return completion[start:end].strip()
+    return completion.strip()
+
+
+def grade_completion(
+    completion: str,
+    hidden_tests: list[str],
+    *,
+    timeout_s: float = 5.0,
+) -> dict[str, Any]:
+    code = extract_python_code(completion)
+    grade = run_hidden_tests(code, hidden_tests, timeout_s=timeout_s)
+    grade["extracted_code"] = code
+    return grade
+
+
+def compute_reward(
+    grade: dict[str, Any],
+    *,
+    format_ok: bool = True,
+    ran_execute: bool = False,
+) -> dict[str, float]:
+    components = {
+        "pass_frac": float(grade.get("pass_frac", 0.0)),
+        "format_bonus": 0.05 if format_ok and grade.get("extracted_code") else 0.0,
+        "tool_bonus": 0.05 if ran_execute else 0.0,
+        "penalties": 0.0,
+    }
+    components["reward_total"] = (
+        components["pass_frac"]
+        + components["format_bonus"]
+        + components["tool_bonus"]
+        - components["penalties"]
+    )
+    return components
+
+
+def grade_solution_file(
+    solution_path: Path,
+    hidden_tests: list[str],
+    *,
+    timeout_s: float = 5.0,
+) -> dict[str, Any]:
+    if not solution_path.exists():
+        return {
+            "pass_frac": 0.0,
+            "passed": 0,
+            "total": len(hidden_tests),
+            "per_test": [],
+            "error": "solution.py missing",
+            "extracted_code": "",
+        }
+    code = solution_path.read_text(encoding="utf-8")
+    grade = run_hidden_tests(code, hidden_tests, timeout_s=timeout_s)
+    grade["extracted_code"] = code
+    grade["error"] = None
+    return grade

src/deepagents_rl/train/__init__.py

@@ -0,0 +1,3 @@
+from .resolve_model import resolve_train_model
+
+__all__ = ["resolve_train_model"]

src/deepagents_rl/train/eval_checkpoint.py

@@ -0,0 +1,34 @@
+"""Held-out eval used BEFORE every checkpoint save."""
+from __future__ import annotations
+
+from typing import Any, Callable
+
+from deepagents_rl.rewards.unit_test import compute_reward, grade_completion
+
+
+def eval_on_problems(
+    generate_fn: Callable[[str], str],
+    problems: list[dict[str, Any]],
+) -> list[dict[str, Any]]:
+    """
+    generate_fn(prompt) -> completion text.
+    Returns per-problem grade/reward rows for Langfuse + W&B.
+    """
+    rows: list[dict[str, Any]] = []
+    for p in problems:
+        prompt = p["prompt"]
+        completion = generate_fn(prompt)
+        grade = grade_completion(completion, list(p["hidden_tests"]))
+        reward = compute_reward(grade, format_ok=True, ran_execute=False)
+        rows.append(
+            {
+                "task_id": p["task_id"],
+                "prompt": prompt,
+                "completion": completion,
+                "pass_frac": reward["pass_frac"],
+                "reward_total": reward["reward_total"],
+                "passed": grade["passed"],
+                "total": grade["total"],
+            }
+        )
+    return rows

src/deepagents_rl/train/resolve_model.py

@@ -0,0 +1,83 @@
+"""Resolve a Kaggle-fitting HF base model (never GGUF)."""
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+DEFAULT_CANDIDATES = [
+    "Qwen/Qwen2.5-Coder-1.5B-Instruct",
+    "Qwen/Qwen2.5-Coder-0.5B-Instruct",
+    "Qwen/Qwen2.5-1.5B-Instruct",
+    "HuggingFaceTB/SmolLM2-1.7B-Instruct",
+    "google/gemma-2-2b-it",
+]
+
+
+def _probe(repo_id: str, token: str | None) -> bool:
+    from huggingface_hub import HfApi
+
+    api = HfApi(token=token)
+    try:
+        info = api.model_info(repo_id)
+    except Exception:
+        return False
+    siblings = [s.rfilename for s in (info.siblings or [])]
+    # Reject pure-GGUF repos
+    has_weight = any(
+        f.endswith((".safetensors", ".bin", ".pt"))
+        or f == "pytorch_model.bin.index.json"
+        or f.endswith(".safetensors.index.json")
+        for f in siblings
+    )
+    only_gguf = siblings and all(
+        f.endswith(".gguf") or f.endswith(".md") or "license" in f.lower() or f.startswith(".")
+        for f in siblings
+    )
+    if only_gguf and not has_weight:
+        return False
+    return has_weight or True  # allow repos that list via remote code
+
+
+def resolve_train_model(
+    preferred: list[str] | None = None,
+    *,
+    override: str | None = None,
+    token: str | None = None,
+    out_path: Path | None = None,
+) -> dict[str, Any]:
+    token = token or os.getenv("HF_TOKEN")
+    override = override or os.getenv("TRAIN_MODEL_ID") or None
+    candidates = []
+    if override:
+        candidates.append(override)
+    candidates.extend(preferred or DEFAULT_CANDIDATES)
+
+    seen: set[str] = set()
+    errors: list[dict[str, str]] = []
+    for repo_id in candidates:
+        if repo_id in seen:
+            continue
+        seen.add(repo_id)
+        try:
+            if _probe(repo_id, token):
+                result = {
+                    "train_model_id": repo_id,
+                    "candidates_tried": list(seen),
+                    "errors": errors,
+                }
+                if out_path:
+                    out_path.parent.mkdir(parents=True, exist_ok=True)
+                    out_path.write_text(json.dumps(result, indent=2), encoding="utf-8")
+                return result
+        except Exception as e:  # noqa: BLE001
+            errors.append({"repo_id": repo_id, "error": f"{type(e).__name__}: {e}"})
+            continue
+        errors.append({"repo_id": repo_id, "error": "probe_failed"})
+
+    raise RuntimeError(
+        "No trainable HF model resolved. Tried: "
+        + ", ".join(seen)
+        + f" errors={errors}"
+    )