zk-data-agent/skills/online-replay-diff/scripts/process_replay_diff.py

#!/usr/bin/env python3
"""Process online traffic replay diff Excel files and export labeling workbooks."""

from __future__ import annotations

import argparse
import json
import math
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import pandas as pd


SCRIPT_DIR = Path(__file__).resolve().parent
SKILL_DIR = SCRIPT_DIR.parent
DEFAULT_EQUIV_CONFIG = SKILL_DIR / "knowledge" / "equivalence_groups.json"

REQUIRED_DOMAIN_COLUMNS = ["Preview环境domain", "Ptr环境domain"]

KEEP_COLUMNS = [
    "Query",
    "PV",
    "diff结果",
    "reviewer",
    "GSB",
    "Preview环境requestId",
    "Ptr环境requestId",
    "Preview环境toSpeak",
    "Ptr环境toSpeak",
    "Preview环境domain",
    "Ptr环境domain",
    "Preview环境normCode",
    "Ptr环境normCode",
    "Preview环境funcCategory",
    "Ptr环境funcCategory",
    "Preview环境funcName",
    "Ptr环境funcName",
    "Preview环境agentType",
    "Ptr环境agentType",
    "Preview环境copilotCode",
    "Ptr环境copilotCode",
    "Preview环境promptString",
    "Ptr环境promptString",
    "Preview环境promptModel",
    "Ptr环境promptModel",
    "baseRequestId",
    "回放时间",
    "is_same",
]


@dataclass
class FileResult:
    source: Path
    sheet_name: str
    total: int
    same: int
    diff: int
    diff_rate: float
    normalized_path: Path


def load_equivalence_config(path: Path) -> tuple[list[list[str]], list[tuple[str, str]]]:
    if not path.exists():
        raise FileNotFoundError(f"等价规则不存在: {path}")
    payload = json.loads(path.read_text(encoding="utf-8"))
    groups = payload.get("groups", [])
    pairs = payload.get("pairs", [])
    if not isinstance(groups, list) or not isinstance(pairs, list):
        raise ValueError("等价规则 JSON 必须包含 list 类型的 groups / pairs")
    normalized_groups = [[str(item).strip() for item in group if str(item).strip()] for group in groups]
    normalized_pairs = []
    for pair in pairs:
        if not isinstance(pair, list) or len(pair) != 2:
            raise ValueError(f"pairs 中每项必须是长度为 2 的数组: {pair}")
        normalized_pairs.append((str(pair[0]).strip(), str(pair[1]).strip()))
    return normalized_groups, normalized_pairs


def build_is_same(groups: list[list[str]], pairs: list[tuple[str, str]]):
    parent: dict[str, str] = {}

    def find(x: str) -> str:
        parent.setdefault(x, x)
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    def union(a: str, b: str) -> None:
        ra, rb = find(a), find(b)
        if ra != rb:
            parent[rb] = ra

    for group in groups:
        if len(group) < 2:
            continue
        for item in group[1:]:
            union(group[0], item)
    for left, right in pairs:
        union(left, right)

    def normalize(value: Any) -> str | None:
        if value is None:
            return None
        if isinstance(value, float) and math.isnan(value):
            return None
        text = str(value).strip()
        return text if text else None

    def is_same(left: Any, right: Any) -> bool:
        left_text = normalize(left)
        right_text = normalize(right)
        if left_text is None or right_text is None:
            return False
        if left_text == right_text:
            return True
        return find(left_text) == find(right_text)

    return is_same


def collect_input_files(args: argparse.Namespace) -> list[Path]:
    files: list[Path] = []
    for item in args.inputs or []:
        path = Path(item).expanduser().resolve()
        if path.is_dir():
            files.extend(sorted(path.glob(args.glob)))
        else:
            files.append(path)
    if args.input_dir:
        files.extend(sorted(Path(args.input_dir).expanduser().resolve().glob(args.glob)))
    unique: list[Path] = []
    seen = set()
    for path in files:
        if path.name.startswith("~$"):
            continue
        if path.suffix.lower() not in {".xls", ".xlsx"}:
            continue
        if path not in seen:
            unique.append(path)
            seen.add(path)
    return unique


def find_header_row(raw: pd.DataFrame) -> int:
    max_scan = min(10, len(raw))
    required = set(["Query", *REQUIRED_DOMAIN_COLUMNS])
    for idx in range(max_scan):
        row_values = {str(value).strip() for value in raw.iloc[idx].tolist() if not pd.isna(value)}
        if required.issubset(row_values):
            return idx
    raise ValueError("前 10 行未找到包含 Query / Preview环境domain / Ptr环境domain 的表头")


def make_unique_columns(columns: list[Any]) -> list[str]:
    seen: dict[str, int] = {}
    result = []
    for col in columns:
        name = str(col).strip()
        if not name or name.lower() == "nan":
            name = "Unnamed"
        count = seen.get(name, 0)
        seen[name] = count + 1
        result.append(name if count == 0 else f"{name}.{count}")
    return result


def read_replay_excel(path: Path) -> pd.DataFrame:
    engine = "xlrd" if path.suffix.lower() == ".xls" else "openpyxl"
    raw = pd.read_excel(path, header=None, engine=engine)
    header_row = find_header_row(raw)
    df = raw.iloc[header_row + 1 :].copy()
    df.columns = make_unique_columns(raw.iloc[header_row].tolist())
    df = df.dropna(how="all").reset_index(drop=True)
    missing = [col for col in REQUIRED_DOMAIN_COLUMNS if col not in df.columns]
    if missing:
        raise ValueError(f"{path.name} 缺少必要列: {missing}")
    return df


def sanitize_sheet_name(name: str) -> str:
    name = re.sub(r"[\[\]:*?/\\]", "_", name).strip()
    return name[:31] if name else "Sheet"


def make_unique_sheet_name(name: str, used: set[str]) -> str:
    base = sanitize_sheet_name(name)
    candidate = base
    index = 1
    while candidate in used:
        suffix = f"_{index}"
        candidate = base[: 31 - len(suffix)] + suffix
        index += 1
    used.add(candidate)
    return candidate


def is_false_value(value: Any) -> bool:
    if pd.isna(value):
        return False
    if isinstance(value, bool):
        return value is False
    if isinstance(value, (int, float)):
        return value == 0
    return str(value).strip().lower() in {"false", "0"}


def write_markdown_summary(path: Path, results: list[FileResult]) -> None:
    total = sum(item.total for item in results)
    diff = sum(item.diff for item in results)
    same = sum(item.same for item in results)
    overall_rate = diff / total if total else 0.0
    lines = [
        "# 线上流量回放 diff 处理结果",
        "",
        "## 处理口径",
        "",
        "- 自动识别回放 Excel 中的真实表头，兼容首行 `DiffCase数据` 标题行。",
        "- 按 `Preview环境domain` 与 `Ptr环境domain` 计算 `is_same`。",
        "- 最终 diff 按 `is_same=False` 统计，待标注数据也只保留这些行。",
        "",
        "## diff 率",
        "",
        "| 文件 | 总行数 | same 数 | diff 数 | diff 率 |",
        "|---|---:|---:|---:|---:|",
    ]
    for item in results:
        lines.append(
            f"| {item.sheet_name} | {item.total:,} | {item.same:,} | {item.diff:,} | {item.diff_rate:.4%} |"
        )
    lines.append(f"| 合计 | {total:,} | {same:,} | {diff:,} | {overall_rate:.4%} |")
    lines.append("")
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def process_files(args: argparse.Namespace) -> dict[str, Any]:
    inputs = collect_input_files(args)
    if not inputs:
        raise FileNotFoundError("未找到输入 Excel 文件")

    output_dir = Path(args.output_dir).expanduser().resolve()
    normalized_dir = output_dir / "normalized"
    output_dir.mkdir(parents=True, exist_ok=True)
    normalized_dir.mkdir(parents=True, exist_ok=True)

    groups, pairs = load_equivalence_config(Path(args.equiv_config).expanduser().resolve())
    is_same = build_is_same(groups, pairs)

    used_sheet_names: set[str] = set()
    processed: list[tuple[str, pd.DataFrame]] = []
    results: list[FileResult] = []

    for input_path in inputs:
        df = read_replay_excel(input_path)
        normalized_path = normalized_dir / f"{input_path.stem}.xlsx"
        df.to_excel(normalized_path, index=False, engine="openpyxl")

        df["is_same"] = df.apply(
            lambda row: is_same(row["Preview环境domain"], row["Ptr环境domain"]),
            axis=1,
        )
        for col in KEEP_COLUMNS:
            if col not in df.columns:
                df[col] = ""
        result_df = df[KEEP_COLUMNS].copy()
        sheet_name = make_unique_sheet_name(input_path.stem, used_sheet_names)
        total = len(result_df)
        diff = int(result_df["is_same"].apply(is_false_value).sum())
        same = int((result_df["is_same"] == True).sum())
        rate = diff / total if total else 0.0
        processed.append((sheet_name, result_df))
        results.append(
            FileResult(
                source=input_path,
                sheet_name=sheet_name,
                total=total,
                same=same,
                diff=diff,
                diff_rate=rate,
                normalized_path=normalized_path,
            )
        )

    summary_workbook = output_dir / f"{args.prefix}_汇总结果.xlsx"
    false_workbook = output_dir / f"{args.prefix}_待标注_仅False.xlsx"
    summary_json = output_dir / f"{args.prefix}_summary.json"
    summary_csv = output_dir / f"{args.prefix}_summary.csv"
    summary_md = output_dir / f"{args.prefix}_summary.md"

    with pd.ExcelWriter(summary_workbook, engine="openpyxl") as writer:
        for sheet_name, df in processed:
            df.to_excel(writer, sheet_name=sheet_name, index=False)

    with pd.ExcelWriter(false_workbook, engine="openpyxl") as writer:
        kept = 0
        for sheet_name, df in processed:
            filtered_df = df[df["is_same"].apply(is_false_value)].copy()
            if not filtered_df.empty:
                filtered_df.to_excel(writer, sheet_name=sheet_name, index=False)
                kept += 1
        if kept == 0:
            pd.DataFrame({"说明": ["所有 sheet 都没有 is_same=False 的数据"]}).to_excel(
                writer,
                sheet_name="result",
                index=False,
            )

    summary_rows = [
        {
            "file": str(item.source),
            "sheet": item.sheet_name,
            "total": item.total,
            "same": item.same,
            "diff": item.diff,
            "diff_rate": item.diff_rate,
            "normalized_file": str(item.normalized_path),
        }
        for item in results
    ]
    total = sum(item.total for item in results)
    diff = sum(item.diff for item in results)
    same = sum(item.same for item in results)
    payload = {
        "total": total,
        "same": same,
        "diff": diff,
        "diff_rate": diff / total if total else 0.0,
        "files": summary_rows,
        "outputs": {
            "summary_workbook": str(summary_workbook),
            "false_workbook": str(false_workbook),
            "summary_json": str(summary_json),
            "summary_csv": str(summary_csv),
            "summary_md": str(summary_md),
        },
    }
    summary_json.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
    pd.DataFrame(summary_rows).to_csv(summary_csv, index=False, encoding="utf-8-sig")
    write_markdown_summary(summary_md, results)
    return payload


def main() -> int:
    parser = argparse.ArgumentParser(description="处理线上流量回放 diff Excel，计算 diff 率并输出待标注数据")
    parser.add_argument("inputs", nargs="*", help="输入 .xls/.xlsx 文件或目录")
    parser.add_argument("--input-dir", help="输入目录；会按 --glob 扫描")
    parser.add_argument("--glob", default="*.xls*", help="目录扫描通配符，默认 *.xls*")
    parser.add_argument("--output-dir", default="output", help="输出目录，默认 ./output")
    parser.add_argument("--prefix", default="replay_diff", help="输出文件名前缀")
    parser.add_argument("--equiv-config", default=str(DEFAULT_EQUIV_CONFIG), help="domain 等价规则 JSON")
    args = parser.parse_args()

    payload = process_files(args)
    print(json.dumps(payload, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())