Tolan/backend/init_database.py

"""
数据库表初始化脚本
运行: python init_database.py
"""
import asyncio
import aiomysql
from config import settings
import sys


# 用户表
CREATE_USERS_TABLE = """
CREATE TABLE IF NOT EXISTS users (
    id BIGINT PRIMARY KEY AUTO_INCREMENT,
    user_id VARCHAR(64) UNIQUE NOT NULL COMMENT '用户唯一标识(UUID)',
    
    -- 基本信息
    username VARCHAR(100) NOT NULL COMMENT '用户名/昵称',
    user_type ENUM('guest', 'registered') DEFAULT 'guest' COMMENT '用户类型',
    bot_name VARCHAR(100) COMMENT '用户给机器人取的名字',
    
    -- 注册信息（可选）
    phone VARCHAR(20) UNIQUE COMMENT '手机号',
    wechat_openid VARCHAR(100) UNIQUE COMMENT '微信OpenID',
    google_email VARCHAR(200) UNIQUE COMMENT 'Google邮箱',
    
    -- 认证信息
    is_verified BOOLEAN DEFAULT FALSE COMMENT '是否已认证',
    registered_at TIMESTAMP NULL COMMENT '注册时间',
    
    -- 统计信息
    total_messages INT DEFAULT 0 COMMENT '总消息数',
    total_rounds INT DEFAULT 0 COMMENT '总对话轮次',
    total_speech_duration_ms BIGINT DEFAULT 0 COMMENT '累积说话时长（毫秒）',
    last_active_at TIMESTAMP NULL COMMENT '最后活跃时间',
    
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
    
    INDEX idx_user_id (user_id),
    INDEX idx_phone (phone),
    INDEX idx_wechat_openid (wechat_openid),
    INDEX idx_google_email (google_email),
    INDEX idx_user_type (user_type)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
COMMENT='用户表'
"""

# 消息表
CREATE_MESSAGES_TABLE = """
CREATE TABLE IF NOT EXISTS messages (
    id BIGINT PRIMARY KEY AUTO_INCREMENT,
    message_id VARCHAR(64) UNIQUE NOT NULL COMMENT '消息ID',
    user_id VARCHAR(64) NOT NULL COMMENT '用户ID',
    
    -- 消息内容
    role ENUM('user', 'assistant', 'system') NOT NULL COMMENT '角色',
    content TEXT NOT NULL COMMENT '消息内容',
    
    -- 轮次信息
    round_number INT NOT NULL COMMENT '对话轮次(从1开始)',
    
    -- 扩展信息
    emotion VARCHAR(50) COMMENT '情感标签',
    message_metadata JSON COMMENT '扩展信息(音频路径、识别耗时等)',
    token_count INT COMMENT 'Token数量',
    
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    
    INDEX idx_user_id (user_id),
    INDEX idx_user_round (user_id, round_number),
    INDEX idx_created_at (created_at)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
COMMENT='消息表'
"""

# 对话总结表
CREATE_SUMMARIES_TABLE = """
CREATE TABLE IF NOT EXISTS conversation_summaries (
    id BIGINT PRIMARY KEY AUTO_INCREMENT,
    summary_id VARCHAR(64) UNIQUE NOT NULL COMMENT '总结ID',
    user_id VARCHAR(64) NOT NULL COMMENT '用户ID',
    
    -- 总结范围
    start_round INT NOT NULL COMMENT '起始轮次',
    end_round INT NOT NULL COMMENT '结束轮次',
    
    -- 总结内容
    summary_text TEXT NOT NULL COMMENT '总结内容',
    summary_type ENUM('auto', 'manual') DEFAULT 'auto' COMMENT '总结类型',
    
    -- 统计信息
    message_count INT COMMENT '总结的消息数量',
    token_count INT COMMENT '总结文本Token数',
    
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    
    INDEX idx_user_id (user_id),
    INDEX idx_user_end_round (user_id, end_round),
    UNIQUE KEY uk_user_round_range (user_id, start_round, end_round)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
COMMENT='对话总结表'
"""

# LLM提示词记录表
CREATE_LLM_PROMPTS_TABLE = """
CREATE TABLE IF NOT EXISTS llm_prompts (
    id BIGINT PRIMARY KEY AUTO_INCREMENT,
    prompt_id VARCHAR(64) UNIQUE NOT NULL COMMENT '提示词记录ID',
    
    -- 关联消息表（外键关联到messages.id）
    message_id BIGINT NOT NULL COMMENT '关联的消息ID',
    
    -- 用户信息
    user_id VARCHAR(64) NOT NULL COMMENT '用户ID',
    round_number INT NOT NULL COMMENT '对话轮次',
    
    -- 提示词内容
    system_prompt TEXT NOT NULL COMMENT '系统提示词（完整内容）',
    runtime_context TEXT NULL COMMENT '运行时上下文',
    messages_json JSON NOT NULL COMMENT '完整的messages数组（JSON格式）',
    
    -- LLM调用参数
    model_name VARCHAR(100) NOT NULL COMMENT '模型名称',
    temperature VARCHAR(20) NULL COMMENT '温度参数',
    max_tokens INT NULL COMMENT '最大Token数',
    stream BOOLEAN DEFAULT TRUE COMMENT '是否流式输出',
    
    -- 统计信息
    prompt_token_count INT NULL COMMENT '提示词Token数',
    total_token_count INT NULL COMMENT '总Token数（包含响应）',
    
    -- 扩展信息
    metadata_json JSON NULL COMMENT '扩展信息（调用耗时、错误信息等）',
    
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    
    INDEX idx_prompt_id (prompt_id),
    INDEX idx_llm_prompt_message (message_id),
    INDEX idx_llm_prompt_user_round (user_id, round_number),
    INDEX idx_llm_prompt_created (created_at),
    
    FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE CASCADE
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
COMMENT='LLM提示词记录表'
"""

# 任务类型表
CREATE_MISSION_TYPES_TABLE = """
CREATE TABLE IF NOT EXISTS mission_types (
    id BIGINT PRIMARY KEY AUTO_INCREMENT,
    type_key VARCHAR(64) UNIQUE NOT NULL COMMENT '任务类型唯一标识',
    display_name VARCHAR(100) NOT NULL COMMENT '任务显示名称',
    description VARCHAR(255) NULL COMMENT '类型描述',
    is_repeatable BOOLEAN DEFAULT FALSE COMMENT '是否可重复触发',
    is_active BOOLEAN DEFAULT TRUE COMMENT '是否启用',
    metadata JSON NULL COMMENT '扩展配置',
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
    INDEX idx_mission_type_active (is_active)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
COMMENT='任务类型配置表'
"""

# 任务能量石表
CREATE_MISSION_ENERGY_STONES_TABLE = """
CREATE TABLE IF NOT EXISTS mission_energy_stones (
    id BIGINT PRIMARY KEY AUTO_INCREMENT,
    stone_id VARCHAR(64) UNIQUE NOT NULL COMMENT '能量石ID',
    user_id VARCHAR(64) NOT NULL COMMENT '用户ID',
    request_id VARCHAR(64) NULL COMMENT '触发的请求ID',
    appear_reason VARCHAR(64) NOT NULL DEFAULT 'first_dialog' COMMENT '触发原因',
    mission_type VARCHAR(64) NOT NULL DEFAULT 'first_dialog' COMMENT '任务类型',
    appear_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT '出现时间',
    clicked_at TIMESTAMP NULL COMMENT '点击时间',
    auto_removed BOOLEAN DEFAULT FALSE COMMENT '是否自动消失',
    metadata JSON NULL COMMENT '扩展信息',
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
    INDEX idx_mission_stone_user (user_id, appear_at),
    INDEX idx_mission_type_user (mission_type, user_id)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
COMMENT='任务能量石表'
"""

# 对话能量石表
CREATE_CONVERSATION_ENERGY_STONES_TABLE = """
CREATE TABLE IF NOT EXISTS conversation_energy_stones (
    id BIGINT PRIMARY KEY AUTO_INCREMENT,
    stone_id VARCHAR(64) UNIQUE NOT NULL COMMENT '能量石ID',
    user_id VARCHAR(64) NOT NULL COMMENT '用户ID',
    request_id VARCHAR(64) NULL COMMENT '触发的请求ID',
    threshold_index INT NOT NULL COMMENT '达成阈值的序号(1=30秒,2=60秒...)',
    total_duration_ms BIGINT NOT NULL COMMENT '累计说话时长（毫秒）',
    duration_increment_ms INT NOT NULL COMMENT '本次新增时长（毫秒）',
    appear_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT '出现时间',
    clicked_at TIMESTAMP NULL COMMENT '点击时间',
    auto_removed BOOLEAN DEFAULT FALSE COMMENT '是否自动消失',
    metadata JSON NULL COMMENT '扩展信息',
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
    UNIQUE KEY uk_conversation_stone_threshold (user_id, threshold_index),
    INDEX idx_conversation_stone_user_threshold (user_id, threshold_index)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
COMMENT='对话能量石表'
"""

# 用户赠送对话能量石表（活动/回馈等，展示时与 conversation_energy_stones 数量相加）
CREATE_USER_BONUS_ENERGY_STONES_TABLE = """
CREATE TABLE IF NOT EXISTS user_bonus_energy_stones (
    user_id VARCHAR(64) PRIMARY KEY COMMENT '用户ID(关联 users.user_id)',
    bonus_conversation_count INT NOT NULL DEFAULT 0 COMMENT '赠送的对话能量石个数',
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间'
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
COMMENT='用户赠送对话能量石数量表'
"""


async def init_tables():
    """初始化数据库表"""
    
    try:
        # 创建数据库连接
        conn = await aiomysql.connect(
            host=settings.MYSQL_HOST,
            port=settings.MYSQL_PORT,
            user=settings.MYSQL_USER,
            password=settings.MYSQL_PASSWORD,
            db=settings.MYSQL_DATABASE,
            charset='utf8mb4'
        )
        
        async with conn.cursor() as cursor:
            # 创建用户表
            await cursor.execute(CREATE_USERS_TABLE)
            
            # 创建消息表
            await cursor.execute(CREATE_MESSAGES_TABLE)
            
            # 创建总结表
            await cursor.execute(CREATE_SUMMARIES_TABLE)

            # 创建LLM提示词记录表
            await cursor.execute(CREATE_LLM_PROMPTS_TABLE)

            # 创建任务类型表
            await cursor.execute(CREATE_MISSION_TYPES_TABLE)

            # 创建任务能量石表
            await cursor.execute(CREATE_MISSION_ENERGY_STONES_TABLE)

            # 创建对话能量石表
            await cursor.execute(CREATE_CONVERSATION_ENERGY_STONES_TABLE)

            # 创建用户赠送对话能量石表
            await cursor.execute(CREATE_USER_BONUS_ENERGY_STONES_TABLE)

            # 写入默认任务类型
            await cursor.execute(
                """
                INSERT INTO mission_types (type_key, display_name, description, is_repeatable)
                VALUES (%s, %s, %s, %s)
                ON DUPLICATE KEY UPDATE
                    display_name = VALUES(display_name),
                    description = VALUES(description),
                    is_repeatable = VALUES(is_repeatable)
                """,
                (
                    "first_dialog",
                    "首次对话奖励",
                    "首次进入对话时发放的任务能量石",
                    False,
                ),
            )
            
            await conn.commit()
        
        conn.close()
        
        return True
        
    except Exception as e:
        import traceback
        traceback.print_exc()
        return False


async def check_tables():
    """检查表是否存在"""
    try:
        conn = await aiomysql.connect(
            host=settings.MYSQL_HOST,
            port=settings.MYSQL_PORT,
            user=settings.MYSQL_USER,
            password=settings.MYSQL_PASSWORD,
            db=settings.MYSQL_DATABASE,
            charset='utf8mb4'
        )
        
        async with conn.cursor() as cursor:
            await cursor.execute("SHOW TABLES")
            tables = await cursor.fetchall()
        
        conn.close()
        
    except Exception as e:
        pass


if __name__ == "__main__":
    # 确认是否继续
    if len(sys.argv) > 1 and sys.argv[1] == "--force":
        confirm = "y"
    else:
        confirm = input("\n是否继续初始化数据库表? (y/n): ")
    
    if confirm.lower() != 'y':
        sys.exit(0)
    
    # 执行初始化
    success = asyncio.run(init_tables())
    
    if success:
        # 检查创建的表
        asyncio.run(check_tables())
    else:
        sys.exit(1)