RAG 生产环境 12 大坑及解决方案

from dataclasses import dataclass
from typing import List
import re

@dataclass
class Chunk:
    text: str
    metadata: dict
    token_count: int = 0

class SemanticChunker:
    """基于语义边界的智能分块器"""
    
    def __init__(
        self,
        target_size: int = 400,      # 目标块大小（tokens）
        min_size: int = 100,          # 最小块大小
        max_size: int = 600,          # 最大块大小
        overlap: int = 50,            # 重叠区间
    ):
        self.target_size = target_size
        self.min_size = min_size
        self.max_size = max_size
        self.overlap = overlap
    
    def chunk_document(self, text: str, source: str = "") -> List[Chunk]:
        """分块主流程"""
        # Step 1: 按结构边界切分
        sections = self._split_by_structure(text)
        
        # Step 2: 对每个 section 按 paragraph 切分
        paragraphs = []
        for section in sections:
            paragraphs.extend(self._split_by_paragraph(section))
        
        # Step 3: 合并过小的段落，拆分过大的段落
        chunks = self._merge_and_split(paragraphs)
        
        # Step 4: 添加重叠
        chunks = self._add_overlap(chunks)
        
        # Step 5: 附加元数据
        return [
            Chunk(
                text=c,
                metadata={"source": source, "chunk_index": i, "total_chunks": len(chunks)},
                token_count=len(c) // 2,  # 粗略估算
            )
            for i, c in enumerate(chunks)
        ]
    
    def _split_by_structure(self, text: str) -> List[str]:
        """按标题、分隔符等结构边界切分"""
        # 按 Markdown 标题切分
        pattern = r'(?=^#{1,6}\s)'
        sections = re.split(pattern, text, flags=re.MULTILINE)
        return [s.strip() for s in sections if s.strip()]
    
    def _split_by_paragraph(self, text: str) -> List[str]:
        """按段落（双换行）切分"""
        paras = text.split("\n\n")
        return [p.strip() for p in paras if p.strip()]
    
    def _merge_and_split(self, paragraphs: List[str]) -> List[str]:
        """合并过小段落，拆分过大段落"""
        chunks = []
        buffer = []
        buffer_size = 0
        
        for para in paragraphs:
            para_size = len(para) // 2  # 粗略 token 估算
            
            if buffer_size + para_size > self.max_size and buffer:
                chunks.append("\n\n".join(buffer))
                buffer = []
                buffer_size = 0
            
            if para_size > self.max_size:
                # 单段落过长，按句子切分
                if buffer:
                    chunks.append("\n\n".join(buffer))
                    buffer = []
                    buffer_size = 0
                sentences = re.split(r'(?<=[。！？.!?])\s+', para)
                sent_buffer = []
                sent_size = 0
                for sent in sentences:
                    if sent_size + len(sent) // 2 > self.max_size and sent_buffer:
                        chunks.append(" ".join(sent_buffer))
                        sent_buffer = []
                        sent_size = 0
                    sent_buffer.append(sent)
                    sent_size += len(sent) // 2
                if sent_buffer:
                    chunks.append(" ".join(sent_buffer))
            else:
                buffer.append(para)
                buffer_size += para_size
                
                if buffer_size >= self.target_size:
                    chunks.append("\n\n".join(buffer))
                    buffer = []
                    buffer_size = 0
        
        if buffer:
            chunks.append("\n\n".join(buffer))
        
        return chunks
    
    def _add_overlap(self, chunks: List[str]) -> List[str]:
        """为相邻块添加重叠"""
        if self.overlap <= 0 or len(chunks) <= 1:
            return chunks
        
        result = [chunks[0]]
        for i in range(1, len(chunks)):
            prev_text = chunks[i - 1]
            overlap_text = prev_text[-self.overlap * 2:]  # 粗略取后半段
            result.append(overlap_text + " " + chunks[i])
        
        return result


# 使用示例
chunker = SemanticChunker(target_size=400, overlap=50)
document = open("knowledge_base/product_manual.md").read()
chunks = chunker.chunk_document(document, source="product_manual.md")
print(f"分块完成: {len(chunks)} 个块")

# 统一使用同一生态的模型，或充分评测跨模型组合
EMBEDDING_MODEL_CONFIGS = {
    "openai_3_large": {
        "model": "text-embedding-3-large",
        "dim": 3072,
        "best_with": ["gpt-4o", "gpt-4o-mini"],
        "cost_per_1k": 0.00013,
    },
    "openai_3_small": {
        "model": "text-embedding-3-small",
        "dim": 1536,
        "best_with": ["gpt-4o-mini"],
        "cost_per_1k": 0.00002,
    },
    "cohere_multilingual": {
        "model": "embed-multilingual-v3",
        "dim": 1024,
        "best_with": ["command-r-plus"],
        "cost_per_1k": 0.0001,
    },
}

from rank_bm25 import BM25Okapi
import numpy as np

class HybridRetriever:
    """混合检索：向量 + BM25"""
    
    def __init__(self, documents: list[Chunk], embeddings: np.ndarray, weight: float = 0.5):
        self.documents = documents
        self.embeddings = embeddings
        self.vector_weight = weight
        self.keyword_weight = 1 - weight
        
        # 构建 BM25 索引
        tokenized_docs = [doc.text.lower().split() for doc in documents]
        self.bm25 = BM25Okapi(tokenized_docs)
    
    def search(self, query: str, query_embedding: np.ndarray, top_k: int = 5) -> list[dict]:
        # 向量检索
        vector_scores = self._cosine_scores(query_embedding)
        
        # BM25 检索
        tokenized_query = query.lower().split()
        bm25_scores = self.bm25.get_scores(tokenized_query)
        
        # 归一化
        vector_scores = self._normalize(vector_scores)
        bm25_scores = self._normalize(bm25_scores)
        
        # 融合
        combined = (
            self.vector_weight * vector_scores +
            self.keyword_weight * bm25_scores
        )
        
        # 取 Top-K
        top_indices = np.argsort(combined)[-top_k:][::-1]
        
        return [
            {
                "chunk": self.documents[i],
                "score": combined[i],
                "vector_score": vector_scores[i],
                "bm25_score": bm25_scores[i],
            }
            for i in top_indices
        ]
    
    def _cosine_scores(self, query_emb: np.ndarray) -> np.ndarray:
        return self.embeddings @ query_emb / (
            np.linalg.norm(self.embeddings, axis=1) * np.linalg.norm(query_emb) + 1e-8
        )
    
    def _normalize(self, scores: np.ndarray) -> np.ndarray:
        min_val, max_val = scores.min(), scores.max()
        if max_val - min_val < 1e-8:
            return np.zeros_like(scores)
        return (scores - min_val) / (max_val - min_val)


# 检索效果对比

class QueryRewriter:
    """查询重写器"""
    
    def __init__(self, llm_client):
        self.client = llm_client
    
    def rewrite(self, query: str, conversation_history: list[dict] = None) -> list[str]:
        """生成多个重写查询以提高召回"""
        
        prompt = """将用户查询重写为 3 个不同的检索查询。
要求：
1. 补充隐含的关键词
2. 使用不同的表述方式
3. 包含同义词和相关术语

对话历史:
{history}

原始查询: {query}

输出格式（JSON）:
{{"queries": ["查询1", "查询2", "查询3"]}}
""".format(
            history=conversation_history or "无",
            query=query,
        )
        
        response = self.client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt}],
            response_format={"type": "json_object"},
        )
        
        import json
        result = json.loads(response.choices[0].message.content)
        return [query] + result["queries"]  # 原始查询 + 重写查询

class ContextWindowManager:
    """RAG 上下文窗口管理器"""
    
    def __init__(self, max_context_tokens: int = 6000):
        self.max_context = max_context_tokens
    
    def build_context(
        self,
        system_prompt: str,
        conversation: list[dict],
        retrieved_chunks: list[dict],
    ) -> list[dict]:
        """在 token 预算内构建最优上下文"""
        
        budget = self.max_context
        system_tokens = self._estimate_tokens(system_prompt)
        budget -= system_tokens
        
        # 预留生成空间
        generation_reserve = 1024
        budget -= generation_reserve
        
        # 对话历史（保留最近 N 轮）
        conversation_tokens = sum(self._estimate_tokens(m["content"]) for m in conversation[-6:])
        budget -= conversation_tokens
        
        if budget <= 0:
            # 严重不足，只保留系统提示和最后一条消息
            return [
                {"role": "system", "content": system_prompt},
                *conversation[-2:],
            ]
        
        # 填充检索结果
        selected_chunks = []
        for chunk_data in retrieved_chunks:
            chunk_tokens = self._estimate_tokens(chunk_data["chunk"].text)
            if chunk_tokens <= budget:
                selected_chunks.append(chunk_data)
                budget -= chunk_tokens
            else:
                # 尝试截断
                truncated = chunk_data["chunk"].text[:budget * 2]
                selected_chunks.append({**chunk_data, "chunk": Chunk(
                    text=truncated + "...[截断]",
                    metadata=chunk_data["chunk"].metadata,
                )})
                break
        
        # 组装最终上下文
        context_text = "\n\n---\n\n".join([
            f"[文档 {i+1}] (相关度: {c['score']:.2f})\n{c['chunk'].text}"
            for i, c in enumerate(selected_chunks)
        ])
        
        return [
            {"role": "system", "content": system_prompt},
            {"role": "system", "content": f"参考文档:\n{context_text}"},
            *conversation[-6:],
        ]
    
    def _estimate_tokens(self, text: str) -> int:
        return len(text) // 2  # 中英混合粗略估算

class Reranker:
    """使用 Cross-Encoder 重排序"""
    
    def __init__(self, model_name: str = "bge-reranker-v2-m3"):
        from sentence_transformers import CrossEncoder
        self.model = CrossEncoder(model_name)
    
    def rerank(self, query: str, documents: list[str], top_k: int = 3) -> list[dict]:
        """对检索结果重排序"""
        pairs = [(query, doc) for doc in documents]
        scores = self.model.predict(pairs)
        
        ranked = sorted(
            [{"document": doc, "score": float(score), "rank": i}
             for i, (doc, score) in enumerate(zip(documents, scores))],
            key=lambda x: x["score"],
            reverse=True,
        )
        
        # 只返回 Top-K
        for i, item in enumerate(ranked[:top_k]):
            item["rank"] = i + 1
        
        return ranked[:top_k]

class TableAwareChunker:
    """表格感知的分块器"""
    
    def chunk(self, text: str) -> list[str]:
        chunks = []
        lines = text.split("\n")
        i = 0
        
        while i < len(lines):
            line = lines[i]
            
            # 检测表格开始
            if "|" in line and i + 1 < len(lines) and "---" in lines[i + 1]:
                # 收集完整表格
                table_lines = [line]
                i += 1
                while i < len(lines) and "|" in lines[i]:
                    table_lines.append(lines[i])
                    i += 1
                # 表格作为整体块
                chunks.append("\n".join(table_lines))
            else:
                # 普通文本行
                chunks.append(line)
                i += 1
        
        # 合并相邻文本块
        return self._merge_text_chunks(chunks)
    
    def _merge_text_chunks(self, chunks: list[str]) -> list[str]:
        merged = []
        text_buffer = []
        
        for chunk in chunks:
            if "|" in chunk and "---" in chunk:
                if text_buffer:
                    merged.append("\n".join(text_buffer))
                    text_buffer = []
                merged.append(chunk)
            else:
                text_buffer.append(chunk)
        
        if text_buffer:
            merged.append("\n".join(text_buffer))
        
        return merged

class MultiLanguageHandler:
    """多语言处理器"""
    
    def __init__(self):
        self.language_routes = {
            "zh": {"embedding_model": "text-embedding-3-large", "index": "idx_zh"},
            "en": {"embedding_model": "text-embedding-3-large", "index": "idx_en"},
            "ja": {"embedding_model": "text-embedding-3-large", "index": "idx_ja"},
        }
    
    def detect_language(self, text: str) -> str:
        """简单语言检测"""
        chinese_chars = sum(1 for c in text if '\u4e00' <= c <= '\u9fff')
        if chinese_chars / max(len(text), 1) > 0.3:
            return "zh"
        return "en"
    
    def translate_query(self, query: str, target_lang: str) -> str:
        """查询翻译（使用 LLM）"""
        # 对于多语言知识库，统一翻译为英文检索
        # 或分别用原语言和翻译语言检索，合并结果
        pass
    
    def search_multilingual(self, query: str, top_k: int = 5) -> list[dict]:
        """多语言搜索"""
        lang = self.detect_language(query)
        
        # 1. 用原始语言检索
        results_original = self._search(query, self.language_routes[lang]["index"])
        
        # 2. 翻译后检索
        translated = self.translate_query(query, "en" if lang != "en" else "zh")
        results_translated = self._search(translated, self.language_routes["en" if lang != "en" else "zh"]["index"])
        
        # 3. 合并去重
        merged = self._merge_results(results_original, results_translated)
        return merged[:top_k]

class CitationTracker:
    """答案引用追踪"""
    
    def build_prompt_with_citations(
        self, query: str, retrieved_chunks: list[dict]
    ) -> str:
        """构建带引用标记的 Prompt"""
        
        doc_section = "参考文档：\n"
        for i, chunk in enumerate(retrieved_chunks):
            doc_section += f"\n[文档{i+1}] (来源: {chunk['source']}, 相关度: {chunk['score']:.2f})\n"
            doc_section += chunk["text"] + "\n"
        
        prompt = f"""{doc_section}

基于上述参考文档回答用户问题。要求：
1. 只使用参考文档中的信息
2. 在每个关键陈述后标注来源，格式：[文档X]
3. 如果参考文档中没有答案，明确说"提供的文档中没有相关信息"

用户问题: {query}
"""
        return prompt
    
    def extract_citations(self, answer: str) -> list[dict]:
        """从答案中提取引用"""
        import re
        citations = []
        pattern = r'\[文档(\d+)\]'
        
        for match in re.finditer(pattern, answer):
            citations.append({
                "doc_id": int(match.group(1)),
                "position": match.start(),
            })
        
        return citations

class IncrementalIndexer:
    """增量索引管理器"""
    
    def __init__(self, vector_store, document_store):
        self.vector_store = vector_store
        self.doc_store = document_store
    
    def update_document(self, doc_id: str, new_content: str):
        """更新单个文档"""
        # 1. 删除旧向量
        self.vector_store.delete(filter={"doc_id": doc_id})
        
        # 2. 更新文档存储
        self.doc_store.update(doc_id, new_content)
        
        # 3. 重新分块
        chunks = self.chunker.chunk_document(new_content, source=doc_id)
        
        # 4. 生成新向量
        embeddings = self.embedding_model.encode([c.text for c in chunks])
        
        # 5. 写入向量存储
        self.vector_store.upsert(
            ids=[f"{doc_id}_chunk_{i}" for i in range(len(chunks))],
            embeddings=embeddings,
            metadatas=[{**c.metadata, "doc_id": doc_id, "version": "latest"} for c in chunks],
        )
    
    def delete_document(self, doc_id: str):
        """安全删除文档及其所有向量"""
        # 先标记，再删除
        self.doc_store.mark_deleted(doc_id)
        self.vector_store.delete(filter={"doc_id": doc_id})
        self.doc_store.actual_delete(doc_id)

@dataclass
class RAGEvaluation:
    """RAG 系统评估"""
    
    def evaluate_retrieval(self, test_set: list[dict]) -> dict:
        """评估检索质量"""
        metrics = {"hit_rate": 0, "mrr": 0, "ndcg": 0}
        
        for case in test_set:
            retrieved = self.retrieve(case["query"], top_k=5)
            relevant_ids = set(case["relevant_doc_ids"])
            retrieved_ids = [r["doc_id"] for r in retrieved]
            
            # Hit Rate: Top-K 中是否包含相关文档
            hit = int(bool(set(retrieved_ids) & relevant_ids))
            
            # MRR: 第一个相关文档的倒数排名
            mrr = 0
            for rank, rid in enumerate(retrieved_ids, 1):
                if rid in relevant_ids:
                    mrr = 1 / rank
                    break
            
            metrics["hit_rate"] += hit
            metrics["mrr"] += mrr
        
        n = len(test_set)
        return {k: v / n for k, v in metrics.items()}
    
    def evaluate_generation(self, test_set: list[dict]) -> dict:
        """评估生成质量"""
        metrics = {"faithfulness": 0, "answer_relevance": 0, "context_precision": 0}
        
        # 使用 RAGAS 或类似框架
        # faithfulness: 答案是否忠实于检索到的上下文
        # answer_relevance: 答案是否回应了用户问题
        # context_precision: 检索到的上下文是否相关
        
        return metrics

class IterativeRetriever:
    """迭代检索：多步推理"""
    
    def __init__(self, retriever, llm_client, max_iterations: int = 3):
        self.retriever = retriever
        self.llm = llm_client
        self.max_iterations = max_iterations
    
    def search(self, query: str) -> dict:
        """多步迭代检索"""
        all_contexts = []
        reasoning_chain = []
        
        current_query = query
        
        for i in range(self.max_iterations):
            # 检索
            results = self.retriever.search(current_query, top_k=3)
            all_contexts.extend(results)
            
            # 判断是否需要继续检索
            decomposition_prompt = f"""基于已有信息，判断是否需要进一步检索。

原始问题: {query}
已检索信息: {[r['chunk'].text[:200] for r in all_contexts]}

如果需要更多信息，生成下一个检索查询。如果信息足够，返回 "SUFFICIENT"。
"""
            response = self.llm.chat.completions.create(
                model="gpt-4o-mini",
                messages=[{"role": "user", "content": decomposition_prompt}],
                max_tokens=200,
            )
            
            next_action = response.choices[0].message.content.strip()
            
            if "SUFFICIENT" in next_action:
                reasoning_chain.append({"step": i, "action": "sufficient", "query": current_query})
                break
            
            current_query = next_action
            reasoning_chain.append({"step": i, "action": "search", "query": current_query})
        
        # 去重
        unique_contexts = self._deduplicate(all_contexts)
        
        return {
            "contexts": unique_contexts,
            "reasoning_chain": reasoning_chain,
            "iterations": len(reasoning_chain),
        }
    
    def _deduplicate(self, contexts: list[dict]) -> list[dict]:
        seen = set()
        unique = []
        for ctx in contexts:
            text_hash = hash(ctx["chunk"].text[:100])
            if text_hash not in seen:
                seen.add(text_hash)
                unique.append(ctx)
        return unique