Maas-Model-platform/backend/scripts/init_price_data.py

"""
价格数据初始化脚本

从model_prices.json文件导入价格数据到数据库
使用model_price.id（主键）作为关联字段
"""
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from sqlalchemy import text
from app.database import SessionLocal


def init_price_data():
    """初始化模型价格数据"""
    json_path = Path(__file__).parent / "model_prices.json"
    
    with open(json_path, 'r', encoding='utf-8') as f:
        prices_data = json.load(f)
    
    db = SessionLocal()
    price_inserted = 0
    tier_inserted = 0
    model_updated = 0
    
    for item in prices_data:
        model_name = item['model_name']
        pricing_mode = item['pricing_mode']
        unit = item.get('unit', 'tokens')
        currency = item.get('currency', 'CNY')
        
        # 检查模型是否存在
        model_row = db.execute(
            text("SELECT id, price_id FROM aigcspace.models WHERE title = :title"),
            {"title": model_name}
        ).fetchone()
        
        if not model_row:
            print(f"警告: 模型 {model_name} 不存在，跳过")
            continue
        
        model_id = model_row[0]
        existing_price_id = model_row[1]
        
        if pricing_mode == 'simple':
            input_price = item['input_price']
            output_price = item['output_price']
        else:
            first_tier = item['tiers'][0]
            input_price = first_tier['input_price']
            output_price = first_tier['output_price']
        
        if existing_price_id:
            # 更新现有价格
            db.execute(
                text("""UPDATE aigcspace.model_price 
                SET input_price = :input_price, output_price = :output_price,
                    pricing_mode = :pricing_mode, unit = :unit, currency = :currency,
                    updated_at = CURRENT_TIMESTAMP
                WHERE id = :id"""),
                {"input_price": input_price, "output_price": output_price,
                 "pricing_mode": pricing_mode, "unit": unit, "currency": currency, 
                 "id": existing_price_id}
            )
            
            if pricing_mode == 'tier':
                db.execute(
                    text("DELETE FROM aigcspace.model_price_tier WHERE price_id = :price_id"),
                    {"price_id": existing_price_id}
                )
                for tier in item['tiers']:
                    db.execute(
                        text("""INSERT INTO aigcspace.model_price_tier 
                        (price_id, tier_min, tier_max, input_price, output_price)
                        VALUES (:price_id, :tier_min, :tier_max, :input_price, :output_price)"""),
                        {"price_id": existing_price_id, "tier_min": tier['tier_min'], 
                         "tier_max": tier.get('tier_max'),
                         "input_price": tier['input_price'], "output_price": tier['output_price']}
                    )
                    tier_inserted += 1
            model_updated += 1
        else:
            # 创建新价格记录，使用RETURNING获取自增ID
            result = db.execute(
                text("""INSERT INTO aigcspace.model_price 
                (input_price, output_price, pricing_mode, unit, currency)
                VALUES (:input_price, :output_price, :pricing_mode, :unit, :currency)
                RETURNING id"""),
                {"input_price": input_price, "output_price": output_price,
                 "pricing_mode": pricing_mode, "unit": unit, "currency": currency}
            )
            new_price_id = result.fetchone()[0]
            
            if pricing_mode == 'tier':
                for tier in item['tiers']:
                    db.execute(
                        text("""INSERT INTO aigcspace.model_price_tier 
                        (price_id, tier_min, tier_max, input_price, output_price)
                        VALUES (:price_id, :tier_min, :tier_max, :input_price, :output_price)"""),
                        {"price_id": new_price_id, "tier_min": tier['tier_min'], 
                         "tier_max": tier.get('tier_max'),
                         "input_price": tier['input_price'], "output_price": tier['output_price']}
                    )
                    tier_inserted += 1
            
            # 更新模型的price_id为新创建的价格记录ID
            db.execute(
                text("UPDATE aigcspace.models SET price_id = :price_id WHERE id = :model_id"),
                {"price_id": new_price_id, "model_id": model_id}
            )
            price_inserted += 1
    
    db.commit()
    db.close()
    
    print(f"价格数据导入完成:")
    print(f"  - 新增价格记录: {price_inserted} 条")
    print(f"  - 更新价格记录: {model_updated} 条")
    print(f"  - 阶梯价格记录: {tier_inserted} 条")


if __name__ == "__main__":
    init_price_data()