main.py

import argparse
import asyncio
import json
import logging
import os
import re
from datetime import datetime
from pathlib import Path

import litellm
from dotenv import load_dotenv

from benchmark.v1 import BenchmarkV1
from benchmark.v2 import BenchmarkV2, BenchmarkV2Reevaluator

# Настройка окружения
litellm.ssl_verify = False
load_dotenv()

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

# Отключаем verbose логи от litellm
for logger_name in ['litellm', 'LiteLLM']:
    litellm_logger = logging.getLogger(logger_name)
    litellm_logger.setLevel(logging.WARNING)
    litellm_logger.propagate = False

def get_dataset_log_name(dataset: str) -> str:
    dataset_path = Path(dataset)
    if dataset_path.suffix.lower() == ".json":
        return dataset_path.stem
    return dataset


def find_existing_runs(timestamp: str, logs_dir: Path, dataset: str | None = None) -> tuple:
    """Находит существующие прогоны с указанным timestamp"""
    dataset_log_name = get_dataset_log_name(dataset) if dataset is not None else None
    pattern = f"benchmark_{timestamp}_run_*_{dataset_log_name}.json" if dataset_log_name else f"benchmark_{timestamp}_run_*.json"
    
    existing_files = list(logs_dir.glob(pattern))
    
    if not existing_files:
        if dataset_log_name:
            raise ValueError(f"Не найдено логов с timestamp {timestamp} и dataset {dataset_log_name}")
        raise ValueError(f"Не найдено логов с timestamp {timestamp}")

    if dataset_log_name is None:
        detected_datasets = set()
        for file in existing_files:
            match = re.match(rf"benchmark_{re.escape(timestamp)}_run_\d+_(.+)\.json$", file.name)
            if match:
                detected_datasets.add(match.group(1))

        if not detected_datasets:
            raise ValueError(f"Не удалось определить датасет для timestamp {timestamp}")

        if len(detected_datasets) > 1:
            datasets_str = ", ".join(sorted(detected_datasets))
            raise ValueError(
                f"Для timestamp {timestamp} найдено несколько датасетов: {datasets_str}. Укажите dataset явно."
            )

        dataset_log_name = detected_datasets.pop()
    
    # Извлекаем номера run
    run_numbers = []
    for file in existing_files:
        match = re.search(r'_run_(\d+)_', file.name)
        if match:
            run_numbers.append(int(match.group(1)))
    
    if not run_numbers:
        raise ValueError(f"Не удалось извлечь номера run из файлов")
    
    max_run = max(run_numbers)
    
    # Читаем config и results из первого прогона
    first_run_file = logs_dir / f"benchmark_{timestamp}_run_1_{dataset_log_name}.json"
    with open(first_run_file, 'r', encoding='utf-8') as f:
        data = json.load(f)
        config = data.get("config", {})
        results = data.get("results", [])
    
    logger.info(f"Найдено {len(run_numbers)} существующих прогонов (максимальный: run_{max_run})")
    
    return max_run, config, results


def parse_reevaluation_log_files(files_arg: str, logs_dir: Path) -> list[Path]:
    """Разбирает список файлов для переоценки и ищет короткие имена в logs_v2."""
    raw_paths = [part.strip() for part in files_arg.split(",") if part.strip()]
    if not raw_paths:
        raise ValueError("Для переоценки нужно указать хотя бы один JSON-файл")

    resolved_files = []
    for raw_path in raw_paths:
        requested_path = Path(raw_path)
        candidates = [requested_path]
        if not requested_path.is_absolute():
            candidates.append(logs_dir / requested_path)

        for candidate in candidates:
            if candidate.exists() and candidate.is_file():
                resolved_files.append(candidate)
                break
        else:
            raise FileNotFoundError(f"Файл для переоценки не найден: {raw_path}")

    return resolved_files


def find_max_run_number_for_timestamp(timestamp: str, logs_dir: Path) -> int:
    """Находит максимальный номер run для продолжения переоценки без привязки к датасету."""
    existing_files = list(logs_dir.glob(f"benchmark_{timestamp}_run_*.json"))
    if not existing_files:
        raise ValueError(f"Не найдено логов с timestamp {timestamp}")

    run_numbers = []
    for file in existing_files:
        match = re.match(rf"benchmark_{re.escape(timestamp)}_run_(\d+)_.+\.json$", file.name)
        if match:
            run_numbers.append(int(match.group(1)))

    if not run_numbers:
        raise ValueError(f"Не удалось извлечь номера run для timestamp {timestamp}")

    return max(run_numbers)


def load_benchmark_log_file(log_file: Path) -> dict:
    """Загружает JSON-лог бенчмарка и проверяет минимальную структуру."""
    with open(log_file, 'r', encoding='utf-8') as f:
        data = json.load(f)

    if not isinstance(data, dict):
        raise ValueError(f"Файл {log_file.as_posix()} должен содержать JSON-объект")
    results = data.get("results")
    if not isinstance(results, list):
        raise ValueError(f"Файл {log_file.as_posix()} не содержит список results")
    if any(not isinstance(result, dict) for result in results):
        raise ValueError(f"Файл {log_file.as_posix()} содержит некорректные элементы results")

    return data


def get_reevaluation_dataset_log_name(log_data: dict, log_file: Path) -> str:
    """Определяет имя датасета для имени нового файла переоценки."""
    config = log_data.get("config", {})
    dataset = config.get("dataset")
    if dataset:
        return get_dataset_log_name(dataset)

    match = re.match(r"benchmark_.+_run_\d+_(.+)\.json$", log_file.name)
    if match:
        return match.group(1)
    return log_file.stem


def build_reevaluation_config(
    source_config: dict,
    source_log_file: Path,
    dataset_log_name: str,
    reevaluator: BenchmarkV2Reevaluator,
    base_timestamp: str,
    run_number: int,
    total_runs: int,
    verbose_name: str | None = None,
) -> dict:
    """Формирует config нового лога переоценки, сохраняя метаданные исходного прогона."""
    config = dict(source_config)
    config.update({
        "dataset": source_config.get("dataset", dataset_log_name),
        "judge_model": reevaluator.judge_model_name,
        "judge_max_workers": reevaluator.judge_max_workers,
        "judge_transport": reevaluator.judge_transport_type,
        "timestamp": base_timestamp,
        "run_number": run_number,
        "total_runs": total_runs,
        "no_regenerate": True,
        "reevaluation": True,
        "reevaluation_source": {
            "file": source_log_file.as_posix(),
            "timestamp": source_config.get("timestamp"),
            "run_number": source_config.get("run_number"),
            "total_runs": source_config.get("total_runs"),
            "judge_model": source_config.get("judge_model"),
        },
    })

    if reevaluator.judge_extra_body:
        config["judge_extra_body"] = reevaluator.judge_extra_body
    else:
        config.pop("judge_extra_body", None)

    if verbose_name:
        config["verbose_name"] = verbose_name

    return config


async def run_v2_reevaluation(
    log_files: list[Path],
    continue_timestamp: str | None,
    num_runs: int,
    judge_model_name: str | None,
    verbose_name: str | None,
    debug_logs: int,
) -> list[Path]:
    """Повторно оценивает готовые v2-логи и сохраняет новые benchmark-файлы."""
    if num_runs < 1:
        raise ValueError("--num-runs должен быть не меньше 1")

    logs_dir = Path("logs_v2")
    logs_dir.mkdir(exist_ok=True)
    source_logs = [(log_file, load_benchmark_log_file(log_file)) for log_file in log_files]

    if continue_timestamp:
        base_timestamp = continue_timestamp
        start_run_number = find_max_run_number_for_timestamp(base_timestamp, logs_dir)
    else:
        base_timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        start_run_number = 0

    total_new_runs = len(source_logs) * num_runs
    total_runs = start_run_number + total_new_runs
    current_run_number = start_run_number
    output_files = []

    reevaluator = BenchmarkV2Reevaluator(
        judge_model_name=judge_model_name,
        debug_logs=debug_logs,
    )

    logger.info("="*70)
    if continue_timestamp:
        logger.info(f"ПРОДОЛЖЕНИЕ ПЕРЕОЦЕНКИ V2 (добавление {total_new_runs} прогонов)")
    else:
        logger.info(f"ПЕРЕОЦЕНКА V2 (новых прогонов: {total_new_runs})")
    logger.info(f"Timestamp: {base_timestamp}")
    logger.info(f"Файлов-источников: {len(source_logs)}, повторов на файл: {num_runs}")
    logger.info("="*70)

    for _ in range(num_runs):
        for source_log_file, source_log_data in source_logs:
            current_run_number += 1
            source_config = source_log_data.get("config", {})
            source_results = source_log_data["results"]
            dataset_log_name = get_reevaluation_dataset_log_name(source_log_data, source_log_file)
            log_filename = logs_dir / f"benchmark_{base_timestamp}_run_{current_run_number}_{dataset_log_name}.json"

            if log_filename.exists():
                raise FileExistsError(f"Файл {log_filename.as_posix()} уже существует, переоценка не будет его перезаписывать")

            logger.info("")
            logger.info(f"{'='*70}")
            logger.info(f"ПЕРЕОЦЕНКА RUN {current_run_number}/{total_runs}: {source_log_file.as_posix()}")
            logger.info(f"{'='*70}")

            results, summary = await reevaluator.reevaluate_results(
                source_results,
                run_number=current_run_number,
            )
            log_data = {
                "config": build_reevaluation_config(
                    source_config,
                    source_log_file,
                    dataset_log_name,
                    reevaluator,
                    base_timestamp,
                    current_run_number,
                    total_runs,
                    verbose_name,
                ),
                "results": results,
                "summary": summary,
            }

            with open(log_filename, 'w', encoding='utf-8') as f:
                json.dump(log_data, f, ensure_ascii=False, indent=2)

            output_files.append(log_filename)
            logger.info(f"Переоценка run {current_run_number}: Обработано {summary['successful_dialogs']}/{summary['total_dialogs']} диалогов")
            logger.info(f"Переоценка run {current_run_number}: Результаты сохранены в {log_filename.as_posix()}")

    return output_files

def main():
    parser = argparse.ArgumentParser(description="RuQualBench - бенчмарк качества русского языка")
    subparsers = parser.add_subparsers(dest="command", help="Версия бенчмарка")

    # Парсер для v1 (по умолчанию)
    parser_v1 = subparsers.add_parser("v1", help="Запуск версии v1")
    
    # Парсер для v2
    parser_v2 = subparsers.add_parser("v2", help="Запуск версии v2")

    # Общие аргументы
    for p in [parser_v1, parser_v2]:
        p.add_argument(
            "--model",
            type=str,
            help="Имя тестируемой модели (для v2 обязательно вместе с --test-provider)"
        )
        p.add_argument(
            "--judge-model",
            type=str,
            help="Переопределить модель-оценщик из .env"
        )
        p.add_argument(
            "--extra-body",
            type=str,
            help='JSON объект для extra_body параметра тестируемой модели (например: \'{"temperature": 0.7}\')'
        )
        p.add_argument(
            "-n", "--num-runs",
            type=int,
            default=1,
            help="Количество прогонов бенчмарка для вычисления средних значений и погрешности (по умолчанию: 1)"
        )
        p.add_argument(
            "-v", "--verbose-name",
            type=str,
            help="Красивое имя модели для отображения в лидерборде (опционально)"
        )

    # Аргументы только для v1
    parser_v1.add_argument(
        "--dataset",
        type=str,
        default="lite",
        choices=["debug", "lite", "base", "large"],
        help="Выбор датасета (по умолчанию: lite)"
    )
    parser_v1.add_argument(
        "--continue",
        type=str,
        dest="continue_timestamp",
        help="Продолжить существующую серию прогонов (указать timestamp, например: 2025-10-17_15-17-05)"
    )
    parser_v1.add_argument(
        "--no-regenerate",
        action="store_true",
        help="Генерировать ответы от модели только один раз, оценивать судьей N раз (работает с -n)"
    )

    # Аргументы только для v2
    parser_v2.add_argument(
        "--dataset-path",
        type=str,
        default="v2_lite.json",
        help="Путь к JSON файлу датасета (по умолчанию: v2_lite.json)"
    )
    parser_v2.add_argument(
        "--continue",
        type=str,
        dest="continue_timestamp",
        help="Продолжить существующую серию прогонов (указать timestamp, например: 2025-10-17_15-17-05)"
    )
    parser_v2.add_argument(
        "--no-regenerate",
        action="store_true",
        help="Генерировать ответы от модели только один раз, оценивать судьей N раз (работает с -n)"
    )
    parser_v2.add_argument(
        "--debug-logs",
        type=int,
        default=0,
        choices=[0, 1, 2],
        help="Уровень отладки логов: 0 - стандартный, 1 - добавить сырой ответ судьи, 2 - добавить сырой ответ и промпт"
    )
    parser_v2.add_argument(
        "--test-provider",
        type=str,
        help="Провайдер тестируемой модели из providers.toml/providers.default.toml"
    )
    parser_v2.add_argument(
        "--providers-file",
        type=Path,
        help="Явный путь к TOML-файлу провайдеров вместо автоматического выбора"
    )
    parser_v2.add_argument(
        "reevaluate_files",
        nargs="?",
        help="Один JSON-лог v2 или список JSON-логов через запятую для повторной оценки без генерации"
    )

    # Если аргументы не переданы или первый аргумент не v1/v2, считаем что это v1
    import sys
    if len(sys.argv) == 1 or (len(sys.argv) > 1 and sys.argv[1] not in ["v1", "v2"]):
        # Вставляем 'v1' первым аргументом
        sys.argv.insert(1, "v1")

    dataset_path_explicitly_set = "--dataset-path" in sys.argv

    args = parser.parse_args()
    
    extra_body = None
    if args.extra_body:
        extra_body = json.loads(args.extra_body)
    
    if args.command == "v2":
        if args.reevaluate_files:
            try:
                reevaluation_log_files = parse_reevaluation_log_files(
                    args.reevaluate_files,
                    Path("logs_v2"),
                )
                asyncio.run(run_v2_reevaluation(
                    log_files=reevaluation_log_files,
                    continue_timestamp=args.continue_timestamp,
                    num_runs=args.num_runs,
                    judge_model_name=args.judge_model,
                    verbose_name=args.verbose_name,
                    debug_logs=args.debug_logs,
                ))
            except (FileNotFoundError, FileExistsError, ValueError, json.JSONDecodeError) as e:
                parser.error(str(e))
        elif args.continue_timestamp:
            max_run, existing_config, existing_results = find_existing_runs(
                args.continue_timestamp,
                Path("logs_v2")
            )

            dataset_path = existing_config.get("dataset", args.dataset_path)
            if dataset_path_explicitly_set and args.dataset_path != dataset_path:
                logger.warning(
                    "ВНИМАНИЕ: --dataset-path=%s отличается от датасета продолжаемой серии (%s)",
                    args.dataset_path,
                    dataset_path,
                )
                logger.warning("Будет использован датасет из исходной серии")

            verbose_name = args.verbose_name or existing_config.get("verbose_name")
            model_name = args.model or existing_config.get("model")
            test_provider = args.test_provider or existing_config.get("test_provider")
            existing_no_regenerate = existing_config.get("no_regenerate", False)

            if not model_name:
                parser.error("Для v2 нужно указать --model")
            if not test_provider:
                parser.error("Для v2 нужно указать --test-provider (в старом логе провайдер не сохранен)")

            if args.no_regenerate and not existing_no_regenerate:
                logger.warning("ВНИМАНИЕ: --no-regenerate указан, но исходная серия запускалась без этого флага")
                logger.warning("Будет использован режим --no-regenerate с ответами из run_1 исходной серии")
                use_no_regenerate = True
            elif existing_no_regenerate:
                logger.info("Исходная серия использовала --no-regenerate, продолжаем в том же режиме")
                use_no_regenerate = True
            else:
                use_no_regenerate = False

            answer_results = None
            if use_no_regenerate and existing_results:
                answer_results = [
                    {
                        "dialog_id": r["dialog_id"],
                        "dialog": r["dialog"],
                        "answer": r["answer"],
                        "tokens": r["tokens"],
                        "error": r["error"]
                    }
                    for r in existing_results
                ]

            try:
                benchmark = BenchmarkV2(
                    dataset_path=dataset_path,
                    model_name=model_name,
                    judge_model_name=args.judge_model,
                    extra_body=extra_body,
                    verbose_name=verbose_name,
                    debug_logs=args.debug_logs,
                    test_provider=test_provider,
                    providers_file=args.providers_file,
                )
            except (FileNotFoundError, ValueError) as e:
                parser.error(str(e))

            previous_model = existing_config.get("model")
            previous_test_provider = existing_config.get("test_provider")
            previous_judge_model = existing_config.get("judge_model")
            if previous_model and benchmark.model_name != previous_model:
                logger.warning(
                    "Тестируемая модель отличается от исходной серии: было %s, сейчас %s",
                    previous_model,
                    benchmark.model_name,
                )
            if previous_test_provider and benchmark.test_provider != previous_test_provider:
                logger.warning(
                    "Провайдер тестируемой модели отличается от исходной серии: было %s, сейчас %s",
                    previous_test_provider,
                    benchmark.test_provider,
                )
            if previous_judge_model and benchmark.judge_model_name != previous_judge_model:
                logger.warning(
                    "Модель-оценщик отличается от исходной серии: было %s, сейчас %s",
                    previous_judge_model,
                    benchmark.judge_model_name,
                )

            asyncio.run(benchmark.run_multiple_benchmarks(
                num_runs=args.num_runs,
                continue_timestamp=args.continue_timestamp,
                start_run_number=max_run,
                no_regenerate=use_no_regenerate,
                existing_answer_results=answer_results
            ))
        else:
            if not args.model:
                parser.error("Для v2 нужно указать --model")
            if not args.test_provider:
                parser.error("Для v2 нужно указать --test-provider")

            try:
                benchmark = BenchmarkV2(
                    dataset_path=args.dataset_path,
                    model_name=args.model,
                    judge_model_name=args.judge_model,
                    extra_body=extra_body,
                    verbose_name=args.verbose_name,
                    debug_logs=args.debug_logs,
                    test_provider=args.test_provider,
                    providers_file=args.providers_file,
                )
            except (FileNotFoundError, ValueError) as e:
                parser.error(str(e))
            asyncio.run(benchmark.run_multiple_benchmarks(
                args.num_runs,
                no_regenerate=args.no_regenerate
            ))
    else:
        # Логика v1
        if args.continue_timestamp:
            max_run, existing_config, existing_results = find_existing_runs(
                args.continue_timestamp,
                Path("logs"),
                args.dataset,
            )
            
            # Берем параметры из существующих логов
            model = existing_config.get("model")
            judge_model = existing_config.get("judge_model")
            verbose_name = existing_config.get("verbose_name")
            existing_no_regenerate = existing_config.get("no_regenerate", False)
            
            logger.info(f"Продолжение серии с моделью: {model}")
            if verbose_name:
                logger.info(f"Verbose name: {verbose_name}")
            
            # Определяем режим no_regenerate
            if args.no_regenerate and not existing_no_regenerate:
                logger.warning("ВНИМАНИЕ: --no-regenerate указан, но исходная серия запускалась без этого флага")
                logger.warning("Будет использован режим --no-regenerate с ответами из run_1 исходной серии")
                use_no_regenerate = True
            elif existing_no_regenerate:
                logger.info("Исходная серия использовала --no-regenerate, продолжаем в том же режиме")
                use_no_regenerate = True
            else:
                use_no_regenerate = False
            
            # Извлекаем только ответы без оценок судьи из существующих результатов
            answer_results = None
            if use_no_regenerate and existing_results:
                answer_results = [
                    {
                        "dialog_id": r["dialog_id"],
                        "dialog": r["dialog"],
                        "answer": r["answer"],
                        "tokens": r["tokens"],
                        "error": r["error"]
                    }
                    for r in existing_results
                ]
            
            benchmark = BenchmarkV1(
                dataset_name=args.dataset,
                model_name=model,
                judge_model_name=judge_model,
                extra_body=extra_body,
                verbose_name=verbose_name
            )
            
            asyncio.run(benchmark.run_multiple_benchmarks(
                num_runs=args.num_runs,
                continue_timestamp=args.continue_timestamp,
                start_run_number=max_run,
                no_regenerate=use_no_regenerate,
                existing_answer_results=answer_results
            ))
        else:
            benchmark = BenchmarkV1(
                dataset_name=args.dataset,
                model_name=args.model,
                judge_model_name=args.judge_model,
                extra_body=extra_body,
                verbose_name=args.verbose_name
            )
            
            asyncio.run(benchmark.run_multiple_benchmarks(
                num_runs=args.num_runs,
                no_regenerate=args.no_regenerate
            ))

if __name__ == "__main__":
    main()