feat: 实现 BM25 稀疏 + 稠密向量混合检索功能
Some checks failed
构建并部署 AI Agent 服务 / deploy (push) Has been cancelled
Some checks failed
构建并部署 AI Agent 服务 / deploy (push) Has been cancelled
This commit is contained in:
@@ -1,170 +1,379 @@
|
||||
"""
|
||||
Qdrant 向量检索器模块
|
||||
Qdrant 混合检索器模块
|
||||
|
||||
提供基于 Qdrant 的混合检索（Dense + Sparse）功能。
|
||||
提供基于 Qdrant 的混合检索（Dense + Sparse）功能，包括：
|
||||
- 纯混合检索（无子父文档）
|
||||
- 父子文档混合检索（先检索子文档，再返回父文档）
|
||||
|
||||
核心原理：
|
||||
- 使用 Qdrant 原生混合检索（langchain-qdrant 的 RetrievalMode.HYBRID）
|
||||
- 同时存储稠密向量和稀疏向量
|
||||
- 语义理解 + 关键词匹配，效果最优
|
||||
|
||||
使用示例：
|
||||
>>> from app.rag.retriever import create_hybrid_retriever
|
||||
>>> retriever = create_hybrid_retriever(collection_name="rag_documents")
|
||||
>>> docs = retriever.invoke("什么是 RAG？")
|
||||
- 使用 Qdrant 原生 Fusion API (RRF) 做分数融合
|
||||
- 同时使用稠密向量（语义）和稀疏向量（BM25 关键词）
|
||||
"""
|
||||
|
||||
from typing import Dict, Any, Optional
|
||||
from typing import Dict, Any, Optional, List
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.http.exceptions import UnexpectedResponse
|
||||
from langchain_qdrant import (
|
||||
QdrantVectorStore,
|
||||
RetrievalMode,
|
||||
FastEmbedSparse,
|
||||
from qdrant_client.http.models import (
|
||||
SearchRequest, Fusion, FusionProtocol, NamedVector, NamedSparseVector
|
||||
)
|
||||
from langchain_core.documents import Document
|
||||
from langchain_core.embeddings import Embeddings
|
||||
from langchain_core.retrievers import BaseRetriever
|
||||
from langchain_core.retrievers import BaseRetriever, RetrieverOutput
|
||||
|
||||
from rag_core import QDRANT_URL, QDRANT_API_KEY
|
||||
from rag_core import QdrantVectorStore, get_sparse_embedder, create_docstore
|
||||
from rag_core.client import create_qdrant_client as create_core_qdrant_client
|
||||
from app.model_services import get_embedding_service
|
||||
from app.config import SPARSE_MODEL_PATH, SPARSE_MODEL_NAME
|
||||
from app.logger import info, warning
|
||||
from app.logger import info, warning, debug
|
||||
|
||||
|
||||
# 模块级常量
|
||||
DEFAULT_SEARCH_K = 20
|
||||
DEFAULT_SCORE_THRESHOLD = 0.3
|
||||
DEFAULT_PARENT_SEARCH_K = 5
|
||||
|
||||
|
||||
def create_base_retriever(
|
||||
collection_name: str,
|
||||
search_kwargs: Dict[str, Any] | None = None,
|
||||
client: QdrantClient | None = None,
|
||||
embeddings: Embeddings | None = None,
|
||||
) -> BaseRetriever:
|
||||
class HybridRetriever(BaseRetriever):
|
||||
"""
|
||||
创建基础向量检索器（仅稠密向量检索）
|
||||
|
||||
Args:
|
||||
collection_name: Qdrant 集合名称
|
||||
search_kwargs: 搜索参数
|
||||
client: 可选的 Qdrant 客户端
|
||||
embeddings: 可选的嵌入模型（默认使用 get_embedding_service()）
|
||||
|
||||
Returns:
|
||||
LangChain 兼容的检索器
|
||||
混合检索器：稠密向量 + BM25 稀疏向量 RRF 分数融合
|
||||
|
||||
直接使用 Qdrant 原生 Fusion API，性能最优。
|
||||
"""
|
||||
# 默认使用统一嵌入服务（已内置降级机制）
|
||||
if embeddings is None:
|
||||
embeddings = get_embedding_service()
|
||||
info("✅ 使用统一嵌入服务（本地 llama.cpp → 智谱 API 自动降级）")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
collection_name: str,
|
||||
vector_store: QdrantVectorStore,
|
||||
search_k: int = DEFAULT_SEARCH_K,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
collection_name: Qdrant 集合名称
|
||||
vector_store: QdrantVectorStore 实例
|
||||
search_k: 检索返回结果数
|
||||
"""
|
||||
self.collection_name = collection_name
|
||||
self.vector_store = vector_store
|
||||
self.search_k = search_k
|
||||
self.client = vector_store.get_qdrant_client()
|
||||
self.sparse_embedder = get_sparse_embedder()
|
||||
|
||||
def _get_relevant_documents(
|
||||
self, query: str, *, run_manager: Optional[Any] = None
|
||||
) -> List[Document]:
|
||||
"""
|
||||
同步检索相关文档
|
||||
|
||||
Args:
|
||||
query: 查询字符串
|
||||
run_manager: LangChain 运行管理器（可选）
|
||||
|
||||
Returns:
|
||||
相关文档列表
|
||||
"""
|
||||
# 生成双向量
|
||||
dense_query = self.vector_store.embeddings.embed_query(query)
|
||||
sparse_query = self.sparse_embedder.embed_query(query)
|
||||
|
||||
# 构建双检索请求
|
||||
searches = [
|
||||
# 稠密检索
|
||||
SearchRequest(
|
||||
vector=NamedVector(name="dense", vector=dense_query),
|
||||
limit=self.search_k,
|
||||
with_payload=True
|
||||
),
|
||||
# 稀疏检索
|
||||
SearchRequest(
|
||||
vector=NamedSparseVector(name="sparse", vector=sparse_query),
|
||||
limit=self.search_k,
|
||||
with_payload=True
|
||||
)
|
||||
]
|
||||
|
||||
# RRF 分数融合
|
||||
fused_results = self.client.fusion(
|
||||
collection_name=self.collection_name,
|
||||
requests=searches,
|
||||
fusion=Fusion(fusion=FusionProtocol.RRF)
|
||||
)
|
||||
|
||||
# 转换为 Document 格式
|
||||
results = []
|
||||
for point in fused_results.points:
|
||||
doc = Document(
|
||||
page_content=point.payload.pop("text", ""),
|
||||
metadata=point.payload
|
||||
)
|
||||
results.append(doc)
|
||||
|
||||
debug(f"混合检索返回 {len(results)} 个文档")
|
||||
return results
|
||||
|
||||
async def _aget_relevant_documents(
|
||||
self, query: str, *, run_manager: Optional[Any] = None
|
||||
) -> List[Document]:
|
||||
"""异步检索（当前调用同步版本）"""
|
||||
# Qdrant 客户端没有原生 async，这里用同步版本
|
||||
return self._get_relevant_documents(query, run_manager=run_manager)
|
||||
|
||||
# 合并默认搜索参数
|
||||
merged_search_kwargs = {"k": DEFAULT_SEARCH_K}
|
||||
if search_kwargs:
|
||||
merged_search_kwargs.update(search_kwargs)
|
||||
|
||||
# 创建或复用 Qdrant 客户端
|
||||
if client is None:
|
||||
client = create_core_qdrant_client()
|
||||
|
||||
# 验证集合是否存在
|
||||
try:
|
||||
client.get_collection(collection_name)
|
||||
except UnexpectedResponse as e:
|
||||
if e.status_code == 404:
|
||||
warning(f"⚠️ Qdrant 集合 '{collection_name}' 不存在，请先创建并索引文档")
|
||||
raise ValueError(f"Qdrant 集合 '{collection_name}' 不存在")
|
||||
raise
|
||||
|
||||
# 构建向量存储
|
||||
vector_store = QdrantVectorStore(
|
||||
client=client,
|
||||
collection_name=collection_name,
|
||||
embedding=embeddings,
|
||||
)
|
||||
|
||||
return vector_store.as_retriever(search_kwargs=merged_search_kwargs)
|
||||
class ParentHybridRetriever(BaseRetriever):
|
||||
"""
|
||||
父子文档混合检索器：
|
||||
|
||||
1. 先用混合检索找到相关子文档
|
||||
2. 根据子文档的 parent_id 找到对应的父文档
|
||||
3. 去重并返回父文档
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
collection_name: str,
|
||||
vector_store: QdrantVectorStore,
|
||||
search_k: int = DEFAULT_PARENT_SEARCH_K,
|
||||
docstore: Optional[Any] = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
collection_name: Qdrant 集合名称
|
||||
vector_store: QdrantVectorStore 实例
|
||||
search_k: 最终返回的父文档数
|
||||
docstore: 文档存储（如果父文档在 PostgreSQL），可选
|
||||
"""
|
||||
self.collection_name = collection_name
|
||||
self.vector_store = vector_store
|
||||
self.search_k = search_k
|
||||
self.client = vector_store.get_qdrant_client()
|
||||
self.sparse_embedder = get_sparse_embedder()
|
||||
self.docstore = docstore
|
||||
|
||||
def _get_relevant_documents(
|
||||
self, query: str, *, run_manager: Optional[Any] = None
|
||||
) -> List[Document]:
|
||||
"""
|
||||
同步检索相关父文档
|
||||
|
||||
Args:
|
||||
query: 查询字符串
|
||||
run_manager: LangChain 运行管理器（可选）
|
||||
|
||||
Returns:
|
||||
相关父文档列表
|
||||
"""
|
||||
# 1. 生成查询双向量
|
||||
dense_query = self.vector_store.embeddings.embed_query(query)
|
||||
sparse_query = self.sparse_embedder.embed_query(query)
|
||||
|
||||
# 2. 多取一些子文档，避免去重后数量不足
|
||||
search_limit = self.search_k * 2
|
||||
searches = [
|
||||
# 稠密检索
|
||||
SearchRequest(
|
||||
vector=NamedVector(name="dense", vector=dense_query),
|
||||
limit=search_limit,
|
||||
with_payload=True
|
||||
),
|
||||
# 稀疏检索
|
||||
SearchRequest(
|
||||
vector=NamedSparseVector(name="sparse", vector=sparse_query),
|
||||
limit=search_limit,
|
||||
with_payload=True
|
||||
)
|
||||
]
|
||||
|
||||
# 3. RRF 分数融合，拿到子文档命中结果
|
||||
fused_results = self.client.fusion(
|
||||
collection_name=self.collection_name,
|
||||
requests=searches,
|
||||
fusion=Fusion(fusion=FusionProtocol.RRF)
|
||||
)
|
||||
|
||||
if not fused_results.points:
|
||||
debug("混合检索未找到任何文档")
|
||||
return []
|
||||
|
||||
# 4. 收集 parent_id 和对应最高得分
|
||||
parent_score_map = {}
|
||||
parent_ids = set()
|
||||
child_point_map = {} # 保存子文档点用于降级
|
||||
|
||||
for point in fused_results.points:
|
||||
parent_id = point.payload.get("parent_id", point.id)
|
||||
score = point.score
|
||||
|
||||
# 同一个 parent_id 只保留最高得分
|
||||
if parent_id not in parent_score_map or score > parent_score_map[parent_id]:
|
||||
parent_score_map[parent_id] = score
|
||||
parent_ids.add(parent_id)
|
||||
child_point_map[parent_id] = point
|
||||
|
||||
# 5. 批量查询父文档
|
||||
# 首先尝试从 Qdrant 直接查询（因为父文档可能也存在 Qdrant 中）
|
||||
parent_docs = []
|
||||
found_parent_ids = set()
|
||||
|
||||
try:
|
||||
parent_points = self.client.retrieve(
|
||||
collection_name=self.collection_name,
|
||||
ids=list(parent_ids),
|
||||
with_payload=True
|
||||
)
|
||||
|
||||
# 处理找到的父文档
|
||||
for point in parent_points:
|
||||
doc = Document(
|
||||
page_content=point.payload.pop("text", ""),
|
||||
metadata=point.payload
|
||||
)
|
||||
parent_docs.append(doc)
|
||||
found_parent_ids.add(point.id)
|
||||
|
||||
except Exception as e:
|
||||
warning(f"从 Qdrant 查询父文档失败: {e}")
|
||||
|
||||
# 6. 如果有 docstore，尝试从 docstore 查询剩余的父文档
|
||||
if self.docstore and len(found_parent_ids) < len(parent_ids):
|
||||
missing_parent_ids = parent_ids - found_parent_ids
|
||||
try:
|
||||
docstore_docs = self.docstore.mget(missing_parent_ids)
|
||||
for doc_id, doc in zip(missing_parent_ids, docstore_docs):
|
||||
if doc is not None:
|
||||
parent_docs.append(doc)
|
||||
found_parent_ids.add(doc_id)
|
||||
except Exception as e:
|
||||
warning(f"从 docstore 查询父文档失败: {e}")
|
||||
|
||||
# 7. 降级：对于仍未找到的父文档，用子文档本身代替
|
||||
missing_parent_ids = parent_ids - found_parent_ids
|
||||
if missing_parent_ids:
|
||||
warning(f"以下 parent_id 未找到对应的父文档，将返回子文档本身: {missing_parent_ids}")
|
||||
for parent_id in missing_parent_ids:
|
||||
child_point = child_point_map.get(parent_id)
|
||||
if child_point:
|
||||
doc = Document(
|
||||
page_content=child_point.payload.pop("text", ""),
|
||||
metadata=child_point.payload
|
||||
)
|
||||
parent_docs.append(doc)
|
||||
|
||||
# 8. 按照得分降序排序，返回前 k 个
|
||||
parent_docs_with_scores = [
|
||||
(doc, parent_score_map.get(doc.metadata.get("id", doc.id), 0.0))
|
||||
for doc in parent_docs
|
||||
]
|
||||
parent_docs_with_scores.sort(key=lambda x: x[1], reverse=True)
|
||||
|
||||
final_docs = [doc for doc, _ in parent_docs_with_scores[:self.search_k]]
|
||||
debug(f"父子文档混合检索返回 {len(final_docs)} 个父文档")
|
||||
|
||||
return final_docs
|
||||
|
||||
async def _aget_relevant_documents(
|
||||
self, query: str, *, run_manager: Optional[Any] = None
|
||||
) -> List[Document]:
|
||||
"""异步检索（当前调用同步版本）"""
|
||||
return self._get_relevant_documents(query, run_manager=run_manager)
|
||||
|
||||
|
||||
def create_hybrid_retriever(
|
||||
collection_name: str,
|
||||
dense_k: int = 10,
|
||||
sparse_k: int = 10,
|
||||
score_threshold: float | None = DEFAULT_SCORE_THRESHOLD,
|
||||
client: QdrantClient | None = None,
|
||||
embeddings: Embeddings | None = None,
|
||||
search_k: int = DEFAULT_SEARCH_K,
|
||||
embeddings: Optional[Embeddings] = None,
|
||||
) -> BaseRetriever:
|
||||
"""
|
||||
创建混合检索器（稠密向量 + BM25 稀疏向量，Qdrant 原生实现）。
|
||||
|
||||
创建混合检索器（稠密向量 + BM25 稀疏向量）。
|
||||
|
||||
这是默认推荐的检索方式，效果最优。
|
||||
|
||||
Args:
|
||||
collection_name: Qdrant 集合名称。
|
||||
dense_k: 稠密向量检索返回数量，默认 10。
|
||||
sparse_k: 稀疏向量检索返回数量，默认 10。
|
||||
score_threshold: 相似度阈值，默认 0.3。
|
||||
client: 可选的 Qdrant 客户端实例。
|
||||
collection_name: Qdrant 集合名称
|
||||
search_k: 检索返回结果数
|
||||
embeddings: 可选的嵌入模型实例。若未提供，将自动获取统一嵌入服务。
|
||||
|
||||
|
||||
Returns:
|
||||
BaseRetriever 实例，配置了混合搜索参数。
|
||||
HybridRetriever 实例
|
||||
"""
|
||||
total_k = dense_k + sparse_k
|
||||
|
||||
search_kwargs = {
|
||||
"k": total_k,
|
||||
"search_type": "similarity_score_threshold",
|
||||
"score_threshold": score_threshold,
|
||||
}
|
||||
|
||||
# 默认使用统一嵌入服务（已内置降级机制）
|
||||
# 默认使用统一嵌入服务
|
||||
if embeddings is None:
|
||||
embeddings = get_embedding_service()
|
||||
info("✅ 使用统一嵌入服务（本地 llama.cpp → 智谱 API 自动降级）")
|
||||
|
||||
# 创建或复用 Qdrant 客户端
|
||||
if client is None:
|
||||
client = create_core_qdrant_client()
|
||||
|
||||
|
||||
# 创建向量存储
|
||||
vector_store = QdrantVectorStore(collection_name=collection_name, embeddings=embeddings)
|
||||
|
||||
# 验证集合是否存在
|
||||
try:
|
||||
client.get_collection(collection_name)
|
||||
vector_store.get_client().get_collection(collection_name)
|
||||
except UnexpectedResponse as e:
|
||||
if e.status_code == 404:
|
||||
warning(f"⚠️ Qdrant 集合 '{collection_name}' 不存在，请先创建并索引文档")
|
||||
raise ValueError(f"Qdrant 集合 '{collection_name}' 不存在")
|
||||
raise
|
||||
|
||||
# 初始化稀疏嵌入（使用本地缓存目录）
|
||||
sparse_embeddings = FastEmbedSparse(
|
||||
model_name=SPARSE_MODEL_NAME,
|
||||
cache_dir=SPARSE_MODEL_PATH
|
||||
)
|
||||
info(f"✅ FastEmbedSparse 初始化成功 (cache_dir={SPARSE_MODEL_PATH})")
|
||||
|
||||
# 创建混合模式的 QdrantVectorStore
|
||||
vector_store = QdrantVectorStore(
|
||||
client=client,
|
||||
|
||||
info(f"✅ Qdrant 混合检索器初始化成功（search_k={search_k}）")
|
||||
return HybridRetriever(
|
||||
collection_name=collection_name,
|
||||
embedding=embeddings,
|
||||
sparse_embedding=sparse_embeddings,
|
||||
retrieval_mode=RetrievalMode.HYBRID,
|
||||
vector_store=vector_store,
|
||||
search_k=search_k
|
||||
)
|
||||
|
||||
info(f"✅ Qdrant 原生混合检索器初始化成功 (k={total_k})")
|
||||
return vector_store.as_retriever(search_kwargs=search_kwargs)
|
||||
|
||||
|
||||
# 可选：提供异步友好的辅助函数
|
||||
async def acreate_base_retriever(
|
||||
def create_parent_hybrid_retriever(
|
||||
collection_name: str,
|
||||
search_kwargs: Dict[str, Any] | None = None,
|
||||
client: QdrantClient | None = None,
|
||||
search_k: int = DEFAULT_PARENT_SEARCH_K,
|
||||
embeddings: Optional[Embeddings] = None,
|
||||
use_docstore: bool = True,
|
||||
) -> BaseRetriever:
|
||||
"""
|
||||
异步创建基础向量检索器（与同步版本功能相同）。
|
||||
|
||||
适用于需要异步初始化的场景（例如在 FastAPI 启动事件中）。
|
||||
创建父子文档混合检索器（默认推荐）。
|
||||
|
||||
检索流程：
|
||||
1. 混合检索找到相关子文档
|
||||
2. 根据 parent_id 找到对应的父文档
|
||||
3. 去重并返回父文档
|
||||
|
||||
Args:
|
||||
collection_name: Qdrant 集合名称
|
||||
search_k: 最终返回的父文档数
|
||||
embeddings: 可选的嵌入模型实例
|
||||
use_docstore: 是否使用 PostgreSQL docstore 存储父文档
|
||||
|
||||
Returns:
|
||||
ParentHybridRetriever 实例
|
||||
"""
|
||||
# 由于 QdrantVectorStore 初始化本身是同步的，这里直接调用同步版本即可
|
||||
return create_base_retriever(collection_name, search_kwargs, client)
|
||||
# 默认使用统一嵌入服务
|
||||
if embeddings is None:
|
||||
embeddings = get_embedding_service()
|
||||
info("✅ 使用统一嵌入服务（本地 llama.cpp → 智谱 API 自动降级）")
|
||||
|
||||
# 创建向量存储
|
||||
vector_store = QdrantVectorStore(collection_name=collection_name, embeddings=embeddings)
|
||||
|
||||
# 验证集合是否存在
|
||||
try:
|
||||
vector_store.get_client().get_collection(collection_name)
|
||||
except UnexpectedResponse as e:
|
||||
if e.status_code == 404:
|
||||
warning(f"⚠️ Qdrant 集合 '{collection_name}' 不存在，请先创建并索引文档")
|
||||
raise ValueError(f"Qdrant 集合 '{collection_name}' 不存在")
|
||||
raise
|
||||
|
||||
# 创建 docstore（如果需要）
|
||||
docstore = None
|
||||
if use_docstore:
|
||||
try:
|
||||
docstore, _ = create_docstore()
|
||||
info("✅ 文档存储初始化成功（PostgreSQL）")
|
||||
except Exception as e:
|
||||
warning(f"⚠️ 文档存储初始化失败，将不使用 docstore: {e}")
|
||||
|
||||
info(f"✅ Qdrant 父子文档混合检索器初始化成功（search_k={search_k}）")
|
||||
return ParentHybridRetriever(
|
||||
collection_name=collection_name,
|
||||
vector_store=vector_store,
|
||||
search_k=search_k,
|
||||
docstore=docstore
|
||||
)
|
||||
|
||||
|
||||
# 别名：默认就是父子文档混合检索
|
||||
create_retriever = create_parent_hybrid_retriever
|
||||
|
||||
Reference in New Issue
Block a user