给本地大模型加知识库:sqlite-vec 混合检索完整实现
给本地大模型加知识库:sqlite-vec 混合检索完整实现
站内搜索
直接问 AI

给本地大模型加知识库:sqlite-vec 混合检索完整实现

这是「AI 模型本地部署」系列的最后一篇。前三篇做出了一个能从外网访问、已经接进日常工具的推理服务。这一篇给它加上知识库:让它能基于你自己的文档回答问题,并标注引用来源。

整套东西跑在同一台机器上,不依赖任何外部服务。会重点讲两件容易做错的事:怎么切分文档,以及中文检索为什么不能只靠向量

本地大模型部署系列(共 4 篇)① 选型与环境② 对外暴露③ 客户端接入④ 知识库 RAG。本文是第 ④ 篇。

整体结构

入库:文档 → 按标题切片 → 向量化 → sqlite-vec + FTS5
检索:查询 → 向量召回 + 全文召回 → RRF 融合 → 前 k 片塞进 system → 模型
组件 选型 理由
embedding bge-m3 FP16(1024 维) 多语言、支持 8K 输入;跑 CPU,不占显存
向量库 sqlite-vec 单文件,好备份好迁移,少一个要守护的进程
全文 SQLite FTS5 + trigram 内置,不用额外服务
融合 RRF 按名次融合,不用给两种不可比的分数做归一化

为什么 embedding 跑 CPU:生成模型已经把显存占得差不多了,而 embedding 比生成轻得多——查询时只需要嵌入一句话,感知不到延迟;批量入库时慢一点也无所谓。实测 bge-m3 FP16 在 CPU 上占 1.6 GB 内存,显存一点不动

第一步:起 embedding 服务

推理引擎的一个进程不能同时做生成和 embedding,需要第二个实例。

curl -L --retry 20 -C - \
  -o /opt/llm/models/bge-m3-FP16.gguf \
  "https://modelscope.cn/models/gpustack/bge-m3-GGUF/resolve/master/bge-m3-FP16.gguf"

embedding 对量化比生成敏感,而且反正跑 CPU、内存够用,直接上 FP16 不要压。

[Unit]
Description=bge-m3 embedding 服务(知识库用,跑 CPU)
After=network-online.target

[Service]
Type=exec
User=youruser
Environment=LD_LIBRARY_PATH=/opt/llama.cpp
ExecStart=/opt/llama.cpp/llama-server \
  --model /opt/llm/models/bge-m3-FP16.gguf \
  --alias bge-m3 \
  --embeddings \
  --pooling cls \
  --host 127.0.0.1 --port 8001 \
  --ctx-size 8192 \
  --batch-size 8192 --ubatch-size 8192 \
  --n-gpu-layers 0 \
  --threads 8 \
  --no-webui \
  --api-key sk-你的key
Restart=on-failure

[Install]
WantedBy=multi-user.target

三个关键参数:

  • --n-gpu-layers 0 强制走 CPU,不和生成模型抢显存。
  • --ubatch-size 8192 必须调大。默认是 512 token,而一个文档切片很容易超过——超了会直接返回 500:input (573 tokens) is too large to process。bge-m3 支持 8192 输入,调到和 ctx 一样即可。
  • --pooling cls 是 bge 系列的正确池化方式。

验证:

curl -s 127.0.0.1:8001/v1/embeddings \
  -H "Authorization: Bearer $KEY" -H 'Content-Type: application/json' \
  -d '{"model":"bge-m3","input":["测试文本"]}' \
| python3 -c "
import json,sys
v = json.load(sys.stdin)['data'][0]['embedding']
v = v[0] if v and isinstance(v[0], list) else v
print(f'{len(v)} 维')"

应该输出 1024 维注意有些版本会把单条结果再包一层 [[...]],代码里要处理这种情况。

第二步:建库

pip install sqlite-vec
import sqlite3, sqlite_vec

SCHEMA = """
CREATE TABLE IF NOT EXISTS documents (
  id INTEGER PRIMARY KEY, name TEXT NOT NULL,
  bytes INTEGER NOT NULL DEFAULT 0,
  n_chunks INTEGER NOT NULL DEFAULT 0,
  added_at INTEGER NOT NULL
);

CREATE TABLE IF NOT EXISTS chunks (
  id INTEGER PRIMARY KEY,
  doc_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE,
  ordinal INTEGER NOT NULL, heading TEXT, text TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_chunks_doc ON chunks(doc_id);

-- 全文索引。trigram 才能切中文;默认的 unicode61 会把整句当成一个 token
CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
  text, content='chunks', content_rowid='id', tokenize='trigram'
);
"""

def connect(path):
    con = sqlite3.connect(path, timeout=30)
    con.row_factory = sqlite3.Row
    con.execute("PRAGMA journal_mode=WAL")
    con.execute("PRAGMA foreign_keys=ON")
    con.enable_load_extension(True)
    sqlite_vec.load(con)
    con.enable_load_extension(False)
    return con

with connect("kb.db") as con:
    con.executescript(SCHEMA)
    con.execute(
        "CREATE VIRTUAL TABLE IF NOT EXISTS vec_chunks USING vec0("
        "chunk_id INTEGER PRIMARY KEY, embedding float[1024])")

tokenize='trigram' 是中文能用的前提。SQLite FTS5 默认的 unicode61 不做中文分词,一整句会变成一个 token,中文全文检索直接废掉。trigram 按三字滑窗切,中英混排都能用。

第三步:切分文档

这一步的顺序决定了检索质量。

直觉做法是”按字数递归切分,然后给每片打个标题标签”。这样做有两个问题:小文档会整篇变成一个切片(检索粒度全丢),而且一个切片跨越多节时,第一节的内容会被标上最后一节的标题——引用张冠李戴,而且很难发现

正确顺序是两级:先按标题分节(语义边界),节内超长再递归切

import re

CHUNK_TARGET = 1100      # 每片目标字符数
CHUNK_OVERLAP = 150      # 相邻片重叠,避免答案正好断在边界
MAX_CHUNK = 2000

_HEADING_RE = re.compile(r"^(#{1,6})\s+(.+?)\s*$", re.M)
_FENCE_RE = re.compile(r"^```.*?^```", re.M | re.S)

_SEPARATORS = ["\n# ", "\n## ", "\n### ", "\n\n", "\n",
               "。", "!", "?", ". ", ",", " "]


def _split_sections(text):
    """先按 markdown 标题切成小节,返回 [(面包屑标题, 正文), ...]"""
    # 排除围栏代码块 —— shell 里的 `# 注释` 和标题长得一模一样
    spans = [(m.start(), m.end()) for m in _FENCE_RE.finditer(text)]
    marks = [m for m in _HEADING_RE.finditer(text)
             if not any(a <= m.start() < b for a, b in spans)]
    if not marks:
        return [(None, text)]

    sections, trail = [], {}
    if marks[0].start() > 0:
        lead = text[:marks[0].start()].strip()
        if lead:
            sections.append((None, lead))

    for i, m in enumerate(marks):
        level, title = len(m.group(1)), m.group(2).strip()
        trail = {k: v for k, v in trail.items() if k < level}
        trail[level] = title

        end = marks[i + 1].start() if i + 1 < len(marks) else len(text)
        body = text[m.start():end].strip()
        content = body[m.end() - m.start():].strip()
        # 只有标题没有正文的(比如 H1 后面直接跟 H2)不单独成片 ——
        # 那种十几个字的空壳切片会在检索里当噪音顶掉真正的答案
        if len(content) < 20:
            continue
        sections.append((" › ".join(trail[k] for k in sorted(trail)), body))
    return sections

两个细节值得单独说:

  • 必须排除围栏代码块。shell 脚本里的 # 注释和 markdown 标题完全同形。不排除的话,引用面包屑会变成「README.md › 3. 下完记得撤掉——路由重启即失效…」这种莫名其妙的东西,那其实是 “`bash 块里的一行注释。
  • 空壳章节要跳过。只有标题没有正文的节切出来是个十几字的片段,它会在检索里当噪音——实测这种片段会排到第一位,把真正有答案的段落挤下去。

节内再按大小切,并加重叠:

def chunk_text(text):
    text = text.replace("\r\n", "\n").strip()
    chunks = []
    for heading, body in _split_sections(text):
        parts = (_split_recursive(body, CHUNK_TARGET, _SEPARATORS)
                 if len(body) > CHUNK_TARGET else [body])
        for i, part in enumerate(parts):
            piece = part.strip()
            if not piece:
                continue
            # 同一节内相邻片加重叠;跨节不加 —— 那会把别的章节内容
            # 混进来,反而污染引用
            if i > 0 and CHUNK_OVERLAP:
                tail = parts[i - 1][-CHUNK_OVERLAP:].strip()
                if tail:
                    piece = tail + "\n" + piece
            # 节内被拆开的片,把标题补回开头,否则第二片就丢了上下文
            if heading and not piece.lstrip().startswith("#"):
                piece = f"{heading}\n{piece}"
            chunks.append({"ordinal": len(chunks),
                           "text": piece[:MAX_CHUNK], "heading": heading})
    return chunks

第四步:混合检索

这是全篇最重要的一节。纯向量检索在中文技术文档上会失灵得比你想象的频繁。

两种检索方式的盲区是互补的:

  • 问「怎么让翻译不思考」——语义问题,只有向量能找到(原文里可能一个”不思考”都没有)。
  • reasoning_budget、某个 .gguf 文件名、某个型号——向量对这类标识符很不敏感,实测排不到正确段落;全文能精确命中。

所以两路都要,用 RRF(Reciprocal Rank Fusion)融合:

import json

def search(con, query, embed_fn, top_k=5, pool=30):
    qvec = embed_fn([query])[0]

    vec_hits = [r["chunk_id"] for r in con.execute(
        "SELECT chunk_id FROM vec_chunks "
        "WHERE embedding MATCH ? AND k = ? ORDER BY distance",
        (json.dumps(qvec), pool))]

    fts_hits = []
    for expr in _fts_queries(query):
        try:
            fts_hits = [r["rowid"] for r in con.execute(
                "SELECT rowid FROM chunks_fts WHERE chunks_fts MATCH ? "
                "ORDER BY rank LIMIT ?", (expr, pool))]
        except sqlite3.OperationalError:
            continue          # 这种写法 FTS5 不认,换下一种
        if fts_hits:
            break

    # RRF:按名次而不是分数融合,省得给两种不可比的分数做归一化
    K, score = 60, {}
    for rank, cid in enumerate(vec_hits):
        score[cid] = score.get(cid, 0.0) + 1.0 / (K + rank + 1)
    for rank, cid in enumerate(fts_hits):
        score[cid] = score.get(cid, 0.0) + 1.0 / (K + rank + 1)

    return sorted(score.items(), key=lambda kv: -kv[1])[:top_k]


def _fts_queries(query):
    """由严到松的一串 FTS5 表达式,逐个试到有命中为止。

    trigram 下短语匹配就是子串匹配,问句原样查基本命中不了 ——
    实测「reasoning_budget 为什么无效」一条都匹配不上,
    全文这一路等于白搭。拆词后做 OR 才能让标识符发挥作用。
    """
    q = query.replace('"', " ").strip()
    if not q:
        return []
    out = [f'"{q}"']                                  # 1. 整句子串
    terms = [t for t in re.split(r"[\s,,。??!!、::;;()()\[\]]+", q)
             if len(t) >= 3]                          # trigram 要求 ≥3 字符
    if terms:
        out.append(" OR ".join(f'"{t}"' for t in terms[:8]))
    return out

RRF 的好处是不需要调权重:向量的余弦距离和 FTS5 的 BM25 分数量纲完全不同,直接加权求和需要反复调参;按名次融合天然可比。K=60 是论文里的经验值,基本不用改。

第五步:拼上下文

def build_context(hits, max_chars=6000):
    """把检索结果拼成给模型看的上下文,并返回引用清单。

    有意控制在 6000 字符左右:小模型塞满 32K 反而会稀释注意力,
    3-5 段精准的片段比一大堆相关片段效果好。
    """
    blocks, cites, used = [], [], 0
    for i, h in enumerate(hits, 1):
        label = h["doc_name"] + (f" › {h['heading']}" if h.get("heading") else "")
        block = f"[{i}] 来源:{label}\n{h['text']}"
        if used + len(block) > max_chars:
            break
        blocks.append(block); used += len(block)
        cites.append({"n": i, "doc_name": h["doc_name"],
                      "heading": h.get("heading")})
    return "\n\n---\n\n".join(blocks), cites


RAG_PROMPT = (
    "下面是从知识库检索到的资料。回答时优先依据这些资料,"
    "并在用到某段时标注它的编号,例如 [1]。\n"
    "如果资料里没有能回答问题的内容,就直接说资料里没有,不要编造。\n\n"
)

最后一句很重要——明确允许模型说”资料里没有”,能显著减少它硬凑答案。

第六步:包成 OpenAI 兼容端点

做成一个和 /v1 用法完全一样的路径,任何 OpenAI 客户端换个 base_url 就能用上知识库:

@app.post("/rag/v1/chat/completions")
async def rag_chat(request: Request):
    body = await request.json()
    messages = body.get("messages") or []

    # 用最后一条 user 消息去检索
    last = next((m for m in reversed(messages) if m.get("role") == "user"), None)
    query = ""
    if last:
        c = last.get("content")
        query = c if isinstance(c, str) else " ".join(
            p.get("text", "") for p in (c or []) if p.get("type") == "text")

    cites = []
    if query.strip():
        try:
            hits = search(con, query, embed_fn, int(body.pop("rag_top_k", 5)))
            if hits:
                ctx, cites = build_context(hits)
                messages = [{"role": "system",
                             "content": RAG_PROMPT + ctx}] + messages
                body["messages"] = messages
        except Exception as exc:
            # 知识库出问题不该让对话直接失败 —— 退化成普通问答,
            # 把原因带在响应头里
            return await forward(request, body,
                                 {"X-RAG-Error": str(exc)[:200], "X-RAG-Hits": "0"})

    # 检索场景要的是照着资料答,不是自由发挥
    body.setdefault("chat_template_kwargs", {}).setdefault("enable_thinking", False)

    return await forward(request, body, {
        "X-RAG-Hits": str(len(cites)),
        "X-RAG-Sources": json.dumps([c["doc_name"] for c in cites],
                                    ensure_ascii=False)[:500],
    })

三个设计决定:

  • 检索失败要降级,不要报错。知识库挂了、embedding 服务没起来,对话应该照常能用,只是没有资料支撑。原因放在响应头里,客户端不用改也能看到。
  • 默认关思考。检索场景要的是照着资料答。
  • 把命中情况放响应头。X-RAG-HitsX-RAG-Sources 让你不改客户端就能看到用了哪些文档,调试时非常有用。

验证

入库几篇你自己的文档,然后问一个只有文档里才有答案的问题:

curl -sS -D /tmp/h.txt https://你的域名/rag/v1/chat/completions \
  -H "Authorization: Bearer $KEY" -H 'Content-Type: application/json' \
  -d '{"model":"gemma4","max_tokens":2048,
       "messages":[{"role":"user","content":"你文档里那个具体数字是多少?"}]}' \
| python3 -c "import json,sys; print(json.load(sys.stdin)['choices'][0]['message']['content'])"

grep -i "x-rag" /tmp/h.txt

好的结果应该是:答案里出现文档中的准确数字,并带 [1] 这样的引用标记,响应头里能看到命中数和来源文件名。

单独测检索(不生成)对调参更方便——可以直接看每条命中是靠向量还是全文来的,快速判断是切分有问题还是检索策略有问题。

已知限制

  • 不解析 PDF / Word。这是有意的:宁可不支持,也不要悄悄把乱码入库——那种问题在检索结果变差时极难定位。要用先自己转成文本。
  • 没有 reranker。top-k 直接进上下文。想更准可以再加一个 bge-reranker 模型,代价是又一份内存和一次额外推理。
  • 上下文控制在 6000 字符。小模型塞满长上下文反而稀释注意力,精准的 3–5 段胜过一堆相关片段。

系列回顾

四篇下来,你有了一套完整的本地推理栈:

  1. 选型与部署——按显存选对模型和引擎,起成系统服务
  2. 对外暴露——隧道和直连两条路,含证书和路由
  3. 接入工具——翻译扩展、编辑器、自己的程序
  4. 知识库——本篇

全程踩到的坑我另外整理了一篇按现象索引的排查记录,遇到”应该能跑但就是不对”的情况可以对照着看。

最后一条经验,比任何具体配置都有用:凡是”成功也可能什么都没做”的步骤,后面都要加一道独立的事实校验。下载命令返回 0 不代表文件下全了,参数写了不代表引擎采纳了,服务在监听不代表包能到达。养成看启动日志里”它实际算出来的值”、而不是”我以为我设了什么”的习惯,能省掉大量排查时间。

发表回复

向下探索