alpha_tools/simple72/main.py

import shutil
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
from fastapi.templating import Jinja2Templates
import sys
import os
import json
import uuid
import zipfile
from datetime import datetime
import logging
import openai
from Tranformer.ace_lib import SingleSession, get_instrument_type_region_delay, get_datafields, get_datasets
from Tranformer.Transformer import run_transformer
import asyncio
import pandas as pd

# 配置日志
logger = logging.getLogger("alpha_transformer")
if not logger.handlers:
    handler = logging.StreamHandler(sys.stdout)
    handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
    logger.addHandler(handler)
    logger.setLevel(logging.INFO)

# 创建 FastAPI 应用实例
app = FastAPI(title="Alpha Transformer", version="1.0.0")

# Chrome DevTools 健康检查端点（可选，阻止 404 日志）
@app.get("/.well-known/appspecific/com.chrome.devtools.json")
async def chrome_devtools_check():
    """Chrome DevTools 健康检查"""
    return {"status": "ok"}

# 提供模板目录中的 CSS 和 JS 文件
@app.get("/styles.css")
async def get_styles():
    """提供 styles.css 文件"""
    from fastapi.responses import FileResponse
    return FileResponse("templates/styles.css")

@app.get("/app.js")
async def get_app_js():
    """提供 app.js 文件"""
    from fastapi.responses import FileResponse
    return FileResponse("templates/app.js")

# 配置 Jinja2 模板引擎，用于渲染 HTML 页面
templates = Jinja2Templates(directory="templates")

# 程序启动前清空 Tranformer/output 目录
if os.path.exists("Tranformer/output"):
    for filename in os.listdir("Tranformer/output"):
        file_path = os.path.join("Tranformer/output", filename)
        try:
            if os.path.isfile(file_path) or os.path.islink(file_path):
                os.unlink(file_path)
            elif os.path.isdir(file_path):
                shutil.rmtree(file_path)
        except Exception as e:
            logger.error(f"⚠ 清空目录 {file_path} 时出错: {e}")



# 存储正在运行的任务信息（task_id -> 任务状态）
transformer_tasks = {}

# 全局配置变量
app_config = {}

def load_config():
    """
    加载配置文件
    读取 config.json 文件，如果有则加载，否则返回空字典
    """
    config_path = os.path.join(os.path.dirname(__file__), 'config.json')
    if os.path.exists(config_path):
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                config = json.load(f)
                logger.info(f"✓ 已加载配置文件: {config_path}")
                return config
        except Exception as e:
            logger.warning(f"⚠ 加载配置文件失败: {e}")
            return {}
    else:
        logger.warning(f"⚠ 配置文件不存在: {config_path}")
        return {}

# 启动时加载配置
app_config = load_config()


@app.get("/", response_class=HTMLResponse)
async def home():
    # 读取并返回前端首页 HTML
    with open("templates/index.html", "r", encoding="utf-8") as f:
        return f.read()


@app.get("/api/config/defaults")
async def get_config_defaults():
    """
    获取默认配置
    返回 config.json 中的配置作为表单默认值
    """
    return JSONResponse(content={
        "success": True,
        "config": app_config
    })


@app.post("/api/config/save")
async def save_config(request: Request):
    """
    保存配置到 config.json
    用于在页面上修改配置后保存
    """
    try:
        data = await request.json()
        
        # 更新全局配置
        global app_config
        app_config = data
        
        # 写入文件
        config_path = os.path.join(os.path.dirname(__file__), 'config.json')
        with open(config_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, indent=4, ensure_ascii=False)
        
        return JSONResponse(content={
            "success": True,
            "message": "配置已保存"
        })
    except Exception as e:
        return JSONResponse(
            status_code=500,
            content={"success": False, "error": str(e)}
        )


@app.post("/api/generate")
async def generate_alpha(request: Request):
    """
    生成 Alpha 变种的 API 端点
    接收前端表单数据，启动 Transformer 脚本执行 Alpha 生成任务
    """
    logger.info("=" * 50)
    logger.info("收到生成变种请求")
    
    try:
        # 解析请求数据
        data = await request.json()
        logger.info(f"请求数据: alpha_id={data.get('alpha_id')}, llm_model={data.get('llm_model')}")
        
        # 生成唯一任务 ID
        task_id = str(uuid.uuid4())
        logger.info(f"生成任务 ID: {task_id}") 

        # 定义必须提交的字段
        required_fields = [
            "alpha_id",
            "llm_api_key",
            "llm_base_url",
            "llm_model",
            "brain_username",
            "brain_password"
        ]

        # 检查必填字段是否完整
        for field in required_fields:
            if not data.get(field):
                logger.warning(f"缺少必填字段: {field}")
                return JSONResponse(
                    status_code=400,
                    content={"success": False, "error": f"Missing required field: {field}"}
                )

        # 获取脚本所在目录和 Transformer 子目录
        script_dir = os.path.dirname(os.path.abspath(__file__))
        transformer_dir = os.path.join(script_dir, 'Tranformer')
        logger.info(f"Transformer 目录: {transformer_dir}")

        # 构建传递给 Transformer 脚本的配置
        config = {
            "LLM_model_name": data.get('llm_model'),
            "LLM_API_KEY": data.get('llm_api_key'),
            "llm_base_url": data.get('llm_base_url'),
            "username": data.get('brain_username'),
            "password": data.get('brain_password'),
            "alpha_id": data.get('alpha_id'),
            "top_n_datafield": int(data.get('top_n_datafield', 50)),
            "user_region": data.get('user_region'),
            "user_universe": data.get('user_universe'),
            "user_delay": int(data.get('user_delay')) if data.get('user_delay') else None,
            "user_category": data.get('user_category'),
            "user_data_type": data.get('user_data_type', 'MATRIX'),
            "max_retries": int(data.get('max_retries', 20))
        }
        logger.info(f"配置已构建: LLM_model={config['LLM_model_name']}, alpha_id={config['alpha_id']}, max_retries={config['max_retries']}")

        # 检查数据字段缓存是否存在，不存在则自动下载
        dataset_dir = os.path.join(script_dir, 'dataset')
        os.makedirs(dataset_dir, exist_ok=True)

        user_category = data.get('user_category', [])
        categories_to_check = user_category if isinstance(user_category, list) else [user_category] if user_category else []

        # 输出用户选择的类别
        logger.info(f"[Generate] 用户选择的类别: {categories_to_check}")
        logger.info(f"[Generate] Region: {data.get('user_region')}, Delay: {data.get('user_delay')}, Universe: {data.get('user_universe')}")

        # 创建会话来获取数据集信息
        session = SingleSession()
        session.auth = (data.get('brain_username'), data.get('brain_password'))
        brain_api_url = "https://api.worldquantbrain.com"
        auth_response = session.post(brain_api_url + "/authentication")

        if auth_response.status_code != 201:
            return JSONResponse(
                status_code=401,
                content={"success": False, "error": "Authentication failed"}
            )

        # 获取所有可用的数据集
        datasets_df = get_datasets(
            s=session,
            instrument_type="EQUITY",
            region=data.get('user_region', 'USA'),
            delay=int(data.get('user_delay', 1)),
            universe=data.get('user_universe', 'TOP3000'),
            theme="ALL"
        )

        if datasets_df.empty:
            return JSONResponse(
                status_code=404,
                content={"success": False, "error": "No datasets found"}
            )

        # 使用 category_name 列（category 列是字典类型）
        category_column = 'category_name' if 'category_name' in datasets_df.columns else 'category'
        
        # 如果没有指定类别，使用所有类别
        if not categories_to_check:
            if category_column == 'category':
                # category 是字典，提取 name
                categories_to_check = datasets_df[category_column].apply(lambda x: x['name'] if isinstance(x, dict) and 'name' in x else x).unique().tolist()
            else:
                categories_to_check = datasets_df[category_column].unique().tolist()
            logger.info(f"未指定类别，将使用所有 {len(categories_to_check)} 个类别: {categories_to_check}")

        # 构建类别到 dataset_id 的映射
        category_to_datasets = {}
        for cat in categories_to_check:
            if category_column == 'category':
                # category 是字典，需要匹配 name 字段
                cat_datasets = datasets_df[datasets_df[category_column].apply(lambda x: x.get('name') if isinstance(x, dict) else x) == cat]
            else:
                cat_datasets = datasets_df[datasets_df[category_column] == cat]
            if not cat_datasets.empty:
                category_to_datasets[cat] = cat_datasets['id'].tolist()

        # 检查每个类别的缓存文件
        missing_categories = []
        for cat in categories_to_check:
            cache_filename = f"datafields_cache_{data.get('user_region', 'USA')}_{data.get('user_universe', 'TOP3000')}_D{data.get('user_delay', 1)}_{cat}.csv"
            cache_path = os.path.join(dataset_dir, cache_filename)
            if not os.path.exists(cache_path):
                missing_categories.append(cat)

        # 如果有缺失的缓存，检查是否都能对应到数据集
        if missing_categories:
            # 检查是否有类别在数据集中不存在
            invalid_categories = [cat for cat in missing_categories if cat not in category_to_datasets]
            if invalid_categories:
                error_msg = f"以下类别在数据集中不存在，无法下载: {', '.join(invalid_categories)}。请检查类别名称或先点击'下载数据字段缓存'按钮。"
                logger.error(error_msg)
                return JSONResponse(
                    status_code=400,
                    content={"success": False, "error": error_msg}
                )

            logger.info(f"检测到缺失的数据字段缓存: {missing_categories}，自动下载中...")

            downloaded_count = 0
            failed_categories = []

            for cat in missing_categories:
                dataset_ids = category_to_datasets[cat]
                all_datafields = []

                try:
                    # 下载该类别下的所有数据集
                    for dataset_id in dataset_ids:
                        logger.info(f"下载数据集 {dataset_id} (类别: {cat})...")
                        datafields_df = get_datafields(
                            s=session,
                            instrument_type="EQUITY",
                            region=data.get('user_region', 'USA'),
                            delay=int(data.get('user_delay', 1)),
                            universe=data.get('user_universe', 'TOP3000'),
                            dataset_id=dataset_id,
                            data_type=data.get('user_data_type', 'MATRIX')
                        )

                        if not datafields_df.empty:
                            all_datafields.append(datafields_df)
                            logger.info(f"数据集 {dataset_id} 包含 {len(datafields_df)} 个字段")

                    if all_datafields:
                        # 合并该类别下的所有数据字段
                        combined_df = pd.concat(all_datafields, ignore_index=True)
                        combined_df.drop_duplicates(subset=['id'], inplace=True)

                        cache_filename = f"datafields_cache_{data.get('user_region', 'USA')}_{data.get('user_universe', 'TOP3000')}_D{data.get('user_delay', 1)}_{cat}.csv"
                        cache_path = os.path.join(dataset_dir, cache_filename)
                        combined_df.to_csv(cache_path, index=False)
                        downloaded_count += 1
                        logger.info(f"自动下载完成: {cache_filename} ({len(combined_df)} 个字段)")
                    else:
                        failed_categories.append(cat)
                        logger.error(f"类别 '{cat}' 没有获取到任何数据字段")

                except Exception as e:
                    failed_categories.append(cat)
                    logger.error(f"下载类别 '{cat}' 失败: {e}")

            # 如果有下载失败的类别，返回错误
            if failed_categories:
                error_msg = f"以下类别下载失败: {', '.join(failed_categories)}。请检查网络连接或手动下载。"
                logger.error(error_msg)
                return JSONResponse(
                    status_code=500,
                    content={"success": False, "error": error_msg}
                )

            logger.info(f"自动下载完成，共下载 {downloaded_count}/{len(missing_categories)} 个类别")

        # 直接调用 Transformer 模块（替代 subprocess）
        logger.info(f"启动 Transformer 执行...")
        try:
            transformer_result = await run_transformer(config)
            logger.info(f"Transformer 执行完成: {transformer_result}")

            # 定义输出文件路径
            output_file = os.path.join(transformer_dir, 'output', 'Alpha_generated_expressions_success.json')
            candidates_file = os.path.join(transformer_dir, 'output', 'Alpha_candidates.json')
            error_file = os.path.join(transformer_dir, 'output', 'Alpha_generated_expressions_error.json')

            # 构建响应数据
            result = {
                "success": True,
                "alpha_id": data.get('alpha_id'),
                "stdout": "",
                "stderr": "",
                "return_code": 0
            }

            # 读取成功生成的表达式
            if os.path.exists(output_file):
                logger.info(f"读取成功表达式文件: {output_file}")
                with open(output_file, 'r', encoding='utf-8') as f:
                    result['expressions_success'] = json.load(f)
            else:
                logger.info(f"成功表达式文件不存在: {output_file}")
                result['expressions_success'] = []

            # 读取候选表达式
            if os.path.exists(candidates_file):
                logger.info(f"读取候选表达式文件: {candidates_file}")
                with open(candidates_file, 'r', encoding='utf-8') as f:
                    result['candidates'] = json.load(f)
            else:
                logger.info(f"候选表达式文件不存在: {candidates_file}")
                result['candidates'] = []

            # 读取生成失败的表达式
            if os.path.exists(error_file):
                logger.info(f"读取错误表达式文件: {error_file}")
                with open(error_file, 'r', encoding='utf-8') as f:
                    result['expressions_error'] = json.load(f)
            else:
                logger.info(f"错误表达式文件不存在: {error_file}")
                result['expressions_error'] = []

            logger.info(f"成功: {len(result['expressions_success'])} 个, 候选: {len(result['candidates'])} 个, 错误: {len(result['expressions_error'])} 个")
            logger.info("=" * 50)
            return JSONResponse(content=result)

        except Exception as transformer_error:
            logger.error(f"Transformer 执行失败: {transformer_error}")
            raise
    except Exception as e:
        logger.warning(f"执行异常: {str(e)}")
        return JSONResponse(
                status_code=500,
                content={"success": False, "error": str(e)}
            )


@app.post("/api/transformer/login-and-fetch-options")
async def login_and_fetch_options(request: Request):
    """
    登录 BRAIN 并获取地区、Delay、Universe、类别等选项
    用于填充高级选项表单
    """
    try:
        data = await request.json()
        username = data.get('username')
        password = data.get('password')
        
        if not username or not password:
            return JSONResponse(
                status_code=400,
                content={'success': False, 'error': 'Username and password are required'}
            )
        
        # 添加 Transformer 目录到 sys.path
        script_dir = os.path.dirname(os.path.abspath(__file__))
        transformer_dir = os.path.join(script_dir, 'Tranformer')
        if transformer_dir not in sys.path:
            sys.path.append(transformer_dir)
        
        # 创建新的会话实例
        session = SingleSession()
        session.auth = (username, password)
        
        brain_api_url = "https://api.worldquantbrain.com"
        response = session.post(brain_api_url + "/authentication")
        
        if response.status_code == 201:
            # 认证成功
            pass
        elif response.status_code == 401:
            return JSONResponse(
                status_code=401,
                content={'success': False, 'error': 'Authentication failed: Invalid credentials'}
            )
        else:
            return JSONResponse(
                status_code=400,
                content={'success': False, 'error': f'Authentication failed: {response.status_code}'}
            )
        
        # 获取 region/delay/universe 选项
        df = get_instrument_type_region_delay(session)
        
        # 获取数据类别
        categories_resp = session.get(brain_api_url + "/data-categories")
        categories = []
        if categories_resp.status_code == 200:
            categories_data = categories_resp.json()
            if isinstance(categories_data, list):
                categories = categories_data
            elif isinstance(categories_data, dict):
                categories = categories_data.get('results', [])
        
        # 转换 DataFrame 为前端需要的嵌套字典结构
        # 结构: Region -> Delay -> Universe
        df_equity = df[df['InstrumentType'] == 'EQUITY']
        
        options = {}
        for _, row in df_equity.iterrows():
            region = row['Region']
            delay = row['Delay']
            universes = row['Universe']  # 这是一个列表
            
            if region not in options:
                options[region] = {}
            
            # 将 delay 转换为字符串作为字典的键
            delay_str = str(delay)
            if delay_str not in options[region]:
                options[region][delay_str] = universes
        
        return JSONResponse(content={
            'success': True,
            'options': options,
            'categories': categories
        })
        
    except Exception as e:
        logger.warning(f"登录获取选项失败: {str(e)}")
        return JSONResponse(
            status_code=500,
            content={'success': False, 'error': str(e)}
        )


@app.get("/api/health")
async def health_check():
    """健康检查端点，用于验证服务是否正常运行"""
    return {"status": "healthy", "service": "alpha-transformer"}


@app.post("/api/test-llm")
async def test_llm_connection(request: Request):
    """
    测试 LLM 连接
    接收 LLM 配置，尝试连接并返回测试结果
    对 529 状态码会自动重试，默认重试 3 次
    """
    try:
        data = await request.json()
        
        api_key = data.get('llm_api_key')
        base_url = data.get('llm_base_url')
        model = data.get('llm_model')
        max_retries = data.get('max_retries', 3)
        
        if not api_key or not base_url or not model:
            return JSONResponse(
                status_code=400,
                content={"success": False, "error": "Missing required LLM configuration"}
            )
        
        logger.info(f"测试 LLM 连接: {base_url}, 模型: {model}, 最大重试次数: {max_retries}")
        
        # 创建客户端
        client = openai.AsyncOpenAI(
            api_key=api_key,
            base_url=base_url
        )

        # 重试机制
        last_error = None
        for attempt in range(1, max_retries + 1):
            try:
                response = await client.chat.completions.create(
                    model=model,
                    messages=[{"role": "user", "content": "Hello, this is a connection test. Reply with 'OK' only."}],
                    max_tokens=10,
                    timeout=30
                )

                # 检查响应
                if response and response.choices and len(response.choices) > 0:
                    content = response.choices[0].message.content
                    logger.info(f"✓ LLM 连接测试成功: {content[:50]}...")
                    return JSONResponse(content={
                        "success": True,
                        "message": "LLM 连接成功",
                        "response": content[:100]
                    })
                else:
                    return JSONResponse(
                        status_code=500,
                        content={"success": False, "error": "Empty response from LLM"}
                    )
                    
            except Exception as e:
                error_msg = str(e)
                last_error = e
                
                # 检查是否是 529 错误（MiniMax 过载错误）
                if "529" in error_msg or "overloaded" in error_msg.lower():
                    if attempt < max_retries:
                        logger.warning(f"⚠ LLM 返回 529 错误，第 {attempt}/{max_retries} 次尝试失败，2 秒后重试...")
                        await asyncio.sleep(2)
                        continue
                    else:
                        logger.error(f"✗ LLM 连接测试失败，已重试 {max_retries} 次，仍然返回 529 错误")
                        return JSONResponse(
                            status_code=529,
                            content={"success": False, "error": f"LLM service overloaded (529) after {max_retries} retries"}
                        )
                else:
                    # 其他错误直接抛出
                    raise
                    
    except Exception as e:
        error_msg = str(e)
        logger.warning(f"✗ LLM 连接测试失败: {error_msg}")
        return JSONResponse(
            status_code=500,
            content={"success": False, "error": error_msg}
        )


@app.post("/api/download-datafields")
async def download_datafields(request: Request):
    """
    下载数据字段缓存
    全量获取指定 region/delay/universe/category 的数据字段并保存到 CSV
    """
    try:
        data = await request.json()
        username = data.get('username')
        password = data.get('password')
        region = data.get('region')
        delay = data.get('delay')
        universe = data.get('universe')
        data_type = data.get('data_type', 'MATRIX')
        category = data.get('category', [])  # 类别列表

        if not all([username, password, region, delay, universe]):
            return JSONResponse(
                status_code=400,
                content={"success": False, "error": "Missing required parameters"}
            )

        # 创建会话
        session = SingleSession()
        session.auth = (username, password)

        brain_api_url = "https://api.worldquantbrain.com"
        response = session.post(brain_api_url + "/authentication")

        if response.status_code != 201:
            return JSONResponse(
                status_code=401,
                content={"success": False, "error": "Authentication failed"}
            )

        # 获取所有数据集信息
        datasets_df = get_datasets(
            s=session,
            instrument_type="EQUITY",
            region=region,
            delay=delay,
            universe=universe,
            theme="ALL"
        )

        if datasets_df.empty:
            return JSONResponse(
                status_code=404,
                content={"success": False, "error": "No datasets found"}
            )

        # 处理类别参数（前端总是传递类别列表）
        categories_to_download = category if isinstance(category, list) else [category] if category else []

        # 如果没有指定类别，使用所有类别
        if not categories_to_download:
            if 'category_id' in datasets_df.columns:
                categories_to_download = datasets_df['category_id'].unique().tolist()
            elif 'category_name' in datasets_df.columns:
                categories_to_download = datasets_df['category_name'].unique().tolist()
            else:
                # category 是字典，提取 id
                categories_to_download = datasets_df['category'].apply(lambda x: x['id'] if isinstance(x, dict) and 'id' in x else x).unique().tolist()
            logger.info(f"未指定类别，将使用所有 {len(categories_to_download)} 个类别")

        # 构建类别到 dataset_id 的映射（使用 category_id 匹配）
        category_to_datasets = {}
        for cat in categories_to_download:
            if 'category_id' in datasets_df.columns:
                cat_datasets = datasets_df[datasets_df['category_id'] == cat]
            elif 'category_name' in datasets_df.columns:
                cat_datasets = datasets_df[datasets_df['category_name'] == cat]
            else:
                cat_datasets = datasets_df[datasets_df['category'].apply(lambda x: x.get('id') if isinstance(x, dict) else x) == cat]
            if not cat_datasets.empty:
                category_to_datasets[cat] = cat_datasets['id'].tolist()

        # 保存到 CSV (dataset 文件夹)
        dataset_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'dataset')
        os.makedirs(dataset_dir, exist_ok=True)

        # 检查哪些类别已经存在缓存文件
        categories_to_download_filtered = []
        already_exists = []
        for cat in categories_to_download:
            cache_filename = f"datafields_cache_{region}_{universe}_D{delay}_{cat}.csv"
            cache_path = os.path.join(dataset_dir, cache_filename)
            if os.path.exists(cache_path):
                already_exists.append(cat)
                logger.info(f"类别 '{cat}' 的缓存文件已存在，跳过下载: {cache_filename}")
            else:
                categories_to_download_filtered.append(cat)

        if already_exists:
            logger.info(f"以下类别缓存已存在，将跳过: {already_exists}")

        if not categories_to_download_filtered:
            logger.info("所有类别的缓存文件都已存在，无需下载")
            return JSONResponse(content={
                "success": True,
                "count": 0,
                "message": "所有类别的缓存文件都已存在",
                "skipped_categories": already_exists
            })

        categories_to_download = categories_to_download_filtered

        logger.info(f"开始下载数据字段: region={region}, delay={delay}, universe={universe}, type={data_type}, categories={categories_to_download}")

        total_count = 0
        saved_files = []

        # 每个类别单独下载并保存
        for cat in categories_to_download:
            if cat not in category_to_datasets:
                logger.warning(f"类别 '{cat}' 在数据集中不存在，跳过")
                continue

            dataset_ids = category_to_datasets[cat]
            all_datafields = []

            # 下载该类别下的所有数据集
            for dataset_id in dataset_ids:
                logger.info(f"下载数据集 {dataset_id} (类别: {cat})...")
                datafields_df = get_datafields(
                    s=session,
                    instrument_type="EQUITY",
                    region=region,
                    delay=delay,
                    universe=universe,
                    dataset_id=dataset_id,
                    data_type=data_type
                )

                if not datafields_df.empty:
                    all_datafields.append(datafields_df)
                    logger.info(f"数据集 {dataset_id} 包含 {len(datafields_df)} 个字段")

            if all_datafields:
                # 合并该类别下的所有数据字段
                combined_df = pd.concat(all_datafields, ignore_index=True)
                combined_df.drop_duplicates(subset=['id'], inplace=True)

                cache_filename = f"datafields_cache_{region}_{universe}_D{delay}_{cat}.csv"
                cache_path = os.path.join(dataset_dir, cache_filename)
                combined_df.to_csv(cache_path, index=False)
                total_count += len(combined_df)
                saved_files.append(cache_filename)
                logger.info(f"类别 '{cat}' 数据字段已保存: {cache_path}, 共 {len(combined_df)} 个字段")
            else:
                logger.warning(f"类别 '{cat}' 未找到数据字段")

        if total_count == 0:
            return JSONResponse(
                status_code=404,
                content={"success": False, "error": "No datafields found for any category"}
            )

        return JSONResponse(content={
            "success": True,
            "count": total_count,
            "cache_files": saved_files
        })

    except Exception as e:
        logger.error(f"下载数据字段失败: {e}")
        return JSONResponse(
            status_code=500,
            content={"success": False, "error": str(e)}
        )


@app.get("/api/download/{alpha_id}")
async def download_results(alpha_id: str):
    """
    下载生成结果的 zip 压缩包
    包含三个 JSON 文件：success, candidates, error
    """
    try:
        transformer_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'Tranformer')
        output_dir = os.path.join(transformer_dir, 'output')
        
        # 检查文件是否存在
        files_to_zip = {
            'Alpha_generated_expressions_success.json': os.path.join(output_dir, 'Alpha_generated_expressions_success.json'),
            'Alpha_candidates.json': os.path.join(output_dir, 'Alpha_candidates.json'),
            'Alpha_generated_expressions_error.json': os.path.join(output_dir, 'Alpha_generated_expressions_error.json')
        }
        
        # 生成时间戳
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        zip_filename = f"{alpha_id}_{timestamp}.zip"
        
        # 保存到 save_zip 文件夹（不会被清空）
        save_zip_dir = os.path.join(transformer_dir, 'save_zip')
        os.makedirs(save_zip_dir, exist_ok=True)
        zip_path = os.path.join(save_zip_dir, zip_filename)
        
        # 创建 zip 文件
        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
            for arcname, filepath in files_to_zip.items():
                if os.path.exists(filepath):
                    zipf.write(filepath, arcname)
        
        # 返回文件
        return FileResponse(
            zip_path,
            media_type='application/zip',
            filename=zip_filename
        )
        
    except Exception as e:
        return JSONResponse(
            status_code=500,
            content={"success": False, "error": str(e)}
        )


if __name__ == "__main__":
    # 启动 FastAPI 应用，使用 uvicorn 作为 ASGI 服务器
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)