ailine/backend/app/rag/pipeline.py

"""
RAG 检索流水线
流程: 检索子文档 → 重排 → 获取父文档 → 置信度评估 → 返回
"""

import asyncio
import re
from dataclasses import dataclass
from typing import List
from langchain_core.documents import Document
from langchain_core.language_models import BaseLanguageModel

from backend.app.logger import info, warning
from ..model_services import get_small_llm_service
from ..rag.rerank import create_document_reranker
from ..rag.query_transform import MultiQueryGenerator
from ..rag.fusion import reciprocal_rank_fusion
from ..rag.retriever import create_parent_hybrid_retriever


@dataclass
class RAGResult:
    """RAG 检索结果（包含置信度）"""
    content: str                       # 格式化后的上下文
    documents: List[Document]         # 原始文档
    confidence: float                 # 综合置信度 0.0-1.0
    scores: dict                     # 各维度分数 {embedding, rerank, llm, final}
    is_useful: bool                  # 是否可用（confidence >= 0.6）


class RAGPipeline:
    def __init__(
        self,
        retriever=None,
        llm: BaseLanguageModel | str = "default_small",
        num_queries: int = 3,
        rerank_top_n: int = 5,
        collection_name: str = "rag_documents",
        use_rerank: bool = True,
        return_parent_docs: bool = True,
        confidence_threshold: float = 0.6,
    ):
        self.retriever = retriever or create_parent_hybrid_retriever(
            collection_name=collection_name, search_k=rerank_top_n * 4
        )
        self.num_queries = num_queries
        self.rerank_top_n = rerank_top_n
        self.use_rerank = use_rerank
        self.return_parent_docs = return_parent_docs
        self.confidence_threshold = confidence_threshold
        self._last_docs: List[Document] = []
        self._last_scores: List[dict] = []

        if llm == "default_small":
            try:
                self.llm = get_small_llm_service()
            except Exception:
                self.llm = None
        else:
            self.llm = llm if llm else None

        self.query_generator = MultiQueryGenerator(self.llm, num_queries) if self.llm else None
        self.reranker = create_document_reranker() if use_rerank else None
        info(f"[Pipeline] init: rerank={use_rerank}, return_parent={return_parent_docs}, threshold={confidence_threshold}")

    @property
    def last_docs(self) -> List[Document]:
        return self._last_docs

    @property
    def last_scores(self) -> List[dict]:
        return self._last_scores

    async def aretrieve(self, query: str) -> List[Document]:
        """原接口，保持向后兼容"""
        docs = await self._do_retrieve(query)
        self._last_docs = docs
        self._last_scores = self._extract_scores(docs)
        return docs

    async def aretrieve_with_confidence(self, query: str, original_query: str = "") -> RAGResult:
        """
        带置信度评估的检索

        Args:
            query: 检索查询
            original_query: 原始用户问题（用于置信度评估）

        Returns:
            RAGResult: 包含内容和置信度的结构化结果
        """
        info(f"[Pipeline] aretrieve_with_confidence: query={query[:50]}...")

        # 1. 执行检索
        docs = await self._do_retrieve(query)
        self._last_docs = docs
        self._last_scores = self._extract_scores(docs)

        # 2. 格式化内容
        content = self.format_context(docs)

        if not docs or not content:
            info(f"[Pipeline] 无检索结果，置信度=0")
            return RAGResult(
                content="",
                documents=[],
                confidence=0.0,
                scores={"embedding": 0.0, "rerank": 0.0, "llm": 0.0, "final": 0.0},
                is_useful=False
            )

        # 3. 评估置信度（三维度）
        scores = await self._evaluate_confidence(
            query=original_query or query,
            docs=docs,
            content=content
        )

        confidence = scores["final"]
        is_useful = confidence >= self.confidence_threshold

        info(f"[Pipeline] 置信度评估完成: confidence={confidence:.3f}, is_useful={is_useful}")

        return RAGResult(
            content=content,
            documents=docs,
            confidence=confidence,
            scores=scores,
            is_useful=is_useful
        )

    async def _do_retrieve(self, query: str) -> List[Document]:
        """执行检索流程"""
        # Step 1: 检索
        child_docs = await self._retrieve(query)

        # Step 1.5: 向量初筛
        vector_top_n = 20
        if len(child_docs) > vector_top_n:
            child_docs = child_docs[:vector_top_n]

        # Step 2: 重排
        if self.reranker:
            try:
                child_docs = self.reranker.compress_documents(child_docs, query, self.rerank_top_n)
            except Exception as e:
                warning(f"[Pipeline] 重排失败: {e}")
                child_docs = child_docs[:self.rerank_top_n]

        # Step 3: 获取父文档
        if self.return_parent_docs:
            return await self._get_parents(child_docs)
        return child_docs

    async def _evaluate_confidence(
        self,
        query: str,
        docs: List[Document],
        content: str
    ) -> dict:
        """
        三维度置信度评估

        Returns:
            {
                "embedding": float,  # 向量相似度 (0-1)
                "rerank": float,     # 重排分数 (0-1)
                "llm": float,        # LLM判断 (0-1)
                "final": float       # 综合分数 (0-1)
            }
        """
        scores = {"embedding": 0.0, "rerank": 0.0, "llm": 0.5, "final": 0.0}

        # 1. 向量相似度
        if docs:
            embedding_scores = []
            for doc in docs:
                score = doc.metadata.get("embedding_score", doc.metadata.get("score", 0.0))
                # 归一化（如果分数 > 1）
                if score > 1.0:
                    score = min(score / 10.0, 1.0)
                embedding_scores.append(score)
            scores["embedding"] = max(embedding_scores) if embedding_scores else 0.0

        info(f"[Confidence] embedding={scores['embedding']:.3f}")

        # 2. 重排分数
        if docs:
            rerank_scores = []
            for doc in docs:
                score = doc.metadata.get("rerank_score", 0.0)
                # 归一化（假设满分 10）
                if score > 1.0:
                    score = min(score / 10.0, 1.0)
                rerank_scores.append(score)
            scores["rerank"] = max(rerank_scores) if rerank_scores else 0.0

        info(f"[Confidence] rerank={scores['rerank']:.3f}")

        # 3. LLM 判断
        if self.llm and content:
            llm_score = await self._get_llm_confidence(query, content)
            scores["llm"] = llm_score
            info(f"[Confidence] llm={scores['llm']:.3f}")

        # 4. 综合得分（加权平均）
        scores["final"] = (
            scores["embedding"] * 0.25 +
            scores["rerank"] * 0.25 +
            scores["llm"] * 0.50
        )

        info(f"[Confidence] final={scores['final']:.3f}")

        return scores

    async def _get_llm_confidence(self, query: str, context: str) -> float:
        """使用 LLM 评估检索结果相关性"""
        try:
            prompt = f"""评估以下检索结果与用户问题的相关性，返回 0.0-1.0 的分数：
- 1.0 = 完全相关，能直接回答问题
- 0.7 = 高度相关，有很大参考价值
- 0.5 = 部分相关，有一定参考价值
- 0.3 = 低度相关，参考价值有限
- 0.0 = 完全不相关，无法回答问题

用户问题：{query}

检索结果：{context[:1500]}

只返回一个数字（0.0-1.0）："""

            response = await self.llm.ainvoke(prompt)
            content = response.content.strip()

            match = re.search(r'(\d+\.?\d*)', content)
            if match:
                score = float(match.group(1))
                return max(0.0, min(1.0, score))

        except Exception as e:
            info(f"[Confidence] LLM评估失败: {e}")

        return 0.5  # 默认中等置信度

    def _extract_scores(self, docs: List[Document]) -> List[dict]:
        """提取文档的分数信息"""
        scores = []
        for doc in docs:
            scores.append({
                "embedding_score": doc.metadata.get("embedding_score", doc.metadata.get("score", 0.0)),
                "rerank_score": doc.metadata.get("rerank_score", 0.0),
            })
        return scores

    async def _retrieve(self, query: str) -> List[Document]:
        if self.query_generator:
            queries = await self.query_generator.agenerate(query)
            queries = [query] + [q for q in queries if q != query]
            doc_lists = await asyncio.gather(*[self.retriever.ainvoke(q) for q in queries])
            return reciprocal_rank_fusion(doc_lists)
        return await self.retriever.ainvoke(query)

    async def _get_parents(self, child_docs: List[Document]) -> List[Document]:
        parent_map = {}
        for doc in child_docs:
            pid = doc.metadata.get("parent_id")
            if pid and pid not in parent_map:
                embedding_score = doc.metadata.get("score", 0.0)
                rerank_score = doc.metadata.get("rerank_score", 0.0)
                parent_map[pid] = (embedding_score, rerank_score)

        if not parent_map:
            return child_docs

        try:
            from backend.rag_core import create_docstore
            docstore, _ = create_docstore()
            parent_docs = await docstore.amget(list(parent_map.keys()))

            result = []
            for doc in parent_docs:
                if doc:
                    pid = doc.metadata.get("id")
                    scores = parent_map.get(pid, (0.0, 0.0))
                    doc.metadata["embedding_score"] = scores[0]
                    doc.metadata["rerank_score"] = scores[1]
                    # 综合分数，rerank 权重更高
                    result.append((doc, scores[0] + scores[1] * 2))

            result.sort(key=lambda x: x[1], reverse=True)
            return [d for d, _ in result]
        except Exception as e:
            warning(f"[Pipeline] 获取父文档失败: {e}")
            return child_docs

    def format_context(self, documents: List[Document]) -> str:
        if not documents:
            return ""
        parts = []
        for i, doc in enumerate(documents, 1):
            source = doc.metadata.get("source", "未知来源")
            parts.append(f"【资料 {i}】来源：{source}\n{doc.page_content}\n---\n")
        return "\n".join(parts)


def create_rag_pipeline(**kwargs) -> RAGPipeline:
    return RAGPipeline(**kwargs)