Neo-ZQYY/scripts/ops/ai_backfill.py

# -*- coding: utf-8 -*-
"""历史数据回填脚本。

用法：
    cd C:\\Project\\NeoZQYY
    uv run python scripts/ops/ai_backfill.py [--dry-run] [--batch-size 10] [--interval 5]

功能：
    1. 查询半年内活跃会员（消费/备注/任务变更三者并集）
    2. 分批执行 AI 调用链（每批 10 人，间隔 5 秒）
    3. 断点续跑（已完成 member_id 记录到本地文件）
    4. App8 写 member_retention_clue（DELETE + INSERT 事务）

需求: D1.1, D1.2, D1.3, D1.4, D1.5, D1.6, D1.7, D2.1, D2.3
"""
from __future__ import annotations

import argparse
import asyncio
import json
import logging
import os
import sys
from datetime import datetime, timezone
from pathlib import Path

from dotenv import load_dotenv

# ── 加载根 .env ──────────────────────────────────────
_ROOT = Path(__file__).resolve().parent.parent.parent
load_dotenv(_ROOT / ".env")

# 必需环境变量校验
_APP_DB_DSN = os.environ.get("APP_DB_DSN")
if not _APP_DB_DSN:
    sys.exit("ERROR: APP_DB_DSN 环境变量未设置，请检查根 .env 文件")

logger = logging.getLogger("ai_backfill")

# ── 常量 ─────────────────────────────────────────────
SITE_ID = 2790685415443269
CHECKPOINT_FILE = "scripts/ops/_ai_backfill_checkpoint.json"
BATCH_SIZE = 10
BATCH_INTERVAL = 5  # 秒
DATE_FROM = "2025-09-21"
DATE_TO = "2026-03-21"

# 预估每次调用消耗 token 数（用于 print_estimate）
AVG_TOKENS_PER_CALL = 2000
# 每个会员平均调用 App 数（App3+App8+App7 = 3，加上可能的 App4+App5+App6 ≈ 5）
AVG_APPS_PER_MEMBER = 5


# ── 活跃会员查询 SQL ─────────────────────────────────
# 半年内消费 ∪ 备注 ∪ 任务变更三者并集
ACTIVE_MEMBERS_SQL = """
SELECT DISTINCT member_id FROM (
    -- 消费记录
    SELECT member_id
    FROM biz.settlement_orders
    WHERE site_id = %(site_id)s
      AND member_id > 0
      AND settle_time >= %(date_from)s
      AND settle_time < %(date_to)s

    UNION

    -- 备注记录
    SELECT member_id
    FROM biz.member_notes
    WHERE site_id = %(site_id)s
      AND member_id > 0
      AND created_at >= %(date_from)s
      AND created_at < %(date_to)s

    UNION

    -- 任务变更（助教任务分配）
    SELECT member_id
    FROM biz.assistant_tasks
    WHERE site_id = %(site_id)s
      AND member_id > 0
      AND updated_at >= %(date_from)s
      AND updated_at < %(date_to)s
) AS active
ORDER BY member_id
"""


class AIBackfillRunner:
    """历史数据回填执行器。

    对半年内活跃会员分批执行 AI 调用链，支持断点续跑。
    """

    def __init__(
        self,
        dry_run: bool = False,
        batch_size: int = BATCH_SIZE,
        interval: int = BATCH_INTERVAL,
    ):
        self.dry_run = dry_run
        self.batch_size = batch_size
        self.interval = interval
        self.completed: set[int] = set()
        self._checkpoint_path = Path(CHECKPOINT_FILE)

    # ── checkpoint 读写 ──────────────────────────────

    def load_checkpoint(self) -> None:
        """从本地 JSON 文件加载已完成的 member_id 集合。

        文件不存在或格式异常时，completed 初始化为空集。
        """
        if not self._checkpoint_path.exists():
            self.completed = set()
            return
        try:
            data = json.loads(self._checkpoint_path.read_text(encoding="utf-8"))
            self.completed = set(data.get("completed_member_ids", []))
        except (json.JSONDecodeError, TypeError, KeyError):
            logger.warning("checkpoint 文件格式异常，从头开始")
            self.completed = set()

    def save_checkpoint(self) -> None:
        """保存已完成的 member_id 到本地 JSON 文件。"""
        data = {
            "completed_member_ids": sorted(self.completed),
            "last_updated": datetime.now(timezone.utc).isoformat(),
        }
        self._checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
        self._checkpoint_path.write_text(
            json.dumps(data, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )

    # ── 数据库查询 ───────────────────────────────────

    async def query_active_members(self) -> list[int]:
        """查询半年内活跃会员（消费 ∪ 备注 ∪ 任务变更三者并集）。

        使用 APP_DB_DSN 连接业务库。
        """
        import psycopg2

        conn = psycopg2.connect(_APP_DB_DSN)
        try:
            with conn.cursor() as cur:
                cur.execute(
                    ACTIVE_MEMBERS_SQL,
                    {
                        "site_id": SITE_ID,
                        "date_from": DATE_FROM,
                        "date_to": DATE_TO,
                    },
                )
                rows = cur.fetchall()
            return [row[0] for row in rows]
        finally:
            conn.close()

    # ── 单会员调用链 ─────────────────────────────────

    async def run_member(self, member_id: int) -> None:
        """对单个会员执行调用链。

        调用链：
        1. App3（维客线索）→ App8（线索整理）→ App7（客户分析）
        2. 如有助教关联：→ App4（关系分析）→ App5（话术参考）
        3. 如有备注：→ App6（备注分析）→ App8（再次整合）

        App8 写入 member_retention_clue 使用 DELETE + INSERT 事务包裹。
        验证 App5 输出包含分类字段（覆盖 P2-1 修复）。

        注意：实际 AI 调用依赖后端 AIDispatcher，此处为占位实现。
        正式使用时需要导入并调用 AIDispatcher 的相关方法。
        """
        if self.dry_run:
            logger.info("[DRY-RUN] 跳过会员 %d 的调用链执行", member_id)
            return

        # TODO: 接入后端 AIDispatcher 实现实际调用链
        # 以下为占位逻辑，标明调用顺序和关键校验点

        # 步骤 1: App3 → App8 → App7（基础链路）
        # app3_result = await dispatcher.run_app("app3_clue", member_id, SITE_ID)
        # app8_result = await dispatcher.run_app("app8_clue_consolidated", member_id, SITE_ID)
        # _write_retention_clue(member_id, app8_result)  # DELETE + INSERT 事务
        # app7_result = await dispatcher.run_app("app7_customer_analysis", member_id, SITE_ID)

        # 步骤 2: 如有助教关联 → App4 → App5
        # if has_coach_relation(member_id):
        #     app4_result = await dispatcher.run_app("app4_analysis", member_id, SITE_ID)
        #     app5_result = await dispatcher.run_app("app5_tactics", member_id, SITE_ID)
        #     # P2-1 修复验证：App5 输出必须包含分类字段
        #     if "category" not in app5_result:
        #         logger.warning("App5 输出缺少分类字段，member_id=%d", member_id)

        # 步骤 3: 如有备注 → App6 → App8（再次整合）
        # if has_notes(member_id):
        #     app6_result = await dispatcher.run_app("app6_note_analysis", member_id, SITE_ID)
        #     app8_result2 = await dispatcher.run_app("app8_clue_consolidated", member_id, SITE_ID)
        #     _write_retention_clue(member_id, app8_result2)  # DELETE + INSERT 事务

        logger.info("会员 %d 调用链执行完成（占位）", member_id)

    # ── 预估输出 ─────────────────────────────────────

    def print_estimate(self, member_count: int) -> None:
        """输出预估信息：会员数、预估调用次数、预估 Token 消耗。"""
        estimated_calls = member_count * AVG_APPS_PER_MEMBER
        estimated_tokens = estimated_calls * AVG_TOKENS_PER_CALL
        batch_count = (member_count + self.batch_size - 1) // self.batch_size
        estimated_time_sec = batch_count * self.interval

        print("=" * 60)
        print("回填预估信息")
        print("=" * 60)
        print(f"  门店 ID:        {SITE_ID}")
        print(f"  日期范围:       {DATE_FROM} ~ {DATE_TO}")
        print(f"  活跃会员数:     {member_count}")
        print(f"  已完成:         {len(self.completed)}")
        print(f"  待处理:         {member_count - len(self.completed)}")
        print(f"  预估调用次数:   {estimated_calls}")
        print(f"  预估 Token:     {estimated_tokens:,}")
        print(f"  批次大小:       {self.batch_size}")
        print(f"  批间间隔:       {self.interval}s")
        print(f"  预估批次数:     {batch_count}")
        print(f"  预估耗时:       ~{estimated_time_sec}s")
        if self.dry_run:
            print("  模式:           DRY-RUN（不执行实际调用）")
        print("=" * 60)

    # ── 主入口 ───────────────────────────────────────

    async def run(self) -> None:
        """主入口：加载断点 → 查询会员 → 分批执行 → 保存断点。

        错误处理：
        - 单个会员失败：记录日志，跳过，继续下一个；不写入 checkpoint
        - 数据库连接中断：保存当前 checkpoint，退出
        - 预算超限：保存 checkpoint，输出提示，退出
        """
        self.load_checkpoint()

        try:
            members = await self.query_active_members()
        except Exception as exc:
            logger.error("查询活跃会员失败: %s", exc)
            self.save_checkpoint()
            return

        # 过滤已完成的会员
        pending = [m for m in members if m not in self.completed]

        self.print_estimate(len(members))

        if not pending:
            logger.info("所有会员已完成，无需回填")
            return

        logger.info("开始回填：%d 个待处理会员", len(pending))

        # 分批执行
        for batch_idx in range(0, len(pending), self.batch_size):
            batch = pending[batch_idx : batch_idx + self.batch_size]
            batch_num = batch_idx // self.batch_size + 1
            total_batches = (len(pending) + self.batch_size - 1) // self.batch_size
            logger.info("批次 %d/%d：处理 %d 个会员", batch_num, total_batches, len(batch))

            for member_id in batch:
                try:
                    await self.run_member(member_id)
                    self.completed.add(member_id)
                except ConnectionError as exc:
                    # 数据库连接中断：保存 checkpoint 并退出
                    logger.error("数据库连接中断 (member_id=%d): %s", member_id, exc)
                    self.save_checkpoint()
                    sys.exit(1)
                except Exception as exc:
                    # 单个会员失败：记录日志，跳过，不写入 checkpoint
                    logger.error("会员 %d 调用链失败: %s", member_id, exc)
                    continue

            # 每批完成后保存 checkpoint
            self.save_checkpoint()

            # 批间间隔（最后一批不等待）
            if batch_idx + self.batch_size < len(pending):
                logger.info("等待 %ds 后继续下一批...", self.interval)
                await asyncio.sleep(self.interval)

        logger.info("回填完成：共处理 %d 个会员", len(self.completed))


# ── CLI 入口 ─────────────────────────────────────────


def main() -> None:
    parser = argparse.ArgumentParser(
        description="AI 历史数据回填脚本（半年内活跃会员）",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="试运行模式，不执行实际 AI 调用",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=BATCH_SIZE,
        help=f"每批处理会员数（默认 {BATCH_SIZE}）",
    )
    parser.add_argument(
        "--interval",
        type=int,
        default=BATCH_INTERVAL,
        help=f"批间间隔秒数（默认 {BATCH_INTERVAL}）",
    )
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    )

    runner = AIBackfillRunner(
        dry_run=args.dry_run,
        batch_size=args.batch_size,
        interval=args.interval,
    )
    asyncio.run(runner.run())


if __name__ == "__main__":
    main()