331 lines
11 KiB
Python
331 lines
11 KiB
Python
"""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)
|