ailine/backend/app/rag/retriever.py

"""
Qdrant 混合检索器模块（完全异步）

提供基于 Qdrant 的混合检索（Dense + Sparse）功能，包括：
- 纯混合检索（无子父文档）
- 父子文档混合检索（先检索子文档，再返回父文档）

核心原理：
- 使用 Qdrant Universal Query API (query_points)
- 使用 Prefetch 并行检索多个源
- 使用 RRF 分数融合
"""

from typing import Dict, Any, Optional, List
from qdrant_client import AsyncQdrantClient, models
from qdrant_client.http.exceptions import UnexpectedResponse
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.retrievers import BaseRetriever
from pydantic import Field, PrivateAttr

from backend.rag_core import QdrantHybridStore, get_sparse_embedder, create_docstore
from backend.rag_core.client import create_async_qdrant_client
from ..model_services import get_embedding_service
from backend.app.logger import info, warning, debug


# 模块级常量
DEFAULT_SEARCH_K = 20
DEFAULT_PARENT_SEARCH_K = 5


class HybridRetriever(BaseRetriever):
    """
    混合检索器：稠密向量 + BM25 稀疏向量 RRF 分数融合（异步）
    
    使用 Qdrant Universal Query API (query_points)
    """
    collection_name: str = Field(description="Qdrant 集合名称")
    search_k: int = Field(default=DEFAULT_SEARCH_K, description="检索返回结果数")
    
    _vector_store: Any = PrivateAttr()
    _client: Any = PrivateAttr()
    _sparse_embedder: Any = PrivateAttr()
    
    def __init__(
        self,
        collection_name: str,
        vector_store: QdrantHybridStore,
        search_k: int = DEFAULT_SEARCH_K,
    ):
        """
        Args:
            collection_name: Qdrant 集合名称
            vector_store: QdrantHybridStore 实例
            search_k: 检索返回结果数
        """
        super().__init__(
            collection_name=collection_name,
            search_k=search_k
        )
        self._vector_store = vector_store
        self._client = vector_store.get_async_qdrant_client()
        self._sparse_embedder = get_sparse_embedder()
    
    def _get_relevant_documents(
        self, query: str, *, run_manager: Any = None
    ) -> List[Document]:
        """
        同步检索（不推荐使用，仅供兼容性）
        
        注意：在异步环境中请使用 _aget_relevant_documents 或 ainvoke
        """
        import asyncio
        try:
            loop = asyncio.get_running_loop()
            # 已有事件循环，使用 create_task
            task = loop.create_task(self._aget_relevant_documents(query))
            return loop.run_until_complete(task)
        except RuntimeError:
            # 没有事件循环，创建新的
            return asyncio.run(self._aget_relevant_documents(query))
    
    async def _aget_relevant_documents(
        self, query: str, *, run_manager: Any = None
    ) -> List[Document]:
        """
        异步混合检索相关文档
        """
        # 1. 生成查询向量
        dense_query = await self._vector_store.aembed_query(query)
        sparse_query = self._sparse_embedder.embed_query(query)
        sparse_vec = models.SparseVector(
            indices=sparse_query["indices"],
            values=sparse_query["values"]
        )
        
        # 2. 使用 Qdrant 的 query_points API
        response = await self._client.query_points(
            collection_name=self.collection_name,
            prefetch=[
                models.Prefetch(
                    query=dense_query,
                    using="dense",
                    limit=self.search_k
                ),
                models.Prefetch(
                    query=sparse_vec,
                    using="sparse",
                    limit=self.search_k
                )
            ],
            query=models.FusionQuery(fusion=models.Fusion.RRF),
            limit=self.search_k,
            with_payload=True
        )
        
        # 3. 转换结果
        results = []
        for point in response.points:
            doc = Document(
                page_content=point.payload.pop("page_content", point.payload.pop("text", "")),
                metadata=point.payload
            )
            results.append(doc)
        
        debug(f"混合检索返回 {len(results)} 个文档")
        return results


class ParentHybridRetriever(BaseRetriever):
    """
    父子文档混合检索器（异步）：

    1. 先用混合检索找到相关子文档
    2. 根据子文档的 parent_id 找到对应的父文档
    3. 去重并返回父文档
    """

    collection_name: str = Field(description="Qdrant 集合名称")
    search_k: int = Field(default=DEFAULT_PARENT_SEARCH_K, description="检索返回结果数")

    _vector_store: Any = PrivateAttr()
    _client: Any = PrivateAttr()
    _sparse_embedder: Any = PrivateAttr()
    _docstore: Any = PrivateAttr()

    def __init__(
        self,
        collection_name: str,
        vector_store: QdrantHybridStore,
        search_k: int = DEFAULT_PARENT_SEARCH_K,
        docstore: Optional[Any] = None,
    ):
        """
        Args:
            collection_name: Qdrant 集合名称
            vector_store: QdrantHybridStore 实例
            search_k: 最终返回的父文档数量
            docstore: 文档存储（如果父文档在 PostgreSQL），可选
        """
        super().__init__(
            collection_name=collection_name,
            search_k=search_k
        )
        self._vector_store = vector_store
        self._client = vector_store.get_async_qdrant_client()
        self._sparse_embedder = get_sparse_embedder()
        self._docstore = docstore
    
    def _get_relevant_documents(
        self, query: str, *, run_manager: Any = None
    ) -> List[Document]:
        """
        同步检索（不推荐使用，仅供兼容性）
        
        注意：在异步环境中请使用 _aget_relevant_documents 或 ainvoke
        """
        import asyncio
        try:
            loop = asyncio.get_running_loop()
            task = loop.create_task(self._aget_relevant_documents(query))
            return loop.run_until_complete(task)
        except RuntimeError:
            return asyncio.run(self._aget_relevant_documents(query))
    
    async def _aget_relevant_documents(
        self, query: str, *, run_manager: Any = None
    ) -> List[Document]:
        """
        异步检索相关子文档
        """
        # 1. 生成查询向量
        dense_query = await self._vector_store.aembed_query(query)
        sparse_query = self._sparse_embedder.embed_query(query)
        sparse_vec = models.SparseVector(
            indices=sparse_query["indices"],
            values=sparse_query["values"]
        )

        # 2. 多取一些子文档，避免去重后数量不足
        search_limit = self.search_k * 2

        # 3. 使用 query_points API 进行混合检索
        response = await self._client.query_points(
            collection_name=self.collection_name,
            prefetch=[
                models.Prefetch(
                    query=dense_query,
                    using="dense",
                    limit=search_limit
                ),
                models.Prefetch(
                    query=sparse_vec,
                    using="sparse",
                    limit=search_limit
                )
            ],
            query=models.FusionQuery(fusion=models.Fusion.RRF),
            limit=search_limit,
            with_payload=True
        )

        if not response.points:
            debug("混合检索未找到任何文档")
            return []

        # 4. 构建子文档列表
        child_docs = []
        for point in response.points:
            payload_copy = point.payload.copy()
            doc = Document(
                page_content=payload_copy.pop("page_content", payload_copy.pop("text", "")),
                metadata={
                    **payload_copy,
                    "child_id": point.id,
                    "score": point.score
                }
            )
            child_docs.append(doc)

        debug(f"父子文档混合检索返回 {len(child_docs)} 个子文档")
        return child_docs


def create_hybrid_retriever(
    collection_name: str,
    search_k: int = DEFAULT_SEARCH_K,
    embeddings: Optional[Embeddings] = None,
) -> BaseRetriever:
    """
    创建混合检索器（稠密向量 + BM25 稀疏向量）- 异步版本。

    这是默认推荐的检索方式，效果最优。

    Args:
        collection_name: Qdrant 集合名称
        search_k: 检索返回结果数
        embeddings: 可选的嵌入模型实例。若未提供，将自动获取统一嵌入服务。

    Returns:
        HybridRetriever 实例
    """
    if embeddings is None:
        embeddings = get_embedding_service()
        info("使用统一嵌入服务（本地 llama.cpp → 智谱 API 自动降级）")

    vector_store = QdrantHybridStore(collection_name=collection_name)

    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

    info(f"✅ Qdrant 混合检索器初始化成功（search_k={search_k}）")
    return HybridRetriever(
        collection_name=collection_name,
        vector_store=vector_store,
        search_k=search_k
    )


def create_parent_hybrid_retriever(
    collection_name: str,
    search_k: int = DEFAULT_PARENT_SEARCH_K,
    embeddings: Optional[Embeddings] = None,
    use_docstore: bool = True,
) -> BaseRetriever:
    """
    创建父子文档混合检索器（默认推荐）- 异步版本。

    检索流程：
    1. 混合检索找到相关子文档
    2. 根据 parent_id 找到对应的父文档
    3. 去重并返回父文档

    Args:
        collection_name: Qdrant 集合名称
        search_k: 最终返回的父文档数量
        embeddings: 可选的嵌入模型实例
        use_docstore: 是否使用 PostgreSQL docstore 存储父文档

    Returns:
        ParentHybridRetriever 实例
    """
    if embeddings is None:
        embeddings = get_embedding_service()
        info("使用统一嵌入服务（本地 llama.cpp → 智谱 API 自动降级）")

    vector_store = QdrantHybridStore(collection_name=collection_name)

    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 = None
    if use_docstore:
        try:
            docstore, _ = create_docstore()
            info("✅ 文档存储初始化成功（PostgreSQL）")
        except Exception as e:
            warning(f"⚠️ 文档存储初始化失败，将不使用 docstore: %s", e)

    info(f"✅ Qdrant 父子文档混合检索器初始化成功（search_k={search_k}）")
    return ParentHybridRetriever(
        collection_name=collection_name,
        vector_store=vector_store,
        search_k=search_k,
        docstore=docstore
    )


def create_base_retriever(
    collection_name: str,
    search_k: int = DEFAULT_SEARCH_K,
    embeddings: Optional[Embeddings] = None,
) -> BaseRetriever:
    """
    创建基础检索器（向后兼容）- 实际上返回混合检索器。
    """
    return create_hybrid_retriever(collection_name, search_k, embeddings)


# 别名：默认就是父子文档混合检索
create_retriever = create_parent_hybrid_retriever