intl_news/llm/pipeline.py

"""批量翻译 + 事件抽取管道。

输入:  data/deduped/{YYYYMMDD}/uniques/{url_hash}.json（M3 去重后唯一条目）
输出:  data/events/{YYYYMMDD}/{url_hash}.json
"""

import json
import logging
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from pathlib import Path

import yaml

from crawler.utils import get_news_day
from extractor.models import ProcessedArticle
from llm.client import LLMConfig, load_llm_config, make_sync_client
from llm.extractor import PromptTemplate, translate_and_extract
from llm.models import EnTranslatedArticle, LLMCallError

logger = logging.getLogger(__name__)

# 默认并发数
_DEFAULT_CONCURRENCY = 3


def _load_concurrency() -> int:
    """从 system.yaml 读取 LLM 并发配置。"""
    config_path = Path("configs/system.yaml")
    if config_path.exists():
        try:
            with open(config_path, encoding="utf-8") as f:
                raw = yaml.safe_load(f)
            return int(raw.get("llm", {}).get("concurrency", _DEFAULT_CONCURRENCY))
        except Exception:
            pass
    return _DEFAULT_CONCURRENCY


def _load_deduped_articles(
    date_str: str,
) -> list[ProcessedArticle]:
    """加载指定日期的去重后唯一条目。

    Args:
        date_str: 日期 YYYYMMDD

    Returns:
        ProcessedArticle 列表
    """
    base_dir = Path(f"data/deduped/{date_str}/uniques")
    if not base_dir.exists():
        return []

    articles: list[ProcessedArticle] = []
    for json_file in sorted(base_dir.glob("*.json")):
        try:
            data = json.loads(json_file.read_text(encoding="utf-8"))
            article = ProcessedArticle(**data)
            # 防御性过滤：旧版数据里可能残留 M2 no_content 唯一篇，
            # 避免用空正文调用 LLM。
            if article.status != "success" or not article.content.strip():
                continue
            articles.append(article)
        except (json.JSONDecodeError, Exception) as e:
            logger.warning("解析去重文章失败 %s: %s", json_file, e)

    return articles


def _process_one(
    article: ProcessedArticle,
    client,
    config: LLMConfig,
    template: PromptTemplate,
) -> EnTranslatedArticle | None:
    """处理单篇文章的翻译 + 事件抽取。返回 None 表示失败。"""
    try:
        return translate_and_extract(
            client=client,
            config=config,
            article=article,
            template=template,
        )
    except LLMCallError as e:
        logger.error(
            "翻译抽取最终失败 url=%s: %s", article.url, e.reason
        )
        return None
    except Exception as e:
        logger.exception("翻译抽取未预期异常 url=%s: %s", article.url, e)
        return None


def translate_all_deduped(
    date_str: str | None = None,
    *,
    provider: str | None = None,
    model: str | None = None,
    concurrency: int | None = None,
) -> dict:
    """对去重后的所有唯一条目执行翻译 + 事件抽取。

    Args:
        date_str: 日期 YYYYMMDD，默认当前新闻日
        provider: LLM provider，默认从 system.yaml 读取
        model: LLM model，默认从 system.yaml 读取
        concurrency: 并发数，默认从 system.yaml 读取

    Returns:
        统计 dict
    """
    if date_str is None:
        date_str = get_news_day()

    concurrency = concurrency or _load_concurrency()

    logger.info("══════ 开始翻译+事件抽取，日期: %s，并发: %d ══════",
                date_str, concurrency)

    # 加载去重后的文章
    articles = _load_deduped_articles(date_str)
    if not articles:
        logger.warning("去重目录无文章: data/deduped/%s/uniques/", date_str)
        return {"date": date_str, "total": 0, "success": 0, "failed": 0, "elapsed_sec": 0}

    # 初始化 LLM 客户端（translation 场景配置见 system.yaml llm_scenes.translation）
    config = load_llm_config(provider=provider, model=model, scene="translation")
    client = make_sync_client(config)
    logger.info(
        "AI 大模型（场景 translation）: provider=%s model=%s",
        config.provider, config.model,
    )
    template = PromptTemplate()

    # 输出目录
    out_dir = Path(f"data/events/{date_str}")
    out_dir.mkdir(parents=True, exist_ok=True)

    start_time = datetime.now()
    success = 0
    failed = 0

    # 增量：跳过已翻译的文章
    new_articles = []
    skipped = 0
    for a in articles:
        if (out_dir / f"{a.url_hash}.json").exists():
            skipped += 1
        else:
            new_articles.append(a)
    if skipped > 0:
        logger.info("增量跳过 %d 篇已翻译，剩余 %d 篇待处理", skipped, len(new_articles))
    articles = new_articles

    # 并发处理
    with ThreadPoolExecutor(max_workers=concurrency) as executor:
        futures = {
            executor.submit(_process_one, article, client, config, template): article
            for article in articles
        }

        for future in as_completed(futures):
            article = futures[future]
            try:
                result = future.result()
            except Exception as e:
                logger.error("并发任务异常 url=%s: %s", article.url, e)
                failed += 1
                continue

            if result is None:
                failed += 1
                continue

            # 写入输出文件
            out_file = out_dir / f"{result.url_hash}.json"
            out_file.write_text(
                result.model_dump_json(indent=2, ensure_ascii=False),
                encoding="utf-8",
            )
            success += 1
            logger.info(
                "[%s] ✅ %s → %s (%d events, %d zh chars)",
                result.source_id,
                result.title[:40],
                result.title_zh[:30],
                len(result.events),
                result.word_count_zh,
            )

    elapsed = (datetime.now() - start_time).total_seconds()

    # 写入事件索引
    _write_event_index(date_str, success, failed, elapsed, config)

    logger.info(
        "══════ 翻译+事件抽取完成: 成功 %d / 失败 %d / 总计 %d，耗时 %.1f 秒 ══════",
        success, failed, len(articles), elapsed,
    )

    return {
        "date": date_str,
        "total": len(articles),
        "success": success,
        "failed": failed,
        "elapsed_sec": elapsed,
        "provider": config.provider,
        "model": config.model,
    }


def _write_event_index(
    date_str: str,
    success: int,
    failed: int,
    elapsed_sec: float,
    config: LLMConfig,
) -> None:
    """写入事件索引文件。

    Args:
        date_str: 日期
        success: 成功数
        failed: 失败数
        elapsed_sec: 耗时（秒）
        config: LLM 配置
    """
    out_dir = Path(f"data/events/{date_str}")
    out_dir.mkdir(parents=True, exist_ok=True)

    index_data = {
        "date": date_str,
        "success": success,
        "failed": failed,
        "elapsed_sec": round(elapsed_sec, 1),
        "provider": config.provider,
        "model": config.model,
        "generated_at": datetime.now().isoformat(),
    }

    index_path = out_dir / "index.json"
    index_path.write_text(
        json.dumps(index_data, indent=2, ensure_ascii=False),
        encoding="utf-8",
    )
    logger.info("事件索引已写入: %s", index_path)