intl_news/vectorstore/pipeline.py

"""Qdrant 入库管道 + 语义搜索。

输入:  data/embeddings/{YYYYMMDD}/{url_hash}.json（M5 向量） + data/events/（M4 元数据）
动作:  upsert 到 Qdrant collection
搜索:  embed query → Qdrant query → 返回 SearchResult 列表
"""

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

from crawler.utils import get_news_day
from embedding.client import (
    embed_batch,
    load_embedding_config,
    make_embedding_client,
)
from vectorstore.client import VectorStore, make_qdrant_client
from vectorstore.models import SearchFilter, SearchResult

logger = logging.getLogger(__name__)


def _load_embedding_files(date_str: str) -> list[dict]:
    """加载指定日期的嵌入向量文件（含对应的 M4 元数据）。

    Args:
        date_str: 日期 YYYYMMDD

    Returns:
        dict 列表，含 url_hash / vector / article 信息
    """
    embedding_dir = Path(f"data/embeddings/{date_str}")
    event_dir = Path(f"data/events/{date_str}")

    if not embedding_dir.exists():
        return []

    items: list[dict] = []
    for emb_file in sorted(embedding_dir.glob("*.json")):
        if emb_file.name == "index.json":
            continue
        try:
            emb_data = json.loads(emb_file.read_text(encoding="utf-8"))

            # 加载对应的 M4 事件文章获取元数据
            event_file = event_dir / emb_file.name
            article_data = {}
            if event_file.exists():
                article_data = json.loads(event_file.read_text(encoding="utf-8"))

            items.append({
                "url_hash": emb_data["url_hash"],
                "source_id": emb_data.get("source_id", ""),
                "vector": emb_data["vector"],
                "article": article_data,
            })
        except (json.JSONDecodeError, Exception) as e:
            logger.warning("加载嵌入文件失败 %s: %s", emb_file, e)

    return items


def _build_payload(article_data: dict) -> dict:
    """从 M4 文章数据构建 Qdrant payload。

    Args:
        article_data: EnTranslatedArticle 的 dict

    Returns:
        payload dict
    """
    content_zh = article_data.get("content_zh", "")
    return {
        "title": article_data.get("title", ""),
        "title_zh": article_data.get("title_zh", ""),
        "url": article_data.get("url", ""),
        "source_id": article_data.get("source_id", ""),
        "source_name": article_data.get("source_name", ""),
        "publish_time": article_data.get("publish_time", ""),
        "events": article_data.get("events", []),
        "word_count": article_data.get("word_count", 0),
        "word_count_zh": article_data.get("word_count_zh", 0),
        "content_zh_preview": content_zh[:300] if content_zh else "",
    }


def ingest_all_embeddings(
    date_str: str | None = None,
    *,
    recreate: bool = False,
) -> dict:
    """将所有 M5 向量入库 Qdrant。

    Args:
        date_str: 日期 YYYYMMDD，默认当前新闻日
        recreate: 是否重建 collection

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

    logger.info("══════ 开始 Qdrant 入库，日期: %s ══════", date_str)

    items = _load_embedding_files(date_str)
    if not items:
        logger.warning("嵌入目录无数据: data/embeddings/%s/", date_str)
        return {"date": date_str, "total": 0, "ingested": 0, "failed": 0, "elapsed_sec": 0}

    start_time = datetime.now()

    # 构造 Qdrant 客户端
    client = make_qdrant_client()
    store = VectorStore(client)

    try:
        # 初始化 collection
        store.init_collection(recreate=recreate)

        # 构建 points
        points: list[dict] = []
        for item in items:
            payload = _build_payload(item["article"])
            points.append({
                "id": item["url_hash"],
                "vector": item["vector"],
                "payload": payload,
            })

        # 批量写入
        ingested = store.upsert(points)
        failed = len(points) - ingested

    finally:
        store.close()

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

    logger.info(
        "══════ Qdrant 入库完成: %d 条，耗时 %.1f 秒 ══════",
        ingested, elapsed,
    )

    return {
        "date": date_str,
        "total": len(items),
        "ingested": ingested,
        "failed": failed,
        "elapsed_sec": elapsed,
    }


def search_news(
    query: str,
    *,
    top_k: int = 10,
    search_filter: SearchFilter | None = None,
    score_threshold: float | None = None,
) -> list[SearchResult]:
    """语义搜索新闻。

    流程:
        1. 将查询文本向量化（使用 M5 Embedding 服务）
        2. Qdrant 语义检索

    Args:
        query: 中文搜索查询
        top_k: 返回条数
        search_filter: 可选过滤条件
        score_threshold: 最低相似度阈值

    Returns:
        SearchResult 列表
    """
    # 1. 向量化查询
    emb_config = load_embedding_config()
    emb_client = make_embedding_client(emb_config)

    try:
        vectors = embed_batch(emb_client, emb_config, [query])
        if not vectors:
            logger.error("查询向量化失败")
            return []
        query_vector = vectors[0]
    finally:
        emb_client.close()

    # 2. Qdrant 检索
    qdrant = make_qdrant_client()
    store = VectorStore(qdrant)

    try:
        results = store.query(
            query_vector=query_vector,
            top_k=top_k,
            search_filter=search_filter,
            score_threshold=score_threshold,
        )
    finally:
        store.close()

    return results


def get_collection_info() -> dict:
    """获取 Qdrant collection 信息。"""
    client = make_qdrant_client()
    store = VectorStore(client)
    try:
        info = store.info()
        return info.model_dump()
    finally:
        store.close()