ailine/backend/app/rag/retriever.py

"""
Qdrant 向量检索器模块

提供基于 Qdrant 的基础向量检索和混合检索（Dense + Sparse）功能。

核心原理：
- 使用 langchain-qdrant 的 RetrievalMode
- Qdrant 原生混合检索（如果集合已配置 sparse_vectors）
- 如果集合未配置，优雅回退到纯稠密检索
- 完全兼容现有代码，无接口改动

使用示例：
    >>> from app.rag.retriever import create_hybrid_retriever
    >>> retriever = create_hybrid_retriever(collection_name="my_docs")
    >>> docs = retriever.invoke("什么是 RAG？")
"""

from typing import Dict, Any, Optional
from qdrant_client import QdrantClient
from qdrant_client.http.exceptions import UnexpectedResponse
from langchain_qdrant import (
    QdrantVectorStore,
    RetrievalMode,
)
from langchain_core.embeddings import Embeddings
from langchain_core.retrievers import BaseRetriever

from rag_core import QDRANT_URL, QDRANT_API_KEY
from rag_core.client import create_qdrant_client as create_core_qdrant_client
from app.model_services import get_embedding_service
from app.logger import info, warning

# 模块级常量
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,
    embeddings: Embeddings | None = None,
) -> BaseRetriever:
    """
    创建基础向量检索器（仅稠密向量检索）

    Args:
        collection_name: Qdrant 集合名称
        search_kwargs: 搜索参数
        client: 可选的 Qdrant 客户端
        embeddings: 可选的嵌入模型（默认使用 get_embedding_service()）

    Returns:
        LangChain 兼容的检索器
    """
    # 默认使用统一嵌入服务（已内置降级机制）
    if embeddings is None:
        embeddings = get_embedding_service()
        info("✅ 使用统一嵌入服务（本地 llama.cpp → 智谱 API 自动降级）")

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


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,
) -> BaseRetriever:
    """
    创建混合检索器（使用 Qdrant 自身的 RetrievalMode.HYBRID）。

    ⚡️ Qdrant 原生混合检索：
    - 如果 Qdrant 集合已配置 sparse_vectors：启用 Qdrant 原生混合检索
    - 如果未配置：优雅回退到纯稠密检索
    - 完全兼容现有代码，接口不变

    Args:
        collection_name: Qdrant 集合名称。
        dense_k: 稠密向量检索返回数量，默认 10。
        sparse_k: 稀疏向量检索返回数量，默认 10。
        score_threshold: 相似度阈值，默认 0.3。
        client: 可选的 Qdrant 客户端实例。
        embeddings: 可选的嵌入模型实例。若未提供，将自动获取统一嵌入服务。

    Returns:
        BaseRetriever 实例，配置了混合搜索参数。
    """
    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()

    # 验证集合是否存在
    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

    # 检查 Qdrant 集合是否有稀疏向量配置
    sparse_available = False
    try:
        collection_info = client.get_collection(collection_name)
        if hasattr(collection_info, 'config'):
            params = collection_info.config.params
            if hasattr(params, 'sparse_vectors') and params.sparse_vectors:
                sparse_available = True
                info("✅ 检测到 Qdrant 集合有稀疏向量配置，启用 Qdrant 原生混合检索")
    except Exception as e:
        warning(f"⚠️ 检查 Qdrant 集合稀疏向量配置失败: {e}")

    # 如果有稀疏向量配置，用 Qdrant 原生混合检索
    if sparse_available:
        try:
            vector_store = QdrantVectorStore(
                client=client,
                collection_name=collection_name,
                embedding=embeddings,
                retrieval_mode=RetrievalMode.HYBRID,
            )
            info(f"✅ Qdrant 原生混合检索器初始化成功 (k={total_k})")
            return vector_store.as_retriever(search_kwargs=search_kwargs)
        except Exception as e:
            warning(f"⚠️ Qdrant 原生混合检索初始化失败: {e}，回退到纯稠密检索")

    # 如果没有稀疏向量配置，回退到纯稠密检索
    info("ℹ️ Qdrant 集合未配置稀疏向量，使用纯稠密检索（完全兼容）")
    return create_base_retriever(
        collection_name=collection_name,
        search_kwargs=search_kwargs,
        client=client,
        embeddings=embeddings,
    )


# 可选：提供异步友好的辅助函数
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)