> ## Content Index
> Fetch the complete content index at: https://www.hubwiz.com/blog/llms.txt
> Use this file to discover other available public pages before exploring further.

# 构建实时语义代码搜索
- URL: https://www.hubwiz.com/blog/build-a-real-time-semantic-code-search/
- Published: 2026-07-21T03:28:49.000Z
- Updated: 2026-10-16T11:31:47.000Z
- Description: 为你的代码库索引语义搜索，支持增量更新——只重新处理更改的内容，由原生Tree-sitter解析驱动。
- Author: admin
- Tags: APPLICATION, TREE-SITTER, COCOINDEX, #Import 2026-10-14 14:41

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

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

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

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

在本文中，我们将构建一个实用的[CocoIndex](https://github.com/cocoindex-io/cocoindex?ref=hubwiz.com)管道，它：

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

[完整源代码可在GitHub上找到](https://github.com/cocoindex-io/cocoindex/tree/main/examples/code%5Fembedding?ref=hubwiz.com)。

## 1、设置

- 安装Postgres，按照[安装指南](https://cocoindex.io/docs/getting%5Fstarted/installation?ref=hubwiz.com#-install-postgres)操作。
- 安装CocoIndex

```bash
pip install -U cocoindex

```

## 2、流程定义

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

### 2.1 将代码库添加为源

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

```python
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识别函数定义、类声明和模块级语句。这可以防止块在函数中间尴尬地分割。

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

```python
@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`参数。

```python
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](http://twitter.com/cocoindex?ref=hubwiz.com).transform\_flow()装饰器确保索引和查询时间使用相同的嵌入模型。

```python
@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`收集器中。

```python
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自动处理架构创建和索引管理。

```python
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、查询索引

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

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

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

```python
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、运行管道

开始很简单：

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

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

## 6、支持的语言

[CocoIndex](https://github.com/cocoindex-io/cocoindex?ref=hubwiz.com)的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](https://freedium-mirror.cfd/https://medium.com/towards-artificial-intelligence/building-real-time-semantic-code-search-with-tree-sitter-and-vector-embeddings-b9b1fc0a94f3?ref=hubwiz.com)

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

> **这是「RAG / 数据 / 文档智能」专题的一部分。**完整地图见：[RAG / 数据 / 文档智能完全指南](https://www.hubwiz.com/blog/rag-data-stack/)

**相关阅读**

- [向量数据库原理：索引与 ANN 近似检索](https://www.hubwiz.com/blog/vector-indexing-and-ann/)
- [RAG 是什么：检索增强生成一次讲清](https://www.hubwiz.com/blog/rag-2026/)