EvalHarness/evalharness/viz/renderers/excel.py

"""Excel renderer: multi-sheet workbook from EvalReports (xlsxwriter).

Sheets:
  1. Summary   -- one row per benchmark: identity + score + quality + perf dashboard
  2. Perf      -- detailed latency/ttft/tpot/token columns
  3. Categories-- per-benchmark category breakdown
  4. Samples   -- per-sample drill-down (first N)

    render([rep1, rep2], style='excel')            -> bytes/str path via CLI
    evalharness viz show r1.json r2.json --style excel  # writes xlsx next to inputs
"""

import json
from typing import Dict, List, Union

from ...eval.record import EvalReport
from .. import register_renderer

# dashboard column spec: (header, source-key, formatter)
_SUMMARY_COLS = [
    ('benchmark', None, None), ('model', None, None), ('metric', None, None),
    ('score', None, 'pct'), ('num_samples', None, 'int'),
    ('extract_fail', None, 'int'), ('time_h', None, 'f2'),
    ('success_rate', 'success_rate', 'pct'), ('latency_mean_s', 'latency_mean_s', 'f3'),
    ('output_tps', 'output_tps', 'f2'), ('request_qps', 'request_qps', 'f4'),
    ('input_tokens_mean', 'input_tokens_mean', 'f1'), ('output_tokens_mean', 'output_tokens_mean', 'f1'),
    ('total_tokens', 'total_tokens', 'int'),
    ('ttft_mean_s', 'ttft_mean_s', 'f3'), ('ttft_p90_s', 'ttft_p90_s', 'f3'),
    ('ttft_p99_s', 'ttft_p99_s', 'f3'),
    ('tpot_mean_s', 'tpot_mean_s', 'f4'), ('tpot_p90_s', 'tpot_p90_s', 'f4'),
    ('tpot_p99_s', 'tpot_p99_s', 'f4'),
    ('retry_rate', 'retry_rate', 'pct'),
]


def _row_for(rep: EvalReport) -> dict:
    perf = rep.metric_groups.get('perf') or {}
    primary = next((k for k in rep.metrics if k != 'extraction_failure_rate'), '')
    secs = sum(float((s.usage or {}).get('latency_s', 0) or 0) for s in rep.samples)
    return {
        'benchmark': rep.dataset, 'model': rep.model or '?', 'metric': primary,
        'score': rep.metrics.get(primary, 0), 'num_samples': rep.num_samples,
        'extract_fail': rep.num_failed_extractions,
        'time_h': round(secs / 3600, 2),
        **{k: v for k, v in perf.items() if v is not None},
    }


@register_renderer('excel')
def excel_workbook(target: Union['EvalReport', List['EvalReport']], opts: Dict) -> str:
    import xlsxwriter

    reports = [r for r in (target if isinstance(target, list) else [target])
               if isinstance(r, EvalReport)]
    if not reports:
        return '(no reports)'
    out_path = opts.get('out') or str(opts.get('dir', '.')) + f'/evalharness_report.xlsx'

    wb = xlsxwriter.Workbook(out_path)
    wb.set_properties({'title': 'EvalHarness Report',
                       'comments': 'generated by evalharness viz --style excel'})

    # formats
    f_hdr = wb.add_format({'bold': True, 'bg_color': '#1F2937', 'font_color': 'white',
                           'border': 1, 'align': 'center', 'valign': 'vcenter'})
    f_pct = wb.add_format({'num_format': '0.0%'})
    f_int = wb.add_format({'num_format': '#,##0'})
    f_f1 = wb.add_format({'num_format': '0.0'})
    f_f2 = wb.add_format({'num_format': '0.00'})
    f_f3 = wb.add_format({'num_format': '0.000'})
    f_f4 = wb.add_format({'num_format': '0.0000'})
    fmt_map = {'pct': f_pct, 'int': f_int, 'f1': f_f1, 'f2': f_f2, 'f3': f_f3, 'f4': f_f4}

    # ---- sheet 1: Summary ----
    ws = wb.add_worksheet('Summary')
    ws.freeze_panes(1, 2)
    for c, (hdr, _, _) in enumerate(_SUMMARY_COLS):
        ws.write(0, c, hdr, f_hdr)
    rows = [_row_for(r) for r in reports]
    for ri, row in enumerate(rows, start=1):
        for c, (hdr, key, fmt) in enumerate(_SUMMARY_COLS):
            val = row.get(hdr)
            if val is None:
                ws.write(ri, c, '')
            elif fmt:
                ws.write_number(ri, c, float(val), fmt_map[fmt])
            else:
                ws.write(ri, c, val)
    ws.autofilter(0, 0, len(rows), len(_SUMMARY_COLS) - 1)
    for c, (hdr, _, _) in enumerate(_SUMMARY_COLS):
        ws.set_column(c, c, max(12, min(22, len(hdr) + 4)))

    # ---- sheet 2: Perf detail ----
    perf_keys = ['n_requests', 'latency_mean_s', 'latency_p50_s', 'latency_p90_s',
                 'latency_p95_s', 'latency_p99_s', 'ttft_mean_s', 'ttft_p50_s',
                 'ttft_p90_s', 'ttft_p99_s', 'tpot_mean_s', 'tpot_p90_s', 'tpot_p99_s',
                 'itl_mean_s', 'output_tps', 'request_qps', 'input_tokens',
                 'output_tokens', 'input_tokens_mean', 'output_tokens_mean',
                 'total_tokens', 'success_rate', 'retry_rate', 'wall_latency_s']
    ws2 = wb.add_worksheet('Perf')
    ws2.freeze_panes(1, 1)
    ws2.write(0, 0, 'benchmark', f_hdr)
    for c, k in enumerate(perf_keys, start=1):
        ws2.write(0, c, k, f_hdr)
    for ri, rep in enumerate(reports, start=1):
        ws2.write(ri, 0, rep.dataset)
        perf = rep.metric_groups.get('perf') or {}
        for c, k in enumerate(perf_keys, start=1):
            v = perf.get(k)
            if isinstance(v, (int, float)):
                ws2.write_number(ri, c, v)
            else:
                ws2.write(ri, c, '' if v is None else str(v))
    ws2.autofilter(0, 0, len(reports), len(perf_keys))

    # ---- sheet 3: Categories ----
    ws3 = wb.add_worksheet('Categories')
    ws3.freeze_panes(1, 1)
    ws3.write(0, 0, 'benchmark', f_hdr)
    ws3.write(0, 1, 'group', f_hdr)
    ws3.write(0, 2, 'subgroup', f_hdr)
    ws3.write(0, 3, 'score', f_hdr)
    r3 = 1
    for rep in reports:
        for gname, groups in rep.metric_groups.items():
            if gname in ('perf', 'run_info') or gname.startswith('agg_error') \
                    or not isinstance(groups, dict):
                continue
            for g, v in groups.items():
                if not isinstance(v, (int, float)):
                    continue
                ws3.write(r3, 0, rep.dataset)
                ws3.write(r3, 1, gname)
                ws3.write(r3, 2, str(g)[:80])
                ws3.write_number(r3, 3, v, f_pct)
                r3 += 1
    ws3.autofilter(0, 0, max(r3 - 1, 1), 3)

    # ---- sheet 4: Samples (first 300) ----
    ws4 = wb.add_worksheet('Samples')
    hdr4 = ['benchmark', 'sample_id', 'correct', 'score', 'latency_s', 'ttft_s',
            'output_tokens', 'retries', 'extract_ok', 'extracted', 'target']
    for c, h in enumerate(hdr4):
        ws4.write(0, c, h, f_hdr)
    r4 = 1
    for rep in reports:
        primary = next((k for k in rep.metrics if k != 'extraction_failure_rate'), '')
        for s in rep.samples[:300]:
            u = s.usage or {}
            ws4.write(r4, 0, rep.dataset)
            ws4.write(r4, 1, s.sample_id if s.sample_id is not None else r4)
            ws4.write(r4, 2, 1 if s.scores.get(primary, 0) >= 1 else 0)
            ws4.write_number(r4, 3, s.scores.get(primary, 0), f_f3)
            ws4.write_number(r4, 4, float(u.get('latency_s', 0) or 0), f_f3)
            tt = u.get('ttft_s')
            ws4.write(r4, 5, tt if tt is not None else '')
            ws4.write_number(r4, 6, int(u.get('output_tokens', 0) or 0), f_int)
            ws4.write_number(r4, 7, int(u.get('retries', 0) or 0), f_int)
            ws4.write(r4, 8, 1 if s.extraction_ok else 0)
            ws4.write(r4, 9, str(s.extracted_prediction)[:120])
            ws4.write(r4, 10, str(s.target)[:80])
            r4 += 1
    ws4.autofilter(0, 0, max(r4 - 1, 1), len(hdr4) - 1)
    ws4.set_column(9, 10, 40)

    wb.close()
    return out_path