nex_docus/backend/app/services/rag_service.py

"""
知识库 RAG 服务 - ZVec 向量检索 + 大模型生成
"""
import re
import math
import logging
from typing import AsyncIterator, List, Dict, Any
from pathlib import Path

from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
import jieba

from app.models.llm_model_config import LLMModelConfig
from app.core.config import settings
from app.services.llm_provider_service import LLMProviderService
from app.services.zvec_service import zvec_service
from app.services.storage import storage_service

logger = logging.getLogger(__name__)

# 单个命中分块注入上下文时，围绕命中位置向前后各扩展的字符数。
# 分块后每次只注入命中段落及其上下文，而非整篇文档，从根本上避免长文档被截断。
CHUNK_CONTEXT_WINDOW = 1200
MAX_CHUNKS_PER_DOCUMENT = 3
HISTORY_USER_QUESTION_LIMIT = 3
HISTORY_USER_QUESTION_MAX_CHARS = 300

# 查询关键词提取时过滤的常见提问词/虚词，避免命中高亮时噪声过大
QUERY_STOPWORDS = {
    "的", "了", "吗", "呢", "啊", "哦", "嗯", "吧", "是", "在", "有", "和",
    "与", "或", "及", "都", "也", "很", "这", "那", "你", "我", "他", "她",
    "它", "们", "什么", "怎么", "如何", "怎样", "哪些", "哪个", "哪几",
    "是否", "为什么", "一个", "这个", "那个", "你们", "我们", "他们",
    "请问", "一下", "可以", "没有", "不是", "就是", "还是", "或者",
    "包括", "包含", "关于", "对于", "以及", "其中", "请",
    "what", "which", "who", "whom", "whose", "when", "where", "why",
    "how", "is", "are", "was", "were", "do", "does", "did", "will",
    "would", "should", "could", "can", "may", "might", "the", "a",
    "an", "of", "to", "in", "for", "on", "with", "and", "or", "not",
    "please", "tell", "show", "list", "about",
}

# 引用支撑句对齐参数
CITATION_QUOTE_MIN_SCORE = 0.45
CITATION_QUOTE_MAX_PER_REF = 3
CITATION_QUOTE_MAX_PER_OCCURRENCE = 1
CITATION_CANDIDATES_PER_REF = 4
CITATION_CLAIM_MAX_CHARS = 200


class RAGService:
    """RAG 知识库检索和生成服务"""

    @staticmethod
    def _extract_query_terms(query: str, max_terms: int = 8) -> List[str]:
        """从用户提问中提取用于命中高亮的关键词（jieba 分词 + 停用词过滤）。"""
        text = (query or "").strip()
        if not text:
            return []
        terms: List[str] = []
        seen = set()
        for word in jieba.cut_for_search(text):
            word = word.strip()
            if not word or word in QUERY_STOPWORDS or word.lower() in QUERY_STOPWORDS:
                continue
            if len(word) < 2:
                continue
            if not re.search(r"[\u4e00-\u9fffA-Za-z0-9]", word):
                continue
            if word not in seen:
                seen.add(word)
                terms.append(word)
            if len(terms) >= max_terms:
                break
        return terms

    @staticmethod
    async def retrieve_documents(
        db: AsyncSession,
        project_id: int,
        query: str,
        top_k: int = 5,
    ) -> List[Dict[str, Any]]:
        """通过 ZVec 向量检索相关文档，并读取文档内容"""
        # 多取一些分块，再按文件去重，避免高频命中文件挤掉其他相关文档。
        matched = await zvec_service.search_similar(db, project_id, query, top_k * 4)

        if not matched:
            return []

        hit_terms = RAGService._extract_query_terms(query)

        from app.models.project import Project
        stmt = select(Project).where(Project.id == project_id)
        result = await db.execute(stmt)
        project = result.scalar_one_or_none()
        if not project:
            return []

        # 读取文件内容做缓存，避免同一文件多个分块命中时重复读盘
        file_content_cache: Dict[str, str] = {}

        async def _read_content(path: str) -> str:
            if path in file_content_cache:
                return file_content_cache[path]
            full_path = storage_service.get_secure_path(project.storage_key, path)
            text = await storage_service.read_file(full_path)
            file_content_cache[path] = text
            return text

        docs_by_path: Dict[str, Dict[str, Any]] = {}
        for item in matched:
            file_path = item["file_path"]
            chunk_text = item.get("chunk_text") or ""
            try:
                content = await _read_content(file_path)
                excerpt = RAGService._resolve_excerpt(
                    content=content,
                    chunk_text=chunk_text,
                    chunk_index=item.get("chunk_index", 0),
                )
                snippet = RAGService._extract_chunk_context(content, chunk_text, excerpt)
                doc = {
                    "file_path": file_path,
                    "file_name": Path(file_path).name,
                    "chunk_index": item.get("chunk_index", 0),
                    "anchor_text": RAGService._build_anchor(excerpt or chunk_text),
                    "score": item.get("score", 0.0),
                    "excerpt": excerpt,
                    "content": snippet,
                    "hit_terms": hit_terms,
                    "_merged_chunks": 1,
                }
            except Exception:
                doc = {
                    "file_path": file_path,
                    "file_name": Path(file_path).name,
                    "chunk_index": item.get("chunk_index", 0),
                    "anchor_text": RAGService._build_anchor(chunk_text),
                    "score": item.get("score", 0.0),
                    "excerpt": "",
                    "content": "",
                    "hit_terms": hit_terms,
                    "_merged_chunks": 1,
                }

            existing = docs_by_path.get(file_path)
            if existing is None:
                docs_by_path[file_path] = doc
                continue
            if existing["_merged_chunks"] >= MAX_CHUNKS_PER_DOCUMENT:
                continue

            existing["score"] = max(existing.get("score", 0.0), doc.get("score", 0.0))
            existing["excerpt"] = RAGService._merge_text_blocks(
                existing.get("excerpt", ""), doc.get("excerpt", "")
            )
            existing["content"] = RAGService._merge_text_blocks(
                existing.get("content", ""), doc.get("content", "")
            )
            existing["_merged_chunks"] += 1

        docs_with_content = list(docs_by_path.values())
        docs_with_content = docs_with_content[:top_k]
        for citation_id, doc in enumerate(docs_with_content, 1):
            doc.pop("_merged_chunks", None)
            doc["citation_id"] = citation_id

        return docs_with_content

    @staticmethod
    def _merge_text_blocks(existing: str, incoming: str) -> str:
        """合并同一文件的命中内容，避免重复分块生成多个引用。"""
        current = (existing or "").strip()
        candidate = (incoming or "").strip()
        if not candidate or candidate in current:
            return current
        if not current or current in candidate:
            return candidate
        return f"{current}\n\n{candidate}"

    @staticmethod
    def _resolve_excerpt(content: str, chunk_text: str, chunk_index: int) -> str:
        """返回引用预览片段。

        document_vector.chunk_text 是向量命中的分块文本，应直接作为引用预览。
        仅在 chunk_text 为空时，才按 chunk_index 从原文反推作为兜底。
        """
        text = (chunk_text or "").strip()
        if text:
            return text
        return RAGService._extract_chunk_by_index(
            content,
            chunk_index,
            settings.CHUNK_SIZE,
            settings.CHUNK_OVERLAP,
        ) or text

    @staticmethod
    def _build_anchor(text: str, max_chars: int = 120) -> str:
        for raw_line in (text or "").splitlines():
            line = raw_line.strip()
            if line:
                return line[:max_chars]
        return (text or "").strip()[:max_chars]

    @staticmethod
    def _extract_chunk_by_index(
        content: str,
        chunk_index: int,
        chunk_size: int,
        overlap: int,
    ) -> str:
        """按向量化时的分块参数，从原文切出命中的原始分块。"""
        text = content or ""
        if not text:
            return ""
        chunk_size = max(1, int(chunk_size or 1))
        overlap = max(0, min(int(overlap or 0), chunk_size - 1))
        step = chunk_size - overlap
        start = max(0, int(chunk_index or 0)) * step
        if start >= len(text):
            return ""
        return text[start:start + chunk_size]

    @staticmethod
    def _extract_chunk_context(content: str, anchor: str, fallback: str = "") -> str:
        """围绕命中锚点截取上下文窗口。

        用锚点在原文中定位命中位置，向前后各扩展 CHUNK_CONTEXT_WINDOW 字符，
        避免注入整篇长文档。定位失败时回退为命中的原始分块。
        """
        text = content or ""
        if not text:
            return ""
        pos = text.find(anchor) if anchor else -1
        if pos < 0:
            return fallback or text[:CHUNK_CONTEXT_WINDOW * 2]
        start = max(0, pos - CHUNK_CONTEXT_WINDOW)
        end = min(len(text), pos + len(anchor) + CHUNK_CONTEXT_WINDOW)
        return text[start:end]

    @staticmethod
    def _split_sentences(text: str) -> List[str]:
        """按句末标点与换行切分句子，过滤过短片段。"""
        parts = re.split(r"(?<=[。！？!?；;])\s*|\r?\n+", text or "")
        sentences = []
        for part in parts:
            part = part.strip()
            if len(part) >= 4:
                sentences.append(part)
        return sentences

    @staticmethod
    def _extract_citation_claims(answer: str) -> Dict[int, List[str]]:
        """按出现顺序提取答案中每个 [n] 标记前的一句作为待对齐的论断文本。

        保留全部出现（不做去重），保证列表下标与答案中标记的出现顺序一一对应，
        供前端按出现位置精确定位每条引用的支撑句。
        """
        claims: Dict[int, List[str]] = {}
        answer = answer or ""
        for match in re.finditer(r"\[\d+\]", answer):
            citation_id = int(match.group()[1:-1])
            prefix = answer[:match.start()]
            # 从上一个引用标记或最近的句末标点之后取论断句，支持句中标注
            delimiter_matches = list(re.finditer(r"[。！？!?；;\n]", prefix))
            last_delim_end = delimiter_matches[-1].end() if delimiter_matches else 0
            last_marker_end = prefix.rfind("]") + 1 if "]" in prefix else 0
            start = max(last_delim_end, last_marker_end)
            claim = prefix[start:].strip(" \t，,、:：")
            if len(claim) < 4:
                tail = re.sub(r"\[\d+\]", "", prefix).strip()
                claim = tail[-CITATION_CLAIM_MAX_CHARS:] if tail else ""
            if not claim:
                continue
            claim = re.sub(r"\[\d+\]", "", claim).strip()[:CITATION_CLAIM_MAX_CHARS]
            bucket = claims.setdefault(citation_id, [])
            bucket.append(claim)
        return claims

    @staticmethod
    def _sentence_candidates(
        ref_text: str,
        claims: List[str],
        max_candidates: int = CITATION_CANDIDATES_PER_REF,
    ) -> List[str]:
        """从引用文本中筛选候选支撑句：优先取与论断词重叠多的句子。"""
        sentences = RAGService._split_sentences(ref_text)
        if not sentences:
            return []

        claim_terms = set()
        for claim in claims or []:
            for word in jieba.cut_for_search(claim):
                word = word.strip()
                if len(word) >= 2 and word not in QUERY_STOPWORDS:
                    claim_terms.add(word)
        if not claim_terms:
            return sentences[:max_candidates]

        scored = []
        for sentence in sentences:
            terms = set()
            for word in jieba.cut_for_search(sentence):
                word = word.strip()
                if len(word) >= 2 and word not in QUERY_STOPWORDS:
                    terms.add(word)
            overlap = len(terms & claim_terms)
            if overlap > 0:
                scored.append((overlap, sentence))
        scored.sort(key=lambda item: item[0], reverse=True)
        chosen = [sentence for _, sentence in scored[:max_candidates]]
        # 词重叠为空时兜底取前几句，交给语义对齐判断
        return chosen or sentences[:max_candidates]

    @staticmethod
    def _cosine_similarity(a: List[float], b: List[float]) -> float:
        if not a or not b or len(a) != len(b):
            return 0.0
        dot = sum(x * y for x, y in zip(a, b))
        norm_a = math.sqrt(sum(x * x for x in a))
        norm_b = math.sqrt(sum(y * y for y in b))
        if norm_a == 0 or norm_b == 0:
            return 0.0
        return dot / (norm_a * norm_b)

    @classmethod
    async def align_citation_quotes(
        cls,
        db: AsyncSession,
        answer: str,
        refs: List[Dict[str, Any]],
    ) -> List[Dict[str, Any]]:
        """为引用回填支撑句：答案论断句 ↔ 原文候选句做向量对齐。

        检索命中是语义匹配，答案又是 LLM 重写的，因此无法用词法精确对应。
        这里把答案中每个 [n] 出现位置前的答案句子与对应引用分块内的句子
        分别向量化，取相似度最高且超过阈值的原文句子作为「支撑句」。

        同文件多处引用会合并为同一个 citation_id，因此支撑句必须按「出现
        顺序」逐一回填：ref.quote_occurrences[k] 对应答案中该引用编号的第
        k 次出现，供前端按出现位置精确展示；ref.quotes 保留聚合结果作兜底。
        """
        if not refs:
            return refs
        try:
            claims_by_id = cls._extract_citation_claims(answer)
            if not claims_by_id:
                return refs

            ref_by_id = {
                int(ref.get("citation_id")): ref
                for ref in refs
                if ref.get("citation_id") is not None
            }
            if not ref_by_id:
                return refs

            # 去重后用于向量化，occurrence 级别的对齐仍按下标一一对应
            all_claims: List[str] = []
            claim_index: Dict[str, int] = {}
            for bucket in claims_by_id.values():
                for claim in bucket:
                    if claim not in claim_index:
                        claim_index[claim] = len(all_claims)
                        all_claims.append(claim)

            candidates_by_id: Dict[int, List[str]] = {}
            for citation_id, ref in ref_by_id.items():
                ref_text = ref.get("excerpt") or ref.get("content") or ""
                candidates = cls._sentence_candidates(
                    ref_text, claims_by_id.get(citation_id, [])
                )
                if candidates:
                    candidates_by_id[citation_id] = candidates
            if not candidates_by_id:
                return refs

            texts_to_embed: List[str] = list(all_claims)
            for candidates in candidates_by_id.values():
                for candidate in candidates:
                    if candidate not in texts_to_embed:
                        texts_to_embed.append(candidate)

            vectors = await zvec_service.generate_embeddings(db, texts_to_embed)
            if not vectors or len(vectors) != len(texts_to_embed):
                logger.warning(
                    "Citation quote alignment skipped: embedding unavailable"
                )
                return refs

            text_index = {text: idx for idx, text in enumerate(texts_to_embed)}
            for citation_id, ref in ref_by_id.items():
                occurrences = []
                all_quotes = []
                candidates = candidates_by_id.get(citation_id, [])
                for claim in claims_by_id.get(citation_id, []):
                    claim_vec = vectors[claim_index[claim]]
                    scored = []
                    for candidate in candidates:
                        score = cls._cosine_similarity(
                            claim_vec, vectors[text_index[candidate]]
                        )
                        if score >= CITATION_QUOTE_MIN_SCORE:
                            scored.append((score, candidate))
                    scored.sort(key=lambda item: item[0], reverse=True)
                    occ_quotes = []
                    for score, candidate_text in scored[:CITATION_QUOTE_MAX_PER_OCCURRENCE]:
                        if not any(q["text"] == candidate_text for q in occ_quotes):
                            occ_quotes.append({
                                "text": candidate_text,
                                "score": round(score, 4),
                            })
                        if not any(q["text"] == candidate_text for q in all_quotes):
                            all_quotes.append({
                                "text": candidate_text,
                                "score": round(score, 4),
                            })
                    occurrences.append({"claim": claim, "quotes": occ_quotes})
                all_quotes.sort(key=lambda q: q["score"], reverse=True)
                ref["quote_occurrences"] = occurrences
                ref["quotes"] = all_quotes[:CITATION_QUOTE_MAX_PER_REF]
            return refs
        except Exception as exc:  # noqa: BLE001
            logger.warning(f"Citation quote alignment failed: {exc}")
            return refs

    @staticmethod
    async def generate_response(
        db: AsyncSession,
        query: str,
        project_id: int,
        llm_config_id: int,
        retrieved_docs: List[Dict[str, Any]],
        conversation_history: List[Dict[str, str]],
    ) -> str:
        """基于检索文档生成对话回复"""
        llm_config, system_prompt, messages = await RAGService._prepare_generation(
            db, query, llm_config_id, retrieved_docs, conversation_history
        )

        response_text = await LLMProviderService.generate_text(
            provider=llm_config.provider,
            endpoint_url=llm_config.endpoint_url,
            api_key=llm_config.api_key,
            llm_model_name=llm_config.llm_model_name,
            timeout=llm_config.llm_timeout,
            temperature=float(llm_config.llm_temperature),
            top_p=float(llm_config.llm_top_p),
            max_tokens=llm_config.llm_max_tokens,
            system_prompt=system_prompt,
            messages=messages,
        )

        return response_text

    @staticmethod
    async def generate_response_stream(
        db: AsyncSession,
        query: str,
        project_id: int,
        llm_config_id: int,
        retrieved_docs: List[Dict[str, Any]],
        conversation_history: List[Dict[str, str]],
    ) -> AsyncIterator[str]:
        """基于检索文档流式生成对话回复"""
        llm_config, system_prompt, messages = await RAGService._prepare_generation(
            db, query, llm_config_id, retrieved_docs, conversation_history
        )

        async for chunk in LLMProviderService.generate_text_stream(
            provider=llm_config.provider,
            endpoint_url=llm_config.endpoint_url,
            api_key=llm_config.api_key,
            llm_model_name=llm_config.llm_model_name,
            timeout=llm_config.llm_timeout,
            temperature=float(llm_config.llm_temperature),
            top_p=float(llm_config.llm_top_p),
            max_tokens=llm_config.llm_max_tokens,
            system_prompt=system_prompt,
            messages=messages,
        ):
            yield chunk

    @staticmethod
    async def _prepare_generation(
        db: AsyncSession,
        query: str,
        llm_config_id: int,
        retrieved_docs: List[Dict[str, Any]],
        conversation_history: List[Dict[str, str]],
    ):
        stmt = select(LLMModelConfig).where(LLMModelConfig.config_id == llm_config_id)
        result = await db.execute(stmt)
        llm_config = result.scalar_one_or_none()

        if not llm_config:
            raise ValueError(f"LLM配置不存在: {llm_config_id}")

        context_text = RAGService._build_context(retrieved_docs)
        citation_hint = RAGService._build_citation_hint(retrieved_docs)
        history_hint = RAGService._build_history_hint(conversation_history)
        history_section = (
            "\n\n先前用户问题（仅用于理解当前问题中的指代，不是待回答任务）:\n"
            f"{history_hint}"
            if history_hint
            else ""
        )
        system_prompt = f"""你是一个知识库助手。基于用户提供的知识库文档，回答用户的问题。

如果知识库中没有相关信息，请明确说明。

知识库文档内容:
{context_text}

引用规则:
{citation_hint}{history_section}

回答要求:
1. 仅基于知识库内容回答，不要编造。
2. 每使用一条来自文档的知识，都必须在该句末尾紧跟标注引用编号，格式为半角方括号数字（如 [1]），多个来源连写（如 [1][3]）。
3. 引用编号必须与「知识库文档内容」中的 [编号] 一一对应，只标注实际引用的文档。
4. 如果没有找到依据，请直接说明未检索到相关内容，此时不要标注引用。
5. 只回答消息列表中最后一个用户问题；先前用户问题仅用于理解指代。
6. 禁止复述、总结或继续回答先前问题，也不要重复先前助手的答案。

请基于以上知识库内容，用中文回答用户的问题。"""

        # 历史助手答案不能进入模型消息，否则部分兼容模型会在新回答中复述它。
        messages = [{"role": "user", "content": query}]
        return llm_config, system_prompt, messages

    @staticmethod
    def _build_history_hint(conversation_history: List[Dict[str, str]]) -> str:
        """仅保留最近用户问题作为指代提示，不注入历史助手答案。"""
        questions = []
        for message in conversation_history or []:
            if message.get("role") != "user":
                continue
            content = str(message.get("content") or "").strip()
            if content:
                questions.append(content[:HISTORY_USER_QUESTION_MAX_CHARS])

        recent_questions = questions[-HISTORY_USER_QUESTION_LIMIT:]
        return "\n".join(f"- {question}" for question in recent_questions)

    @staticmethod
    def _build_context(retrieved_docs: List[Dict[str, Any]]) -> str:
        """构建上下文文本"""
        if not retrieved_docs:
            return "（未找到相关知识库内容）"

        context_parts = []
        for i, doc_info in enumerate(retrieved_docs, 1):
            content = doc_info.get("content", "").strip()
            if not content:
                continue
            citation_id = doc_info.get("citation_id", i)
            context_parts.append(
                f"--- [{citation_id}] 文档: {doc_info['file_path']} ---\n{content}"
            )

        return "\n\n".join(context_parts) if context_parts else "（未找到相关知识库内容）"

    @staticmethod
    def _build_citation_hint(retrieved_docs: List[Dict[str, Any]]) -> str:
        """构建引用提示"""
        if not retrieved_docs:
            return "未检索到文档时不要标注引用。"

        return "\n".join(
            f"[{doc.get('citation_id', index)}] {doc.get('file_name') or doc.get('file_path')}"
            for index, doc in enumerate(retrieved_docs, 1)
        )


rag_service = RAGService()