code/self_healing_finetuning/orchestrator/promotion_blocker_controller.py

from __future__ import annotations

import json
import re
from collections import Counter
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from config.profile_loader import load_strategy_profile
from data_pipeline.pairwise_preference_memory import candidate_critical_failed, candidate_lost, failure_bucket_for_prediction
from n21.config import write_json
from n21.settings import SHFT_WORKSPACE_ROOT


CONTROLLER_SCHEMA_VERSION = "shft_promotion_blocker_controller_v1"


def utc_now() -> str:
    return datetime.now(UTC).replace(microsecond=0).isoformat().replace("+00:00", "Z")


def read_json(path: Path) -> dict[str, Any] | None:
    if not path.exists():
        return None
    return json.loads(path.read_text(encoding="utf-8-sig"))


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    if not path.exists():
        return []
    rows: list[dict[str, Any]] = []
    with path.open("r", encoding="utf-8-sig") as handle:
        for line in handle:
            if line.strip():
                rows.append(json.loads(line))
    return rows


def preference_round_count(run_id: str) -> int:
    return len(re.findall(r"_pref_", run_id))


def _number(value: Any, default: float = 0.0) -> float:
    try:
        return float(value)
    except (TypeError, ValueError):
        return default


def _failed_checks(gate: dict[str, Any] | None) -> list[str]:
    if not gate:
        return ["model_quality_gate_missing"]
    checks = gate.get("checks", {})
    failed: list[str] = []
    if isinstance(checks, dict):
        for name, payload in checks.items():
            if isinstance(payload, dict) and payload.get("ok") is False:
                failed.append(str(name))
    if not failed:
        for error in gate.get("errors", []):
            failed.append(str(error).split(":", 1)[0])
    return failed


def _bucket_impact(
    *,
    predictions: list[dict[str, Any]],
    human_review: dict[str, Any] | None,
    profile: dict[str, Any],
) -> list[dict[str, Any]]:
    weights = profile.get("promotion_impact_weights", {})
    human_failure_ids = set()
    if human_review:
        for key in ("critical_failure_ids", "failed_sample_ids", "failure_ids"):
            values = human_review.get(key)
            if isinstance(values, list):
                human_failure_ids.update(str(item) for item in values)

    impact: dict[str, dict[str, Any]] = {}
    for prediction in predictions:
        bucket = failure_bucket_for_prediction(prediction)
        item = impact.setdefault(
            bucket,
            {
                "failure_bucket": bucket,
                "pairwise_loss_count": 0,
                "critical_regression_count": 0,
                "human_failure_count": 0,
                "sample_ids": [],
                "promotion_impact_score": 0.0,
            },
        )
        pred_id = str(prediction.get("id") or prediction.get("prompt") or "")
        if candidate_lost(prediction):
            item["pairwise_loss_count"] += 1
            item["promotion_impact_score"] += _number(weights.get("pairwise_loss"), 3.0)
            item["sample_ids"].append(pred_id)
        baseline_critical = None
        baseline_score = prediction.get("baseline_score")
        if isinstance(baseline_score, dict):
            baseline_critical = baseline_score.get("critical_pass")
        if candidate_critical_failed(prediction) and baseline_critical is not False:
            item["critical_regression_count"] += 1
            item["promotion_impact_score"] += _number(weights.get("critical_regression"), 5.0)
            if pred_id not in item["sample_ids"]:
                item["sample_ids"].append(pred_id)
        if pred_id in human_failure_ids:
            item["human_failure_count"] += 1
            item["promotion_impact_score"] += _number(weights.get("human_failure"), 4.0)
            if pred_id not in item["sample_ids"]:
                item["sample_ids"].append(pred_id)

    rows = [row for row in impact.values() if row["promotion_impact_score"] > 0]
    rows.sort(key=lambda row: (-row["promotion_impact_score"], row["failure_bucket"]))
    for row in rows:
        row["promotion_impact_score"] = round(float(row["promotion_impact_score"]), 4)
        row["sample_ids"] = row["sample_ids"][:25]
    return rows


def choose_strategy(
    *,
    failed_checks: list[str],
    paired_eval: dict[str, Any] | None,
    human_review: dict[str, Any] | None,
    preference_rounds: int,
    max_preference_rounds: int,
    bucket_impacts: list[dict[str, Any]],
) -> tuple[str, list[str], bool]:
    reasons: list[str] = []
    failset = set(failed_checks)
    improvement = (paired_eval or {}).get("improvement", {})
    aggregate_delta = _number(improvement.get("aggregate_abs"))
    critical_delta = _number(improvement.get("critical_pass_rate_abs"))
    loss_rate = _number(improvement.get("pairwise_loss_rate"))
    win_rate = _number(improvement.get("pairwise_win_rate"))
    human_critical_failures = int(_number((human_review or {}).get("critical_failures"), 0))

    if preference_rounds >= max_preference_rounds:
        reasons.append(f"preference round cap reached: {preference_rounds}/{max_preference_rounds}")
        return "hold", reasons, True
    if "critical_pass_not_regressed" in failset or critical_delta < 0:
        reasons.append(f"critical pass regressed: {critical_delta:.4f}")
        return "critical_safety_repair", reasons, False
    if human_critical_failures > 0 or "human_review_critical_failures" in failset:
        reasons.append(f"human review critical failures: {human_critical_failures}")
        return "human_failure_repair", reasons, False
    if "pairwise_loss_rate" in failset or loss_rate > 0.10:
        reasons.append(f"pairwise loss rate too high: {loss_rate:.4f}")
        return "hard_negative_dpo", reasons, False
    if "pairwise_win_rate" in failset or (paired_eval and win_rate < 0.40):
        reasons.append(f"pairwise win rate too low: {win_rate:.4f}")
        return "answer_quality_repair", reasons, False
    if "aggregate_delta_absolute" in failset or (paired_eval and aggregate_delta < 0.05):
        reasons.append(f"aggregate delta shortfall: {aggregate_delta:.4f}")
        return "domain_sft", reasons, False
    if not paired_eval:
        reasons.append("paired eval is missing")
        return "eval_only", reasons, False
    if bucket_impacts:
        reasons.append(f"residual impacted bucket: {bucket_impacts[0]['failure_bucket']}")
        return "hard_negative_dpo", reasons, False
    reasons.append("no actionable blocker found")
    return "hold", reasons, True


def build_controller_decision(
    *,
    run_id: str,
    release_id: str | None = None,
    asset_class: str,
    role: str,
    max_preference_rounds: int = 5,
    output_path: Path | None = None,
) -> dict[str, Any]:
    run_path = SHFT_WORKSPACE_ROOT / "runs" / run_id
    eval_dir = run_path / "eval"
    profile = load_strategy_profile(asset_class, role)
    gate = read_json(eval_dir / "model_quality_gate.json")
    paired_eval = read_json(eval_dir / "paired_eval_report.json")
    human_review = read_json(eval_dir / "human_spot_check_report.json")
    predictions = read_jsonl(eval_dir / "paired_predictions.jsonl")
    failed = _failed_checks(gate)
    rounds = preference_round_count(run_id)
    bucket_impacts = _bucket_impact(predictions=predictions, human_review=human_review, profile=profile)
    next_strategy, reasons, should_hold = choose_strategy(
        failed_checks=failed,
        paired_eval=paired_eval,
        human_review=human_review,
        preference_rounds=rounds,
        max_preference_rounds=max_preference_rounds,
        bucket_impacts=bucket_impacts,
    )
    result = {
        "ok": True,
        "schema_version": CONTROLLER_SCHEMA_VERSION,
        "run_id": run_id,
        "release_id": release_id,
        "asset_class": asset_class,
        "role": role,
        "profile_id": profile.get("profile_id"),
        "profile_path": profile.get("profile_path"),
        "quality_gate_ok": bool(gate and gate.get("ok") is True),
        "failed_promotion_checks": failed,
        "next_strategy": next_strategy,
        "should_hold": should_hold,
        "hold_reason": "; ".join(reasons) if should_hold else None,
        "strategy_reasons": reasons,
        "preference_round_count": rounds,
        "max_preference_rounds": max_preference_rounds,
        "promotion_impact": {
            "top_failure_buckets": bucket_impacts[:10],
            "pairwise_losses": int(_number((paired_eval or {}).get("improvement", {}).get("losses"), 0)),
            "pairwise_loss_rate": _number((paired_eval or {}).get("improvement", {}).get("pairwise_loss_rate")),
            "pairwise_win_rate": _number((paired_eval or {}).get("improvement", {}).get("pairwise_win_rate")),
            "critical_pass_delta": _number((paired_eval or {}).get("improvement", {}).get("critical_pass_rate_abs")),
            "human_critical_failures": int(_number((human_review or {}).get("critical_failures"), 0)),
        },
        "allowed_training_modes": profile.get("training_modes"),
        "created_at": utc_now(),
    }
    out = output_path or run_path / "autopilot" / "promotion_blocker_decision.json"
    result["output_path"] = str(out)
    if should_hold:
        result["hold_report_path"] = str(run_path / "autopilot" / "autopilot_hold_report.json")
    write_json(out, result)
    if should_hold:
        write_json(run_path / "autopilot" / "autopilot_hold_report.json", result)
    return result