ailine/backend/app/rag/retriever.py

"""
Qdrant 向量检索器模块

提供基于 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 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, LlamaCppEmbedder
from rag_core.client import create_qdrant_client as create_core_qdrant_client

# 模块级常量
DEFAULT_SEARCH_K = 20
DEFAULT_SCORE_THRESHOLD = 0.3


def create_base_retriever(
    collection_name: str,
    search_kwargs: Dict[str, Any] | None = None,
    client: QdrantClient | None = None,
) -> BaseRetriever:
    """
    创建基础向量检索器（仅稠密向量检索）。

    该检索器使用嵌入模型将查询转为向量，在 Qdrant 集合中执行 ANN 搜索，
    返回语义上最相似的文档块。

    Args:
        collection_name: Qdrant 集合名称（需预先创建并索引）。
        search_kwargs: 搜索参数，可包含：
            - k (int): 返回的文档数量，默认 20。
            - score_threshold (float): 相似度阈值，仅返回高于此分数的文档。
            - filter (dict): Qdrant 过滤条件。
            若为 None，则使用默认值 {"k": 20}。
        client: 可选的 Qdrant 客户端实例。若未提供，将自动创建。

    Returns:
        BaseRetriever 实例，可直接调用 .invoke(query) 或 .ainvoke(query) 检索。

    Raises:
        ValueError: 如果集合不存在或嵌入模型无效。
    """
    # 嵌入模型
    embedder = LlamaCppEmbedder()
    embeddings = embedder.as_langchain_embeddings()

    # 合并默认搜索参数
    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:
            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)


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,
) -> BaseRetriever:
    """
    创建混合检索器（稠密向量 + BM25 稀疏向量）。

    混合检索结合了语义相似度（Dense）和关键词匹配（Sparse），
    能够更好地处理专有名词、精确匹配等场景。

    注意：此功能要求 Qdrant 集合已配置稀疏向量字段并生成了 BM25 索引。
    若集合未配置稀疏向量，将回退到纯稠密检索（不会报错，但检索效果降级）。

    Args:
        collection_name: Qdrant 集合名称。
        dense_k: 稠密向量检索返回数量，默认 10。
        sparse_k: 稀疏向量检索返回数量，默认 10。
        score_threshold: 相似度阈值，默认 0.3。
        client: 可选的 Qdrant 客户端实例。

    Returns:
        BaseRetriever 实例，配置了混合搜索参数。
    """
    total_k = dense_k + sparse_k

    search_kwargs = {
        "k": total_k,
    }
    if score_threshold is not None:
        search_kwargs["score_threshold"] = score_threshold

    # 复用基础检索器创建逻辑，只需调整搜索参数
    return create_base_retriever(
        collection_name=collection_name,
        search_kwargs=search_kwargs,
        client=client,
    )


# 可选：提供异步友好的辅助函数
async def acreate_base_retriever(
    collection_name: str,
    search_kwargs: Dict[str, Any] | None = None,
    client: QdrantClient | None = None,
) -> BaseRetriever:
    """
    异步创建基础向量检索器（与同步版本功能相同）。

    适用于需要异步初始化的场景（例如在 FastAPI 启动事件中）。
    """
    # 由于 QdrantVectorStore 初始化本身是同步的，这里直接调用同步版本即可
    return create_base_retriever(collection_name, search_kwargs, client)