RAG 的上下文窗口管理：检索到 20 个文档，但 LLM 只能吃下 5 个

def two_stage_retrieve(query: str, vector_store, embedder, reranker,
                       top_k_initial: int = 20, top_k_final: int = 5):
    # Stage 1: 向量检索（粗筛）
    query_vector = embedder.embed(query)
    candidates = vector_store.search(query_vector, top_k=top_k_initial)
    
    # Stage 2: 重排序（精排）
    pairs = [(query, doc.text) for doc in candidates]
    scores = reranker.score(pairs)
    
    # 按重排序得分取前 N 个
    ranked = sorted(zip(candidates, scores), key=lambda x: x[1], reverse=True)
    return [doc for doc, score in ranked[:top_k_final]]

def deduplicate_chunks(chunks: list[Chunk], similarity_threshold: float = 0.85) -> list[Chunk]:
    """去除语义高度重复的 chunk，保留最相关的一个"""
    selected = []
    for chunk in sorted(chunks, key=lambda c: c.relevance_score, reverse=True):
        is_duplicate = False
        for selected_chunk in selected:
            sim = cosine_similarity(chunk.embedding, selected_chunk.embedding)
            if sim > similarity_threshold:
                is_duplicate = True
                break
        if not is_duplicate:
            selected.append(chunk)
    return selected

def filter_low_density_chunks(chunks: list[Chunk], min_info_density: float = 0.3) -> list[Chunk]:
    """过滤信息密度低的 chunk"""
    def info_density(text: str) -> float:
        # 简单指标：去重后的有效信息占比
        words = text.split()
        unique_words = set(words)
        if len(words) == 0:
            return 0
        # 信息密度 = 唯一词比例 × 内容长度（归一化）
        uniqueness = len(unique_words) / len(words)
        # 太短的 chunk 大概率信息密度低
        length_factor = min(len(words) / 50, 1.0)
        return uniqueness * length_factor
    
    return [c for c in chunks if info_density(c.text) >= min_info_density]

def select_by_metadata(chunks: list[Chunk], max_chunks: int = 5) -> list[Chunk]:
    """
    基于元数据选择 chunk，保证：
    1. 优先选择高相关性（重排序得分）
    2. 同一文档来源不超过 2 个
    3. 优先选择更高层级的标题（如 h1 > h2 > h3）
    """
    # 按文档来源分组
    from collections import defaultdict
    source_groups = defaultdict(list)
    for chunk in chunks:
        source_groups[chunk.metadata.source].append(chunk)
    
    # 每个来源最多选 2 个
    selected = []
    for source, group in source_groups.items():
        group.sort(key=lambda c: c.relevance_score, reverse=True)
        selected.extend(group[:2])
    
    # 按相关性排序，取前 max_chunks
    selected.sort(key=lambda c: c.relevance_score, reverse=True)
    return selected[:max_chunks]

Round 1: 给 LLM 最相关的 Top-3 chunk + 问题 → 生成答案
如果 LLM 认为信息不足，标记缺失内容
Round 2: 根据缺失内容，从剩余 chunk 中补充 Top-2 → 再次生成
最多 N 轮

def progressive_reading(query: str, ranked_chunks: list[Chunk], 
                        llm, max_rounds: int = 3, chunk_per_round: int = 3):
    """渐进式阅读：分批给 LLM 提供信息"""
    context = []
    pointer = 0
    full_answer = ""
    
    for round in range(max_rounds):
        # 取下一批 chunk
        batch = ranked_chunks[pointer:pointer + chunk_per_round]
        if not batch:
            break
        pointer += chunk_per_round
        context.extend(batch)
        
        # 构建 prompt
        prompt = f"""你正在基于以下文档片段回答问题。

已有文档片段：
{format_chunks(context)}

问题：{query}

请回答以上问题。如果你认为当前文档片段中的信息不足以给出完整准确的回答，
请明确说明"信息不足"以及缺少哪些信息。
否则，请直接给出完整回答。"""
        
        answer = llm.generate(prompt)
        
        # 判断 LLM 是否认为信息足够
        if "信息不足" not in answer:
            return answer
        
        # 记录部分答案和缺失信息，继续下一轮
        full_answer = answer
    
    # 如果达到最大轮次仍有信息不足，返回当前最佳答案
    return full_answer

def extractive_compress(chunks: list[Chunk], query: str, llm, 
                        max_tokens: int = 2000) -> str:
    """从多个 chunk 中提取与 query 最相关的句子，控制在 max_tokens 内"""
    combined_text = "\n".join(c.text for c in chunks)
    
    prompt = f"""请从以下文档中提取与问题最相关的句子，去除无关内容。
要求：
1. 只保留与问题直接相关的句子
2. 保持原文的表述方式，不要改写
3. 提取的内容长度不超过 {max_tokens} token

问题：{query}

文档：
{combined_text}

提取的内容："""
    
    compressed = llm.generate(prompt)
    return compressed

def summarize_chunks(chunks: list[Chunk], query: str, llm) -> str:
    """对每个 chunk 做面向 query 的摘要"""
    summaries = []
    for chunk in chunks:
        prompt = f"""请为以下文档片段生成一个面向问题"{query}"的摘要。
要求：保留关键事实和数据，去除无关细节，控制在 100 字以内。

文档片段：{chunk.text}

面向问题的摘要："""
        summary = llm.generate(prompt)
        summaries.append(summary)
    
    return "\n".join(summaries)

class ContextManager:
    """RAG 上下文管理器"""
    
    def build_context(self, query: str, raw_chunks: list[Chunk], 
                      max_context_tokens: int = 4000) -> str:
        # Step 1: 重排序（策略一）
        ranked = self.reranker.rerank(query, raw_chunks)
        
        # Step 2: 去重与过滤（策略二）
        deduped = self.deduplicate(ranked)
        filtered = self.filter_low_density(deduped)
        
        # Step 3: 渐进式选择（策略三）
        if self.use_progressive:
            return self.progressive_read(query, filtered, max_context_tokens)
        
        # Step 4: 截断 + 压缩（策略四）
        truncated = filtered[:self.top_k_final]
        compressed = self.compress(truncated, query, max_context_tokens)
        
        return compressed