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

src/forge_lite/runner.py
200 строк7 KB
wearetyomsmnv
add 10 attack, model combinations, upgrade of UI
14 апр 2026, 18:08
14 апр 2026, 18:08423255c
from __future__ import annotations

import asyncio
import logging
from collections import Counter
from contextlib import nullcontext
from dataclasses import dataclass
from typing import Any, Optional

from rich.console import Console
from rich.progress import (
    BarColumn,
    MofNCompleteColumn,
    Progress,
    SpinnerColumn,
    TextColumn,
    TimeElapsedColumn,
    TimeRemainingColumn,
)

from .clients import ClientFactory
from .clients.base import ChatMessage, RunConfig
from .judge import JudgeFactory
from .recipe import RecipeSpec

logger = logging.getLogger(__name__)
console = Console()


@dataclass
class RunResult:
    prompt: str
    parameters: dict[str, Any]
    modified_prompt: str
    messages_sent: list[dict]
    response: str
    finish_reason: str
    judge_result: str = "not_evaluated"
    judge_reasoning: Optional[str] = None
    error: Optional[str] = None


def _normalize_message(msg: str | list[dict]) -> tuple[list[ChatMessage], str]:
    if isinstance(msg, str):
        return [ChatMessage(role="user", content=msg)], msg
    elif isinstance(msg, list):
        messages = [
            ChatMessage(
                role=m["role"],
                content=m["content"],
                attachments=m.get("attachments", []),
            )
            for m in msg
        ]
        user_parts = [m["content"] for m in msg if m.get("role") == "user"]
        return messages, " | ".join(user_parts)
    else:
        raise ValueError(f"generate_message must return str or list[dict], got {type(msg)}")


def make_progress(console: Console = console) -> Progress:
    return Progress(
        SpinnerColumn(),
        TextColumn("[bold cyan]{task.description:<42}[/]"),
        BarColumn(bar_width=30),
        MofNCompleteColumn(),
        TextColumn("[green]{task.percentage:>5.1f}%[/]"),
        TimeElapsedColumn(),
        TimeRemainingColumn(),
        console=console,
        transient=False,
    )


def _error_key(msg: str) -> str:
    """Extract the short, human-readable part of an exception message."""
    line = msg.split("\n")[0].strip()
    # Trim giant aiohttp / requests chains: keep up to first parenthesis
    for sep in (" (", ": Max retries", ": HTTPSConnectionPool"):
        if sep in line:
            line = line[: line.index(sep)]
    return line[:90] if line else msg[:90]


def _print_error_summary(errors: list[str], phase: str) -> None:
    if not errors:
        return
    counts = Counter(_error_key(e) for e in errors)
    total = len(errors)
    top = counts.most_common(2)
    parts = "  |  ".join(f"{msg} [dim](×{n})[/]" for msg, n in top)
    console.print(f"[yellow]  ⚠  {total} {phase} error{'s' if total != 1 else ''}:[/]  {parts}")


class AttackRunner:
    def __init__(self, recipe: RecipeSpec, limit: Optional[int] = None) -> None:
        self._recipe = recipe
        self._prompts = recipe.prompts[:limit] if limit is not None else recipe.prompts

    async def run_async(self, progress: Optional[Progress] = None) -> list[RunResult]:
        recipe = self._recipe
        label = recipe.model_config.model

        # Phase 1: Generate all messages (sync, fast)
        tasks_input: list[tuple[str, dict, list[ChatMessage], str]] = []
        for combo in recipe.parameter_combos:
            for prompt in self._prompts:
                raw = recipe.generate_fn(prompt=prompt, **combo)
                messages, modified = _normalize_message(raw)
                tasks_input.append((prompt, combo, messages, modified))

        total = len(tasks_input)
        client = ClientFactory.get(recipe.model_config)
        run_cfg = RunConfig(
            model=recipe.model_config.model,
            temperature=recipe.model_config.temperature,
            max_tokens=recipe.model_config.max_tokens,
            concurrency=recipe.model_config.concurrency,
        )

        # Phase 2: Execute attacks
        async def _call(prompt, combo, messages, modified) -> RunResult:
            try:
                resp = await client.complete(messages, run_cfg)
                return RunResult(
                    prompt=prompt,
                    parameters=combo,
                    modified_prompt=modified,
                    messages_sent=[{"role": m.role, "content": m.content} for m in messages],
                    response=resp.content,
                    finish_reason=resp.finish_reason,
                )
            except Exception as e:
                return RunResult(
                    prompt=prompt,
                    parameters=combo,
                    modified_prompt=modified,
                    messages_sent=[{"role": m.role, "content": m.content} for m in messages],
                    response="",
                    finish_reason="failed",
                    error=str(e),
                )

        own_progress = progress is None
        atk_progress = make_progress() if own_progress else progress
        atk_ctx = atk_progress if own_progress else nullcontext()

        with atk_ctx:
            task_id = atk_progress.add_task(f"Attacking  {label}", total=total)

            async def _tracked(args):
                result = await _call(*args)
                atk_progress.advance(task_id)
                return result

            results: list[RunResult] = list(
                await asyncio.gather(*[_tracked(t) for t in tasks_input])
            )

        if hasattr(client, "close"):
            await client.close()

        _print_error_summary([r.error for r in results if r.error], "request")

        # Phase 3: Judge
        if recipe.judge_config is not None:
            judge = JudgeFactory.get(recipe.judge_config)
            successful = [r for r in results if r.error is None]

            if successful:
                jdg_progress = make_progress() if own_progress else progress
                jdg_ctx = jdg_progress if own_progress else nullcontext()

                with jdg_ctx:
                    task_id = jdg_progress.add_task(f"Judging    {label}", total=len(successful))

                    async def _judge_one(r: RunResult) -> tuple[str, Optional[str]]:
                        verdict = await judge._evaluate_one(r.prompt, r.response)
                        jdg_progress.advance(task_id)
                        return verdict

                    verdicts = list(
                        await asyncio.gather(*[_judge_one(r) for r in successful])
                    )

                for r, (verdict, reasoning) in zip(successful, verdicts):
                    r.judge_result = verdict
                    r.judge_reasoning = reasoning

            judge_errors = sum(1 for r in results if r.judge_result == "error")
            if judge_errors:
                console.print(
                    f"[yellow]  ⚠  {judge_errors} judge error{'s' if judge_errors != 1 else ''} "
                    f"(network / auth)[/]"
                )

            if hasattr(judge._client, "close"):
                await judge._client.close()

        return results