跳到主要内容

Khoj 检索管线 — 一次语义搜索的完整链路

本章讲出库:一个查询字符串怎么变成一串排好序的命中片段。核心两件事:pgvector 余弦召回,和**"双编码器召回 + 交叉编码器重排"的两段式排序**。

1. 为什么要"两段式"(先建直觉)​

检索质量和速度天生打架:

  • 双编码器(bi-encoder)——文档向量提前算好,查询来了只算一次距离。快,但粗:它是"隔空比对"两个独立向量,细微的相关性差异会丢。
  • 交叉编码器(cross-encoder)——把"查询 + 一个候选片段"拼在一起过一遍模型,直接输出相关分。准,但慢:每个候选都要过一次模型,不能预计算。

Khoj 的策略是经典的"召回-重排"两段式:先用便宜的双编码器从全库快速捞出 top-10 候选,再用昂贵的交叉编码器只对这 10 个精排。用快的做粗筛,用准的做精排,兼顾两头。

全库(可能几十万块)
      │  ① 双编码器召回(pgvector 余弦距离,便宜)
      ▼
  top-10 候选
      │  ② 交叉编码器重排(只对 10 个,昂贵但可控)
      ▼
  精排后取前 n 条

2. 入口:text_search.query​

检索的核心函数是 query(src/khoj/search_type/text_search.py:99)。它做三件事:定阈值、编码查询、调数据库召回。

# 真实源码风 —— text_search.py:112-135(节选)
search_model = get_default_search_model()
if not max_distance:
    max_distance = search_model.bi_encoder_confidence_threshold or math.inf  # 默认 0.18

if question_embedding is None:
    question_embedding = state.embeddings_model[search_model.name].embed_query(query)  # 编码查询

hits = EntryAdapters.search_with_embeddings(       # 真正的向量查询
    raw_query=raw_query, embeddings=question_embedding,
    max_results=10, max_distance=max_distance, user=user, agent=agent,
).all()

三点值得注意:

  1. 查询用的是同一个双编码器 embed_query——和入库用的 embed_documents 是同一个模型,只有"文档/查询"模式差别。这保证查询向量和文档向量在同一个空间里可比。
  2. top_k = 10 是写死的(text_search.py:125),即无论用户要几条,底层召回都先取 10 个候选再精排。
  3. max_distance 默认取模型的置信度阈值 0.18(SearchModelConfig.bi_encoder_confidence_threshold,models/__init__.py:582),超过这个余弦距离的块被认为"不够相关",直接剪掉。

3. 召回:pgvector 的余弦距离查询​

真正的向量查询在 EntryAdapters.search_with_embeddings(src/khoj/database/adapters/__init__.py:2144)。它把"找最近向量"这件事交给 Postgres 的 pgvector 扩展,用 Django ORM 表达:

# 真实源码 —— adapters/__init__.py:2163-2172(节选)
relevant_entries = EntryAdapters.apply_filters(user, raw_query, file_type_filter, agent)
relevant_entries = relevant_entries.filter(owner_filter).annotate(
    distance=CosineDistance("embeddings", embeddings)      # ① 每块算与查询的余弦距离
)
relevant_entries = relevant_entries.filter(distance__lte=max_distance)  # ② 剪掉太远的
relevant_entries = relevant_entries.order_by("distance")   # ③ 按距离升序
return relevant_entries[:max_results]                      # ④ 取前 10

怎么读这段:CosineDistance("embeddings", embeddings) 让数据库对每个候选块计算它的向量与查询向量的余弦距离(pgvector 的 <=> 操作符),annotate 把这个距离挂成一个可排序、可过滤的字段;然后 distance__lte=max_distance 剪枝、order_by("distance") 升序、切片取 top-10。距离越小 = 意义越接近。

owner_filter 保证只搜当前用户(或指定 agent)的数据(adapters/__init__.py:2153-2161),这是多租户隔离。apply_filters 则把查询里夹带的词/文件/日期过滤翻译成 ORM 条件——详见 03-rag-and-boundaries.md §2。

重要边界:没有 ANN 索引。 全仓 grep 不到 HnswIndex / IvfflatIndex(pgvector 的近似最近邻索引),Entry 也没在 embeddings 列上建向量索引。也就是说这个 order_by(distance) 是一次精确的顺序扫描——每次查询都会算"该用户全部块"的距离。个人规模(几千到几万块)完全够用,但它不是为百万级向量的亚线性检索设计的。(inferred:结论由"无向量索引迁移"这一事实推出。)

4. 重排:交叉编码器精排​

召回回来的 10 个候选(此时按双编码器距离排序),若开启重排,交给交叉编码器再打一遍分。组织逻辑在 rerank_and_sort_results(text_search.py:192):

# 真实源码风 —— text_search.py:195-204(节选)
# 三个条件任一满足才重排:显式要求 r=True、或配了推理服务器、且候选 > 1 条
rank_results = (rank_results or state.cross_encoder_model[name].inference_server_enabled()) and len(hits) > 1
if rank_results:
    hits = cross_encoder_score(query, hits, name)   # 用交叉编码器打分
hits = sort_results(rank_results=rank_results, hits=hits)

交叉编码器模型默认 mixedbread-ai/mxbai-rerank-xsmall-v1(SearchModelConfig.cross_encoder,models/__init__.py:566)。打分在 CrossEncoderModel.predict(embeddings.py:135):把 [查询, 候选compiled] 成对喂模型,过 Sigmoid 得 0-1 的相关分(embeddings.py:144-145)。

一个巧妙的坐标统一: 交叉编码器输出的是"越大越相关"的分,而双编码器用的是"越小越相关"的距离。为了让两套分能一起排序,cross_encoder_score 把交叉分转成距离:

# 真实源码 —— text_search.py:246
hits[idx]["cross_score"] = 1 - cross_scores[idx]   # 相关分 → 距离(越小越好)

然后排序时先按双编码器距离排,再按交叉分距离排(sort_results,text_search.py:251-257,两次 sort 利用 Python 排序稳定性,后一次为主序):

# 真实源码 —— text_search.py:254-256
hits.sort(key=lambda x: x["score"])           # 先按 bi-encoder 距离
if rank_results:
    hits.sort(key=lambda x: x["cross_score"])  # 再按 cross-encoder 距离(主序)

容错细节: 如果交叉编码器走的是远程推理端点且 HTTP 失败,cross_encoder_score 会把所有分置 0 而不是让整个查询崩掉(text_search.py:240-242)——检索降级为"只有双编码器"的结果。

5. 组装与去重​

数据库返回的是 Django 对象,collate_results(text_search.py:141)把它们转成对外的 SearchResponse,并顺带去重:

# 真实源码风 —— text_search.py:144-150(节选)
for hit in hits:
    if dedupe and (hit.hashed_value in hit_hashes or hit.corpus_id in hit_ids):
        continue          # 同指纹 或 同源条目(corpus_id)只留一个
    hit_hashes.add(hit.hashed_value); hit_ids.add(hit.corpus_id)
    yield SearchResponse.model_validate({... "score": hit.distance ...})

两级去重键: hashed_value(内容完全相同的块)和 corpus_id(切自同一原条目的兄弟块)。后者尤其有用——避免同一篇笔记的多个相邻块霸占前几名。

还有一个更靠上层的去重 deduplicated_search_responses(text_search.py:167),按 compiled 文本去重,用在"多个子查询各自检索后合并"的 RAG 场景(见 03-rag-and-boundaries.md)。

上面这些零件由 execute_search(src/khoj/routers/helpers.py:1485)编排成端到端流程。挑几个整体设计点:

  1. 查询缓存: 用 "查询-n-类型-r-距离-dedupe" 做 key,命中直接返回(helpers.py:1517-1520),重复查询零成本。
  2. 先剥过滤器再编码: 把 dt>/+词/file: 从查询里 defilter 掉,只对"干净查询"编码(helpers.py:1523-1531)——过滤语法不该污染语义向量。
  3. 多类型并行: 若搜 All,对各内容类型的查询用 asyncio.gather 并行跑(helpers.py:1533-1558)。
  4. 合并后统一重排取前 n: 各类型结果 collate 后合到一起,rerank_and_sort_results(...)[:results_count] 统一精排截断(helpers.py:1559-1566)。
execute_search 全景

 q ─▶ 查缓存(命中即返回)
   ─▶ defilter 剥离 dt/+词/file 过滤语法
   ─▶ embed_query(只编码干净查询)
   ─▶ 各内容类型并行 text_search.query ──▶ 各自 top-10
   ─▶ collate_results(转换 + 两级去重)
   ─▶ rerank_and_sort_results(交叉编码器精排)
   ─▶ 取前 n ─▶ 写缓存 ─▶ 返回

小结: 出库 = 编码查询 → pgvector CosineDistance 精确扫描召回 top-10(按 0.18 阈值剪枝)→ 交叉编码器精排(把相关分转成 1 - score 距离统一排序)→ 两级去重 → 取前 n。最值得带走的是两段式召回-重排,以及**用 1 - score 把两套异号分数统一到"越小越好"**这个小技巧。

下一章看这套检索怎么接进聊天,变成真正的 RAG。