YG_FT_Platform/src/main.py

"""
远光软件微调平台 - Flask 后端 API
"""
import os
import sys
import json
import pymysql
import yaml
import time
from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS
from werkzeug.utils import secure_filename

# 导入API蓝图
from api import register_blueprints

# 获取项目根目录
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, PROJECT_ROOT)

# 加载配置
CONFIG_PATH = os.path.join(PROJECT_ROOT, 'config.yaml')


def load_config():
    """加载配置文件"""
    with open(CONFIG_PATH, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)


CONFIG = load_config()


def get_db_connection():
    """获取数据库连接"""
    db_config = CONFIG['database']
    return pymysql.connect(
        host=db_config['host'],
        port=db_config['port'],
        user=db_config['username'],
        password=db_config['password'],
        database=db_config['name'],
        charset=db_config.get('charset', 'utf8mb4'),
        cursorclass=pymysql.cursors.DictCursor
    )


def init_database():
    """初始化数据库表"""
    print("正在初始化数据库...")
    try:
        conn = get_db_connection()
        cursor = conn.cursor()

        tables = [
            # 精调训练表
            """CREATE TABLE IF NOT EXISTS fine_tune (
                id INT AUTO_INCREMENT PRIMARY KEY,
                name VARCHAR(255) NOT NULL,
                base_model VARCHAR(255),
                train_type VARCHAR(50),
                train_method VARCHAR(50),
                dataset_id INT,
                valid_split VARCHAR(50),
                valid_ratio INT DEFAULT 10,
                output_model_name VARCHAR(255),
                status VARCHAR(50) DEFAULT 'pending',
                progress INT DEFAULT 0,
                create_time DATETIME DEFAULT CURRENT_TIMESTAMP,
                update_time DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4""",

            # 我的模型表
            """CREATE TABLE IF NOT EXISTS my_models (
                id INT AUTO_INCREMENT PRIMARY KEY,
                name VARCHAR(255) NOT NULL,
                type VARCHAR(100),
                version VARCHAR(50),
                description TEXT,
                create_time DATETIME DEFAULT CURRENT_TIMESTAMP,
                update_time DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4""",

            # 模型评测表
            """CREATE TABLE IF NOT EXISTS model_eval (
                id INT AUTO_INCREMENT PRIMARY KEY,
                model_name VARCHAR(255) NOT NULL,
                dataset VARCHAR(255),
                metric VARCHAR(100),
                score DECIMAL(10, 4),
                status VARCHAR(50) DEFAULT 'completed',
                create_time DATETIME DEFAULT CURRENT_TIMESTAMP
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4""",

            # 模型部署表
            """CREATE TABLE IF NOT EXISTS model_deploy (
                id INT AUTO_INCREMENT PRIMARY KEY,
                model_name VARCHAR(255) NOT NULL,
                endpoint VARCHAR(255),
                instance VARCHAR(100),
                status VARCHAR(50) DEFAULT 'running',
                create_time DATETIME DEFAULT CURRENT_TIMESTAMP,
                update_time DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4""",

            # 数据集管理表
            """CREATE TABLE IF NOT EXISTS dataset_manage (
                id INT AUTO_INCREMENT PRIMARY KEY,
                name VARCHAR(255) NOT NULL,
                type VARCHAR(100),
                size VARCHAR(50),
                count INT,
                description TEXT,
                file_path VARCHAR(500),
                file_count INT DEFAULT 0,
                storage_type VARCHAR(50) DEFAULT 'local',
                minio_config TEXT,
                create_time DATETIME DEFAULT CURRENT_TIMESTAMP,
                update_time DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4""",

            # 数据集文件表
            """CREATE TABLE IF NOT EXISTS dataset_files (
                id INT AUTO_INCREMENT PRIMARY KEY,
                dataset_id INT NOT NULL,
                file_name VARCHAR(255) NOT NULL,
                file_path VARCHAR(500) NOT NULL,
                file_size BIGINT DEFAULT 0,
                file_type VARCHAR(50),
                create_time DATETIME DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (dataset_id) REFERENCES dataset_manage(id) ON DELETE CASCADE
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4""",

            # 数据生成表
            """CREATE TABLE IF NOT EXISTS data_generate (
                id INT AUTO_INCREMENT PRIMARY KEY,
                name VARCHAR(255) NOT NULL,
                template VARCHAR(255),
                count INT DEFAULT 0,
                status VARCHAR(50) DEFAULT 'pending',
                create_time DATETIME DEFAULT CURRENT_TIMESTAMP,
                update_time DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4""",

            # 权限管理表
            """CREATE TABLE IF NOT EXISTS permission (
                id INT AUTO_INCREMENT PRIMARY KEY,
                username VARCHAR(100) NOT NULL,
                role VARCHAR(50) DEFAULT 'user',
                permissions TEXT,
                create_time DATETIME DEFAULT CURRENT_TIMESTAMP,
                update_time DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4""",

            # 模型管理表
            """CREATE TABLE IF NOT EXISTS model_manage (
                id INT AUTO_INCREMENT PRIMARY KEY,
                name VARCHAR(255) NOT NULL,
                type VARCHAR(100),
                model_source VARCHAR(20) DEFAULT 'local',
                path VARCHAR(500),
                api_url VARCHAR(500),
                api_key VARCHAR(500),
                model_name VARCHAR(255),
                description TEXT,
                create_time DATETIME DEFAULT CURRENT_TIMESTAMP,
                update_time DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4""",

            # 系统配置表
            """CREATE TABLE IF NOT EXISTS sys_config (
                id INT AUTO_INCREMENT PRIMARY KEY,
                config_key VARCHAR(100) NOT NULL UNIQUE,
                config_value TEXT,
                description VARCHAR(255),
                update_time DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4""",

            # 用户表
            """CREATE TABLE IF NOT EXISTS users (
                id INT AUTO_INCREMENT PRIMARY KEY,
                username VARCHAR(100) NOT NULL UNIQUE,
                password VARCHAR(255) NOT NULL,
                role VARCHAR(50) DEFAULT 'user',
                create_time DATETIME DEFAULT CURRENT_TIMESTAMP
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4"""
        ]

        for i, table_sql in enumerate(tables):
            try:
                cursor.execute(table_sql)
                print(f"  表 {i+1}/{len(tables)} 创建/检查成功")
            except Exception as e:
                print(f"  表 {i+1} 创建失败: {e}")

        # 插入默认管理员用户
        cursor.execute("SELECT * FROM users WHERE username = 'admin'")
        if not cursor.fetchone():
            cursor.execute("INSERT INTO users (username, password, role) VALUES ('admin', 'admin', 'admin')")
            print("  默认管理员用户创建成功")

        conn.commit()
        cursor.close()
        conn.close()
        print("数据库初始化完成")
    except Exception as e:
        print(f"数据库初始化失败: {e}")
        raise


app = Flask(__name__)
app.config['SECRET_KEY'] = CONFIG['secret_key']
CORS(app, resources={r"/api/*": {"origins": "*"}})

# 注册蓝图
register_blueprints(app)


# ============ 健康检查 ============
@app.route('/api/health', methods=['GET'])
def health_check():
    """健康检查接口"""
    return jsonify({'status': 'ok', 'code': 0})


# ============ 通用 CRUD 操作 ============
def generic_get_all(table_name, order_by='create_time DESC'):
    """通用查询所有"""
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute(f"SELECT * FROM {table_name} ORDER BY {order_by}")
    result = cursor.fetchall()
    cursor.close()
    conn.close()
    return result


def generic_get_by_id(table_name, id_val):
    """通用按ID查询"""
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute(f"SELECT * FROM {table_name} WHERE id = %s", (id_val,))
    result = cursor.fetchone()
    cursor.close()
    conn.close()
    return result


def generic_create(table_name, data):
    """通用创建"""
    conn = get_db_connection()
    cursor = conn.cursor()
    columns = ', '.join(data.keys())
    placeholders = ', '.join(['%s'] * len(data))
    sql = f"INSERT INTO {table_name} ({columns}) VALUES ({placeholders})"
    cursor.execute(sql, list(data.values()))
    conn.commit()
    new_id = cursor.lastrowid
    cursor.close()
    conn.close()
    return new_id


def generic_update(table_name, id_val, data):
    """通用更新"""
    conn = get_db_connection()
    cursor = conn.cursor()
    set_clause = ', '.join([f"{k} = %s" for k in data.keys()])
    sql = f"UPDATE {table_name} SET {set_clause} WHERE id = %s"
    values = list(data.values()) + [id_val]
    cursor.execute(sql, values)
    conn.commit()
    cursor.close()
    conn.close()


def generic_delete(table_name, id_val):
    """通用删除"""
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute(f"DELETE FROM {table_name} WHERE id = %s", (id_val,))
    conn.commit()
    cursor.close()
    conn.close()


# ============ 登录接口 ============
@app.route('/api/login', methods=['POST'])
def login():
    data = request.json
    username = data.get('username')
    password = data.get('password')

    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM users WHERE username = %s AND password = %s", (username, password))
    user = cursor.fetchone()
    cursor.close()
    conn.close()

    if user:
        return jsonify({'code': 0, 'message': '登录成功', 'data': {'username': user['username'], 'role': user['role']}})
    return jsonify({'code': 1, 'message': '用户名或密码错误'})


# ============ 精调训练接口 ============
@app.route('/api/fine-tune', methods=['GET'])
def get_fine_tune():
    return jsonify({'code': 0, 'data': generic_get_all('fine_tune')})


@app.route('/api/fine-tune', methods=['POST'])
def create_fine_tune():
    data = request.json
    new_id = generic_create('fine_tune', data)
    return jsonify({'code': 0, 'message': '创建成功', 'id': new_id})


@app.route('/api/fine-tune/<int:id>', methods=['PUT'])
def update_fine_tune(id):
    data = request.json
    generic_update('fine_tune', id, data)
    return jsonify({'code': 0, 'message': '更新成功'})


@app.route('/api/fine-tune/<int:id>', methods=['DELETE'])
def delete_fine_tune(id):
    generic_delete('fine_tune', id)
    return jsonify({'code': 0, 'message': '删除成功'})


# ============ 我的模型接口 ============
@app.route('/api/my-models', methods=['GET'])
def get_my_models():
    return jsonify({'code': 0, 'data': generic_get_all('my_models')})


@app.route('/api/my-models', methods=['POST'])
def create_my_model():
    data = request.json
    new_id = generic_create('my_models', data)
    return jsonify({'code': 0, 'message': '创建成功', 'id': new_id})


@app.route('/api/my-models/<int:id>', methods=['PUT'])
def update_my_model(id):
    data = request.json
    generic_update('my_models', id, data)
    return jsonify({'code': 0, 'message': '更新成功'})


@app.route('/api/my-models/<int:id>', methods=['DELETE'])
def delete_my_model(id):
    generic_delete('my_models', id)
    return jsonify({'code': 0, 'message': '删除成功'})


# ============ 模型评测接口 ============
@app.route('/api/model-eval', methods=['GET'])
def get_model_eval():
    return jsonify({'code': 0, 'data': generic_get_all('model_eval')})


@app.route('/api/model-eval', methods=['POST'])
def create_model_eval():
    data = request.json
    new_id = generic_create('model_eval', data)
    return jsonify({'code': 0, 'message': '创建成功', 'id': new_id})


@app.route('/api/model-eval/<int:id>', methods=['PUT'])
def update_model_eval(id):
    data = request.json
    generic_update('model_eval', id, data)
    return jsonify({'code': 0, 'message': '更新成功'})


@app.route('/api/model-eval/<int:id>', methods=['DELETE'])
def delete_model_eval(id):
    generic_delete('model_eval', id)
    return jsonify({'code': 0, 'message': '删除成功'})


# ============ 模型部署接口 ============
@app.route('/api/model-deploy', methods=['GET'])
def get_model_deploy():
    return jsonify({'code': 0, 'data': generic_get_all('model_deploy')})


@app.route('/api/model-deploy', methods=['POST'])
def create_model_deploy():
    data = request.json
    new_id = generic_create('model_deploy', data)
    return jsonify({'code': 0, 'message': '创建成功', 'id': new_id})


@app.route('/api/model-deploy/<int:id>', methods=['PUT'])
def update_model_deploy(id):
    data = request.json
    generic_update('model_deploy', id, data)
    return jsonify({'code': 0, 'message': '更新成功'})


@app.route('/api/model-deploy/<int:id>', methods=['DELETE'])
def delete_model_deploy(id):
    generic_delete('model_deploy', id)
    return jsonify({'code': 0, 'message': '删除成功'})


# ============ 数据生成接口 ============
@app.route('/api/data-generate', methods=['GET'])
def get_data_generate():
    return jsonify({'code': 0, 'data': generic_get_all('data_generate')})


@app.route('/api/data-generate', methods=['POST'])
def create_data_generate():
    data = request.json
    new_id = generic_create('data_generate', data)
    return jsonify({'code': 0, 'message': '创建成功', 'id': new_id})


@app.route('/api/data-generate/<int:id>', methods=['PUT'])
def update_data_generate(id):
    data = request.json
    generic_update('data_generate', id, data)
    return jsonify({'code': 0, 'message': '更新成功'})


@app.route('/api/data-generate/<int:id>', methods=['DELETE'])
def delete_data_generate(id):
    generic_delete('data_generate', id)
    return jsonify({'code': 0, 'message': '删除成功'})


# ============ 权限管理接口 ============
@app.route('/api/permission', methods=['GET'])
def get_permission():
    return jsonify({'code': 0, 'data': generic_get_all('permission')})


@app.route('/api/permission', methods=['POST'])
def create_permission():
    data = request.json
    new_id = generic_create('permission', data)
    return jsonify({'code': 0, 'message': '创建成功', 'id': new_id})


@app.route('/api/permission/<int:id>', methods=['PUT'])
def update_permission(id):
    data = request.json
    generic_update('permission', id, data)
    return jsonify({'code': 0, 'message': '更新成功'})


@app.route('/api/permission/<int:id>', methods=['DELETE'])
def delete_permission(id):
    generic_delete('permission', id)
    return jsonify({'code': 0, 'message': '删除成功'})


# ============ 系统配置接口 ============
@app.route('/api/sys-config', methods=['GET'])
def get_sys_config():
    return jsonify({'code': 0, 'data': generic_get_all('sys_config')})


@app.route('/api/sys-config', methods=['POST'])
def create_sys_config():
    data = request.json
    new_id = generic_create('sys_config', data)
    return jsonify({'code': 0, 'message': '创建成功', 'id': new_id})


@app.route('/api/sys-config/<int:id>', methods=['PUT'])
def update_sys_config(id):
    data = request.json
    generic_update('sys_config', id, data)
    return jsonify({'code': 0, 'message': '更新成功'})


@app.route('/api/sys-config/<int:id>', methods=['DELETE'])
def delete_sys_config(id):
    generic_delete('sys_config', id)
    return jsonify({'code': 0, 'message': '删除成功'})


if __name__ == '__main__':
    # 启动前先初始化数据库
    init_database()
    app_config = CONFIG['app']
    app.run(host=app_config['host'], port=app_config['port'], debug=app_config.get('debug', True))