RAG 检索服务的工程化

import re
from dataclasses import dataclass

@dataclass
class Chunk:
    text: str
    start: int
    end: int

def recursive_split(text: str, chunk_size: int = 512,
                    overlap: int = 64) -> list[Chunk]:
    # 按段落 -> 句子 -> 字符逐级回退，尽量在自然边界处切分
    separators = ["\n\n", "\n", "。", "！", "？", ". ", " "]
    pieces: list[str] = [text]
    for sep in separators:
        merged: list[str] = []
        for piece in pieces:
            if len(piece) <= chunk_size:
                merged.append(piece)
            else:
                parts = [p + sep for p in piece.split(sep) if p]
                merged.extend(parts or [piece])
        pieces = merged
    chunks, buf = [], ""
    for piece in pieces:
        if len(buf) + len(piece) > chunk_size and buf:
            chunks.append(buf)
            buf = buf[-overlap:] if overlap else ""
        buf += piece
    if buf.strip():
        chunks.append(buf)
    out, cursor = [], 0
    for c in chunks:
        idx = text.find(c[:32], cursor)
        idx = idx if idx >= 0 else cursor
        out.append(Chunk(text=c, start=idx, end=idx + len(c)))
        cursor = idx
    return out

collection:
  name: docs_zh
  vectors:
    size: 1024                 # bge-m3 维度
    distance: Cosine
  hnsw_config:
    m: 32                      # 连接数，内存换精度
    ef_construct: 256          # 建索引候选，离线可大
  optimizers_config:
    memmap_threshold: 20000    # 超过 2 万条转 mmap 省内存
  quantization_config:
    scalar:
      type: int8               # 标量量化，内存降 4 倍，精度损失约 1%

import time

def sweep_ef_search(index, queries, ground_truth, ef_values, top_k: int = 10):
    rows = []
    for ef in ef_values:
        index.set_ef(ef)                       # Milvus/Qdrant 均支持动态设置
        hits, latency = 0, []
        for q, truth in zip(queries, ground_truth):
            t0 = time.perf_counter()
            res = index.search(q, limit=top_k)
            latency.append((time.perf_counter() - t0) * 1000)
            hits += len(set(r[0] for r in res) & set(truth))
        recall = hits / (len(queries) * top_k)
        p99 = sorted(latency)[int(len(latency) * 0.99) - 1]
        rows.append({"ef_search": ef, "recall": round(recall, 4),
                     "p99_ms": round(p99, 2)})
    return rows

if __name__ == "__main__":
    for row in sweep_ef_search(None, [], [], [32, 64, 128, 256, 512]):
        print(row)

from collections import defaultdict

def rrf_fuse(ranked_lists: list[list[str]], k: int = 60) -> list[tuple[str, float]]:
    scores: dict[str, float] = defaultdict(float)
    for ranked in ranked_lists:
        for rank, doc_id in enumerate(ranked, start=1):
            scores[doc_id] += 1.0 / (k + rank)
    return sorted(scores.items(), key=lambda kv: kv[1], reverse=True)

if __name__ == "__main__":
    vector_hits = ["d1", "d2", "d3", "d4"]
    bm25_hits = ["d3", "d5", "d1", "d6"]
    for doc_id, score in rrf_fuse([vector_hits, bm25_hits])[:4]:
        print(doc_id, round(score, 5))

from collections import defaultdict

def bm25_search(query: str, corpus: dict[str, str], top_n: int = 50) -> list[str]:
    # 简化 BM25：按词频打分，生产环境用 rank_bm25 或 ES
    terms = [t for t in query.lower().split() if t]
    scored = []
    for doc_id, text in corpus.items():
        tl = text.lower()
        score = sum(tl.count(t) for t in terms)
        if score > 0:
            scored.append((doc_id, score))
    scored.sort(key=lambda kv: kv[1], reverse=True)
    return [d for d, _ in scored[:top_n]]

def vector_search(query: str, embed_fn, index, top_n: int = 50) -> list[str]:
    qvec = embed_fn(query)
    hits = index.search(qvec, limit=top_n)      # 返回 (doc_id, score)
    return [doc_id for doc_id, _ in hits]

def rrf_fuse(ranked_lists: list[list[str]], k: int = 60) -> list[str]:
    scores: dict[str, float] = defaultdict(float)
    for ranked in ranked_lists:
        for rank, doc_id in enumerate(ranked, start=1):
            scores[doc_id] += 1.0 / (k + rank)
    return [d for d, _ in sorted(scores.items(), key=lambda kv: kv[1], reverse=True)]

def rerank(query: str, doc_ids: list[str], corpus: dict[str, str],
           rerank_fn, top_k: int = 3) -> list[tuple[str, float]]:
    pairs = [(query, corpus[d]) for d in doc_ids]
    scores = rerank_fn(pairs)                    # bge-reranker-v2-m3 打分
    ranked = sorted(zip(doc_ids, scores), key=lambda kv: kv[1], reverse=True)
    return ranked[:top_k]

def rag_retrieve(query: str, corpus, embed_fn, index, rerank_fn,
                 recall_n: int = 50, top_k: int = 3) -> list[tuple[str, float]]:
    bm25_hits = bm25_search(query, corpus, recall_n)
    vec_hits = vector_search(query, embed_fn, index, recall_n)
    fused = rrf_fuse([bm25_hits, vec_hits])[:recall_n]
    return rerank(query, fused, corpus, rerank_fn, top_k)

if __name__ == "__main__":
    corpus = {"d1": "错误码 E1042 表示鉴权失败", "d2": "鉴权流程说明", "d3": "E1042 排查步骤"}
    def fake_embed(q): return [0.1, 0.2, 0.3]
    class FakeIndex:
        def search(self, qvec, limit): return [("d2", 0.9), ("d3", 0.8), ("d1", 0.7)]
    def fake_rerank(pairs): return [0.3, 0.95, 0.6][:len(pairs)]
    print(rag_retrieve("E1042 怎么排查", corpus, fake_embed, FakeIndex(),
                       fake_rerank))

import hashlib
import json

def retrieval_cache_key(query: str, filters: dict, top_k: int,
                        index_version: str) -> str:
    payload = json.dumps({
        "q": query.strip().lower(),
        "filters": filters,
        "top_k": top_k,
        "idx": index_version,        # 索引更新后键自动失效
    }, sort_keys=True, ensure_ascii=False)
    return "rag:retrieval:" + hashlib.sha256(payload.encode()).hexdigest()

if __name__ == "__main__":
    k1 = retrieval_cache_key("E1042 排查", {"tenant": "a"}, 50, "v3")
    k2 = retrieval_cache_key("E1042 排查", {"tenant": "a"}, 50, "v4")
    print(k1 != k2)   # 索引版本变化 -> 键变化 -> 缓存自动失效