第 14 章 RAG 系统的基础设施

from sentence_transformers import SentenceTransformer
 
model = SentenceTransformer("BAAI/bge-large-zh-v1.5")
embeddings = model.encode(["你好世界", "Hello world"], normalize_embeddings=True)
# embeddings.shape: (2, 1024)

docker run -d --gpus all \
  -v /data/models:/data \
  -p 8080:80 \
  ghcr.io/huggingface/text-embeddings-inference:1.5 \
  --model-id BAAI/bge-large-zh-v1.5 \
  --max-batch-tokens 16384 \
  --max-concurrent-requests 128

def batch_embed_sorted(texts: list[str], model, batch_size: int = 64):
    """按长度排序后分批 embedding，减少 padding"""
    indexed = sorted(enumerate(texts), key=lambda x: len(x[1]))
    embeddings = [None] * len(texts)
 
    for i in range(0, len(indexed), batch_size):
        batch = indexed[i:i + batch_size]
        batch_texts = [t for _, t in batch]
        batch_embs = model.encode(batch_texts, normalize_embeddings=True)
        for (orig_idx, _), emb in zip(batch, batch_embs):
            embeddings[orig_idx] = emb
 
    return embeddings

数据量 < 100 万条 + 已有 PostgreSQL → pgvector
  - 零额外运维成本
  - PG 15+ 性能已经够用
  - 用 pgvector 0.7+ 支持 HNSW 索引

数据量 < 1000 万条 + 需要 Hybrid Search → Qdrant
  - 单机部署简单，Docker 一行搞定
  - Rust 实现，性能好，内存效率高
  - 原生支持 named vectors（同时存 dense 和 sparse）

数据量 > 1000 万条 + 需要分布式 → Milvus
  - 专为大规模设计
  - 支持 DiskANN 索引，十亿级数据不全放内存
  - 运维复杂度高，需要 etcd + MinIO + 多组件

快速原型 / 本地开发 → Chroma
  - pip install chromadb 即用
  - 不适合生产

docker run -d \
  -p 6333:6333 \
  -p 6334:6334 \
  -v /data/qdrant:/qdrant/storage \
  qdrant/qdrant:v1.12.5

from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct
 
client = QdrantClient(host="localhost", port=6333)
 
# 创建 collection
client.create_collection(
    collection_name="documents",
    vectors_config=VectorParams(size=1024, distance=Distance.COSINE),
)
 
# 插入向量
client.upsert(
    collection_name="documents",
    points=[
        PointStruct(
            id=1,
            vector=embedding_vector,  # list[float], 长度 1024
            payload={"text": "原始文本", "source": "doc1.pdf", "page": 3},
        ),
    ],
)
 
# 检索
results = client.query_points(
    collection_name="documents",
    query=query_vector,
    limit=10,
)

def fixed_size_chunk(text: str, chunk_size: int = 500, overlap: int = 100) -> list[str]:
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        start = end - overlap
    return chunks

SEPARATORS = ["\n\n", "\n", "。", "！", "？", "；", "，", " ", ""]
 
def recursive_split(text: str, chunk_size: int = 500, separators=None) -> list[str]:
    if separators is None:
        separators = SEPARATORS
 
    if len(text) <= chunk_size:
        return [text] if text.strip() else []
 
    sep = separators[0]
    remaining_seps = separators[1:]
 
    parts = text.split(sep)
    chunks = []
    current = ""
 
    for part in parts:
        candidate = current + sep + part if current else part
        if len(candidate) <= chunk_size:
            current = candidate
        else:
            if current:
                chunks.append(current)
            if len(part) > chunk_size and remaining_seps:
                chunks.extend(recursive_split(part, chunk_size, remaining_seps))
            else:
                current = part
 
    if current:
        chunks.append(current)
    return chunks

import re
 
def markdown_chunk(text: str, max_chunk_size: int = 1000) -> list[dict]:
    """按 Markdown 标题分段"""
    sections = re.split(r'(^#{1,3}\s+.+$)', text, flags=re.MULTILINE)
    chunks = []
    current_header = ""
    current_content = ""
 
    for section in sections:
        if re.match(r'^#{1,3}\s+', section):
            if current_content.strip():
                chunks.append({
                    "header": current_header,
                    "content": current_content.strip(),
                })
            current_header = section.strip()
            current_content = ""
        else:
            current_content += section
 
    if current_content.strip():
        chunks.append({"header": current_header, "content": current_content.strip()})
 
    # 对超长段落二次切分
    result = []
    for chunk in chunks:
        if len(chunk["content"]) > max_chunk_size:
            sub_chunks = recursive_split(chunk["content"], max_chunk_size)
            for sc in sub_chunks:
                result.append({"header": chunk["header"], "content": sc})
        else:
            result.append(chunk)
    return result

def reciprocal_rank_fusion(
    results_list: list[list[dict]],
    k: int = 60,
    top_n: int = 10,
) -> list[dict]:
    """
    RRF 融合多个检索结果列表
    results_list: 多个排序结果，每个元素是 [{"id": ..., "score": ...}, ...]
    k: RRF 参数，通常取 60
    """
    scores = {}
    for results in results_list:
        for rank, item in enumerate(results):
            doc_id = item["id"]
            if doc_id not in scores:
                scores[doc_id] = {"id": doc_id, "score": 0, "payload": item.get("payload", {})}
            scores[doc_id]["score"] += 1.0 / (k + rank + 1)
 
    sorted_results = sorted(scores.values(), key=lambda x: x["score"], reverse=True)
    return sorted_results[:top_n]

from qdrant_client.models import (
    Distance, VectorParams, SparseVectorParams,
    NamedVector, NamedSparseVector, SparseVector,
    SearchRequest, FusionQuery, Fusion,
)
 
# 创建支持 Hybrid Search 的 collection
client.create_collection(
    collection_name="hybrid_docs",
    vectors_config={
        "dense": VectorParams(size=1024, distance=Distance.COSINE),
    },
    sparse_vectors_config={
        "sparse": SparseVectorParams(),
    },
)
 
# BGE-M3 同时产出 dense 和 sparse 向量
from FlagEmbedding import BGEM3FlagModel
model = BGEM3FlagModel("BAAI/bge-m3", use_fp16=True)
output = model.encode("查询文本", return_dense=True, return_sparse=True)
 
dense_vector = output["dense_vecs"]
sparse_dict = output["lexical_weights"]  # {token_id: weight}
 
# Qdrant 的 query 接口直接支持 RRF 融合
results = client.query_points(
    collection_name="hybrid_docs",
    prefetch=[
        SearchRequest(
            query=NamedVector(name="dense", vector=dense_vector),
            limit=20,
        ),
        SearchRequest(
            query=NamedSparseVector(
                name="sparse",
                vector=SparseVector(
                    indices=list(sparse_dict.keys()),
                    values=list(sparse_dict.values()),
                ),
            ),
            limit=20,
        ),
    ],
    query=FusionQuery(fusion=Fusion.RRF),
    limit=10,
)

Embedding 查询文本:  20-50ms
向量检索:           10-30ms
Reranking:          50-200ms
LLM 生成 (TTFT):   200-2000ms
LLM 生成 (Decode):  2-30s

import hashlib
import json
import redis
 
r = redis.Redis()
 
def cached_embed(text: str, model, ttl: int = 3600) -> list[float]:
    """带 Redis 缓存的 embedding"""
    cache_key = f"emb:{hashlib.md5(text.encode()).hexdigest()}"
    cached = r.get(cache_key)
    if cached:
        return json.loads(cached)
 
    embedding = model.encode(text, normalize_embeddings=True).tolist()
    r.setex(cache_key, ttl, json.dumps(embedding))
    return embedding

from sentence_transformers import CrossEncoder
 
reranker = CrossEncoder("BAAI/bge-reranker-v2-m3", max_length=1024)
 
def rerank(query: str, documents: list[str], top_n: int = 5) -> list[tuple[int, float]]:
    """对检索结果重排序"""
    pairs = [(query, doc) for doc in documents]
    scores = reranker.predict(pairs)
 
    ranked = sorted(enumerate(scores), key=lambda x: x[1], reverse=True)
    return ranked[:top_n]

import asyncio
 
async def parallel_retrieve(
    query: str,
    collections: list[str],
    client: QdrantClient,
    top_k: int = 10,
) -> list[dict]:
    """并行检索多个 collection"""
    query_vector = embed(query)
 
    async def search_one(collection: str):
        return client.query_points(
            collection_name=collection,
            query=query_vector,
            limit=top_k,
        )
 
    tasks = [search_one(c) for c in collections]
    all_results = await asyncio.gather(*tasks)
 
    # 合并并去重
    merged = []
    seen_ids = set()
    for results in all_results:
        for r in results.points:
            if r.id not in seen_ids:
                seen_ids.add(r.id)
                merged.append(r)
 
    return merged

def agentic_rag(query: str, max_retries: int = 3) -> str:
    """带自纠正能力的 RAG"""
    current_query = query
 
    for attempt in range(max_retries):
        # 1. 检索
        docs = retrieve(current_query, top_k=10)
        docs = rerank(current_query, docs, top_n=5)
 
        # 2. LLM 评估检索质量
        relevance = llm_judge_relevance(query, docs)
 
        if relevance.score > 0.7:
            # 检索结果够好，直接生成
            return llm_generate(query, docs)
 
        # 3. 检索结果不行，让 LLM 改写查询
        current_query = llm_rewrite_query(
            original_query=query,
            failed_query=current_query,
            feedback=relevance.feedback,  # "检索结果主要在讲 X，但用户问的是 Y"
        )
 
    # 兜底：用最后一次检索结果硬生成
    return llm_generate(query, docs)
 
 
def llm_judge_relevance(query: str, docs: list[str]) -> RelevanceResult:
    """用 LLM 判断检索结果是否和问题相关"""
    prompt = f"""判断以下检索结果是否能回答用户的问题。
 
用户问题：{query}
 
检索结果：
{format_docs(docs)}
 
请给出：
1. 相关性分数（0-1）
2. 如果不相关，说明为什么不相关，以及建议用什么关键词重新检索"""
 
    return call_llm(prompt, response_format=RelevanceResult)