_viz_worker.py

"""
独立进程运行 GIF 生成，避免 pydantic 的 warnings.warn patch 干扰 matplotlib。
用法: python _viz_worker.py <log_file> <out_gif>
"""
import json
import re
import sys

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.animation import FuncAnimation, PillowWriter
from matplotlib import rcParams

# ── 字体 ──────────────────────────────────────────────────────────────────────
CJK_FONT = "PingFang HK"
rcParams["font.family"]        = CJK_FONT
rcParams["axes.unicode_minus"] = False

log_file = sys.argv[1]
out_gif  = sys.argv[2]

with open(log_file, encoding="utf-8") as f:
    entries = [json.loads(l) for l in f if l.strip()]

# ── Agent 定义 ────────────────────────────────────────────────────────────────
AGENTS = [
    "Benchmark Runner",
    "Benchmark Data Parser",
    "Statistical Analyst",
    "QA Detail Reviewer",
    "Comparative Research Analyst",
    "Report Publisher",
]
AGENT_SHORT  = ["Runner", "Parser", "Analyst", "QA", "Comparator", "Publisher"]
AGENT_COLORS = ["#80CBC4", "#4FC3F7", "#81C784", "#FFD54F", "#FFB74D", "#CE93D8"]
BG, FG, GRID, RED = "#1E1E2E", "#CDD6F4", "#313244", "#F38BA8"
N = len(AGENTS)

# ── 工具名映射（英文缩写，避免中文在 monospace 区域乱码）────────────────────
TOOL_SHORT = {
    "benchmark_runner_tool":   "BenchmarkRunner",
    "json_parse_tool":         "JSONParse",
    "stats_tool":              "Stats",
    "qa_detail_tool":          "QADetail",
    "big_bench_baseline_tool": "BaselineLookup",
    "visualize_tool":          "Visualize",
    "file_writer_tool":        "FileWriter",
}

def agent_idx(name: str) -> int:
    nl = name.lower()
    for i, a in enumerate(AGENTS):
        if a.lower() in nl or nl in a.lower():
            return i
    return -1

total_t = max(e["ts"] for e in entries) + 3.0

# ── task spans ────────────────────────────────────────────────────────────────
task_spans: list[tuple[int, float, float]] = []
pending: dict[int, float] = {}
for e in entries:
    ti, ts, ev = e.get("task_index", -1), e["ts"], e["event"]
    if ev == "task_start" and ti >= 0:
        pending[ti] = ts
    elif ev == "task_end" and ti >= 0 and ti in pending:
        task_spans.append((ti, pending.pop(ti), ts))

# ── 右侧摘要：每个 agent 第一条 tool_result 的关键内容 ───────────────────────
STAGE_DEFAULT = [
    "Check / run inference",
    "Parse results.jsonl",
    "Compute accuracy & stats",
    "Per-question comparison",
    "Fetch public baselines",
    "Write report + GIF",
]
stage_result_text: dict[int, str] = {}
for e in entries:
    if e["event"] != "tool_result":
        continue
    ai = agent_idx(e.get("agent", ""))
    if ai < 0 or ai in stage_result_text:
        continue
    raw = e.get("data", {}).get("result", "")
    lines = [l.strip() for l in raw.split("\n") if l.strip()]
    picked = []
    for l in lines:
        if any(k in l for k in ["准确率", "mean=", "pass@1", "错误数", "正确",
                                   "accuracy", "GIF", "已写入", "OK", "results"]):
            picked.append(l[:60])
        if len(picked) >= 2:
            break
    stage_result_text[ai] = "  ".join(picked) if picked else (lines[0][:60] if lines else "")

point_events = [e for e in entries
                if e["event"] in ("tool_call", "tool_result", "agent_final")]

# ═══════════════════════════════════════════════════════════════════════════════
# 布局：左侧时间线（72%) + 右侧摘要（26%）
# 上下：时间线（56%）+ 日志（30%）
# ═══════════════════════════════════════════════════════════════════════════════
FIG_W, FIG_H = 20, 11        # inches  → 3000×1650 px @ 150 dpi
DPI          = 150

fig = plt.figure(figsize=(FIG_W, FIG_H), facecolor=BG)

# axes 坐标 [left, bottom, width, height]
ax_t = fig.add_axes([0.07, 0.36, 0.65, 0.57])   # timeline
ax_l = fig.add_axes([0.02, 0.03, 0.70, 0.28])   # log
ax_r = fig.add_axes([0.75, 0.03, 0.23, 0.92])   # stage results

for ax in (ax_t, ax_l, ax_r):
    ax.set_facecolor(BG)
    for sp in ax.spines.values():
        sp.set_color(GRID)
ax_r.spines["left"].set_visible(True)

# ── 时间线 ────────────────────────────────────────────────────────────────────
ax_t.set_xlim(0, total_t)
ax_t.set_ylim(-0.65, N - 0.35)
ax_t.set_yticks(range(N))
ax_t.set_yticklabels(AGENT_SHORT, color=FG, fontsize=13)
ax_t.tick_params(colors=FG, labelsize=11)
ax_t.set_xlabel("Elapsed time (s)", color=FG, fontsize=12, labelpad=6)
ax_t.set_title(
    "CrewAI Pipeline  |  llama3.1:8b  |  BIG-Bench Hard causal_judgement",
    color=FG, fontsize=14, pad=10, fontweight="bold",
)
ax_t.grid(axis="x", color=GRID, linewidth=0.7, linestyle="--", alpha=0.7)
for i in range(N):
    ax_t.axhline(i, color=GRID, linewidth=0.5, zorder=0)

patches = [mpatches.Patch(color=c, label=s, alpha=0.85)
           for c, s in zip(AGENT_COLORS, AGENT_SHORT)]
ax_t.legend(handles=patches, loc="upper left",
            facecolor=BG, edgecolor=GRID, labelcolor=FG,
            fontsize=10, ncol=3, framealpha=0.85,
            handlelength=1.2, handleheight=1.2)

# ── 日志区 ────────────────────────────────────────────────────────────────────
ax_l.set_xlim(0, 1); ax_l.set_ylim(0, 1); ax_l.axis("off")
ax_l.text(0.0, 1.04, "Live Event Log", transform=ax_l.transAxes,
          color=FG, fontsize=12, va="bottom", fontweight="bold")
log_obj = ax_l.text(
    0.01, 0.97, "",
    transform=ax_l.transAxes,
    va="top", ha="left", color=FG,
    fontsize=9, fontfamily="monospace",
    linespacing=1.5,
)

# ── 右侧摘要面板 ──────────────────────────────────────────────────────────────
ax_r.set_xlim(0, 1); ax_r.set_ylim(0, 1); ax_r.axis("off")
ax_r.text(0.5, 0.985, "Stage Results", transform=ax_r.transAxes,
          color=FG, fontsize=13, va="top", ha="center", fontweight="bold")
ax_r.plot([0, 1], [0.972, 0.972], color=GRID, lw=0.8, transform=ax_r.transAxes)

ROW_H   = 1.0 / (N + 1)          # 每个 agent 占的高度比例
result_objs: list = []
for i in range(N):
    # y_top：从顶部往下排列
    y_top = 0.955 - i * (0.940 / N)
    col   = AGENT_COLORS[i]

    # 彩色小方块 + agent 名
    ax_r.add_patch(mpatches.FancyBboxPatch(
        (0.02, y_top - 0.018), 0.06, 0.032,
        boxstyle="round,pad=0.004",
        facecolor=col, edgecolor="none",
        transform=ax_r.transAxes, zorder=3,
    ))
    ax_r.text(0.11, y_top - 0.002, AGENT_SHORT[i],
              transform=ax_r.transAxes,
              color=col, fontsize=11, fontweight="bold", va="center")

    # 结果摘要文字（初始灰色省略号）
    obj = ax_r.text(
        0.04, y_top - 0.055, "waiting...",
        transform=ax_r.transAxes,
        color="#666888", fontsize=9, va="top",
        fontfamily=CJK_FONT,
    )
    result_objs.append(obj)

    # 分隔线
    ax_r.plot([0.02, 0.98], [y_top - 0.105, y_top - 0.105],
              color=GRID, lw=0.6, transform=ax_r.transAxes)

# ── 游标 ──────────────────────────────────────────────────────────────────────
cursor = ax_t.axvline(0, color=RED, linewidth=2.0, zorder=10, alpha=0.9)

# 时间戳标注（游标顶部）
time_label = ax_t.text(0, N - 0.2, "0.0s", color=RED, fontsize=9,
                        ha="center", va="bottom", zorder=11)

# ── 动画参数 ──────────────────────────────────────────────────────────────────
FPS        = 8
FRAMES     = 100
LOG_LINES  = 10
times      = [i * total_t / (FRAMES - 1) for i in range(FRAMES)]

drawn_spans, drawn_points, filled_results = set(), set(), set()


def update(frame):
    t = times[frame]
    cursor.set_xdata([t, t])
    time_label.set_x(min(t, total_t * 0.97))
    time_label.set_text(f"{t:.0f}s")

    # task bars
    for key, (ai, t0, t1) in enumerate(task_spans):
        if t1 <= t and key not in drawn_spans:
            col = AGENT_COLORS[ai] if ai < N else "#90A4AE"
            ax_t.barh(ai, t1 - t0, left=t0, height=0.55,
                      color=col, alpha=0.78, zorder=3)
            dur  = t1 - t0
            label = f"{dur:.0f}s"
            ax_t.text((t0 + t1) / 2, ai, label,
                      ha="center", va="center",
                      fontsize=11, color="white",
                      fontweight="bold", zorder=5)
            drawn_spans.add(key)

            # 右侧摘要填充
            if ai < N and ai not in filled_results:
                txt = stage_result_text.get(ai, STAGE_DEFAULT[ai] if ai < len(STAGE_DEFAULT) else "")
                result_objs[ai].set_text(txt[:70])
                result_objs[ai].set_color(FG)
                filled_results.add(ai)

    # 点事件（tool_call / tool_result / agent_final）
    for key, e in enumerate(point_events):
        if e["ts"] <= t and key not in drawn_points:
            ai  = agent_idx(e.get("agent", ""))
            ev  = e["event"]
            col = AGENT_COLORS[ai] if 0 <= ai < N else "#90A4AE"
            if 0 <= ai < N:
                mk  = "D" if ev == "tool_call" else ("*" if ev == "agent_final" else "^")
                sz  = 80 if ev == "agent_final" else 55
                ax_t.scatter(e["ts"], ai, s=sz, marker=mk,
                             color=col, edgecolors="white",
                             linewidths=0.7, zorder=7)
            drawn_points.add(key)

    # 日志滚动
    visible = [e for e in entries if e["ts"] <= t][-LOG_LINES:]
    lines = []
    for e in visible:
        icon = {"task_start": ">>", "task_end": "OK",
                "tool_call": " T", "tool_result": " R",
                "agent_final": " *"}.get(e["event"], "  ")
        ag_s = e.get("agent", "?")
        for full, short in zip(AGENTS, AGENT_SHORT):
            ag_s = ag_s.replace(full, short)
        ag_s = f"{ag_s[:10]:<10}"
        det = ""
        if e["event"] == "tool_call":
            raw_tool = e.get("data", {}).get("tool", "")
            det = TOOL_SHORT.get(raw_tool, raw_tool)[:30]
        elif e["event"] == "tool_result":
            raw = re.sub(r'\s+', ' ', e.get("data", {}).get("result", ""))
            det = re.sub(r'[^\x00-\x7F]+', '', raw).strip()[:36]
        elif e["event"] == "agent_final":
            det = "==> Final Answer"
        elif e["event"] == "task_end":
            det = "==> DONE"
        lines.append(f"{e['ts']:7.1f}s {icon} [{ag_s}] {e['event']:<16} {det}")
    log_obj.set_text("\n".join(lines))

    return [cursor, time_label, log_obj] + result_objs


anim = FuncAnimation(fig, update, frames=FRAMES,
                     interval=1000 // FPS, blit=False)
anim.save(out_gif, writer=PillowWriter(fps=FPS), dpi=DPI)
plt.close(fig)
print(f"OK:{out_gif}")