Initial commit

2026-07-18 15:51:01 +08:00
commit f2c80c5a9c
799 changed files with 133475 additions and 0 deletions
@@ -0,0 +1,326 @@
+"""Qdrant 客户端封装 (M6)。
+
+核心:
+    - 连接:本地文件模式(默认,嵌入运行无需 Docker)或 HTTP 远程模式
+    - 初始化 Collection:1024 维 / 余弦距离
+    - upsert:幂等写入(url_hash 做 point ID)
+    - query:语义检索 + 结构化过滤
+    - info / delete / count:运维辅助
+"""
+
+from __future__ import annotations
+
+import os
+import uuid
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+from loguru import logger
+from qdrant_client import QdrantClient
+from qdrant_client.http.models import (
+    DatetimeRange,
+    Distance,
+    FieldCondition,
+    Filter,
+    MatchAny,
+    MatchValue,
+    PointStruct,
+    Range,
+    VectorParams,
+)
+
+from .models import CollectionInfo, SearchFilter, SearchResult
+
+# 默认配置
+DEFAULT_HOST = "localhost"
+DEFAULT_PORT = 6333
+DEFAULT_COLLECTION = "a_share_news"
+DEFAULT_VECTOR_DIM = 1024
+DEFAULT_DISTANCE = Distance.COSINE
+
+# UUID namespace for url_hash -> UUID conversion (确定性,便于幂等 upsert)
+_UUID_NAMESPACE = uuid.UUID("a1b2c3d4-e5f6-7890-abcd-ef1234567890")
+
+
+def url_hash_to_uuid(url_hash: str) -> str:
+    """把 16 位 hex url_hash 转为 UUID 字符串(point ID 要求)。
+
+    使用 uuid5 保证确定性——相同 url_hash 总是得到相同 UUID。
+    """
+    return str(uuid.uuid5(_UUID_NAMESPACE, url_hash))
+
+
+def _read_env(key: str, default: str | None = None) -> str | None:
+    val = os.environ.get(key)
+    if val is None or val.strip() == "":
+        return default
+    return val.strip()
+
+
+def make_qdrant_client(
+    host: str | None = None,
+    port: int | None = None,
+    *,
+    memory: bool = False,
+    path: str | None = None,
+) -> QdrantClient:
+    """构造 QdrantClient。
+
+    模式优先级:
+        1. memory=True  -> 内存模式(测试用)
+        2. path 非空    -> 本地文件模式(嵌入运行,无需 Docker,默认 data/qdrant_storage)
+        3. host/port    -> 远程 HTTP 模式(需要单独 Qdrant 服务)
+
+    树莓派 5 ARM64 的 Docker Qdrant 不兼容 16K 页内核,
+    推荐默认用本地文件模式。
+    """
+    if memory:
+        logger.debug("Qdrant 内存模式")
+        return QdrantClient(location=":memory:")
+
+    # 本地文件模式:显式 path 或 host 未指定时默认走文件
+    if path is not None or (host is None and port is None):
+        use_path = path or str(DEFAULT_STORAGE_PATH)
+        logger.debug("Qdrant 本地文件模式: {}", use_path)
+        return QdrantClient(path=use_path)
+
+    h = host or _read_env("QDRANT_HOST", DEFAULT_HOST) or DEFAULT_HOST
+    p = int(port or int(_read_env("QDRANT_PORT", str(DEFAULT_PORT)) or DEFAULT_PORT))  # type: ignore[arg-type]
+    api_key = _read_env("QDRANT_API_KEY") or None
+    url = f"http://{h}:{p}"
+    logger.debug("Qdrant HTTP {} (key={})", url, "yes" if api_key else "no")
+    return QdrantClient(url=url, api_key=api_key, timeout=10)
+
+# 默认持久化目录
+DEFAULT_STORAGE_PATH = Path("data/qdrant_storage")
+
+
+# --------------------------------------------------------------------------- #
+# 客户端封装
+# --------------------------------------------------------------------------- #
+
+class VectorStore:
+    """Qdrant 向量知识库封装。
+
+    线程不安全,批处理串行使用即可。
+    """
+
+    def __init__(
+        self,
+        client: QdrantClient,
+        collection_name: str | None = None,
+        vector_dim: int = DEFAULT_VECTOR_DIM,
+    ) -> None:
+        self._c = client
+        self.collection_name = collection_name or (
+            _read_env("QDRANT_COLLECTION", DEFAULT_COLLECTION) or DEFAULT_COLLECTION
+        )
+        self.vector_dim = vector_dim
+
+    # ------------------------------------------------------------------ #
+    # Collection 管理
+    # ------------------------------------------------------------------ #
+
+    def init_collection(self, *, recreate: bool = False) -> None:
+        """创建 collection(已存在时若 recreate 则重建)。
+
+        幂等:已存在且非 recreate 时直接返回。
+        """
+        exists = self._c.collection_exists(self.collection_name)
+        if exists and not recreate:
+            logger.debug("Collection {} 已存在,跳过初始化", self.collection_name)
+            return
+        if exists and recreate:
+            logger.warning("重建 collection {}", self.collection_name)
+            self._c.delete_collection(self.collection_name)
+            exists = False
+
+        self._c.create_collection(
+            collection_name=self.collection_name,
+            vectors_config=VectorParams(
+                size=self.vector_dim,
+                distance=DEFAULT_DISTANCE,
+            ),
+        )
+        logger.info(
+            "已创建 collection {} (dim={} distance={})",
+            self.collection_name, self.vector_dim, DEFAULT_DISTANCE.name,
+        )
+
+    def delete_collection(self) -> None:
+        if self._c.collection_exists(self.collection_name):
+            self._c.delete_collection(self.collection_name)
+            logger.info("已删除 collection {}", self.collection_name)
+
+    def info(self) -> CollectionInfo:
+        exists = self._c.collection_exists(self.collection_name)
+        if not exists:
+            return CollectionInfo(name=self.collection_name, exists=False, vectors_count=0)
+        c_info = self._c.get_collection(self.collection_name)
+        return CollectionInfo(
+            name=self.collection_name,
+            exists=True,
+            vectors_count=c_info.points_count or 0,
+            indexed_vectors_count=getattr(c_info, "indexed_vectors_count", None),
+            segments_count=getattr(c_info, "segments_count", None),
+        )
+
+    def count(self) -> int:
+        try:
+            return self._c.count(self.collection_name).count
+        except Exception:  # noqa: BLE001 - collection 不存在时优雅退化
+            return 0
+
+    # ------------------------------------------------------------------ #
+    # 数据写入(幂等 upsert)
+    # ------------------------------------------------------------------ #
+
+    def upsert(
+        self,
+        points: list[dict[str, Any]],
+        *,
+        batch_size: int = 100,
+    ) -> int:
+        """批量幂等写入。
+
+        参数:
+            points: 每个 dict 包含:
+                    id (str)               point ID(用 url_hash)
+                    vector (list[float])   嵌入向量
+                    payload (dict)        任意结构化数据
+            batch_size: 每批写入条数。
+
+        返回: 写入条数。
+        """
+        structs = [
+            PointStruct(
+                id=url_hash_to_uuid(p["id"]),
+                vector=p["vector"],
+                payload={"url_hash": p["id"], **(p.get("payload") or {})},
+            )
+            for p in points
+        ]
+        total = len(structs)
+        for i in range(0, total, batch_size):
+            chunk = structs[i : i + batch_size]
+            self._c.upsert(collection_name=self.collection_name, points=chunk)
+            logger.debug("upsert 批 {}/{} ({} 条)", i // batch_size + 1, (total + batch_size - 1) // batch_size, len(chunk))
+        logger.info("upsert 完成: {} 条 -> collection {}", total, self.collection_name)
+        return total
+
+    # ------------------------------------------------------------------ #
+    # 检索
+    # ------------------------------------------------------------------ #
+
+    def query(
+        self,
+        query_vector: list[float],
+        *,
+        top_k: int = 10,
+        filter: SearchFilter | None = None,
+        score_threshold: float | None = None,
+    ) -> list[SearchResult]:
+        """语义检索 + 可选结构化过滤。
+
+        参数:
+            query_vector:    嵌入向量(需与 collection 维度一致)。
+            top_k:           返回条数。
+            filter:          结构化过滤(AND 关系)。
+            score_threshold: 最低余弦相似度。
+
+        返回: 列表按 score 降序。
+        """
+        # 构建 Qdrant Filter
+        q_filter = _build_filter(filter)
+        hits = self._c.query_points(
+            collection_name=self.collection_name,
+            query=query_vector,
+            query_filter=q_filter,
+            limit=top_k,
+            score_threshold=score_threshold,
+            with_payload=True,
+            with_vectors=False,
+        )
+        results: list[SearchResult] = []
+        for p in hits.points:
+            payload = p.payload or {}
+            pt_raw = payload.get("publish_time")
+            publish_time = (
+                datetime.fromisoformat(pt_raw)
+                if isinstance(pt_raw, str) and pt_raw
+                else None
+            )
+            results.append(SearchResult(
+                url_hash=payload.get("url_hash") or "",
+                score=p.score if p.score is not None else 0.0,
+                title=payload.get("title") or "",
+                url=payload.get("url") or "",
+                source_id=payload.get("source_id") or "",
+                publish_time=publish_time,
+                event=payload.get("event"),
+                char_count=payload.get("char_count"),
+                word_count=payload.get("word_count"),
+            ))
+        logger.debug(
+            "检索完成 top_k={} filter={} -> {} 条",
+            top_k, filter, len(results),
+        )
+        return results
+
+    def close(self) -> None:
+        self._c.close()
+
+    def __enter__(self) -> VectorStore:
+        return self
+
+    def __exit__(self, *_: object) -> None:
+        self.close()
+
+
+# --------------------------------------------------------------------------- #
+# Filter 构建
+# --------------------------------------------------------------------------- #
+
+def _build_filter(f: SearchFilter | None) -> Filter | None:
+    """把 SearchFilter 转换为 Qdrant Filter。"""
+    if f is None:
+        return None
+    conditions: list[FieldCondition] = []
+
+    if f.source_id:
+        conditions.append(FieldCondition(key="source_id", match=MatchValue(value=f.source_id)))
+    if f.source_ids:
+        conditions.append(FieldCondition(key="source_id", match=MatchAny(any=f.source_ids)))
+    if f.stock_codes:
+        conditions.append(FieldCondition(key="event.stock_codes", match=MatchAny(any=f.stock_codes)))
+    if f.company_names:
+        conditions.append(FieldCondition(key="event.company_names", match=MatchAny(any=f.company_names)))
+    if f.industries:
+        conditions.append(FieldCondition(key="event.industries", match=MatchAny(any=f.industries)))
+    if f.sentiment:
+        conditions.append(FieldCondition(key="event.sentiment", match=MatchValue(value=f.sentiment)))
+    if f.importance_min is not None:
+        conditions.append(FieldCondition(key="event.importance", range=Range(gte=f.importance_min)))
+    if f.event_types:
+        conditions.append(FieldCondition(key="event.event_type", match=MatchAny(any=f.event_types)))
+    if f.publish_date_from or f.publish_date_to:
+        try:
+            range_kwargs: dict[str, datetime] = {}
+            if f.publish_date_from:
+                range_kwargs["gte"] = datetime.fromisoformat(f.publish_date_from + "T00:00:00")
+            if f.publish_date_to:
+                range_kwargs["lte"] = datetime.fromisoformat(f.publish_date_to + "T23:59:59")
+            conditions.append(FieldCondition(
+                key="publish_time",
+                range=DatetimeRange(**range_kwargs),
+            ))
+        except ValueError:
+            logger.warning(
+                "filter 日期格式错误 from={!r} to={!r},跳过日期过滤",
+                f.publish_date_from, f.publish_date_to,
+            )
+
+    if not conditions:
+        return None
+    return Filter(must=conditions)