RAG 检索的工程取舍：切分、重排与上下文预算

文档 → 切分 → 向量化 → 索引 → 检索 → 重排 → 拼装上下文 → 生成
                      ↑                        ↑
                  离线批处理              在线低延迟

import re
from typing import Iterable
 
HEADING = re.compile(r"^(#{1,6})\s+(.*)$")
 
 
def split_markdown(text: str, max_chars: int = 800) -> list[str]:
    """先按标题分节，节内按空行合并，超长节再按段落切开。"""
    sections: list[str] = []
    current: list[str] = []
    buffer: list[str] = []
 
    def flush_section() -> None:
        if not buffer:
            return
        chunk = "\n".join(buffer).strip()
        if chunk:
            sections.append(chunk)
        buffer.clear()
 
    for line in text.split("\n"):
        if HEADING.match(line):
            flush_section()
        buffer.append(line)
        if sum(len(x) for x in buffer) > max_chars:
            flush_section()
    flush_section()
 
    # 单节仍然超长时，按段落边界继续切
    result: list[str] = []
    for section in sections:
        if len(section) <= max_chars:
            result.append(section)
            continue
        para = ""
        for block in section.split("\n\n"):
            if len(para) + len(block) > max_chars and para:
                result.append(para.strip())
                para = block
            else:
                para = f"{para}\n\n{block}" if para else block
        if para.strip():
            result.append(para.strip())
    return result

原始 chunk：
  "这个参数控制最大递归深度，超过就抛异常。"
 
加上标题路径后：
  "文档标题：配置项说明 > 递归参数 > 最大深度
   这个参数控制最大递归深度，超过就抛异常。"

def retrieve(query: str, k: int = 20, top_n: int = 5):
    # 第一阶段：向量检索召回更多候选
    candidates = vector_store.similarity_search(query, k=k)
 
    # 第二阶段：交叉编码器精排，只保留少量
    pairs = [(query, doc.page_content) for doc in candidates]
    scores = reranker.predict(pairs)
    ranked = sorted(zip(candidates, scores), key=lambda x: -x[1])
    return [doc for doc, _ in ranked[:top_n]]

def build_context(reranked, max_tokens: int = 4000, per_doc_cap: int = 800) -> str:
    used = 0
    blocks = []
    for i, doc in enumerate(reranked):
        text = doc.page_content[:per_doc_cap]
        tokens = count_tokens(text)
        if used + tokens > max_tokens:
            break
        used += tokens
        blocks.append(f"[来源 {i+1}｜{doc.metadata.get('title', '未知')}]\n{text}")
    return "\n\n".join(blocks)

[来源 1｜文档 A]
...
 
[来源 2｜文档 B]
...
 
请基于以上资料回答，并标注引用了哪些来源编号。

@dataclass
class EvalCase:
    question: str
    answer: str
    expected_sources: list[str]
 
 
def evaluate(cases, pipeline) -> dict:
    """检索指标 + 生成指标分开看。"""
    hit_rate = []
    for case in cases:
        retrieved = pipeline.retrieve(case.question)
        hit = any(
            src in [r.metadata["id"] for r in retrieved]
            for src in case.expected_sources
        )
        hit_rate.append(hit)
    return {"recall_at_k": sum(hit_rate) / len(hit_rate)}