app/run.py

import os
import sys
import json
import time
import threading
import random
from concurrent.futures import ThreadPoolExecutor, as_completed
from loguru import logger
from typing import Any, Dict, List, Tuple
import copy
from flask import Flask, request, jsonify
from pathlib import Path

project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(project_root)

from core.base import Router
from core.path_utils import get_project_root
from core.engine import Engine

# 尝试导入func_timeout，如果没有则提供警告
try:
    from func_timeout import func_timeout, FunctionTimedOut

    HAS_FUNC_TIMEOUT = True
except ImportError:
    HAS_FUNC_TIMEOUT = False
    logger.warning(
        "func_timeout not installed. Timeout control will be disabled. Install with: pip install func-timeout")

from app.evaluation_helper import evaluate_sql_execution, calculate_score_from_evaluation

USE_LLM_EVALUATION = False
if USE_LLM_EVALUATION:
    from evaluation_helper import evaluate_sql_by_llm
    with open("data/ground_truth_res.json", "r", encoding="utf-8") as f:
        ground_truth_res = json.load(f)
    available_instance_ids = list(ground_truth_res.keys())
    print(f"### Available instance ids:\n {available_instance_ids}")
else:
    ground_truth_res = None
    available_instance_ids = None


CONFIG_DIR = "app_config.json"

# 超时常量定义
TASK_MAX_WAIT_TIME = 1200  # 5分钟 = 300秒，任务执行最大等待时间
SQL_MAX_WAIT_TIME = 150  # 60秒，SQL评估最大等待时间

# LLM评估概率（可调节的超参数）
LLM_EVALUATION_PROBABILITY = -1  # 30%的概率使用LLM评估

# Global router & engine
# We assume the data_source and schema_source is pre_defined in config file, and support by our system benchmarks
# This facilitate the dynamic data loading in the running process
router = Router(config_path=CONFIG_DIR)

# The initialization of the engine
engine = Engine(router)
dataloader = engine.dataloader
# We assume the data_source and schema_source is fixed, because all samples are from the same datasets.
data_source = dataloader.get_data_source_index()
schema_source = dataloader.get_schema_source_index()[0]

dataset = dataloader.generate_dataset(data_source, schema_source)
# dataset.db_path = str(get_project_root() / "benchmarks" / "MergeRL" / "database")
dataset.db_path = str(get_project_root() / "benchmarks" / "FinalRL" / "database")

app = Flask(__name__)


def parse_task_from_id(task_id: str):
    if not task_id:
        return None
    if task_id.endswith("Generator"):
        return "GenerateTask"
    elif task_id.endswith("Decomposer"):
        return "DecomposeTask"
    elif task_id.endswith("Optimizer"):
        return "OptimizeTask"
    elif task_id.endswith("Parser"):
        return "ParseTask"
    elif task_id.endswith("Reducer"):
        return "ReduceTask"
    elif task_id.endswith("Scaler"):
        return "ScaleTask"
    elif task_id.endswith("Selector"):
        return "SelectTask"
    else:
        return None


def task_type_mapping(task_type: str):
    if task_type == "GenerateTask":
        return "generate_type"
    elif task_type == "DecomposeTask":
        return "decompose_type"
    elif task_type == "OptimizeTask":
        return "optimize_type"
    elif task_type == "ParseTask":
        return "parse_type"
    elif task_type == "ReduceTask":
        return "reduce_type"
    elif task_type == "ScaleTask":
        return "scale_type"
    elif task_type == "SelectTask":
        return "select_type"
    else:
        return None


def format_task_id(task_id, instance_id):
    return f"{task_id}_{instance_id}"


def init_tasks(task_id: str, instance_id: str, **kwargs):
    # 首先判断任务是否存在，并统一返回 (task, formatted_id)
    formatted_id = format_task_id(task_id, instance_id)
    if task_id in engine.task_ids or formatted_id in engine.task_ids:
        existing = engine.tasks.get(formatted_id) or engine.tasks.get(task_id)
        return existing, formatted_id

    # Task Type
    task_type = parse_task_from_id(task_id)
    flag, task_type = engine.check_task_type(0, task_type)
    if not flag:
        return None, formatted_id

    if dataloader is None:
        return None, formatted_id

    cpy_dataset = copy.deepcopy(dataset)
    sub_data = [row for row in cpy_dataset._dataset if row["instance_id"] == instance_id]
    if not sub_data:
        return None, formatted_id
    cpy_dataset._dataset = sub_data
    # generate task
    generate_args = {
        "task_id": formatted_id,
        "dataset": cpy_dataset,
        "task_type": task_type,
        "open_parallel": True,
        "max_workers": 1,
        "llm_args": kwargs.get("llm", {}),
        "actor_args": kwargs.get("actor", {})
    }
    task_label = task_type_mapping(task_type)
    if task_label:
        generate_args.update({task_label: task_id})

    task = engine.generate_task(**generate_args)

    return task, formatted_id


def parse_task_lis_from_origin(task_lis: List):
    rtn_lis = []
    for item in task_lis:
        if isinstance(item, str):
            rtn_lis.append(item)
        elif isinstance(item, (list, tuple)):
            rtn_lis.extend(list(item))
    return rtn_lis


def init_complex_tasks(task_list: List[str], instance_id: str):
    from datetime import datetime

    # 初始化所有 Meta Tasks
    new_task_lis = parse_task_lis_from_origin(task_list)
    new_task_lis = [init_tasks(id_, instance_id) for id_ in new_task_lis]
    all_tasks = {id_: task for task, id_ in new_task_lis}

    for ind, row in enumerate(task_list):
        if isinstance(row, str):
            task_list[ind] = format_task_id(row, instance_id)
        elif isinstance(row, List):
            for iind, item in enumerate(row):
                row[iind] = format_task_id(item, instance_id)

    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    cpx_task_id = f"{data_source}_{instance_id}_{timestamp}"
    cpx_task_args = {
        "task_id": cpx_task_id,
        "task_lis": task_list,
        "meta_tasks": all_tasks,
        "eval_type": ["execute_accuracy"],
        "is_save_dataset": False,
    }
    cpx_task = engine.generate_complex_task(**cpx_task_args)

    all_tasks.update({
        cpx_task_id: cpx_task
    })
    engine.tasks.update(all_tasks)

    return cpx_task, cpx_task_id


@app.post("/api/run")
def run_complex_actor():
    # It works for single requests, but may have issues in concurrent scenarios.
    payload = request.get_json(force=True, silent=True) or {}
    instance_id = payload.get("instance_id") or request.args.get("instance_id")
    task_lis = payload.get("task_lis")
    logger.info(f"Payload parameter parsing: instance_id {instance_id}, task_lis {task_lis}")

    if not instance_id or not isinstance(task_lis, list):
        return jsonify({"error": "instance_id and task_lis are required"}), 400

    logger.info(f"Begin to initialize complex tasks")
    cpx_task, cpx_task_id = init_complex_tasks(task_lis, instance_id)
    logger.info(f"Finish to initialize complex tasks")

    engine.exec_process = [cpx_task_id]

    started = time.perf_counter()
    try:
        engine.execute()
    finally:
        duration = time.perf_counter() - started
        engine.tasks.pop(cpx_task_id, None)

    eval_res = cpx_task.eval(force=True) or {}
    execute_accuracy = eval_res.get("execute_accuracy", {}).get("avg")

    return jsonify({
        "duration_seconds": duration,
        "execute_accuracy": execute_accuracy
    })


@app.post("/api/run_batch")
def run_batch():
    """
    批量执行任务并返回评分
    
    评分规则:
    1. 任务执行阶段（TASK_MAX_WAIT_TIME=5分钟）:
       - 超时: score = -0.5，跳过后续所有评估
       - 在时间内完成: score += 0.5
    
    2. SQL评估阶段（SQL_MAX_WAIT_TIME=60秒）:
       - 超时或无法评估: score -= 1，跳过后续所有评估
       - SQL可执行（无语法错误）: score += 1，继续下一阶段
       - SQL不可执行（有语法错误）: score -= 1，跳过后续所有评估
    
    3. SQL正确性阶段（仅当SQL可执行时才评估）:
       - 查询结果正确: score += 1.5
       - 查询结果错误: score -= 1.5
    
    4. 时间奖励阶段（仅当SQL正确时才计算）:
       - score += 0.5 * (SQL_MAX_WAIT_TIME - eval_time) / SQL_MAX_WAIT_TIME
       - 评估越快，额外奖励越高（最多+0.5分）
    
    Example of payload:
    {
        "val_0": [[DINSQLGenerator],[DINSQLGenerator],[DINSQLGenerator],[MACSQLGenerator]],
        "val_1": [[DINSQLGenerator],[MACSQLGenerator]],
    }
    return:
    {
        "val_0": [3.0, 3.0, 3.0, 0],
        "val_1": [-0.5, 0.5]
    }
    """
    
    payload = request.get_json(force=True, silent=True) or {}
    print(payload)
    if not payload or not isinstance(payload, dict):
        return jsonify({"error": "instance_id and task_lis are required"}), 400

    if not HAS_FUNC_TIMEOUT:
        logger.warning("func_timeout not available, timeout control disabled")

    def normalize_task_signature(task_lis: List[Any]) -> Tuple[Any, ...]:
        """Convert task_lis into an immutable, comparable structure."""

        def _normalize(item: Any):
            if isinstance(item, str):
                return ("str", item)
            if isinstance(item, (list, tuple)):
                return ("list", tuple(_normalize(sub_item) for sub_item in item))
            try:
                return ("val", json.dumps(item, ensure_ascii=False, sort_keys=True))
            except (TypeError, ValueError):
                return ("val", repr(item))

        return tuple(_normalize(entry) for entry in task_lis)

    # 构建执行计划
    instance_plan: Dict[str, List[Tuple[str, Tuple[Any, ...]]]] = {}
    signature_to_task: Dict[Tuple[str, Tuple[Any, ...]], Dict[str, Any]] = {}
    exec_plan: List[str] = []
    result_cache: Dict[Tuple[str, Tuple[Any, ...]], float] = {}
    result_cache_lock = threading.Lock()

    def cleanup_tasks(restore_callbacks: bool = False):
        """清理创建的任务"""
        for info in signature_to_task.values():
            if restore_callbacks and "original_callback" in info:
                if original_cb := info["original_callback"]:
                    info["task"].call_back = original_cb
            if cpx_id := info.get("cpx_task_id"):
                engine.tasks.pop(cpx_id, None)

    # 解析payload并初始化任务
    for instance_id, task_group in payload.items():
        if not isinstance(task_group, list):
            cleanup_tasks()
            return jsonify({"error": f"task_lis for `{instance_id}` must be a list"}), 400

        instance_plan[instance_id] = []
        for index, task_lis in enumerate(task_group):
            if not isinstance(task_lis, list) or not task_lis:
                cleanup_tasks()
                return jsonify(
                    {"error": f"task_lis at position {index} for `{instance_id}` must be a non-empty list"}
                ), 400
            # if isinstance(task_lis, list) and len(task_lis) > 0:
            #     last_item = task_lis[-1]
            #     if isinstance(last_item, str):
            #         # 如果是字符串且以selector结尾，直接替换
            #         if last_item.lower().endswith("selector"):
            #             task_lis[-1] = "ChaseSelector"
            #     elif isinstance(last_item, list):
            #         if len(last_item) == 1 and isinstance(last_item[0], str) and last_item[0].lower().endswith("selector"):
            #             task_lis[-1][0] = "ChaseSelector"
            signature = (instance_id, normalize_task_signature(task_lis))
            instance_plan[instance_id].append(signature)

            if signature in signature_to_task:
                continue

            # 尝试初始化任务，如果失败则设置分数为 -0.5 并跳过
            try:
                cpx_task, cpx_task_id = init_complex_tasks(copy.deepcopy(task_lis), instance_id)
                if cpx_task is None:
                    logger.warning(f"Failed to initialize task_lis `{index}` for `{instance_id}`: returned None")
                    with result_cache_lock:
                        result_cache[signature] = -0.5
                    engine.tasks.pop(cpx_task_id, None)
                    continue
                
                signature_to_task[signature] = {
                    "task": cpx_task,
                    "cpx_task_id": cpx_task_id,
                    "original_task_lis": task_lis,  # 保存原始task_lis用于LLM评估
                    "instance_id": instance_id,  # 保存instance_id
                }
                exec_plan.append(cpx_task_id)
            except Exception as exc:
                logger.exception(f"Failed to initialize task_lis `{index}` for `{instance_id}`: {exc}")
                with result_cache_lock:
                    result_cache[signature] = -0.5
                continue

    if not exec_plan:
        return jsonify({instance_id: [] for instance_id in instance_plan})

    logger.info(f"run_batch received {len(payload)} instances, executing {len(exec_plan)} unique task groups.")

    # 决定每个instance_id是否使用LLM评估（30%概率）
    use_llm_eval_for_instance: Dict[str, bool] = {}
    if USE_LLM_EVALUATION:
        for instance_id in instance_plan.keys():
            # 只有当instance_id在available_instance_ids中时才可能使用LLM评估
            if available_instance_ids and instance_id in available_instance_ids:
                use_llm_eval = random.random() < LLM_EVALUATION_PROBABILITY
                use_llm_eval_for_instance[instance_id] = use_llm_eval
                if use_llm_eval:
                    logger.info(f"Instance {instance_id} will use LLM evaluation")
            else:
                use_llm_eval_for_instance[instance_id] = False
    else:
        for instance_id in instance_plan.keys():
            use_llm_eval_for_instance[instance_id] = False

    # 记录每个任务的执行状态
    task_execution_status: Dict[Tuple[str, Tuple[Any, ...]], Dict[str, Any]] = {}

    def execute_single_task(cpx_task):
        """执行单个任务"""
        cpx_task.run()
        cpx_task.end()

    def evaluate_sql(cpx_task):
        """评估SQL并返回结果"""
        if not cpx_task.dataset or len(cpx_task.dataset) == 0:
            return None, "No dataset available for evaluation"

        row = cpx_task.dataset[0]
        eval_result = evaluate_sql_execution(
            pred_sql=row.get("pred_sql", ""),
            gold_sql=row.get("query", ""),
            db_id=row.get("db_id", ""),
            db_type=row.get("db_type", "sqlite"),
            db_path=dataset.db_path,
            db_credential=dataset.credential
        )
        return eval_result, None

    # 并行执行任务并评分（每个任务独立超时控制）
    def process_single_signature(signature, info):
        """处理单个signature的任务执行和评估"""
        cpx_task = info["task"]
        cpx_task_id = info["cpx_task_id"]
        instance_id = info["instance_id"]
        original_task_lis = info["original_task_lis"]
        score = 0.0

        # 检查是否使用LLM评估
        if use_llm_eval_for_instance.get(instance_id, False):
            logger.info(f"Using LLM evaluation for {cpx_task_id} (instance: {instance_id})")
            try:
                # 调用LLM评估
                llm_success, llm_score = evaluate_sql_by_llm(instance_id, original_task_lis)
                
                if llm_success:
                    # LLM评估成功，直接使用返回的score
                    logger.info(f"LLM evaluation succeeded for {cpx_task_id}, score: {llm_score}")
                    with result_cache_lock:
                        result_cache[signature] = llm_score
                    return
                else:
                    # LLM评估返回False，继续执行原来的流程
                    logger.info(f"LLM evaluation returned False for {cpx_task_id}, falling back to standard evaluation")
            except Exception as exc:
                logger.exception(f"LLM evaluation error for {cpx_task_id}: {exc}, falling back to standard evaluation")
                # 发生异常，继续执行原来的流程

        # 标准评估流程
        # 1. 执行任务（带超时控制）
        task_timeout = False
        task_start_time = time.time()
        try:
            if HAS_FUNC_TIMEOUT:
                try:
                    func_timeout(1500, execute_single_task, args=(cpx_task,))
                except FunctionTimedOut:
                    task_timeout = True
                    logger.warning(f"Task {cpx_task_id} exceeded {TASK_MAX_WAIT_TIME}s timeout")
            else:
                execute_single_task(cpx_task)
        except Exception as exc:
            logger.exception(f"Task {cpx_task_id} execution failed: {exc}")
            task_timeout = True  # 执行失败也当作超时处理
        task_time = time.time() - task_start_time
        # 2. 任务超时，分数为-0.5，跳过评估
        if task_timeout:
            with result_cache_lock:
                result_cache[signature] = -0.5
            return

        # 3. 任务完成，基础分 +0.5
        score = 0.5

        # 4. SQL评估（带超时控制）
        eval_timeout = False
        eval_result = None
        eval_error = None
        eval_start_time = time.time()

        try:
            if HAS_FUNC_TIMEOUT:
                try:
                    eval_result, eval_error = func_timeout(SQL_MAX_WAIT_TIME, evaluate_sql, args=(cpx_task,))
                except FunctionTimedOut:
                    eval_timeout = True
                    logger.warning(f"SQL evaluation for task {cpx_task_id} exceeded {SQL_MAX_WAIT_TIME}s timeout")
            else:
                eval_result, eval_error = evaluate_sql(cpx_task)
        except Exception as exc:
            eval_error = str(exc)
            logger.exception(f"SQL evaluation error for task {cpx_task_id}: {exc}")
        
        eval_time = time.time() - eval_start_time

        # 5. 评估超时或失败，score -= 1
        if eval_timeout or eval_error or eval_result is None:
            with result_cache_lock:
                result_cache[signature] = score - 1
            return

        # 6. SQL不可执行，score -= 1
        if not eval_result.can_execute:
            with result_cache_lock:
                result_cache[signature] = score - 1
            return

        # 7. SQL可执行，score += 1
        score += 1

        # 8. 正确性评估
        if eval_result.is_correct:
            score += 1.5
            # 额外的时间奖励：评估越快，奖励越高
            score += 0.5 * max((TASK_MAX_WAIT_TIME - task_time) / TASK_MAX_WAIT_TIME, 0)
        else:
            score -= 1.5
        with result_cache_lock:
            result_cache[signature] = score

    # 使用线程池并行执行所有任务
    max_workers = min(len(signature_to_task), os.cpu_count() or 4)
    logger.info(f"Starting parallel execution with {max_workers} workers for {len(signature_to_task)} tasks")
    
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        # 提交所有任务
        futures = {
            executor.submit(process_single_signature, signature, info): signature
            for signature, info in signature_to_task.items()
        }
        
        # 等待所有任务完成
        for future in as_completed(futures):
            signature = futures[future]
            try:
                future.result()  # 获取结果（如果有异常会在这里抛出）
            except Exception as exc:
                logger.exception(f"Unexpected error processing signature {signature}: {exc}")
                # 如果发生意外错误，确保该signature有个默认分数
                with result_cache_lock:
                    if signature not in result_cache:
                        result_cache[signature] = -0.5

    # 清理任务
    try:
        cleanup_tasks()
    except Exception as exc:
        logger.warning(f"Error during cleanup: {exc}")
    
    res = {
        instance_id: [result_cache.get(sig, -0.5) for sig in signatures]
        for instance_id, signatures in instance_plan.items()
    }
    print(res)
    # 构建最终响应
    return jsonify(res)


@app.get("/healthz")
def healthz():
    return jsonify({"status": "ok"}), 200


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=int(os.environ.get("PORT", "6517")), debug=False)