agents.py

from crewai import Agent, LLM
from crewai_tools import FileWriterTool
from tools.stats_tool import parse_results, compute_stats
from tools.baseline_tool import get_bigbench_baselines
from tools.benchmark_tool import run_or_load_benchmark
from tools.qa_tool import generate_qa_detail
from tools.viz_tool import generate_pipeline_gif

LOCAL_LLM = LLM(
    model="ollama/llama3.1:8b",
    base_url="http://localhost:11434",
    temperature=0.0,
)


def create_agents():
    file_tool = FileWriterTool()

    # ── 1. Benchmark Runner ────────────────────────────────────────────────────
    runner = Agent(
        role="Benchmark Runner",
        goal="确保 {model_name} 在 {benchmark_name} 上的推理结果文件已存在，"
             "若不存在则自动运行推理并写入 data/results.jsonl",
        backstory=(
            "你是负责 LLM 推理的工程师。"
            "你的任务是检查结果文件是否就绪，没有就立刻触发推理流程。"
        ),
        tools=[run_or_load_benchmark],
        llm=LOCAL_LLM,
        verbose=True,
        max_iter=3,
    )

    # ── 2. Parser ──────────────────────────────────────────────────────────────
    parser = Agent(
        role="Benchmark Data Parser",
        goal="准确读取 {model_name} 在 {benchmark_name} 上的评测结果文件，"
             "提取所有关键字段和数据分布",
        backstory=(
            "你是一位数据工程师，专门处理 LLM benchmark 评测结果。"
            "你只描述数据，不做主观判断。输出结构清晰、数字准确。"
        ),
        tools=[parse_results],
        llm=LOCAL_LLM,
        verbose=True,
        max_iter=4,
    )

    # ── 3. Analyst ────────────────────────────────────────────────────────────
    analyst = Agent(
        role="Statistical Analyst",
        goal="对 {model_name} 的评测数据做完整统计分析，"
             "计算 accuracy、pass@1、错误分布",
        backstory=(
            "你是一位 NLP benchmark 评测研究员，熟悉 BIG-Bench 评测体系。"
            "你只输出数字和客观分析，不做建议。"
        ),
        tools=[compute_stats],
        llm=LOCAL_LLM,
        verbose=True,
        max_iter=4,
    )

    # ── 4. QA Reviewer ────────────────────────────────────────────────────────
    qa_reviewer = Agent(
        role="QA Detail Reviewer",
        goal="对每道题目逐一分析 {model_name} 的回答，"
             "生成完整的逐题对照表并识别错误模式",
        backstory=(
            "你是一位细心的 NLP 评测分析师，专注于逐题分析模型的错误。"
            "你能识别模型的系统性偏差，例如 Yes/No 偏向、因果推理盲点等。"
        ),
        tools=[generate_qa_detail],
        llm=LOCAL_LLM,
        verbose=True,
        max_iter=3,
    )

    # ── 5. Comparator ─────────────────────────────────────────────────────────
    comparator = Agent(
        role="Comparative Research Analyst",
        goal="查找 {benchmark_name} 上已有模型的公开分数，"
             "将 {model_name} 与之横向对比定位",
        backstory=(
            "你是一位 AI 评测领域研究分析师，熟悉 BIG-Bench 排行榜。"
            "你善于从已知基线数据中找到权威数字进行对比分析。"
        ),
        tools=[get_bigbench_baselines],
        llm=LOCAL_LLM,
        verbose=True,
        max_iter=4,
    )

    # ── 6. Report Publisher ───────────────────────────────────────────────────
    publisher = Agent(
        role="Report Publisher",
        goal="将所有分析师的结论整合为专业 Markdown 报告并保存，"
             "同时生成 pipeline 可视化 GIF",
        backstory=(
            "你是一位擅长学术写作和数据可视化的 AI 研究员。"
            "你的报告结构清晰，数字精确，并附带可视化图表说明。"
        ),
        tools=[file_tool, generate_pipeline_gif],
        llm=LOCAL_LLM,
        verbose=True,
        max_iter=5,
    )

    return runner, parser, analyst, qa_reviewer, comparator, publisher