intl_news/embedding/pipeline.py

"""批量向量生成管道。

输入:  data/events/{YYYYMMDD}/{url_hash}.json（M4 翻译+事件输出）
输出:  data/embeddings/{YYYYMMDD}/{url_hash}.json
"""

import json
import logging
from datetime import datetime
from pathlib import Path

from crawler.utils import get_news_day
from embedding.client import (
    EmbeddingConfig,
    load_embedding_config,
    make_embedding_client,
)
from embedding.embedder import embed_articles
from embedding.models import EmbeddingError
from llm.models import EnTranslatedArticle

logger = logging.getLogger(__name__)


def _load_event_articles(date_str: str) -> list[EnTranslatedArticle]:
    """加载指定日期的翻译+事件文章。

    Args:
        date_str: 日期 YYYYMMDD

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

    articles: list[EnTranslatedArticle] = []
    for json_file in sorted(base_dir.glob("*.json")):
        if json_file.name == "index.json":
            continue
        try:
            data = json.loads(json_file.read_text(encoding="utf-8"))
            articles.append(EnTranslatedArticle(**data))
        except (json.JSONDecodeError, Exception) as e:
            logger.warning("解析事件文章失败 %s: %s", json_file, e)

    return articles


def embed_all_events(
    date_str: str | None = None,
    *,
    model: str | None = None,
) -> dict:
    """对所有 M4 输出的文章执行向量化。

    Args:
        date_str: 日期 YYYYMMDD，默认当前新闻日
        model: Embedding 模型名，默认从 system.yaml 读取

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

    logger.info("══════ 开始向量生成，日期: %s ══════", date_str)

    # 加载文章
    articles = _load_event_articles(date_str)
    if not articles:
        logger.warning("事件目录无文章: data/events/%s/", date_str)
        return {"date": date_str, "total": 0, "success": 0, "failed": 0, "elapsed_sec": 0}

    # 初始化 Embedding 客户端
    config = load_embedding_config(model=model)
    client = make_embedding_client(config)

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

    # 增量：跳过已向量化的文章
    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

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

    # 批量嵌入（按 batch_size 分块，每批输出进度）
    batch_size = config.batch_size
    total = len(articles)
    logger.info("开始向量化 %d 篇文章（batch_size=%d, model=%s）",
                total, batch_size, config.model)

    for batch_start in range(0, total, batch_size):
        batch_end = min(batch_start + batch_size, total)
        batch = articles[batch_start:batch_end]

        try:
            results = embed_articles(client, config, batch)
            for result in results:
                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("  [%d/%d] ✅ %d 篇 → %d 维向量",
                        batch_end, total, len(results), config.dimension)

        except EmbeddingError as e:
            failed += len(batch)
            logger.error("批量嵌入失败 [%d-%d]: %s", batch_start, batch_end, e.reason)
        except Exception as e:
            failed += len(batch)
            logger.exception("批量嵌入异常 [%d-%d]: %s", batch_start, batch_end, e)

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

    # 写入索引
    _write_embedding_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,
        "dimension": config.dimension,
    }


def _write_embedding_index(
    date_str: str,
    success: int,
    failed: int,
    elapsed_sec: float,
    config: EmbeddingConfig,
) -> None:
    """写入向量索引文件。"""
    out_dir = Path(f"data/embeddings/{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,
        "dimension": config.dimension,
        "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)