用代码图优化 AI Agent 的代码理解
通过构建代码依赖图解决 AI agent 盲目搜索问题,显著提升代码映射和重构精度。
通过构建代码依赖图解决 AI agent 盲目搜索问题,显著提升代码映射和重构精度。
我的 AI 代码 agent 曾经在每个会话中都要从头开始重新发现代码库的结构——grep 各处、打开文件"只是检查一下",有时候在重构期间还会漏掉某个调用点。我构建了一个小的代码图(函数和类作为节点,调用和导入作为边),并给了 agent 一些有针对性的查询,而不是原始的 grep。以下是什么有效、什么失效以及我会如何做不同的地方。
我大多数天都会针对一个中等规模的代码库(几百个文件、多个服务)运行 Claude Code。长期以来,每次会话都以相同的方式开始:agent 会 grep 一个符号、打开三四个文件"保险起见"、用稍微不同的模式再 grep 一次,最终建立起对事物如何连接的心理模型——然后在会话结束的那一刻就把这个模型抛弃。
这个重新发现的成本以两种方式表现出来:
浪费的 token 和时间。像"重命名这个函数并更新所有调用点"这样的任务在 agent 甚至开始编辑之前就会变成五六轮的 grep-读取。
漏掉调用点。Grep 是文本搜索,而不是结构性搜索。它会漏掉通过别名、重新导出和间接引用的调用。两次,一个"安全的"重命名被部署了但带有一个 grep 根本没有找到的损坏调用者,我只在 CI 中才发现它。
第二个问题才是真正推动我修复这个的原因。一个工具如果无法可靠地告诉你"谁调用这个函数",就不安全把重构工作交给 agent 无人监管。
让这个问题变得不可回避的具体事件是:我要求 agent 重命名一个在几个模块中使用的辅助函数。它 grep 了,找到了四个调用点,更新了所有四个,为它正在编辑的模块运行了本地测试,并报告了成功。20分钟后 CI 失败了——第五个调用点隐藏在一个模块的重新导出后面,grep 模式没有匹配到,因为导入使用了不同的本地别名。在 agent 拥有的工具给定的条件下,agent 的过程没有什么错误。Grep 只是不是"谁调用这个"这个问题的正确工具,没有多少更聪明的正则表达式可以修复这一类漏洞。
我也想坦率地说这不是一个新想法——IDE 二十年来就一直有"查找引用"功能,语言服务器早已为人类做这个了。有趣的部分不是这个概念,而是让它足够便宜和可靠,agent 默认会使用它而不是出于习惯而回退到 grep。
修复方案不是一个高级的静态分析平台——而是一个小的、无趣的索引,它可以很好地回答少数几个特定问题。
第 1 步:将源解析为结构索引,而不仅仅是文本。
我使用 tree-sitter 将每个文件解析成 AST,并提取出对导航很重要的部分:函数和类定义、它们的位置,以及它们内部的调用/导入表达式。简化后,提取看起来像这样:
from tree_sitter import Parser
from tree_sitter_languages import get_language
def extract_symbols(file_path: str, parser: Parser):
source = open(file_path, "rb").read()
tree = parser.parse(source)
symbols = []
def walk(node):
if node.type in ("function_definition", "class_definition"):
name_node = node.child_by_field_name("name")
symbols.append({
"name": name_node.text.decode(),
"kind": node.type,
"file": file_path,
"start_line": node.start_point[0] + 1,
"end_line": node.end_point[0] + 1,
})
for child in node.children:
walk(child)
walk(tree.root_node)
return symbols
调用和导入以相同的方式提取,按它们所在的外层函数作为键。那就是整个图:节点是符号,边是"调用"和"导入"。
第 2 步:将其存储在某个可查询的地方,并使其增量化。
第一个版本在每次运行时重建整个图。这在几百个文件上还不错,但一旦我把它指向一个更大的单仓库就变得无法接受了——在 agent 能做任何事之前要花 40+ 秒。修复方案是对每个文件的内容进行哈希处理,将哈希与其提取的符号一起存储,并只重新解析哈希改变的文件:
import hashlib
def needs_reindex(file_path: str, stored_hashes: dict) -> bool:
content = open(file_path, "rb").read()
current_hash = hashlib.sha256(content).hexdigest()
return stored_hashes.get(file_path) != current_hash
这将一个完整的重建时间从一个典型的单文件编辑时的 40+ 秒降到了不到一秒,因为只有改变的文件(及其直接邻居)需要重新解析。
有一个转折:仅对文件进行哈希处理是不够的。如果文件 A 改变了,任何从其他地方指向文件 A 的边仍然有效,但文件 A 内部的任何边都需要被删除并重建。我最终存储的是以它们的源文件作为键的边,所以失效就是"删除所有源文件 = A 的边,重新提取 A,重新插入"。便宜,而且它避免了在文件缩小后意外留下陈旧的边。
第 3 步:公开查询,而不是查询语言。
我最初的直觉是给 agent 一个通用的图查询接口——想象"写你自己的 Cypher 式查询"。那是我稍后会详细讲的一个错误。真正有效的是一个小的、固定的高级动词集合:
find_definition(symbol_name) — 这个是在哪定义的?
find_callers(symbol_name) — 谁调用这个,在整个仓库中?
trace_call_chain(symbol_name, depth) — 遍历 N 层调用者/被调用者
get_architecture_overview(directory) — 主要模块是什么,它们如何连接?
find_callers 的一个简化例子:
def find_callers(graph, symbol_name: str):
target = graph.get_node(symbol_name)
if not target:
return []
return [
{"caller": edge.source.name, "file": edge.source.file, "line": edge.line}
for edge in graph.incoming_edges(target)
if edge.kind == "calls"
]
agent 调用 find_callers("process_payment") 而不是在整个仓库中 grep process_payment,然后手动过滤掉注释、字符串字面量和不相关的匹配。
以下是整个事物的大致形状:
graph LR
A[Agent] -->|find_callers / trace_call_chain| B[Query Layer]
B --> C[Code Graph: nodes + edges]
C --> D[Incremental Indexer]
D -->|parses via tree-sitter| E[Source Files]
Grep 优先的 agent 在每个会话中都要付出重新发现的成本。没有持久的结构,agent 每次都从头推导"什么调用什么",即使对于一个它"昨天还在里面工作过"的代码库。一个便宜的索引在多个会话中分摊成本,而不是重复支付。
图不能替代阅读代码——它对阅读的位置进行分类。早期我让 agent 在实际打开文件之前直接作用于图查询结果,进行了一个跨越约 15 个调用点的重命名。其中一个结果指向的行由于陈旧的索引条目而略微移动了,编辑落在了错误的地方。现在图在工具描述中被明确框架化为"这是要查看的地方",而不是"这是基本事实"——agent 在编辑之前仍然打开文件。
陈旧性快速破坏信任,所以从第一天就构建增量重新索引。上面的漏洞只发生在我还没有连接哈希处理的时候。一个陈旧索引的事件足以让我暂时不相信这个工具,这会破坏目的。基于内容哈希的增量更新修复了它,但我应该从一开始就内置它,而不是在被烧伤后才匆匆缝补。
查询设计比图的完整性更重要。我的第一个版本公开了一个通用查询接口,以便 agent 可以"问任何事"。实际上它几乎从不很好地使用它——它要么写过于宽泛的查询返回噪声,要么放弃并 grep。四个起名良好、范围窄的动词(find_definition、find_callers、trace_call_chain、get_architecture_overview)在第一次尝试时几乎每次都被正确使用。范围窄且标签清晰的优于范围广且强大的,至少对 LLM 实际上如何获取工具来说。
真正的 ROI 不是速度——而是重构期间的正确性。我预期的收益是"花在 grep 上的 token 更少"。更大的收益结果是定性的:接触许多调用点的重构现在实际上会找到所有调用点。那就是船一个重命名和船一个重命名加一个两个文件外破坏的构建之间的区别。
这最好作为第一选择,而不是唯一选择。当查询返回空时,我仍然让 agent 回退到 grep——有时图确实不覆盖一个案例(一个字符串模板化的调用,一个在解析器还不支持的语言中定义的符号)。将图视为"先试试这个"而不是"完全相信这个"使其即使在无法回答某些东西的日子里仍然有用。
跨服务跟踪。 现在图在进程边界处停止。从一个服务中的 HTTP 处理程序跟踪一个调用到它最终在另一个服务中调用的函数意味着跟踪基于字符串的路由名称,这是 tree-sitter 单独无法解决的。
动态分派和鸭子类型。 在运行时解决的任何东西(依赖注入、getattr 式分派、基于接口的多态)对静态 AST 遍历是不可见的。我正在寻找一个轻量级的运行时跟踪作为这些确切情况的回退,而不是尝试静态建模每个动态模式。
对图无法建模的东西的回退。 对于图确实无法很好地表示的代码,我想要一个辅助搜索路径(语义/嵌入基础的),而不是静默回退到普通 grep。
如果你的 agent 在每个会话中花费一半的时间重新发现你的代码库是如何组合在一起的,一个小的结构索引会快速为自己买单——而且它不需要很花哨。四个查询动词和一个基于哈希的增量解析器为我获得了大部分价值。
如果你已经构建过类似的东西——或找到了更好的方式来处理静态图中的动态分派——我真的很想听听。在 Dev.to 上关注我以获取更多这些构建日志,如果你正在用 Claude Code 进行代理工作流实验,在你自己的仓库中尝试这个模式。