    },
]

def sample_task(level: int) -> Dict[str, Any]:
    import random

    # Primitive-only mode: skip foundry/task templates entirely
    try:
        prim_only = str(os.environ.get("SELFPLAY_PRIMITIVE_ONLY", "")).strip() not in ("", "0", "false", "False")
        prim_only = prim_only or (str(os.environ.get("SELFPLAY_BLOCK_IMPORTS", "")).strip() not in ("", "0", "false", "False"))
    except Exception:
        prim_only = False

    # Prefer foundry templates when available and not in primitive-only mode
    try:
        if _FOUNDRY_REG and (not prim_only) and level >= 0:
            # Optional tag filter from env
            tag_env = str(os.environ.get("SELFPLAY_TASK_TAGS", "")).strip()
            want_tags = [s.strip() for s in tag_env.split(",") if s.strip()] if tag_env else []
            # Filter by difficulty: allow within level+2
            def _ok(t: _FT) -> bool:  # type: ignore
                diff_ok = getattr(t, "difficulty", 0) <= min(5, int(level) + 2)
                tag_ok = (not want_tags) or any((tg in getattr(t, "tags", [])) for tg in want_tags)
                return bool(diff_ok and tag_ok)
            cand = [t for t in _FOUNDRY_REG if _ok(t)]
            if not cand:
                cand = list(_FOUNDRY_REG)
            # Seeded RNG for reproducibility when SELFPLAY_SEED is set
            try:
                # Episode-scoped seed has priority for reproducibility across steps
                seed_env = os.environ.get("SELFPLAY_SEED_EP")
                if seed_env is None or str(seed_env).strip() == "":
                    seed_env = os.environ.get("SELFPLAY_SEED")
                seed_val = int(seed_env) if seed_env is not None and str(seed_env).strip() != "" else None
            except Exception:
                seed_val = None
            rng = random.Random(seed_val) if seed_val is not None else random.Random()
            T = random.choice(cand)
            inputs, oracle = T.sampler(rng)
            # Base tests
            tests = _FOUNDRY_MK(inputs, oracle)
            # Metamorphic expansion
            try:
                for mr in getattr(T, "metamorphics", []) or []:


from __future__ import annotations

from .common_imports import *
from .generation import generate_text
from .healing import heal_code
from .sandbox import _worker_run_one_subproc

def _extract_code_block(text: str) -> str | None:
    s = text or ""
    if s.count("```") >= 2:
        try:
            parts = s.split("```")
            inner = max((seg for seg in parts[1::2]), key=lambda u: len(u), default="")
            if inner:
                lines = inner.splitlines()
                if lines and re.match(r"^[a-zA-Z0-9_+-]+$", lines[0].strip()):
                    inner = "\n".join(lines[1:])
                s = inner.strip()
        except Exception:
            s = s
    # Find def solve block
    lines = s.splitlines()
    idx = -1
    for i, ln in enumerate(lines):
        if re.match(r"^\s*def\s+solve\s*\(.*\):\s*$", ln):
            idx = i; break
    if idx == -1:
        return None
    body = [lines[idx]]
    for ln in lines[idx+1:]:
        if not ln.strip():
            body.append(ln); continue
        if re.match(r"^\S", ln):
            break
        body.append(ln)
    code = "\n".join(body).rstrip()
    return code if code.strip() else None

def model_generate_task(tok, model, level: int, target_sr: Tuple[float, float], log_fn: Callable[[Dict[str, Any]], None] | None = None) -> Dict[str, Any]:
    """Generate a model-driven task: ask for a reference function, derive tests via sandbox.
    Returns a spec dict: {name, signature, tests}.
    """
    # Guidance
    sr_lo, sr_hi = float(target_sr[0]), float(target_sr[1])
    guidance = (
        "Output one fenced Python code block containing exactly one function.\n"
        "The code block should be:\n"
        "```python\n"
        "def solve(n:int)->int:\n"
        "# implementation\n"
        "```\n"
        f"Aim for difficulty with success rate in [{sr_lo:.2f},{sr_hi:.2f}]."
    )
    txt = generate_text(tok, model, guidance, max_new_tokens=256)
    if log_fn:
        try:
            log_fn({"taskgen_raw": (txt[:800] if isinstance(txt, str) else "")})
            if isinstance(txt, str):
                stripped = (txt or "").strip()
                head_line = stripped.splitlines()[0] if stripped else ""
                log_fn({"taskgen_summary": {"head": head_line}})
        except Exception:
            pass
    code = _extract_code_block(str(txt))
    if not code:
        healed = heal_code(tok, model, guidance, str(txt), {"err": "no_code"}, max_new_tokens=256)
        if log_fn:
            try:
                log_fn({"taskgen_healed_raw": (healed[:800] if isinstance(healed, str) else "")})
            except Exception:
                pass
        code = _extract_code_block(str(healed))
    if not code:
        raise ValueError("no_ref_code")

    # Optional cross-validation
    try:
        do_cv = bool(int(os.environ.get("SELFPLAY_REF_CV", "1") or 1))
    except Exception:
        do_cv = True
    if do_cv:
        txt2 = generate_text(tok, model, guidance, max_new_tokens=256)
        code2 = _extract_code_block(str(txt2))
        if not code2:
            healed2 = heal_code(tok, model, guidance, str(txt2), {"err": "no_code"}, max_new_tokens=256)
            code2 = _extract_code_block(str(healed2))
        if code2:
            try:
                test_timeout = float(os.environ.get("SELFPLAY_TEST_TIMEOUT_SEC", "2.5") or 2.5)
            except Exception:
                test_timeout = 2.5
            K = 20
            agree = 0
            tried = 0
            import random as _r  # type: ignore
            _r.seed(level + 17)
            while tried < K:
                n = int(_r.randint(-max(10, 50 * max(1, level)), max(10, 50 * max(1, level))))
                okA, outA = _worker_run_one_subproc(code, "def solve(n:int)->int:", [n], timeout_sec=test_timeout)
                okB, outB = _worker_run_one_subproc(code2, "def solve(n:int)->int:", [n], timeout_sec=test_timeout)
                if (not okA) or (not okB):
                    tried += 1
                    continue
                agree += 1 if outA == outB else 0
                tried += 1
            try:
                thr = float(os.environ.get("SELFPLAY_REF_CV_THRESH", "0.9") or 0.9)
            except Exception:
                thr = 0.9
            if tried > 0 and (float(agree) / float(tried)) < thr:
                raise ValueError("ref_disagree")

    sig = "def solve(n:int)->int:"

    # Name inference: simple fallback to hash if not provided
    try:
        import hashlib as _hl  # type: ignore
        name = f"gen_{_hl.sha1(code.encode('utf-8')).hexdigest()[:8]}"
    except Exception:
        name = "gen_task"

    # Build tests via sandbox execution
    def _rand_n(level: int) -> int:
        span = max(10, 50 * max(1, level))
        return int(__import__("random").randint(-span, span))

    try:
        test_timeout = float(os.environ.get("SELFPLAY_TEST_TIMEOUT_SEC", "2.5") or 2.5)
    except Exception:
        test_timeout = 2.5
    tests: List[Dict[str, Any]] = []
    num_cases = max(12, min(48, 12 + 2 * max(1, int(level))))
    tried: set[int] = set()
    while len(tests) < num_cases and len(tried) < (num_cases * 5):
        n = _rand_n(level)
        if n in tried:
            continue
        tried.add(n)
        ok_ref1, out_ref1 = _worker_run_one_subproc(code, sig, [n], timeout_sec=test_timeout)
        ok_ref2, out_ref2 = _worker_run_one_subproc(code, sig, [n], timeout_sec=test_timeout) if ok_ref1 else (ok_ref1, out_ref1)
        if ok_ref1 and ok_ref2 and out_ref1 == out_ref2:
            try:
                outn = int(out_ref1) if isinstance(out_ref1, (int, bool)) else int(out_ref1) if str(out_ref1).isdigit() else 0
            except Exception:
                outn = 0
            tests.append({"inp": [int(n)], "out": outn})

    # Identity guard for trivial refs
    try:
        guard_on = bool(int(os.environ.get("SELFPLAY_REF_IDENTITY_GUARD", "1") or 1))
    except Exception:
        guard_on = True
    if guard_on and sig.endswith("int):"):
        if tests:
            eq = int(sum(1 for t in tests if isinstance(t.get("inp"), list) and len(t["inp"]) == 1 and int(t["inp"][0]) == int(t.get("out", 1_000_000))))
            ratio = float(eq) / float(len(tests))
            try:
                thr_eq = float(os.environ.get("SELFPLAY_REF_IDENTITY_THRESH", "0.8") or 0.8)
            except Exception:
                thr_eq = 0.8
            if ratio > thr_eq:
                raise ValueError("ref_trivial_identity")

    if not tests:
        raise ValueError("no_ref_tests")

    spec = {"name": name, "signature": sig, "tests": tests}
    if log_fn:
        try:
            log_fn({"taskgen_parsed": {"name": name, "num_tests": len(spec.get("tests", []))}, "taskgen_entry": {"name": name, "sig": sig, "has_gen": True, "has_ref": True}})
        except Exception:
            pass
    return spec


from __future__ import annotations

from .common_imports import *
from .generation import generate_text
from .healing import heal_code
from .sandbox import _worker_run_one_subproc

def _extract_code_block(text: str) -> str | None:
    s = text or ""
    if s.count("```") >= 2:
        try:
            parts = s.split("```")
            inner = max((seg for seg in parts[1::2]), key=lambda u: len(u), default="")
            if inner:
                lines = inner.splitlines()
                if lines and re.match(r"^[a-zA-Z0-9_+-]+$", lines[0].strip()):
                    inner = "\n".join(lines[1:])
                s = inner.strip()
        except Exception:
            s = s
    # Find def solve block
    lines = s.splitlines()
    idx = -1
    for i, ln in enumerate(lines):
        if re.match(r"^\s*def\s+solve\s*\(.*\):\s*$", ln):
            idx = i; break
    if idx == -1:
        return None
    body = [lines[idx]]
    for ln in lines[idx+1:]:
        if not ln.strip():
            body.append(ln); continue
        if re.match(r"^\S", ln):
            break
        body.append(ln)
    code = "\n".join(body).rstrip()
    return code if code.strip() else None

def model_generate_task(tok, model, level: int, target_sr: Tuple[float, float], log_fn: Callable[[Dict[str, Any]], None] | None = None) -> Dict[str, Any]:
    """Generate a model-driven task: ask for a reference function, derive tests via sandbox.
    Returns a spec dict: {name, signature, tests}.
    """
    # Guidance
    sr_lo, sr_hi = float(target_sr[0]), float(target_sr[1])
    guidance = (
        "Output one fenced Python code block containing exactly one function.\n"
        "The code block should be:\n"
        "```python\n"
        "def solve(n:int)->int:\n"
        "# implementation\n"
        "```\n"
        f"Aim for difficulty with success rate in [{sr_lo:.2f},{sr_hi:.2f}]."
    )
    txt = generate_text(tok, model, guidance, max_new_tokens=256)
    if log_fn:
        try:


            s = s
    # Find def solve block
    lines = s.splitlines()
    idx = -1
    for i, ln in enumerate(lines):
        if re.match(r"^\s*def\s+solve\s*\(.*\):\s*$", ln):
            idx = i; break
    if idx == -1:
        return None
    body = [lines[idx]]
    for ln in lines[idx+1:]:
        if not ln.strip():
            body.append(ln); continue
        if re.match(r"^\S", ln):
            break
        body.append(ln)
    code = "\n".join(body).rstrip()
    return code if code.strip() else None

def model_generate_task(tok, model, level: int, target_sr: Tuple[float, float], log_fn: Callable[[Dict[str, Any]], None] | None = None) -> Dict[str, Any]:
    """Generate a model-driven task: ask for a reference function, derive tests via sandbox.
    Returns a spec dict: {name, signature, tests}.
    """
    # Guidance
    sr_lo, sr_hi = float(target_sr[0]), float(target_sr[1])
    guidance = (
        "Output one fenced Python code block containing exactly one function.\n"
        "The code block should be:\n"
        "```python\n"
        "def solve(n:int)->int:\n"
        "# implementation\n"
        "```\n"
        f"Aim for difficulty with success rate in [{sr_lo:.2f},{sr_hi:.2f}]."
    )
    txt = generate_text(tok, model, guidance, max_new_tokens=256)
    if log_fn:
        try:
            log_fn({"taskgen_raw": (txt[:800] if isinstance(txt, str) else "")})
            if isinstance(txt, str):
                stripped = (txt or "").strip()
                head_line = stripped.splitlines()[0] if stripped else ""
                log_fn({"taskgen_summary": {"head": head_line}})
        except Exception:
            pass
    code = _extract_code_block(str(txt))
    if not code:
        healed = heal_code(tok, model, guidance, str(txt), {"err": "no_code"}, max_new_tokens=256)
        if log_fn:
            try:
                log_fn({"taskgen_healed_raw": (healed[:800] if isinstance(healed, str) else "")})
            except Exception:
                pass
        code = _extract_code_block(str(healed))
    if not code:
        raise ValueError("no_ref_code")

    # Optional cross-validation
    try:
        do_cv = bool(int(os.environ.get("SELFPLAY_REF_CV", "1") or 1))
    except Exception:
        do_cv = True
    if do_cv:
        txt2 = generate_text(tok, model, guidance, max_new_tokens=256)
        code2 = _extract_code_block(str(txt2))
        if not code2:
            healed2 = heal_code(tok, model, guidance, str(txt2), {"err": "no_code"}, max_new_tokens=256)
            code2 = _extract_code_block(str(healed2))
        if code2:
            try:
                test_timeout = float(os.environ.get("SELFPLAY_TEST_TIMEOUT_SEC", "2.5") or 2.5)
            except Exception:
                test_timeout = 2.5
            K = 20
            agree = 0
            tried = 0
            import random as _r  # type: ignore
            _r.seed(level + 17)
            while tried < K:
                n = int(_r.randint(-max(10, 50 * max(1, level)), max(10, 50 * max(1, level))))
                okA, outA = _worker_run_one_subproc(code, "def solve(n:int)->int:", [n], timeout_sec=test_timeout)
                okB, outB = _worker_run_one_subproc(code2, "def solve(n:int)->int:", [n], timeout_sec=test_timeout)
                if (not okA) or (not okB):
                    tried += 1
                    continue
                agree += 1 if outA == outB else 0
                tried += 1
            try:
                thr = float(os.environ.get("SELFPLAY_REF_CV_THRESH", "0.9") or 0.9)
            except Exception:
                thr = 0.9
            if tried > 0 and (float(agree) / float(tried)) < thr:
                raise ValueError("ref_disagree")

    sig = "def solve(n:int)->int:"

    # Name inference: simple fallback to hash if not provided
    try:
        import hashlib as _hl  # type: ignore
        name = f"gen_{_hl.sha1(code.encode('utf-8')).hexdigest()[:8]}"
    except Exception:
        name = "gen_task"

    # Build tests via sandbox execution
    def _rand_n(level: int) -> int:
        span = max(10, 50 * max(1, level))
        return int(__import__("random").randint(-span, span))

    try:
        test_timeout = float(os.environ.get("SELFPLAY_TEST_TIMEOUT_SEC", "2.5") or 2.5)
    except Exception:
        test_timeout = 2.5
    tests: List[Dict[str, Any]] = []
    num_cases = max(12, min(48, 12 + 2 * max(1, int(level))))
    tried: set[int] = set()
    while len(tests) < num_cases and len(tried) < (num_cases * 5):
        n = _rand_n(level)
        if n in tried:
            continue
        tried.add(n)
        ok_ref1, out_ref1 = _worker_run_one_subproc(code, sig, [n], timeout_sec=test_timeout)
        ok_ref2, out_ref2 = _worker_run_one_subproc(code, sig, [n], timeout_sec=test_timeout) if ok_ref1 else (ok_ref1, out_ref1)
        if ok_ref1 and ok_ref2 and out_ref1 == out_ref2:
            try:
                outn = int(out_ref1) if isinstance(out_ref1, (int, bool)) else int(out_ref1) if str(out_ref1).isdigit() else 0
            except Exception:
                outn = 0
            tests.append({"inp": [int(n)], "out": outn})

    # Identity guard for trivial refs
    try:
        guard_on = bool(int(os.environ.get("SELFPLAY_REF_IDENTITY_GUARD", "1") or 1))
    except Exception:
        guard_on = True
    if guard_on and sig.endswith("int):"):
        if tests:
            eq = int(sum(1 for t in tests if isinstance(t.get("inp"), list) and len(t["inp"]) == 1 and int(t["inp"][0]) == int(t.get("out", 1_000_000))))
            ratio = float(eq) / float(len(tests))
            try:
                thr_eq = float(os.environ.get("SELFPLAY_REF_IDENTITY_THRESH", "0.8") or 0.8)
            except Exception:
                thr_eq = 0.8
            if ratio > thr_eq:
                raise ValueError("ref_trivial_identity")

    if not tests:
        raise ValueError("no_ref_tests")

    spec = {"name": name, "signature": sig, "tests": tests}
    if log_fn:
        try:
            log_fn({"taskgen_parsed": {"name": name, "num_tests": len(spec.get("tests", []))}, "taskgen_entry": {"name": name, "sig": sig, "has_gen": True, "has_ref": True}})
        except Exception:
            pass
    return spec


                    continue
                agree += 1 if outA == outB else 0
                tried += 1
            try:
                thr = float(os.environ.get("SELFPLAY_REF_CV_THRESH", "0.9") or 0.9)
            except Exception:
                thr = 0.9
            if tried > 0 and (float(agree) / float(tried)) < thr:
                raise ValueError("ref_disagree")

    sig = "def solve(n:int)->int:"

    # Name inference: simple fallback to hash if not provided
    try:
        import hashlib as _hl  # type: ignore
        name = f"gen_{_hl.sha1(code.encode('utf-8')).hexdigest()[:8]}"
    except Exception:
        name = "gen_task"

    # Build tests via sandbox execution
    def _rand_n(level: int) -> int:
        span = max(10, 50 * max(1, level))
        return int(__import__("random").randint(-span, span))

    try:
        test_timeout = float(os.environ.get("SELFPLAY_TEST_TIMEOUT_SEC", "2.5") or 2.5)
    except Exception:
        test_timeout = 2.5
    tests: List[Dict[str, Any]] = []
    num_cases = max(12, min(48, 12 + 2 * max(1, int(level))))
    tried: set[int] = set()
    while len(tests) < num_cases and len(tried) < (num_cases * 5):
        n = _rand_n(level)
        if n in tried:
            continue
        tried.add(n)
        ok_ref1, out_ref1 = _worker_run_one_subproc(code, sig, [n], timeout_sec=test_timeout)
        ok_ref2, out_ref2 = _worker_run_one_subproc(code, sig, [n], timeout_sec=test_timeout) if ok_ref1 else (ok_ref1, out_ref1)
        if ok_ref1 and ok_ref2 and out_ref1 == out_ref2:
            try:
                outn = int(out_ref1) if isinstance(out_ref1, (int, bool)) else int(out_ref1) if str(out_ref1).isdigit() else 0
            except Exception:
                outn = 0


from .common_imports import *

def _loraify_linear(module: torch.nn.Module, r: int = 8, alpha: int = 16) -> None:
    if not isinstance(module, torch.nn.Linear):
        return
    # Idempotent: if already has LoRA params, skip
    try:
        if hasattr(module, "A") and hasattr(module, "B"):
            return
    except Exception:
        pass
    w = module.weight
    in_f, out_f = w.shape[1], w.shape[0]
    # Create and register LoRA params/buffers on the same device/dtype as base weight
    A = torch.nn.Parameter(torch.zeros(r, in_f, device=w.device, dtype=w.dtype))
    B = torch.nn.Parameter(torch.zeros(out_f, r, device=w.device, dtype=w.dtype))
    module.register_parameter("A", A)
    module.register_parameter("B", B)
    # Register scaling as a buffer to ensure device moves with the module
    module.register_buffer("lora_scaling", torch.tensor(float(alpha / r), dtype=w.dtype, device=w.device))
    module.weight.requires_grad_(False)
    # Mark that LoRA is enabled for this module so we can toggle if needed
    try:
        module._lora_enabled = True  # type: ignore[attr-defined]
    except Exception:
        pass
    # Ensure a pre-hook keeps buffers on the right device/dtype
    def _prehook(_mod, _inp):  # type: ignore
        try:
            if hasattr(_mod, "lora_scaling") and _mod.lora_scaling.device != _mod.weight.device:  # type: ignore[attr-defined]
                _mod.lora_scaling = _mod.lora_scaling.to(_mod.weight.device)  # type: ignore[attr-defined]
        except Exception:
            pass
    try:
        if not hasattr(module, "_lora_hook_added"):
            module.register_forward_pre_hook(_prehook)  # type: ignore[arg-type]
            module._lora_hook_added = True  # type: ignore[attr-defined]
    except Exception:
        pass
    # Single-module only
    return module

def _inject_lora(model: torch.nn.Module, r: int = 8, alpha: int = 16) -> torch.nn.Module:
    # Avoid injecting LoRA into obvious output heads to reduce risk with tied weights
    skip_name_substrings = ("lm_head", "embed_out", "output", "score")
    # Build a reference count map for parameters to detect shared/tied weights
    try:
        param_ref_counts: Dict[int, int] = Counter(int(id(p)) for p in model.parameters())  # type: ignore[assignment]
    except Exception:
        param_ref_counts = {}
    for name, m in model.named_modules():
        if not isinstance(m, torch.nn.Linear):
            continue
        if any(s in name for s in skip_name_substrings):
            continue
        try:
            # Skip modules whose weight is shared/tied
            if hasattr(m, "weight"):
                rc = int(param_ref_counts.get(int(id(m.weight)), 1))
                if rc > 1:
                    continue
        except Exception:
            pass
        _loraify_linear(m, r=r, alpha=alpha)
    return model

def _forward_with_lora_delta(module: torch.nn.Linear, orig_forward, *args, **kwargs):
    """Preserve original forward signature; add LoRA delta to base output.

    The input tensor is expected as the first positional argument for Linear.
    """
    y = orig_forward(*args, **kwargs)
    try:
        if getattr(module, "_lora_enabled", False) and hasattr(module, "A") and hasattr(module, "B"):
            # Extract input tensor from args
            x = args[0] if len(args) > 0 else None
            if isinstance(x, torch.Tensor):
                # Prefer buffer scaling when present
                _scale = getattr(module, "lora_scaling", None)
                scale = float(_scale.detach().item()) if torch.is_tensor(_scale) else float(getattr(module, "scaling", 1.0))
                delta = scale * torch.nn.functional.linear(
                    torch.nn.functional.linear(x, module.A), module.B
                )
                y = y + delta
    except Exception:
        # If anything goes wrong, fall back to base output
        return y
    return y

def _set_lora_enabled(model: torch.nn.Module, enabled: bool) -> None:
    for _name, _m in model.named_modules():
        try:
            if isinstance(_m, torch.nn.Linear) and hasattr(_m, "A"):
                _m._lora_enabled = bool(enabled)  # type: ignore[attr-defined]
        except Exception:
            pass

def _patch_forwards(model: torch.nn.Module) -> torch.nn.Module:
    # Skip patching obvious output heads
    skip_name_substrings = ("lm_head", "embed_out", "output", "score")
    for name, m in model.named_modules():
        if isinstance(m, torch.nn.Linear) and hasattr(m, "A"):
            if any(s in name for s in skip_name_substrings):
                continue
            try:
                if not hasattr(m, "_orig_fwd"):
                    m._orig_fwd = m.forward  # type: ignore[attr-defined]
                orig = m._orig_fwd  # type: ignore[attr-defined]
                def _wrapped_forward(*args, _m=m, _orig=orig, **kwargs):
                    return _forward_with_lora_delta(_m, _orig, *args, **kwargs)
                m.forward = _wrapped_forward  # type: ignore[assignment]
            except Exception:
                pass
    return model


from .common_imports import *

def _loraify_linear(module: torch.nn.Module, r: int = 8, alpha: int = 16) -> None:
    if not isinstance(module, torch.nn.Linear):
        return
    # Idempotent: if already has LoRA params, skip
    try:
        if hasattr(module, "A") and hasattr(module, "B"):
            return
    except Exception:
        pass
    w = module.weight
    in_f, out_f = w.shape[1], w.shape[0]
    # Create and register LoRA params/buffers on the same device/dtype as base weight
    A = torch.nn.Parameter(torch.zeros(r, in_f, device=w.device, dtype=w.dtype))
    B = torch.nn.Parameter(torch.zeros(out_f, r, device=w.device, dtype=w.dtype))
    module.register_parameter("A", A)
    module.register_parameter("B", B)
    # Register scaling as a buffer to ensure device moves with the module
    module.register_buffer("lora_scaling", torch.tensor(float(alpha / r), dtype=w.dtype, device=w.device))
    module.weight.requires_grad_(False)
    # Mark that LoRA is enabled for this module so we can toggle if needed
    try:
        module._lora_enabled = True  # type: ignore[attr-defined]
    except Exception:
        pass
    # Ensure a pre-hook keeps buffers on the right device/dtype
    def _prehook(_mod, _inp):  # type: ignore
        try:
            if hasattr(_mod, "lora_scaling") and _mod.lora_scaling.device != _mod.weight.device:  # type: ignore[attr-defined]
                _mod.lora_scaling = _mod.lora_scaling.to(_mod.weight.device)  # type: ignore[attr-defined]
        except Exception:
            pass
    try:
        if not hasattr(module, "_lora_hook_added"):
            module.register_forward_pre_hook(_prehook)  # type: ignore[arg-type]
            module._lora_hook_added = True  # type: ignore[attr-defined]
    except Exception:
        pass
    # Single-module only
    return module

def _inject_lora(model: torch.nn.Module, r: int = 8, alpha: int = 16) -> torch.nn.Module:
    # Avoid injecting LoRA into obvious output heads to reduce risk with tied weights
    skip_name_substrings = ("lm_head", "embed_out", "output", "score")
    # Build a reference count map for parameters to detect shared/tied weights
    try:
        param_ref_counts: Dict[int, int] = Counter(int(id(p)) for p in model.parameters())  # type: ignore[assignment]
    except Exception:
        param_ref_counts = {}
    for name, m in model.named_modules():
        if not isinstance(m, torch.nn.Linear):
            continue
        if any(s in name for s in skip_name_substrings):
            continue
        try:
            # Skip modules whose weight is shared/tied
            if hasattr(m, "weight"):
                rc = int(param_ref_counts.get(int(id(m.weight)), 1))
                if rc > 1:
                    continue


    try:
        module._lora_enabled = True  # type: ignore[attr-defined]
    except Exception:
        pass
    # Ensure a pre-hook keeps buffers on the right device/dtype
    def _prehook(_mod, _inp):  # type: ignore
        try:
            if hasattr(_mod, "lora_scaling") and _mod.lora_scaling.device != _mod.weight.device:  # type: ignore[attr-defined]
                _mod.lora_scaling = _mod.lora_scaling.to(_mod.weight.device)  # type: ignore[attr-defined]
        except Exception:
            pass
    try:
        if not hasattr(module, "_lora_hook_added"):
            module.register_forward_pre_hook(_prehook)  # type: ignore[arg-type]
            module._lora_hook_added = True  # type: ignore[attr-defined]
    except Exception:
        pass
    # Single-module only
    return module

def _inject_lora(model: torch.nn.Module, r: int = 8, alpha: int = 16) -> torch.nn.Module:
    # Avoid injecting LoRA into obvious output heads to reduce risk with tied weights
    skip_name_substrings = ("lm_head", "embed_out", "output", "score")
    # Build a reference count map for parameters to detect shared/tied weights
    try:
        param_ref_counts: Dict[int, int] = Counter(int(id(p)) for p in model.parameters())  # type: ignore[assignment]
    except Exception:
        param_ref_counts = {}
    for name, m in model.named_modules():
        if not isinstance(m, torch.nn.Linear):
            continue
        if any(s in name for s in skip_name_substrings):
            continue
        try:
            # Skip modules whose weight is shared/tied
            if hasattr(m, "weight"):
                rc = int(param_ref_counts.get(int(id(m.weight)), 1))
                if rc > 1:
                    continue
        except Exception:
            pass
        _loraify_linear(m, r=r, alpha=alpha)
    return model

def _forward_with_lora_delta(module: torch.nn.Linear, orig_forward, *args, **kwargs):
    """Preserve original forward signature; add LoRA delta to base output.

    The input tensor is expected as the first positional argument for Linear.
    """
    y = orig_forward(*args, **kwargs)
    try:
        if getattr(module, "_lora_enabled", False) and hasattr(module, "A") and hasattr(module, "B"):
            # Extract input tensor from args
            x = args[0] if len(args) > 0 else None
            if isinstance(x, torch.Tensor):
                # Prefer buffer scaling when present
                _scale = getattr(module, "lora_scaling", None)
                scale = float(_scale.detach().item()) if torch.is_tensor(_scale) else float(getattr(module, "scaling", 1.0))
                delta = scale * torch.nn.functional.linear(
                    torch.nn.functional.linear(x, module.A), module.B
                )
                y = y + delta
    except Exception:


    try:
        param_ref_counts: Dict[int, int] = Counter(int(id(p)) for p in model.parameters())  # type: ignore[assignment]
    except Exception:
        param_ref_counts = {}
    for name, m in model.named_modules():
        if not isinstance(m, torch.nn.Linear):
            continue
        if any(s in name for s in skip_name_substrings):
            continue
        try:
            # Skip modules whose weight is shared/tied
            if hasattr(m, "weight"):
                rc = int(param_ref_counts.get(int(id(m.weight)), 1))
                if rc > 1:
                    continue
        except Exception:
            pass
        _loraify_linear(m, r=r, alpha=alpha)
    return model

def _forward_with_lora_delta(module: torch.nn.Linear, orig_forward, *args, **kwargs):
    """Preserve original forward signature; add LoRA delta to base output.

    The input tensor is expected as the first positional argument for Linear.
    """
    y = orig_forward(*args, **kwargs)
    try:
        if getattr(module, "_lora_enabled", False) and hasattr(module, "A") and hasattr(module, "B"):
            # Extract input tensor from args
            x = args[0] if len(args) > 0 else None
            if isinstance(x, torch.Tensor):
                # Prefer buffer scaling when present
                _scale = getattr(module, "lora_scaling", None)
                scale = float(_scale.detach().item()) if torch.is_tensor(_scale) else float(getattr(module, "scaling", 1.0))
                delta = scale * torch.nn.functional.linear(
                    torch.nn.functional.linear(x, module.A), module.B
                )
                y = y + delta
    except Exception:
        # If anything goes wrong, fall back to base output
        return y
    return y

def _set_lora_enabled(model: torch.nn.Module, enabled: bool) -> None:
    for _name, _m in model.named_modules():
        try:
            if isinstance(_m, torch.nn.Linear) and hasattr(_m, "A"):
                _m._lora_enabled = bool(enabled)  # type: ignore[attr-defined]
        except Exception:
            pass

def _patch_forwards(model: torch.nn.Module) -> torch.nn.Module:
    # Skip patching obvious output heads
    skip_name_substrings = ("lm_head", "embed_out", "output", "score")
    for name, m in model.named_modules():
        if isinstance(m, torch.nn.Linear) and hasattr(m, "A"):
            if any(s in name for s in skip_name_substrings):
                continue
            try:
                if not hasattr(m, "_orig_fwd"):
                    m._orig_fwd = m.forward  # type: ignore[attr-defined]
                orig = m._orig_fwd  # type: ignore[attr-defined]


    The input tensor is expected as the first positional argument for Linear.
    """
    y = orig_forward(*args, **kwargs)
    try:
        if getattr(module, "_lora_enabled", False) and hasattr(module, "A") and hasattr(module, "B"):
            # Extract input tensor from args
            x = args[0] if len(args) > 0 else None
            if isinstance(x, torch.Tensor):
                # Prefer buffer scaling when present
                _scale = getattr(module, "lora_scaling", None)
                scale = float(_scale.detach().item()) if torch.is_tensor(_scale) else float(getattr(module, "scaling", 1.0))
                delta = scale * torch.nn.functional.linear(
                    torch.nn.functional.linear(x, module.A), module.B
                )
                y = y + delta
    except Exception:
        # If anything goes wrong, fall back to base output
        return y
    return y

def _set_lora_enabled(model: torch.nn.Module, enabled: bool) -> None:
    for _name, _m in model.named_modules():
        try:
            if isinstance(_m, torch.nn.Linear) and hasattr(_m, "A"):
                _m._lora_enabled = bool(enabled)  # type: ignore[attr-defined]
        except Exception:
            pass

def _patch_forwards(model: torch.nn.Module) -> torch.nn.Module:
    # Skip patching obvious output heads
    skip_name_substrings = ("lm_head", "embed_out", "output", "score")
    for name, m in model.named_modules():
        if isinstance(m, torch.nn.Linear) and hasattr(m, "A"):
            if any(s in name for s in skip_name_substrings):
                continue
            try:
                if not hasattr(m, "_orig_fwd"):
                    m._orig_fwd = m.forward  # type: ignore[attr-defined]
                orig = m._orig_fwd  # type: ignore[attr-defined]
                def _wrapped_forward(*args, _m=m, _orig=orig, **kwargs):
                    return _forward_with_lora_delta(_m, _orig, *args, **kwargs)
                m.forward = _wrapped_forward  # type: ignore[assignment]
            except Exception:
                pass
    return model


                # Prefer buffer scaling when present
                _scale = getattr(module, "lora_scaling", None)
                scale = float(_scale.detach().item()) if torch.is_tensor(_scale) else float(getattr(module, "scaling", 1.0))
                delta = scale * torch.nn.functional.linear(
                    torch.nn.functional.linear(x, module.A), module.B
                )
                y = y + delta
    except Exception:
        # If anything goes wrong, fall back to base output
        return y
    return y

def _set_lora_enabled(model: torch.nn.Module, enabled: bool) -> None:
    for _name, _m in model.named_modules():
        try:
            if isinstance(_m, torch.nn.Linear) and hasattr(_m, "A"):
                _m._lora_enabled = bool(enabled)  # type: ignore[attr-defined]
        except Exception:
            pass

def _patch_forwards(model: torch.nn.Module) -> torch.nn.Module:
    # Skip patching obvious output heads
    skip_name_substrings = ("lm_head", "embed_out", "output", "score")
    for name, m in model.named_modules():
        if isinstance(m, torch.nn.Linear) and hasattr(m, "A"):
            if any(s in name for s in skip_name_substrings):
                continue
            try:
                if not hasattr(m, "_orig_fwd"):
                    m._orig_fwd = m.forward  # type: ignore[attr-defined]
                orig = m._orig_fwd  # type: ignore[attr-defined]
                def _wrapped_forward(*args, _m=m, _orig=orig, **kwargs):
                    return _forward_with_lora_delta(_m, _orig, *args, **kwargs)
                m.forward = _wrapped_forward  # type: ignore[assignment]
            except Exception:
                pass
    return model


        if hasattr(module, "A") and hasattr(module, "B"):
            return
    except Exception:
        pass
    w = module.weight
    in_f, out_f = w.shape[1], w.shape[0]
    # Create and register LoRA params/buffers on the same device/dtype as base weight
    A = torch.nn.Parameter(torch.zeros(r, in_f, device=w.device, dtype=w.dtype))
    B = torch.nn.Parameter(torch.zeros(out_f, r, device=w.device, dtype=w.dtype))
    module.register_parameter("A", A)
    module.register_parameter("B", B)
    # Register scaling as a buffer to ensure device moves with the module
    module.register_buffer("lora_scaling", torch.tensor(float(alpha / r), dtype=w.dtype, device=w.device))
    module.weight.requires_grad_(False)
    # Mark that LoRA is enabled for this module so we can toggle if needed
    try:
        module._lora_enabled = True  # type: ignore[attr-defined]
    except Exception:
        pass
    # Ensure a pre-hook keeps buffers on the right device/dtype
    def _prehook(_mod, _inp):  # type: ignore
        try:
            if hasattr(_mod, "lora_scaling") and _mod.lora_scaling.device != _mod.weight.device:  # type: ignore[attr-defined]
                _mod.lora_scaling = _mod.lora_scaling.to(_mod.weight.device)  # type: ignore[attr-defined]
        except Exception:
            pass
    try:
        if not hasattr(module, "_lora_hook_added"):
            module.register_forward_pre_hook(_prehook)  # type: ignore[arg-type]
            module._lora_hook_added = True  # type: ignore[attr-defined]
    except Exception:
        pass
    # Single-module only
    return module

def _inject_lora(model: torch.nn.Module, r: int = 8, alpha: int = 16) -> torch.nn.Module:
    # Avoid injecting LoRA into obvious output heads to reduce risk with tied weights
    skip_name_substrings = ("lm_head", "embed_out", "output", "score")
    # Build a reference count map for parameters to detect shared/tied weights
    try:
        param_ref_counts: Dict[int, int] = Counter(int(id(p)) for p in model.parameters())  # type: ignore[assignment]
    except Exception:
        param_ref_counts = {}
    for name, m in model.named_modules():
        if not isinstance(m, torch.nn.Linear):
            continue


def _set_lora_enabled(model: torch.nn.Module, enabled: bool) -> None:
    for _name, _m in model.named_modules():
        try:
            if isinstance(_m, torch.nn.Linear) and hasattr(_m, "A"):
                _m._lora_enabled = bool(enabled)  # type: ignore[attr-defined]
        except Exception:
            pass

def _patch_forwards(model: torch.nn.Module) -> torch.nn.Module:
    # Skip patching obvious output heads
    skip_name_substrings = ("lm_head", "embed_out", "output", "score")
    for name, m in model.named_modules():
        if isinstance(m, torch.nn.Linear) and hasattr(m, "A"):
            if any(s in name for s in skip_name_substrings):
                continue
            try:
                if not hasattr(m, "_orig_fwd"):
                    m._orig_fwd = m.forward  # type: ignore[attr-defined]
                orig = m._orig_fwd  # type: ignore[attr-defined]
                def _wrapped_forward(*args, _m=m, _orig=orig, **kwargs):
                    return _forward_with_lora_delta(_m, _orig, *args, **kwargs)
                m.forward = _wrapped_forward  # type: ignore[assignment]
            except Exception:
                pass
    return model


#!/usr/bin/env python3
from __future__ import annotations

import argparse
import collections as C
import json
import os
import re
from typing import Dict, Any

def pct(a: int, b: int) -> float:
    return (float(a) / float(max(1, b))) * 100.0

def normalize_first_failure(text: str, max_len: int = 200) -> str:
    try:
        s = re.sub(r"\s+", " ", str(text)).strip()
        return s[:max_len]
    except Exception:
        return str(text)[:max_len]

def analyze(path: str) -> str:
    per_level: Dict[int | None, Dict[str, int]] = C.defaultdict(lambda: {"pass": 0, "total": 0, "timeouts": 0, "rejects": 0})
    reject_reasons: C.Counter[str] = C.Counter()
    first_fail: C.Counter[str] = C.Counter()
    level_changes: list[tuple[int | None, int | None]] = []
    evaluated = 0
    last_level: int | None = None

    with open(path, "r", encoding="utf-8") as f:
        for ln in f:
            try:
                ev: Dict[str, Any] = json.loads(ln)
            except Exception:
                continue
            mode = ev.get("mode")
            if mode == "arena":
                lvl = ev.get("level")
                rep = ev.get("rep")
                if isinstance(rep, dict):
                    evaluated += 1
                    per_level[lvl]["pass"] += int(rep.get("passed", 0) or 0)
                    per_level[lvl]["total"] += int(rep.get("total", 0) or 0)
                    per_level[lvl]["timeouts"] += int(rep.get("timeouts", 0) or 0)
                    ff = rep.get("first_failure")
                    if ff:
                        first_fail[normalize_first_failure(ff)] += 1
                elif ev.get("status") == "reject":
                    per_level[lvl]["rejects"] += 1
                    pr = ev.get("policy_reason") or ev.get("reason") or "reject"
                    reject_reasons[str(pr)] += 1
            elif mode == "arena_summary":
                lvl = ev.get("level")
                if last_level is None:
                    last_level = lvl
                elif lvl != last_level:
                    level_changes.append((last_level, lvl))
                    last_level = lvl

    lines: list[str] = []
    # Overall
    total_pass = sum(d["pass"] for d in per_level.values())
    total_total = sum(d["total"] for d in per_level.values())
    total_timeouts = sum(d["timeouts"] for d in per_level.values())
    total_rejects = sum(d["rejects"] for d in per_level.values())
    lines.append(f"overall: pass_rate={pct(total_pass, total_total):.1f}%  evaluated={evaluated}  rejects={total_rejects}  timeouts={total_timeouts}")
    # Per-level
    for lvl in sorted(per_level, key=lambda x: (-1 if x is None else int(x))):
        d = per_level[lvl]
        lines.append(
            f"level {lvl}: pass_rate={pct(d['pass'], d['total']):.1f}%  pass={d['pass']}/{d['total']}  rejects={d['rejects']}  timeouts={d['timeouts']}"
        )
    # Rejects
    if reject_reasons:
        lines.append("reject_reasons:")
        for k, v in reject_reasons.most_common(10):
            lines.append(f"  {k}: {v}")
    # First failures
    if first_fail:
        lines.append("first_failure_top:")
        for k, v in first_fail.most_common(10):
            lines.append(f"  {v}x: {k}")
    # Level transitions
    if level_changes:
        lines.append("level_transitions:")
        for a, b in level_changes:
            lines.append(f"  {a} -> {b}")
    return "\n".join(lines)

def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--path", default=str(os.environ.get("SELFPLAY_PROGRESS", "/data/agiattempt/agi_dw/data/selfplay/progress.jsonl")))
    args = ap.parse_args()
    print(analyze(args.path))

if __name__ == "__main__":
    main()


pct

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import collections as C
import json
import os
import re
from typing import Dict, Any

def pct(a: int, b: int) -> float:
    return (float(a) / float(max(1, b))) * 100.0

def normalize_first_failure(text: str, max_len: int = 200) -> str:
    try:
        s = re.sub(r"\s+", " ", str(text)).strip()
        return s[:max_len]
    except Exception:
        return str(text)[:max_len]

def analyze(path: str) -> str:
    per_level: Dict[int | None, Dict[str, int]] = C.defaultdict(lambda: {"pass": 0, "total": 0, "timeouts": 0, "rejects": 0})
    reject_reasons: C.Counter[str] = C.Counter()
    first_fail: C.Counter[str] = C.Counter()
    level_changes: list[tuple[int | None, int | None]] = []
    evaluated = 0
    last_level: int | None = None

    with open(path, "r", encoding="utf-8") as f:
        for ln in f:


#!/usr/bin/env python3
from __future__ import annotations

import argparse
import collections as C
import json
import os
import re
from typing import Dict, Any

def pct(a: int, b: int) -> float:
    return (float(a) / float(max(1, b))) * 100.0

def normalize_first_failure(text: str, max_len: int = 200) -> str:
    try:
        s = re.sub(r"\s+", " ", str(text)).strip()
        return s[:max_len]
    except Exception:
        return str(text)[:max_len]

def analyze(path: str) -> str:
    per_level: Dict[int | None, Dict[str, int]] = C.defaultdict(lambda: {"pass": 0, "total": 0, "timeouts": 0, "rejects": 0})
    reject_reasons: C.Counter[str] = C.Counter()
    first_fail: C.Counter[str] = C.Counter()
    level_changes: list[tuple[int | None, int | None]] = []
    evaluated = 0
    last_level: int | None = None

    with open(path, "r", encoding="utf-8") as f:
        for ln in f:
            try:
                ev: Dict[str, Any] = json.loads(ln)
            except Exception:
                continue
            mode = ev.get("mode")
            if mode == "arena":
                lvl = ev.get("level")
                rep = ev.get("rep")


analyze

import argparse
import collections as C
import json
import os
import re
from typing import Dict, Any

def pct(a: int, b: int) -> float:
    return (float(a) / float(max(1, b))) * 100.0

def normalize_first_failure(text: str, max_len: int = 200) -> str:
    try:
        s = re.sub(r"\s+", " ", str(text)).strip()
        return s[:max_len]
    except Exception:
        return str(text)[:max_len]

def analyze(path: str) -> str:
    per_level: Dict[int | None, Dict[str, int]] = C.defaultdict(lambda: {"pass": 0, "total": 0, "timeouts": 0, "rejects": 0})
    reject_reasons: C.Counter[str] = C.Counter()
    first_fail: C.Counter[str] = C.Counter()
    level_changes: list[tuple[int | None, int | None]] = []
    evaluated = 0
    last_level: int | None = None

    with open(path, "r", encoding="utf-8") as f:
        for ln in f:
            try:
                ev: Dict[str, Any] = json.loads(ln)
            except Exception:
                continue
            mode = ev.get("mode")
            if mode == "arena":
                lvl = ev.get("level")
                rep = ev.get("rep")
                if isinstance(rep, dict):
                    evaluated += 1
                    per_level[lvl]["pass"] += int(rep.get("passed", 0) or 0)
                    per_level[lvl]["total"] += int(rep.get("total", 0) or 0)
                    per_level[lvl]["timeouts"] += int(rep.get("timeouts", 0) or 0)
                    ff = rep.get("first_failure")
                    if ff:
                        first_fail[normalize_first_failure(ff)] += 1
                elif ev.get("status") == "reject":
                    per_level[lvl]["rejects"] += 1
                    pr = ev.get("policy_reason") or ev.get("reason") or "reject"
                    reject_reasons[str(pr)] += 1
            elif mode == "arena_summary":
                lvl = ev.get("level")
                if last_level is None:
                    last_level = lvl
                elif lvl != last_level:
                    level_changes.append((last_level, lvl))
                    last_level = lvl

    lines: list[str] = []
    # Overall
    total_pass = sum(d["pass"] for d in per_level.values())
    total_total = sum(d["total"] for d in per_level.values())
    total_timeouts = sum(d["timeouts"] for d in per_level.values())
    total_rejects = sum(d["rejects"] for d in per_level.values())
    lines.append(f"overall: pass_rate={pct(total_pass, total_total):.1f}%  evaluated={evaluated}  rejects={total_rejects}  timeouts={total_timeouts}")
    # Per-level
    for lvl in sorted(per_level, key=lambda x: (-1 if x is None else int(x))):
        d = per_level[lvl]
        lines.append(
            f"level {lvl}: pass_rate={pct(d['pass'], d['total']):.1f}%  pass={d['pass']}/{d['total']}  rejects={d['rejects']}  timeouts={d['timeouts']}"
        )
    # Rejects
    if reject_reasons:
        lines.append("reject_reasons:")
        for k, v in reject_reasons.most_common(10):
            lines.append(f"  {k}: {v}")
    # First failures
    if first_fail:
        lines.append("first_failure_top:")
        for k, v in first_fail.most_common(10):
            lines.append(f"  {v}x: {k}")
    # Level transitions
    if level_changes:
        lines.append("level_transitions:")
        for a, b in level_changes:
            lines.append(f"  {a} -> {b}")
    return "\n".join(lines)

def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--path", default=str(os.environ.get("SELFPLAY_PROGRESS", "/data/agiattempt/agi_dw/data/selfplay/progress.jsonl")))
    args = ap.parse_args()
    print(analyze(args.path))

if __name__ == "__main__":
    main()


main

            f"level {lvl}: pass_rate={pct(d['pass'], d['total']):.1f}%  pass={d['pass']}/{d['total']}  rejects={d['rejects']}  timeouts={d['timeouts']}"
        )
    # Rejects
    if reject_reasons:
        lines.append("reject_reasons:")
        for k, v in reject_reasons.most_common(10):
            lines.append(f"  {k}: {v}")
    # First failures
    if first_fail:
        lines.append("first_failure_top:")
        for k, v in first_fail.most_common(10):
            lines.append(f"  {v}x: {k}")
    # Level transitions
    if level_changes:
        lines.append("level_transitions:")
        for a, b in level_changes:
            lines.append(f"  {a} -> {b}")
    return "\n".join(lines)

def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--path", default=str(os.environ.get("SELFPLAY_PROGRESS", "/data/agiattempt/agi_dw/data/selfplay/progress.jsonl")))
    args = ap.parse_args()
    print(analyze(args.path))

if __name__ == "__main__":
    main()


import os
import sys
import json
from pathlib import Path

# Ensure project root is on sys.path so 'agi_dw' package imports work when run directly
try:
    ROOT = Path(__file__).resolve().parents[4]
    if str(ROOT) not in sys.path:
        sys.path.insert(0, str(ROOT))
except Exception:
    pass

from agi_dw.scripts.selfplay.modules.model_io import load_seed
from agi_dw.scripts.selfplay.modules.tasks import sample_task, render_prompt
from agi_dw.scripts.selfplay.modules.generation import sample
from agi_dw.scripts.selfplay.modules.sandbox import run_subset_tests, run_and_test
from agi_dw.scripts.selfplay.modules.healing import heal_code

def _bool_env(name: str, default: bool) -> bool:
    v = os.environ.get(name)
    if v is None:
        return bool(default)
    s = str(v).strip().lower()
    return s in ("1", "true", "yes", "y", "on")

def main() -> None:
    # Configs via env
    n_cands = int(os.environ.get("PIPE_N_SAMPLES", os.environ.get("SELFPLAY_N_SAMPLES", "3") or 3))
    max_new = int(os.environ.get("PIPE_MAX_NEW", os.environ.get("SELFPLAY_MAX_TOKENS", "256") or 256))
    pre_k = int(os.environ.get("PIPE_PRECHECK_K", os.environ.get("SELFPLAY_PRECHECK_K", "3") or 3))
    do_heal = _bool_env("PIPE_HEAL", True)

    # Load model and task
    cfg_seed = {
        "model_name": os.environ.get("SELFPLAY_MODEL", "meta-llama/Llama-3.1-8B-Instruct"),
        "direct_model": False,
        "adapter_dir": os.environ.get("SELFPLAY_ADAPTER_DIR", "") or "",
    }
    tok, model = load_seed(cfg_seed)
    spec = sample_task(0)
    prompt = render_prompt(spec)

    # Trace output file
    out_path = ROOT / "scripts" / "data" / "traces" / "pipeline.jsonl"
    out_path.parent.mkdir(parents=True, exist_ok=True)

    def write(obj):
        with out_path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(obj, ensure_ascii=False) + "\n")

    write({
        "stage": "prompt",
        "task": spec.get("name"),
        "signature": spec.get("signature"),
        "prompt_head": (prompt or "")[:1200],
        "n_samples": n_cands,
        "max_new_tokens": max_new,
    })

    cands = sample(tok, model, prompt, n=n_cands, max_new_tokens=max_new)
    write({"stage": "sample_done", "num_candidates": len(cands)})

    passes = 0
    for idx, code in enumerate(cands):
        rec = {
            "stage": "candidate_eval",
            "cand_idx": idx,
            "code_head": (code or "")[:1200],
        }
        try:
            pr_sub, rep_sub = run_subset_tests(code, spec, k=pre_k)
            rec["subset"] = {"pr": pr_sub, **rep_sub}
        except Exception as e:
            rec["subset_err"] = str(e)
        try:
            ok, rep = run_and_test(code, spec)
            rec["full"] = {"ok": ok, **rep}
            if ok:
                passes += 1
            elif do_heal:
                healed = heal_code(tok, model, prompt, code, rep, max_new_tokens=max_new)
                ok2, rep2 = run_and_test(healed, spec)
                rec["heal"] = {"ok": ok2, **rep2, "healed_head": (healed or "")[:1200]}
                if ok2:
                    passes += 1
        except Exception as e:
            rec["full_err"] = str(e)
        write(rec)

    write({"stage": "summary", "num_candidates": len(cands), "passes": passes})
    print(str(out_path))

if __name__ == "__main__":
    main()


import os
import sys
import json
from pathlib import Path

# Ensure project root is on sys.path so 'agi_dw' package imports work when run directly
try:
    ROOT = Path(__file__).resolve().parents[4]
    if str(ROOT) not in sys.path:
        sys.path.insert(0, str(ROOT))
except Exception:
    pass

from agi_dw.scripts.selfplay.modules.model_io import load_seed
from agi_dw.scripts.selfplay.modules.tasks import sample_task, render_prompt
from agi_dw.scripts.selfplay.modules.generation import sample
from agi_dw.scripts.selfplay.modules.sandbox import run_subset_tests, run_and_test
from agi_dw.scripts.selfplay.modules.healing import heal_code

def _bool_env(name: str, default: bool) -> bool:
    v = os.environ.get(name)
    if v is None:
        return bool(default)
    s = str(v).strip().lower()
    return s in ("1", "true", "yes", "y", "on")

def main() -> None:
    # Configs via env
    n_cands = int(os.environ.get("PIPE_N_SAMPLES", os.environ.get("SELFPLAY_N_SAMPLES", "3") or 3))
    max_new = int(os.environ.get("PIPE_MAX_NEW", os.environ.get("SELFPLAY_MAX_TOKENS", "256") or 256))
    pre_k = int(os.environ.get("PIPE_PRECHECK_K", os.environ.get("SELFPLAY_PRECHECK_K", "3") or 3))
    do_heal = _bool_env("PIPE_HEAL", True)

    # Load model and task
    cfg_seed = {
        "model_name": os.environ.get("SELFPLAY_MODEL", "meta-llama/Llama-3.1-8B-Instruct"),
        "direct_model": False,
        "adapter_dir": os.environ.get("SELFPLAY_ADAPTER_DIR", "") or "",
    }
    tok, model = load_seed(cfg_seed)
    spec = sample_task(0)
    prompt = render_prompt(spec)

    # Trace output file


    if str(ROOT) not in sys.path:
        sys.path.insert(0, str(ROOT))
except Exception:
    pass

from agi_dw.scripts.selfplay.modules.model_io import load_seed
from agi_dw.scripts.selfplay.modules.tasks import sample_task, render_prompt
from agi_dw.scripts.selfplay.modules.generation import sample
from agi_dw.scripts.selfplay.modules.sandbox import run_subset_tests, run_and_test
from agi_dw.scripts.selfplay.modules.healing import heal_code

def _bool_env(name: str, default: bool) -> bool:
    v = os.environ.get(name)
    if v is None:
        return bool(default)
    s = str(v).strip().lower()
    return s in ("1", "true", "yes", "y", "on")

def main() -> None:
    # Configs via env
    n_cands = int(os.environ.get("PIPE_N_SAMPLES", os.environ.get("SELFPLAY_N_SAMPLES", "3") or 3))
    max_new = int(os.environ.get("PIPE_MAX_NEW", os.environ.get("SELFPLAY_MAX_TOKENS", "256") or 256))
    pre_k = int(os.environ.get("PIPE_PRECHECK_K", os.environ.get("SELFPLAY_PRECHECK_K", "3") or 3))
    do_heal = _bool_env("PIPE_HEAL", True)

    # Load model and task
    cfg_seed = {
        "model_name": os.environ.get("SELFPLAY_MODEL", "meta-llama/Llama-3.1-8B-Instruct"),
        "direct_model": False,
        "adapter_dir": os.environ.get("SELFPLAY_ADAPTER_DIR", "") or "",
    }
    tok, model = load_seed(cfg_seed)
    spec = sample_task(0)
    prompt = render_prompt(spec)

    # Trace output file
    out_path = ROOT / "scripts" / "data" / "traces" / "pipeline.jsonl"
    out_path.parent.mkdir(parents=True, exist_ok=True)

    def write(obj):
        with out_path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(obj, ensure_ascii=False) + "\n")

    write({
        "stage": "prompt",
        "task": spec.get("name"),
        "signature": spec.get("signature"),
        "prompt_head": (prompt or "")[:1200],
        "n_samples": n_cands,
        "max_new_tokens": max_new,
    })

    cands = sample(tok, model, prompt, n=n_cands, max_new_tokens=max_new)
    write({"stage": "sample_done", "num_candidates": len(cands)})

    passes = 0
    for idx, code in enumerate(cands):
        rec = {
            "stage": "candidate_eval",
            "cand_idx": idx,
            "code_head": (code or "")[:1200],
        }
        try:
            pr_sub, rep_sub = run_subset_tests(code, spec, k=pre_k)
            rec["subset"] = {"pr": pr_sub, **rep_sub}
        except Exception as e:
            rec["subset_err"] = str(e)
        try:
            ok, rep = run_and_test(code, spec)
            rec["full"] = {"ok": ok, **rep}
            if ok:
                passes += 1
            elif do_heal:
                healed = heal_code(tok, model, prompt, code, rep, max_new_tokens=max_new)
                ok2, rep2 = run_and_test(healed, spec)
                rec["heal"] = {"ok": ok2, **rep2, "healed_head": (healed or "")[:1200]}
                if ok2:
                    passes += 1
        except Exception as e:
            rec["full_err"] = str(e)
        write(rec)

    write({"stage": "summary", "num_candidates": len(cands), "passes": passes})
    print(str(out_path))

if __name__ == "__main__":
    main()


write

    # Configs via env
    n_cands = int(os.environ.get("PIPE_N_SAMPLES", os.environ.get("SELFPLAY_N_SAMPLES", "3") or 3))
    max_new = int(os.environ.get("PIPE_MAX_NEW", os.environ.get("SELFPLAY_MAX_TOKENS", "256") or 256))
    pre_k = int(os.environ.get("PIPE_PRECHECK_K", os.environ.get("SELFPLAY_PRECHECK_K", "3") or 3))
    do_heal = _bool_env("PIPE_HEAL", True)

    # Load model and task
    cfg_seed = {
        "model_name": os.environ.get("SELFPLAY_MODEL", "meta-llama/Llama-3.1-8B-Instruct"),
        "direct_model": False,
        "adapter_dir": os.environ.get("SELFPLAY_ADAPTER_DIR", "") or "",
    }
    tok, model = load_seed(cfg_seed)
    spec = sample_task(0)
    prompt = render_prompt(spec)

    # Trace output file
    out_path = ROOT / "scripts" / "data" / "traces" / "pipeline.jsonl"
    out_path.parent.mkdir(parents=True, exist_ok=True)

    def write(obj):
        with out_path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(obj, ensure_ascii=False) + "\n")

    write({
        "stage": "prompt",
        "task": spec.get("name"),
        "signature": spec.get("signature"),
        "prompt_head": (prompt or "")[:1200],
        "n_samples": n_cands,
        "max_new_tokens": max_new,
    })

    cands = sample(tok, model, prompt, n=n_cands, max_new_tokens=max_new)
    write({"stage": "sample_done", "num_candidates": len(cands)})

    passes = 0
    for idx, code in enumerate(cands):
        rec = {
            "stage": "candidate_eval",
            "cand_idx": idx,
            "code_head": (code or "")[:1200],
        }


import os
import sys
import json
import itertools
from pathlib import Path

# Ensure project root on sys.path
try:
    ROOT = Path(__file__).resolve().parents[4]
    if str(ROOT) not in sys.path:
        sys.path.insert(0, str(ROOT))
except Exception:
    pass

from agi_dw.scripts.selfplay.modules.model_io import load_seed
from agi_dw.scripts.selfplay.modules.tasks import sample_task, render_prompt
from agi_dw.scripts.selfplay.modules.generation import sample
from agi_dw.scripts.selfplay.modules.sandbox import run_subset_tests

def main() -> None:
    # Sweep grids
    temps = [0.2, 0.3, 0.4]
    topps = [0.8, 0.9, 0.95]
    n_vals = [4, 8]
    pre_k = int(os.environ.get("SWEEP_PRECHECK_K", os.environ.get("SELFPLAY_PRECHECK_K", "3") or 3))
    trials = int(os.environ.get("SWEEP_TRIALS", "3") or 3)  # number of tasks per setting

    # Load seed
    cfg_seed = {
        "model_name": os.environ.get("SELFPLAY_MODEL", "meta-llama/Llama-3.1-8B-Instruct"),
        "direct_model": False,
    }
    tok, model = load_seed(cfg_seed)

    out_path = ROOT / "scripts" / "data" / "traces" / "sweep_compliance.jsonl"
    out_path.parent.mkdir(parents=True, exist_ok=True)

    def write(obj):
        with out_path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(obj, ensure_ascii=False) + "\n")

    best = None
    for temp, topp, n in itertools.product(temps, topps, n_vals):
        # set decoding env for generation.sample
        os.environ["SAMPLING_TEMPERATURE"] = str(temp)
        os.environ["SAMPLING_TOP_P"] = str(topp)
        stats = {"candidates": 0, "valid_def": 0, "subset_ok": 0}
        for _ in range(trials):
            spec = sample_task(0)
            prompt = render_prompt(spec)
            cands = sample(tok, model, prompt, n=n, max_new_tokens=int(os.environ.get("SELFPLAY_MAX_TOKENS", "256") or 256))
            for code in cands:
                stats["candidates"] += 1
                has_def = bool(code and ("def solve(" in code))
                if has_def:
                    stats["valid_def"] += 1
                    pr, rep = run_subset_tests(code, spec, k=pre_k)
                    if pr > 0.0:
                        stats["subset_ok"] += 1
        result = {"temp": temp, "top_p": topp, "n": n, **stats}
        result["valid_def_rate"] = stats["valid_def"] / max(1, stats["candidates"])
        result["subset_ok_rate"] = stats["subset_ok"] / max(1, stats["candidates"])
        write(result)
        if (best is None) or (result["subset_ok_rate"] > best["subset_ok_rate"]) or (
            result["subset_ok_rate"] == best["subset_ok_rate"] and result["valid_def_rate"] > best["valid_def_rate"]
        ):
            best = result

    print("Best:", json.dumps(best, ensure_ascii=False))
    print(str(out_path))

if __name__ == "__main__":
    main()


def main() -> None:
    # Sweep grids
    temps = [0.2, 0.3, 0.4]
    topps = [0.8, 0.9, 0.95]
    n_vals = [4, 8]
    pre_k = int(os.environ.get("SWEEP_PRECHECK_K", os.environ.get("SELFPLAY_PRECHECK_K", "3") or 3))
    trials = int(os.environ.get("SWEEP_TRIALS", "3") or 3)  # number of tasks per setting

    # Load seed
    cfg_seed = {
        "model_name": os.environ.get("SELFPLAY_MODEL", "meta-llama/Llama-3.1-8B-Instruct"),
        "direct_model": False,
    }
    tok, model = load_seed(cfg_seed)

    out_path = ROOT / "scripts" / "data" / "traces" / "sweep_compliance.jsonl"
    out_path.parent.mkdir(parents=True, exist_ok=True)

    def write(obj):
        with out_path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(obj, ensure_ascii=False) + "\n")

    best = None
    for temp, topp, n in itertools.product(temps, topps, n_vals):
        # set decoding env for generation.sample
        os.environ["SAMPLING_TEMPERATURE"] = str(temp)
        os.environ["SAMPLING_TOP_P"] = str(topp)
        stats = {"candidates": 0, "valid_def": 0, "subset_ok": 0}
        for _ in range(trials):
            spec = sample_task(0)
            prompt = render_prompt(spec)
            cands = sample(tok, model, prompt, n=n, max_new_tokens=int(os.environ.get("SELFPLAY_MAX_TOKENS", "256") or 256))
            for code in cands:
                stats["candidates"] += 1
                has_def = bool(code and ("def solve(" in code))
                if has_def:
                    stats["valid_def"] += 1
                    pr, rep = run_subset_tests(code, spec, k=pre_k)
                    if pr > 0.0:
                        stats["subset_ok"] += 1
        result = {"temp": temp, "top_p": topp, "n": n, **stats}


from __future__ import annotations
import json
from pathlib import Path
import os
from typing import Dict, Any

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer  # type: ignore

from agi_dw.scripts.selfplay.modules.paths import data_path
from agi_dw.scripts.selfplay.modules.train import dpo_step, kto_step

def _env_str(name: str, default: str) -> str:
    v = os.environ.get(name)
    return str(v) if v is not None and str(v).strip() != "" else str(default)

def _env_float(name: str, default: float) -> float:
    try:
        v = os.environ.get(name)
        return float(v) if v is not None and str(v).strip() != "" else float(default)
    except Exception:
        return float(default)

def main() -> None:
    model_name = _env_str("PREF_MODEL", _env_str("SELFPLAY_MODEL", "meta-llama/Llama-3.1-8B-Instruct"))
    lr = _env_float("PREF_LR", 1e-4)
    algo = _env_str("PREF_ALGO", "dpo").lower()  # dpo | kto
    beta = _env_float("PREF_DPO_BETA", 0.1)
    tau = _env_float("PREF_KTO_TAU", 1.0)
    max_pairs = int(os.environ.get("PREF_MAX_PAIRS", "0") or 0)
    pairs_path = Path(os.environ.get("PREF_PAIRS_PATH", str(data_path("selfplay", "pref_pairs.jsonl"))))

    tok = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
    # Prefer BF16/FP16 when possible for speed
    if torch.cuda.is_available():
        try:
            model = model.to(torch.bfloat16)
        except Exception:
            try:
                model = model.to(torch.float16)
            except Exception:
                pass

    seen = 0
    with pairs_path.open("r", encoding="utf-8") as f:
        for line in f:
            try:
                obj = json.loads(line)
            except Exception:
                continue
            prompt = str(obj.get("prompt_head", ""))
            good = str(obj.get("good", ""))
            bad = str(obj.get("bad", ""))
            if not (prompt and good and bad and good != bad):
                continue
            if algo == "kto":
                _ = kto_step(tok, model, prompt, good, bad, lr=lr, tau=tau)
            else:
                _ = dpo_step(tok, model, prompt, good, bad, lr=lr, beta=beta)
            seen += 1
            if max_pairs > 0 and seen >= max_pairs:
                break
    # Optionally save adapters if using LoRA / PEFT
    out = os.environ.get("PREF_SAVE_ADAPTERS", "")
    if out:
        try:
            from agi_dw.scripts.selfplay.modules.model_io import save_adapters
            save_adapters(model, out)
        except Exception:
            pass

if __name__ == "__main__":
    main()


from __future__ import annotations
import json
from pathlib import Path
import os
from typing import Dict, Any

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer  # type: ignore

from agi_dw.scripts.selfplay.modules.paths import data_path
from agi_dw.scripts.selfplay.modules.train import dpo_step, kto_step

def _env_str(name: str, default: str) -> str:
    v = os.environ.get(name)
    return str(v) if v is not None and str(v).strip() != "" else str(default)

def _env_float(name: str, default: float) -> float:
    try:
        v = os.environ.get(name)
        return float(v) if v is not None and str(v).strip() != "" else float(default)
    except Exception:
        return float(default)

def main() -> None:
    model_name = _env_str("PREF_MODEL", _env_str("SELFPLAY_MODEL", "meta-llama/Llama-3.1-8B-Instruct"))
    lr = _env_float("PREF_LR", 1e-4)
    algo = _env_str("PREF_ALGO", "dpo").lower()  # dpo | kto
    beta = _env_float("PREF_DPO_BETA", 0.1)
    tau = _env_float("PREF_KTO_TAU", 1.0)
    max_pairs = int(os.environ.get("PREF_MAX_PAIRS", "0") or 0)
    pairs_path = Path(os.environ.get("PREF_PAIRS_PATH", str(data_path("selfplay", "pref_pairs.jsonl"))))

    tok = AutoTokenizer.from_pretrained(model_name)


from __future__ import annotations
import json
from pathlib import Path
import os
from typing import Dict, Any

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer  # type: ignore

from agi_dw.scripts.selfplay.modules.paths import data_path
from agi_dw.scripts.selfplay.modules.train import dpo_step, kto_step

def _env_str(name: str, default: str) -> str:
    v = os.environ.get(name)
    return str(v) if v is not None and str(v).strip() != "" else str(default)

def _env_float(name: str, default: float) -> float:
    try:
        v = os.environ.get(name)
        return float(v) if v is not None and str(v).strip() != "" else float(default)
    except Exception:
        return float(default)

def main() -> None:
    model_name = _env_str("PREF_MODEL", _env_str("SELFPLAY_MODEL", "meta-llama/Llama-3.1-8B-Instruct"))
    lr = _env_float("PREF_LR", 1e-4)
    algo = _env_str("PREF_ALGO", "dpo").lower()  # dpo | kto
    beta = _env_float("PREF_DPO_BETA", 0.1)
    tau = _env_float("PREF_KTO_TAU", 1.0)
    max_pairs = int(os.environ.get("PREF_MAX_PAIRS", "0") or 0)
    pairs_path = Path(os.environ.get("PREF_PAIRS_PATH", str(data_path("selfplay", "pref_pairs.jsonl"))))

    tok = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
    # Prefer BF16/FP16 when possible for speed
    if torch.cuda.is_available():
        try:
            model = model.to(torch.bfloat16)
        except Exception:
            try:
                model = model.to(torch.float16)


main

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer  # type: ignore

from agi_dw.scripts.selfplay.modules.paths import data_path
from agi_dw.scripts.selfplay.modules.train import dpo_step, kto_step

def _env_str(name: str, default: str) -> str:
    v = os.environ.get(name)
    return str(v) if v is not None and str(v).strip() != "" else str(default)

def _env_float(name: str, default: float) -> float:
    try:
        v = os.environ.get(name)
        return float(v) if v is not None and str(v).strip() != "" else float(default)
    except Exception:
        return float(default)

def main() -> None:
    model_name = _env_str("PREF_MODEL", _env_str("SELFPLAY_MODEL", "meta-llama/Llama-3.1-8B-Instruct"))
    lr = _env_float("PREF_LR", 1e-4)
    algo = _env_str("PREF_ALGO", "dpo").lower()  # dpo | kto
    beta = _env_float("PREF_DPO_BETA", 0.1)
    tau = _env_float("PREF_KTO_TAU", 1.0)
    max_pairs = int(os.environ.get("PREF_MAX_PAIRS", "0") or 0)
    pairs_path = Path(os.environ.get("PREF_PAIRS_PATH", str(data_path("selfplay", "pref_pairs.jsonl"))))

    tok = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
    # Prefer BF16/FP16 when possible for speed
    if torch.cuda.is_available():
        try:
            model = model.to(torch.bfloat16)
        except Exception:
            try:
                model = model.to(torch.float16)
            except Exception:
                pass

    seen = 0
    with pairs_path.open("r", encoding="utf-8") as f:
        for line in f:
            try:
                obj = json.loads(line)
            except Exception:
                continue
            prompt = str(obj.get("prompt_head", ""))
            good = str(obj.get("good", ""))
            bad = str(obj.get("bad", ""))
            if not (prompt and good and bad and good != bad):
                continue
            if algo == "kto":
                _ = kto_step(tok, model, prompt, good, bad, lr=lr, tau=tau)
            else:
                _ = dpo_step(tok, model, prompt, good, bad, lr=lr, beta=beta)
            seen += 1
            if max_pairs > 0 and seen >= max_pairs:
                break
    # Optionally save adapters if using LoRA / PEFT
    out = os.environ.get("PREF_SAVE_ADAPTERS", "")
    if out:
        try:
            from agi_dw.scripts.selfplay.modules.model_io import save_adapters
            save_adapters(model, out)
        except Exception:
            pass

if __name__ == "__main__":
    main()


import os
import sys
import json
from pathlib import Path

# Ensure project root on sys.path
try:
    ROOT = Path(__file__).resolve().parents[4]
    if str(ROOT) not in sys.path:
        sys.path.insert(0, str(ROOT))
except Exception:
    pass

from agi_dw.core.llm.hf_client import HFClient  # type: ignore
from agi_dw.scripts.selfplay.modules.tasks import sample_task, render_prompt
from agi_dw.scripts.selfplay.modules.generation import sample
from agi_dw.scripts.selfplay.modules.sandbox import run_subset_tests

def main() -> None:
    model_name = os.environ.get("HF_PROBE_MODEL", os.environ.get("SELFPLAY_MODEL", "meta-llama/Llama-3.1-8B-Instruct"))
    n = int(os.environ.get("HF_PROBE_N", "8") or 8)
    trials = int(os.environ.get("HF_PROBE_TRIALS", "5") or 5)
    temp = float(os.environ.get("HF_PROBE_TEMPERATURE", os.environ.get("SAMPLING_TEMPERATURE", "0.3") or 0.3))
    top_p = float(os.environ.get("HF_PROBE_TOP_P", os.environ.get("SAMPLING_TOP_P", "0.95") or 0.95))
    max_new = int(os.environ.get("HF_PROBE_MAX_NEW", os.environ.get("SELFPLAY_MAX_TOKENS", "256") or 256))

    # Set decoding for generation.sample
    os.environ["SAMPLING_TEMPERATURE"] = str(temp)
    os.environ["SAMPLING_TOP_P"] = str(top_p)

    client = HFClient.get_cached(model_name)
    tok, model = client.tokenizer, client.model
    # Ensure pad/eos interop
    try:
        if getattr(tok, "pad_token_id", None) is None and getattr(tok, "eos_token_id", None) is not None:
            tok.pad_token_id = tok.eos_token_id  # type: ignore[attr-defined]
    except Exception:
        pass

    out_path = ROOT / "scripts" / "data" / "traces" / "hf_direct_probe.jsonl"
    out_path.parent.mkdir(parents=True, exist_ok=True)

    def write(obj):
        with out_path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(obj, ensure_ascii=False) + "\n")

    stats = {"candidates": 0, "valid_def": 0, "syntax_ok": 0, "subset_ok": 0}
    for _ in range(trials):
        spec = sample_task(0)
        prompt = render_prompt(spec)
        codes = sample(tok, model, prompt, n=n, max_new_tokens=max_new)
        for code in codes:
            stats["candidates"] += 1
            has_def = bool(code and ("def solve(" in code))
            if has_def:
                stats["valid_def"] += 1
            pr, rep = run_subset_tests(code, spec, k=int(os.environ.get("HF_PROBE_PRECHECK_K", "3") or 3))
            if pr > 0.0:
                stats["subset_ok"] += 1
            # syntax ok if subset runner didn't throw parse error; approximate via absence of 'invalid syntax'
            first_fail = str(rep.get("first_failure", ""))
            if "invalid syntax" not in first_fail and "expected an indented block" not in first_fail:
                stats["syntax_ok"] += 1
            write({
                "task": spec.get("name"),
                "code_head": (code or "")[:800],
                "valid_def": has_def,
                "subset_pr": pr,
                "first_failure": rep.get("first_failure"),
            })

    summary = {
        **stats,
        "valid_def_rate": stats["valid_def"] / max(1, stats["candidates"]),
        "syntax_ok_rate": stats["syntax_ok"] / max(1, stats["candidates"]),
        "subset_ok_rate": stats["subset_ok"] / max(1, stats["candidates"]),
        "model": model_name,
        "n": n,
        "trials": trials,
        "temperature": temp,
        "top_p": top_p,
    }
    write({"summary": summary})
    print(json.dumps(summary, ensure_ascii=False))
    print(str(out_path))

if __name__ == "__main__":
    main()


write

    temp = float(os.environ.get("HF_PROBE_TEMPERATURE", os.environ.get("SAMPLING_TEMPERATURE", "0.3") or 0.3))
    top_p = float(os.environ.get("HF_PROBE_TOP_P", os.environ.get("SAMPLING_TOP_P", "0.95") or 0.95))
    max_new = int(os.environ.get("HF_PROBE_MAX_NEW", os.environ.get("SELFPLAY_MAX_TOKENS", "256") or 256))

    # Set decoding for generation.sample
    os.environ["SAMPLING_TEMPERATURE"] = str(temp)
    os.environ["SAMPLING_TOP_P"] = str(top_p)

    client = HFClient.get_cached(model_name)
    tok, model = client.tokenizer, client.model
    # Ensure pad/eos interop
    try:
        if getattr(tok, "pad_token_id", None) is None and getattr(tok, "eos_token_id", None) is not None:
            tok.pad_token_id = tok.eos_token_id  # type: ignore[attr-defined]
    except Exception:
        pass

    out_path = ROOT / "scripts" / "data" / "traces" / "hf_direct_probe.jsonl"
    out_path.parent.mkdir(parents=True, exist_ok=True)

    def write(obj):
        with out_path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(obj, ensure_ascii=False) + "\n")

    stats = {"candidates": 0, "valid_def": 0, "syntax_ok": 0, "subset_ok": 0}
    for _ in range(trials):
        spec = sample_task(0)
        prompt = render_prompt(spec)
        codes = sample(tok, model, prompt, n=n, max_new_tokens=max_new)
        for code in codes:
            stats["candidates"] += 1
            has_def = bool(code and ("def solve(" in code))
            if has_def:
                stats["valid_def"] += 1
            pr, rep = run_subset_tests(code, spec, k=int(os.environ.get("HF_PROBE_PRECHECK_K", "3") or 3))
            if pr > 0.0:
                stats["subset_ok"] += 1
            # syntax ok if subset runner didn't throw parse error; approximate via absence of 'invalid syntax'
            first_fail = str(rep.get("first_failure", ""))
            if "invalid syntax" not in first_fail and "expected an indented block" not in first_fail:
                stats["syntax_ok"] += 1
            write({
                "task": spec.get("name"),


import os
import sys
from pathlib import Path

# Ensure project root is on sys.path so 'agi_dw' package imports work when run directly
try:
    ROOT = Path(__file__).resolve().parents[4]
    if str(ROOT) not in sys.path:
        sys.path.insert(0, str(ROOT))
except Exception:
    pass

from agi_dw.scripts.selfplay.modules.model_io import load_seed
from agi_dw.scripts.selfplay.modules.tasks import sample_task, render_prompt
from agi_dw.scripts.selfplay.modules.generation import sample, generate_text

def main() -> None:
    cfg = {
        "model_name": os.environ.get("SELFPLAY_MODEL", "meta-llama/Llama-3.1-8B-Instruct"),
        "direct_model": False,
    }
    tok, model = load_seed(cfg)
    spec = sample_task(0)
    prompt = render_prompt(spec)
    # run one sample and one generate_text to emit traces
    _ = sample(tok, model, prompt, n=1, max_new_tokens=int(os.environ.get("TRACE_MAX_NEW", "128") or 128))
    _ = generate_text(tok, model, prompt, max_new_tokens=int(os.environ.get("TRACE_MAX_NEW", "128") or 128))

if __name__ == "__main__":
    main()


#!/usr/bin/env python3
import os
from pathlib import Path
import git

def is_git_repo(path):
    """Check if a directory is a git repository."""
    try:
        git.Repo(path)
        return True
    except git.exc.InvalidGitRepositoryError:
        return False

def list_repos(inspiration_dir, output_file="inspiration_repos.txt"):
    """List all git repositories in the inspiration directory and save to a file."""
    inspiration_path = Path(inspiration_dir)

    # Prepare the output lines
    output_lines = []

    if not inspiration_path.exists():


downloadarxiv.py

#!/usr/bin/env python3
import argparse
import os
import time
from urllib.parse import urljoin, urlparse

import requests
from bs4 import BeautifulSoup

LIST_URL = "https://arxiv.org/list/cs/recent"

def get_all_page(url: str) -> str:
    resp = requests.get(url, timeout=30)
    resp.raise_for_status()
    soup = BeautifulSoup(resp.text, "html.parser")
    all_link = None
    for a in soup.select("a"):
        if a.get_text(strip=True).lower() == "all" and a.get("href"):
            all_link = urljoin(url, a["href"])
            break
    # Fallback to show=5000 if "all" not found


#!/usr/bin/env python3
from __future__ import annotations

import argparse
import os
from pathlib import Path
from typing import Dict, Tuple, List

DEFAULT_EXCLUDES = [
    ".git",
    "__pycache__",
    ".mypy_cache",
    ".pytest_cache",
    "node_modules",
    "agi_dw.egg-info",
    "venv",
    ".venv",
    "env",
    "build",
    "dist",


repo/b.py

__all__ = ["util", "Base"]

def util():
    return 42

def _secret():
    return 0

class Base:
    def base_method(self):
        return "base"


repo/a.py

from b import *

class A(Base):
    def n(self):
        return "n"

    def m(self):
        self.n()
        super().base_method()

def top():
    util()
    _secret()  # should remain unresolved due to star import skipping private names


from a import A, top

def test_top():
    assert top() is None

class TestA:
    def test_methods(self):
        a = A()
        assert a.n() == "n"
        assert a.m() is None


import logging
"""
Setup script for agi_dw package.
"""

from setuptools import setup, find_packages

setup(
	name="agi_dw",
	version="0.1.0",
	description="AGI Development Workflow",
	packages=find_packages(),
	python_requires=">=3.8",
	install_requires=[
		"pytest>=7.0.0",
		"pytest-cov>=4.0.0",
		"pathlib",
	],
	extras_require={
		"dev": [
			"pytest>=7.0.0",


import logging
"""
AGI Development Workflow package.
"""

__version__ = "0.1.0"


import logging
from pathlib import Path
from agi_dw.tools.patch_actuator import apply_unified_diff

def _make_repo(tmp_path: Path) -> Path:
	repo = tmp_path / "repo"
	repo.mkdir(parents=True, exist_ok=True)
	(repo / ".git").mkdir(exist_ok=True)  # minimal marker
	(repo / "a.py").write_text("print('a')\n", encoding="utf-8")
	return repo

def test_apply_respects_max_files(tmp_path: Path):
	repo = _make_repo(tmp_path)
	diff = (
		"diff --git a/a.py b/a.py\n"
		"index 0000000..0000000 100644\n"
		"--- a/a.py\n"
		"+++ b/a.py\n"
		"@@ -1,1 +1,1 @@\n"


from __future__ import annotations
import logging

from agi_dw.tools.patch_actuator import validate_diff_policy

def test_validate_diff_policy_allows_py_and_limits_hunk():
	diff = """--- a/x.py
+++ b/x.py
@@ -1,1 +1,2 @@
-print('a')
+print('a')
+print('b')
"""
	res = validate_diff_policy(diff, allowed_exts=[".py"], max_hunk_lines=10)
	assert res["ok"] is True
	assert not res.get("issues")

def test_validate_diff_policy_blocks_extension():
	diff = """--- a/secret.pem


import logging
from agi_dw.tools.patch_safety import validate_unified_diff_schema, check_python_syntax
from pathlib import Path

def test_validate_unified_diff_schema_ok():
	diff = (
		"diff --git a/foo.py b/foo.py\n"
		"index 0000000..0000000 100644\n"
		"--- a/foo.py\n"
		"+++ b/foo.py\n"
		"@@ -1,1 +1,1 @@\n"
		"-print('a')\n"
		"+print('b')\n"
	)
	ok, why = validate_unified_diff_schema(diff)
	assert ok, why

def test_validate_unified_diff_schema_blocks_binary_and_modes():
	diff = (


import logging
import json
from pathlib import Path
from agi_dw.scripts.eval.ci_assert_safe_edits import main as ci_main

def test_ci_assert_safe_edits_tmp(tmp_path: Path, monkeypatch):
	log = tmp_path / "scheduler_runs.jsonl"
	# Write two lines: one safe, one unsafe
	line_ok = {"stdout_tail": "", "stderr_tail": ""}
	line_bad = {"stdout_tail": "diff_validation_failed:too_many_files"}
	log.write_text("\n".join([json.dumps(line_ok), json.dumps(line_bad)]), encoding="utf-8")

	# Point script to temp log path
	from importlib import reload
	import agi_dw.scripts.eval.ci_assert_safe_edits as mod
	def _fake_args():
		class A:
			logs = str(log)
		return A()
	monkeypatch.setattr(mod, "argparse", type("X", (), {"ArgumentParser": lambda *a, **k: type("Y", (), {"add_argument": lambda *a, **k: None, "parse_args": staticmethod(_fake_args)})()})())


from __future__ import annotations
import logging

from pathlib import Path

from agi_dw.scripts.devtools.devtools_orchestrator import make_flake8_whitespace_fix_diff

def test_whitespace_autofix_diff(tmp_path: Path):
	p = tmp_path / "a.py"
	p.write_text("print('x')  \n", encoding="utf-8")  # trailing spaces
	diff = make_flake8_whitespace_fix_diff(tmp_path)
	assert "--- a/" in diff and "+++ b/" in diff
	# Apply diff via a simple check: ensure it proposes removing trailing spaces
	assert "-print('x')  " in diff
	assert "+print('x')" in diff


import logging

from agi_dw.tools.index_rank import rank_index_candidates

def test_rank_index_candidates_basic():
	obs = {"content": "count lines", "meta": {"cwd": "/tmp"}}
	idx = {
		"functions": {
			"/tmp/app.py": [
				{"name": "count_lines"},
				{"name": "helper"},
			]
		},
		"classes": {
			"/tmp/app.py": [
				{"name": "Counter"}
			]
		}
	}
	out = rank_index_candidates(obs, idx, 2)


from __future__ import annotations
import logging

import tempfile
from pathlib import Path

def test_dual_route_inject_idempotent():
    src = """
def new_api(x):
    return x + 1

def old_api(x):
    return x
"""
    with tempfile.TemporaryDirectory() as td:
        p = Path(td) / "mod.py"
        p.write_text(src, encoding="utf-8")
        # First inject
        import subprocess
        rc = subprocess.run(["python", "scripts/codemods/dual_route_inject.py", str(p), "new_api", "old_api", "NEW_API"], capture_output=True, text=True).returncode


from __future__ import annotations
import logging

from agi_dw.core.utils.bench_utils import (
    strip_fences,
    precheck_code,
    retry_with_backoff,
    load_code_index,
    inject_similar_functions,
)

def test_strip_fences_basic():
    code = """```
print('hi')
```"""
    assert strip_fences(code) == "print('hi')"

def test_strip_fences_no_fences():
    code = "print('ok')"


from __future__ import annotations
import logging

from pathlib import Path

from agi_dw.tools.patch_safety import validate_unified_diff_schema, check_python_syntax, count_files_in_diff

def test_validate_unified_diff_schema_ok() -> None:
	diff = (
		"diff --git a/foo.py b/foo.py\n"
		"index 0000000..0000000 100644\n"
		"--- a/foo.py\n"
		"+++ b/foo.py\n"
		"@@ -1,1 +1,1 @@\n"
		"-x = 1\n"
		"+x = 2\n"
	)
	ok, why = validate_unified_diff_schema(diff)
	assert ok, f"expected ok, got reason: {why}"


from __future__ import annotations

import logging
import os
from typing import Optional

class Flags:
	"""Minimal feature flag SDK with env override.

	- is_enabled("NAME") checks environment variable AGI_FLAG_NAME.
	  Values considered true: 1, true, yes, on (case-insensitive).
	- fallback default can be provided per-call.
	"""

	def __init__(self) -> None:
		self._cache: dict[str, bool] = {}

	@staticmethod
	def _normalize_flag_env_name(name: str) -> str:
		return f"AGI_FLAG_{str(name or '').strip().upper()}"


from .sdk import FLAGS, Flags  # re-export for convenience

__all__ = ["FLAGS", "Flags"]


import logging
"""
Core AGI components.
"""


from __future__ import annotations

import re
from typing import Any, Dict, List, Tuple

def _count_bytes(text: str) -> int:
	try:
		return len(text.encode("utf-8"))
	except Exception:
		return len(text)

def assert_diff_safe(old: str, new: str, rules: Dict[str, Any] | None = None) -> Tuple[bool, Dict[str, Any]]:
	"""Assess risk of a text replacement using simple heuristics.

	Rules (all optional):
	- max_bytes_added
	- max_bytes_deleted
	- forbid_patterns: list[str] regex
	"""


from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple

from agi_dw.core.ops.tracing import trace_span  # type: ignore
from agi_dw.core.planner.llm_planner import emit_plan

# Domain-agnostic callable: given (obs, plan) returns a quick action dict for WM rollout
ActionProvider = Callable[[Dict[str, Any], Dict[str, Any]], Dict[str, Any]]

@dataclass
class PlannerConfig:
	model: str
	backend: str = "hf"
	timeout_sec: int = 30
	adapter_dir: Optional[str] = None


import logging
import json
import os
from typing import Any, Dict
from pathlib import Path
from datetime import datetime

try:
	import yaml  # type: ignore
except Exception:
	yaml = None

import torch

from agi_dw.core.llm.hf_client import HFClient
from agi_dw.core.utils.prompt_logger import PromptLogger, get_prompt_logger  # type: ignore
from agi_dw.core.ops.tracing import trace_span, meter_cost  # type: ignore
from agi_dw.tools.artifact_cache import is_enabled as cache_enabled, get_cached_text as cache_get, set_cached_text as cache_set  # type: ignore

def emit_plan(


from __future__ import annotations

import os
import time
from pathlib import Path
from typing import Optional

def _locks_dir() -> Path:
	try:
		return Path(__file__).resolve().parents[2] / "data" / "locks"
	except Exception:
		return Path.cwd() / "data" / "locks"

def _lock_path(key: str) -> Path:
	k = str(key).strip()
	if not k:
		raise ValueError("empty lease key")
	return _locks_dir() / (k.replace("/", "_") + ".lock")


from __future__ import annotations

import base64
import hashlib
import hmac
import os
from dataclasses import dataclass
from typing import Any, Optional, Tuple

@dataclass
class SecretHandle:
	name: str
	source: str
	_ref: str
	_value: Optional[bytes] = None

	def materialize(self) -> bytes:
		if self._value is not None:
			return self._value
		# Env source


from __future__ import annotations

import json
import os
import threading
import time
import uuid
from contextlib import contextmanager
from pathlib import Path
from typing import Any, Dict, Optional

_SPAN_LOCAL = threading.local()

def _now_iso() -> str:
	try:
		from datetime import datetime, timezone
		return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ")
	except Exception:
		return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())


from __future__ import annotations

import json
import time
import uuid
from pathlib import Path
from typing import Any, Dict, Optional

def _jobs_dir() -> Path:
	try:
		return Path(__file__).resolve().parents[2] / "data" / "jobs"
	except Exception:
		return Path.cwd() / "data" / "jobs"

def job_enqueue(task_spec: Dict[str, Any], delay_sec: int = 0, schedule: Optional[str] = None) -> str:
	jid = uuid.uuid4().hex
	p = _jobs_dir() / f"{jid}.json"
	p.parent.mkdir(parents=True, exist_ok=True)
	rec = {


from __future__ import annotations

import hashlib
import json
import os
import time
from pathlib import Path
from typing import Any, Dict, Optional, Tuple, Union

def _root() -> Path:
	try:
		return Path(__file__).resolve().parents[2] / "data" / "artifacts"
	except Exception:
		return Path.cwd() / "data" / "artifacts"

def _safe_key_path(key: str) -> Path:
	# Normalize key into nested directories based on sha256 of key
	ks = str(key).strip()
	if not ks:


from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

@dataclass
class MemoryAugmentConfig:
	# Episodic memory
	use_memory: bool = False
	mem_path: Optional[str] = None
	mem_topk: int = 3
	mem_recency: float = 0.0
	mem_query: Optional[str] = None
	# Code index (planner context)
	index_k: int = 0
	index_path: Optional[str] = None
	# Policy injections
	inject_dom_policy: bool = True


from __future__ import annotations
import logging

import json
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Any, Dict, List, Optional

@dataclass
class Skill:
	id: str
	domain: str  # "cli" | "dom" | "code" | ...
	description: str
	signature: Dict[str, Any]  # minimal inputs needed to reuse (e.g., {"tool":"grep"} or {"url": "..."})
	action: Dict[str, Any]  # canonical action payload (tool+args)
	metrics: Dict[str, Any]  # success_count, attempts, success_rate, last_used
	adapters: Dict[str, str] | None = None  # optional PEFT adapters per role: {"planner": dir, "verifier": dir}
	version: str = "0.1"


"""Code-specific memory operations for HumanEval and similar tasks."""

from __future__ import annotations
import logging
from typing import Any, Dict, List, Optional
from pathlib import Path
import json

from .hybrid import HybridMemory, MemoryConfig

# Optional utilities; provide robust fallbacks if unavailable
try:
    from ..utils.code_utils import extract_function_signature, normalize_code  # type: ignore
except Exception:  # Fallbacks to avoid hard dependency
    try:
        # Reuse existing helper to extract target function name
        from ..utils.bench_utils import extract_target_function_name  # type: ignore
    except Exception:
        extract_target_function_name = None  # type: ignore

    def extract_function_signature(prompt: str):  # type: ignore


from __future__ import annotations
import json
from pathlib import Path
from typing import Dict, Any, List

class MemoryRegistry:
    """
    Provides unified access to episodic/procedural/conceptual/contextual memory.
    """

    def __init__(self, root: Path) -> None:
        self.root = root

    # Episodic: task traces
    def list_traces(self) -> List[Path]:
        base = self.root / "data" / "traces"
        return sorted(base.glob("*.jsonl")) if base.exists() else []

    # Procedural: curated skills datasets
    def list_skills(self) -> List[Path]:


"""Hybrid memory system combining episodic and product key memories."""

from __future__ import annotations
import logging
from typing import Any, Dict, List, Optional
from dataclasses import dataclass

from .episodic import EpisodicMemory, MemoryItem
from .product_key import ProductKeyMemory

@dataclass
class MemoryConfig:
    """Configuration for hybrid memory system."""
    embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2"
    query_dim: int = 512
    key_dim: int = 64
    value_dim: int = 64
    num_heads: int = 4
    device: Optional[str] = None

class MemoryRouter:


from __future__ import annotations
import logging

import json
import os
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Tuple

import numpy as np

def _try_imports():
	try:
		import faiss  # type: ignore
	except Exception as e:
		raise RuntimeError("faiss is required. pip install faiss-cpu") from e
	try:
		from sentence_transformers import SentenceTransformer  # type: ignore
	except Exception as e:


"""Product Key Memory implementation."""

from __future__ import annotations
import logging
from typing import Any, Dict, List, Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from pathlib import Path
import json
try:
    import numpy as np  # type: ignore
except Exception:  # optional
    np = None  # type: ignore
try:
    from sentence_transformers import SentenceTransformer  # type: ignore
except Exception:  # optional
    SentenceTransformer = None  # type: ignore

class ProductKeyMemory(nn.Module):
    """Product Key Memory optimized for code storage and retrieval."""


from __future__ import annotations

import logging
import os
from pathlib import Path
from datetime import datetime
from typing import Optional

def _should_redact() -> bool:
	try:
		v = os.environ.get("AGI_LLM_REDACT_LOGS", "1").strip()
		return v not in ("0", "false", "False")
	except Exception:
		return True

def _redact(text: str) -> str:
	if not _should_redact():
		return text
	try:


from __future__ import annotations
import logging

import re
from typing import Tuple

EMAIL_RE = re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}")
PHONE_RE = re.compile(r"(?:(?:\+\d{1,3}[\s-]?)?(?:\(\d{2,4}\)[\s-]?)?\d{3,4}[\s-]?\d{3,4})")

def redact_text(text: str) -> Tuple[str, int]:
	"""Redact emails and phone-like numbers. Returns (redacted_text, num_replacements)."""
	count = 0
	def _email_sub(_m: re.Match) -> str:
		nonlocal count
		count += 1
		return "[REDACTED_EMAIL]"

	def _phone_sub(_m: re.Match) -> str:
		nonlocal count


from __future__ import annotations
import logging

import yaml
from typing import Dict, List, Literal, Optional, Tuple

from agi_dw.core.llm.hf_client import HFClient

class CodeCritic:
    """Code review critic for filtering risky completions.

    Uses a small LLM to analyze code and return structured feedback.
    """

    def __init__(self, model: str = "meta-llama/Llama-3.1-8B-Instruct") -> None:
        self.llm = HFClient.get_cached(model)

    def review(self, code: str, lints: str = "") -> Tuple[bool, Dict[str, List[str] | Literal["low", "medium", "high"]]]:
        """Review code and return (is_safe, feedback).


from __future__ import annotations
import logging

import json
import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, TypeVar

T = TypeVar("T")
R = TypeVar("R")

def ensure_safe_env() -> None:
	os.environ.setdefault("PYTEST_DISABLE_PLUGIN_AUTOLOAD", "1")
	os.environ.setdefault("PYTHONDONTWRITEBYTECODE", "1")
	# Avoid HF tokenizers fork warnings/deadlocks in worker pools
	os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")


from __future__ import annotations
import logging

from pathlib import Path
from typing import Any, Dict, Optional
import math

class WorldModelPrior:
	"""Lightweight wrapper around a joblib-packed WM (TF-IDF + linear models).

	Expects a joblib.dump of a dict with keys: {"vec", "clf", "reg"}.
	- vec: vectorizer with .transform
	- clf: classifier with .predict_proba
	- reg: regressor with .predict
	"""

	def __init__(self, vec: Any, clf: Any, reg: Any, cal: Any | None = None, risk_std: float | None = None) -> None:
		self.vec = vec
		self.clf = clf
		self.reg = reg


from __future__ import annotations
import logging

from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from .api import WorldModelPrior
from .rollout_api import RolloutAPI, RolloutConfig
from .next_state import NextStatePredictor

class WorldModelService:
	"""Central service for accessing world model priors and rollouts.

	Encapsulates model loading, prior prediction, and short-horizon rollout
	to provide a uniform interface to callers.
	"""

	def __init__(self, wm: Optional[WorldModelPrior], next_state_predictor: Optional[NextStatePredictor] = None) -> None:
		self._wm = wm
		self._nsp = next_state_predictor or NextStatePredictor()


from __future__ import annotations
import logging

from typing import Dict, Any, List, Optional

from .rollout_api import RolloutAPI, RolloutConfig

class ShortHorizonRollout:
	"""[DEPRECATED] Use WorldModelService instead.

	This class is deprecated and will be removed in a future version.
	Please use WorldModelService from service.py which provides a consolidated
	interface for world model operations including rollouts.

	Legacy docstring:
	k-step abstract action rollout using the world model prior.
	Extends scoring by producing a synthetic next observation per step via a
	simple next-state predictor. Designed for planning-only simulations.
	"""

	def __init__(self, wm_prior, next_state_predictor: Any | None = None) -> None:


from __future__ import annotations
import logging

from typing import Dict, Any, List, Optional

class OffPolicyProposer:
	"""Use WorldModelPrior to rank and propose low-risk actions off-policy.

	For CLI: proposes between IL NN and T5 candidates.
	For DOM: proposes between NN and T5 candidates.
	"""

	def __init__(self, wm) -> None:
		self.wm = wm

	def propose_cli(self, obs: Dict[str, Any], plan: Dict[str, Any], il_data_path: str, t5_model_path: str) -> Dict[str, Any]:
		from agi_dw.core.actuator.service import ActuatorConfig, RouterExtras, select_action  # type: ignore
		# Build NN and T5 candidates via centralized service
		cfg_nn = ActuatorConfig(mode="nn", il_path=str(il_data_path))
		extra = RouterExtras(domain="cli")