初始化
This commit is contained in:
@@ -0,0 +1,330 @@
|
||||
"""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)
|
||||
Reference in New Issue
Block a user