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

#!/usr/bin/env python3
"""
导入爬虫数据到新的数据库表

支持新格式（{.txt）：
  每个 model 有 prices 字段，结构为：
    { "label": { "输入": { price, unit, price_original? }, "输出": {...} } }
  或简单格式：
    { "label": { price, unit, price_original? } }

- 导入模型数据到 models_new 表
- 导入价格数据到 model_price_new 表
"""
import json
import sys
from pathlib import Path
from datetime import datetime
from decimal import Decimal

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

from app.database import SessionLocal
from app.models.model import ModelNew, ModelPriceNew, infer_categories


def _compute_discount_label(discount_rate: float) -> str | None:
    """将折扣率转换为折扣标签，如 0.1 -> '1折'，1.0 -> None"""
    if discount_rate >= 1.0 or discount_rate <= 0:
        return None
    tenths = round(discount_rate * 10)
    if tenths >= 10:
        return None
    return f"{tenths}折"


def _parse_unit_multiplier(unit: str) -> int:
    """根据单位字符串推断 display_multiplier"""
    u = unit.lower()
    if '百万' in u or 'million' in u or 'tokens' in u:
        return 1_000_000
    if '千' in u or '1k' in u:
        return 1_000
    return 1


def _upsert_price_row(db, model_code: str, label: str, scraped_at: datetime,
                      input_orig: float, output_orig: float,
                      input_disc: float, output_disc: float,
                      unit: str, currency: str = 'CNY',
                      tier_min=None, tier_max=None, tier_unit: str | None = None,
                      source_url: str | None = None):
    """插入或更新一条价格记录（同一 model_code+label 只保留一条 is_active=true）"""
    # 计算折扣率
    if input_orig and input_orig > 0 and abs(input_disc - input_orig) > 1e-9:
        discount_rate = round(input_disc / input_orig, 4)
    elif output_orig and output_orig > 0 and abs(output_disc - output_orig) > 1e-9:
        discount_rate = round(output_disc / output_orig, 4)
    else:
        discount_rate = 1.0

    discount_label = _compute_discount_label(discount_rate)
    multiplier = _parse_unit_multiplier(unit)

    # 将旧的 is_active 记录标记为 false
    db.query(ModelPriceNew).filter(
        ModelPriceNew.model_code == model_code,
        ModelPriceNew.label == label,
        ModelPriceNew.is_active == True,
    ).update({"is_active": False}, synchronize_session=False)

    price = ModelPriceNew(
        model_code=model_code,
        label=label,
        tier_min=tier_min,
        tier_max=tier_max,
        tier_unit=tier_unit,
        input_price_original=input_orig,
        output_price_original=output_orig,
        discount_rate=discount_rate,
        discount_label=discount_label,
        input_price_discounted=input_disc,
        output_price_discounted=output_disc,
        currency=currency,
        unit=unit,
        display_multiplier=multiplier,
        source_url=source_url,
        crawled_at=scraped_at,
        is_active=True,
    )
    db.add(price)
    return price


def _parse_prices_new_format(model_data: dict, db, scraped_at: datetime):
    """
    解析新格式 prices 字段：
    {
      "label": {
        "输入": { "price": 2.5, "unit": "元/每百万tokens", "price_original": 10.0 },
        "输出": { ... },
        ...
      }
    }
    或简单格式（无子键）：
    {
      "label": { "price": 2.0, "unit": "元/每万字符" }
    }
    """
    model_code = model_data.get('model_name')
    prices_raw = model_data.get('prices', {})
    source_url = model_data.get('url')
    rows = []

    for label, label_data in prices_raw.items():
        if not isinstance(label_data, dict):
            continue

        # 判断是否为简单格式（直接有 price 字段）
        if 'price' in label_data:
            # 简单格式：整个 label 只有一个价格
            price_val = float(label_data.get('price', 0) or 0)
            price_orig = label_data.get('price_original')
            orig_val = float(price_orig) if price_orig is not None else price_val
            unit = label_data.get('unit', '')
            currency = label_data.get('currency', 'CNY')
            rows.append((label, orig_val, 0.0, price_val, 0.0, unit, currency))
        else:
            # 复合格式：label_data 的 key 是子项名（输入/输出/显式缓存创建/...）
            # 找输入和输出价格
            input_entry = label_data.get('输入') or label_data.get('input') or {}
            output_entry = label_data.get('输出') or label_data.get('output') or {}

            if not input_entry and not output_entry:
                # 没有输入/输出，把所有子项各自作为独立 label
                for sub_key, sub_val in label_data.items():
                    if not isinstance(sub_val, dict) or 'price' not in sub_val:
                        continue
                    sub_label = f"{label}_{sub_key}"
                    price_val = float(sub_val.get('price', 0) or 0)
                    price_orig = sub_val.get('price_original')
                    orig_val = float(price_orig) if price_orig is not None else price_val
                    unit = sub_val.get('unit', '')
                    currency = sub_val.get('currency', 'CNY')
                    rows.append((sub_label, orig_val, 0.0, price_val, 0.0, unit, currency))
                continue

            input_price = float(input_entry.get('price', 0) or 0)
            input_orig_raw = input_entry.get('price_original')
            input_orig = float(input_orig_raw) if input_orig_raw is not None else input_price

            output_price = float(output_entry.get('price', 0) or 0)
            output_orig_raw = output_entry.get('price_original')
            output_orig = float(output_orig_raw) if output_orig_raw is not None else output_price

            unit = (input_entry.get('unit') or output_entry.get('unit') or '')
            currency = (input_entry.get('currency') or output_entry.get('currency') or 'CNY')
            rows.append((label, input_orig, output_orig, input_price, output_price, unit, currency))

    price_objs = []
    for label, in_orig, out_orig, in_disc, out_disc, unit, currency in rows:
        p = _upsert_price_row(
            db, model_code, label, scraped_at,
            in_orig, out_orig, in_disc, out_disc,
            unit, currency, source_url=source_url
        )
        price_objs.append(p)
    return price_objs


def import_crawl_data(crawl_file: str):
    with open(crawl_file, 'r', encoding='utf-8') as f:
        data = json.load(f)

    db = SessionLocal()
    try:
        models_list = data if isinstance(data, list) else data.get('models', [])
        version = data.get('version') if isinstance(data, dict) else None

        print(f"开始导入，共 {len(models_list)} 个模型，版本号: {version}")

        model_count = 0
        price_count = 0

        for model_data in models_list:
            model_code = model_data.get('model_name')
            if not model_code:
                continue

            model_info = model_data.get('model_info', {})
            input_mods = model_info.get('input_modalities') or []
            output_mods = model_info.get('output_modalities') or []
            display_tags = model_info.get('display_tags') or []
            cats = infer_categories(input_mods, output_mods, model_code, display_tags)

            scraped_at_str = model_data.get('scraped_at')
            scraped_at = datetime.fromisoformat(scraped_at_str.replace('Z', '+00:00')) if scraped_at_str else datetime.now()

            existing = db.query(ModelNew).filter(ModelNew.model_code == model_code).first()
            if existing:
                # 更新爬虫字段（不覆盖管理员配置字段）
                existing.description = model_info.get('description') or existing.description
                existing.display_tags = display_tags or existing.display_tags
                existing.input_modalities = input_mods or existing.input_modalities
                existing.output_modalities = output_mods or existing.output_modalities
                existing.features = model_info.get('features') or existing.features
                existing.rate_limits = model_data.get('rate_limits') or existing.rate_limits
                existing.tool_call_prices = model_data.get('tool_prices') or existing.tool_call_prices
                existing.raw_prices = model_data.get('prices') or existing.raw_prices
                existing.source_url = model_data.get('url') or existing.source_url
                existing.crawled_at = scraped_at
                existing.categories = cats
                existing.updated_at = datetime.now()
            else:
                model = ModelNew(
                    model_code=model_code,
                    description=model_info.get('description'),
                    display_tags=display_tags,
                    input_modalities=input_mods,
                    output_modalities=output_mods,
                    features=model_info.get('features'),
                    rate_limits=model_data.get('rate_limits'),
                    tool_call_prices=model_data.get('tool_prices'),
                    raw_prices=model_data.get('prices'),
                    source_url=model_data.get('url'),
                    crawled_at=scraped_at,
                    display_name=model_code,
                    supplier='Qwen',
                    categories=cats,
                    keywords=','.join(display_tags),
                    is_show_enabled=True,
                    is_api_enabled=True,
                )
                db.add(model)
                model_count += 1

            # 导入价格
            price_objs = _parse_prices_new_format(model_data, db, scraped_at)
            price_count += len(price_objs)

        db.commit()
        print(f"完成：新增模型 {model_count} 个，价格记录 {price_count} 条")

        # 记录同步日志
        if version is not None:
            from app.models.model import CrawlerSyncLog
            log = CrawlerSyncLog(
                crawler_version=version,
                model_count=model_count,
                price_count=price_count,
                status='success',
            )
            db.add(log)
            db.commit()

    except Exception as e:
        db.rollback()
        print(f"导入失败: {e}")
        raise
    finally:
        db.close()


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="导入爬虫数据到数据库")
    parser.add_argument(
        "--file",
        type=str,
        default=str(Path(__file__).parent.parent.parent / "{.txt"),
        help="爬虫数据文件路径（默认：项目根目录的 {.txt）",
    )
    args = parser.parse_args()

    crawl_file = args.file
    if not Path(crawl_file).exists():
        print(f"文件不存在: {crawl_file}")
        sys.exit(1)

    print(f"导入数据从: {crawl_file}")
    import_crawl_data(crawl_file)
    print("导入完成！")