初始化
This commit is contained in:
@@ -0,0 +1,37 @@
|
||||
"""Qdrant 向量知识库模块 (M6)。
|
||||
|
||||
公共 API:
|
||||
- VectorStore: Qdrant 封装（init / upsert / query / info / count）
|
||||
- SearchFilter / SearchResult / CollectionInfo: 数据模型
|
||||
- make_qdrant_client: 工厂（内存/本地文件/远程 HTTP）
|
||||
- ingest_all_embeddings / search_news / get_collection_info: 管道
|
||||
"""
|
||||
|
||||
from vectorstore.client import (
|
||||
DEFAULT_COLLECTION,
|
||||
DEFAULT_VECTOR_DIM,
|
||||
VectorStore,
|
||||
make_qdrant_client,
|
||||
)
|
||||
from vectorstore.models import CollectionInfo, SearchFilter, SearchResult
|
||||
from vectorstore.pipeline import (
|
||||
get_collection_info,
|
||||
ingest_all_embeddings,
|
||||
search_news,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# 客户端
|
||||
"DEFAULT_COLLECTION",
|
||||
"DEFAULT_VECTOR_DIM",
|
||||
"VectorStore",
|
||||
"make_qdrant_client",
|
||||
# 模型
|
||||
"CollectionInfo",
|
||||
"SearchFilter",
|
||||
"SearchResult",
|
||||
# 管道
|
||||
"get_collection_info",
|
||||
"ingest_all_embeddings",
|
||||
"search_news",
|
||||
]
|
||||
@@ -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)
|
||||
@@ -0,0 +1,52 @@
|
||||
"""Qdrant 向量存储数据模型 (M6)。"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class SearchFilter(BaseModel):
|
||||
"""可选检索过滤条件，全部为 AND 关系。"""
|
||||
|
||||
source_id: str | None = None
|
||||
source_ids: list[str] | None = None
|
||||
stock_codes: list[str] | None = Field(default=None, description="match any")
|
||||
sentiment: str | None = None # positive / neutral / negative
|
||||
importance_min: int | None = None # >= N
|
||||
event_types: list[str] | None = None # match any
|
||||
publish_date_from: str | None = None # YYYY-MM-DD
|
||||
publish_date_to: str | None = None # YYYY-MM-DD
|
||||
|
||||
|
||||
class SearchResult(BaseModel):
|
||||
"""单条检索结果（含双语信息）。"""
|
||||
|
||||
url_hash: str
|
||||
score: float
|
||||
title: str = ""
|
||||
title_zh: str = ""
|
||||
url: str = ""
|
||||
source_id: str = ""
|
||||
publish_time: str = ""
|
||||
events: list[dict[str, Any]] = Field(default_factory=list)
|
||||
content_zh_preview: str = "" # 中文正文前 300 字
|
||||
|
||||
def short_summary(self) -> str:
|
||||
codes = set()
|
||||
for ev in self.events:
|
||||
codes.update(ev.get("stock_codes", []))
|
||||
codes_str = ",".join(sorted(codes)[:5]) or "-"
|
||||
return (
|
||||
f"[{self.source_id}] score={self.score:.4f} "
|
||||
f"《{self.title_zh[:40] or self.title[:40]}》 {codes_str}"
|
||||
)
|
||||
|
||||
|
||||
class CollectionInfo(BaseModel):
|
||||
"""Collection 概览信息。"""
|
||||
|
||||
name: str
|
||||
exists: bool
|
||||
vectors_count: int = 0
|
||||
indexed_vectors_count: int | None = None
|
||||
segments_count: int | None = None
|
||||
@@ -0,0 +1,216 @@
|
||||
"""Qdrant 入库管道 + 语义搜索。
|
||||
|
||||
输入: data/embeddings/{YYYYMMDD}/{url_hash}.json（M5 向量） + data/events/（M4 元数据）
|
||||
动作: upsert 到 Qdrant collection
|
||||
搜索: embed query → Qdrant query → 返回 SearchResult 列表
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
from crawler.utils import get_news_day
|
||||
from embedding.client import (
|
||||
embed_batch,
|
||||
load_embedding_config,
|
||||
make_embedding_client,
|
||||
)
|
||||
from vectorstore.client import VectorStore, make_qdrant_client
|
||||
from vectorstore.models import SearchFilter, SearchResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _load_embedding_files(date_str: str) -> list[dict]:
|
||||
"""加载指定日期的嵌入向量文件（含对应的 M4 元数据）。
|
||||
|
||||
Args:
|
||||
date_str: 日期 YYYYMMDD
|
||||
|
||||
Returns:
|
||||
dict 列表，含 url_hash / vector / article 信息
|
||||
"""
|
||||
embedding_dir = Path(f"data/embeddings/{date_str}")
|
||||
event_dir = Path(f"data/events/{date_str}")
|
||||
|
||||
if not embedding_dir.exists():
|
||||
return []
|
||||
|
||||
items: list[dict] = []
|
||||
for emb_file in sorted(embedding_dir.glob("*.json")):
|
||||
if emb_file.name == "index.json":
|
||||
continue
|
||||
try:
|
||||
emb_data = json.loads(emb_file.read_text(encoding="utf-8"))
|
||||
|
||||
# 加载对应的 M4 事件文章获取元数据
|
||||
event_file = event_dir / emb_file.name
|
||||
article_data = {}
|
||||
if event_file.exists():
|
||||
article_data = json.loads(event_file.read_text(encoding="utf-8"))
|
||||
|
||||
items.append({
|
||||
"url_hash": emb_data["url_hash"],
|
||||
"source_id": emb_data.get("source_id", ""),
|
||||
"vector": emb_data["vector"],
|
||||
"article": article_data,
|
||||
})
|
||||
except (json.JSONDecodeError, Exception) as e:
|
||||
logger.warning("加载嵌入文件失败 %s: %s", emb_file, e)
|
||||
|
||||
return items
|
||||
|
||||
|
||||
def _build_payload(article_data: dict) -> dict:
|
||||
"""从 M4 文章数据构建 Qdrant payload。
|
||||
|
||||
Args:
|
||||
article_data: EnTranslatedArticle 的 dict
|
||||
|
||||
Returns:
|
||||
payload dict
|
||||
"""
|
||||
content_zh = article_data.get("content_zh", "")
|
||||
return {
|
||||
"title": article_data.get("title", ""),
|
||||
"title_zh": article_data.get("title_zh", ""),
|
||||
"url": article_data.get("url", ""),
|
||||
"source_id": article_data.get("source_id", ""),
|
||||
"source_name": article_data.get("source_name", ""),
|
||||
"publish_time": article_data.get("publish_time", ""),
|
||||
"events": article_data.get("events", []),
|
||||
"word_count": article_data.get("word_count", 0),
|
||||
"word_count_zh": article_data.get("word_count_zh", 0),
|
||||
"content_zh_preview": content_zh[:300] if content_zh else "",
|
||||
}
|
||||
|
||||
|
||||
def ingest_all_embeddings(
|
||||
date_str: str | None = None,
|
||||
*,
|
||||
recreate: bool = False,
|
||||
) -> dict:
|
||||
"""将所有 M5 向量入库 Qdrant。
|
||||
|
||||
Args:
|
||||
date_str: 日期 YYYYMMDD，默认当前新闻日
|
||||
recreate: 是否重建 collection
|
||||
|
||||
Returns:
|
||||
统计 dict
|
||||
"""
|
||||
if date_str is None:
|
||||
date_str = get_news_day()
|
||||
|
||||
logger.info("══════ 开始 Qdrant 入库，日期: %s ══════", date_str)
|
||||
|
||||
items = _load_embedding_files(date_str)
|
||||
if not items:
|
||||
logger.warning("嵌入目录无数据: data/embeddings/%s/", date_str)
|
||||
return {"date": date_str, "total": 0, "ingested": 0, "failed": 0, "elapsed_sec": 0}
|
||||
|
||||
start_time = datetime.now()
|
||||
|
||||
# 构造 Qdrant 客户端
|
||||
client = make_qdrant_client()
|
||||
store = VectorStore(client)
|
||||
|
||||
try:
|
||||
# 初始化 collection
|
||||
store.init_collection(recreate=recreate)
|
||||
|
||||
# 构建 points
|
||||
points: list[dict] = []
|
||||
for item in items:
|
||||
payload = _build_payload(item["article"])
|
||||
points.append({
|
||||
"id": item["url_hash"],
|
||||
"vector": item["vector"],
|
||||
"payload": payload,
|
||||
})
|
||||
|
||||
# 批量写入
|
||||
ingested = store.upsert(points)
|
||||
failed = len(points) - ingested
|
||||
|
||||
finally:
|
||||
store.close()
|
||||
|
||||
elapsed = (datetime.now() - start_time).total_seconds()
|
||||
|
||||
logger.info(
|
||||
"══════ Qdrant 入库完成: %d 条，耗时 %.1f 秒 ══════",
|
||||
ingested, elapsed,
|
||||
)
|
||||
|
||||
return {
|
||||
"date": date_str,
|
||||
"total": len(items),
|
||||
"ingested": ingested,
|
||||
"failed": failed,
|
||||
"elapsed_sec": elapsed,
|
||||
}
|
||||
|
||||
|
||||
def search_news(
|
||||
query: str,
|
||||
*,
|
||||
top_k: int = 10,
|
||||
search_filter: SearchFilter | None = None,
|
||||
score_threshold: float | None = None,
|
||||
) -> list[SearchResult]:
|
||||
"""语义搜索新闻。
|
||||
|
||||
流程:
|
||||
1. 将查询文本向量化（使用 M5 Embedding 服务）
|
||||
2. Qdrant 语义检索
|
||||
|
||||
Args:
|
||||
query: 中文搜索查询
|
||||
top_k: 返回条数
|
||||
search_filter: 可选过滤条件
|
||||
score_threshold: 最低相似度阈值
|
||||
|
||||
Returns:
|
||||
SearchResult 列表
|
||||
"""
|
||||
# 1. 向量化查询
|
||||
emb_config = load_embedding_config()
|
||||
emb_client = make_embedding_client(emb_config)
|
||||
|
||||
try:
|
||||
vectors = embed_batch(emb_client, emb_config, [query])
|
||||
if not vectors:
|
||||
logger.error("查询向量化失败")
|
||||
return []
|
||||
query_vector = vectors[0]
|
||||
finally:
|
||||
emb_client.close()
|
||||
|
||||
# 2. Qdrant 检索
|
||||
qdrant = make_qdrant_client()
|
||||
store = VectorStore(qdrant)
|
||||
|
||||
try:
|
||||
results = store.query(
|
||||
query_vector=query_vector,
|
||||
top_k=top_k,
|
||||
search_filter=search_filter,
|
||||
score_threshold=score_threshold,
|
||||
)
|
||||
finally:
|
||||
store.close()
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def get_collection_info() -> dict:
|
||||
"""获取 Qdrant collection 信息。"""
|
||||
client = make_qdrant_client()
|
||||
store = VectorStore(client)
|
||||
try:
|
||||
info = store.info()
|
||||
return info.model_dump()
|
||||
finally:
|
||||
store.close()
|
||||
Reference in New Issue
Block a user