This commit is contained in:
@@ -1,39 +1,83 @@
|
||||
"""
|
||||
Qdrant 向量检索器
|
||||
Qdrant 向量检索器模块
|
||||
|
||||
提供基础向量检索、混合检索（Dense + BM25）功能。
|
||||
提供基于 Qdrant 的基础向量检索和混合检索（Dense + Sparse）功能。
|
||||
|
||||
核心原理：
|
||||
- 基础检索：将查询文本转换为向量，在 Qdrant 中进行近似最近邻（ANN）搜索，
|
||||
使用余弦相似度返回最相似的 k 个文档。
|
||||
- 混合检索：结合稠密向量检索（语义相似）和 BM25 稀疏向量检索（关键词匹配），
|
||||
通过加权或分数融合提高召回精度。
|
||||
|
||||
使用示例：
|
||||
>>> from rag_core import LlamaCppEmbedder
|
||||
>>> embedder = LlamaCppEmbedder()
|
||||
>>> embeddings = embedder.as_langchain_embeddings()
|
||||
>>>
|
||||
>>> # 创建基础检索器
|
||||
>>> retriever = create_base_retriever(
|
||||
... collection_name="my_docs",
|
||||
... embeddings=embeddings,
|
||||
... search_kwargs={"k": 10}
|
||||
... )
|
||||
>>>
|
||||
>>> # 执行检索
|
||||
>>> docs = retriever.invoke("什么是 RAG？")
|
||||
"""
|
||||
|
||||
from typing import List, Dict, Any, Optional
|
||||
from langchain_qdrant import QdrantVectorStore
|
||||
from langchain.embeddings.base import Embeddings
|
||||
# from langchain.retrievers import EnsembleRetriever
|
||||
from typing import Optional, Dict, Any
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.http.exceptions import UnexpectedResponse
|
||||
from langchain_qdrant import QdrantVectorStore
|
||||
from langchain_core.embeddings import Embeddings
|
||||
from langchain_core.retrievers import BaseRetriever
|
||||
|
||||
from rag_core import QDRANT_URL, QDRANT_API_KEY
|
||||
|
||||
# 模块级常量
|
||||
DEFAULT_SEARCH_K = 20
|
||||
DEFAULT_SCORE_THRESHOLD = 0.3
|
||||
|
||||
|
||||
def create_qdrant_client(
|
||||
url: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
timeout: int = 30,
|
||||
) -> QdrantClient:
|
||||
"""
|
||||
创建 Qdrant 客户端
|
||||
创建并返回一个配置好的 Qdrant 客户端。
|
||||
|
||||
优先使用传入参数，若未提供则回退到环境变量 QDRANT_URL 和 QDRANT_API_KEY。
|
||||
|
||||
Args:
|
||||
url: Qdrant 服务地址，默认从环境变量 QDRANT_URL 读取
|
||||
api_key: API 密钥，默认从环境变量 QDRANT_API_KEY 读取
|
||||
url: Qdrant 服务地址，例如 "http://localhost:6333"。
|
||||
默认从环境变量 QDRANT_URL 读取。
|
||||
api_key: API 密钥（若 Qdrant 启用了认证）。
|
||||
默认从环境变量 QDRANT_API_KEY 读取。
|
||||
timeout: 请求超时时间（秒），默认 30 秒。
|
||||
|
||||
Returns:
|
||||
QdrantClient 实例
|
||||
配置好的 QdrantClient 实例。
|
||||
|
||||
Raises:
|
||||
ValueError: 如果 url 为空且环境变量也未设置。
|
||||
"""
|
||||
url = url or QDRANT_URL
|
||||
api_key = api_key or QDRANT_API_KEY
|
||||
effective_url = url or QDRANT_URL
|
||||
if not effective_url:
|
||||
raise ValueError(
|
||||
"Qdrant URL 未提供，请设置参数 url 或环境变量 QDRANT_URL"
|
||||
)
|
||||
|
||||
client_args = {"url": url}
|
||||
if api_key:
|
||||
client_args["api_key"] = api_key
|
||||
effective_api_key = api_key or QDRANT_API_KEY
|
||||
|
||||
return QdrantClient(**client_args)
|
||||
client_kwargs = {
|
||||
"url": effective_url,
|
||||
"timeout": timeout,
|
||||
}
|
||||
if effective_api_key:
|
||||
client_kwargs["api_key"] = effective_api_key
|
||||
|
||||
return QdrantClient(**client_kwargs)
|
||||
|
||||
|
||||
def create_base_retriever(
|
||||
@@ -41,33 +85,57 @@ def create_base_retriever(
|
||||
embeddings: Embeddings,
|
||||
search_kwargs: Optional[Dict[str, Any]] = None,
|
||||
client: Optional[QdrantClient] = None,
|
||||
) -> QdrantVectorStore:
|
||||
) -> BaseRetriever:
|
||||
"""
|
||||
创建基础向量检索器
|
||||
创建基础向量检索器（仅稠密向量检索）。
|
||||
|
||||
该检索器使用嵌入模型将查询转为向量，在 Qdrant 集合中执行 ANN 搜索，
|
||||
返回语义上最相似的文档块。
|
||||
|
||||
Args:
|
||||
collection_name: Qdrant 集合名称
|
||||
embeddings: 嵌入模型
|
||||
search_kwargs: 搜索参数，默认 {"k": 20}
|
||||
client: Qdrant 客户端，如果为 None 则自动创建
|
||||
collection_name: Qdrant 集合名称（需预先创建并索引）。
|
||||
embeddings: LangChain 兼容的嵌入模型实例。
|
||||
search_kwargs: 搜索参数，可包含：
|
||||
- k (int): 返回的文档数量，默认 20。
|
||||
- score_threshold (float): 相似度阈值，仅返回高于此分数的文档。
|
||||
- filter (dict): Qdrant 过滤条件。
|
||||
若为 None，则使用默认值 {"k": 20}。
|
||||
client: 可选的 Qdrant 客户端实例。若未提供，将自动创建。
|
||||
|
||||
Returns:
|
||||
QdrantVectorStore 检索器实例
|
||||
"""
|
||||
search_kwargs = search_kwargs or {"k": 20}
|
||||
BaseRetriever 实例，可直接调用 .invoke(query) 或 .ainvoke(query) 检索。
|
||||
|
||||
# 创建 Qdrant 客户端
|
||||
Raises:
|
||||
ValueError: 如果集合不存在或嵌入模型无效。
|
||||
"""
|
||||
# 合并默认搜索参数
|
||||
merged_search_kwargs = {"k": DEFAULT_SEARCH_K}
|
||||
if search_kwargs:
|
||||
merged_search_kwargs.update(search_kwargs)
|
||||
|
||||
# 创建或复用 Qdrant 客户端
|
||||
if client is None:
|
||||
client = create_qdrant_client()
|
||||
|
||||
# 使用 QdrantVectorStore 创建向量存储
|
||||
# 验证集合是否存在（可选，便于提前发现问题）
|
||||
try:
|
||||
client.get_collection(collection_name)
|
||||
except UnexpectedResponse as e:
|
||||
if e.status_code == 404:
|
||||
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=search_kwargs)
|
||||
# 返回检索器
|
||||
return vector_store.as_retriever(search_kwargs=merged_search_kwargs)
|
||||
|
||||
|
||||
def create_hybrid_retriever(
|
||||
@@ -75,64 +143,57 @@ def create_hybrid_retriever(
|
||||
embeddings: Embeddings,
|
||||
dense_k: int = 10,
|
||||
sparse_k: int = 10,
|
||||
score_threshold: Optional[float] = DEFAULT_SCORE_THRESHOLD,
|
||||
client: Optional[QdrantClient] = None,
|
||||
) -> QdrantVectorStore:
|
||||
) -> BaseRetriever:
|
||||
"""
|
||||
创建混合检索器（Dense Vector + BM25）
|
||||
创建混合检索器（稠密向量 + BM25 稀疏向量）。
|
||||
|
||||
混合检索结合了语义相似度（Dense）和关键词匹配（Sparse），
|
||||
能够更好地处理专有名词、精确匹配等场景。
|
||||
|
||||
注意：此功能要求 Qdrant 集合已配置稀疏向量字段并生成了 BM25 索引。
|
||||
若集合未配置稀疏向量，将回退到纯稠密检索（不会报错，但检索效果降级）。
|
||||
|
||||
Args:
|
||||
collection_name: Qdrant 集合名称
|
||||
embeddings: 嵌入模型
|
||||
dense_k: 向量检索返回数量
|
||||
sparse_k: BM25 检索返回数量
|
||||
client: Qdrant 客户端
|
||||
collection_name: Qdrant 集合名称。
|
||||
embeddings: 嵌入模型（用于稠密向量）。
|
||||
dense_k: 稠密向量检索返回数量，默认 10。
|
||||
sparse_k: 稀疏向量检索返回数量，默认 10。
|
||||
score_threshold: 相似度阈值，默认 0.3。
|
||||
client: 可选的 Qdrant 客户端实例。
|
||||
|
||||
Returns:
|
||||
混合检索器
|
||||
BaseRetriever 实例，配置了混合搜索参数。
|
||||
"""
|
||||
# 创建 Qdrant 客户端
|
||||
if client is None:
|
||||
client = create_qdrant_client()
|
||||
|
||||
# 使用 QdrantVectorStore 创建向量存储
|
||||
vector_store = QdrantVectorStore(
|
||||
client=client,
|
||||
collection_name=collection_name,
|
||||
embedding=embeddings,
|
||||
)
|
||||
total_k = dense_k + sparse_k
|
||||
|
||||
search_kwargs = {
|
||||
"k": dense_k + sparse_k,
|
||||
"score_threshold": 0.3,
|
||||
"k": total_k,
|
||||
}
|
||||
if score_threshold is not None:
|
||||
search_kwargs["score_threshold"] = score_threshold
|
||||
|
||||
return vector_store.as_retriever(search_kwargs=search_kwargs)
|
||||
# 复用基础检索器创建逻辑，只需调整搜索参数
|
||||
return create_base_retriever(
|
||||
collection_name=collection_name,
|
||||
embeddings=embeddings,
|
||||
search_kwargs=search_kwargs,
|
||||
client=client,
|
||||
)
|
||||
|
||||
|
||||
# def create_ensemble_retriever(
|
||||
# retrievers: List[Any],
|
||||
# weights: Optional[List[float]] = None,
|
||||
# c: int = 60,
|
||||
# ) -> EnsembleRetriever:
|
||||
# """
|
||||
# 创建集成检索器，支持倒数排名融合 (RRF)
|
||||
#
|
||||
# Args:
|
||||
# retrievers: 检索器列表
|
||||
# weights: 检索器权重
|
||||
# c: RRF 常数（通常为60）
|
||||
#
|
||||
# Returns:
|
||||
# 集成检索器
|
||||
# """
|
||||
# if weights is None:
|
||||
# weights = [1.0 / len(retrievers)] * len(retrievers)
|
||||
#
|
||||
# ensemble = EnsembleRetriever(
|
||||
# retrievers=retrievers,
|
||||
# weights=weights,
|
||||
# c=c,
|
||||
# search_type="rrf",
|
||||
# )
|
||||
#
|
||||
# return ensemble
|
||||
# 可选：提供异步友好的辅助函数
|
||||
async def acreate_base_retriever(
|
||||
collection_name: str,
|
||||
embeddings: Embeddings,
|
||||
search_kwargs: Optional[Dict[str, Any]] = None,
|
||||
client: Optional[QdrantClient] = None,
|
||||
) -> BaseRetriever:
|
||||
"""
|
||||
异步创建基础向量检索器（与同步版本功能相同）。
|
||||
|
||||
适用于需要异步初始化的场景（例如在 FastAPI 启动事件中）。
|
||||
"""
|
||||
# 由于 QdrantVectorStore 初始化本身是同步的，这里直接调用同步版本即可
|
||||
return create_base_retriever(collection_name, embeddings, search_kwargs, client)
|
||||
Reference in New Issue
Block a user