构建实时语义代码搜索

每个开发人员都经历过这种挫折：你知道仓库中存在一些代码可以完全满足你的需求，但你不记得在哪里。Grep适用于精确匹配，但如果你想找到与你所寻找的代码在语义上相似的代码呢？

想象一下能够像这样查询你的代码库：

"查找所有处理HTTP认证的函数"
"显示与此错误处理模式相似的代码"
"数据库连接池逻辑在哪里？"

这就是语义代码搜索大放异彩的地方。通过将代码转换为向量嵌入并存储在数据库中，你可以解锁传统文本搜索永远无法实现的强大基于相似性的查询。

在本文中，我们将构建一个实用的CocoIndex管道，它：

从本地文件系统读取代码文件
使用Tree-sitter解析将代码分割成语义块
为每个块生成向量嵌入
使用pgvector将所有内容存储在Postgres中以进行相似性搜索
仅在源文件更改时自动更新

完整源代码可在GitHub上找到。

1、设置

安装Postgres，按照安装指南操作。
安装CocoIndex

pip install -U cocoindex

2、流程定义

让我们分解这个CocoIndex管道的关键组件。

2.1 将代码库添加为源

我们使用LocalFile源来摄取代码文件。管道包含常见代码扩展名（.py、.rs、.toml、.md、.mdx）的文件模式，并排除node_modules和target等目录。

import os

@cocoindex.flow_def(name="CodeEmbedding")
def code_embedding_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope):
    data_scope["files"] = flow_builder.add_source(
        cocoindex.sources.LocalFile(path=os.path.join('..', '..'),
                                    included_patterns=["*.py", "*.rs", "*.toml", "*.md", "*.mdx"],
                                    excluded_patterns=[".*", "target", "**/node_modules"]))
    code_embeddings = data_scope.add_collector()

flow_builder.add_source将创建一个包含子字段（filename、content）的表。

2.2 将代码分割成语义块

与简单的文本分割不同，CocoIndex使用Tree-sitter进行语言感知分块。Tree-sitter理解数十种编程语言的语法，确保块尊重函数边界、类定义和其他语义单元。

例如，当处理Python代码时，Tree-sitter识别函数定义、类声明和模块级语句。这可以防止块在函数中间尴尬地分割。

让我们定义一个函数，在处理每个文件时提取文件名的扩展名。

@cocoindex.op.function()
def extract_extension(filename: str) -> str:
    """Extract the extension of a filename."""
    return os.path.splitext(filename)[1]

我们使用SplitRecursively函数将文件分割成块。SplitRecursively是CocoIndex构建块，与Tree-sitter原生集成。如果处理代码，需要将语言传递给language参数。

with data_scope["files"].row() as file:
    # Extract the extension of the filename.
    file["extension"] = file["filename"].transform(extract_extension)
    file["chunks"] = file["content"].transform(
          cocoindex.functions.SplitRecursively(),
          language=file["extension"], chunk_size=1000, chunk_overlap=300)

2.3 生成嵌入

我们使用SentenceTransformers（特别是all-MiniLM-L6-v2）将每个代码块转换为384维向量。@cocoindex.transform_flow()装饰器确保索引和查询时间使用相同的嵌入模型。

@cocoindex.transform_flow()
def code_to_embedding(text: cocoindex.DataSlice[str]) -> cocoindex.DataSlice[list[float]]:
    return text.transform(
        cocoindex.functions.SentenceTransformerEmbed(
            model="sentence-transformers/all-MiniLM-L6-v2"))

然后对于每个块，我们将使用code_to_embedding函数对其进行嵌入，并将嵌入收集到code_embeddings收集器中。

with data_scope["files"].row() as file:
    with file["chunks"].row() as chunk:
        chunk["embedding"] = chunk["text"].call(code_to_embedding)
        code_embeddings.collect(filename=file["filename"], location=chunk["location"],
                                code=chunk["text"], embedding=chunk["embedding"])

2.4 使用pgvector存储在Postgres中

嵌入使用pgvector扩展存储在Postgres中，以进行高效的相似性搜索。CocoIndex自动处理架构创建和索引管理。

code_embeddings.export(
    "code_embeddings",
    cocoindex.storages.Postgres(),
    primary_key_fields=["filename", "location"],
    vector_indexes=[cocoindex.VectorIndex("embedding", cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)])

3、增量处理的力量

这就是CocoIndex真正闪耀的地方。传统索引方法在每次更新时都会重新处理整个代码库。对于拥有数千个文件的大型代码库，这变得极其昂贵。

CocoIndex的增量处理意味着：

只有更改的文件触发重新处理
未更改的文件保留其缓存的嵌入
数据库接收最小的变更
更新在接近实时的情况下发生

对于每天代码流失率为1%的企业，每次更新周期只有1%的文件触及嵌入模型。其余99%的文件永远不会触及昂贵的计算路径。

4、查询索引

一旦索引构建完成，查询就很简单了。搜索函数：

使用相同的模型将你的自然语言查询转换为嵌入
对存储的向量执行余弦相似性搜索
返回具有相似性分数的最相关代码块

查询重用索引流程中定义的相同嵌入计算，确保索引和检索之间的一致性。

def search(pool: ConnectionPool, query: str, top_k: int = 5):
    # Get the table name, for the export target in the code_embedding_flow above.
    table_name = cocoindex.utils.get_target_storage_default_name(code_embedding_flow, "code_embeddings")
    # Evaluate the transform flow defined above with the input query, to get the embedding.
    query_vector = code_to_embedding.eval(query)
    # Run the query and get the results.
    with pool.connection() as conn:
        with conn.cursor() as cur:
            cur.execute(f"""
                SELECT filename, code, embedding <=> %s::vector AS distance
                FROM {table_name} ORDER BY distance LIMIT %s
            """, (query_vector, top_k))
            return [
                {"filename": row[0], "code": row[1], "score": 1.0 - row[2]}
                for row in cur.fetchall()
            ]

5、运行管道

开始很简单：

安装依赖项：pip install -e .
更新索引：cocoindex update main
测试查询：python main.py

启动查询界面后，你可以输入自然语言查询，如"spec"或"flow definition"，并接收包含文件名、代码片段和相似性分数的排名结果。

6、支持的语言

CocoIndex的SplitRecursively函数对所有主要编程语言都有原生Tree-sitter支持，包括Python、Rust、JavaScript、TypeScript、Go、Java、C、C++、Ruby等。语言自动从文件扩展名检测，因此不需要额外配置。

7、扩展到代码之外

这种模式远远超出了代码仓库：

文档搜索——使用语义搜索功能索引技术文档
基础设施即代码——使SRE团队能够搜索Terraform、Kubernetes配置和部署脚本
配置文件分析——跟踪配置漂移并跨环境搜索特定模式
多仓库索引——在组织中的多个仓库中构建统一搜索

8、结束语

构建语义代码搜索不一定很复杂。借助CocoIndex的增量处理、Tree-sitter集成和向量存储，你可以在不到100行Python代码中创建一个生产就绪的代码搜索系统。

关键要点：

Tree-sitter提供尊重代码语义的语言感知分块
增量处理大大降低了大型代码库的计算成本
向量嵌入实现了传统搜索无法匹配的强大基于相似性的检索
相同的管道模式适用于任何文本语料库，而不仅仅是代码

无论你是在为AI编码代理构建上下文、驱动内部代码搜索工具，还是启用自动化代码审查，这种架构都提供了一个随代码库扩展的坚实基础。

原文链接: Building Real-Time Semantic Code Search With Tree-sitter and Vector Embeddings

汇智网翻译整理，转载请标明出处