OBLITERATUS/obliteratus/community.py

"""Community contribution system for crowdsourced paper data.

Enables users to contribute anonymized experiment results to the shared
paper dataset. Unlike telemetry (which is fire-and-forget to a remote
endpoint), contributions are saved as local JSON files that can be
submitted via pull request to the community results repository.

Usage:
    from obliteratus.community import save_contribution

    # After running a pipeline:
    path = save_contribution(
        pipeline,
        model_name="meta-llama/Llama-2-7b-chat-hf",  # public model ID
        notes="Ran on A100 with default prompts",
    )
    # Generates: community_results/llama2-7b_advanced_20260227_143052.json
"""

from __future__ import annotations

import hashlib
import json
import logging
import re
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from obliteratus.telemetry import (
    _direction_stats,
    _extract_excise_details,
    _extract_prompt_counts,
    _extract_stage_durations,
    _get_peak_vram,
    _safe_float,
    build_report,
)

logger = logging.getLogger(__name__)

# Schema version for community contributions (extends telemetry schema v2)
CONTRIBUTION_SCHEMA_VERSION = 1

# Default output directory for contributions
DEFAULT_CONTRIB_DIR = "community_results"


def _model_short_name(model_name: str) -> str:
    """Extract a filesystem-safe short name from a HuggingFace model ID."""
    # "meta-llama/Llama-2-7b-chat-hf" -> "llama-2-7b-chat-hf"
    name = model_name.split("/")[-1].lower()
    name = re.sub(r"[^a-z0-9\-]", "-", name)
    name = re.sub(r"-+", "-", name).strip("-")
    return name[:60]  # cap length


def _config_fingerprint(config: dict[str, Any]) -> str:
    """Deterministic short hash of the method configuration."""
    canonical = json.dumps(config, sort_keys=True, default=str)
    return hashlib.sha256(canonical.encode()).hexdigest()[:8]


def save_contribution(
    pipeline,
    *,
    model_name: str,
    notes: str = "",
    output_dir: str | Path = DEFAULT_CONTRIB_DIR,
    informed_report=None,
) -> Path:
    """Save a contribution record from a completed pipeline run.

    Unlike telemetry, this:
    - Includes the public model name (for aggregation by model)
    - Saves locally (not sent remotely)
    - Uses a human-readable filename
    - Includes a config fingerprint for deduplication
    - Is always explicit (no silent opt-in)

    Args:
        pipeline: A completed AbliterationPipeline instance.
        model_name: HuggingFace model ID (e.g., "meta-llama/Llama-2-7b-chat-hf").
        notes: Optional free-text notes about the run.
        output_dir: Directory to save contribution files.
        informed_report: Optional InformedPipelineReport for informed pipeline runs.

    Returns:
        Path to the saved contribution JSON file.
    """
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Build the base telemetry report (reuse existing schema)
    summary = pipeline.handle.summary()

    config_keys = [
        "n_directions", "norm_preserve", "regularization",
        "refinement_passes", "project_biases", "use_chat_template",
        "use_whitened_svd", "true_iterative_refinement",
        "use_jailbreak_contrast", "layer_adaptive_strength",
        "attention_head_surgery", "safety_neuron_masking",
        "per_expert_directions", "use_sae_features", "invert_refusal",
        "project_embeddings", "embed_regularization",
        "activation_steering", "steering_strength",
        "expert_transplant", "transplant_blend",
        "reflection_strength",
    ]
    method_config = {}
    for key in config_keys:
        val = getattr(pipeline, key, None)
        if val is not None:
            method_config[key] = val

    # Extract analysis insights if informed report is available
    analysis_insights = None
    informed_extras = None
    if informed_report is not None:
        try:
            from obliteratus.telemetry import _extract_analysis_insights
            analysis_insights = _extract_analysis_insights(informed_report)
            informed_extras = {}
            if hasattr(informed_report, "ouroboros_passes"):
                informed_extras["ouroboros_passes"] = informed_report.ouroboros_passes
            if hasattr(informed_report, "final_refusal_rate"):
                informed_extras["final_refusal_rate"] = _safe_float(
                    informed_report.final_refusal_rate
                )
        except Exception:
            logger.debug("Failed to extract analysis insights from informed report", exc_info=True)

    base_report = build_report(
        architecture=summary.get("architecture", "unknown"),
        num_layers=summary.get("num_layers", 0),
        num_heads=summary.get("num_heads", 0),
        hidden_size=summary.get("hidden_size", 0),
        total_params=summary.get("total_params", 0),
        method=pipeline.method,
        method_config=method_config,
        quality_metrics=pipeline._quality_metrics,
        stage_durations=_extract_stage_durations(pipeline),
        strong_layers=pipeline._strong_layers,
        direction_stats=_direction_stats(pipeline),
        excise_details=_extract_excise_details(pipeline),
        prompt_counts=_extract_prompt_counts(pipeline),
        gpu_memory=_get_peak_vram(),
        analysis_insights=analysis_insights,
        informed_extras=informed_extras,
    )

    # Wrap in community contribution envelope
    timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    contribution = {
        "contribution_schema_version": CONTRIBUTION_SCHEMA_VERSION,
        "timestamp": timestamp,
        "model_name": model_name,
        "config_fingerprint": _config_fingerprint(method_config),
        "notes": notes,
        "telemetry": base_report,
    }

    # Generate filename
    short_name = _model_short_name(model_name)
    method = pipeline.method
    ts_short = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
    filename = f"{short_name}_{method}_{ts_short}.json"
    filepath = output_dir / filename

    filepath.write_text(json.dumps(contribution, indent=2, default=str))
    logger.info("Community contribution saved: %s", filepath)
    return filepath


def load_contributions(
    contrib_dir: str | Path = DEFAULT_CONTRIB_DIR,
) -> list[dict[str, Any]]:
    """Load all contribution records from a directory.

    Args:
        contrib_dir: Directory containing contribution JSON files.

    Returns:
        List of parsed contribution records, sorted by timestamp.
    """
    contrib_dir = Path(contrib_dir)
    if not contrib_dir.exists():
        return []

    records = []
    for path in sorted(contrib_dir.glob("*.json")):
        try:
            data = json.loads(path.read_text())
            if "contribution_schema_version" in data:
                data["_source_file"] = str(path)
                records.append(data)
        except (json.JSONDecodeError, OSError) as e:
            logger.warning("Skipping invalid contribution file %s: %s", path, e)

    records.sort(key=lambda r: r.get("timestamp", ""))
    return records


def aggregate_results(
    records: list[dict[str, Any]],
) -> dict[str, dict[str, Any]]:
    """Aggregate contribution records into per-model, per-method summaries.

    Groups results by (model_name, method) and computes summary statistics
    for key metrics (refusal_rate, perplexity, coherence).

    Returns:
        Nested dict: {model_name: {method: {metric: {mean, std, n, values}}}}
    """
    import statistics

    groups: dict[tuple[str, str], list[dict]] = {}

    for record in records:
        model = record.get("model_name", "unknown")
        telemetry = record.get("telemetry", {})
        method = telemetry.get("method", "unknown")
        metrics = telemetry.get("quality_metrics", {})

        key = (model, method)
        if key not in groups:
            groups[key] = []
        groups[key].append(metrics)

    results: dict[str, dict[str, Any]] = {}
    for (model, method), metric_list in groups.items():
        if model not in results:
            results[model] = {}

        summary: dict[str, Any] = {"n_runs": len(metric_list)}

        for metric_name in ["refusal_rate", "perplexity", "coherence"]:
            values = [
                m[metric_name]
                for m in metric_list
                if metric_name in m and m[metric_name] is not None
            ]
            if values:
                summary[metric_name] = {
                    "mean": round(statistics.mean(values), 4),
                    "std": round(statistics.stdev(values), 4) if len(values) > 1 else 0.0,
                    "n": len(values),
                    "min": round(min(values), 4),
                    "max": round(max(values), 4),
                }

        results[model][method] = summary

    return results


def generate_latex_table(
    aggregated: dict[str, dict[str, Any]],
    methods: list[str] | None = None,
    metric: str = "refusal_rate",
) -> str:
    """Generate a LaTeX table from aggregated community results.

    Args:
        aggregated: Output of aggregate_results().
        methods: Methods to include (default: all found).
        metric: Which metric to display (default: refusal_rate).

    Returns:
        LaTeX table source string.
    """
    if methods is None:
        all_methods: set[str] = set()
        for model_data in aggregated.values():
            all_methods.update(model_data.keys())
        methods = sorted(all_methods)

    # Build header
    method_cols = " & ".join(f"\\textbf{{{m}}}" for m in methods)
    header = f"\\textbf{{Model}} & {method_cols} \\\\"

    lines = [
        "\\begin{tabular}{@{}l" + "c" * len(methods) + "@{}}",
        "\\toprule",
        header,
        "\\midrule",
    ]

    for model in sorted(aggregated.keys()):
        model_data = aggregated[model]
        short = model.split("/")[-1] if "/" in model else model

        cells = []
        for method in methods:
            if method in model_data and metric in model_data[method]:
                stats = model_data[method][metric]
                mean = stats["mean"]
                n = stats["n"]
                if stats["std"] > 0 and n > 1:
                    cells.append(f"{mean:.1f}$\\pm${stats['std']:.1f} ({n})")
                else:
                    cells.append(f"{mean:.1f} ({n})")
            else:
                cells.append("---")

        row = f"{short} & " + " & ".join(cells) + " \\\\"
        lines.append(row)

    lines.extend(["\\bottomrule", "\\end{tabular}"])
    return "\n".join(lines)