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