跳到主要内容

代码知识图谱(CKG):tree-sitter 建图 + SQLite 查询 + 快照缓存

30 秒导读: 大模型改代码时老要问"某某函数长啥样、定义在哪"。CKG(Code Knowledge Graph,代码知识图谱)预先把整个代码库的函数和类用 tree-sitter 解析成 AST、抽成一张 SQLite 表;agent 给一个符号名(search_function foo),就能秒回它的文件、行号、函数体——不用 grep 全库,也不用把整份代码塞进上下文。

本章讲 Trae Agent 里工程含量最高的一支工具。它是 02-tools.md 里"内置工具"的一员(注册名 ckg),但复杂度远超其它工具,值得单开一章。

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

一句话定义: CKG 是一个"按函数名/类名查代码"的本地索引——把代码库里所有函数和类抽出来,存进 SQLite,给 agent 当"代码目录"用。

它解决谁的什么问题? 想象 agent 要修一个几万行的项目。它读到某处调用了 parse_config,想看这个函数怎么写的。没有 CKG,它只能:

  • 用 grep 全库找 parse_config —— 命中一堆无关的字符串、注释、调用点,还得自己数行号;
  • 或者把相关文件整份读进上下文 —— 又贵又占窗口。

有了 CKG,agent 只需一条命令,直接拿到"这个符号在哪、第几行到第几行、函数体是什么"。

它能做什么(三条命令):

命令查什么
search_function顶层函数(不属于任何类)
search_class类(附带字段、方法签名列表)
search_class_method类里的方法

用起来什么样。 工具是给 LLM 调的,一次调用长这样(参数见 get_parameters,ckg_tool.py:51):

command    = "search_function"
path       = "/repo"
identifier = "parse_config"
print_body = true          # 默认 true,连函数体一起返回

返回(见 _search_function,ckg_tool.py:135):

Found 1 functions named parse_config:
1. /repo/src/config.py:12-30
def parse_config(path):
    ...

一句话直觉/类比: 把 CKG 当成一本自动生成的"代码电话簿"——你报一个名字(符号),它翻到那一页(文件:行号)并把内容念给你听。而"电话簿"本身是用 tree-sitter 读懂每份源码的语法结构后编出来的。

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

CKG 分三个阶段:建图(解析源码 → 落库)、查图(符号 → 记录)、缓存(用快照哈希决定复用还是重建)。

怎么读下图: 从上到下是一次工具调用的生命周期;左侧是"第一次遇到这个代码库"要走的建图路径,右侧是"命中缓存"的快路径。

        agent 调用 ckg 工具 (path, identifier, command)
                        │
                        ▼
                 CKGTool.execute            ← ckg_tool.py:80
                        │
          按 path 找已建好的 CKGDatabase?
                 ┌──────┴───────┐
              有 │              │ 没有 → new CKGDatabase(path)
                 │              ▼
                 │      get_folder_snapshot_hash(path)   ← ckg_database.py:97
                 │              │  (git 状态 / 文件元数据 → 一个哈希串)
                 │        ┌─────┴─────┐
                 │  哈希已有 .db      哈希是新的
                 │   复用 SQLite      _construct_ckg() 重建   ← ckg_database.py:534
                 │        │           (tree-sitter 遍历所有源码 → 落库)
                 └────────┴─────┬─────┘
                                ▼
                   query_function / query_class     ← ckg_database.py:648 / 695
                   (SELECT ... WHERE name = ?)
                                ▼
                   拼输出 + MAX_RESPONSE_LEN 截断    ← ckg_tool.py:155

部件一句话职责:

部件干什么在哪
CKGTool对外三命令、参数校验、缓存 DB 实例、拼输出+截断tools/ckg_tool.py:14
CKGDatabase建图与查图的核心;持有一个 SQLite 连接tools/ckg/ckg_database.py:148
FunctionEntry/ClassEntry承载一条函数/类记录的 dataclasstools/ckg/base.py:9 :24
extension_to_language文件后缀 → tree-sitter 语言名tools/ckg/base.py:40
get_folder_snapshot_hash给代码库算一个快照哈希,判断要不要重建tools/ckg/ckg_database.py:97
clear_older_ckg清理一周前的旧 .db,agent 启动时调一次tools/ckg/ckg_database.py:107

主线走一遍(高层): agent 传入 (path, identifier, command) → CKGTool 若没为该 path 建过库,就 new CKGDatabase → 后者用快照哈希查缓存,命中就复用 .db,否则 tree-sitter 遍历所有源码建表 → 三命令各自走 query_function/query_class 做一次 SELECT ... WHERE name = ? → 结果拼成文本,超过 16000 字符就截断。

3. 核心原理(逐个机制,由浅入深)​

3.1 对外的三命令与截断(CKGTool)​

它要解决的小问题: 把"查函数/查类/查方法"这三种意图,变成三条简单命令,并保证再多的命中也不会撑爆上下文。

入口 execute。 先做参数校验(command/path/identifier 缺一不可,path 必须是存在的目录,ckg_tool.py:82-112),然后按 path 缓存 CKGDatabase 实例——同一个代码库只建一次库,后续调用直接复用:

ckg_database = self._ckg_databases.get(codebase_path)
if ckg_database is None:
    ckg_database = CKGDatabase(codebase_path)   # 第一次才建
    self._ckg_databases[codebase_path] = ckg_database

（ckg_tool.py:114-117,_ckg_databases 是个 dict[Path, CKGDatabase])

三命令 → 两个查询函数的映射(ckg_tool.py:119-133):

命令底层调用关键区别
search_functionquery_function(id, entry_type="function")只要 parent_class is None 的
search_class_methodquery_function(id, entry_type="class_method")只要 parent_class is not None 的
search_classquery_class(id)查 classes 表

注意前两个共用同一张 functions 表和同一个查询函数,靠 parent_class 字段区分"顶层函数"还是"类方法"(见 §3.4)。

输出与 print_body。 每条命中打印 序号. 文件:起行-止行,print_body=True(默认)时附上完整函数体;search_class 还会额外打印 Fields: 和 Methods: 两块(ckg_tool.py:179-184)。search_class_method 的行里多一句 within class <parent_class>(ckg_tool.py:211)。

MAX_RESPONSE_LEN 截断(= 16000,来自 tools/run.py:18)。 三个 _search_* 都在循环里累加输出,一旦超长就砍掉尾巴并附提示,还诚实告诉 agent 剩几条没显示:

if len(output) > MAX_RESPONSE_LEN:
    output = (
        output[:MAX_RESPONSE_LEN]
        + f"\n<response clipped> {len(entries) - index + 1} more entries not shown"
    )
    break

（ckg_tool.py:155-160,三处结构相同）

3.2 建图:tree-sitter 遍历 AST 抽符号(_construct_ckg)​

它要解决的小问题: 怎么从一份源码文本里,可靠地抽出"这里有个函数、名字叫 X、从第 12 行到第 30 行、body 是这段"——而且要跨 Python/Java/C/C++/TS/JS 六种语言。

思路/直觉: 用正则或字符串匹配去认函数,遇到嵌套、装饰器、多行签名就会翻车。正确做法是把源码解析成 AST(抽象语法树,源码的结构化树形表示),再在树上找 function_definition、class_definition 这类节点——语法层面认符号,天然精确。Trae 用 tree-sitter(一个增量式多语言解析器)干这活。

建图主循环 _construct_ckg。 遍历代码库所有文件,跳过隐藏文件/隐藏目录,按后缀查 extension_to_language 决定语言(不认识的后缀直接跳过),懒加载对应的 parser,解析成树,再按语言分派到对应的 _recursive_visit_*:

for file in self._codebase_path.glob("**/*"):
    ...
    extension = file.suffix
    if extension not in extension_to_language:   # 只认已知语言
        continue
    language = extension_to_language[extension]
    language_parser = language_to_parser.get(language) or get_parser(language)  # 懒加载
    tree = language_parser.parse(file.read_bytes())
    match language:
        case "python": self._recursive_visit_python(tree.root_node, ...)
        case "java":   self._recursive_visit_java(...)
        ...

（简化自 ckg_database.py:534-574;get_parser 来自 tree_sitter_languages,ckg_database.py:13）

支持的语言(extension_to_language,base.py:40):

语言后缀
python.py
java.java
cpp.cpp .hpp .c++ .cxx .cc
c.c .h
typescript.ts .tsx
javascript.js .jsx

以 Python 为例看抽取逻辑 _recursive_visit_python。 它在 AST 上递归下降,认两种节点:

  • 碰到 function_definition → 取 name 字段建 FunctionEntry,行号用 start_point[0] + 1(tree-sitter 行号从 0 起,这里 +1 转成人类行号),body 用节点整段文本;
  • 碰到 class_definition → 建 ClassEntry,并顺手遍历类体,把每个方法的"名字+参数+返回类型"拼成一行 methods 摘要(ckg_database.py:250-271),连装饰过的方法(decorated_definition)也认。
if root_node.type == "function_definition":
    function_name_node = root_node.child_by_field_name("name")
    if function_name_node:
        function_entry = FunctionEntry(
            name=function_name_node.text.decode(),
            file_path=file_path,
            body=root_node.text.decode(),
            start_line=root_node.start_point[0] + 1,   # 0-based → 1-based
            end_line=root_node.end_point[0] + 1,
        )
        ...
        self._insert_entry(function_entry)
        parent_function = function_entry   # 供更深层区分嵌套函数 vs 方法

（ckg_database.py:213-237）

父子归属的判定(精巧处)。 递归时携带 parent_class/parent_function 两个上下文。当一个函数同时有父函数和父类时,它用行号区间包含关系判断这函数到底是"类里的方法"还是"函数里的嵌套函数":

if parent_function and parent_class:
    if (parent_function.start_line >= parent_class.start_line
            and parent_function.end_line <= parent_class.end_line):
        function_entry.parent_function = parent_function.name   # 嵌套在函数里
    else:
        function_entry.parent_class = parent_class.name         # 是类的方法

（ckg_database.py:223-235）这个 parent_class 字段后面决定了它被 search_function 还是 search_class_method 查到(见 §3.4)。

其它语言的 _recursive_visit_* 结构类似,细节按语言语法调整:Java 认 class_declaration/method_declaration(ckg_database.py:279),C/C++ 从 function_definition 的 declarator 里再挖一层拿函数名(ckg_database.py:379-393),C 不建类只建函数(ckg_database.py:399)。

3.3 落库:两张 SQLite 表​

它要解决的小问题: 抽出来的符号得存在一个能按名字快速查、又能跨进程复用的地方。SQLite 单文件、零配置、自带索引查询,正合适。

建表 _construct_ckg 前置(ckg_database.py:122-145)。 两张表,SQL_LIST 里定义:

表关键列说明
functionsname, file_path, body, start_line, end_line, parent_function, parent_class顶层函数与类方法共用此表,靠 parent_class 区分
classesname, file_path, body, fields, methods, start_line, end_linefields/methods 是预先拼好的文本摘要

落库 _insert_entry(ckg_database.py:576) 按 dataclass 类型分派到 _insert_function 或 _insert_class,各做一条参数化 INSERT 并 commit。用参数化占位符 ?,避免符号名里的引号破坏 SQL。

3.4 查图:一条 SELECT + parent_class 分流(query_function)​

它要解决的小问题: 一个符号名,可能既是某处的顶层函数、又是别处的类方法。工具要能只返回 agent 想要的那一类。

巧妙处:两命令共用一次查询,靠 entry_type 在 Python 侧过滤。 query_function 先无差别捞出所有同名记录,再按 parent_class(记录第 7 列,record[6])是否为空分流:

records = ... SELECT ... FROM functions WHERE name = ?
for record in records:
    match entry_type:
        case "function":
            if record[6] is None:        # parent_class 为空 = 顶层函数
                function_entries.append(FunctionEntry(...))
        case "class_method":
            if record[6] is not None:    # parent_class 非空 = 类方法
                function_entries.append(FunctionEntry(...))

（ckg_database.py:660-693）

query_class 更直接:SELECT ... FROM classes WHERE name = ?,每条记录还原成一个 ClassEntry(ckg_database.py:695-722)。

关键细节/坑: 查询用的是精确相等 name = ?,不是 LIKE 或模糊匹配——符号名必须一字不差。这也是工具描述里那句"CKG 未必找得全"(ckg_tool.py:47)的一部分原因:名字对不上就是零命中。

3.5 快照与缓存:什么时候该重建这张图​

它要解决的小问题: 建一次图要遍历+解析整个代码库,不便宜。同一个库反复调用工具,不该每次重建;但代码一旦改了,又必须重建,否则查到的是过时的行号和 body。

核心思路:给代码库算一个"快照哈希"当缓存键。 哈希不变 → 复用 .db;哈希变了 → 删旧建新。逻辑在 CKGDatabase.__init__(ckg_database.py:149-196):它从 storage_info.json 读出该 path 上次的哈希,和当前哈希比;相等就复用,不等就删旧库、建新库、更新 storage_info.json,并调 _construct_ckg 重建。

哈希怎么算 get_folder_snapshot_hash(ckg_database.py:97) 分两种情况:

是 git 仓库?  ──是──▶  get_git_status_hash   (基于 commit + 未提交改动)
      │
      否
      ▼
get_file_metadata_hash   (基于所有文件的 name+mtime+size)

git 仓库(get_git_status_hash,ckg_database.py:51)区分干净/脏:

工作区状态哈希形如含义
无未提交改动git-clean-<HEAD>只用 commit 哈希,同一 commit 稳定复用
有未提交改动git-dirty-<HEAD>-<md5前8位>把 git status --porcelain 输出算 md5 掺进去,改动一变哈希就变

git 命令失败时回退到文件元数据哈希(ckg_database.py:78-80)。

非 git(get_file_metadata_hash,ckg_database.py:83): 遍历所有非隐藏文件,把每个文件的文件名 + 修改时间 + 大小喂进一个 md5,得到 metadata-<md5>。任一文件被改动(mtime/size 变化)哈希就变,触发重建。

存哪: 数据库落在 LOCAL_STORAGE_PATH / "ckg",即 ~/.trae-agent/ckg/(LOCAL_STORAGE_PATH 见 utils/constants.py:6),每个快照一个 <hash>.db 文件。

过期清理 clear_older_ckg(ckg_database.py:107): 遍历 ckg/ 下的 .db,删掉 mtime 早于"一周前"的(CKG_DATABASE_EXPIRY_TIME = 60*60*24*7,ckg_database.py:20)。这个清理在 agent 每次启动时触发一次——BaseAgent.__init__ 末尾直接调用它(agent/base_agent.py:81,见 01-agent-loop.md),防止 ~/.trae-agent/ckg/ 无限膨胀。

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

  • AST 而非正则抽符号。 用 tree-sitter 在语法树上认 function_definition/class_definition,天然处理多行签名、装饰器、嵌套——比 grep/正则鲁棒得多(_recursive_visit_python,ckg_database.py:205)。
  • 顶层函数与类方法共表、查询时按 parent_class 分流。 省一张表,两条命令复用同一次 SELECT(ckg_database.py:660-693)。
  • 行号区间判嵌套归属。 用 parent_function 是否落在 parent_class 行区间内,区分"方法"与"函数内嵌套函数"(ckg_database.py:223-231)——纯靠位置关系,不需额外的语义分析。
  • 快照哈希做缓存键,还区分 git 干净/脏。 clean 只认 commit(跨会话稳定命中缓存),dirty 把未提交 diff 的 md5 掺进键(改一行就失效重建),兼顾复用与新鲜(get_git_status_hash,ckg_database.py:51)。
  • 懒加载 parser。 只在真遇到某语言文件时才 get_parser,纯 Python 项目不会白白加载 Java/C++ 的语法(ckg_database.py:552-555)。

5. 边界与局限(诚实)​

源码顶部的 docstring(ckg_database.py:23-28)自己列了三条 Known issues,加上阅读代码可见的其它限制:

局限说明出处
子目录重复建图已索引过的代码库,若又对它的某个子目录建 CKG,会为该子目录重新建一遍图(不复用父目录成果)Known issue 1,ckg_database.py:25
不做增量代码一改,快照哈希变,就整库重建,不是只重建被改的文件Known issue 2,ckg_database.py:26;update() 也是全量重跑 ckg_database.py:201
JS/TS AST 不完整匿名函数、箭头函数等不被解析——JS/TS 项目里大量此类符号查不到Known issue 3,ckg_database.py:27;_recursive_visit_javascript 只认 class_declaration/method_definition,ckg_database.py:487-516
精确匹配、无模糊查询用 WHERE name = ?,符号名不一字不差就零命中ckg_database.py:661
同名符号无消歧全库同名函数会全部返回,靠输出里的 file:line 让 agent 自行辨别ckg_tool.py:145-153
只索引已知后缀六种语言之外(Go/Rust/Ruby…)的文件被静默跳过extension_to_language,base.py:40;ckg_database.py:548

工具描述里也直白写了 "The CKG is not completely accurate, and may not be able to find all functions or classes"(ckg_tool.py:47)——把不确定性提前告诉调用它的 LLM。

6. 横向对比(同库其它章)​

CKG 是 Trae 工具体系里的一员,但机制自成一套:

  • 它作为工具怎么注册、怎么被 executor 调度,见 02-tools.md;其 MAX_RESPONSE_LEN 截断常量与 bash/编辑工具共享(tools/run.py)。
  • clear_older_ckg 的调用时机嵌在主循环启动处,见 01-agent-loop.md。
  • 存储路径 LOCAL_STORAGE_PATH 同时也是配置/轨迹相关目录,配置分层见 06-observability.md。

7. 代码地图(导航索引)​

主题文件符号
工具对外三命令 + 参数 + 截断trae_agent/tools/ckg_tool.pyCKGTool、CKGToolCommands、_search_function、_search_class、_search_class_method
建图主循环(遍历+分派)trae_agent/tools/ckg/ckg_database.pyCKGDatabase._construct_ckg
Python AST 抽取(父子归属判定)trae_agent/tools/ckg/ckg_database.py_recursive_visit_python
其它语言抽取trae_agent/tools/ckg/ckg_database.py_recursive_visit_java、_recursive_visit_cpp、_recursive_visit_c、_recursive_visit_typescript、_recursive_visit_javascript
落库trae_agent/tools/ckg/ckg_database.py_insert_entry、_insert_function、_insert_class、SQL_LIST
查图(按 parent_class 分流)trae_agent/tools/ckg/ckg_database.pyquery_function、query_class
快照哈希 / 缓存决策trae_agent/tools/ckg/ckg_database.pyget_folder_snapshot_hash、get_git_status_hash、get_file_metadata_hash、CKGDatabase.__init__
过期清理trae_agent/tools/ckg/ckg_database.pyclear_older_ckg、CKG_DATABASE_EXPIRY_TIME
数据模型 / 语言映射trae_agent/tools/ckg/base.pyFunctionEntry、ClassEntry、extension_to_language
存储根路径trae_agent/utils/constants.pyLOCAL_STORAGE_PATH
启动时触发清理trae_agent/agent/base_agent.pyBaseAgent.__init__(第 81 行 clear_older_ckg())