初始化

vectorstore/pipeline.py

@@ -0,0 +1,216 @@
+"""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()