code/local_eval/gates.py

from __future__ import annotations

import re
from dataclasses import asdict, dataclass, field

from albedo_eval_service.judge_core import amputated_thinking, reserved_token_leak
from albedo_eval_service.shared.loop_check import loop_verdict
from albedo_eval_service.shared.observation_format import first_bash_block, is_truncated
from sanity_service.checks import check_one

from .live_protocol import is_live_submit

_EDIT_RE = re.compile(
    r"sed\s+-i|tee\s+[\w./-]|cat\s*>|str_replace|git apply|patch\s+-p|applypatch|"
    r"cp\s+[\w./-]|mv\s+[\w./-]|(?<![-\d&])>>?\s*(?!/dev/)[\w.][\w./-]*"
)


@dataclass
class GateReport:
    sample_id: str
    passed: bool
    fatal: list[str] = field(default_factory=list)
    warnings: list[str] = field(default_factory=list)
    n_cmds: int = 0
    n_edits: int = 0
    submitted: bool = False
    looped: bool = False
    amputated_thinking: bool = False
    truncated: bool = False
    dup_cmd_ratio: float = 0.0
    max_cmd_run: int = 0
    proxy_score: float = 0.0

    def as_dict(self) -> dict:
        return asdict(self)


def _assistant_texts(turns: list[dict] | None, document: str) -> list[str]:
    texts = [
        str(turn.get("content") or "")
        for turn in (turns or [])
        if turn.get("role") == "assistant" and turn.get("score_target")
    ]
    return texts or ([document] if document else [])


def _commands(texts: list[str]) -> list[str]:
    return [first_bash_block(text) for text in texts if first_bash_block(text)]


def _submitted(sample, texts: list[str]) -> bool:
    command = getattr(sample, "submit_command", "") or ""
    marker = getattr(sample, "submit_marker", "") or ""
    return any(is_live_submit(text, command=command, marker=marker) for text in texts)


def evaluate_side(sample, document: str, turns: list[dict] | None, *, truncated: bool) -> GateReport:
    texts = _assistant_texts(turns, document)
    commands = _commands(texts)
    edits = [cmd for cmd in commands if _EDIT_RE.search(cmd)]
    looped = loop_verdict(texts)
    heuristic = check_one(texts[0]) if texts else check_one("")
    leak = reserved_token_leak(document)
    amputated = amputated_thinking(document)
    submitted = _submitted(sample, texts)

    fatal: list[str] = []
    warnings: list[str] = []
    if truncated or any(is_truncated(text) for text in texts):
        fatal.append("truncated")
    if not heuristic.passed:
        fatal.append(heuristic.reason)
    if looped.looped:
        fatal.append("; ".join(looped.reasons) or "looped")
    if leak:
        fatal.append(f"reserved_token_leak:{leak}")
    if not commands:
        fatal.append("no bash command")
    if amputated:
        warnings.append("amputated_thinking (official score x0.5)")
    if not edits:
        warnings.append("no edit command")
    if not submitted:
        warnings.append("did not issue submit command")

    # Local proxy only — not the official GLM checklist.
    # Survive gates first, then reward edit + submit + command diversity.
    score = 0.0 if fatal else 0.35
    if not fatal:
        score += 0.25 if edits else 0.0
        score += 0.20 if submitted else 0.0
        score += 0.10 if len(set(commands)) >= 3 else 0.0
        score += 0.10 if looped.dup_cmd_ratio < 0.25 else 0.0
        if amputated:
            score *= 0.5

    return GateReport(
        sample_id=getattr(sample, "sample_id", ""),
        passed=not fatal,
        fatal=fatal,
        warnings=warnings,
        n_cmds=len(commands),
        n_edits=len(edits),
        submitted=submitted,
        looped=looped.looped,
        amputated_thinking=amputated,
        truncated=truncated or any(is_truncated(text) for text in texts),
        dup_cmd_ratio=round(looped.dup_cmd_ratio, 3),
        max_cmd_run=looped.max_cmd_run,
        proxy_score=round(score, 4),
    )