FinDataHub/Scripts/bin/backfill_ai.py

#!/usr/bin/env python3
"""
历史快讯 AI 回填脚本 — 对 importance=0 的历史快讯做完整 AI 分析

分阶段策略：
  阶段一（评分）：对所有 importance=0 的快讯做 AI 评分，更新 importance 字段
  阶段二（增强）：对 importance>=2 的快讯做内容增强 + 摘要 + 标签

同时更新分表和 news_flash 主表，保持数据一致。

用法:
  python3 backfill_ai.py                      # 全流程：先评分后增强
  python3 backfill_ai.py --phase score        # 仅评分
  python3 backfill_ai.py --phase enrich       # 仅增强（评分已完成）
  python3 backfill_ai.py --table news_sina    # 仅处理指定表
  python3 backfill_ai.py --limit 100          # 限制每表处理数量
  python3 backfill_ai.py --batch-size 10      # 批次大小
  python3 backfill_ai.py --dry-run            # 只查询不执行 AI
"""
import asyncio
import json
import os
import sys
import time
from pathlib import Path

# ── 路径设置 + .env 加载 ─────────────────────────────────
from common import setup_kuaixun_paths, load_env, KXS_HOME, MYSQL_CFG
setup_kuaixun_paths()
load_env(f"{KXS_HOME}/.env")

import pymysql
from extend.scorer import batch_score_items
from extend.summarizer import summarize_items
from extend.tagger import batch_tag_items
from extend.content_enricher import batch_enrich_items

# 需要回填的分表（news_hotlist 已全部评分，跳过）
SOURCE_TABLES = [
    "news_sina", "news_cls", "news_wallstreetcn", "news_36kr",
    "news_gelonghui", "news_gnews", "news_newsapi",
]

BATCH_SIZE = 20
MIN_IMPORTANCE_FOR_ENRICH = 2  # 增强阶段的最小重要性阈值


# ═══════════════════════════════════════════════════════════
# 数据库操作
# ═══════════════════════════════════════════════════════════

def get_unscored(table: str, limit: int) -> list:
    """查询 importance=0 的快讯（待评分）"""
    conn = pymysql.connect(**MYSQL_CFG)
    cur = conn.cursor(pymysql.cursors.DictCursor)
    cur.execute(f"""
        SELECT id, source_news_id, title, content, url
        FROM {table}
        WHERE importance = 0
          AND title IS NOT NULL AND title != ''
        ORDER BY published_at DESC
        LIMIT %s
    """, (limit,))
    rows = cur.fetchall()
    cur.close()
    conn.close()
    return rows


def get_need_enrich(table: str, limit: int, min_imp: int, force: bool = False) -> list:
    """查询 importance>=阈值 的快讯（待增强+摘要+标签）

    force=False 时只查 summary 为空的；force=True 时查所有 importance>=阈值的
    """
    conn = pymysql.connect(**MYSQL_CFG)
    cur = conn.cursor(pymysql.cursors.DictCursor)
    if force:
        cur.execute(f"""
            SELECT id, source_news_id, title, content, url
            FROM {table}
            WHERE importance >= %s
              AND title IS NOT NULL AND title != ''
            ORDER BY importance DESC, published_at DESC
            LIMIT %s
        """, (min_imp, limit))
    else:
        cur.execute(f"""
            SELECT id, source_news_id, title, content, url
            FROM {table}
            WHERE importance >= %s
              AND (summary IS NULL OR summary = '')
              AND title IS NOT NULL AND title != ''
            ORDER BY importance DESC, published_at DESC
            LIMIT %s
        """, (min_imp, limit))
    rows = cur.fetchall()
    cur.close()
    conn.close()
    return rows


def get_need_tags(table: str, limit: int, min_imp: int) -> list:
    """查询 importance>=阈值 且 tags 为空的快讯（待打标签）"""
    conn = pymysql.connect(**MYSQL_CFG)
    cur = conn.cursor(pymysql.cursors.DictCursor)
    cur.execute(f"""
        SELECT id, source_news_id, title, content, url
        FROM {table}
        WHERE importance >= %s
          AND (tags IS NULL OR tags = '' OR tags = '[]')
          AND title IS NOT NULL AND title != ''
        ORDER BY importance DESC, published_at DESC
        LIMIT %s
    """, (min_imp, limit))
    rows = cur.fetchall()
    cur.close()
    conn.close()
    return rows


def update_ai_fields(table: str, items: list, fields: list):
    """批量更新 AI 字段到分表 + news_flash 主表

    Args:
        table: 分表名
        items: 快讯列表（每项含 id, source_news_id 及 AI 字段）
        fields: 要更新的字段名列表，如 ['importance'], ['summary','tags']
    """
    if not items:
        return
    conn = pymysql.connect(**MYSQL_CFG)
    cur = conn.cursor()

    for item in items:
        nid = item.get("source_news_id")
        rid = item.get("id")
        if not nid or not rid:
            continue

        # 构建 SET 子句
        set_parts = []
        params = []
        for f in fields:
            if f == "tags":
                val = json.dumps(item.get("tags", []), ensure_ascii=False)
            elif f == "importance":
                val = item.get("importance", 0)
            else:
                val = item.get(f, "")
            set_parts.append(f"{f} = %s")
            params.append(val)

        set_clause = ", ".join(set_parts)

        # 更新分表
        try:
            cur.execute(
                f"UPDATE {table} SET {set_clause} WHERE id = %s",
                (*params, rid),
            )
        except Exception:
            pass

        # 更新 news_flash 主表（用 source_news_id 关联）
        try:
            cur.execute(
                f"UPDATE news_flash SET {set_clause} WHERE source_news_id = %s",
                (*params, nid),
            )
        except Exception:
            pass

    conn.commit()
    cur.close()
    conn.close()


def count_unscored(table: str) -> int:
    conn = pymysql.connect(**MYSQL_CFG)
    cur = conn.cursor()
    cur.execute(f"SELECT COUNT(*) FROM {table} WHERE importance = 0 AND title IS NOT NULL AND title != ''")
    cnt = cur.fetchone()[0]
    cur.close()
    conn.close()
    return cnt


def count_need_enrich(table: str, min_imp: int, force: bool = False) -> int:
    conn = pymysql.connect(**MYSQL_CFG)
    cur = conn.cursor()
    if force:
        cur.execute(f"SELECT COUNT(*) FROM {table} WHERE importance >= %s AND title IS NOT NULL AND title!=''", (min_imp,))
    else:
        cur.execute(f"SELECT COUNT(*) FROM {table} WHERE importance >= %s AND (summary IS NULL OR summary='') AND title IS NOT NULL AND title!=''", (min_imp,))
    cnt = cur.fetchone()[0]
    cur.close()
    conn.close()
    return cnt


# ═══════════════════════════════════════════════════════════
# 阶段一：评分
# ═══════════════════════════════════════════════════════════

async def phase_score(table: str, max_items: int, batch_size: int) -> dict:
    """对指定表的 importance=0 快讯做 AI 评分"""
    total = count_unscored(table)
    if max_items > 0:
        total = min(total, max_items)
    print(f"\n{'='*50}")
    print(f"阶段一·评分 | 表: {table} | 待处理: {total} 条")
    print(f"{'='*50}", flush=True)

    if total == 0:
        print("  ✅ 无待评分数据", flush=True)
        return {"table": table, "scored": 0, "errors": 0}

    processed = 0
    scored = 0
    errors = 0
    start = time.time()

    while processed < total:
        batch = get_unscored(table, batch_size)
        if not batch:
            break

        items = [
            {"id": r["id"], "source_news_id": r["source_news_id"],
             "title": r["title"], "content": r.get("content") or "", "url": r.get("url") or ""}
            for r in batch
        ]

        try:
            results = await batch_score_items(items)
            scored_count = sum(1 for r in results if r.get("importance", 0) > 0)
            scored += scored_count
            update_ai_fields(table, results, ["importance"])
            processed += len(batch)
            print(f"  批次: {len(batch)} 条, 评分成功: {scored_count}, 累计: {processed}/{total}", flush=True)
        except Exception as e:
            errors += len(batch)
            print(f"  ⚠️ 批次失败: {e}", flush=True)

        await asyncio.sleep(0.3)

    elapsed = time.time() - start
    print(f"  📊 {table}: 评分 {processed} 条, 成功 {scored}, 失败 {errors}, 耗时 {elapsed:.1f}s", flush=True)
    return {"table": table, "scored": scored, "errors": errors, "elapsed": elapsed}


# ═══════════════════════════════════════════════════════════
# 阶段二：内容增强 + 摘要 + 标签
# ═══════════════════════════════════════════════════════════

async def phase_enrich(table: str, max_items: int, batch_size: int, min_imp: int, force: bool = False) -> dict:
    """对指定表 importance>=阈值的快讯做内容增强 + 摘要 + 标签"""
    total = count_need_enrich(table, min_imp, force)
    if max_items > 0:
        total = min(total, max_items)
    print(f"\n{'='*50}")
    print(f"阶段二·增强+摘要+标签 | 表: {table} | 待处理: {total} 条 | 最低重要性: {min_imp} | 强制重生成: {force}")
    print(f"{'='*50}", flush=True)

    if total == 0:
        print("  ✅ 无待增强数据", flush=True)
        return {"table": table, "enriched": 0, "errors": 0}

    processed = 0
    enriched = 0
    summarized = 0
    tagged = 0
    errors = 0
    start = time.time()

    while processed < total:
        batch = get_need_enrich(table, batch_size, min_imp, force)
        if not batch:
            break

        items = [
            {"id": r["id"], "source_news_id": r["source_news_id"],
             "title": r["title"], "content": r.get("content") or "", "url": r.get("url") or ""}
            for r in batch
        ]

        try:
            # 1. 内容增强（抓原文 + SearXNG 搜索）
            try:
                items = await batch_enrich_items(items, concurrency=3)
                enriched += sum(1 for i in items if i.get("enriched_content"))
            except Exception as ee:
                print(f"  ⚠️ 内容增强异常（降级）: {ee}", flush=True)

            # 2. AI 摘要（用增强后的全文）
            items = await summarize_items(items)
            summarized += sum(1 for i in items if i.get("summary"))

            # 3. AI 打标签（复用已增强内容）
            items = await batch_tag_items(items)
            tagged += sum(1 for i in items if i.get("tags"))

            # 更新 summary + tags 到数据库
            update_ai_fields(table, items, ["summary", "tags"])

            processed += len(batch)
            print(f"  批次: {len(batch)} 条, 增强: {enriched}, 摘要: {summarized}, 标签: {tagged}, 累计: {processed}/{total}", flush=True)
        except Exception as e:
            errors += len(batch)
            print(f"  ⚠️ 批次失败: {e}", flush=True)

        await asyncio.sleep(0.5)

    elapsed = time.time() - start
    print(f"  📊 {table}: 处理 {processed} 条, 增强 {enriched}, 摘要 {summarized}, 标签 {tagged}, 失败 {errors}, 耗时 {elapsed:.1f}s", flush=True)
    return {"table": table, "enriched": enriched, "summarized": summarized, "tagged": tagged, "errors": errors, "elapsed": elapsed}


# ═══════════════════════════════════════════════════════════
# 主入口
# ═══════════════════════════════════════════════════════════

async def main():
    import argparse
    parser = argparse.ArgumentParser(description="历史快讯 AI 回填")
    parser.add_argument("--phase", choices=["score", "enrich", "all"], default="all", help="执行阶段（默认 all）")
    parser.add_argument("--table", default="", help="仅处理指定表（默认全部）")
    parser.add_argument("--limit", type=int, default=0, help="限制每表处理数量（0=不限）")
    parser.add_argument("--batch-size", type=int, default=BATCH_SIZE, help="批次大小")
    parser.add_argument("--min-importance", type=int, default=MIN_IMPORTANCE_FOR_ENRICH, help="增强阶段最低重要性")
    parser.add_argument("--force", action="store_true", help="强制重新生成摘要和标签（即使已有）")
    parser.add_argument("--dry-run", action="store_true", help="只统计不执行")
    args = parser.parse_args()

    tables = [args.table] if args.table else SOURCE_TABLES

    scorer_model = os.getenv("AI_SCORER_MODEL") or os.getenv("AI_MODEL", "qwen2.5:0.5b")
    summarizer_model = os.getenv("AI_SUMMARIZER_MODEL") or os.getenv("AI_MODEL", "qwen2.5:0.5b")
    tagger_model = os.getenv("AI_TAGGER_MODEL") or os.getenv("AI_MODEL", "qwen2.5:0.5b")

    print(f"{'='*60}")
    print(f"  历史快讯 AI 回填")
    print(f"  阶段: {args.phase} | 表: {tables}")
    print(f"  评分模型: {scorer_model}")
    print(f"  摘要模型: {summarizer_model}")
    print(f"  标签模型: {tagger_model}")
    print(f"  批次大小: {args.batch_size} | 增强阈值: {args.min_importance}")
    print(f"{'='*60}", flush=True)

    if args.dry_run:
        print("\n[Dry Run] 统计待处理数据:", flush=True)
        for t in tables:
            unscored = count_unscored(t)
            need_enrich = count_need_enrich(t, args.min_importance, args.force)
            print(f"  {t}: 待评分 {unscored} 条, 待增强 {need_enrich} 条", flush=True)
        return

    grand = {"scored": 0, "enriched": 0, "summarized": 0, "tagged": 0, "errors": 0}
    start = time.time()

    for table in tables:
        # 阶段一：评分
        if args.phase in ("score", "all"):
            r = await phase_score(table, args.limit, args.batch_size)
            grand["scored"] += r.get("scored", 0)
            grand["errors"] += r.get("errors", 0)

        # 阶段二：增强 + 摘要 + 标签
        if args.phase in ("enrich", "all"):
            r = await phase_enrich(table, args.limit, args.batch_size, args.min_importance, args.force)
            grand["enriched"] += r.get("enriched", 0)
            grand["summarized"] += r.get("summarized", 0)
            grand["tagged"] += r.get("tagged", 0)
            grand["errors"] += r.get("errors", 0)

    elapsed = time.time() - start
    print(f"\n{'='*60}")
    print(f"  全部完成 | 耗时 {elapsed:.1f}s")
    print(f"  评分: {grand['scored']} | 增强: {grand['enriched']} | 摘要: {grand['summarized']} | 标签: {grand['tagged']} | 错误: {grand['errors']}")
    print(f"{'='*60}", flush=True)


if __name__ == "__main__":
    asyncio.run(main())