tps-dashboard/ai-labeling/scripts/gen_content_digest.py

"""
gen_content_digest.py - 批量生成 22 期节目内容摘要卡
用法:
  cd E:\tps-dashboard\ai-labeling
  python scripts/gen_content_digest.py

功能:
  - 读取 prompt4_content_digest.md 作为 system prompt
  - 遍历 doco/deliverables/ 下所有融合A稿 .docx 文件
  - 用 python-docx 提取文稿文本，调用 MiMo API 生成结构化摘要卡
  - 支持断点续跑（已存在的 digest 文件自动跳过）
"""

import sys
sys.stdout.reconfigure(encoding='utf-8')
sys.stderr.reconfigure(encoding='utf-8')

import os
import re
import json
import time
from pathlib import Path
from datetime import datetime
from openai import OpenAI
from dotenv import load_dotenv
from docx import Document

# 加载 .env（优先加载 ai-labeling 目录下的 .env）
SCRIPT_DIR = Path(__file__).parent
BASE_DIR = SCRIPT_DIR.parent  # ai-labeling/
load_dotenv(BASE_DIR / ".env")
load_dotenv()  # 也尝试加载项目根目录的 .env

# 目录配置
DELIVERABLES_DIR = BASE_DIR.parent / "doco" / "deliverables"
PROMPTS_DIR = BASE_DIR / "prompts"
EXPERIMENTS_DIR = BASE_DIR / "experiments" / "content_digests"

# MiMo API 配置（与 run_labeling.py 一致）
MIMO_CONFIG = {
    "base_url": "https://api.xiaomimimo.com/v1",
    "model_name": "mimo-v2.5-pro",
    "api_key_env": "MIMO_API_KEY",
}


def load_system_prompt() -> str:
    """加载 prompt4_content_digest.md 作为 system prompt。"""
    prompt_file = PROMPTS_DIR / "prompt4_content_digest.md"
    if not prompt_file.exists():
        raise FileNotFoundError(f"找不到 prompt 文件: {prompt_file}")
    return prompt_file.read_text(encoding="utf-8")


def parse_filename(filename: str) -> dict:
    """
    从文件名解析元信息。
    文件名格式: 第02期_20260113_武器进化论：海战颠覆者_付天雨_融合A稿.docx
    返回: {"ep": 2, "date": "2026-01-13", "title": "...", "editor": "..."}
    """
    name = filename.replace(".docx", "")
    parts = name.split("_")
    if len(parts) < 4:
        return None
    # 解析期号
    ep_match = re.search(r'第(\d+)期', parts[0])
    if not ep_match:
        return None
    ep = int(ep_match.group(1))
    # 解析日期（YYYYMMDD -> YYYY-MM-DD）
    raw_date = parts[1]
    if len(raw_date) == 8 and raw_date.isdigit():
        date = f"{raw_date[:4]}-{raw_date[4:6]}-{raw_date[6:8]}"
    else:
        date = raw_date
    title = parts[2]
    editor = parts[3]
    return {"ep": ep, "date": date, "title": title, "editor": editor}


def extract_docx_text(filepath: Path) -> str:
    """用 python-docx 提取 .docx 文件的全部文本，逐段拼接。"""
    doc = Document(str(filepath))
    paragraphs = []
    for para in doc.paragraphs:
        text = para.text.strip()
        if text:
            paragraphs.append(text)
    return "\n".join(paragraphs)


def build_user_message(meta: dict, body: str) -> str:
    """构造 user message：元信息 + 文稿全文。"""
    return (
        f"期号：第{meta['ep']:02d}期\n"
        f"播出日期：{meta['date']}\n"
        f"节目名：{meta['title']}\n"
        f"编导：{meta['editor']}\n"
        f"\n以下是节目文稿全文：\n\n{body}"
    )


def extract_json_from_response(raw: str) -> dict:
    """从模型响应中提取 JSON，兼容推理模型的<think>...输出。"""
    # 先去掉<think>...标签及其内容
    text = re.sub(r'<think>.*?</think>', '', raw, flags=re.DOTALL)
    text = text.strip()
    # 去掉 markdown 代码块
    text = re.sub(r'^```(?:json)?\s*', '', text)
    text = re.sub(r'\s*```$', '', text)
    text = text.strip()
    # 从第一个 { 开始，到最后一个 } 结束
    first_brace = text.find('{')
    last_brace = text.rfind('}')
    if first_brace != -1 and last_brace != -1 and last_brace >= first_brace:
        json_str = text[first_brace:last_brace + 1]
        return json.loads(json_str)
    # 兜底：直接尝试解析
    return json.loads(text)


def call_mimo(system_prompt: str, user_prompt: str) -> dict:
    """调用 MiMo API 生成摘要卡，返回解析后的 JSON dict。"""
    api_key = os.environ.get(MIMO_CONFIG["api_key_env"])
    if not api_key:
        raise EnvironmentError(
            f"环境变量 {MIMO_CONFIG['api_key_env']} 未设置，请检查 ai-labeling/.env 或根目录 .env 文件"
        )
    client = OpenAI(
        api_key=api_key,
        base_url=MIMO_CONFIG["base_url"],
    )
    response = client.chat.completions.create(
        model=MIMO_CONFIG["model_name"],
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        temperature=0.0,
        # 关闭 thinking（与 run_labeling.py 一致）
        extra_body={"thinking": {"type": "disabled"}},
    )
    raw = response.choices[0].message.content
    return extract_json_from_response(raw)


def collect_docx_files() -> list[dict]:
    """收集所有 .docx 文件并按期号排序。"""
    files = []
    for f in sorted(DELIVERABLES_DIR.iterdir()):
        if not f.name.endswith(".docx"):
            continue
        meta = parse_filename(f.name)
        if meta is None:
            print(f"⚠️  跳过无法解析的文件: {f.name}")
            continue
        meta["filepath"] = f
        files.append(meta)
    # 按期号排序
    files.sort(key=lambda x: x["ep"])
    return files


def main():
    # 确保输出目录存在
    EXPERIMENTS_DIR.mkdir(parents=True, exist_ok=True)

    # 加载 system prompt
    system_prompt = load_system_prompt()
    print(f"✅ 已加载 system prompt: prompt4_content_digest.md")

    # 收集 docx 文件
    docx_files = collect_docx_files()
    print(f"✅ 找到 {len(docx_files)} 个融合A稿 docx 文件\n")

    all_digests = []
    skipped = 0
    success = 0
    failed = 0

    for i, meta in enumerate(docx_files):
        ep = meta["ep"]
        out_file = EXPERIMENTS_DIR / f"ep{ep:02d}_digest.json"

        # 断点续跳：已存在则跳过
        if out_file.exists():
            print(f"[{i+1}/{len(docx_files)}] ep{ep:02d} 已存在，跳过")
            try:
                existing = json.loads(out_file.read_text(encoding="utf-8"))
                all_digests.append(existing)
            except Exception:
                pass
            skipped += 1
            continue

        print(f"[{i+1}/{len(docx_files)}] 正在处理 ep{ep:02d} - {meta['title']} ...")

        try:
            # 提取 docx 文本
            body = extract_docx_text(meta["filepath"])
            if not body.strip():
                print(f"  ⚠️  ep{ep:02d} 文稿内容为空，跳过")
                failed += 1
                continue

            # 构造 user message
            user_msg = build_user_message(meta, body)

            # 调用 MiMo
            result = call_mimo(system_prompt, user_msg)

            # 构造输出
            output = {
                "ep": ep,
                "date": meta["date"],
                "title": meta["title"],
                "editor": meta["editor"],
                "filename": meta["filepath"].name,
                "digest": result,
                "generated_at": datetime.now().isoformat(),
            }

            # 写入单期文件
            out_file.write_text(
                json.dumps(output, ensure_ascii=False, indent=2),
                encoding="utf-8",
            )
            all_digests.append(output)
            print(f"  ✅ ep{ep:02d} 摘要卡已生成 -> {out_file.name}")
            success += 1

        except json.JSONDecodeError as e:
            print(f"  ⚠️  ep{ep:02d} LLM 返回的不是合法 JSON: {e}")
            failed += 1
        except Exception as e:
            print(f"  ❌ ep{ep:02d} 处理失败: {e}")
            failed += 1

        # 限流保护：每期之间 sleep 1 秒
        if i < len(docx_files) - 1:
            time.sleep(1)

    # 写入汇总文件
    summary_file = EXPERIMENTS_DIR / "_all_digests.json"
    summary_file.write_text(
        json.dumps(all_digests, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )

    print("\n" + "=" * 60)
    print(f"📊 处理完成:")
    print(f"   成功: {success}")
    print(f"   跳过(已存在): {skipped}")
    print(f"   失败: {failed}")
    print(f"   汇总文件: {summary_file}")
    print("=" * 60)


if __name__ == "__main__":
    main()