这篇能做出什么
做完这篇教程，你会得到一个可以本地跑的代码检索流水线：输入一句自然语言，比如"哪个函数负责把用户上传的文件转成缩略图"，它能在你的仓库里找出对应的函数块，而不只是匹配同名关键字。
具体产出四样东西：
1. 一个按符号切块的索引器，把 Python 仓库切成函数级、类级的知识块；
2. 一个代码嵌入 + 向量检索的召回层，附带 BM25 关键词召回做混合检索；
3. 一个重排层，对召回的候选做精排；
4. 一个Recall@k 评测脚本，能回答"换一个嵌入模型，召回率到底涨了没有"这种问题。
整套流程跑通后，把最后一步接上大模型生成答案，它就是一个标准的代码库 RAG 问答系统。检索层质量决定了 RAG 的上限，所以这篇把力气都花在检索和评测上。
前置条件清单
- Python 3.10 以上（具体版本以官方文档为准），熟悉函数和 dataclass 即可；
- 一个待检索的代码仓库，建议先用 500～5000 个文件的仓库练手，太大不利于快速迭代；
- 安装依赖：
```bash
pip install sentence-transformers transformers rank-bm25 numpy tqdm
```
如果仓库规模上万文件，再加一个向量索引库（faiss-cpu、hnswlib、chromadb 都可以，按仓库规模选）。
- 硬件：有 GPU 会快不少，纯 CPU 也能跑，只是建索引阶段要等一会儿；
- 评测集：50～200 条查询就够了，后面会讲怎么半自动地造出来。
目录结构建议这样：
```text
semsearch/
chunker.py
build_index.py
search.py
eval.py
data/
chunks.jsonl
emb.npy
eval_queries.jsonl
```
分步骤
第 1 步：按语法结构分块，而不是按字符数
固定切 512 个字符是很多人的默认做法，用在代码上问题明显：一个函数会被从中间劈开，签名和函数体分到两个块；一个类的方法和它的字段定义分开，语义被打散。代码天然有结构，直接用 Python 自带的 ast 就能拿到边界。
```python
chunker.py
import ast
from dataclasses import dataclass
from pathlib import Path
@dataclass
class Chunk:
path: str
symbol: str # 例如 "ImagePipeline.thumbnail"
kind: str # function / class
start_line: int
end_line: int
text: str
header: str = "" # 注入的上下文，见第 2 步
def _qualified(parent: str, name: str) -> str:
return f"{parent}.{name}" if parent else name
def _walk(body, lines, path, parent=""):
out = []
for node in body:
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
out.append(Chunk(
path=path,
symbol=_qualified(parent, node.name),
kind="function",
start_line=node.lineno,
end_line=node.end_lineno,
text="".join(lines[node.lineno - 1: node.end_lineno]),
))
elif isinstance(node, ast.ClassDef):
类块只取"签名 + 文档字符串"，避免把整个类重复进索引
stop = node.end_lineno
for child in node.body:
if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
stop = child.lineno - 1
break
out.append(Chunk(
path=path,
symbol=_qualified(parent, node.name),
kind="class",
start_line=node.lineno,
end_line=stop,
text="".join(lines[node.lineno - 1: stop]),
))
out += _walk(node.body, lines, path, _qualified(parent, node.name))
return out
def chunk_file(path: Path) -> list[Chunk]:
src = path.read_text(encoding="utf-8", errors="ignore")
lines = src.splitlines(keepends=True)
try:
tree = ast.parse(src)
except SyntaxError:
return []
return _walk(tree.body, lines, str(path))
```
两个补充规则，代码里没写但实际要用：
- 超长函数二次切分：一个函数超过 1500 字符时，按顶层语句边界再切，每一块都拼上原始的函数签名行，这样切开的块仍然知道自己属于谁。
- 非 Python 文件：JavaScript、Go、Java 用 tree-sitter 的查询语句拿到同名节点即可，思路完全一致。
第 2 步：给每个块注入上下文头
嵌入模型只看块文本，而块文本里往往没有文件名、没有它属于哪个类。用户问"缩略图是在哪生成的"，一个内容正确但缺少上下文的函数块可能排不上去。
做法是给每个块拼一行人造头：
```python
def with_header(chunk: Chunk) -> str:
return (
f"# file: {chunk.path}\n"
f"# symbol: {chunk.symbol}\n"
f"# kind: {chunk.kind}\n\n"
f"{chunk.text}"
)
```
这行头会一起送进嵌入模型。检索命中后再把 header 剥掉，返回给用户的是原始代码，不污染展示。
第 3 步：选代码嵌入模型
选型时看这几个维度，不看宣传语：
- 训练语料是否包含代码。纯自然语言模型在标识符上表现一般，因为
parse_user_id这种词在通用语料里几乎不出现。 - 最大序列长度。函数级块通常几百个 token，够用；类块和跨文件块需要更长窗口。
- 是否需要查询前缀。一些模型要求查询侧加指令前缀、文档侧不加，配置错了召回率会明显掉。用之前翻一下模型卡。
- 向量维度。维度高检索精度通常好一点，内存和延迟也更高，权衡着来。
社区有专门的代码检索公开榜单（搜 code information retrieval benchmark 能找到），跑分以官方页面当前版本为准，不要照搬别人文章里的数字。
接入代码大致长这样：
```python
from sentence_transformers import SentenceTransformer
模型名称与是否需要 prompt_name 以官方文档为准
model = SentenceTransformer("your-code-embedding-model")
def embed_texts(texts, is_query=False, batch_size=64):
kwargs = {"normalize_embeddings": True, "batch_size": batch_size,
"show_progress_bar": True}
if is_query:
若模型卡要求查询侧加前缀，在这里指定
kwargs["prompt_name"] = "query"
pass
return model.encode(texts, **kwargs)
```
normalize_embeddings=True 很关键：向量归一化之后，内积就等于AI 词典：余弦相似度">余弦相似度，后面用 numpy 点积检索即可。
第 4 步：建向量索引
```python
build_index.py
import json
import numpy as np
from pathlib import Path
from chunker import chunk_file
from tqdm import tqdm
REPO = Path("your-repo")
SKIP_DIRS = {".git", "node_modules", "vendor", "dist", "build", "__pycache__"}
SKIP_SUFFIX = {".min.js", ".lock", ".pyc"}
chunks = []
for path in REPO.rglob("*.py"):
if any(part in SKIP_DIRS for part in path.parts):
continue
if path.suffix in SKIP_SUFFIX:
continue
chunks += chunk_file(path)
print(f"chunks: {len(chunks)}")
from build_index import embed_texts # 按你的工程结构调整导入
texts = [c.header or "" for c in chunks] # 这里填入第 2 步的 with_header 结果
embs = np.asarray(embed_texts(texts), dtype="float32")
np.save("data/emb.npy", embs)
with open("data/chunks.jsonl", "w", encoding="utf-8") as f:
for c, t in zip(chunks, texts):
f.write(json.dumps({
"path": c.path, "symbol": c.symbol, "kind": c.kind,
"start_line": c.start_line, "end_line": c.end_line,
"header": c.header, "text": c.text, "indexed_text": t,
}, ensure_ascii=False) + "\n")
```
检索侧就三行：
```python
def vector_search(query_vec, top_k=20):
scores = embs @ query_vec # 向量已归一化，内积即余弦
idx = np.argsort(-scores)[:top_k]
return [(int(i), float(scores[i])) for i in idx]
```
向量数超过十万时换成 FAISS 或 HNSW，接口基本一致。
第 5 步：混合检索，把关键词召回补上
向量检索对语义改写友好，但对精确标识符不敏感。有人直接搜 thumbnail_size，向量可能给出语义相近但名字不同的函数。BM25 恰好相反。两者用 RRF（倒数排名融合）合并，不需要调权重：
```python
search.py
import re
from rank_bm25 import BM25Okapi
def code_tokens(text: str) -> list[str]:
raw = re.findall(r"[A-Za-z_][A-Za-z0-9_]*|\d+", text)
out = []
for tok in raw:
parts = re.split(r"_|(?<=[a-z0-9])(?=[A-Z])", tok)
out += [p.lower() for p in parts if p]
return out
bm25 = BM25Okapi([code_tokens(c["indexed_text"]) for c in corpus])
def bm25_search(query: str, top_k=20):
scores = bm25.get_scores(code_tokens(query))
idx = sorted(range(len(scores)), key=lambda i: -scores[i])[:top_k]
return [int(i) for i in idx]
def rrf(rank_lists, k=60):
fused = {}
for lst in rank_lists:
for rank, doc_id in enumerate(lst):
fused[doc_id] = fused.get(doc_id, 0.0) + 1.0 / (k + rank + 1)
return sorted(fused, key=fused.get, reverse=True)
```
code_tokens 把 parse_user_id 拆成 parse、user、id，把 getUserName 拆成 get、user、name，这一步对代码检索的提升比换模型还明显。
第 6 步：接重排
召回阶段为了不漏，通常取 50～100 个候选；这么多结果直接给用户看太差，加一层 cross-encoder 精排，把查询和候选块拼在一起打分：
```python
from sentence_transformers import CrossEncoder
重排模型名称以官方文档为准
reranker = CrossEncoder("your-rerank-model")
def rerank(query, candidate_ids, corpus, top_n=10):
pairs = [(query, corpus[i]["indexed_text"]) for i in candidate_ids]
scores = reranker.predict(pairs, batch_size=32)
order = sorted(range(len(candidate_ids)), key=lambda i: -scores[i])
return [candidate_ids[i] for i in order[:top_n]]
```
注意重排模型的输入长度一般比嵌入模型短，超长的块会被截断，可以在送进去之前只保留签名 + 前若干行。
完整链路：查询 → 向量召回 50 + BM25 召回 50 → RRF 融合 → 重排取前 10。
第 7 步：用 Recall@k 做召回评测
没有评测就没法判断改动是否有效。造评测集有个省事的办法：拿仓库里的函数文档字符串首句当查询，把该函数对应的块标为正样本。
```python
make_eval.py
import ast, json
from pathlib import Path
rows = []
for path in Path("your-repo").rglob("*.py"):
try:
tree = ast.parse(path.read_text(encoding="utf-8", errors="ignore"))
except SyntaxError:
continue
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
doc = ast.get_docstring(node)
if doc and len(doc.split()) >= 6:
rows.append({
"query": doc.strip().splitlines()[0],
"gold_symbol": node.name,
"path": str(path),
})
with open("data/eval_queries.jsonl", "w", encoding="utf-8") as f:
for r in rows:
f.write(json.dumps(r, ensure_ascii=False) + "\n")
```
自动生成的查询偏"教科书式"，还要人工补 20～30 条真实口吻的提问，比如"图片上传后多久生成缩略图"、"哪里做了重试"。真实提问才是你要优化的分布。
评测脚本：
```python
eval.py
import json
import numpy as np
def recall_at_k(ranked_ids, gold_ids, k):
hit = len(set(ranked_ids[:k]) & set(gold_ids))
return hit / len(gold_ids)
def run_eval(pipeline, eval_rows, corpus, ks=(1, 5, 10, 20)):
totals = {k: 0.0 for k in ks}
for row in eval_rows:
gold_ids = {c["id"] for c in corpus
if c["symbol"] == row["gold_symbol"] and c["path"] == row["path"]}
if not gold_ids:
continue
ranked = pipeline(row["query"]) # 返回排好序的 chunk id 列表
for k in ks:
totals[k] += recall_at_k(ranked, gold_ids, k)
n = len(eval_rows)
return {k: round(v / n, 4) for k, v in totals.items()}
if __name__ == "__main__":
依次替换成不同 pipeline：纯向量 / 向量+BM25 / 向量+BM25+重排
print(run_eval(pipeline, eval_rows, corpus))
```
跑完会得到一张很像样的表：
```text
配置 Recall@1 Recall@5 Recall@10 Recall@20
纯向量 0.31 0.58 0.71 0.80
向量 + BM25 (RRF) 0.38 0.66 0.78 0.86
向量 + BM25 + 重排 0.52 0.79 0.86 0.88
```
上面是格式示例，具体数值取决于你的仓库和模型。重点看两件事：加了重排之后 Recall@1 有没有明显抬升（重排主要影响头部），Recall@20 是否已经接近上限（决定还要不要继续加召回策略）。
常见坑与排错
所有结果都很相似、得分都在 0.7 以上。 大概率忘了归一化，或者向量库里存的还是上一次的旧向量。删掉索引重建，并确认 normalize_embeddings=True。
换了模型，召回率反而掉了。 先检查查询侧前缀。很多模型要求查询加指令、文档不加，两边用错就会系统性变差。再看是不是只换了嵌入模型、没重建索引。
同一个函数出现好几条、挤占结果位。 类块和函数块文本重叠，或者二次切分后多个子块命中同一函数。检索后按 symbol 去重，每个符号只留分数最高的那条。
评测分数高得离谱。 检查评测集和索引有没有重叠泄漏：如果 gold 是"包含该函数的所有块"，而正样本恰好被切成了好几块，Recall 会被虚高。评测时把 gold 限定为唯一那个块，或者明确接受多块命中。
精确标识符搜不到。 说明 BM25 那一路没生效。检查分词函数是否真的把驼峰和蛇形命名拆开了，以及 BM25 语料是不是用了和向量索引同一批块。
重排后结果变差。 看候选长度。重排模型输入被截断，会丢掉块尾部的关键代码。把送进重排的文本限制在签名加前若干行，通常更稳。
文件改了但搜索结果没变。 索引是离线批处理的，需要重建或者做增量：用文件哈希记录已索引文件，只重算变更文件和删除文件的向量。
噪声文件污染索引。 生成代码、压缩后的 JS、依赖目录、测试 fixture 都要排除。这些东西会大量重复，把真正的业务代码挤下去。
下一步建议
1. 接上生成环节。把重排后的前 5 个块连同查询一起送进大模型，让模型引用文件路径和行号作答。返回的块里保留 path 和 start_line，方便直接跳到源码。
2. 扩充评测指标。Recall 只看"有没有捞到"，再加 MRR 或 nDCG 看排序质量，加人工相关性标注校验自动生成的评测集。
3. 把评测接进 CI。每次换模型、改分块策略、调参数都跑一遍 eval.py，对比基线。检索系统的改动很容易凭感觉误判。
4. 扩展到多语言仓库。用 tree-sitter 的查询语言统一切块，前端、后端、基础设施代码放进同一个索引。
5. 做查询改写。用户问"上传图片后缩略图什么时候出来"，可以先让大模型改写成几个候选查询（含可能的函数名片段），分别检索后合并，对模糊提问提升明显。