training/rollout.py

"""Multi-turn rollout function for PM-Ops GRPO training.

History design (no duplicates):
  Each turn record stores the obs that TRIGGERED the completion, not the result.
  build_messages reconstructs:
    [sys] [user:obs_0] [asst:comp_0] [user:obs_1] [asst:comp_1] ... [user:current_obs]
  current_obs is never in history — it becomes the final user message.

Other design decisions:
- Action format B: chain-of-thought reasoning + ```json block
- Truncation: runbook-pinned sliding window (runbook pair always kept)
- Fallback cascade: read_runbook (early) → noop (mid) → finish (late)
- Three-pass JSON extraction: code block → raw JSON → regex
"""
import json
import re
from typing import Any

import torch
import torch.nn.functional as F

from training.dataset import parse_seed_from_prompt
from training.prompts import SYSTEM_PROMPT, format_observation


# ---------------------------------------------------------------------------
# HF model.generate() — replaces generate_rollout_completions (vLLM-only)
# ---------------------------------------------------------------------------

def _get_model_for_generation(trainer):
    """Unwrap the model safely regardless of accelerate/PEFT/DDP wrapping.

    Priority:
      1. accelerator.unwrap_model — handles DDP + PEFT + DeepSpeed
      2. trainer.model.module     — plain DDP wrapping
      3. trainer.model            — unwrapped (local or single-GPU)
    """
    if hasattr(trainer, "accelerator"):
        return trainer.accelerator.unwrap_model(trainer.model)
    if hasattr(trainer.model, "module"):
        return trainer.model.module
    return trainer.model


def _generate_no_vllm(trainer, prompt_text: str, tokenizer,
                      max_new_tokens: int = 512, temperature: float = 1.1) -> dict:
    """Generate one completion using HF model.generate() without vLLM.

    Returns the same dict shape as generate_rollout_completions so the rest
    of rollout_once is unchanged:
        prompt_ids:     list[int]
        completion_ids: list[int]
        logprobs:       list[float]  (per-token log-prob under current policy)
        text:           str
    """
    model = _get_model_for_generation(trainer)

    # Device: prefer accelerator.device, fall back to first param device
    if hasattr(trainer, "accelerator"):
        device = trainer.accelerator.device
    else:
        device = next(model.parameters()).device

    enc = tokenizer(prompt_text, return_tensors="pt").to(device)
    prompt_len = enc["input_ids"].shape[1]

    with torch.no_grad():
        out = model.generate(
            **enc,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=temperature,
            top_p=0.95,
            top_k=50,
            pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
            output_scores=True,
            return_dict_in_generate=True,
        )

    prompt_ids     = enc["input_ids"][0].tolist()
    completion_ids = out.sequences[0][prompt_len:].tolist()

    # Per-token log-probs from output.scores (one score tensor per new token)
    logprobs = [
        F.log_softmax(score[0], dim=-1)[tok_id].item()
        for score, tok_id in zip(out.scores, completion_ids)
    ]

    text = tokenizer.decode(completion_ids, skip_special_tokens=True)
    return {
        "prompt_ids":     prompt_ids,
        "completion_ids": completion_ids,
        "logprobs":       logprobs,
        "text":           text,
    }

MAX_STEPS = 40
# ~3000 tokens at 4 chars/token; leaves room for completion tokens
MAX_PROMPT_CHARS = 12_000
RUNBOOK_RESPONSE_MAX_CHARS = 3_000
HISTORY_PAIRS = 6  # max (user+asst) pairs kept from non-runbook history


# ---------------------------------------------------------------------------
# JSON extraction — three-pass, most-specific first
# ---------------------------------------------------------------------------

def extract_json_action(text: str) -> dict | None:
    # Pass 1: last ```json ... ``` block
    matches = re.findall(r"```json\s*(\{.*?\})\s*```", text, re.DOTALL)
    if matches:
        try:
            return json.loads(matches[-1])
        except json.JSONDecodeError:
            pass

    # Pass 2: entire output is JSON
    stripped = text.strip()
    if stripped.startswith("{"):
        try:
            return json.loads(stripped)
        except json.JSONDecodeError:
            pass

    # Pass 3: any {...action_type...} pattern — allow one level of nested {} (e.g. "args": {})
    matches = re.findall(
        r'\{(?:[^{}]|\{[^{}]*\})*"action_type"\s*:\s*"[^"]*"(?:[^{}]|\{[^{}]*\})*\}',
        text, re.DOTALL,
    )
    if matches:
        try:
            return json.loads(matches[-1])
        except json.JSONDecodeError:
            pass

    return None


def step_aware_fallback(step: int, max_steps: int = MAX_STEPS) -> dict:
    """Safe fallback that degrades gracefully across the episode."""
    if step <= 1:
        return {"action_type": "meta.read_runbook", "args": {}}
    elif step >= max_steps - 3:
        return {"action_type": "meta.finish", "args": {}}
    return {"action_type": "meta.noop", "args": {}}


# ---------------------------------------------------------------------------
# Context builder — runbook-pinned sliding window, no duplicate user turns
# ---------------------------------------------------------------------------

def _chars(messages: list[dict]) -> int:
    return sum(len(m["content"]) for m in messages)


def build_messages(
    turn_history: list[dict],
    current_obs_text: str,
) -> list[dict]:
    """Build prompt messages with runbook-pinned sliding window truncation.

    turn_history entries: {"obs_text": str, "completion": str, "is_runbook": bool}
      obs_text   = observation that triggered this completion (user side)
      completion = model output for that step (assistant side)

    Final conversation shape:
      [sys] [user:obs_0][asst:comp_0] ... [user:obs_k][asst:comp_k] [user:current_obs]
    No entry in turn_history represents current_obs — it's only the final user turn.
    """
    # Task brief is always prepended to current_obs so it stays visible even
    # when old context is truncated.
    messages: list[dict] = [{"role": "system", "content": SYSTEM_PROMPT}]

    # Separate the runbook exchange from the rest
    runbook_turn: dict | None = None
    general: list[dict] = []

    for turn in turn_history:
        if turn["is_runbook"] and runbook_turn is None:
            runbook_turn = turn
        else:
            general.append(turn)

    # Pin runbook exchange (truncate only if absurdly large)
    if runbook_turn is not None:
        rb_comp = runbook_turn["completion"][:RUNBOOK_RESPONSE_MAX_CHARS]
        messages.append({"role": "user", "content": runbook_turn["obs_text"]})
        messages.append({"role": "assistant", "content": rb_comp})

    # Fill budget with most-recent general turns (newest first, then reverse)
    fixed_chars = _chars(messages) + len(current_obs_text)
    budget = MAX_PROMPT_CHARS - fixed_chars
    window: list[dict] = []

    for turn in reversed(general[-HISTORY_PAIRS:]):
        pair_chars = len(turn["obs_text"]) + len(turn["completion"])
        if budget - pair_chars < 0:
            break
        window.append(turn)
        budget -= pair_chars

    for turn in reversed(window):
        messages.append({"role": "user", "content": turn["obs_text"]})
        messages.append({"role": "assistant", "content": turn["completion"]})

    # Current observation is always the final user turn (never stored in history)
    messages.append({"role": "user", "content": current_obs_text})
    return messages


# ---------------------------------------------------------------------------
# Observation normaliser — handles both object and dict forms
# ---------------------------------------------------------------------------

def _obs_to_dict(obs: Any) -> dict:
    if isinstance(obs, dict):
        return obs
    fields = ("task_brief", "last_action_result", "step", "steps_remaining",
              "reward", "done", "token_budget_remaining")
    return {f: getattr(obs, f, None) for f in fields}


def _current_obs_text(obs_dict: dict, step: int, task_brief: str) -> str:
    """Format observation, prepending a task reminder so it survives truncation."""
    reminder = f"**Task reminder:** {task_brief[:200]}\n\n"
    return reminder + format_observation(obs_dict, step, obs_dict.get("last_action_result"))


# ---------------------------------------------------------------------------
# Single-episode rollout
# ---------------------------------------------------------------------------

def rollout_once(
    trainer,
    sync_env,
    tokenizer,
    dataset_prompt: str,
    max_steps: int = 15,
    gen_offset: int = 0,
) -> dict:
    """Play one full PM-Ops episode. Returns trajectory + reward signals.

    gen_offset: added to seed so each GRPO generation explores a different env
    episode even when receiving the same prompt (same base seed).
    """
    seed = parse_seed_from_prompt(dataset_prompt)
    if seed is not None:
        result = sync_env.reset(seed=seed + gen_offset)
    else:
        result = sync_env.reset()

    obs = result.observation if hasattr(result, "observation") else result
    obs_dict = _obs_to_dict(obs)
    task_brief: str = obs_dict.get("task_brief") or dataset_prompt

    # Flat trajectory buffers (TRL expects flat lists across all steps)
    prompt_ids: list = []
    completion_ids: list = []
    logprobs: list = []

    # Turn history for context building
    # Each entry: {"obs_text": str, "completion": str, "is_runbook": bool}
    turn_history: list[dict] = []

    # Rollout accumulators
    valid_action_count = 0
    final_score = 0.0
    step = 0
    done = False

    # Runbook-compliance reward tracking
    read_runbook_done  = False
    valid_labels:     set[str] = set()   # org label_taxonomy values
    valid_priorities: set[str] = set()   # org priority_levels
    valid_teams:      set[str] = set()   # org team_map values
    oncall_channels:  set[str] = set()   # org oncall_channels values
    ticket_label:    str | None = None   # label used in create_ticket
    ticket_priority: str | None = None   # priority used in create_ticket
    assigned_team:   str | None = None   # team from assign_ticket
    posted_channels: list[str] = []      # every channel posted to

    while not done and step < max_steps:
        # obs_text for THIS step — stored in history BEFORE stepping
        obs_text = _current_obs_text(obs_dict, step, task_brief)

        messages = build_messages(turn_history, obs_text)
        prompt_text = tokenizer.apply_chat_template(
            messages,
            add_generation_prompt=True,
            tokenize=False,
            enable_thinking=False,
        )

        rollout_out = _generate_no_vllm(trainer, prompt_text, tokenizer)
        prompt_ids.extend(rollout_out["prompt_ids"])
        completion_ids.extend(rollout_out["completion_ids"])
        logprobs.extend(rollout_out["logprobs"])

        completion_text = rollout_out["text"]

        # Parse action; fall back gracefully on parse failure
        parsed = extract_json_action(completion_text)
        is_valid_json = parsed is not None

        if not is_valid_json:
            # Log the raw output on step 0 to diagnose format failures
            if step == 0 and valid_action_count == 0:
                snippet = repr(completion_text[:300])
                print(f"[rollout] step=0 NO JSON — raw output: {snippet}")
            parsed = step_aware_fallback(step, max_steps)
        else:
            valid_action_count += 1

        action_type: str = parsed.get("action_type", "meta.noop")
        args: dict = parsed.get("args", {})

        if action_type == "meta.read_runbook" and is_valid_json:
            read_runbook_done = True

        if action_type == "ticketing.create_ticket" and is_valid_json and ticket_label is None:
            ticket_label    = args.get("label")
            ticket_priority = args.get("priority")

        if action_type == "ticketing.assign_ticket" and is_valid_json and assigned_team is None:
            assigned_team = args.get("team")

        if action_type == "chat.post_message" and is_valid_json:
            ch = args.get("channel", "")
            if ch:
                posted_channels.append(ch)

        # Store the (obs_text, completion) pair BEFORE stepping the env
        # is_runbook marks this turn for pinning in future context windows
        turn_history.append({
            "obs_text": obs_text,
            "completion": completion_text,
            "is_runbook": (action_type == "meta.read_runbook" and is_valid_json),
        })

        # Step the environment — obs_dict now holds the NEXT state
        result = sync_env.step({"action_type": action_type, "args": args})
        new_obs = result.observation if hasattr(result, "observation") else result
        obs_dict = _obs_to_dict(new_obs)

        # Extract full org_config from runbook response (one step after the call)
        last_result = obs_dict.get("last_action_result") or {}
        if action_type == "meta.read_runbook" and last_result.get("ok"):
            data = last_result.get("data") or {}
            if isinstance(data, dict):
                org = data.get("org_config") or {}
                valid_labels     = set(org.get("label_taxonomy", {}).values())
                valid_priorities = set(org.get("priority_levels", []))
                valid_teams      = set(org.get("team_map", {}).values())
                oncall_channels  = set(org.get("oncall_channels", {}).values())

        done = bool(getattr(result, "done", obs_dict.get("done", False)))
        final_score = float(getattr(result, "reward", obs_dict.get("reward", 0.0)))
        step += 1

    from training.rewards import compute_rollout_reward
    combined = compute_rollout_reward(
        read_runbook_done  = read_runbook_done,
        valid_labels       = valid_labels,
        valid_priorities   = valid_priorities,
        valid_teams        = valid_teams,
        oncall_channels    = oncall_channels,
        ticket_label       = ticket_label,
        ticket_priority    = ticket_priority,
        assigned_team      = assigned_team,
        posted_channels    = posted_channels,
        env_score          = final_score,
        valid_json_count   = valid_action_count,
    )

    print(
        f"[rollout] steps={step} env={final_score:.3f} "
        f"label={'✓' if ticket_label and ticket_label in valid_labels else '✗' if ticket_label else '-'} "
        f"priority={'✓' if ticket_priority and ticket_priority in valid_priorities else '✗' if ticket_priority else '-'} "
        f"team={'✓' if assigned_team and assigned_team in valid_teams else '✗' if assigned_team else '-'} "
        f"channel={'✓' if any(ch in oncall_channels for ch in posted_channels) else '✗' if posted_channels else '-'} "
        f"→ reward={combined:.3f}"
    )
    return {
        "prompt_ids":     prompt_ids,
        "completion_ids": completion_ids,
        "logprobs":       logprobs,
        "reward":         combined,
    }


# ---------------------------------------------------------------------------
# GRPOTrainer-compatible rollout function (factory)
# ---------------------------------------------------------------------------

def make_rollout_func(sync_env, tokenizer, max_steps: int = 15):
    """Bind env + tokenizer; return the function GRPOTrainer calls each batch.

    max_steps per task:
      triage           → 15  (solvable in 5, cap gives room for exploration)
      incident_routing → 20
      release_notes    → 30
      dep_update       → 30
    """
    def rollout_func(prompts: list[str], trainer=None) -> dict:
        out: dict[str, list] = {
            "prompt_ids":     [],
            "completion_ids": [],
            "logprobs":       [],
            "reward":         [],
        }
        # Track how many times each unique prompt has appeared so we can pass
        # a gen_offset — ensures repeated prompts (num_generations > 1) hit
        # different env seeds and produce different rollouts.
        prompt_seen: dict[str, int] = {}
        for prompt_text in prompts:
            gen_offset = prompt_seen.get(prompt_text, 0)
            prompt_seen[prompt_text] = gen_offset + 1
            episode = rollout_once(
                trainer=trainer,
                sync_env=sync_env,
                tokenizer=tokenizer,
                dataset_prompt=prompt_text,
                max_steps=max_steps,
                gen_offset=gen_offset,
            )
            for k in out:
                out[k].append(episode[k])
        return out

    return rollout_func