"""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)