FinDataHub/Scripts/bin/daily_summary.py

#!/usr/bin/env python3
"""
daily_summary.py — 快讯当天总结 + 热度分析
每小时生成一次当天总结，包含：
- 今日快讯总数、来源分布、热度分布
- 高热度快讯 Top10
- AI 生成的当日总结
"""
import json, os, sys, time
from datetime import datetime, timedelta
from pathlib import Path

# 引入 kuaixun-service
from common import setup_kuaixun_paths, KXS_HOME, BIN_DIR, MYSQL_CFG, PidManager, OLLAMA_URL
setup_kuaixun_paths()

import pymysql
import asyncio

# 连接池
from db_pool import get_conn, put_conn

try:
    from paths import pid_file as _pid_file
    PID_FILE = _pid_file("summary")
except ImportError:
    PID_FILE = "/tmp/findatahub-summary.pid"
UPDATE_INTERVAL = 3600  # 3600秒 = 1小时


# ── AI 摘要模块 ──
try:
    from extend.summarizer import summarize
    _AI_AVAILABLE = True
except ImportError:
    _AI_AVAILABLE = False

# 当日总结专用模型（优先级最高）
AI_SUMMARY_MODEL = os.getenv("AI_SUMMARY_MODEL") or os.getenv("AI_MODEL", "qwen2.5:0.5b")

SUMMARY_PROMPT = """你是一个专业的财经新闻分析师。请根据以下今天的财经快讯，生成一份简明扼要的当日总结：

要求：
1. 分3-5个要点，每个要点一句话
2. 优先突出重要性≥4的快讯
3. 涵盖不同来源和主题
4. 语言简洁，适合快速阅读

快讯列表：
{news_text}

当日总结："""


# ── MySQL 表管理 ──

def ensure_table():
    """确保 daily_summary 表存在"""
    conn = get_conn()
    cur = conn.cursor()
    cur.execute("""
        CREATE TABLE IF NOT EXISTS daily_summary (
            id INT AUTO_INCREMENT PRIMARY KEY,
            date DATE NOT NULL,
            summary TEXT,
            stats JSON,
            hot_topics JSON,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
            UNIQUE KEY uk_date (date)
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
    """)
    conn.commit()
    cur.close()
    put_conn(conn)


# ── 数据查询 ──

def get_today_stats():
    """获取今日统计"""
    conn = get_conn()
    cur = conn.cursor()
    today = datetime.now().strftime("%Y-%m-%d")

    # 总数
    cur.execute(
        "SELECT COUNT(*) FROM news_flash WHERE published_at >= UNIX_TIMESTAMP(%s) * 1000 AND published_at < UNIX_TIMESTAMP(DATE_ADD(%s, INTERVAL 1 DAY)) * 1000",
        (today, today),
    )
    total = cur.fetchone()[0]

    # 来源分布
    cur.execute("""
        SELECT s.name AS source_name, COUNT(*) AS count
        FROM news_flash f LEFT JOIN news_sources s ON f.source_id = s.id
        WHERE f.published_at >= UNIX_TIMESTAMP(%s) * 1000 AND f.published_at < UNIX_TIMESTAMP(DATE_ADD(%s, INTERVAL 1 DAY)) * 1000
        GROUP BY s.name ORDER BY count DESC
    """, (today, today))
    sources = [{"source": r[0], "count": r[1]} for r in cur.fetchall()]

    # 热度分布（importance 0-5）
    cur.execute("""
        SELECT importance, COUNT(*) AS count
        FROM news_flash
        WHERE published_at >= UNIX_TIMESTAMP(%s) * 1000 AND published_at < UNIX_TIMESTAMP(DATE_ADD(%s, INTERVAL 1 DAY)) * 1000 AND importance >= 0
        GROUP BY importance ORDER BY importance DESC
    """, (today, today))
    importance_dist = {str(r[0]): r[1] for r in cur.fetchall()}

    # 高热度 Top10
    cur.execute("""
        SELECT f.title, f.summary, f.content, s.name AS source_name, f.importance, f.published_at
        FROM news_flash f LEFT JOIN news_sources s ON f.source_id = s.id
        WHERE f.published_at >= UNIX_TIMESTAMP(%s) * 1000 AND f.published_at < UNIX_TIMESTAMP(DATE_ADD(%s, INTERVAL 1 DAY)) * 1000
        ORDER BY f.importance DESC, f.published_at DESC LIMIT 10
    """, (today, today))
    hot_items = []
    for r in cur.fetchall():
        hot_items.append({
            "title": r[0], "summary": r[1] or "", "content": r[2] or "",
            "source": r[3], "importance": r[4], "published_at": r[5],
        })

    # 频率统计（每小时条数）
    cur.execute("""
        SELECT HOUR(FROM_UNIXTIME(published_at/1000)) AS hour, COUNT(*) AS count
        FROM news_flash
        WHERE published_at >= UNIX_TIMESTAMP(%s) * 1000 AND published_at < UNIX_TIMESTAMP(DATE_ADD(%s, INTERVAL 1 DAY)) * 1000
        GROUP BY hour ORDER BY hour
    """, (today, today))
    hourly_freq = [{"hour": r[0], "count": r[1]} for r in cur.fetchall()]

    put_conn(conn)
    return {
        "total": total,
        "sources": sources,
        "importance_dist": importance_dist,
        "hot_top10": hot_items,
        "hourly_freq": hourly_freq,
    }


# ── AI 总结生成 ──

async def generate_summary(stats):
    """根据今日数据生成 AI 总结（直接调用 Ollama API，使用 AI_SUMMARY_MODEL）"""
    hot_items = stats.get("hot_top10", [])
    if not hot_items:
        return "今日暂无高热度快讯"

    news_text = ""
    for i, item in enumerate(hot_items[:8], 1):
        importance_label = {5: "重大", 4: "重要", 3: "一般"}.get(item["importance"], "")
        news_text += f"{i}. [{importance_label}]{item['source']}: {item['title']}\n"

    prompt = SUMMARY_PROMPT.format(news_text=news_text)

    # 直接调用 Ollama API（使用 AI_SUMMARY_MODEL，qwen3 推理模型需关闭 think）
    try:
        import httpx
        async with httpx.AsyncClient(timeout=120) as client:
            resp = await client.post(
                OLLAMA_URL,
                json={
                    "model": AI_SUMMARY_MODEL,
                    "messages": [{"role": "user", "content": prompt}],
                    "stream": False,
                    "think": False,
                    "options": {"num_predict": 500, "temperature": 0.3}
                }
            )
            return resp.json().get("message", {}).get("content", "").strip()[:800]
    except Exception as e:
        print(f"  ⚠️ 总结生成失败: {e}", flush=True)
        return f"总结生成失败: {str(e)}"


# ── 写入总结 ──

def save_summary(date_str, summary, stats, hot_topics):
    """保存/更新当日总结"""
    conn = get_conn()
    cur = conn.cursor()
    cur.execute("""
        INSERT INTO daily_summary (date, summary, stats, hot_topics)
        VALUES (%s, %s, %s, %s)
        ON DUPLICATE KEY UPDATE summary=%s, stats=%s, hot_topics=%s, updated_at=NOW()
    """, (date_str, summary, json.dumps(stats, ensure_ascii=False),
          json.dumps(hot_topics, ensure_ascii=False),
          summary, json.dumps(stats, ensure_ascii=False),
          json.dumps(hot_topics, ensure_ascii=False)))
    conn.commit()
    put_conn(conn)


# ── 主流程 ──

async def run_summary():
    """执行一次总结生成"""
    ensure_table()
    today = datetime.now().strftime("%Y-%m-%d")
    print(f"[{datetime.now().isoformat()}] 生成 {today} 总结...", flush=True)

    stats = get_today_stats()
    summary = await generate_summary(stats)
    hot_topics = stats.get("hot_top10", [])[:5]

    save_summary(today, summary, stats, hot_topics)
    print(f"  ✅ 总结生成完成（共 {stats['total']} 条快讯）", flush=True)

    return {"date": today, "total": stats["total"], "summary": summary[:100] + "..."}


# ── 定时任务 ──

_pid_mgr = PidManager(PID_FILE)

def get_status():
    return _pid_mgr.status()

def write_pid():
    _pid_mgr.write()

def stop():
    return _pid_mgr.stop()


# ── 查询接口 ──

def query_summary(date_str=None):
    """查询指定日期的总结（默认今天）"""
    if not date_str:
        date_str = datetime.now().strftime("%Y-%m-%d")

    conn = get_conn()
    cur = conn.cursor()
    cur.execute("SELECT summary, stats, hot_topics, updated_at FROM daily_summary WHERE date = %s", (date_str,))
    row = cur.fetchone()
    put_conn(conn)

    if not row:
        return {"date": date_str, "summary": "", "stats": {}, "hot_topics": [], "updated_at": None}

    return {
        "date": date_str,
        "summary": row[0],
        "stats": json.loads(row[1]) if row[1] else {},
        "hot_topics": json.loads(row[2]) if row[2] else [],
        "updated_at": str(row[3]),
    }


# ── 入口 ──

def main():
    import argparse
    parser = argparse.ArgumentParser(description="快讯当日总结")
    parser.add_argument("mode", choices=["start", "stop", "status", "once", "query"], nargs="?", default="status")
    parser.add_argument("--interval", type=int, default=UPDATE_INTERVAL)
    parser.add_argument("--date", help="查询日期（YYYY-MM-DD）")
    args = parser.parse_args()

    if args.mode == "status":
        print(json.dumps(get_status()))
        return

    if args.mode == "stop":
        ok = stop()
        print(json.dumps({"stopped": ok}))
        return

    if args.mode == "query":
        result = query_summary(args.date)
        print(json.dumps(result, ensure_ascii=False, default=str))
        return

    if args.mode == "once":
        result = asyncio.run(run_summary())
        print(json.dumps(result, ensure_ascii=False))
        return

    if args.mode == "start":
        s = get_status()
        if s["running"]:
            print(json.dumps({"error": "already_running", "pid": s["pid"]}))
            return

        write_pid()
        print(json.dumps({"status": "started", "pid": os.getpid()}))
        sys.stdout.flush()

        # 立即执行一次
        asyncio.run(run_summary())

        # 定时执行
        while True:
            now = datetime.now()
            next_hour = (now + timedelta(hours=1)).replace(minute=0, second=0, microsecond=0)
            sleep_sec = (next_hour - now).total_seconds()
            time.sleep(sleep_sec)
            try:
                asyncio.run(run_summary())
            except Exception as e:
                print(f"  ❌ 总结异常: {e}", flush=True)


if __name__ == "__main__":
    main()