code/self_healing_finetuning/orchestrator/stall_breakout.py

from __future__ import annotations

from pathlib import Path
from typing import Any

from data_pipeline.learning_pdf_to_jsonl import convert_learning_tree
from data_pipeline.live_source_discovery import discover_public_sources
from data_pipeline.reasoning_data_generation import generate_grounded_reasoning_examples_from_intake_records
from data_pipeline.source_intake import intake_public_sources, load_policy
from data_pipeline.training_data_validation import validate_training_data
from n21.config import write_json
from n21.settings import CONFIG_ROOT, SHFT_WORKSPACE_ROOT
from observability.audit_log import utc_now


TRAINABLE_SUFFIXES = {".pdf", ".jsonl", ".hf_finetune.jsonl", ".normalized.json"}


def _run_dir(run_id: str) -> Path:
    return SHFT_WORKSPACE_ROOT / "runs" / run_id


def _trainable_paths(records: list[dict[str, Any]]) -> list[str]:
    paths: list[str] = []
    for record in records:
        path = str(record.get("training_path") or "")
        if not path:
            continue
        lower = path.lower()
        if lower.endswith(".hf_finetune.jsonl") or lower.endswith(".normalized.json") or Path(path).suffix.lower() in TRAINABLE_SUFFIXES:
            paths.append(path)
    return paths


def _urls_from_records(records: list[dict[str, Any]]) -> set[str]:
    urls: set[str] = set()
    for record in records:
        url = str(record.get("url") or record.get("source_url") or "").strip().lower()
        if url:
            urls.add(url)
    return urls


def _quality_errors(run_path: Path) -> list[str]:
    path = run_path / "eval" / "model_quality_gate.json"
    if not path.exists():
        return []
    try:
        import json

        report = json.loads(path.read_text(encoding="utf-8-sig"))
    except Exception:
        return []
    return [str(item) for item in report.get("errors", [])]


def run_stall_breakout(
    *,
    run_id: str,
    asset_class: str,
    role: str,
    train_source: Path,
    quality_gate_exit_code: int,
    catalog_path: Path | None = None,
    policy_path: Path | None = None,
    max_sources: int | None = None,
    apply_quarantine: bool = False,
    intake_root: Path | None = None,
    training_root: Path | None = None,
) -> dict[str, Any]:
    """Run the Step 0 data breakout that follows a failed model-quality gate.

    This is intentionally not a certification bypass. It downloads public
    sources according to source_policy.yaml, promotes only train-allowed files
    into data/learning, validates the training corpus, and records whether the
    new material is actually trainable by the current Step 0b converters.
    """

    run_path = _run_dir(run_id)
    breakout_dir = run_path / "stall_breakout"
    validation_dir = breakout_dir / "training_data_validation"
    backup_dir = breakout_dir / "quarantine_backup"
    breakout_dir.mkdir(parents=True, exist_ok=True)

    quality_errors = _quality_errors(run_path)
    policy = load_policy(policy_path)
    intake_manifest = intake_public_sources(
        asset_class=asset_class,
        role=role,
        catalog_path=catalog_path,
        policy_path=policy_path,
        max_sources=max_sources,
        promote_approved=True,
        intake_root=intake_root,
        training_root=training_root,
        quality_errors=quality_errors,
    )
    records = list(intake_manifest.get("records", []))
    trainable_paths = _trainable_paths(records)
    live_discovery_result: dict[str, Any] | None = None
    live_intake_manifest: dict[str, Any] | None = None
    live_attempts: list[dict[str, Any]] = []
    reasoning_generation_report: dict[str, Any] | None = None

    discovery_policy = policy.get("live_discovery", {})
    max_new_sources = int(max_sources or policy.get("stall_breakout", {}).get("max_new_sources_per_round", 10))
    if len(trainable_paths) < max_new_sources:
        if bool(discovery_policy.get("enabled", True)):
            max_attempts = int(
                policy.get("stall_breakout", {}).get(
                    "max_source_discovery_attempts",
                    discovery_policy.get("max_source_discovery_attempts", 4),
                )
            )
            excluded_urls = _urls_from_records(records)
            for attempt in range(1, max_attempts + 1):
                if len(trainable_paths) >= max_new_sources:
                    break
                live_dir_name = "live_discovery" if attempt == 1 else f"live_discovery_attempt_{attempt:02d}"
                live_dir = breakout_dir / live_dir_name
                live_discovery_result = discover_public_sources(
                    asset_class=asset_class,
                    role=role,
                    policy=policy,
                    output_dir=live_dir,
                    exclude_urls=excluded_urls,
                    quality_errors=quality_errors,
                    discovery_attempt=attempt,
                )
                discovered_catalog = live_dir / "live_discovered_public_source_catalog.json"
                remaining = max(1, max_new_sources - len(trainable_paths))
                live_intake_manifest = intake_public_sources(
                    asset_class=asset_class,
                    role=role,
                    catalog_path=discovered_catalog,
                    policy_path=policy_path,
                    max_sources=remaining,
                    promote_approved=True,
                    intake_root=(live_dir / "intake"),
                    training_root=training_root,
                    quality_errors=quality_errors,
                )
                new_records = list(live_intake_manifest.get("records", []))
                records.extend(new_records)
                excluded_urls.update(_urls_from_records(new_records))
                trainable_paths = _trainable_paths(records)
                live_attempts.append(
                    {
                        "attempt": attempt,
                        "discovery_dir": str(live_dir),
                        "candidate_count": (live_discovery_result or {}).get("manifest", {}).get("candidate_count", 0),
                        "catalog_path": (live_discovery_result or {}).get("manifest", {}).get("catalog_path"),
                        "intake_source_count": live_intake_manifest.get("source_count", 0),
                        "intake_downloaded_count": live_intake_manifest.get("downloaded_count", 0),
                        "intake_training_eligible_count": live_intake_manifest.get("training_eligible_count", 0),
                        "intake_verification_eligible_count": live_intake_manifest.get("validation_eligible_count", 0),
                        "ai_training_certified_count": live_intake_manifest.get("ai_training_certified_count", 0),
                        "ai_verification_certified_count": live_intake_manifest.get("ai_verification_certified_count", 0),
                        "ai_rejected_count": live_intake_manifest.get("ai_rejected_count", 0),
                        "content_ai_training_certified_count": live_intake_manifest.get("content_ai_training_certified_count", 0),
                        "content_ai_rejected_count": live_intake_manifest.get("content_ai_rejected_count", 0),
                        "download_failed_count": sum(1 for record in new_records if record.get("decision") == "download_failed"),
                        "already_in_training_count": sum(1 for record in new_records if record.get("already_in_training")),
                        "review_required_count": live_intake_manifest.get("review_required_count", 0),
                        "total_trainable_new_source_count": len(trainable_paths),
                        "errors": (live_discovery_result or {}).get("manifest", {}).get("errors", []),
                    }
                )
                if not new_records and not (live_discovery_result or {}).get("manifest", {}).get("candidate_count", 0):
                    # The next attempt would use broader retry terms, but if the
                    # provider returns no candidates at all after query expansion
                    # is already active, stop and surface that blocker.
                    if attempt > 1:
                        break

    if not trainable_paths:
        reasoning_max_records = int(policy.get("stall_breakout", {}).get("max_breakout_reasoning_records", 300))
        reasoning_output = train_source / f"synthetic_{asset_class}_{role}_critical_reasoning.hf_finetune.jsonl"
        reasoning_generation_report = generate_grounded_reasoning_examples_from_intake_records(
            asset_class=asset_class,
            role=role,
            records=records,
            output_path=reasoning_output,
            max_records=reasoning_max_records,
            policy_path=policy_path,
        )
        if reasoning_generation_report.get("ok"):
            trainable_paths.append(str(reasoning_output))

    conversion_report: dict[str, Any] | None = None
    conversion_error: str | None = None
    if any(not str(path).lower().endswith(".jsonl") for path in trainable_paths):
        try:
            conversion_report = convert_learning_tree(
                asset_class=asset_class,
                role=role,
                source_type="pdf",
                skip_existing=False,
            )
        except Exception as exc:  # pragma: no cover - reported as artifact, not hidden
            conversion_error = str(exc)

    validation_report = validate_training_data(
        source=train_source,
        output_dir=validation_dir,
        backup_dir=backup_dir,
        apply_quarantine=apply_quarantine,
    )

    non_trainable_approved = [
        str(record.get("training_path"))
        for record in records
        if record.get("train_allowed") and record.get("training_path") and str(record.get("training_path")) not in trainable_paths
    ]
    approved_count = sum(1 for record in records if record.get("train_allowed"))
    verification_count = sum(1 for record in records if record.get("verification_allowed"))
    ai_training_certified_count = sum(1 for record in records if (record.get("ai_certification") or {}).get("training_eligible"))
    ai_verification_certified_count = sum(
        1 for record in records if (record.get("ai_certification") or {}).get("verification_eligible")
    )
    ai_rejected_count = sum(1 for record in records if (record.get("ai_certification") or {}).get("intended_use") == "reject")
    content_ai_training_certified_count = sum(
        1 for record in records if (record.get("content_ai_certification") or {}).get("training_eligible")
    )
    content_ai_verification_certified_count = sum(
        1 for record in records if (record.get("content_ai_certification") or {}).get("verification_eligible")
    )
    content_ai_rejected_count = sum(
        1 for record in records if (record.get("content_ai_certification") or {}).get("intended_use") == "reject"
    )
    generated_reasoning_count = int((reasoning_generation_report or {}).get("generated_count") or 0)
    downloaded_count = sum(1 for record in records if record.get("downloaded"))
    review_required_count = sum(1 for record in records if record.get("requires_human_review"))
    blockers: list[str] = []
    if approved_count == 0 and generated_reasoning_count == 0:
        blockers.append("no AI-certified, policy-approved training sources were found for this asset/role")
    if not trainable_paths and generated_reasoning_count == 0:
        blockers.append("no newly AI-certified training sources are trainable by current Step 0b converters")
    if verification_count and not approved_count and generated_reasoning_count == 0:
        blockers.append("only verification material was found; verification material measures quality but does not improve the model")
    if non_trainable_approved:
        blockers.append("some approved downloads need a converter before they can become training records")
    if not validation_report.get("ok"):
        blockers.append("training data validation failed")
    if conversion_error:
        blockers.append("Step 0b conversion failed after breakout intake")

    next_actions = [
        "Review stall_breakout/source_intake_manifest.json and license_manifest.json.",
        "Use only records with train_allowed=true and ai_certification.training_eligible=true for training.",
        "Use records with verification_allowed=true only for evaluation/model-judge/human spot checks.",
        "If no trainable sources were added, add high-signal PDF/HTML/JSONL sources with explicit case-study, checklist, red-flag, or pass/fail reasoning.",
        "Start a fresh run/version after corpus validation passes; do not reuse the failed RUN_ID unless explicitly allowed.",
    ]

    plan = {
        "schema_version": "stall_breakout_plan_v1",
        "run_id": run_id,
        "asset_class": asset_class,
        "role": role,
        "train_source": str(train_source),
        "quality_gate_exit_code": quality_gate_exit_code,
        "policy_path": str(policy_path or (CONFIG_ROOT / "data" / "source_policy.yaml")),
        "catalog_path": str(catalog_path or (CONFIG_ROOT / "data" / "public_source_catalog.json")),
        "policy_summary": {
            "auto_download_public_sources": bool(policy.get("auto_download_public_sources", True)),
            "promote_approved_sources_to_training": bool(policy.get("promote_approved_sources_to_training", True)),
            "approved_license_statuses": policy.get("approved_license_statuses", []),
            "review_license_statuses": policy.get("review_license_statuses", []),
            "blocked_license_statuses": policy.get("blocked_license_statuses", []),
            "source_quality_certification": policy.get("source_quality_certification", {}),
        },
        "intake": {
            "source_count": len(records),
            "catalog_source_count": intake_manifest.get("source_count", 0),
            "downloaded_count": downloaded_count,
            "approved_for_training_count": approved_count,
            "approved_for_verification_count": verification_count,
            "ai_training_certified_count": ai_training_certified_count,
            "ai_verification_certified_count": ai_verification_certified_count,
            "ai_rejected_count": ai_rejected_count,
            "content_ai_training_certified_count": content_ai_training_certified_count,
            "content_ai_verification_certified_count": content_ai_verification_certified_count,
            "content_ai_rejected_count": content_ai_rejected_count,
            "review_required_count": review_required_count,
            "trainable_new_source_count": len(trainable_paths),
            "trainable_new_sources": trainable_paths,
            "non_trainable_approved_sources": non_trainable_approved,
            "generated_reasoning_count": generated_reasoning_count,
        },
        "live_discovery": {
            "attempted": bool(live_attempts),
            "attempt_count": len(live_attempts),
            "attempts": live_attempts,
            "candidate_count": (live_discovery_result or {}).get("manifest", {}).get("candidate_count", 0),
            "catalog_path": (live_discovery_result or {}).get("manifest", {}).get("catalog_path"),
            "intake_source_count": (live_intake_manifest or {}).get("source_count", 0),
            "intake_downloaded_count": (live_intake_manifest or {}).get("downloaded_count", 0),
            "intake_training_eligible_count": (live_intake_manifest or {}).get("training_eligible_count", 0),
            "intake_verification_eligible_count": (live_intake_manifest or {}).get("validation_eligible_count", 0),
            "ai_training_certified_count": (live_intake_manifest or {}).get("ai_training_certified_count", 0),
            "ai_verification_certified_count": (live_intake_manifest or {}).get("ai_verification_certified_count", 0),
            "ai_rejected_count": (live_intake_manifest or {}).get("ai_rejected_count", 0),
            "content_ai_training_certified_count": (live_intake_manifest or {}).get("content_ai_training_certified_count", 0),
            "content_ai_verification_certified_count": (live_intake_manifest or {}).get("content_ai_verification_certified_count", 0),
            "content_ai_rejected_count": (live_intake_manifest or {}).get("content_ai_rejected_count", 0),
            "errors": (live_discovery_result or {}).get("manifest", {}).get("errors", []),
        },
        "conversion": {
            "attempted": any(not str(path).lower().endswith(".jsonl") for path in trainable_paths),
            "ok": conversion_error is None,
            "error": conversion_error,
            "report": conversion_report,
        },
        "reasoning_generation": reasoning_generation_report
        or {
            "attempted": False,
            "ok": False,
            "generated_count": 0,
            "status": "not_needed",
        },
        "validation": {
            "ok": bool(validation_report.get("ok")),
            "record_count": validation_report.get("record_count", 0),
            "schema_error_count": validation_report.get("schema_error_count", 0),
            "conflict_count": validation_report.get("conflict_count", 0),
            "report_path": str(validation_dir / "training_data_validation_report.json"),
        },
        "status": "ready_for_next_training_run" if not blockers else "blocked_after_breakout",
        "blockers": blockers,
        "next_actions": next_actions,
        "created_at": utc_now(),
    }
    write_json(breakout_dir / "stall_breakout_plan.json", plan)
    write_json(breakout_dir / "source_intake_manifest.json", intake_manifest)
    if live_intake_manifest is not None:
        write_json(breakout_dir / "live_source_intake_manifest.json", live_intake_manifest)
    write_json(
        breakout_dir / "license_manifest.json",
        {"records": records, "created_at": utc_now()},
    )
    return plan