src/forge_lite/output.py - master - zhmax/FORGE_LITE

src/forge_lite/output.py
363 строки14 KB
wearetyomsmnv
add openai, and gigaturbo
26 май 2026, 18:25
26 май 2026, 18:25e2591b3
from __future__ import annotations

import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich import box
from rich.text import Text

from .metrics import Metrics
from .recipe import RecipeSpec
from .runner import RunResult

console = Console()


# ---------------------------------------------------------------------------
# Startup banner
# ---------------------------------------------------------------------------

def print_startup_banner(
    recipes: list[RecipeSpec],
    limit: Optional[int] = None,
    workers: Optional[int] = None,
    combinations: Optional[int] = None,
    models: Optional[list[str]] = None,
) -> None:
    # ASCII logo
    _LOGO = (
        "[bold cyan]███████╗ ██████╗ ██████╗  ██████╗ ███████╗   ██╗     ██╗████████╗███████╗[/]\n"
        "[bold cyan]██╔════╝██╔═══██╗██╔══██╗██╔════╝ ██╔════╝   ██║     ██║╚══██╔══╝██╔════╝[/]\n"
        "[bold cyan]█████╗  ██║   ██║██████╔╝██║  ███╗█████╗     ██║     ██║   ██║   █████╗  [/]\n"
        "[bold cyan]██╔══╝  ██║   ██║██╔══██╗██║   ██║██╔══╝     ██║     ██║   ██║   ██╔══╝  [/]\n"
        "[bold cyan]██║     ╚██████╔╝██║  ██║╚██████╔╝███████╗   ███████╗██║   ██║   ███████╗[/]\n"
        "[bold cyan]╚═╝      ╚═════╝ ╚═╝  ╚═╝ ╚═════╝ ╚══════╝   ╚══════╝╚═╝   ╚═╝   ╚══════╝[/]\n"
        "[dim]lightweight llm attack runner[/]"
    )
    console.print()
    console.print(_LOGO)
    console.print()

    # Attack list table
    t = Table(box=None, show_header=True, header_style="dim", padding=(0, 3), show_edge=False)
    t.add_column("#",        justify="right", style="dim",        width=2)
    t.add_column("attack",   style="white")
    t.add_column("combos",   justify="right", style="cyan",       width=6)
    t.add_column("prompts",  justify="right", style="cyan",       width=7)
    t.add_column("requests", justify="right", style="bold white", width=8)

    total_requests = 0
    for i, r in enumerate(recipes, 1):
        prompts_n = min(len(r.prompts), limit) if limit else len(r.prompts)
        combos_n  = len(r.parameter_combos)
        reqs      = prompts_n * combos_n
        total_requests += reqs
        t.add_row(str(i), r.attack_name, str(combos_n), str(prompts_n), str(reqs))

    console.print(t)
    console.print()

    # Run config panel
    first = recipes[0]
    judge_str = (
        f"{first.judge_config.provider}[dim]/[/]{first.judge_config.model}"
        if first.judge_config else "[dim]disabled[/]"
    )

    info = Table(box=None, show_header=False, padding=(0, 2), show_edge=False)
    info.add_column(style="dim", width=11)
    info.add_column()
    info.add_row("provider",  f"[white]{first.model_config.provider}[/]")
    info.add_row("model",     f"[bold white]{first.model_config.model}[/]")
    info.add_row("judge",     judge_str)
    _w = first.model_config.concurrency
    info.add_row("workers",   f"[cyan]{'auto' if _w is None else _w}[/]")
    _keys_n = len(first.model_config.api_key_envs) if first.model_config.api_key_envs else 1
    if _keys_n > 1:
        info.add_row("keys",  f"[cyan]{_keys_n}[/]  [dim]round-robin[/]")
    info.add_row("attacks",   f"[white]{len(recipes)}[/]")
    if models:
        info.add_row("models",  f"[bold magenta]{len(models)}[/]  [dim]{', '.join(models)}[/]")
        info.add_row("total req", f"[bold yellow]{total_requests * len(models)}[/]  [dim]({total_requests} × {len(models)} models)[/]")
    else:
        info.add_row("total req", f"[bold yellow]{total_requests}[/]")
    if limit:
        info.add_row("limit",   f"[dim]{limit} prompts / attack[/]")
    if combinations is not None:
        info.add_row("combos",  f"[dim]{combinations} combinations / attack[/]")

    console.print(Panel(info, title="[bold cyan] run config [/]", border_style="cyan", expand=False))
    console.print()


# ---------------------------------------------------------------------------
# Single attack output
# ---------------------------------------------------------------------------

def build_single_output(
    recipe: RecipeSpec,
    results: list[RunResult],
    metrics: Metrics,
) -> dict[str, Any]:
    return {
        "metadata": {
            "attack_name": recipe.attack_name,
            "model": recipe.model_config.model,
            "provider": recipe.model_config.provider,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "recipe_path": recipe.recipe_path,
            "total_combinations": len(recipe.parameter_combos),
            "total_prompts": len(recipe.prompts),
            "total_runs": len(results),
        },
        "metrics": _metrics_dict(metrics),
        "results": [_result_dict(r) for r in results],
    }


def print_summary(recipe: RecipeSpec, metrics: Metrics) -> None:
    judge_enabled = recipe.judge_config is not None

    table = Table(box=box.SIMPLE_HEAVY, show_header=False, padding=(0, 1))
    table.add_column(style="bold dim", width=12)
    table.add_column()

    table.add_row("Attack", f"[bold white]{recipe.attack_name}[/]")
    table.add_row("Model", f"{recipe.model_config.provider}[dim]/[/]{recipe.model_config.model}")
    table.add_row("Runs", f"{metrics.total}  [dim]({len(recipe.parameter_combos)} combos × {len(recipe.prompts)} prompts)[/]")

    if judge_enabled:
        asr_color = "red" if metrics.asr > 0.3 else "yellow" if metrics.asr > 0.1 else "green"
        table.add_row("ASR", f"[bold {asr_color}]{metrics.asr:.1%}[/]  [dim]{metrics.unsafe_count} unsafe / {metrics.total} total[/]")
        table.add_row("BR", f"{metrics.blacklist_rate:.1%}  [dim]{metrics.blacklisted_count} blacklisted[/]")
    else:
        table.add_row("Judge", "[dim]disabled[/]")

    if metrics.failed_count:
        table.add_row("Failed", f"[red]{metrics.failed_count}[/]")

    console.print(Panel(table, title=f"[bold cyan] Results [/]", border_style="cyan", expand=False))

    if judge_enabled and metrics.per_parameter_asr:
        _print_per_param_table(metrics)


def _print_per_param_table(metrics: Metrics) -> None:
    top = sorted(metrics.per_parameter_asr.items(), key=lambda x: -x[1])[:10]

    t = Table(box=box.SIMPLE_HEAD, show_header=True, header_style="bold cyan", padding=(0, 2), expand=False)
    t.add_column("Parameters", style="white")
    t.add_column("ASR", justify="right", width=8)
    t.add_column("Bar", justify="left", width=22, no_wrap=True)

    for key, asr in top:
        color = "red" if asr > 0.3 else "yellow" if asr > 0.1 else "green"
        filled = round(asr * 20)
        bar = f"[{color}]{'█' * filled}[/][dim]{'░' * (20 - filled)}[/]"
        t.add_row(key, f"[bold {color}]{asr:.1%}[/]", bar)

    console.print(Panel(t, title="[bold cyan] Per-parameter ASR [/]", border_style="cyan", expand=False))


# ---------------------------------------------------------------------------
# Multi-attack output
# ---------------------------------------------------------------------------

def build_multi_output(
    all_results: dict[str, tuple[RecipeSpec, list[RunResult], Metrics]],
) -> dict[str, Any]:
    attacks_out = {}
    for name, (recipe, results, metrics) in all_results.items():
        attacks_out[name] = {
            "metrics": _metrics_dict(metrics),
            "results": [_result_dict(r) for r in results],
        }

    asrs = {name: m.asr for name, (_, _, m) in all_results.items()}
    best_name = max(asrs, key=lambda k: asrs[k]) if asrs else None

    return {
        "metadata": {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "attacks_run": list(all_results.keys()),
            "total_attacks": len(all_results),
        },
        "summary": {
            "best_attack": best_name,
            "best_asr": round(asrs[best_name], 4) if best_name else 0.0,
            "per_attack_asr": {k: round(v, 4) for k, v in asrs.items()},
        },
        "attacks": attacks_out,
    }


def print_multi_summary(
    all_results: dict[str, tuple[RecipeSpec, list[RunResult], Metrics]],
) -> None:
    t = Table(
        title="[bold] Final Summary [/]",
        box=box.ROUNDED,
        show_header=True,
        header_style="bold cyan",
    )
    t.add_column("#", justify="right", style="dim", width=3)
    t.add_column("Attack", style="white")
    t.add_column("ASR", justify="right", width=8)
    t.add_column("Unsafe", justify="right", width=8)
    t.add_column("Total", justify="right", width=8)
    t.add_column("BR", justify="right", width=7)

    rows = sorted(
        [(name, recipe, m) for name, (recipe, _, m) in all_results.items()],
        key=lambda x: -x[2].asr,
    )

    for i, (name, recipe, m) in enumerate(rows, 1):
        color = "red" if m.asr > 0.3 else "yellow" if m.asr > 0.1 else "green"
        t.add_row(
            str(i),
            name,
            f"[bold {color}]{m.asr:.1%}[/]",
            str(m.unsafe_count),
            str(m.total),
            f"{m.blacklist_rate:.1%}",
        )

    console.print()
    console.print(t)


# ---------------------------------------------------------------------------
# Multi-model output
# ---------------------------------------------------------------------------

def build_multi_model_output(
    attack_name: str,
    model_results: dict[str, tuple[RecipeSpec, list[RunResult], "Metrics"]],
) -> dict[str, Any]:
    models_out = {}
    for model_name, (recipe, results, metrics) in model_results.items():
        models_out[model_name] = {
            "metrics": _metrics_dict(metrics),
            "results": [_result_dict(r) for r in results],
        }

    asrs = {m: metrics.asr for m, (_, _, metrics) in model_results.items()}
    best_model = max(asrs, key=lambda k: asrs[k]) if asrs else None
    first_recipe = next(iter(model_results.values()))[0]

    return {
        "metadata": {
            "attack_name": attack_name,
            "models": list(model_results.keys()),
            "provider": first_recipe.model_config.provider,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "total_runs_per_model": sum(
                len(results) for _, results, _ in model_results.values()
            ) // max(len(model_results), 1),
        },
        "summary": {
            "best_model": best_model,
            "best_asr": round(asrs[best_model], 4) if best_model else 0.0,
            "per_model_asr": {m: round(v, 4) for m, v in asrs.items()},
        },
        "models": models_out,
    }


def build_multi_attack_multi_model_output(
    attack_model_results: dict[str, dict[str, tuple[RecipeSpec, list[RunResult], "Metrics"]]],
) -> dict[str, Any]:
    """Output for multiple attacks × multiple models."""
    attacks_out = {}
    for attack_name, model_results in attack_model_results.items():
        models_out = {}
        for model_name, (recipe, results, metrics) in model_results.items():
            models_out[model_name] = {
                "metrics": _metrics_dict(metrics),
                "results": [_result_dict(r) for r in results],
            }
        attacks_out[attack_name] = {"models": models_out}

    return {
        "metadata": {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "attacks_run": list(attack_model_results.keys()),
            "models": list(next(iter(attack_model_results.values())).keys()) if attack_model_results else [],
        },
        "attacks": attacks_out,
    }


def print_model_comparison_summary(
    model_results: dict[str, tuple[RecipeSpec, list[RunResult], "Metrics"]],
    attack_name: str = "",
) -> None:
    title = f"[bold] Model Comparison{': ' + attack_name if attack_name else ''} [/]"
    t = Table(title=title, box=box.ROUNDED, show_header=True, header_style="bold cyan")
    t.add_column("#",      justify="right", style="dim", width=3)
    t.add_column("Model",  style="white")
    t.add_column("ASR",    justify="right", width=8)
    t.add_column("Unsafe", justify="right", width=8)
    t.add_column("Total",  justify="right", width=8)
    t.add_column("BR",     justify="right", width=7)

    rows = sorted(
        [(model_name, recipe, m) for model_name, (recipe, _, m) in model_results.items()],
        key=lambda x: -x[2].asr,
    )
    for i, (model_name, _recipe, m) in enumerate(rows, 1):
        color = "red" if m.asr > 0.3 else "yellow" if m.asr > 0.1 else "green"
        t.add_row(
            str(i),
            model_name,
            f"[bold {color}]{m.asr:.1%}[/]",
            str(m.unsafe_count),
            str(m.total),
            f"{m.blacklist_rate:.1%}",
        )

    console.print()
    console.print(t)


# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------

def _metrics_dict(metrics: Metrics) -> dict:
    return {
        "total": metrics.total,
        "evaluated": metrics.evaluated,
        "unsafe_count": metrics.unsafe_count,
        "safe_count": metrics.safe_count,
        "failed_count": metrics.failed_count,
        "blacklisted_count": metrics.blacklisted_count,
        "not_evaluated_count": metrics.not_evaluated_count,
        "asr": round(metrics.asr, 4),
        "blacklist_rate": round(metrics.blacklist_rate, 4),
        "per_parameter_asr": {k: round(v, 4) for k, v in metrics.per_parameter_asr.items()},
    }


def _result_dict(r: RunResult) -> dict:
    return {
        "prompt": r.prompt,
        "parameters": r.parameters,
        "modified_prompt": r.modified_prompt,
        "messages_sent": r.messages_sent,
        "response": r.response,
        "finish_reason": r.finish_reason,
        "judge_result": r.judge_result,
        "judge_reasoning": r.judge_reasoning,
        "error": r.error,
    }


def write_output(data: dict, path: str | Path) -> None:
    Path(path).write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")