"""Battle viewer data layer (pure, unit-tested). Indexes a playback root that may hold *many* runs/models/scenarios and serves the cascade the UI needs: run_id → model → scenario (pack:level:split @ seed) → episode plus per-turn assembly (`episode_view`) and the comparison pairing rule (B is locked to A's scenario+seed; only run/model vary). Layout written by run_eval: /__//seed/{manifest,turns, messages,score,minimap_turnNNN}. Identity is read from each manifest (run_id/model/scenario/seed), never parsed from the path, so the folder scheme can change without breaking the viewer. """ from __future__ import annotations from dataclasses import dataclass from pathlib import Path from .playback_view import load_episode @dataclass(frozen=True) class EpisodeRef: run_id: str model: str scenario: str # pack:level:split seed: int outcome: str composite: float | None objective: float | None dir: str @property def scenario_key(self) -> str: """The (scenario, seed) identity comparison mode locks on.""" return f"{self.scenario}@{self.seed}" def scan(root: str | Path) -> list[EpisodeRef]: """Every episode under ``root``, newest run first. Tolerant of partial/running episodes (missing score → None).""" out: list[EpisodeRef] = [] for man in Path(root).glob("**/manifest.json"): d = man.parent ep = load_episode(d) m = ep["manifest"] if not m: continue score = {} sp = d / "score.json" if sp.exists(): import json try: score = json.loads(sp.read_text()) except Exception: # noqa: BLE001 score = {} out.append(EpisodeRef( run_id=str(m.get("run_id") or "unknown"), model=str(m.get("model") or "agent"), scenario=str(m.get("scenario") or d.parent.name), seed=int(m.get("seed") or 0), outcome=str(m.get("outcome") or "?"), composite=score.get("composite"), objective=score.get( "objective_progress", m.get("objective_progress") ), dir=str(d), )) out.sort(key=lambda e: (e.run_id, e.model, e.scenario, e.seed), reverse=True) return out def _u(seq): seen, out = set(), [] for x in seq: if x not in seen: seen.add(x) out.append(x) return out def runs(idx: list[EpisodeRef]) -> list[str]: return _u(e.run_id for e in idx) def models(idx: list[EpisodeRef], run_id: str) -> list[str]: return _u(e.model for e in idx if e.run_id == run_id) def scenarios(idx: list[EpisodeRef], run_id: str, model: str) -> list[str]: return _u( e.scenario_key for e in idx if e.run_id == run_id and e.model == model ) def find(idx: list[EpisodeRef], run_id: str, model: str, scenario_key: str) -> EpisodeRef | None: for e in idx: if (e.run_id == run_id and e.model == model and e.scenario_key == scenario_key): return e return None def compare_candidates(idx: list[EpisodeRef], a: EpisodeRef) -> list[EpisodeRef]: """Episodes runnable as B against A: same scenario+seed, any other run/model (A itself excluded).""" return [ e for e in idx if e.scenario_key == a.scenario_key and e.dir != a.dir ] def _text(content) -> str: if isinstance(content, list): return "\n".join( p.get("text", "") for p in content if isinstance(p, dict) and p.get("type") == "text" ) return content or "" def _role_turns(messages: list[dict], role: str) -> list[dict]: return [m for m in messages if m.get("role") == role] def episode_view(ep_dir: str | Path, turn_idx: int) -> dict: """Everything to render one turn of one episode. ``turn_idx`` is 0-based and clamped to range.""" ep = load_episode(ep_dir) turns = ep["turns"] n = len(turns) if n == 0: return {"n_turns": 0, "manifest": ep["manifest"]} i = max(0, min(turn_idx, n - 1)) t = turns[i] msgs = ep["messages"] sys_prompt = next( (m.get("content", "") for m in msgs if m.get("role") == "system"), "" ) if isinstance(sys_prompt, list): sys_prompt = _text(sys_prompt) users = _role_turns(msgs, "user") assts = _role_turns(msgs, "assistant") tools = _role_turns(msgs, "tool") g = t.get("goal") or {} # The model's "debrief" for a turn = the briefing it was handed # (post-interrupt this is the scoped re-prompt). Surface it # explicitly + the interrupt reason, matching the training viewer. briefing = _text(users[i]["content"]) if i < len(users) else "" return { "n_turns": n, "turn_idx": i, "turn": t.get("turn"), "tick": t.get("tick"), "interrupt": t.get("interrupt"), "minimap_png": t.get("minimap_png"), "ascii_minimap": t.get("ascii_minimap", ""), "system_prompt": sys_prompt, "briefing": briefing, "debrief": briefing, # alias: what the model was (re-)briefed with "reasoning": (assts[i].get("reasoning", "") if i < len(assts) else ""), "assistant_text": (_text(assts[i].get("content")) if i < len(assts) else ""), "tool_result": (_text(tools[i].get("content")) if i < len(tools) else ""), "commands": t.get("commands", []), "signals": t.get("signals", {}), "goal": g, "objective_progress": g.get("objective_progress"), "won": g.get("won"), "manifest": ep["manifest"], }