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

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

import asyncio
import os
from pathlib import Path
from typing import Optional

from .metrics import calculate_metrics
from .output import (
    build_multi_attack_multi_model_output,
    build_multi_model_output,
    build_multi_output,
    build_single_output,
    print_model_comparison_summary,
    print_multi_summary,
    print_startup_banner,
    print_summary,
    write_output,
)
from .recipe import RecipeLoader, RecipeSpec
from .runner import AttackRunner


def _find_recipes_dir() -> Path:
    """Locate the recipes/ directory.

    Search order:
    1. FORGE_LITE_RECIPES_DIR env var (explicit override)
    2. recipes/ relative to CWD  (running from repo root)
    3. recipes/ relative to this file (editable install from repo root)
    """
    if env := os.environ.get("FORGE_LITE_RECIPES_DIR"):
        return Path(env)
    cwd_candidate = Path.cwd() / "recipes"
    if cwd_candidate.exists():
        return cwd_candidate
    return Path(__file__).parent.parent.parent / "recipes"


_DEFAULT_ATTACKS_DIR = _find_recipes_dir()


def _resolve_bundles(
    input_path: str | Path,
    attack_filter: Optional[list[str]],
) -> list[Path]:
    """
    Resolve input_path to a list of .attack / .py files.
    - "all"      → all .attack files from the default recipes/ directory
    - File       → [file]
    - Directory  → all .attack files inside, optionally filtered by name
    """
    if str(input_path).lower() == "all":
        p = _DEFAULT_ATTACKS_DIR
        if not p.exists():
            raise FileNotFoundError(
                f"Default attacks directory not found: {p}\n"
                f"Put your .attack bundles there or specify a path explicitly."
            )
    else:
        p = Path(input_path)
        if not p.exists():
            raise FileNotFoundError(f"Not found: {p}")

    if p.is_file():
        return [p]

    # Directory: collect all .attack files
    bundles = sorted(p.glob("*.attack"))
    if not bundles:
        raise ValueError(f"No .attack files found in directory: {p}")

    if attack_filter:
        filtered = [b for b in bundles if b.stem in attack_filter]
        missing = set(attack_filter) - {b.stem for b in filtered}
        if missing:
            raise ValueError(
                f"Attack(s) not found in {p}: {sorted(missing)}. "
                f"Available: {[b.stem for b in bundles]}"
            )
        bundles = filtered

    return bundles


async def run_async(
    input_path: str | Path,
    output: Optional[str | Path] = None,
    limit: Optional[int] = None,
    no_judge: bool = False,
    model: Optional[str] = None,
    models: Optional[list[str]] = None,
    provider: Optional[str] = None,
    attack_filter: Optional[list[str]] = None,
    workers: Optional[int] = None,
    combinations: Optional[int] = None,
    keys: Optional[list[str]] = None,
) -> dict:
    bundles = _resolve_bundles(input_path, attack_filter)

    # Load all recipes upfront so the banner can show full info
    recipes = []
    for b in bundles:
        r = RecipeLoader.load(b)
        _apply_overrides(r, no_judge, model, provider, workers, combinations, keys)
        recipes.append(r)

    print_startup_banner(recipes, limit=limit, workers=workers, combinations=combinations, models=models)

    if models:
        if len(recipes) == 1:
            return await _run_single_recipe_multi_models(recipes[0], models, output, limit)
        return await _run_multi_recipes_multi_models(recipes, models, output, limit)

    if len(recipes) == 1:
        return await _run_single_recipe(recipes[0], output, limit)

    return await _run_multi_recipes(recipes, output, limit)


async def _run_single_recipe(
    recipe: RecipeSpec,
    output: Optional[str | Path],
    limit: Optional[int],
) -> dict:
    results = await AttackRunner(recipe, limit=limit).run_async()
    metrics = calculate_metrics(results)
    data = build_single_output(recipe, results, metrics)

    print_summary(recipe, metrics)
    if output:
        write_output(data, output)
        from .runner import console
        console.print(f"[dim]Results saved → [/][bold]{output}[/]")

    return data


async def _run_multi_recipes(
    recipes: list[RecipeSpec],
    output: Optional[str | Path],
    limit: Optional[int],
) -> dict:
    from rich.console import Console
    _console = Console()

    all_results = {}
    for recipe in recipes:
        _console.print(f"\n[bold cyan]▶ {recipe.attack_name}[/]")
        results = await AttackRunner(recipe, limit=limit).run_async()
        metrics = calculate_metrics(results)
        print_summary(recipe, metrics)
        all_results[recipe.attack_name] = (recipe, results, metrics)

    data = build_multi_output(all_results)
    print_multi_summary(all_results)

    if output:
        write_output(data, output)
        _console.print(f"[dim]All results saved → [/][bold]{output}[/]")

    return data


async def _run_single_recipe_multi_models(
    recipe: RecipeSpec,
    models: list[str],
    output: Optional[str | Path],
    limit: Optional[int],
) -> dict:
    from copy import deepcopy
    from .runner import console, make_progress

    recipes_per_model = []
    for model_name in models:
        r = deepcopy(recipe)
        r.model_config.model = model_name
        if r.judge_config:
            r.judge_config.model = model_name
        recipes_per_model.append(r)

    shared_progress = make_progress()

    async def _run_one(r: RecipeSpec) -> list:
        return await AttackRunner(r, limit=limit).run_async(progress=shared_progress)

    with shared_progress:
        all_results_raw = await asyncio.gather(*[_run_one(r) for r in recipes_per_model])

    model_results = {}
    for r, results in zip(recipes_per_model, all_results_raw):
        metrics = calculate_metrics(results)
        model_results[r.model_config.model] = (r, results, metrics)

    data = build_multi_model_output(recipe.attack_name, model_results)
    print_model_comparison_summary(model_results, attack_name=recipe.attack_name)

    if output:
        write_output(data, output)
        console.print(f"[dim]Results saved → [/][bold]{output}[/]")

    return data


async def _run_multi_recipes_multi_models(
    recipes: list[RecipeSpec],
    models: list[str],
    output: Optional[str | Path],
    limit: Optional[int],
) -> dict:
    from copy import deepcopy
    from rich.console import Console
    from .runner import make_progress

    _console = Console()
    attack_model_results: dict = {}

    for recipe in recipes:
        _console.print(f"\n[bold cyan]▶ {recipe.attack_name}[/]")

        recipes_per_model = []
        for model_name in models:
            r = deepcopy(recipe)
            r.model_config.model = model_name
            if r.judge_config:
                r.judge_config.model = model_name
            recipes_per_model.append(r)

        shared_progress = make_progress()

        async def _run_one(r: RecipeSpec) -> list:
            return await AttackRunner(r, limit=limit).run_async(progress=shared_progress)

        with shared_progress:
            all_results_raw = await asyncio.gather(*[_run_one(r) for r in recipes_per_model])

        model_results = {}
        for r, results in zip(recipes_per_model, all_results_raw):
            metrics = calculate_metrics(results)
            model_results[r.model_config.model] = (r, results, metrics)

        print_model_comparison_summary(model_results, attack_name=recipe.attack_name)
        attack_model_results[recipe.attack_name] = model_results

    data = build_multi_attack_multi_model_output(attack_model_results)

    if output:
        write_output(data, output)
        _console.print(f"[dim]All results saved → [/][bold]{output}[/]")

    return data


_IFT_DEFAULTS = {
    "login_env":  "IFT_LOGIN",
    "api_key_env": "IFT_PASSWORD",
    "auth_url":   "https://gigachat.sberdevices.ru/v1/token",
    "base_url":   "https://gigachat.ift.sberdevices.ru/v1",
}

_GIGACHAT_DEFAULTS = {
    "login_env":  "",
    "api_key_env": "GIGACHAT_API_KEY",
    "auth_url":   None,
    "base_url":   None,
}


def _apply_provider_defaults(cfg, provider: str) -> None:
    """Apply connection defaults when provider is overridden from CLI."""
    defaults = _IFT_DEFAULTS if provider == "gigachat_ift" else _GIGACHAT_DEFAULTS
    cfg.provider  = provider
    cfg.login_env  = defaults["login_env"]
    cfg.api_key_env = defaults["api_key_env"]
    if not cfg.auth_url:
        cfg.auth_url = defaults["auth_url"]
    if not cfg.base_url:
        cfg.base_url = defaults["base_url"]


def _apply_overrides(
    recipe: RecipeSpec,
    no_judge: bool,
    model: Optional[str],
    provider: Optional[str],
    workers: Optional[int] = None,
    combinations: Optional[int] = None,
    keys: Optional[list[str]] = None,
) -> None:
    if no_judge:
        recipe.judge_config = None
    if model:
        recipe.model_config.model = model
        if recipe.judge_config:
            recipe.judge_config.model = model
    if provider:
        _apply_provider_defaults(recipe.model_config, provider)
        if recipe.judge_config:
            _apply_provider_defaults(recipe.judge_config, provider)
    if workers is not None:
        recipe.model_config.concurrency = workers
    if combinations is not None:
        recipe.parameter_combos = recipe.parameter_combos[:combinations]
    if keys:
        recipe.model_config.api_key_envs = keys


def run(
    input_path: str | Path,
    output: Optional[str | Path] = None,
    limit: Optional[int] = None,
    no_judge: bool = False,
    model: Optional[str] = None,
    models: Optional[list[str]] = None,
    provider: Optional[str] = None,
    attack_filter: Optional[list[str]] = None,
    workers: Optional[int] = None,
    combinations: Optional[int] = None,
    keys: Optional[list[str]] = None,
) -> dict:
    return asyncio.run(run_async(
        input_path, output=output, limit=limit,
        no_judge=no_judge, model=model, models=models, provider=provider,
        attack_filter=attack_filter, workers=workers, combinations=combinations,
        keys=keys,
    ))