"""Qdrant 客户端封装 (M6)。 核心: - 连接: 本地文件模式（默认，无需 Docker）或 HTTP 远程模式 - 初始化 Collection: 1024 维 / 余弦距离 - upsert: 幂等写入（url_hash 转 UUID 做 point ID） - query: 语义检索 + 结构化过滤 - info / count: 运维辅助 """ import logging import os import uuid from pathlib import Path import yaml from qdrant_client import QdrantClient from qdrant_client.http.models import ( DatetimeRange, Distance, FieldCondition, Filter, MatchAny, MatchValue, PointStruct, Range, VectorParams, ) from vectorstore.models import CollectionInfo, SearchFilter, SearchResult logger = logging.getLogger(__name__) # 默认配置 DEFAULT_COLLECTION = "en_finance_news" DEFAULT_VECTOR_DIM = 1024 DEFAULT_DISTANCE = Distance.COSINE DEFAULT_STORAGE_PATH = Path("data/qdrant_storage") # UUID namespace for url_hash -> UUID conversion（确定性，便于幂等 upsert） _UUID_NAMESPACE = uuid.UUID("e1d2c3b4-a5f6-7890-abcd-ef1234567890") def url_hash_to_uuid(url_hash: str) -> str: """把 url_hash 转为 UUID 字符串（point ID 要求）。 使用 uuid5 保证确定性——相同 url_hash 总是得到相同 UUID。 """ return str(uuid.uuid5(_UUID_NAMESPACE, url_hash)) def _load_qdrant_config() -> dict: """从 system.yaml 加载 qdrant 段配置。""" 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 raw.get("qdrant", {}) except Exception: pass return {} def make_qdrant_client( *, memory: bool = False, path: str | None = None, ) -> QdrantClient: """构造 QdrantClient。 模式优先级: 1. memory=True → 内存模式（测试用） 2. path 非空 → 本地文件模式（嵌入式运行，无需 Docker） 3. QDRANT_URL + QDRANT_API_KEY 环境变量 → HTTP 远程模式 本地文件模式是默认推荐方式，对 ARM/Raspberry Pi 友好。 """ if memory: logger.debug("Qdrant 内存模式") return QdrantClient(location=":memory:") # 远程模式: 仅在 QDRANT_URL 为非 localhost 且显式指定 path 为 None 时使用 remote_url = os.environ.get("QDRANT_URL", "") if (remote_url and remote_url.startswith("http") and "localhost" not in remote_url and "127.0.0.1" not in remote_url): api_key = os.environ.get("QDRANT_API_KEY") or None logger.info("Qdrant 远程模式: %s", remote_url) return QdrantClient(url=remote_url, api_key=api_key, timeout=10) # 默认本地文件模式 use_path = path or str(DEFAULT_STORAGE_PATH) logger.info("Qdrant 本地文件模式: %s", use_path) return QdrantClient(path=use_path) class VectorStore: """Qdrant 向量知识库封装。 线程不安全，批处理串行使用即可。 """ def __init__( self, client: QdrantClient, collection_name: str | None = None, vector_dim: int = DEFAULT_VECTOR_DIM, ) -> None: self._c = client config = _load_qdrant_config() self.collection_name = collection_name or config.get("collection", DEFAULT_COLLECTION) self.vector_dim = vector_dim # ------------------------------------------------------------------ # # Collection 管理 # ------------------------------------------------------------------ # def init_collection(self, *, recreate: bool = False) -> None: """创建 collection（已存在时若 recreate 则重建）。""" exists = self._c.collection_exists(self.collection_name) if exists and not recreate: logger.debug("Collection %s 已存在，跳过初始化", self.collection_name) return if exists and recreate: logger.warning("重建 collection %s", self.collection_name) self._c.delete_collection(self.collection_name) self._c.create_collection( collection_name=self.collection_name, vectors_config=VectorParams( size=self.vector_dim, distance=DEFAULT_DISTANCE, ), ) logger.info( "已创建 collection %s (dim=%d distance=%s)", self.collection_name, self.vector_dim, DEFAULT_DISTANCE.name, ) def info(self) -> CollectionInfo: """获取 collection 概览信息。""" exists = self._c.collection_exists(self.collection_name) if not exists: return CollectionInfo(name=self.collection_name, exists=False) 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: return 0 # ------------------------------------------------------------------ # # 数据写入（幂等 upsert） # ------------------------------------------------------------------ # def upsert( self, points: list[dict], *, batch_size: int = 100, ) -> int: """批量幂等写入。 Args: points: 每个 dict 包含: id (str) point ID（url_hash） vector (list[float]) 嵌入向量 payload (dict) 任意结构化数据 batch_size: 每批写入条数 Returns: 写入条数 """ 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 批 %d/%d (%d 条)", i // batch_size + 1, (total + batch_size - 1) // batch_size, len(chunk), ) logger.info("upsert 完成: %d 条 → collection %s", total, self.collection_name) return total # ------------------------------------------------------------------ # # 检索 # ------------------------------------------------------------------ # def query( self, query_vector: list[float], *, top_k: int = 10, search_filter: SearchFilter | None = None, score_threshold: float | None = None, ) -> list[SearchResult]: """语义检索 + 可选结构化过滤。 Args: query_vector: 嵌入向量（需与 collection 维度一致） top_k: 返回条数 search_filter: 结构化过滤（AND 关系） score_threshold: 最低余弦相似度 Returns: 列表按 score 降序 """ q_filter = _build_filter(search_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 {} 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 "", title_zh=payload.get("title_zh") or "", url=payload.get("url") or "", source_id=payload.get("source_id") or "", publish_time=payload.get("publish_time") or "", events=payload.get("events") or [], content_zh_preview=payload.get("content_zh_preview") or "", )) logger.debug("检索完成 top_k=%d → %d 条", top_k, 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="events[].stock_codes", match=MatchAny(any=f.stock_codes), ) ) if f.sentiment: conditions.append( FieldCondition( key="events[].sentiment", match=MatchValue(value=f.sentiment), ) ) if f.importance_min is not None: conditions.append( FieldCondition( key="events[].importance", range=Range(gte=f.importance_min), ) ) if f.event_types: conditions.append( FieldCondition( key="events[].event_type", match=MatchAny(any=f.event_types), ) ) if f.publish_date_from or f.publish_date_to: try: range_kwargs: dict = {} if f.publish_date_from: range_kwargs["gte"] = f.publish_date_from + "T00:00:00" if f.publish_date_to: range_kwargs["lte"] = 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)