RAG 流水线优化全攻略：从检索到生成的极致调优

用户查询 → 查询改写 → 多路召回 → 重排序 → 上下文压缩 → LLM 生成 → 流式输出
                ↑                              ↓
              缓存                          引用标注

from langchain.text_splitter import RecursiveCharacterTextSplitter

splitter = RecursiveCharacterTextSplitter(
    chunk_size=512,
    chunk_overlap=64,
    separators=["\n\n", "\n", "。", "！", "？", "，", " ", ""]
)

chunks = splitter.split_text(long_document)

from langchain.text_splitter import SemanticChunker
from langchain.embeddings import HuggingFaceEmbeddings

embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-large-zh-v1.5")
splitter = SemanticChunker(
    embeddings,
    breakpoint_threshold_type="percentile",  # 或 "standard_deviation"
    breakpoint_threshold_amount=95
)

chunks = splitter.split_text(long_document)

from langchain.text_splitter import MarkdownHeaderTextSplitter

headers_to_split_on = [
    ("#", "Header 1"),
    ("##", "Header 2"),
    ("###", "Header 3"),
]

splitter = MarkdownHeaderTextSplitter(headers_to_split_on)
chunks = splitter.split_text(markdown_doc)
# 每个 chunk 保留了所属标题的 metadata

import asyncio

async def multi_retrieve(query: str, top_k: int = 5):
    # 路路1：稠密向量检索
    dense_results = await vector_search(query, top_k=top_k*2)
    
    # 路路2：稀疏检索 (BM25)
    bm25_results = await bm25_search(query, top_k=top_k*2)
    
    # 路路3：混合检索 (bge-m3 sparse)
    sparse_results = await sparse_search(query, top_k=top_k*2)
    
    # 融合去重
    all_docs = {}
    for results in [dense_results, bm25_results, sparse_results]:
        for doc in results:
            doc_id = doc["id"]
            if doc_id not in all_docs:
                all_docs[doc_id] = doc
    
    return list(all_docs.values())[:top_k*3]

def rrf_fusion(result_lists, k=60):
    """Reciprocal Rank Fusion"""
    scores = {}
    for results in result_lists:
        for rank, doc in enumerate(results):
            doc_id = doc["id"]
            scores[doc_id] = scores.get(doc_id, 0) + 1 / (k + rank)
    
    sorted_ids = sorted(scores, key=scores.get, reverse=True)
    return sorted_ids

from FlagEmbedding import FlagReranker

reranker = FlagReranker('BAAI/bge-reranker-v2-m3', use_fp16=True)

# 对召回的 top-20 文档重排
pairs = [[query, doc["text"]] for doc in candidates]
scores = reranker.compute_score(pairs)

# 按分数排序取 top-5
ranked = sorted(zip(candidates, scores), key=lambda x: x[1], reverse=True)
final_docs = [doc for doc, _ in ranked[:5]]

# 用 LLM 生成多个改写版本
rewrite_prompt = """请将以下问题改写为 3 个不同角度的检索查询，用于提升检索效果。

原始问题：{question}

输出格式（JSON 数组）：
["改写1", "改写2", "改写3"]
"""

response = await llm.chat.completions.create(
    model="gpt-4o-mini",
    messages=[{"role": "user", "content": rewrite_prompt.format(question=query)}]
)
rewrites = json.loads(response.choices[0].message.content)

# 对每个改写查询分别检索，再融合
all_results = []
for q in [query] + rewrites:
    results = await vector_search(q, top_k=5)
    all_results.append(results)

fused = rrf_fusion(all_results)

hyde_prompt = """请为以下问题生成一个简短的假设性答案（2-3句话），不需要准确，但要用专业术语。

问题：{question}

假设性答案："""

fake_doc = await llm.generate(hyde_prompt.format(question=query))
# 用 fake_doc 而非原始 query 做向量检索
embedding = embed_model.encode(fake_doc)
results = vector_db.search(embedding, top_k=5)

compress_prompt = """以下是从知识库检索到的文档片段。请提取与用户问题最相关的内容，去除无关信息。

用户问题：{question}

文档片段：
{documents}

请输出压缩后的相关内容（保留关键信息，去除冗余）："""

compressed = await llm.generate(compress_prompt.format(
    question=query,
    documents="\n\n".join(doc["text"] for doc in final_docs)
))

from fastapi import FastAPI
from fastapi.responses import StreamingResponse

app = FastAPI()

@app.post("/chat")
async def chat(req: ChatRequest):
    # 1. 检索
    docs = await multi_retrieve(req.query)
    context = "\n\n".join(d["text"] for d in docs)
    
    # 2. 构造 prompt
    messages = [
        {"role": "system", "content": f"基于以下资料回答问题：\n{context}"},
        {"role": "user", "content": req.query}
    ]
    
    # 3. 流式生成
    async def stream():
        async for chunk in await llm.chat.completions.create(
            model="deepseek-chat",
            messages=messages,
            stream=True
        ):
            if chunk.choices[0].delta.content:
                yield f"data: {chunk.choices[0].delta.content}\n\n"
        yield "data: [DONE]\n\n"
    
    return StreamingResponse(stream(), media_type="text/event-stream")

import redis
import numpy as np

r = redis.Redis()
SIM_THRESHOLD = 0.95

async def cached_answer(query: str, query_embedding: np.ndarray):
    # 遍历缓存中所有 query 向量
    cached_keys = r.keys("query:*")
    for key in cached_keys:
        cached_emb = np.frombuffer(r.get(key), dtype=np.float32)
        sim = np.dot(query_embedding, cached_emb) / (
            np.linalg.norm(query_embedding) * np.linalg.norm(cached_emb)
        )
        if sim > SIM_THRESHOLD:
            # 命中缓存
            cached_answer_key = key.replace("query:", "answer:")
            return r.get(cached_answer_key).decode()
    
    # 未命中，走 RAG 流程
    answer = await rag_pipeline(query)
    
    # 写入缓存
    query_id = f"query:{hash(query)}"
    r.set(query_id, query_embedding.tobytes())
    r.set(f"answer:{hash(query)}", answer)
    
    return answer

import hashlib

def get_or_embed(text: str, embed_model, redis_client):
    doc_hash = hashlib.md5(text.encode()).hexdigest()
    cached = redis_client.get(f"emb:{doc_hash}")
    if cached:
        return np.frombuffer(cached, dtype=np.float32)
    
    embedding = embed_model.encode(text)
    redis_client.set(f"emb:{doc_hash}", embedding.tobytes())
    return embedding