RAG与长文本处理¶

# 分块策略对比
class ChunkingStrategies:
    """
    不同分块策略的实现与对比
    """

    @staticmethod  # @staticmethod无需实例即可调用
    def fixed_size_chunk(text, chunk_size=512, overlap=50):
        """
        固定大小分块 - 最简单，但可能切断语义单元
        """
        chunks = []
        start = 0
        while start < len(text):
            end = start + chunk_size
            chunk = text[start:end]
            chunks.append({
                'content': chunk,
                'start': start,
                'end': end
            })
            start = end - overlap  # 重叠保持上下文
        return chunks

    @staticmethod
    def recursive_character_chunk(text, chunk_size=512, separators=["\n\n", "\n", ". ", " "]):
        """
        递归字符分块 - LangChain默认策略
        优先在语义边界处分割
        """
        chunks = []
        # 实现逻辑：递归尝试不同分隔符
        # 优先使用大分隔符（段落），必要时使用小分隔符（句子、单词）
        return chunks

    @staticmethod
    def semantic_chunk(text, embedding_model, similarity_threshold=0.8):
        """
        语义分块 - 基于语义相似性动态分块
        当相邻句子语义相似度低于阈值时分割
        """
        # 1. 分割为句子
        sentences = sent_tokenize(text)

        # 2. 计算句子嵌入
        embeddings = embedding_model.encode(sentences)

        # 3. 基于相似度聚类
        chunks = []
        current_chunk = [sentences[0]]

        for i in range(1, len(sentences)):
            similarity = cosine_similarity(
                embeddings[i-1:i],
                embeddings[i:i+1]
            )[0][0]

            if similarity < similarity_threshold:
                # 语义断裂，保存当前块
                chunks.append(" ".join(current_chunk))
                current_chunk = [sentences[i]]
            else:
                current_chunk.append(sentences[i])

        if current_chunk:
            chunks.append(" ".join(current_chunk))

        return chunks

    @staticmethod
    def agentic_chunk(text, llm, target_chunk_size=512):
        """
        Agentic分块 - 使用LLM智能分块
        让模型决定最佳分割点，保持主题一致性
        """
        prompt = f"""
        将以下文本分割为语义完整的段落，每个段落约{target_chunk_size}字符。
        确保每个段落围绕单一主题，在段落之间标注[CHUNK_BOUNDARY]。

        文本：
        {text[:4000]}  # 限制输入长度
        """

        response = llm.generate(prompt)
        chunks = response.split("[CHUNK_BOUNDARY]")
        return [c.strip() for c in chunks if c.strip()]  # 链式调用：strip去除空白

from sentence_transformers import SentenceTransformer
import numpy as np

# 加载开源嵌入模型
model = SentenceTransformer('BAAI/bge-large-zh-v1.5')

# 对文档进行编码
documents = [
    "RAG是一种结合检索和生成的技术",
    "Transformer架构是现代NLP的基础",
    "向量数据库用于存储高维嵌入"
]

doc_embeddings = model.encode(documents, normalize_embeddings=True)

# 查询编码
query = "什么是RAG技术？"
query_embedding = model.encode([query], normalize_embeddings=True)

# 计算相似度
similarities = np.dot(doc_embeddings, query_embedding.T).squeeze()
ranked_indices = np.argsort(similarities)[::-1]

print("检索结果：")
for idx in ranked_indices[:3]:
    print(f"  [{similarities[idx]:.4f}] {documents[idx]}")

import faiss
import numpy as np

# 准备数据
dimension = 768  # 嵌入维度
nb = 100000      # 数据库大小
nq = 1000        # 查询数量

# 生成随机数据（实际使用真实嵌入）
xb = np.random.random((nb, dimension)).astype('float32')
xq = np.random.random((nq, dimension)).astype('float32')

# 1. Flat Index - 精确搜索，最慢但最准
index_flat = faiss.IndexFlatL2(dimension)
index_flat.add(xb)
D, I = index_flat.search(xq, k=10)  # 返回距离和索引

# 2. IVF Index - 倒排文件，加速搜索
nlist = 100  # 聚类中心数
quantizer = faiss.IndexFlatL2(dimension)
index_ivf = faiss.IndexIVFFlat(quantizer, dimension, nlist)
index_ivf.train(xb)  # 需要训练
index_ivf.add(xb)
index_ivf.nprobe = 10  # 搜索时访问的聚类数

# 3. HNSW Index - 图索引，高召回率
index_hnsw = faiss.IndexHNSWFlat(dimension, 32)  # 32为M参数
index_hnsw.hnsw.efConstruction = 200
index_hnsw.add(xb)
index_hnsw.hnsw.efSearch = 128  # 搜索时扩展因子

# 4. PQ Index - 乘积量化，内存高效
m = 16  # 子空间数
code_size = 8  # 每个子空间编码位数
index_pq = faiss.IndexPQ(dimension, m, code_size)
index_pq.train(xb)
index_pq.add(xb)

# 5. IVF-PQ - 组合索引，平衡速度与精度
index_ivfpq = faiss.IndexIVFPQ(quantizer, dimension, nlist, m, code_size)
index_ivfpq.train(xb)
index_ivfpq.add(xb)

# 性能对比
print("索引类型对比：")
print("  Flat:    精确搜索，内存占用高，速度慢")
print("  IVF:     近似搜索，速度较快，需要调nprobe")
print("  HNSW:    图索引，高召回率，构建较慢")
print("  PQ:      高压缩，内存小，精度损失较大")
print("  IVF-PQ:  综合方案，大规模数据首选")

class QueryRewriter:
    """
    查询重写技术提升检索质量
    """

    def __init__(self, llm):
        self.llm = llm

    def hyde_rewrite(self, query):
        """
        HyDE (Hypothetical Document Embeddings)
        生成假设文档，用假设文档的嵌入进行检索
        """
        prompt = f"""
        请根据以下问题，写一段可能包含答案的文档段落。
        这段文档应该看起来像是知识库中的真实内容。

        问题：{query}

        文档段落：
        """

        hypothetical_doc = self.llm.generate(prompt)
        return hypothetical_doc

    def query_expansion(self, query, num_expansions=3):
        """
        查询扩展：生成多个相关查询，合并检索结果
        """
        prompt = f"""
        针对以下问题，生成{num_expansions}个不同表述但语义相同的查询。
        每个查询关注问题的不同方面。

        原问题：{query}

        请用JSON格式返回：{{"queries": ["查询1", "查询2", ...]}}
        """

        response = self.llm.generate(prompt)
        expanded_queries = json.loads(response)["queries"]  # json.loads将JSON字符串→Python对象
        expanded_queries.append(query)  # 包含原查询

        return expanded_queries

    def step_back_prompting(self, query):
        """
        Step-back Prompting: 从具体问题抽象到一般原理
        """
        prompt = f"""
        将以下具体问题抽象为更一般的原理或概念。
        用更通用的术语重新表述。

        具体问题：{query}

        一般原理表述：
        """

        step_back_query = self.llm.generate(prompt)
        return [step_back_query, query]  # 同时检索抽象和具体层面

    def decompose_complex_query(self, query):
        """
        复杂查询分解：将多跳问题分解为子问题
        """
        prompt = f"""
        将以下复杂问题分解为2-3个简单的子问题。
        这些子问题的答案组合起来可以回答原问题。

        复杂问题：{query}

        请用JSON格式返回：{{"sub_questions": ["子问题1", "子问题2", ...]}}
        """

        response = self.llm.generate(prompt)
        sub_questions = json.loads(response)["sub_questions"]
        return sub_questions

class Reranker:
    """
    重排序技术提升最终检索质量
    """

    def __init__(self):
        # 加载重排序模型
        from sentence_transformers import CrossEncoder
        self.cross_encoder = CrossEncoder('BAAI/bge-reranker-large')

    def rerank(self, query, documents, top_k=5):
        """
        使用交叉编码器进行精确重排序
        """
        # 构建查询-文档对
        pairs = [[query, doc] for doc in documents]

        # 计算相关性分数
        scores = self.cross_encoder.predict(pairs)

        # 按分数排序
        scored_docs = list(zip(documents, scores))
        scored_docs.sort(key=lambda x: x[1], reverse=True)  # lambda匿名函数

        return scored_docs[:top_k]

    def diversity_rerank(self, query, documents, top_k=5, lambda_param=0.5):
        """
        MMR (Maximal Marginal Relevance) 多样性重排序
        平衡相关性和多样性
        """
        # 计算文档嵌入
        doc_embeddings = self.embedding_model.encode(documents)
        query_embedding = self.embedding_model.encode([query])

        # 计算与查询的相似度
        query_sims = cosine_similarity(query_embedding, doc_embeddings)[0]

        selected = []
        remaining = list(range(len(documents)))

        while len(selected) < top_k and remaining:
            if not selected:
                # 第一个选择最相关的
                best_idx = remaining[np.argmax(query_sims[remaining])]
            else:
                # MMR评分
                mmr_scores = []
                for idx in remaining:
                    # 相关性
                    relevance = query_sims[idx]
                    # 与已选文档的最大相似度（冗余度）
                    redundancy = max([
                        cosine_similarity(
                            doc_embeddings[idx:idx+1],
                            doc_embeddings[s:s+1]
                        )[0][0]
                        for s in selected
                    ]) if selected else 0

                    mmr_score = lambda_param * relevance - (1 - lambda_param) * redundancy
                    mmr_scores.append(mmr_score)

                best_idx = remaining[np.argmax(mmr_scores)]

            selected.append(best_idx)
            remaining.remove(best_idx)

        return [documents[i] for i in selected]

class ContextCompressor:
    """
    上下文压缩技术解决长文档问题
    """

    def __init__(self, llm):
        self.llm = llm

    def map_reduce_compress(self, documents, query, max_tokens=2000):
        """
        Map-Reduce压缩：先分别摘要，再合并
        """
        # Map阶段：分别摘要每个文档
        partial_summaries = []
        for doc in documents:
            prompt = f"""
            针对问题"{query}"，提取以下文档中的相关信息。
            只保留与问题相关的内容，去除冗余信息。

            文档：{doc[:1000]}

            相关信息：
            """
            summary = self.llm.generate(prompt, max_tokens=200)
            partial_summaries.append(summary)

        # Reduce阶段：合并摘要
        combined = "\n".join(partial_summaries)
        if len(combined) > max_tokens:
            prompt = f"""
            将以下信息整合为简洁的参考材料，用于回答问题"{query}"。
            去除重复信息，保持关键事实。

            信息：{combined}

            整合后的参考材料：
            """
            final = self.llm.generate(prompt, max_tokens=max_tokens)
            return final

        return combined

    def selective_context(self, documents, query, embedding_model, top_k_sentences=10):
        """
        选择性上下文：只保留最相关的句子
        """
        # 分割为句子
        all_sentences = []
        for doc in documents:
            sentences = sent_tokenize(doc)
            all_sentences.extend(sentences)

        # 计算与查询的相似度
        query_emb = embedding_model.encode([query])
        sentence_embs = embedding_model.encode(all_sentences)
        similarities = cosine_similarity(query_emb, sentence_embs)[0]

        # 选择最相关的句子
        top_indices = np.argsort(similarities)[-top_k_sentences:]
        selected_sentences = [all_sentences[i] for i in sorted(top_indices)]

        return " ".join(selected_sentences)

    def refine_compression(self, documents, query):
        """
        Refine方法：迭代构建上下文
        """
        context = ""
        for doc in documents:
            prompt = f"""
            基于已有上下文和当前文档，更新回答问题的相关信息。

            问题：{query}
            已有上下文：{context}
            当前文档：{doc[:1000]}

            更新后的相关信息（整合已有和新增）：
            """
            context = self.llm.generate(prompt, max_tokens=500)

        return context

class CitationGenerator:
    """
    生成带引用的回答，提高可验证性
    """

    def generate_with_citations(self, query, documents, llm):
        """
        生成带有明确来源引用的回答
        """
        # 为每个文档添加编号
        context = ""
        for i, doc in enumerate(documents, 1):  # enumerate同时获取索引和元素
            context += f"[文档{i}] {doc}\n\n"

        prompt = f"""
        基于以下文档回答问题。重要规则：
        1. 每个事实后必须标注来源，格式为[^文档编号^]
        2. 如果信息不在文档中，明确说明"根据提供文档无法确定"
        3. 优先使用高相关度文档的信息

        文档：
        {context}

        问题：{query}

        回答（带引用）：
        """

        response = llm.generate(prompt)
        return response

    def verify_citations(self, response, documents):
        """
        验证引用是否正确
        """
        import re

        # 提取引用
        citations = re.findall(r'\[\^(\d+)\^\]', response)  # re.findall正则查找所有匹配项

        verification_results = []
        for citation in set(citations):
            doc_idx = int(citation) - 1
            if doc_idx < 0 or doc_idx >= len(documents):
                verification_results.append({
                    'citation': citation,
                    'valid': False,
                    'reason': '文档编号超出范围'
                })
            else:
                # 提取引用附近的文本
                context = self._extract_citation_context(response, citation)
                # 检查是否确实在文档中
                doc = documents[doc_idx]
                if self._check_containment(context, doc):
                    verification_results.append({
                        'citation': citation,
                        'valid': True
                    })
                else:
                    verification_results.append({
                        'citation': citation,
                        'valid': False,
                        'reason': '引用内容未在文档中找到'
                    })

        return verification_results

# 稀疏注意力模式实现

class SparseAttention:
    """
    稀疏注意力机制降低长序列计算复杂度
    """

    @staticmethod
    def sliding_window_attention(Q, K, V, window_size=512):
        """
        滑动窗口注意力：每个token只关注附近的token
        复杂度：O(n × window_size)
        """
        batch_size, num_heads, seq_len, head_dim = Q.shape

        # 创建滑动窗口掩码
        mask = torch.zeros(seq_len, seq_len, dtype=torch.bool)
        for i in range(seq_len):
            start = max(0, i - window_size // 2)
            end = min(seq_len, i + window_size // 2 + 1)
            mask[i, start:end] = True

        # 应用掩码的注意力
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(head_dim)
        scores = scores.masked_fill(~mask, float('-inf'))
        attn = torch.softmax(scores, dim=-1)
        output = torch.matmul(attn, V)

        return output

    @staticmethod
    def dilated_attention(Q, K, V, dilation_rate=4):
        """
        空洞注意力：间隔采样，关注稀疏位置
        """
        batch_size, num_heads, seq_len, head_dim = Q.shape

        # 对K, V进行间隔采样
        K_dilated = K[:, :, ::dilation_rate, :]
        V_dilated = V[:, :, ::dilation_rate, :]

        # 计算与采样后key的注意力
        scores = torch.matmul(Q, K_dilated.transpose(-2, -1)) / math.sqrt(head_dim)
        attn = torch.softmax(scores, dim=-1)
        output = torch.matmul(attn, V_dilated)

        return output

    @staticmethod
    def global_local_attention(Q, K, V, num_global_tokens=16, local_window=512):
        """
        全局-局部混合注意力：部分token全局可见，其余局部可见
        """
        batch_size, num_heads, seq_len, head_dim = Q.shape

        # 选择全局token（如开头和结尾的特殊token）
        global_indices = list(range(num_global_tokens // 2)) + \
                        list(range(seq_len - num_global_tokens // 2, seq_len))

        # 全局token关注所有token
        # 局部token关注自己和全局token

        outputs = []
        for i in range(seq_len):
            if i in global_indices:
                # 全局注意力
                scores = torch.matmul(Q[:, :, i:i+1], K.transpose(-2, -1)) / math.sqrt(head_dim)
            else:
                # 局部注意力 + 全局token
                local_start = max(0, i - local_window // 2)
                local_end = min(seq_len, i + local_window // 2 + 1)
                attend_indices = global_indices + list(range(local_start, local_end))
                attend_indices = list(set(attend_indices))  # 去重

                K_local = K[:, :, attend_indices, :]
                scores = torch.matmul(Q[:, :, i:i+1], K_local.transpose(-2, -1)) / math.sqrt(head_dim)

            attn = torch.softmax(scores, dim=-1)
            if i in global_indices:
                out = torch.matmul(attn, V)
            else:
                V_local = V[:, :, attend_indices, :]
                out = torch.matmul(attn, V_local)

            outputs.append(out)

        return torch.cat(outputs, dim=2)

class LongContextPositionEncoding:
    """
    支持长上下文的位置编码技术
    """

    @staticmethod
    def apply_ntk_scaling(base=10000, seq_len=32768, training_len=2048, dim=128):
        """
        NTK-aware RoPE缩放
        通过调整base实现长度外推，不降低短文本性能
        """
        # NTK缩放公式
        scaling_factor = seq_len / training_len

        # 调整base
        adjusted_base = base * (scaling_factor ** (dim / (dim - 2)))

        return adjusted_base

    @staticmethod
    def yarn_scaling(seq_len=128000, training_len=4096, beta_fast=32, beta_slow=1):
        """
        YaRN (Yet another RoPE extensioN)
        结合NTK和非均匀插值
        """
        # 计算频率因子
        scale = seq_len / training_len

        # 高频和低频的不同处理
        # 高频（短波长）：直接插值
        # 低频（长波长）：NTK-aware缩放

        def yarn_factor(dim_idx, dim=128):
            freq = 1.0 / (10000 ** (dim_idx / dim))

            # 波长
            wavelength = 2 * math.pi / freq

            # 根据波长决定缩放策略
            if wavelength < beta_fast:
                return 1 / scale  # 直接插值
            elif wavelength > beta_slow * training_len:
                return 1.0  # 不缩放
            else:
                # 平滑过渡
                t = (wavelength - beta_fast) / (beta_slow * training_len - beta_fast)
                return (1 - t) / scale + t

        factors = [yarn_factor(i) for i in range(dim)]
        return factors

    @staticmethod
    def alibi_biases(seq_len, num_heads):
        """
        ALiBi (Attention with Linear Biases)
        通过距离惩罚实现外推，无需位置编码
        """
        # 创建距离矩阵
        distances = torch.arange(seq_len).unsqueeze(0) - torch.arange(seq_len).unsqueeze(1)  # unsqueeze增加一个维度
        distances = distances.abs()

        # 为每个头设置不同的惩罚斜率
        slopes = torch.tensor([
            2 ** (-8 * (head_idx + 1) / num_heads)
            for head_idx in range(num_heads)
        ])

        # 计算偏置
        biases = slopes.unsqueeze(1).unsqueeze(2) * distances.unsqueeze(0)

        return -biases  # 负值作为注意力偏置

# 完整RAG系统实现

import os
from typing import List, Dict, Any
import numpy as np
from sentence_transformers import SentenceTransformer, CrossEncoder
import faiss

class AdvancedRAGSystem:
    """
    高级RAG系统实现，包含所有优化组件
    """

    def __init__(
        self,
        embedding_model_name: str = 'BAAI/bge-large-zh-v1.5',
        reranker_model_name: str = 'BAAI/bge-reranker-large',
        llm_client = None,  # OpenAI或其他LLM客户端
        chunk_size: int = 512,
        chunk_overlap: int = 50,
        top_k_retrieve: int = 20,
        top_k_rerank: int = 5
    ):
        # 初始化模型
        self.embedding_model = SentenceTransformer(embedding_model_name)
        self.reranker = CrossEncoder(reranker_model_name)
        self.llm = llm_client

        # 配置
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.top_k_retrieve = top_k_retrieve
        self.top_k_rerank = top_k_rerank

        # 存储
        self.documents = []
        self.chunks = []
        self.index = None
        self.dimension = self.embedding_model.get_sentence_embedding_dimension()

    def add_documents(self, documents: List[str], metadata: List[Dict] = None):
        """
        添加文档到知识库
        """
        if metadata is None:
            metadata = [{} for _ in documents]

        self.documents.extend(documents)

        # 分块
        for doc_idx, (doc, meta) in enumerate(zip(documents, metadata)):  # zip配对文档与元数据，enumerate添加索引，(doc, meta)解包配对结果
            doc_chunks = self._chunk_document(doc)
            for chunk_idx, chunk in enumerate(doc_chunks):
                self.chunks.append({
                    'content': chunk,
                    'doc_idx': doc_idx,
                    'chunk_idx': chunk_idx,
                    'metadata': meta
                })

        # 构建索引
        self._build_index()

    def _chunk_document(self, document: str) -> List[str]:
        """
        递归字符分块
        """
        chunks = []
        separators = ["\n\n", "\n", "。", ". ", " ", ""]

        def recursive_split(text, sep_idx=0):
            if sep_idx >= len(separators):
                return [text]

            separator = separators[sep_idx]
            parts = text.split(separator)

            result = []
            current_chunk = ""

            for part in parts:
                if separator:
                    part = part + separator

                if len(current_chunk) + len(part) <= self.chunk_size:
                    current_chunk += part
                else:
                    if current_chunk:
                        result.append(current_chunk.strip())

                    # 如果单个部分太长，递归分割
                    if len(part) > self.chunk_size:
                        result.extend(recursive_split(part, sep_idx + 1))
                    else:
                        current_chunk = part

            if current_chunk:
                result.append(current_chunk.strip())

            return result

        return recursive_split(document)

    def _build_index(self):
        """
        构建FAISS索引
        """
        if not self.chunks:
            return

        # 编码所有块
        texts = [chunk['content'] for chunk in self.chunks]
        embeddings = self.embedding_model.encode(
            texts,
            normalize_embeddings=True,
            show_progress_bar=True
        )

        # 创建HNSW索引
        self.index = faiss.IndexHNSWFlat(self.dimension, 32)
        self.index.hnsw.efConstruction = 200
        self.index.add(np.array(embeddings).astype('float32'))
        self.index.hnsw.efSearch = 128

    def query(
        self,
        query: str,
        use_rerank: bool = True,
        use_citation: bool = True
    ) -> Dict[str, Any]:
        """
        执行RAG查询
        """
        # 1. 查询重写（可选）
        expanded_queries = self._expand_query(query)

        # 2. 检索
        retrieved_chunks = self._retrieve(expanded_queries)

        # 3. 重排序
        if use_rerank:
            final_chunks = self._rerank(query, retrieved_chunks)
        else:
            final_chunks = retrieved_chunks[:self.top_k_rerank]

        # 4. 生成回答
        answer = self._generate(query, final_chunks, use_citation)

        return {
            'query': query,
            'answer': answer,
            'sources': final_chunks,
            'num_sources': len(final_chunks)
        }

    def _expand_query(self, query: str) -> List[str]:
        """
        简单的查询扩展
        """
        # 可以集成LLM进行更复杂的扩展
        return [query]

    def _retrieve(self, queries: List[str]) -> List[Dict]:
        """
        多查询检索并去重
        """
        all_indices = set()

        for q in queries:
            query_embedding = self.embedding_model.encode(
                [q],
                normalize_embeddings=True
            )
            distances, indices = self.index.search(
                np.array(query_embedding).astype('float32'),
                self.top_k_retrieve
            )
            all_indices.update(indices[0])

        retrieved = [self.chunks[i] for i in all_indices]
        return retrieved

    def _rerank(self, query: str, chunks: List[Dict]) -> List[Dict]:
        """
        交叉编码器重排序
        """
        pairs = [[query, chunk['content']] for chunk in chunks]
        scores = self.reranker.predict(pairs)

        # 添加分数并排序
        for chunk, score in zip(chunks, scores):  # zip按位置配对多个可迭代对象
            chunk['rerank_score'] = float(score)

        sorted_chunks = sorted(
            chunks,
            key=lambda x: x['rerank_score'],
            reverse=True
        )

        return sorted_chunks[:self.top_k_rerank]

    def _generate(
        self,
        query: str,
        chunks: List[Dict],
        use_citation: bool
    ) -> str:
        """
        使用LLM生成回答
        """
        # 构建上下文
        if use_citation:
            context = ""
            for i, chunk in enumerate(chunks, 1):
                context += f"[^{i}^] {chunk['content']}\n\n"

            prompt = f"""基于以下参考文档回答问题。请遵循以下规则：
1. 每个事实后标注来源，格式为[^编号^]
2. 如果信息不足，明确说明
3. 保持回答简洁准确

参考文档：
{context}

问题：{query}

回答："""
        else:
            context = "\n".join([c['content'] for c in chunks])
            prompt = f"""基于以下信息回答问题：

{context}

问题：{query}

回答："""

        # 调用LLM
        response = self.llm.generate(prompt)
        return response

# 使用示例
if __name__ == "__main__":
    # 初始化RAG系统
    rag = AdvancedRAGSystem(
        llm_client=YourLLMClient(),  # 替换为你的LLM客户端
        chunk_size=512,
        top_k_retrieve=20,
        top_k_rerank=5
    )

    # 添加文档
    documents = [
        "RAG（检索增强生成）是一种结合信息检索和文本生成的技术...",
        "Transformer架构由Vaswani等人在2017年提出...",
        # 更多文档...
    ]
    rag.add_documents(documents)

    # 查询
    result = rag.query("什么是RAG技术？")
    print(f"回答：{result['answer']}")
    print(f"引用来源：{result['num_sources']}个")

class RAGEvaluator:
    """
    RAG系统评估框架
    """

    def __init__(self):
        self.metrics = {}

    def evaluate_retrieval(
        self,
        queries: List[str],
        ground_truth_docs: List[List[int]],
        retrieved_docs: List[List[int]]
    ) -> Dict[str, float]:
        """
        评估检索质量
        """
        metrics = {
            'hit_rate@1': 0,
            'hit_rate@5': 0,
            'mrr': 0,  # Mean Reciprocal Rank
            'recall@5': 0
        }

        for query, gt, retrieved in zip(queries, ground_truth_docs, retrieved_docs):
            # Hit Rate@1
            if any(r in gt for r in retrieved[:1]):  # any()任一为True则返回True
                metrics['hit_rate@1'] += 1

            # Hit Rate@5
            if any(r in gt for r in retrieved[:5]):
                metrics['hit_rate@5'] += 1

            # MRR
            for rank, doc in enumerate(retrieved, 1):
                if doc in gt:
                    metrics['mrr'] += 1.0 / rank
                    break

            # Recall@5
            retrieved_set = set(retrieved[:5])
            gt_set = set(gt)
            metrics['recall@5'] += len(retrieved_set & gt_set) / len(gt_set)

        # 平均化
        n = len(queries)
        return {k: v / n for k, v in metrics.items()}

    def evaluate_generation(
        self,
        predictions: List[str],
        references: List[str]
    ) -> Dict[str, float]:
        """
        评估生成质量
        """
        from rouge import Rouge
        from bert_score import score

        # ROUGE分数
        rouge = Rouge()
        rouge_scores = rouge.get_scores(predictions, references, avg=True)

        # BERTScore
        P, R, F1 = score(predictions, references, lang='zh')

        return {
            'rouge-1': rouge_scores['rouge-1']['f'],
            'rouge-2': rouge_scores['rouge-2']['f'],
            'rouge-l': rouge_scores['rouge-l']['f'],
            'bertscore_f1': F1.mean().item()
        }

    def evaluate_faithfulness(
        self,
        answers: List[str],
        contexts: List[str]
    ) -> float:
        """
        评估回答对上下文的忠实度（幻觉检测）
        """
        # 使用NLI模型或LLM判断
        faithful_count = 0

        for answer, context in zip(answers, contexts):
            # 提取回答中的陈述
            claims = self._extract_claims(answer)

            # 检查每个陈述是否被上下文支持
            supported = 0
            for claim in claims:
                if self._verify_claim(claim, context):
                    supported += 1

            if supported == len(claims):
                faithful_count += 1

        return faithful_count / len(answers)

    def _extract_claims(self, text: str) -> List[str]:
        """从文本中提取事实陈述"""
        # 使用LLM或规则提取
        # 简化实现：按句子分割
        import re
        sentences = re.split(r'[。！？]', text)
        return [s.strip() for s in sentences if len(s.strip()) > 10]

    def _verify_claim(self, claim: str, context: str) -> bool:
        """验证陈述是否被上下文支持"""
        # 使用NLI模型或相似度判断
        # 简化实现：检查关键词重叠
        claim_words = set(claim.lower().split())
        context_words = set(context.lower().split())
        overlap = len(claim_words & context_words) / len(claim_words)
        return overlap > 0.5