喂给索引的两端:多后端 embedding 计算与 AST 感知分块

30 秒导读: 一段文本要进 LEANN 的向量索引,先得过两道工序:切块(chunking)——把长文档/代码切成合适大小的片段;算向量(embedding)——把每个片段变成一串浮点数。本章讲这两道"进索引之前"的工序:分块怎么按代码结构(函数级)聪明地切,向量怎么在多种后端(本地 sentence-transformers、Apple Silicon 的 mlx、OpenAI/Gemini/Ollama 云 API)之间分发,以及为什么 build 时直算、search 时走常驻 server 这个分叉如此重要。

本章是 LEANN 系列的一章。其它章见:全景导览、核心存储魔法、查询端到端、可插拔后端、CLI 与 Agent。本章只讲"喂给索引的两端",不重复它们的内容。

---

1. 这是什么(零基础也能懂)

一句话定义: 这是 LEANN 里把"原始文本"变成"能进向量索引的东西"的两个前置子系统——分块器(把长文本切小)和 embedding 计算器(把小块算成向量)。

为什么需要它们: 向量检索的前提是"每段文本都有一个向量"。但你不能把一整个 10 万行的代码库或一本 PDF 塞给 embedding 模型——模型有 token 上限(常见 512),而且"检索的粒度"太粗会让结果没用。所以必须先切块,再逐块算向量。

两道工序,各解决一个小问题:

工序	输入	输出	解决什么
分块 chunking	一篇文档 / 一个代码文件	一串文本块 {"text", "metadata"}	让每块落进模型 token 上限,且检索粒度合理
embedding 计算	一串文本块	一个向量矩阵 (N, dim)	把文字变成可做相似度检索的浮点向量

用起来什么样(直观感受): 你几乎不会直接调它们——leann build 内部就串起了这两步。但心智模型是这样一条流水线:

文档/代码  ──▶  chunking_utils.create_text_chunks  ──▶  [块1, 块2, 块3, …]
                                                              │
                                                              ▼
                                          embedding_compute.compute_embeddings
                                                              │
                                                              ▼
                                                   向量矩阵 (N, dim) ──▶ 进索引

一句话直觉: 把 embedding 模型想成一台"翻译机"——只能一小口一小口喂(token 上限),所以你得先把大文档切成一口大小的块,再排队喂进去。分块决定"切法",embedding 决定"用哪台翻译机、在哪儿翻"。

---

2. 顶层全景(它大概怎么转)

本章覆盖三个文件、三个子系统。先看它们各自的职责与位置:

子系统	干什么	文件	关键符号
embedding 计算	按 mode 分发到不同后端算向量	packages/leann-core/src/leann/embedding_compute.py	compute_embeddings
embedding server + 守护进程	让 search 时的模型常驻、跨进程复用	packages/leann-core/src/leann/embedding_server_manager.py	EmbeddingServerManager.start_server
AST 感知分块	用 tree-sitter 按函数级切代码	packages/leann-core/src/leann/chunking_utils.py	create_ast_chunks / create_text_chunks

最重要的一条暗线:build 和 search 走的是两条不同的 embedding 路径。 这条分叉贯穿全章,先记住这张图(从左看):

                        api.compute_embeddings(use_server=?)
                                    │
                 use_server=False   │   use_server=True
              (build_index 走这条)   │  (查询/search 走这条)
                        ┌───────────┴───────────┐
                        ▼                        ▼
          embedding_compute            compute_embeddings_via_server
          .compute_embeddings          （ZMQ REQ → 常驻 server）
          (直接在本进程算)                       │
                        │                        ▼
             按 mode 分发到:            EmbeddingServerManager
        sentence-transformers          启动/复用一个持有模型的子进程
        / mlx / openai / gemini        （避免每次查询都重新加载模型）
        / ollama

为什么要分叉? build 是一次性的大批量计算(算完就写盘),在本进程直算最简单;search 则是反复、按需地重算向量(见 02-search-recompute.md 的选择性重算),如果每次查询都重新加载几百 MB 的模型会慢到不能用——所以 search 端把模型放进一个常驻 server,通过 ZMQ 通信,让模型只加载一次、反复复用。

---

3. 核心机制一:embedding 计算的多 mode 分发

3.1 它要解决的小问题

同一个 LEANN,用户可能想用本地小模型(省钱、离线)、也可能想用 OpenAI 云 API(质量高)、Mac 用户想用 Apple Silicon 加速。同一个入口,按一个字符串参数分发到不同实现,这就是 compute_embeddings 的活。

3.2 直觉:一个 mode 字符串的调度台

入口函数签名里,mode 就是那个调度开关(embedding_compute.py:422 compute_embeddings)。它本身不算向量,只做分发:

# 示意,非源码:分发的骨架
def compute_embeddings(texts, model_name, mode="sentence-transformers", ...):
    if mode == "sentence-transformers":      # 默认:本地 HuggingFace 模型
        return compute_embeddings_sentence_transformers(...)
    elif mode == "openai":                    # 云 API(也兼容 Gemini/DashScope 等 OpenAI 协议端点)
        return compute_embeddings_openai(...)
    elif mode == "mlx":                       # Apple Silicon 原生加速
        return compute_embeddings_mlx(...)
    elif mode == "ollama":                    # 本地 Ollama 服务
        return compute_embeddings_ollama(...)
    elif mode == "gemini":                    # Google GenAI SDK
        return compute_embeddings_gemini(...)
    else:
        raise ValueError(f"Unsupported embedding mode: {mode}")

真实分发就在 embedding_compute.py:458-499——一串 if/elif 按 mode 派发,最后 else 抛 ValueError。重点看:每个分支是一个独立函数,互不干扰,新增一种后端只要加一个分支。

3.3 五种后端各自的取舍

mode	实现函数(embedding_compute.py)	跑在哪	关键细节
sentence-transformers(默认)	compute_embeddings_sentence_transformers :502	本地(CUDA/MPS/CPU 自动选)	模型缓存 _model_cache、FP16、CUDA torch.compile、OOM 自动降 batch
openai	compute_embeddings_openai :846	云 API	兼容 Gemini/DashScope 的 OpenAI 端点,按端点自动缩 batch
mlx	compute_embeddings_mlx :974	Apple Silicon 原生	用 mlx.core + mlx_lm,手动分批 + mean pooling
ollama	compute_embeddings_ollama :1051	本地 Ollama	用 /api/embed 批量端点,自己做 L2 归一化
gemini	compute_embeddings_gemini :1374	Google GenAI	task_type="RETRIEVAL_DOCUMENT",batch ≤ 100

sentence-transformers(默认路径)里几个不显然的巧思:

• 模型缓存:cache key 拼进 model_name/device/use_fp16,加载一次后放进 _model_cache,后续复用不再重载(embedding_compute.py:554、:564)。
• 设备自动探测:device="auto" 时按 LEANN_EMBEDDING_DEVICE 环境变量 → CUDA → MPS → CPU 依次探测(embedding_compute.py:536-547)。
• CUDA OOM 自愈:_encode_with_oom_retry 在遇到显存不足时把 batch_size 砍半重试,直到 batch=1 才放弃(embedding_compute.py:383-419)。
• CUDA 按显存封顶 batch:_cap_cuda_batch_by_vram 按剩余显存估算安全 batch(注意力峰值内存 ~O(seq²))(embedding_compute.py:355-380)。

3.4 token 上限:切完还要再兜底截断

分块时已经控制了块大小(见 §5),但云端模型的真实 token 上限各不相同,所以 embedding 时再兜一层截断。get_model_token_limit 用"动态发现 + 注册表兜底"的混合策略(embedding_compute.py:61):

• Ollama:调 /api/show 动态查真实 context_length(_query_ollama_context_limit :202)。
• LM Studio:起一个 Node.js 子进程调 @lmstudio/sdk 查上限(_query_lmstudio_context_limit :235)。
• 兜底:查不到就用内置注册表 EMBEDDING_MODEL_LIMITS(:39,如 text-embedding-3-small = 8192),再兜底默认 2048。
• 查到的结果按 (model_name, base_url) 缓存在 _token_limit_cache,避免重复 API 调用(:58、:81)。

拿到上限后 truncate_to_token_limit 用 tiktoken 的 cl100k_base 逐条编码、超限就截断(embedding_compute.py:136)。OpenAI / Ollama 路径都在批处理前调它(:901、:1223)。

---

4. 核心机制二:embedding server 生命周期与守护进程复用

4.1 它要解决的小问题

search 端要反复重算向量(每次遍历图都要给若干节点算 embedding)。如果每次都在本进程从磁盘加载几百 MB 的模型,冷启动能吃掉几秒——对一次几十毫秒的查询是灾难。把模型放进一个常驻子进程,只加载一次、反复用,这就是 EmbeddingServerManager 的使命。

4.2 直觉:build 直算、search 走 server

回到 §2 那条分叉。api.py 里的 compute_embeddings(use_server=...) 是分水岭(api.py:48):

• build:build_index 里两次调用都写死 use_server=False(api.py:546、:585),走本进程直算——一次性批量,不需要常驻。
• search:走 use_server=True → compute_embeddings_via_server(api.py:93),用 ZMQ REQ socket 把文本块 msgpack 打包发给 server,收回向量矩阵。server 由 EmbeddingServerManager 负责起。

搜索路径上,searcher_base.py:88 的 start_server(...) 负责确保 server 就绪,拿到端口后查询才通过该端口做 ZMQ 通信。

4.3 start_server 的三段决策

EmbeddingServerManager.start_server(embedding_server_manager.py:242)按这个顺序决定"复用还是新起":

start_server(port, model_name, embedding_mode, use_daemon, daemon_ttl_seconds, …)
        │
        ▼  ① 算配置指纹 config_signature（model+mode+passages+distance+provider_options+文件mtime）
        │
        ├─▶ ② 本进程已有同指纹 live server?  ──是──▶ 直接复用（"Reusing in-process server"）
        │
        ├─▶ ③ use_daemon 且非 Colab?  ──▶ 查跨进程 registry 里有没有同指纹的活 daemon
        │        （_adopt_registered_server：校验 pid 活着 + 端口在听）──命中──▶ 复用它
        │
        └─▶ ④ 都没有 ──▶ _start_new_server 起新子进程；若 daemon 模式则写 registry 记录

配置指纹(config signature) 是复用的钥匙:_build_config_signature(:349)把 model_name、embedding_mode、passages_file、distance_metric、provider_options,以及 passages 文件的 mtime 签名(_build_passages_signature :157)一起哈希(_registry_key :789 用 SHA-256)。只有指纹完全一致才复用——换了模型、换了距离度量、或索引文件被改过,就不会误用旧 server。

4.4 跨进程复用:registry 让不同进程共享同一个 daemon

普通 server 随起它的进程消亡;daemon 则跨进程存活并被别的进程"领养"。机制是文件 registry:

• 注册目录:~/.leann/servers/(_registry_dir :724)。
• 写记录:起 daemon 成功后,_write_registry_record(:797)把 {pid, port, backend, ttl, created_at, config_signature} 原子写成 <指纹>.json。
• 领养:别的进程调 _adopt_registered_server(:820),按指纹找到 json,校验 backend_module_name 一致、config_signature 一致、pid 还活着(_pid_is_alive :120)、端口还在听(_check_port :114);全过才复用,否则删掉陈旧记录。
• 子进程标志:_build_server_command(:451)在 use_daemon=True 时给子进程加 --daemon-mode 和 --daemon-ttl <秒> 参数(:475-478),默认 TTL 900 秒。
• 并发安全:领养/新起这段"检查-然后-启动"的临界区被双层锁保护(_registry_lock :728)——进程内用 threading.Lock(每个 registry key 一把),跨进程用 fcntl.flock(POSIX)/ msvcrt(Windows),避免两个进程同时起两个 server。

为什么值得这么麻烦? 省冷启动。你连续跑 leann search foo、leann ask foo,第二条命令是全新进程,但只要索引没变,它能直接领养第一条命令留下的、还没到 TTL 的常驻 daemon——模型加载一次,后续多次查询共享。

4.5 detach 语义与 daemon 列表

• stop 的非对称语义:stop_server(:581)对 daemon 模式只 detach 不杀——把本地引用清空、让 daemon 继续为别的进程服务;非 daemon 的临时 server 才真正 terminate/kill。
• 列举:list_daemons(类方法 :852)遍历 registry 目录,逐条校验 pid + 端口是否还活着,顺手清掉死记录,返回活着的 daemon 列表。CLI 的 leann daemon status 就打印它(cli.py:3232)。
• 批量停:stop_daemons(:877)可按 backend 或 passages 文件过滤,os.kill(pid, 15) 优雅停掉并删记录。

---

5. 核心机制三:AST 感知分块

5.1 它要解决的小问题

传统按字符/句子切文本,遇到代码会把一个函数从中间劈开——半个函数当一块,检索出来毫无意义。AST(抽象语法树)感知分块的思路:先解析代码结构,尽量按函数/类这种语法单元的边界切,让每块是一个"语义完整"的代码片段。

5.2 直觉:先分流,代码走 AST,文本走句子切分

统一入口是 create_text_chunks(chunking_utils.py:346)。当 use_ast_chunking=True 时,它先把文档分成"代码文件"和"普通文本文件"两堆,分别处理:

create_text_chunks(use_ast_chunking=True)
        │
        ▼  detect_code_files() 按扩展名分流
        ├──▶ 代码文件 ──▶ create_ast_chunks（astchunk + tree-sitter，函数级切)
        │                       │（astchunk 缺失或失败)
        │                       └──▶ 回退 create_traditional_chunks（句子切分)
        │
        └──▶ 文本文件 ──▶ create_traditional_chunks（LlamaIndex SentenceSplitter)

分流靠扩展名:detect_code_files(:142)查每个文档的 file_path 后缀,命中 CODE_EXTENSIONS(:131,支持 .py/.java/.cs/.ts/.tsx/.js/.jsx)就打上 is_code=True 和 language 标签。辅助函数 get_language_from_extension(:169)做同样的扩展名 → 语言映射。

5.3 真实实现:create_ast_chunks 用 astchunk 切

create_ast_chunks(chunking_utils.py:175)对每个代码文档:

1. 拿 doc.metadata["language"],没有语言就回退传统切分(:197-201)。
2. 配置 ASTChunkBuilder(来自可选依赖 astchunk,内部用 tree-sitter),传入 max_chunk_size、language、chunk_overlap 等(:219-235)。
3. chunk_builder.chunkify(code_content) 切块(:241),再兼容 astchunk 可能返回的多种格式(.text 属性 / str / {"content", "metadata"} dict)统一成 {"text", "metadata"}(:242-277)。
4. astchunk 不可用就整体回退:文件顶部 try: from astchunk import ASTChunkBuilder 失败时,直接走 _traditional_chunks_as_dicts(:188-193);单个文件切失败也会 except 后回退(:282-285)。

行号导航的小巧思:AST 块若带 start_line_no 元数据,create_text_chunks 会给每行前缀行号(chunking_utils.py:390-398),这样检索出的代码块自带"这段在第几行",方便跳回源码。

5.4 chunk size 的 token/char 换算与 overlap

一个坑:传统切分以 token 计,AST 切分以字符计。 calculate_safe_chunk_size(chunking_utils.py:41)专门处理这个换算:

• 先留安全边际:safe_limit = model_token_limit * safety_factor(默认 0.9,即留 10%)。
• traditional 模式(按 token):chunk_size = safe_limit - overlap_size(:64)。因为最大块 = chunk_size + overlap,反推出 chunk_size。
• AST 模式(按字符):按"代码约 1.2 token/字符"换算——safe_chars = safe_limit / 1.2,overlap 也从 token 换成 overlap * 3 字符(:66-70)。

overlap(块间重叠) 是让相邻块共享一小段内容,避免答案正好卡在两块交界处被切断。传统切分在 create_traditional_chunks(:290)里校正 overlap:负数归零、overlap >= chunk_size 时砍到 chunk_size // 2(:301-304)。

AST 超限预警:create_ast_chunks 会估算 (max_chunk_size + chunk_overlap) * 1.2 token,超 512 就警告一次(用全局 _ast_token_warning_shown 保证每会话只警告一次)(:206-217)。

5.5 最后一道保险:validate_chunk_token_limits

validate_chunk_token_limits(chunking_utils.py:73)是分块侧的最终截断兜底:用 tiktoken 逐块估 token,超 max_tokens(默认 512)就截断,返回 (validated_chunks, num_truncated)。tiktoken 缺失时回退到 max_tokens / 1.2 的字符截断(:108-120)。

注意两层截断的分工:分块侧的 validate_chunk_token_limits 用固定阈值(512),而真正的动态截断已挪到 embedding 时——create_text_chunks 尾部注释明说"Token truncation is now handled at embedding time with dynamic model limits"(:415-416),对应 §3.4 的 get_model_token_limit + truncate_to_token_limit。

---

6. 巧妙之处(可借鉴的技术)

• 归一化模型自动切 cosine。 OpenAI/Voyage/Cohere 这类模型输出的向量已 L2 归一化(范数=1),此时用 MIPS(最大内积)不如 cosine。LeannBuilder 构造时检测:先查白名单集合 normalized_embeddings_models(api.py:429-445),再按名字模式兜底(:461-475);命中且用户没显式设 distance_metric 就自动设成 cosine(api.py:478-479),并 warn 提醒;若用户设了别的度量则只警告不覆盖(:487-495)。借鉴点:根据模型特性自动选对距离度量,省掉用户踩坑。

• build/search 双路径,一个 use_server 布尔切换。 用一个参数把"一次性批量直算"和"常驻 server 反复算"两种截然不同的需求统一在同一入口(api.py:48-90),调用方无需关心底层。

• 配置指纹 = 复用的安全边界。 daemon 复用不是"同名就复用",而是把模型、模式、距离度量、甚至索引文件的 mtime 都哈进指纹(embedding_server_manager.py:349-369),索引一变指纹就变、绝不误用旧向量。

• 多后端 API 端点的 batch 自适应。 compute_embeddings_openai 检测到 Gemini 端点(generativelanguage.googleapis.com)就把 batch 压到 100、DashScope 压到 10(embedding_compute.py:915-931)——同一份 OpenAI 兼容代码适配各家硬限制。

• 失败降级链贯穿始终。 AST 切分失败 → 传统切分;模型本地加载失败 → 网络下载;CUDA OOM → 砍半 batch;daemon 记录陈旧 → 删掉重起。每一层都不让单点失败拖垮整条流水线。

---

7. 边界与局限

• AST 分块只覆盖 7 种扩展名(CODE_EXTENSIONS chunking_utils.py:131)。其余语言/文件一律走传统句子切分——不是"所有代码都函数级切"。
• astchunk 是可选依赖。 没装 astchunk 时 AST 分块静默回退成传统切分(:188-193),质量下降但不报错——用户可能没意识到没吃到 AST 的好处。
• token 估算是近似。 "1.2 token/字符""4 字符/token"都是保守估计(estimate_token_count :18、calculate_safe_chunk_size :68),不是精确值;真正精确的是 embedding 时用 tiktoken 的动态截断。
• daemon 复用有 TTL 窗口。 默认 900 秒(:255),过期后 daemon 自行退出,下次查询要重新冷启动。
• Colab 环境不用 daemon。 _is_colab_environment(:96)为真时走独立的 _start_server_colab 路径(:371),不做跨进程复用。
• cosine 自动检测靠名字模式(api.py:461-475),自定义命名的归一化模型可能漏检——此时需用户手动传 distance_metric="cosine"。

---

8. 代码地图(导航索引)

主题	文件路径	符号名
embedding 多 mode 分发入口	packages/leann-core/src/leann/embedding_compute.py	compute_embeddings
默认后端(本地 HF 模型)	packages/leann-core/src/leann/embedding_compute.py	compute_embeddings_sentence_transformers
OpenAI / 兼容端点(Gemini/DashScope)	packages/leann-core/src/leann/embedding_compute.py	compute_embeddings_openai
Apple Silicon 加速	packages/leann-core/src/leann/embedding_compute.py	compute_embeddings_mlx
Ollama / Gemini 后端	packages/leann-core/src/leann/embedding_compute.py	compute_embeddings_ollama / compute_embeddings_gemini
token 上限动态发现 + 注册表	packages/leann-core/src/leann/embedding_compute.py	get_model_token_limit / EMBEDDING_MODEL_LIMITS
embedding 时截断	packages/leann-core/src/leann/embedding_compute.py	truncate_to_token_limit
CUDA OOM 降 batch 重试	packages/leann-core/src/leann/embedding_compute.py	_encode_with_oom_retry
build/search 双路径分叉	packages/leann-core/src/leann/api.py	compute_embeddings (use_server)
search 端 ZMQ 调用 server	packages/leann-core/src/leann/api.py	compute_embeddings_via_server
归一化模型自动切 cosine	packages/leann-core/src/leann/api.py	LeannBuilder.__init__(normalized_embeddings_models)
server 启动/复用决策	packages/leann-core/src/leann/embedding_server_manager.py	EmbeddingServerManager.start_server
配置指纹	packages/leann-core/src/leann/embedding_server_manager.py	_build_config_signature / _registry_key
跨进程 daemon 领养	packages/leann-core/src/leann/embedding_server_manager.py	_adopt_registered_server / _write_registry_record
daemon 子进程参数	packages/leann-core/src/leann/embedding_server_manager.py	_build_server_command(--daemon-mode)
列举/停止 daemon	packages/leann-core/src/leann/embedding_server_manager.py	list_daemons / stop_daemons
搜索侧起 server 的调用点	packages/leann-core/src/leann/searcher_base.py	start_server(:88)
分块统一入口(代码/文本分流)	packages/leann-core/src/leann/chunking_utils.py	create_text_chunks
AST 感知分块	packages/leann-core/src/leann/chunking_utils.py	create_ast_chunks
代码文件识别 / 语言映射	packages/leann-core/src/leann/chunking_utils.py	detect_code_files / get_language_from_extension / CODE_EXTENSIONS
chunk size 的 token/char 换算	packages/leann-core/src/leann/chunking_utils.py	calculate_safe_chunk_size
分块侧截断兜底	packages/leann-core/src/leann/chunking_utils.py	validate_chunk_token_limits
传统句子切分	packages/leann-core/src/leann/chunking_utils.py	create_traditional_chunks