evalstone/bash/tests/test_summary_upsert.py

"""Upsert behaviour for the project-level CSV/Excel summary."""

from pathlib import Path
import sys

import numpy as np
import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from collect_results import (  # noqa: E402
    OUTPUT_COLUMNS,
    TOTAL_CATEGORY,
    collect_all,
    collect_benchmark,
    upsert_summary_rows,
)


def _row(benchmark: str, score: float, n: int = 1, category: str = '其他') -> dict:
    row = {col: np.nan for col in OUTPUT_COLUMNS}
    row.update({
        '分类': category,
        'Benchmark': benchmark,
        '得分': score,
        '实测时间(h)': 1.0,
        '总样本数': n,
        '累计总tokens': 10,
    })
    return row


def test_upsert_appends_new_benchmark_and_keeps_existing():
    existing = pd.DataFrame([_row('gpqa_diamond', 0.5, category='知识与语言理解')], columns=OUTPUT_COLUMNS)
    existing = pd.concat(
        [existing, pd.DataFrame([{col: np.nan for col in OUTPUT_COLUMNS} | {
            '分类': TOTAL_CATEGORY,
            'Benchmark': '',
            '得分': 0.5,
        }], columns=OUTPUT_COLUMNS)],
        ignore_index=True,
    )

    out = upsert_summary_rows(existing, [_row('terminal_bench_v2_1', 0.0, category='智能体与工具')])
    names = [n for n in out['Benchmark'].tolist() if str(n).strip()]
    assert names == ['gpqa_diamond', 'terminal_bench_v2_1']
    assert out.iloc[-1]['分类'] == TOTAL_CATEGORY
    assert out.iloc[-1]['总样本数'] == 2


def test_upsert_overwrites_matching_benchmark_in_place():
    existing = pd.DataFrame(
        [
            _row('gpqa_diamond', 0.1, n=1, category='知识与语言理解'),
            _row('hle', 0.2, n=2, category='知识与语言理解'),
        ],
        columns=OUTPUT_COLUMNS,
    )

    out = upsert_summary_rows(existing, [_row('gpqa_diamond', 0.9, n=100, category='知识与语言理解')])
    gpqa = out.loc[out['Benchmark'] == 'gpqa_diamond'].iloc[0]
    assert gpqa['得分'] == 0.9
    assert gpqa['总样本数'] == 100
    assert list(out['Benchmark'].tolist()[:-1]) == ['gpqa_diamond', 'hle']
    assert (out['Benchmark'] == 'gpqa_diamond').sum() == 1


def test_upsert_aliases_hle_low_to_hle():
    existing = pd.DataFrame([_row('hle_low', 0.3)], columns=OUTPUT_COLUMNS)
    out = upsert_summary_rows(existing, [_row('hle', 0.8)])
    assert list(out['Benchmark'].tolist()[:-1]) == ['hle']
    assert out.loc[out['Benchmark'] == 'hle'].iloc[0]['得分'] == 0.8


def test_collect_all_merges_into_existing_csv(tmp_path: Path):
    summary_dir = tmp_path / 'results'
    summary_dir.mkdir()
    prior = pd.DataFrame([_row('gpqa_diamond', 0.4, category='知识与语言理解')], columns=OUTPUT_COLUMNS)
    prior = upsert_summary_rows(None, [_row('gpqa_diamond', 0.4, category='知识与语言理解')])
    csv_path = summary_dir / 'mock-model.csv'
    prior.to_csv(csv_path, index=False, encoding='utf-8-sig')

    output_dir = tmp_path / 'output'
    reports = output_dir / 'terminal_bench_v2_1' / 'seed_42' / 'reports'
    reports.mkdir(parents=True)
    (reports / 'terminal_bench_v2_1.json').write_text(
        '{"score": 0.0, "num": 2, "metrics": [{"identity": {"name": "accuracy", "aggregation": "mean", "dimensions": {}}, "score": 0.0, "num": 2}]}',
        encoding='utf-8',
    )

    collect_all(
        output_dir,
        'mock-model',
        out_name='mock-model',
        include_benchmarks=['terminal_bench_v2_1'],
        excel_output_dir=summary_dir,
    )

    df = pd.read_csv(csv_path, encoding='utf-8-sig')
    names = [n for n in df['Benchmark'].fillna('').tolist() if str(n).strip()]
    assert names == ['gpqa_diamond', 'terminal_bench_v2_1']
    tb = df.loc[df['Benchmark'] == 'terminal_bench_v2_1'].iloc[0]
    assert tb['得分'] == 0.0
    assert df.iloc[-1]['分类'] == TOTAL_CATEGORY


def test_agent_bench_uses_report_per_request_perf(tmp_path: Path):
    import json

    output_dir = tmp_path / 'output'
    bench = 'terminal_bench_v2_1'
    reports = output_dir / bench / 'seed_42' / 'reports'
    preds = output_dir / bench / 'seed_42' / 'predictions'
    reports.mkdir(parents=True)
    preds.mkdir(parents=True)
    (reports / f'{bench}.json').write_text(
        json.dumps({
            'score': 0.0,
            'num': 2,
            'metrics': [{
                'identity': {'name': 'accuracy', 'aggregation': 'mean', 'dimensions': {}},
                'score': 0.0,
                'num': 2,
            }],
            'perf_metrics': {
                'summary': {
                    'n_samples': 22,
                    'latency': {'mean': 12.80038},
                    'throughput': {'avg_output_tps': 140.06, 'avg_req_ps': 0.0781},
                    'usage': {
                        'input_tokens': {'mean': 6238.136364},
                        'output_tokens': {'mean': 1792.818182},
                        'total_tokens_count': 176681,
                    },
                    'ttft': {'mean': 0.583404, '90%': 1.482557, '99%': 1.721449},
                    'tpot': {'mean': 0.006934, '90%': 0.008364, '99%': 0.008729},
                }
            },
        }),
        encoding='utf-8',
    )
    (preds / f'{bench}__m.jsonl').write_text(
        json.dumps({
            'index': 0,
            'model_output': {'choices': [{'message': {'content': 'file:///tmp/missing-trial'}}]},
        }) + '\n',
        encoding='utf-8',
    )

    row = collect_benchmark(output_dir, bench, 'm')
    assert row['总样本数'] == 22
    assert row['延迟_mean(s)'] == 12.80038
    assert row['TTFT_mean(s)'] == 0.5834
    assert row['TTFT P90'] == 1.48256
    assert row['TPOT_mean(s)'] == 0.00693
    assert row['输入tokens_mean'] == 6238.14
    assert row['累计总tokens'] == 176681
    assert row['输出TPS'] == 140.06


def test_agent_jsonl_does_not_use_trial_wall_clock(tmp_path: Path):
    import json

    output_dir = tmp_path / 'output'
    bench = 'terminal_bench_v2_1'
    reports = output_dir / bench / 'seed_42' / 'reports'
    preds = output_dir / bench / 'seed_42' / 'predictions'
    trial = output_dir / bench / 'seed_42' / 'trials' / 'task__abc'
    reports.mkdir(parents=True)
    preds.mkdir(parents=True)
    (trial / 'agent').mkdir(parents=True)
    (reports / f'{bench}.json').write_text(
        json.dumps({
            'score': 0.0,
            'num': 1,
            'metrics': [{
                'identity': {'name': 'accuracy', 'aggregation': 'mean', 'dimensions': {}},
                'score': 0.0,
                'num': 1,
            }],
        }),
        encoding='utf-8',
    )
    (trial / 'agent' / 'trajectory.json').write_text(
        json.dumps({
            'steps': [
                {'timestamp': '2026-09-03T00:00:00'},
                {'timestamp': '2026-09-03T00:11:46'},
            ],
            'final_metrics': {
                'total_prompt_tokens': 638286,
                'total_completion_tokens': 154502,
            },
        }),
        encoding='utf-8',
    )
    (preds / f'{bench}__m.jsonl').write_text(
        json.dumps({
            'index': 0,
            'model_output': {
                'choices': [{'message': {'content': f'file://{trial}'}}],
            },
        }) + '\n',
        encoding='utf-8',
    )

    row = collect_benchmark(output_dir, bench, 'm')
    assert row['总样本数'] == 1
    assert np.isnan(row['延迟_mean(s)'])
    assert np.isnan(row['TTFT_mean(s)'])
    assert np.isnan(row['累计总tokens'])