ai-agent-book/chapter9/gaia-experience/run_with_experience.py

"""
Modified GAIA Runner with Learning from Experience
This script extends the original GAIA runner with experience learning capabilities.
"""

import argparse
import json
import logging
import os
import re
import traceback
from pathlib import Path
from typing import Any, Dict, List
from dotenv import load_dotenv

from llm_env import resolve_llm, DEFAULT_MODEL

# The heavy runtime dependencies (AWorld framework + local experience
# components, which pull in sentence-transformers / faiss) are imported lazily
# via _load_runtime() so that `--help` and argument parsing work even when the
# full stack is not installed. The names below are populated on first run.
AgentConfig = TaskConfig = Task = None
system_prompt = None
add_file_path = load_dataset_meta = question_scorer = report_results = None
ExperienceAgent = KnowledgeBase = TrajectorySummarizer = None


def _load_runtime():
    """Import AWorld and experience components lazily (see note above)."""
    global AgentConfig, TaskConfig, Task, system_prompt
    global add_file_path, load_dataset_meta, question_scorer, report_results
    global ExperienceAgent, KnowledgeBase, TrajectorySummarizer

    from AWorld.aworld.config.conf import AgentConfig as _AgentConfig
    from AWorld.aworld.config.conf import TaskConfig as _TaskConfig
    from AWorld.aworld.core.task import Task as _Task
    from AWorld.examples.gaia.prompt import system_prompt as _system_prompt
    from AWorld.examples.gaia.utils import (
        add_file_path as _add_file_path,
        load_dataset_meta as _load_dataset_meta,
        question_scorer as _question_scorer,
        report_results as _report_results,
    )
    from experience_agent import ExperienceAgent as _ExperienceAgent
    from knowledge_base import KnowledgeBase as _KnowledgeBase
    from trajectory_summarizer import TrajectorySummarizer as _TrajectorySummarizer

    AgentConfig, TaskConfig, Task = _AgentConfig, _TaskConfig, _Task
    system_prompt = _system_prompt
    add_file_path, load_dataset_meta = _add_file_path, _load_dataset_meta
    question_scorer, report_results = _question_scorer, _report_results
    ExperienceAgent = _ExperienceAgent
    KnowledgeBase = _KnowledgeBase
    TrajectorySummarizer = _TrajectorySummarizer

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def parse_arguments():
    """解析命令行参数（Parse command line arguments）。"""
    parser = argparse.ArgumentParser(
        prog="run_with_experience.py",
        description=(
            "在 GAIA 基准上运行带“从经验中学习”能力的 Agent（实验 9-2：经验知识文档）。\n"
            "支持两种模式：学习模式（learning-mode）在任务成功后把轨迹提炼成经验并入库；\n"
            "应用模式（apply-experience）在解题前检索最相似的历史经验注入系统提示词。\n"
            "使用 --compare 可自动做 A/B 对照，直观展示“复用经验是否提升 GAIA 成绩”。"
        ),
        epilog=(
            "示例（Examples）：\n"
            "  # 1) 仅解析参数、查看帮助（无需 API/数据集）\n"
            "  python run_with_experience.py --help\n\n"
            "  # 2) 学习模式：从前 10 道题的成功轨迹中沉淀经验\n"
            "  python run_with_experience.py --learning-mode --start 0 --end 10\n\n"
            "  # 3) 应用模式：用已学到的经验解 10~20 题\n"
            "  python run_with_experience.py --apply-experience --start 10 --end 20\n\n"
            "  # 4) A/B 对照：同一批题分别在“无经验/有经验”下各跑一次并对比准确率\n"
            "  #    （先用 learning-mode 在其它题上积累经验，再在未见过的题上对照，避免数据泄漏）\n"
            "  python run_with_experience.py --compare --start 10 --end 20 \\\n"
            "      --experience-db ./learned_experiences.json\n\n"
            "  # 5) 指定主 Agent 模型与结果输出路径\n"
            "  python run_with_experience.py --apply-experience --model gpt-5.6-luna \\\n"
            "      --output ./results/exp_run.json --start 0 --end 5\n"
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    # ---- 任务集选择（Task set selection）----
    parser.add_argument(
        "--start",
        type=int,
        default=0,
        help="数据集起始下标（默认：0）",
    )
    parser.add_argument(
        "--end",
        type=int,
        default=20,
        help="数据集结束下标（不含，默认：20）",
    )
    parser.add_argument(
        "--q",
        type=str,
        help="指定单个题目下标或 task_id；优先级最高，会覆盖 --start/--end",
    )
    parser.add_argument(
        "--skip",
        action="store_true",
        help="若题目此前已处理过则跳过",
    )
    parser.add_argument(
        "--split",
        type=str,
        default="validation",
        help="数据集划分，如 validation、test（默认：validation）",
    )
    parser.add_argument(
        "--blacklist_file_path",
        type=str,
        nargs="?",
        help="黑名单文件路径，如 blacklist.txt（其中的 task_id 会被跳过）",
    )

    # ---- 模型与输出（Model & output）----
    parser.add_argument(
        "--model",
        type=str,
        default=None,
        help="主 Agent 使用的模型名；缺省时读取环境变量 LLM_MODEL_NAME（默认 gpt-5.6-luna）",
    )
    parser.add_argument(
        "--summary-model",
        type=str,
        default=DEFAULT_MODEL,
        help="用于轨迹总结（经验提炼）的模型（默认：gpt-5.6-luna）",
    )
    parser.add_argument(
        "--embedding-model",
        type=str,
        default="all-MiniLM-L6-v2",
        help="用于语义检索的 sentence-transformers 模型（默认：all-MiniLM-L6-v2）",
    )
    parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="结果 JSON 输出路径；缺省写入 $AWORLD_WORKSPACE/experience_results.json",
    )

    # ---- 经验学习开关（Experience learning）----
    parser.add_argument(
        "--learning-mode",
        action="store_true",
        help="启用学习模式：捕获成功轨迹并总结为可复用经验",
    )
    parser.add_argument(
        "--apply-experience",
        action="store_true",
        help="启用应用模式：为新任务检索并注入相关历史经验",
    )
    parser.add_argument(
        "--compare",
        action="store_true",
        help="A/B 对照模式：同一批题分别在“无经验/有经验”下各跑一次并对比准确率",
    )
    parser.add_argument(
        "--preload-kb",
        action="store_true",
        help="从 gaia-validation.jsonl 预加载知识库（注意：勿在评测同一批题上预加载，以免泄漏答案）",
    )
    parser.add_argument(
        "--kb-path",
        type=str,
        default="./kb_index",
        help="知识库索引存储路径（默认：./kb_index）",
    )
    parser.add_argument(
        "--experience-db",
        type=str,
        default="./learned_experiences.json",
        help="已学习经验的存储路径（默认：./learned_experiences.json）",
    )
    parser.add_argument(
        "--validation-file",
        type=str,
        default="gaia-validation.jsonl",
        help="用于预加载的 gaia-validation.jsonl 路径",
    )

    return parser.parse_args()


def setup_logging(args):
    """Setup logging configuration."""
    workspace = os.getenv("AWORLD_WORKSPACE", ".")
    os.makedirs(workspace, exist_ok=True)
    
    log_file_name = f"experience_agent_{args.q}.log" if args.q else f"experience_agent_{args.start}_{args.end}.log"
    file_handler = logging.FileHandler(
        os.path.join(workspace, log_file_name),
        mode="a",
        encoding="utf-8",
    )
    file_handler.setLevel(logging.INFO)
    
    formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
    file_handler.setFormatter(formatter)
    
    logging.getLogger().addHandler(file_handler)


def load_mcp_config(config_path: Path) -> Dict[str, Any]:
    """Load MCP configuration."""
    mcp_config = {}
    available_servers = []
    
    try:
        if config_path.exists():
            with open(config_path, mode="r", encoding="utf-8") as f:
                mcp_config = json.loads(f.read())
                available_servers = list(mcp_config.get("mcpServers", {}).keys())
                logger.info(f"🔧 MCP Available Servers: {available_servers}")
    except json.JSONDecodeError as e:
        logger.error(f"Error loading mcp_collections.json: {e}")
    
    return mcp_config, available_servers


async def run_with_experience(args):
    """Main execution function with experience learning."""
    _load_runtime()

    # Load dataset
    gaia_dataset_path = os.getenv("GAIA_DATASET_PATH", "./gaia_dataset")
    full_dataset = load_dataset_meta(gaia_dataset_path, split=args.split)
    logger.info(f"Total questions: {len(full_dataset)}")
    
    # Load MCP configuration
    mcp_config_path = Path(__file__).parent / "AWorld" / "examples" / "gaia" / "mcp.json"
    mcp_config, available_servers = load_mcp_config(mcp_config_path)
    
    # Setup agent configuration. The main agent model can be overridden via the
    # --model CLI flag; otherwise it falls back to the LLM_MODEL_NAME env var.
    # OpenAI 直连优先，缺 Key 时自动走 OpenRouter 兜底（见 llm_env.resolve_llm）。
    agent_config = AgentConfig(
        **resolve_llm(model_override=args.model),
        llm_temperature=float(os.getenv("LLM_TEMPERATURE", "0.0"))
    )
    
    # Initialize knowledge base
    knowledge_base = None
    if args.apply_experience or args.preload_kb:
        logger.info("Initializing knowledge base...")
        knowledge_base = KnowledgeBase(
            index_path=args.kb_path,
            model_name=args.embedding_model
        )
        
        # Preload validation data if requested
        if args.preload_kb and os.path.exists(args.validation_file):
            logger.info(f"Preloading knowledge base from {args.validation_file}")
            knowledge_base.index_gaia_validation(args.validation_file)
            stats = knowledge_base.get_statistics()
            logger.info(f"Knowledge base statistics: {stats}")
    
    # Initialize trajectory summarizer
    summarizer = None
    if args.learning_mode:
        logger.info("Initializing trajectory summarizer...")
        summarizer = TrajectorySummarizer(
            llm_config=agent_config,
            model_name=resolve_llm(model_override=args.summary_model)["llm_model_name"]
        )
    
    # Create experience agent
    experience_agent = ExperienceAgent(
        conf=agent_config,
        name="gaia_experience_agent",
        system_prompt=system_prompt,
        learning_mode=args.learning_mode,
        apply_experience=args.apply_experience,
        experience_db_path=args.experience_db,
        knowledge_base=knowledge_base,
        summarizer=summarizer,
        mcp_config=mcp_config,
        mcp_servers=available_servers,
    )
    
    logger.info(f"Experience Agent initialized:")
    logger.info(f"  - Learning mode: {args.learning_mode}")
    logger.info(f"  - Apply experience: {args.apply_experience}")
    logger.info(f"  - Knowledge base: {'Yes' if knowledge_base else 'No'}")
    logger.info(f"  - Summarizer: {'Yes' if summarizer else 'No'}")
    
    # Load existing results (path overridable via --output)
    results_file = args.output or os.path.join(
        os.getenv("AWORLD_WORKSPACE", "."), "experience_results.json"
    )
    results_dir = os.path.dirname(results_file)
    if results_dir:
        os.makedirs(results_dir, exist_ok=True)
    if os.path.exists(results_file):
        with open(results_file, "r", encoding="utf-8") as f:
            results = json.load(f)
    else:
        results = []
    
    # Load blacklist
    blacklist = set()
    if args.blacklist_file_path and os.path.exists(args.blacklist_file_path):
        with open(args.blacklist_file_path, "r", encoding="utf-8") as f:
            blacklist = set(f.read().splitlines())
    
    try:
        # Determine dataset slice
        if args.q:
            dataset_slice = [
                record for record in full_dataset 
                if record["task_id"] == args.q
            ]
        else:
            dataset_slice = full_dataset[args.start:args.end]
        
        # Process each question
        for i, dataset_i in enumerate(dataset_slice):
            # Check if should skip
            if not args.q:
                if dataset_i["task_id"] in blacklist:
                    logger.info(f"Skipping blacklisted task: {dataset_i['task_id']}")
                    continue
                
                if args.skip and any(
                    result["task_id"] == dataset_i["task_id"]
                    for result in results
                ):
                    logger.info(f"Skipping already processed task: {dataset_i['task_id']}")
                    continue
            
            try:
                # Log task details
                logger.info(f"{'='*60}")
                logger.info(f"Processing task {i+1}/{len(dataset_slice)}: {dataset_i['task_id']}")
                logger.info(f"Question: {dataset_i['Question']}")
                logger.info(f"Level: {dataset_i['Level']}")
                logger.info(f"Tools: {dataset_i['Annotator Metadata']['Tools']}")
                
                # Prepare question with file paths
                question_data = add_file_path(dataset_i, file_path=gaia_dataset_path, split=args.split)
                question = question_data["Question"]
                
                # Create and execute task
                task = Task(
                    input=question,
                    agent=experience_agent,
                    conf=TaskConfig()
                )
                
                # Execute with experience learning/application
                task_response = await experience_agent.execute_task(task)
                
                # Extract answer
                answer = None
                if task_response and task_response.answer:
                    match = re.search(r"<answer>(.*?)</answer>", task_response.answer)
                    if match:
                        answer = match.group(1)
                
                # Evaluate result
                is_correct = False
                if answer:
                    logger.info(f"Agent answer: {answer}")
                    logger.info(f"Correct answer: {dataset_i['Final answer']}")
                    is_correct = question_scorer(answer, dataset_i["Final answer"])
                    
                    if is_correct:
                        logger.info(f"✓ Question {i} Correct!")
                    else:
                        logger.info(f"✗ Incorrect!")
                else:
                    logger.warning("No answer extracted from response")
                
                # Record result
                new_result = {
                    "task_id": dataset_i["task_id"],
                    "level": dataset_i["Level"],
                    "question": question,
                    "answer": dataset_i["Final answer"],
                    "response": answer or "",
                    "is_correct": is_correct,
                    "learning_mode": args.learning_mode,
                    "applied_experience": args.apply_experience
                }
                
                # Update or append result
                existing_index = next(
                    (idx for idx, result in enumerate(results) 
                     if result["task_id"] == dataset_i["task_id"]),
                    None
                )
                
                if existing_index is not None:
                    results[existing_index] = new_result
                else:
                    results.append(new_result)
                
                # Save intermediate results
                with open(results_file, "w", encoding="utf-8") as f:
                    json.dump(results, f, indent=4, ensure_ascii=False)
                
            except Exception as e:
                logger.error(f"Error processing task {dataset_i['task_id']}: {e}")
                logger.error(traceback.format_exc())
                continue
        
    except KeyboardInterrupt:
        logger.info("Interrupted by user")
    
    finally:
        # Report final results
        report_results(results)
        
        # Save final results
        with open(results_file, "w", encoding="utf-8") as f:
            json.dump(results, f, indent=4, ensure_ascii=False)
        
        # Log experience learning statistics if enabled
        if args.learning_mode:
            num_experiences = len(experience_agent.experiences)
            logger.info(f"\nLearning Statistics:")
            logger.info(f"  - Total experiences learned: {num_experiences}")
            logger.info(f"  - Experience database: {args.experience_db}")
        
        if args.apply_experience and knowledge_base:
            stats = knowledge_base.get_statistics()
            logger.info(f"\nKnowledge Base Statistics:")
            logger.info(f"  - Total indexed documents: {stats['total_documents']}")
            logger.info(f"  - Sources: {stats['sources']}")


async def _evaluate_slice(experience_agent, dataset_slice, gaia_dataset_path, args, pass_label):
    """
    Run one evaluation pass over a dataset slice and return per-task results.

    The agent's ``apply_experience`` flag is read as currently set on the
    ``experience_agent`` instance, which lets the caller toggle experience reuse
    on/off between passes for an apples-to-apples A/B comparison. Correctness is
    computed with GAIA's own ``question_scorer`` — no numbers are fabricated.

    Args:
        experience_agent: The (already constructed) ExperienceAgent.
        dataset_slice: The list of dataset records to evaluate.
        gaia_dataset_path: Root path of the GAIA dataset (for attached files).
        args: Parsed CLI arguments.
        pass_label: Human-readable label for logging (e.g. "baseline").

    Returns:
        A list of per-task result dicts including an ``is_correct`` flag.
    """
    results = []
    for i, dataset_i in enumerate(dataset_slice):
        try:
            logger.info(f"[{pass_label}] Task {i + 1}/{len(dataset_slice)}: {dataset_i['task_id']}")
            question_data = add_file_path(dataset_i, file_path=gaia_dataset_path, split=args.split)
            question = question_data["Question"]

            task = Task(input=question, agent=experience_agent, conf=TaskConfig())
            task_response = await experience_agent.execute_task(task)

            answer = None
            if task_response and task_response.answer:
                match = re.search(r"<answer>(.*?)</answer>", task_response.answer)
                if match:
                    answer = match.group(1)

            is_correct = bool(answer) and question_scorer(answer, dataset_i["Final answer"])
            results.append({
                "task_id": dataset_i["task_id"],
                "level": dataset_i["Level"],
                "question": question,
                "answer": dataset_i["Final answer"],
                "response": answer or "",
                "is_correct": is_correct,
            })
        except Exception as e:
            logger.error(f"[{pass_label}] Error on {dataset_i['task_id']}: {e}")
            results.append({
                "task_id": dataset_i["task_id"],
                "level": dataset_i["Level"],
                "is_correct": False,
                "error": str(e),
            })
    return results


def _accuracy(results):
    """Compute accuracy (correct / total) from a list of result dicts."""
    total = len(results)
    correct = sum(1 for r in results if r.get("is_correct"))
    return correct, total, (correct / total if total else 0.0)


async def run_comparison(args):
    """
    A/B comparison: evaluate the same task slice twice — once WITHOUT experience
    reuse (baseline) and once WITH it — then report the accuracy delta.

    This directly demonstrates the experiment's thesis: reusing accumulated
    experience improves GAIA performance. All reported numbers are computed from
    the actual runs; none are hard-coded.
    """
    _load_runtime()

    gaia_dataset_path = os.getenv("GAIA_DATASET_PATH", "./gaia_dataset")
    full_dataset = load_dataset_meta(gaia_dataset_path, split=args.split)

    mcp_config_path = Path(__file__).parent / "AWorld" / "examples" / "gaia" / "mcp.json"
    mcp_config, available_servers = load_mcp_config(mcp_config_path)

    agent_config = AgentConfig(
        **resolve_llm(model_override=args.model),
        llm_temperature=float(os.getenv("LLM_TEMPERATURE", "0.0")),
    )

    # Build the experience knowledge base for the "with experience" pass.
    knowledge_base = KnowledgeBase(index_path=args.kb_path, model_name=args.embedding_model)
    if args.preload_kb and os.path.exists(args.validation_file):
        logger.warning(
            "--preload-kb indexes gaia-validation.jsonl; if the evaluated tasks are in "
            "that file this leaks their reference solutions. For a fair comparison, "
            "accumulate experiences from OTHER tasks (via --learning-mode) instead."
        )
        knowledge_base.index_gaia_validation(args.validation_file)

    experience_agent = ExperienceAgent(
        conf=agent_config,
        name="gaia_experience_agent",
        system_prompt=system_prompt,
        learning_mode=False,          # keep both passes clean; learn separately
        apply_experience=False,
        experience_db_path=args.experience_db,
        knowledge_base=knowledge_base,
        summarizer=None,
        mcp_config=mcp_config,
        mcp_servers=available_servers,
    )

    num_learned = len(experience_agent.experiences)
    kb_docs = knowledge_base.get_statistics()["total_documents"]
    logger.info(
        f"Comparison ready: {num_learned} learned experiences, {kb_docs} KB documents available."
    )
    if num_learned == 0 and kb_docs == 0:
        logger.warning(
            "No experiences available to reuse. Run --learning-mode first (or pass "
            "--preload-kb) so the 'with experience' pass has something to retrieve."
        )

    dataset_slice = full_dataset[args.start:args.end]

    # Pass 1: baseline (no experience reuse)
    experience_agent.apply_experience = False
    baseline = await _evaluate_slice(experience_agent, dataset_slice, gaia_dataset_path, args, "baseline")

    # Pass 2: with experience reuse
    experience_agent.apply_experience = True
    with_exp = await _evaluate_slice(experience_agent, dataset_slice, gaia_dataset_path, args, "with-experience")

    b_correct, b_total, b_acc = _accuracy(baseline)
    e_correct, e_total, e_acc = _accuracy(with_exp)

    report = {
        "split": args.split,
        "range": [args.start, args.end],
        "num_learned_experiences": num_learned,
        "kb_documents": kb_docs,
        "baseline": {"correct": b_correct, "total": b_total, "accuracy": b_acc, "results": baseline},
        "with_experience": {"correct": e_correct, "total": e_total, "accuracy": e_acc, "results": with_exp},
        "delta_accuracy": e_acc - b_acc,
    }

    out = args.output or os.path.join(os.getenv("AWORLD_WORKSPACE", "."), "comparison_results.json")
    out_dir = os.path.dirname(out)
    if out_dir:
        os.makedirs(out_dir, exist_ok=True)
    with open(out, "w", encoding="utf-8") as f:
        json.dump(report, f, indent=2, ensure_ascii=False)

    # Console summary (computed from real runs)
    print("\n" + "=" * 60)
    print("A/B COMPARISON: experience reuse vs. baseline")
    print("=" * 60)
    print(f"  Tasks evaluated       : {b_total}  (split={args.split}, range=[{args.start}, {args.end}))")
    print(f"  Reusable experiences  : {num_learned} learned, {kb_docs} preloaded")
    print(f"  Baseline accuracy     : {b_correct}/{b_total} = {b_acc:.1%}")
    print(f"  With-experience acc.  : {e_correct}/{e_total} = {e_acc:.1%}")
    print(f"  Delta (with - base)   : {e_acc - b_acc:+.1%}")
    print(f"  Full report written to: {out}")
    print("=" * 60)


def main():
    """Main entry point."""
    # Parse arguments
    args = parse_arguments()

    # Load environment
    load_dotenv()

    # Setup logging
    setup_logging(args)

    # Log configuration
    logger.info("Starting GAIA with Experience Learning")
    logger.info(f"Configuration:")
    logger.info(f"  - Learning mode: {args.learning_mode}")
    logger.info(f"  - Apply experience: {args.apply_experience}")
    logger.info(f"  - Preload KB: {args.preload_kb}")
    logger.info(f"  - Compare mode: {args.compare}")

    # Run async main
    import asyncio
    if args.compare:
        asyncio.run(run_comparison(args))
    else:
        asyncio.run(run_with_experience(args))


if __name__ == "__main__":
    main()