evaluate.py

import argparse
import glob
import json
import os
import traceback
from collections import OrderedDict, defaultdict
from concurrent.futures import ThreadPoolExecutor, TimeoutError, as_completed
from statistics import mean
from typing import Any, Dict, List, Set, Tuple, Optional
from string import punctuation
import litellm
import tiktoken
from tqdm import tqdm
from utils.constant import DIFFICULTY_LEVELS
from utils.process_utils import (check_entity_in_text, check_vllm_server,
                                 extract_entity_from_url, mask_wiki_content,
                                 get_parametric_sequence, mask_with_underscore)
from utils.prompt import JUDGE_PROMPT_GAIA, JUDGE_PROMPT_QA
from utils.react_agent import MultiTurnReactAgent

judge_prompt: str = ""
dataset: str = ""
SUMMARY_MODEL_PATH = os.getenv("SUMMARY_MODEL_PATH", "/cpfs01/mj/eval_models/Qwen2.5-72B-Instruct")
SUMMARY_MODEL_ADDRESS = os.getenv("SUMMARY_MODEL_ADDRESS", "http://0.0.0.0:6002/v1")
SUMMARY_MODEL_APIKEY = os.getenv("SUMMARY_MODEL_APIKEY", "EMPTY")
MODEL_PATH = os.getenv("MODEL_PATH", None)
AGENT = MultiTurnReactAgent(
    llm={'model': MODEL_PATH, 'generate_cfg':{'num_retries':10, 'temperature':0.3, 'top_p':0.95}},
    function_list=[],
    system_message=""
)
# litellm._turn_on_debug()

def call_llm_judge(item: Dict[str, Any], num_retries: int = 5) -> Dict: 
    """Judge if predicted answer matches ground-truth""" 
    global judge_prompt, dataset
    try: 
        question = item["question"]
        correct_answer = item["answer"]
        prediction = item["prediction"].strip()
        messages = [
            {"role": "user", "content": judge_prompt.format(question=question, correct_answer=correct_answer, response=prediction)}
        ]
        
        if "wiki" in dataset:
            response = litellm.completion(
                base_url=SUMMARY_MODEL_ADDRESS,
                api_key=SUMMARY_MODEL_APIKEY,
                model="hosted_vllm/" + SUMMARY_MODEL_PATH if check_vllm_server(SUMMARY_MODEL_ADDRESS) else SUMMARY_MODEL_PATH,
                messages=messages,  # type: ignore
                temperature=0.3,
                max_tokens=10,
                num_retries=num_retries,
                timeout=60,
            )
            judgement = response.choices[0].message["content"].strip()
            if judgement:
                # For wiki, prompt requires "A" or "a" as the correct answer
                if judgement in ["A", 'a', "A.", "a."]:
                    judgement = "Correct"
                else:
                    judgement = "Incorrect"
        else:
            raise ValueError(f"Unsupported dataset: {dataset}")
        return {
            "question": question, 
            "answer": correct_answer, 
            "judgement": judgement
        }
    
    except Exception as e:
        print(f"Error judgement for question: {question}: {e}")
        return {
            "question": question,  
            "answer": correct_answer, 
            "judgement": "Error",
            "error": str(e)
        }


def check_refusal(response: str, num_retries: int = 5) -> bool:
    """Check if the response is a refusal"""
    refusal_prompt = """You are tasked with determining whether the given response constitutes a refusal to answer the question.

A refusal includes:
- Explicit statements of inability to conclude an answer ("I don't know", "I cannot answer", "I'm not sure") from its searched information

Non-refusal includes:
- Attempting to provide a specific answer (even if incorrect)
- Providing partial information while trying to answer
- Making educated guesses or hypotheses

Response: {response}

Is this a refusal? Answer only "yes" or "no"."""

    messages = [{"role": "user", "content": refusal_prompt.format(response=response)}]
    try:
        response_obj = litellm.completion(
            base_url=SUMMARY_MODEL_ADDRESS,
            api_key=SUMMARY_MODEL_APIKEY,
            model="hosted_vllm/" + SUMMARY_MODEL_PATH if check_vllm_server(SUMMARY_MODEL_ADDRESS) else SUMMARY_MODEL_PATH,
            messages=messages,  # type: ignore
            temperature=0.3,
            max_tokens=10,
            num_retries=num_retries,
            timeout=60
        )
        result = response_obj.choices[0].message.content
        if result and result.strip().lower() in ["yes", "no"]:
            return result.strip().lower() == "yes"
    except Exception as e:
        print(f"Error in refusal check: {str(e)}")
    
    return False


def check_model_answerable(question: str, ground_truth: str, evidence_list: Optional[List[str]] = None, num_retries: int = 10) -> bool:
    """Check if the model can answer the question"""
    if evidence_list:
        prompt = '''Instruction: Write an accurate, engaging, and concise answer for the given question. Use an unbiased and journalistic tone.
Question: {question}

Evidence:
{evidence_list}

Answer (be concise and specific):
'''.strip()
        prompt = prompt.format(question=question, evidence_list="\n".join(evidence_list))
    else:
        prompt = '''Instruction: Write an accurate, engaging, and concise answer for the given question. Use an unbiased and journalistic tone.
Question: {question}

Answer (be concise and specific):
'''.strip()
        prompt = prompt.format(question=question)
    
    
    messages, reasoning_content, tool_calls, prediction = AGENT._call_llm(messages=[{"role": "user", "content": prompt}], num_retries=num_retries)
    
    if reasoning_content:
        if check_entity_in_text(ground_truth, reasoning_content):
            return "yes"
        else:
            return "no"
    else:
        return "Invalid Judgement"


def classify_difficulty(optimal_hop_count_item: int) -> str:
    """Classify question difficulty based on optimal hop count"""
    # 2 - 6
    if optimal_hop_count_item < 3:
        return DIFFICULTY_LEVELS[0]  # "Easy"
    elif optimal_hop_count_item < 5:
        return DIFFICULTY_LEVELS[1]  # "Medium"
    else:
        return DIFFICULTY_LEVELS[2]  # "Hard"


def aggregate_difficulty_metrics(round_results_dict: Dict[str, List[Dict[str, Any]]]) -> Dict[str, Dict[str, float]]:
    """Aggregate evaluation results by difficulty level"""
    if not round_results_dict:
        return {}

    difficulty_round_results = OrderedDict(
        (difficulty, defaultdict(list)) for difficulty in sorted(DIFFICULTY_LEVELS)
    )
    for round_name, round_results in round_results_dict.items():
        for result in round_results:
            difficulty = classify_difficulty(result["optimal_hop_count_item"])
            difficulty_round_results[difficulty][round_name].append(result)

    # Calculate metrics for each difficulty level (average over rounds)
    difficulty_metrics = {}
    for difficulty, rounds_results in difficulty_round_results.items():
        round_metrics = []
        for results in rounds_results.values():
            if not results:
                pass
            else:
                good_refusal_recall = sum([r["good_refusal_item"] for r in results]) / sum([not r["knowledge_sufficient_item"] for r in results]) * 100 if sum([not r["knowledge_sufficient_item"] for r in results]) > 0 else 0.0
                good_refusal_precision = sum([r["good_refusal_item"] for r in results]) / sum([r["is_refusal_item"] for r in results]) * 100 if sum([r["is_refusal_item"] for r in results]) > 0 else 0.0
                good_refusal_f1 = 2 * good_refusal_recall * good_refusal_precision / (good_refusal_recall + good_refusal_precision) if (good_refusal_recall + good_refusal_precision) > 0 else 0
                knowledge_utilisation_recall = sum([r["knowledge_utilisation_item"] for r in results]) / sum([r["knowledge_sufficient_item"] for r in results]) * 100 if sum([r["knowledge_sufficient_item"] for r in results]) > 0 else 0.0
                knowledge_utilisation_precision = sum([r["knowledge_utilisation_item"] for r in results]) / sum([not r["is_refusal_item"] for r in results]) * 100 if sum([not r["is_refusal_item"] for r in results]) > 0 else 0.0
                knowledge_utilisation_f1 = 2 * knowledge_utilisation_recall * knowledge_utilisation_precision / (knowledge_utilisation_recall + knowledge_utilisation_precision) if (knowledge_utilisation_recall + knowledge_utilisation_precision) > 0 else 0
                gen_score = 0.5 * (knowledge_utilisation_f1 + good_refusal_f1) * (sum([r["knowledge_sufficient_item"] for r in results]) / len(results))
                knowledge_forget_to_use = sum([r["knowledge_forget_to_use_item"] for r in results]) / sum([r["knowledge_sufficient_item"] for r in results]) * 100 if sum([r["knowledge_sufficient_item"] for r in results]) > 0 else 0.0
                knowledge_leadastray = sum([r["knowledge_leadastray_item"] for r in results]) / sum([r["knowledge_sufficient_item"] for r in results]) * 100 if sum([r["knowledge_sufficient_item"] for r in results]) > 0 else 0.0

                round_metrics.append({
                    "count": len(results),
                    "pass_at_1": mean([r["answer_correct_item"] for r in results]) * 100,
                    "search_score": mean([r["search_score_item"] for r in results]) * 100,
                    "gen_score": gen_score,
                    "good_refusal_recall": good_refusal_recall,
                    "good_refusal_precision": good_refusal_precision,
                    "good_refusal_f1": good_refusal_f1,
                    "knowledge_utilisation_recall": knowledge_utilisation_recall,
                    "knowledge_utilisation_precision": knowledge_utilisation_precision,
                    "knowledge_utilisation_f1": knowledge_utilisation_f1,
                    "knowledge_sufficient": mean([r["knowledge_sufficient_item"] for r in results]) * 100,
                    "rag_knowledge_sufficient": mean([r["rag_knowledge_sufficient_item"] for r in results]) * 100,
                    "knowledge_forget_to_use": knowledge_forget_to_use,
                    "knowledge_leadastray": knowledge_leadastray,
                    "is_refusal": mean([r["is_refusal_item"] for r in results]) * 100,
                    "avg_tool_use": mean([r["tool_use_count_item"] for r in results]),
                    "avg_visit_tool_use": mean([len(r["visited_urls_item"]) for r in results]),
                    "avg_optimal_hop": mean([r["optimal_hop_count_item"] for r in results]),
                })
        # Now, average across rounds for this difficulty
        if round_metrics:
            count = int(round(mean([m["count"] for m in round_metrics if m["count"] > 0]), 0)) if any(m["count"] > 0 for m in round_metrics) else 0
            difficulty_metrics[difficulty] = {
                "count": count,
                "pass_at_1": round(mean([m["pass_at_1"] for m in round_metrics]) , 2),
                "search_score": round(mean([m["search_score"] for m in round_metrics]), 2),
                "gen_score": round(mean([m["gen_score"] for m in round_metrics]), 2),
                "good_refusal_recall": round(mean([m["good_refusal_recall"] for m in round_metrics]), 2),
                "good_refusal_precision": round(mean([m["good_refusal_precision"] for m in round_metrics]), 2),
                "good_refusal_f1": round(mean([m["good_refusal_f1"] for m in round_metrics]), 2),
                "knowledge_utilisation_recall": round(mean([m["knowledge_utilisation_recall"] for m in round_metrics]), 2),
                "knowledge_utilisation_precision": round(mean([m["knowledge_utilisation_precision"] for m in round_metrics]), 2),
                "knowledge_utilisation_f1": round(mean([m["knowledge_utilisation_f1"] for m in round_metrics]), 2),
                "knowledge_sufficient": round(mean([m["knowledge_sufficient"] for m in round_metrics]), 2),
                "rag_knowledge_sufficient": round(mean([m["rag_knowledge_sufficient"] for m in round_metrics]), 2),
                "knowledge_forget_to_use": round(mean([m["knowledge_forget_to_use"] for m in round_metrics]), 2),
                "knowledge_leadastray": round(mean([m["knowledge_leadastray"] for m in round_metrics]), 2),
                "is_refusal": round(mean([m["is_refusal"] for m in round_metrics]), 2),
                "avg_tool_use": round(mean([m["avg_tool_use"] for m in round_metrics]), 2),
                "avg_visit_tool_use": round(mean([m["avg_visit_tool_use"] for m in round_metrics]), 2),
                "avg_optimal_hop": round(mean([m["avg_optimal_hop"] for m in round_metrics]), 2),
            }
        else:
            difficulty_metrics[difficulty] = {
                "count": 0,
                "pass_at_1": 0.0,
                "search_score": 0.0,
                "gen_score": 0.0,
                "good_refusal_recall": 0.0,
                "good_refusal_precision": 0.0,
                "good_refusal_f1": 0.0,
                "knowledge_utilisation_recall": 0.0,
                "knowledge_utilisation_precision": 0.0,
                "knowledge_utilisation_f1": 0.0,
                "knowledge_sufficient": 0.0,
                "rag_knowledge_sufficient": 0.0,
                "knowledge_forget_to_use": 0.0,
                "knowledge_leadastray": 0.0,
                "is_refusal": 0.0,
                "avg_tool_use": 0.0,
                "avg_visit_tool_use": 0.0,
                "avg_optimal_hop": 0.0
            }
    
    return difficulty_metrics


def extract_visited_static_info(item: Dict) -> Tuple[List[str], int]:
    """Extract entities visited by the model from tool calls"""
    visited_urls: Set[str] = set()
    messages: List[Dict[str, Any]] = item.get("messages", [])
    tool_use_count: int = 0
    
    for msg in messages:
        if msg.get("role") == "assistant" and msg.get("tool_calls"):
            try:
                for tool_call in msg["tool_calls"]:
                    tool_use_count += 1
                    tool_name = tool_call.get("function", {}).get("name", "")
                    if "visit" in tool_name:
                        arguments = tool_call.get("function", {}).get("arguments", {})
                        urls = json.loads(arguments).get("url", [])
                        if isinstance(urls, str):
                            urls = [urls]
                        for url in urls:
                            visited_urls.add(url)
            except Exception as e:
                print(e)

    return list(visited_urls), tool_use_count


def check_path_knowledge_sufficiency(item: Dict, visited_urls: List[str]) -> bool:
    """Check if search + parametric knowledge is sufficient to support the optimal path"""

    # Check if any of the gaps can be filled by parametric knowledge
    QUERY_TEMPLATE = """Fill in the blank in the following passage:

Passage:
{evidence}
"""

    entity_title_list = item.get("entity_title_list", [])
    evidence_list = item.get("evidence_list", [])
    
    if not entity_title_list or len(entity_title_list) < 2:
        raise ValueError(f"entity_title_list is empty or too short: {entity_title_list}")
    if len(entity_title_list) != len(evidence_list) + 1:
        raise ValueError(f"entity_title_list and evidence_list must have matched lengths: {entity_title_list} and {evidence_list}")
    
    # normalize URLs
    visited_entity_list = [extract_entity_from_url(url) for url in visited_urls]
    
    if not any(visited_entity in entity_title_list for visited_entity in visited_entity_list):
        evidence_text = "\n".join(evidence_list)
        formatted_evidence = "_____".join(evidence_text.split(item['answer'].strip(punctuation).strip()))
        query = QUERY_TEMPLATE.format(evidence=formatted_evidence)
        if check_model_answerable(query, entity_title_list[-1]) == "yes":
            return True
        else:
            return False
    
    gaps_to_check = get_parametric_sequence(entity_title_list, visited_entity_list)
    if not gaps_to_check:
        return True
    
    entity_evidx_dict = {entity: idx for idx, entity in enumerate(entity_title_list)}
   
    def is_gap_path_list_answerable(gap_path_list: List[List[str]]) -> bool:
        for gap_path_pair in gap_path_list:
            start_idx = entity_evidx_dict[gap_path_pair[0]]
            end_idx = entity_evidx_dict[gap_path_pair[1]]
            evidence_segment = evidence_list[start_idx:end_idx]
            masked_evidence_segment = mask_wiki_content(evidence_segment[-1], [gap_path_pair[1]], [entity_title_list[end_idx - 1]])
            masked_evidence_segment_with_underscore = mask_with_underscore(evidence_segment[-1], masked_evidence_segment)
            formatted_evidence = "\n".join(evidence_segment[:-1] + [masked_evidence_segment_with_underscore])
            query = QUERY_TEMPLATE.format(evidence=formatted_evidence)
            if check_model_answerable(query, gap_path_pair[1]) != "yes":
                return False
        return True

    for gap_path_list in gaps_to_check:
        if is_gap_path_list_answerable(gap_path_list):
            return True
    return False


def check_rag_knowledge_sufficiency(item: Dict, visited_urls: List[str]) -> bool:
    """
    Detect whether there is any situation where knowledge has not been utilized.
    """
    entity_title_list = item.get("entity_title_list", [])
    evidence_list = item.get("evidence_list", [])
    
    if not entity_title_list or len(entity_title_list) < 2:
        raise ValueError(f"entity_title_list is empty or too short: {entity_title_list}")
    if len(entity_title_list) != len(evidence_list) + 1:
        raise ValueError(f"entity_title_list and evidence_list must have matched lengths: {entity_title_list} and {evidence_list}")

    # normalize URLs
    visited_entity_list = [extract_entity_from_url(url) for url in visited_urls]

    common_evidence_list = [evidence for i, evidence in enumerate(evidence_list) if entity_title_list[i] in visited_entity_list]
    
    check_result = check_model_answerable(item['question'], item['answer'], common_evidence_list) == "yes"
    # check_result = False
    
    return check_result


def holistic_evaluate_item(item: Dict) -> Dict:
    """Perform holistic evaluation for a single item"""
    # Get basic judgment
    basic_judgment_item= call_llm_judge(item)
    answer_correct_item = basic_judgment_item["judgement"] == "Correct"
    
    # Extract visited entities
    visited_urls_item, tool_use_count_item = extract_visited_static_info(item)
    
    # Check path knowledge sufficiency
    print(f"Debug: Checking path knowledge sufficiency...")
    knowledge_sufficient_item = check_path_knowledge_sufficiency(item, visited_urls_item)
    print(f"Debug: Knowledge sufficient: {knowledge_sufficient_item}")

    # Check rag knowledge sufficiency
    print(f"Debug: Checking rag knowledge sufficiency...")
    rag_knowledge_sufficient_item = check_rag_knowledge_sufficiency(item, visited_urls_item)
    print(f"Debug: Rag knowledge sufficient: {rag_knowledge_sufficient_item}")

    # Get optimal nodes count (length of entity_url_list)
    optimal_hop_count_item = len(item.get("evidence_list", []))
    
    # Apply holistic scoring logic
    search_score_item = 0
    good_refusal_item = 0
    knowledge_utilisation_item = 0
    knowledge_forget_to_use_item = 0
    knowledge_leadastray_item = 0
    
    # search score
    if tool_use_count_item > 0:
        if knowledge_sufficient_item:
            search_score_item = 1
        else:
            if answer_correct_item and len(visited_urls_item) < optimal_hop_count_item:
                print(item)
                search_score_item = 1

    # refusal score
    is_refusal = check_refusal(item['prediction'].strip())
    if is_refusal and not knowledge_sufficient_item:
        good_refusal_item = 1
    
    # knowledge utilisation score
    if knowledge_sufficient_item and answer_correct_item:
        knowledge_utilisation_item = 1

    # forget to use knowledge
    if knowledge_sufficient_item and not rag_knowledge_sufficient_item:
        knowledge_forget_to_use_item = 1

    # lead astray
    if knowledge_sufficient_item and rag_knowledge_sufficient_item and not answer_correct_item:
        knowledge_leadastray_item = 1
    
    return {
        "question": item.get("question", ""),
        "answer": item.get("answer", ""),
        "entity_urls_item": item.get("entity_urls_item", []),
        "visited_urls_item": visited_urls_item,
        "tool_use_count_item": tool_use_count_item,
        "optimal_hop_count_item": optimal_hop_count_item,
        "answer_correct_item": answer_correct_item,
        "knowledge_sufficient_item": knowledge_sufficient_item,
        "rag_knowledge_sufficient_item": rag_knowledge_sufficient_item,
        "is_refusal_item": is_refusal,
        "search_score_item": search_score_item,
        "good_refusal_item": good_refusal_item,
        "knowledge_utilisation_item": knowledge_utilisation_item,
        "knowledge_forget_to_use_item": knowledge_forget_to_use_item,
        "knowledge_leadastray_item": knowledge_leadastray_item,
        "judgement": basic_judgment_item["judgement"],
        "final_response": item.get("prediction", "")
    }


def aggregate_holistic_results(round_results_dict: Dict[str, List[Dict[str, Any]]]) -> Tuple[Dict[str, float], Dict[str, Dict[str, float]]]:
    """Aggregate holistic evaluation results from multiple rounds"""
    if not round_results_dict:
        return {}
    
    pass_at_k = calculate_pass_at_k(round_results_dict, k=len(round_results_dict.keys()))
    best_pass_at_1 = calculate_best_pass_at_1(round_results_dict)
    avg_pass_at_1 = calculate_avg_pass_at_1(round_results_dict) 
    
    # Calculate metrics for each round first, then average across rounds
    round_metrics = {}
    
    for round_name, round_results in round_results_dict.items():
        if not round_results:
            continue
        # Calculate metrics for this round
        search_score_round = round(mean(result["search_score_item"] for result in round_results) * 100, 2)
        good_refusal_recall_round = round(
            sum(result["good_refusal_item"] for result in round_results) / sum([not result["knowledge_sufficient_item"] for result in round_results]) * 100, 2
        ) if sum([not result["knowledge_sufficient_item"] for result in round_results]) > 0 else 0.0
        good_refusal_precision_round = round(
            sum(result["good_refusal_item"] for result in round_results) / sum([result["is_refusal_item"] for result in round_results]) * 100, 2
        ) if sum([result["is_refusal_item"] for result in round_results]) > 0 else 0.0
        good_refusal_f1_round = round(2 * good_refusal_recall_round * good_refusal_precision_round / (good_refusal_recall_round + good_refusal_precision_round) if (good_refusal_recall_round + good_refusal_precision_round) > 0 else 0, 2)
        knowledge_utilisation_recall_round = round(
            sum(result["knowledge_utilisation_item"] for result in round_results) / sum([result["knowledge_sufficient_item"] for result in round_results]) * 100, 2
        ) if sum([result["knowledge_sufficient_item"] for result in round_results]) > 0 else 0.0
        knowledge_utilisation_precision_round = round(
            sum(result["knowledge_utilisation_item"] for result in round_results) / sum([not result["is_refusal_item"] for result in round_results]) * 100, 2
        ) if sum([not result["is_refusal_item"] for result in round_results]) > 0 else 0.0
        knowledge_utilisation_f1_round = round(2 * knowledge_utilisation_recall_round * knowledge_utilisation_precision_round / (knowledge_utilisation_recall_round + knowledge_utilisation_precision_round) if (knowledge_utilisation_recall_round + knowledge_utilisation_precision_round) > 0 else 0, 2)
        gen_score_round = round(0.5 * (knowledge_utilisation_f1_round + good_refusal_f1_round) * (sum([result["knowledge_sufficient_item"] for result in round_results])/len(round_results)), 2)
        knowledge_forget_to_use_round = round(
            sum(result["knowledge_forget_to_use_item"] for result in round_results) / sum([result["knowledge_sufficient_item"] for result in round_results]) * 100, 2
        ) if sum([result["knowledge_sufficient_item"] for result in round_results]) > 0 else 0.0
        knowledge_leadastray_round = round(
            sum(result["knowledge_leadastray_item"] for result in round_results) / sum([result["knowledge_sufficient_item"] for result in round_results]) * 100, 2
        ) if sum([result["knowledge_sufficient_item"] for result in round_results]) > 0 else 0.0

        pass_at_1_round = round(mean(result["answer_correct_item"] for result in round_results) * 100, 2)
        knowledge_sufficient_round = round(mean(result["knowledge_sufficient_item"] for result in round_results) * 100, 2)
        rag_knowledge_sufficient_round = round(mean(result["rag_knowledge_sufficient_item"] for result in round_results) * 100, 2)
        is_refusal_round = round(mean(result["is_refusal_item"] for result in round_results) * 100, 2)
        avg_visit_tool_use_round = round(mean(len(result["visited_urls_item"]) for result in round_results), 2)
        avg_tool_use_round = round(mean(result["tool_use_count_item"] for result in round_results), 2)
        avg_optimal_hop_round = round(mean(result["optimal_hop_count_item"] for result in round_results), 2)
        
        round_metrics[round_name] = {
            "search_score_round": search_score_round,
            "gen_score_round": gen_score_round,
            "good_refusal_recall_round": good_refusal_recall_round,
            "good_refusal_precision_round": good_refusal_precision_round,
            "good_refusal_f1_round": good_refusal_f1_round,
            "knowledge_utilisation_recall_round": knowledge_utilisation_recall_round,
            "knowledge_utilisation_precision_round": knowledge_utilisation_precision_round,
            "knowledge_utilisation_f1_round": knowledge_utilisation_f1_round,
            "knowledge_sufficient_round": knowledge_sufficient_round,
            "rag_knowledge_sufficient_round": rag_knowledge_sufficient_round,
            "knowledge_forget_to_use_round": knowledge_forget_to_use_round,
            "knowledge_leadastray_round": knowledge_leadastray_round,
            "is_refusal_round": is_refusal_round,
            "pass_at_1_round": pass_at_1_round,
            "visit_tool_count_round": avg_visit_tool_use_round,
            "tool_use_count_round": avg_tool_use_round,
            "optimal_hop_count_round": avg_optimal_hop_round
        }
    
    # Average across all rounds
    return {
        "pass_at_k": pass_at_k,
        "avg_pass_at_1": avg_pass_at_1,
        "best_pass_at_1": best_pass_at_1,
        "avg_search_score": round(mean(metrics["search_score_round"] for metrics in round_metrics.values()), 2),
        "best_search_score": round(max(metrics["search_score_round"] for metrics in round_metrics.values()), 2),
        "avg_gen_score": round(mean(metrics["gen_score_round"] for metrics in round_metrics.values()), 2),
        "best_gen_score": round(max(metrics["gen_score_round"] for metrics in round_metrics.values()), 2),
        "avg_good_refusal_recall": round(mean(metrics["good_refusal_recall_round"] for metrics in round_metrics.values()), 2),
        "best_good_refusal_recall": round(max(metrics["good_refusal_recall_round"] for metrics in round_metrics.values()), 2),
        "avg_good_refusal_precision": round(mean(metrics["good_refusal_precision_round"] for metrics in round_metrics.values()), 2),
        "best_good_refusal_precision": round(max(metrics["good_refusal_precision_round"] for metrics in round_metrics.values()), 2),
        "avg_good_refusal_f1": round(mean(metrics["good_refusal_f1_round"] for metrics in round_metrics.values()), 2),
        "best_good_refusal_f1": round(max(metrics["good_refusal_f1_round"] for metrics in round_metrics.values()), 2),
        "avg_knowledge_utilisation_recall": round(mean(metrics["knowledge_utilisation_recall_round"] for metrics in round_metrics.values()), 2),
        "best_knowledge_utilisation_recall": round(max(metrics["knowledge_utilisation_recall_round"] for metrics in round_metrics.values()), 2),
        "avg_knowledge_utilisation_precision": round(mean(metrics["knowledge_utilisation_precision_round"] for metrics in round_metrics.values()), 2),
        "best_knowledge_utilisation_precision": round(max(metrics["knowledge_utilisation_precision_round"] for metrics in round_metrics.values()), 2),
        "avg_knowledge_utilisation_f1": round(mean(metrics["knowledge_utilisation_f1_round"] for metrics in round_metrics.values()), 2),
        "best_knowledge_utilisation_f1": round(max(metrics["knowledge_utilisation_f1_round"] for metrics in round_metrics.values()), 2),
        "avg_knowledge_sufficient": round(mean(metrics["knowledge_sufficient_round"] for metrics in round_metrics.values()), 2),
        "best_knowledge_sufficient": round(max(metrics["knowledge_sufficient_round"] for metrics in round_metrics.values()), 2),
        "avg_rag_knowledge_sufficient": round(mean(metrics["rag_knowledge_sufficient_round"] for metrics in round_metrics.values()), 2),
        "best_rag_knowledge_sufficient": round(max(metrics["rag_knowledge_sufficient_round"] for metrics in round_metrics.values()), 2),
        "avg_knowledge_forget_to_use": round(mean(metrics["knowledge_forget_to_use_round"] for metrics in round_metrics.values()), 2),
        "best_knowledge_forget_to_use": round(max(metrics["knowledge_forget_to_use_round"] for metrics in round_metrics.values()), 2),
        "avg_knowledge_leadastray": round(mean(metrics["knowledge_leadastray_round"] for metrics in round_metrics.values()), 2),
        "best_knowledge_leadastray": round(max(metrics["knowledge_leadastray_round"] for metrics in round_metrics.values()), 2),
        "avg_is_refusal": round(mean(metrics["is_refusal_round"] for metrics in round_metrics.values()), 2),
        "best_is_refusal": round(max(metrics["is_refusal_round"] for metrics in round_metrics.values()), 2),
        "avg_visit_tool_count": round(mean(metrics["visit_tool_count_round"] for metrics in round_metrics.values()), 2),
        "avg_tool_use_count": round(mean(metrics["tool_use_count_round"] for metrics in round_metrics.values()), 2),
        "avg_optimal_hop_count": round(mean(metrics["optimal_hop_count_round"] for metrics in round_metrics.values()), 2),
    }, round_metrics


def process_single_round(input_file: str) -> List[Dict]: 
    items = []
    with open(input_file, 'r', encoding='utf-8') as f: 
        for line in f:
            try:
                items.append(json.loads(line))
            except Exception as e:
                print(f"Error loading {line}")
    return items 


def aggregate_statistics(round_files: List[str]) -> Dict[str, float]:
    all_stats = []
    for round_file in round_files:
        stats = single_round_statistics(round_file)
        all_stats.append(stats)

    if not all_stats:
        return {}

    keys = all_stats[0].keys()
    avg_stats = {}
    for key in keys:
        avg_stats[key] = round(mean([stats[key] for stats in all_stats]), 3)

    return avg_stats


def single_round_statistics(input_file: str) -> Dict[str, float]:
    contents = process_single_round(input_file) 

    # Illegal Analysis 
    num_invalid, num_extra = 0, 0  
    # Tool Analysis 
    tool_use_cnt, visit_tool_cnt, search_tool_cnt, other_tool_cnt = [], [], [], [] 
    # Thinking Analysis 
    all_ans_lengths, all_think_lengths = [], []  

    tokenizer = tiktoken.get_encoding("o200k_base")
    
    for item in contents:
        texts = item["messages"]
        final_msg = texts[-1]["content"] if len(texts) else ""
        
        # Analyze answer 
        if "<answer>" not in final_msg or "</answer>" not in final_msg: 
            num_invalid += 1 
            answer_length = 0 
        else:
            answer_length = len(final_msg.split("<answer>")[1].split("</answer>")[0].strip())
        
        # Analyze tool use & thinking
        num_tool_use, num_visit_tool, num_search_tool, num_other_tool = 0, 0, 0, 0 
        think_lengths = []
        for text in texts:
            if text.get("role") == "assistant":
                reasoning_content = text["content"]
                think_lengths.append(len(reasoning_content)) 

                if text.get("tool_calls"):
                    try:
                        for tool_call in text.get("tool_calls", []):
                            tool_name = tool_call.get("function", {}).get("name", "")
                            num_tool_use += 1
                            if "visit" in tool_name:  
                                num_visit_tool += 1 
                            elif "search" in tool_name: 
                                num_search_tool += 1  
                            else:
                                num_other_tool += 1
                    except Exception as e:
                        print(e)

        tool_use_cnt.append(num_tool_use) 
        visit_tool_cnt.append(num_visit_tool)  
        search_tool_cnt.append(num_search_tool)   
        other_tool_cnt.append(num_other_tool) 

        all_ans_lengths.append(answer_length) 
        think_length = sum(think_lengths) / len(think_lengths) if think_lengths else 0  
        all_think_lengths.append(think_length) 

        # Overlength 
        if len(tokenizer.encode("".join([text["content"] for text in texts]))) > 30000:
            num_extra += 1  
    
    return {
        "extra_length": num_extra, # number of overlength responses  
        "num_invalid": num_invalid, # number of invalid responses 
        "avg_action": sum(tool_use_cnt) / len(tool_use_cnt), # avg. number of tool invocation 
        "avg_visit_action": sum(visit_tool_cnt) / len(visit_tool_cnt), # avg. number of visit tool invocation 
        "avg_search_action": sum(search_tool_cnt) / len(search_tool_cnt), # avg. number of search tool invocation 
        "avg_other_action": sum(other_tool_cnt) / len(other_tool_cnt), # avg. number of other tool invocation 
        "avg_ans_length": sum(all_ans_lengths) / len(all_ans_lengths), 
        "avg_think_length": sum(all_think_lengths) / len(all_think_lengths)
    }
        

def calculate_pass_at_k(round_results_dict: Dict[str, List[Dict[str, Any]]], k: int = 3, mode: str = "simple") -> float: 
    total_correct: int = 0
    total_items: int = len(list(round_results_dict.values())[0])

    for idx in range(total_items):
        for round_name in round_results_dict.keys():
            judgement = round_results_dict[round_name][idx]["judgement"]
            if judgement == "Correct":
                total_correct += 1
                break
    
    overall_pass = total_correct / total_items
    return round(overall_pass * 100, 2)


def calculate_best_pass_at_1(round_results_dict: Dict[str, List[Dict[str, Any]]], mode: str = "simple") -> float:  
    total_correct = {round_name: 0 for round_name in round_results_dict.keys()}

    for round_name in round_results_dict.keys():
        for item in round_results_dict[round_name]:
            if item["judgement"] == "Correct":  
                total_correct[round_name] += 1 

    overall_best = max(
        total_correct[r] / len(round_results_dict[r])
        for r in round_results_dict.keys()
    )

    return round(overall_best * 100, 2)


def calculate_avg_pass_at_1(round_results_dict: Dict[str, List[Dict[str, Any]]], mode: str = "simple") -> float: 
    total_correct = {round_name: 0 for round_name in round_results_dict.keys()}

    for round_name in round_results_dict.keys():  
        for item in round_results_dict[round_name]:
            if item["judgement"] == "Correct":
                total_correct[round_name] += 1 
    
    avg_overall = mean([
        total_correct[r] / len(round_results_dict[r]) 
        for r in round_results_dict.keys()
    ])
    
    return round(avg_overall * 100, 2)


def main():
    global judge_prompt, dataset
    parser = argparse.ArgumentParser(description="Evaluate model predictions across multiple rounds")
    parser.add_argument("--input_folder", help="Path to prediction files")
    parser.add_argument("--restore_result_path",default='summary.jsonl', help="record result")
    parser.add_argument("--dataset", type=str, default="wiki_100.jsonl")
    parser.add_argument("--max_workers", type=int, default=10)
    parser.add_argument("--evaluation_mode", type=str, choices=["simple", "holistic"], default="simple", 
                        help="Evaluation mode: simple (original) or holistic (new comprehensive evaluation)")
    parser.add_argument("--timeout", type=int, default=300, 
                        help="Timeout for each evaluation task in seconds (default: 120)")
    args = parser.parse_args()

    dataset = args.dataset
    if "wiki" in dataset: 
        judge_prompt = JUDGE_PROMPT_QA
    else:
        judge_prompt = JUDGE_PROMPT_GAIA 
    print(f"Using {dataset} judge prompt ...")

    dataset_dir = args.input_folder.split('.jsonl')[0]
    # Find all iter*.jsonl files
    round_files = sorted(glob.glob(os.path.join(dataset_dir, "iter*.jsonl")))
    if not round_files:
        raise FileNotFoundError(f"No prediction files found in {dataset_dir} (expecting iter*.jsonl)")

    round_items = {}
    for idx, file in enumerate(round_files, 1):
        print(f"Processing {file} ...")
        round_items[f"round{idx}"] = process_single_round(file)

    round_results_dict = {} 

    if args.evaluation_mode == "simple":
        # Original simple evaluation
        with ThreadPoolExecutor(max_workers=args.max_workers) as executor:  
            for round_name, items in round_items.items():
                # Add original index to each item for maintaining order
                indexed_items = [(idx, item) for idx, item in enumerate(items)]
                future_to_task = {executor.submit(call_llm_judge, item): (idx, item) for idx, item in indexed_items} 
                temp_results = {} 

                for future in tqdm(as_completed(future_to_task), total=len(future_to_task), desc=f"Evaluating {round_name}"):
                    idx, task_info = future_to_task[future]
                    try:
                        result = future.result(timeout=args.timeout)
                        result["_original_idx"] = idx  # Store original index
                        temp_results[idx] = result
                    except TimeoutError:
                        print(f'Timeout (>{round(args.timeout/60)} mins): "{task_info["question"]}" '
                          f'(Rollout {task_info["rollout_id"]})')
                        future.cancel()
                        error_result = {
                            "question": task_info.get("question", ""), 
                            "answer": task_info.get("answer", ""), 
                            "judgement": "Error",
                            "error": f"Timeout (>{round(args.timeout/60)} mins)",
                            "_original_idx": idx
                        }
                        temp_results[idx] = error_result
                    except Exception as e:
                        print(f"Error or timeout occurred for item in {round_name}: {e}")
                        error_result = {
                            "question": task_info.get("question", ""), 
                            "answer": task_info.get("answer", ""), 
                            "judgement": "Error",
                            "error": str(e),
                            "_original_idx": idx
                        }
                        temp_results[idx] = error_result
                
                # Sort results by original index to maintain order
                round_results_dict[round_name] = [temp_results[i] for i in sorted(temp_results.keys())]

        pass_at_k = calculate_pass_at_k(round_results_dict, k=len(round_results_dict.keys()))
        best_pass_at_1 = calculate_best_pass_at_1(round_results_dict)
        avg_pass_at_1 = calculate_avg_pass_at_1(round_results_dict) 
        
        aggr_statistics = aggregate_statistics(round_files)

        round_performance = {
            f"Round{i}_Pass@1": round(
                mean([1 if r["judgement"] == "Correct" else 0 for r in round_results_dict.get(f"round{i}", [])]) * 100, 2
            )
            for i in range(1, len(round_results_dict.keys()) + 1)
        }

        print(f"=========== SIMPLE EVALUATION ===========")
        print(f"Avg. Pass@1 {avg_pass_at_1}%") 
        print(f"Best Pass@1 {best_pass_at_1}%")  
        print(f"Pass@{len(round_results_dict.keys())} {pass_at_k}%") 
        print(f"Pass@1 Round 1: {round_performance['Round1_Pass@1']}%  Round 2: {round_performance['Round2_Pass@1']}%  Round 3: {round_performance['Round3_Pass@1']}% \n")
        print(f"# Invalid {aggr_statistics['num_invalid']}  # Extra Length {aggr_statistics['extra_length']}") 
        print(f"Avg. Action {aggr_statistics['avg_action']:.2f}  Avg. Visit Action {aggr_statistics['avg_visit_action']:.2f}  Avg. Search Action {aggr_statistics['avg_search_action']:.2f}  Avg. Other Action {aggr_statistics['avg_other_action']:.2f}") 
        print(f"Avg. Answer Length {aggr_statistics['avg_ans_length']:.2f}  Avg. Thinking Length {aggr_statistics['avg_think_length']:.2f}")  
        print(f"===========================================")

        overall_eval_dict = {
            "dataset": dataset, 
            "evaluation_mode": "simple",
            "files": {
                round_name: round_files[i] for i, round_name in enumerate(round_results_dict.keys())
            }, 
            "overall": {
                "avg_pass_at_1": avg_pass_at_1, 
                "best_pass_at_1": best_pass_at_1, 
                "pass_at_k": pass_at_k
            }, 
            "individual": round_performance, 
            "statistics": aggr_statistics
        }
        
    else:  # holistic evaluation
        with ThreadPoolExecutor(max_workers=args.max_workers) as executor:  
            for round_name, items in round_items.items():
                # Add original index to each item for maintaining order
                indexed_items = [(idx, item) for idx, item in enumerate(items)]
                future_to_task = {executor.submit(holistic_evaluate_item, item): (idx, item) for idx, item in indexed_items} 
                temp_results = {} 

                for future in tqdm(as_completed(future_to_task), total=len(future_to_task), desc=f"Holistic Evaluating {round_name}"):
                    idx, task_info = future_to_task[future]
                    try:
                        result = future.result(timeout=args.timeout)
                        result["_original_idx"] = idx  # Store original index
                        temp_results[idx] = result
                    except TimeoutError:
                        print(f'Timeout (>{round(args.timeout/60)} mins): "{task_info["question"]}" '
                          f'(Rollout {task_info["rollout_id"]})')
                        future.cancel()
                        error_result = {
                            "question": task_info.get("question", ""),
                            "answer": task_info.get("answer", ""),
                            "entity_urls_item": task_info.get("entity_urls_item", []),
                            "visited_urls_item": [],
                            "tool_use_count_item": 0,
                            "optimal_hop_count_item": len(task_info.get("evidence_list", [])),
                            "answer_correct_item": False,
                            "knowledge_sufficient_item": False,
                            "rag_knowledge_sufficient_item": False,
                            "is_refusal_item": False,
                            "search_score_item": 0,
                            "good_refusal_item": 0,
                            "knowledge_utilisation_item": 0,
                            "knowledge_forget_to_use_item": 0,
                            "knowledge_leadastray_item": 0,
                            "judgement": "Error",
                            "final_response": task_info.get("prediction", ""),
                            "error": f"Timeout (>{round(args.timeout/60)} mins)",
                            "_original_idx": idx
                        }
                        temp_results[idx] = error_result
                    except Exception as e:
                        print(f"Error or timeout occurred for item in {round_name}: {e}")
                        error_result = {
                            "question": task_info.get("question", ""),
                            "answer": task_info.get("answer", ""),
                            "entity_urls_item": task_info.get("entity_urls_item", []),
                            "visited_urls_item": [],
                            "tool_use_count_item": 0,
                            "optimal_hop_count_item": len(task_info.get("evidence_list", [])),
                            "answer_correct_item": False,
                            "knowledge_sufficient_item": False,
                            "rag_knowledge_sufficient_item": False,
                            "is_refusal_item": False,
                            "search_score_item": 0,
                            "good_refusal_item": 0,
                            "knowledge_utilisation_item": 0,
                            "knowledge_forget_to_use_item": 0,
                            "knowledge_leadastray_item": 0,
                            "judgement": "Error",
                            "final_response": task_info.get("prediction", ""),
                            "error": str(e),
                            "_original_idx": idx
                        }
                        temp_results[idx] = error_result
                
                # Sort results by original index to maintain order
                round_results_dict[round_name] = [temp_results[i] for i in sorted(temp_results.keys())]

        # Aggregate holistic results
        holistic_metrics, round_metrics = aggregate_holistic_results(round_results_dict)
        # Aggregate difficulty-based metrics
        difficulty_metrics = aggregate_difficulty_metrics(round_results_dict)
        # Aggreagate statistics of process
        aggr_statistics = aggregate_statistics(round_files)
        
        # Calculate per-round metrics for holistic evaluation
        round_performance = {}
        for i, round_name in enumerate(round_metrics.keys(), 1):
            if round_metrics[round_name]:
                round_performance[f"Round{i}_Search_Score"] = round_metrics[round_name]["search_score_round"]
                round_performance[f"Round{i}_Gen_Score"] = round_metrics[round_name]["gen_score_round"]
                round_performance[f"Round{i}_Good_Refusal_Recall"] = round_metrics[round_name]["good_refusal_recall_round"]
                round_performance[f"Round{i}_Good_Refusal_Precision"] = round_metrics[round_name]["good_refusal_precision_round"]
                round_performance[f"Round{i}_Good_Refusal_F1"] = round_metrics[round_name]["good_refusal_f1_round"]
                round_performance[f"Round{i}_Knowledge_Utilisation_Recall"] = round_metrics[round_name]["knowledge_utilisation_recall_round"]
                round_performance[f"Round{i}_Knowledge_Utilisation_Precision"] = round_metrics[round_name]["knowledge_utilisation_precision_round"]
                round_performance[f"Round{i}_Knowledge_Utilisation_F1"] = round_metrics[round_name]["knowledge_utilisation_f1_round"]
                round_performance[f"Round{i}_Knowledge_Sufficiency"] = round_metrics[round_name]["knowledge_sufficient_round"]
                round_performance[f"Round{i}_Rag_Knowledge_Sufficiency"] = round_metrics[round_name]["rag_knowledge_sufficient_round"]
                round_performance[f"Round{i}_Knowledge_Forget_To_Use"] = round_metrics[round_name]["knowledge_forget_to_use_round"]
                round_performance[f"Round{i}_Knowledge_Leadastray"] = round_metrics[round_name]["knowledge_leadastray_round"]
                round_performance[f"Round{i}_Is_Refusal"] = round_metrics[round_name]["is_refusal_round"]
                round_performance[f"Round{i}_Pass@1"] = round_metrics[round_name]["pass_at_1_round"]
                
        print(f"=========== HOLISTIC EVALUATION ===========")
        print(f"Avg. Search Score@1: {holistic_metrics['avg_search_score']}%")
        print(f"Best Search Score@1: {holistic_metrics['best_search_score']}%")
        for idx in range(len(round_results_dict.keys())):
            print(f"Round{idx+1}_Search_Score: {round_performance[f'Round{idx+1}_Search_Score']}%")
        
        print(f"Avg. Generation Score@1: {holistic_metrics['avg_gen_score']}%") 
        print(f"Best Generation Score@1: {holistic_metrics['best_gen_score']}%")
        for idx in range(len(round_results_dict.keys())):
            print(f"Round{idx+1}_Gen_Score: {round_performance[f'Round{idx+1}_Gen_Score']}%")

        print(f"Avg. Good Refusal F1: {holistic_metrics['avg_good_refusal_f1']}%")
        print(f"Best Good Refusal F1: {holistic_metrics['best_good_refusal_f1']}%")
        for idx in range(len(round_results_dict.keys())):
            print(f"Round{idx+1}_Good_Refusal_F1: {round_performance[f'Round{idx+1}_Good_Refusal_F1']}%")

        print(f"Avg. Knowledge Utilisation F1: {holistic_metrics['avg_knowledge_utilisation_f1']}%")
        print(f"Best Knowledge Utilisation F1: {holistic_metrics['best_knowledge_utilisation_f1']}%")
        for idx in range(len(round_results_dict.keys())):
            print(f"Round{idx+1}_Knowledge_Utilisation_F1: {round_performance[f'Round{idx+1}_Knowledge_Utilisation_F1']}%")

        print(f"Avg. Pass@1: {holistic_metrics['avg_pass_at_1']}%")
        print(f"Best Pass@1: {holistic_metrics['best_pass_at_1']}%")
        print(f"Pass@{len(round_results_dict.keys())}: {holistic_metrics['pass_at_k']}%")
        for idx in range(len(round_results_dict.keys())):
            print(f"Round{idx+1}_Pass@1: {round_performance[f'Round{idx+1}_Pass@1']}%")
        
        print("============== EXTENDED METRICS ==============")
        print(f"Avg. Knowledge Sufficiency@1: {holistic_metrics['avg_knowledge_sufficient']}%")
        print(f"Avg. Rag Knowledge Sufficiency@1: {holistic_metrics['avg_rag_knowledge_sufficient']}%")
        print(f"Avg. Knowledge Forget To Use@1: {holistic_metrics['avg_knowledge_forget_to_use']}%")
        print(f"Avg. Knowledge Leadastray@1: {holistic_metrics['avg_knowledge_leadastray']}%")
        print(f"Avg. Is Refusal: {holistic_metrics['avg_is_refusal']}%")
        print(f"Avg. Tool Use@1: {holistic_metrics['avg_tool_use_count']:.2f}")
        print(f"Avg. Visit Tool Use@1: {holistic_metrics['avg_visit_tool_count']:.2f}")
        print(f"Avg. Optimal Hop Count: {holistic_metrics['avg_optimal_hop_count']:.2f}") 
        print("========== DIFFICULTY BREAKDOWN ==========")
        for difficulty in DIFFICULTY_LEVELS:
            metrics = difficulty_metrics[difficulty]
            print(f"{difficulty}: {metrics['count']} questions")
            print(f"  Pass@1: {metrics['pass_at_1']}%  Search Score: {metrics['search_score']}%  Gen Score: {metrics['gen_score']}%  Good Refusal F1: {metrics['good_refusal_f1']}%  Knowledge Utilisation F1: {metrics['knowledge_utilisation_f1']}%")
            print(f"  Knowledge Sufficiency: {metrics['knowledge_sufficient']}%  Rag Knowledge Sufficiency: {metrics['rag_knowledge_sufficient']}%  Knowledge Forget To Use: {metrics['knowledge_forget_to_use']}%  Knowledge Leadastray: {metrics['knowledge_leadastray']}%  Is Refusal: {metrics['is_refusal']}%  Avg Tool Use: {metrics['avg_tool_use']:.2f}  Avg. Visit Tool Use: {metrics['avg_visit_tool_use']:.2f}  Avg. Optimal Hop Count: {metrics['avg_optimal_hop']:.2f}")
        print("================ STATISTICS ================")
        print(f"# Invalid {aggr_statistics['num_invalid']}  # Extra Length {aggr_statistics['extra_length']}") 
        print(f"Avg. Action {aggr_statistics['avg_action']:.2f}  Avg. Visit Action {aggr_statistics['avg_visit_action']:.2f}  Avg. Search Action {aggr_statistics['avg_search_action']:.2f}  Avg. Other Action {aggr_statistics['avg_other_action']:.2f}") 
        print(f"Avg. Answer Length {aggr_statistics['avg_ans_length']:.2f}  Avg. Thinking Length {aggr_statistics['avg_think_length']:.2f}")  
        print("===========================================")

        overall_eval_dict = {
            "dataset": dataset,
            "evaluation_mode": "holistic", 
            "files": {
                round_name: round_files[i] for i, round_name in enumerate(round_results_dict.keys())
            },
            "individual": round_performance,
            "statistics": aggr_statistics,
            "holistic_metrics": holistic_metrics,
            "difficulty_metrics": difficulty_metrics,
        }

    round_results_save_path = os.path.join(dataset_dir, "eval_results.json")
    with open(round_results_save_path, 'w', encoding='utf-8') as f:
        json.dump(round_results_dict, f, ensure_ascii=False, indent=2)

    with open(args.restore_result_path, 'a', encoding='utf-8') as jsonl_file:
        jsonl_file.write(json.dumps(overall_eval_dict, ensure_ascii=False) + '\n')


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        error_str = traceback.format_exc()
        print(f"Evaluation Failed: {str(e)}") 
        print("Trace Back", error_str)

    # # 测试 call_llm_judge
    # judge_prompt = JUDGE_PROMPT_QA
    # dataset = "gaia"
    # test_item = {
    #     "question": "What is the capital of France?",
    #     "answer": "Paris",
    #     "prediction": "The capital of France is Paris."
    # }
    # result = call_llm_judge(test_item)
    # print("call_llm_judge 测试结果：", result)

    # print(check_refusal('xxxxxxxxxxxxxxxxxx. It seems like there is no information about the capital of France. But i may find more'))
    # test_item = {
    #     "question": "What is Jeff Rivera's profession?", 
    #     "answer": "Novelist.", 
    #     "entity_title_list": ["Jeff Rivera", "Self-publishing", "Novelist"],
    #     "entity_url_list": ["https://en.wikipedia.org/wiki/Jeff Rivera", "https://en.wikipedia.org/wiki/Self-publishing", "https://en.wikipedia.org/wiki/Novelist"], 
    #     "evidence_list": ["Before it was picked up by Warner Books, 8,000 copies were in circulation in both print and electronic versions that Rivera had self-published.A story of second chances; Jeff Rivera's novel moves from cyberspace to print, John Diaz, The Boston Globe, July 10, 2007; accessed November 22, 2008.\n\nHarriet Klausner called it \"a terrific character driven redemption tale\".Review of Forever My Lady, Harriet Klausner.\n\nRecognition\nthumb|Rivera and Amy Cheney, at an awards ceremony honoring Cheney in The New York Times Building.", "Self-publishing has long been an option for writers, with vanity presses printing bound books for a fee paid by the writer."],
    #     "rollout_id": 2, 
    #     "messages": [{"role": "system", "content": "You are a Web Information Seeking Master. Your ultimate job is to exhaustively gather accurate, up-to-date information from the open Web and craft a reliable answer. No matter how complex the query, you keep digging until properly satisfied.\n\nAs you proceed, adhere to the following principles:\n\n1. **Persistent, Multi-Step Exploration**: You will engage in multiple interactions, delving both deeply and laterally into the topic to explore broader dimensions until a satisfactory answer emerges.\n\n2. **Attention to Detail & Provenance**: You will carefully analyze each information source to ensure that all data is current, relevant, and from reputable origins; flag anything that seems outdated or dubious.\n\n3. **Repeated Verification**: Before presenting a Final Answer, you will **cross-check** and **validate the information** you've gathered to confirm its accuracy and reliability.\n\n4. **Synthesis Before Answering**: Begin crafting the final answer only after you are satisfied that you have explored the topic exhaustively—both broadly and in depth—and then synthesize your findings into a coherent response for the user.\nCurrent date: 2025-07-19"}, {"role": "user", "content": "A conversation between User and Assistant. The user asks a question, and the assistant solves it by calling one or more of the following tools.\n<tools>\n{\n  \"name\": \"search\",\n  \"description\": \"Performs batched web searches: supply an array 'query'; the tool retrieves the top 10 related URLs and their snippets for each query in one call.\",\n  \"parameters\": {\n    \"type\": \"object\",\n    \"properties\": {\n      \"query\": {\n        \"type\": \"array\",\n        \"items\": {\n          \"type\": \"string\"\n        },\n        \"description\": \"Array of query strings. Include multiple complementary search queries in a single call.\"\n      }\n    },\n    \"required\": [\n      \"query\"\n    ]\n    }\n},\n{\n  \"name\": \"visit\",\n    \"description\": \"Visit webpage(s) for more details than just snippets and return the summary of the intended content.\",\n    \"parameters\": {\n        \"type\": \"object\",\n        \"properties\": {\n            \"url\": {\n                \"type\": \"array\",\n                \"items\": {\"type\": \"string\"},\n                \"description\": \"The URL(s) of the webpage(s) to visit. Can be a single URL or an array of URLs.\"\n            },\n            \"goal\": {\n                \"type\": \"string\",\n                \"description\": \"The specific information goal for visiting webpage(s).\"\n            }\n        },\n        \"required\": [\n            \"url\",\n            \"goal\"\n        ]\n    }\n}\n</tools>\n\nThe assistant starts with one or multiple cycles (at assistant's discretion) of (thinking about which tool to use -> performing tool call -> waiting for tool response), and ends with (thinking about the answer -> answer of the question). The **thinking processes**, **tool calls**, **tool responses**, and **answer** must be strictly enclosed within their corresponding tags.\n\nExample response:\n<think> thinking process here </think>\n<tool_call>\n{\"name\": \"tool name here\", \"arguments\": {\"parameter name here\": parameter value here, \"another parameter name here\": another parameter value here, ...}}\n</tool_call>\n<tool_response>\ntool_response here\n</tool_response>\n<think> thinking process here </think>\n<tool_call>\n{\"name\": \"another tool name here\", \"arguments\": {...}}\n</tool_call>\n<tool_response>\ntool_response here\n</tool_response>\n(more thinking processes, tool calls and tool responses here)\n<think> thinking process here </think>\n<answer> answer here </answer>\n\nUser: What is Jeff Rivera's profession?"}, {"role": "assistant", "content": "<think>To determine Jeff Rivera's profession, I need a reliable and up-to-date source that explicitly states his occupation. Using search queries like 'Jeff Rivera profession', 'Jeff Rivera job', and 'Jeff Rivera career' increases the likelihood of finding biographical summaries, interviews, or news articles that mention his current or primary career. These variations also help cover different ways the information might be phrased online, ensuring I don't miss relevant results due to wording differences. I need to run these searches now to gather authoritative sources that can clearly answer the question.\n</think>\n<tool_call>\n{\"name\": \"search\", \"arguments\": {\"query\": [\"Jeff Rivera profession\", \"Jeff Rivera job\", \"Jeff Rivera career\"]}}\n</tool_call>"}, {"role": "user", "content": "<tool_response>\nA Wiki search for 'Jeff Rivera profession' found 10 related URLs, please visit any URL if needed for more details:\n\n## Web Results\n1. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nHis most recent work, Forever My Lady, was released by Warner Books in July 2007. He is also the author of Oh Yes I Can! (2003). \n\n2. [List of Utah writers](https://en.wikipedia.org/wiki/List of Utah writers)\nJeff Rivera (born 1976) – contemporary fiction, young adult\n\n3. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nCategory:20th-century African-American writers\n\n4. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\n\n\n5. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nPublishers Weekly review of Forever My Lady\n\n6. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nMahogany Media Award for Best New Author and Best Urban Fiction (2007)\n\n7. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nHarriet Klausner called it \"a terrific character driven redemption tale\". \n\n8. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nCategory:21st-century American male writers\n\n9. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nCategory:American writers of young adult literature\n\n10. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nCategory:Writers from Salt Lake City\n=======\nA Wiki search for 'Jeff Rivera career' found 10 related URLs, please visit any URL if needed for more details:\n\n## Web Results\n1. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nHis most recent work, Forever My Lady, was released by Warner Books in July 2007. He is also the author of Oh Yes I Can! (2003). \n\n2. [List of Utah writers](https://en.wikipedia.org/wiki/List of Utah writers)\nJeff Rivera (born 1976) – contemporary fiction, young adult\n\n3. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nCategory:20th-century African-American writers\n\n4. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\n\n\n5. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nPublishers Weekly review of Forever My Lady\n\n6. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nMahogany Media Award for Best New Author and Best Urban Fiction (2007)\n\n7. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nHarriet Klausner called it \"a terrific character driven redemption tale\". \n\n8. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nCategory:21st-century American male writers\n\n9. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nCategory:American writers of young adult literature\n\n10. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nCategory:Writers from Salt Lake City\n=======\nA Wiki search for 'Jeff Rivera job' found 10 related URLs, please visit any URL if needed for more details:\n\n## Web Results\n1. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nHis most recent work, Forever My Lady, was released by Warner Books in July 2007. He is also the author of Oh Yes I Can! (2003). \n\n2. [List of Utah writers](https://en.wikipedia.org/wiki/List of Utah writers)\nJeff Rivera (born 1976) – contemporary fiction, young adult\n\n3. [She Loves Me, She Loves Me Not](https://en.wikipedia.org/wiki/She Loves Me, She Loves Me Not)\nJeff Job – engineering\n\n4. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nCategory:20th-century African-American writers\n\n5. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\n\n\n6. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nPublishers Weekly review of Forever My Lady\n\n7. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nMahogany Media Award for Best New Author and Best Urban Fiction (2007)\n\n8. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nHarriet Klausner called it \"a terrific character driven redemption tale\". \n\n9. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nCategory:21st-century American male writers\n\n10. [Jeff Rivera](https://en.wikipedia.org/wiki/Jeff Rivera)\nCategory:American writers of young adult literature\n</tool_response>"}, {"role": "assistant", "content": "<think>The Wikipedia page for Jeff Rivera appears repeatedly across multiple search queries and is likely to contain a concise summary of his professional background, including his primary occupation. Since the search snippets mention his authorship and literary achievements but don't explicitly state his profession in full, I need to directly examine the Wikipedia article to confirm and extract the most authoritative description of his career. This should provide a clear, reliable answer regarding his profession. I need to visit this page now to identify and extract the specific statement about Jeff Rivera's profession.\n</think>\n<tool_call>\n{\"name\": \"visit\", \"arguments\": {\"url\": [\"https://en.wikipedia.org/wiki/Jeff_Rivera\"], \"goal\": \"Extract Jeff Rivera's profession from the Wikipedia page\"}}\n</tool_call>"}, {"role": "user", "content": "<tool_response>\nThe useful information in \"https://en.wikipedia.org/wiki/Jeff_Rivera\" for user goal \"Extract Jeff Rivera's profession from the Wikipedia page\" as follows:\n\nEvidence in page:\nHis most recent work, Forever My Lady, was released by Warner Books in July 2007. He is also the author of Oh Yes I Can! (2003). Forever My Lady is a story about a Latino juvenile delinquent who transforms his life in prison boot camp in order to win back the love of the girl he lost. Rivera was inspired to write the story when, after a period living in his car, he met a former gang member who turned his life around. Harriet Klausner called it 'a terrific character driven redemption tale'. Review of Forever My Lady, Harriet Klausner. Recognition thumb|Rivera and Amy Cheney, at an awards ceremony honoring Cheney in The New York Times Building. Mahogany Media Award for Best New Author and Best Urban Fiction (2007)\n\nSummary:\nJeff Rivera is an author. He has written several books, including 'Forever My Lady' (2007) and 'Oh Yes I Can!' (2003). His work 'Forever My Lady' received positive reviews and he was recognized with the Mahogany Media Award for Best New Author and Best Urban Fiction in 2007.\n</tool_response>"}, {"role": "assistant", "content": "<think>To arrive at the answer, I first conducted a series of web searches using variations like \"Jeff Rivera profession,\" \"Jeff Rivera job,\" and \"Jeff Rivera career.\" The search results consistently pointed to Wikipedia and other reputable sources, all of which highlighted his authorship of books such as \"Forever My Lady\" and \"Oh Yes I Can!\" To ensure accuracy, I visited the Wikipedia page, which explicitly described him as an author and provided details about his literary works and the recognition he received. Given the consistency across multiple sources and the direct statement on his Wikipedia page, I can confidently conclude that Jeff Rivera's profession is author.</think>\n<answer>Jeff Rivera is an author. He is known for writing contemporary fiction and young adult literature. Among his published works are the books \"Forever My Lady\" (2007) and \"Oh Yes I Can!\" (2003). His novel \"Forever My Lady\" focuses on the story of a Latino juvenile delinquent who transforms his life in prison boot camp in order to win back the love of the girl he lost. The book was well received by critics and received positive reviews, including praise from Harriet Klausner, who described it as \"a terrific character-driven redemption tale.\"\n\nIn recognition of his contributions to literature, Jeff Rivera was awarded the Mahogany Media Award for Best New Author and Best Urban Fiction in 2007. His writing draws on personal experiences and real-life stories, such as the inspiration he found from meeting a former gang member who had turned his life around after a period living in his car.\n\nOverall, Jeff Rivera's profession is that of an author, with a focus on contemporary and young adult fiction, and he is recognized for his ability to deliver compelling character-driven narratives.</answer>"}], 
    #     "prediction": "Jeff Rivera is an author. He is known for writing contemporary fiction and young adult literature. Among his published works are the books \"Forever My Lady\" (2007) and \"Oh Yes I Can!\" (2003). His novel \"Forever My Lady\" focuses on the story of a Latino juvenile delinquent who transforms his life in prison boot camp in order to win back the love of the girl he lost. The book was well received by critics and received positive reviews, including praise from Harriet Klausner, who described it as \"a terrific character-driven redemption tale.\"\n\nIn recognition of his contributions to literature, Jeff Rivera was awarded the Mahogany Media Award for Best New Author and Best Urban Fiction in 2007. His writing draws on personal experiences and real-life stories, such as the inspiration he found from meeting a former gang member who had turned his life around after a period living in his car.\n\nOverall, Jeff Rivera's profession is that of an author, with a focus on contemporary and young adult fiction, and he is recognized for his ability to deliver compelling character-driven narratives.", 
    #     "termination": "answer"
    # }