构建代码库分析AI代理

我决定构建一个能够读取、理解和解释代码库的AI代理。不是简单的ChatGPT包装器,而是一个能够理解代码结构、关系和上下文的系统。

构建代码库分析AI代理

几个月前,我接手了一个50,000行的Ruby代码库,没有任何文档。之前的团队已经消失,留下了一个由相互连接的服务、自定义gem和分散在几十个文件中的业务逻辑组成的迷宫。听起来熟悉吗?

与其花几周时间手动解析代码,我决定构建一个原型:一个能够读取、理解和解释代码库的AI代理。不是简单的ChatGPT包装器,而是一个能够理解代码结构、关系和上下文的系统。

这个原型的效果超出了预期。它显著减少了理解代码库结构和定位相关功能所需的时间。

以下是我在项目中构建的内容以及可以应用于你自己的项目的经验教训。

1、代码理解的真正问题

我们都有过这样的经历。你盯着一个调用三个其他函数的函数,这些函数从四个不同的模块导入,这些模块依赖于两个外部服务。你的大脑开始构建一个心理地图,但它是脆弱的。一旦被打断,你就回到了起点。

传统的代码搜索有助于回答“这个函数在哪里定义的?”,但在回答“这个系统实际上是如何工作的?”时却失败了。我们需要一种能理解意图而不是仅仅语法的东西。

如果你可以问:

  • “用户认证如何在整个系统中流动?”
  • “当API请求失败时会发生什么?”
  • “在哪里添加日志来跟踪用户行为?”

这不仅仅是取代开发人员。而是增强我们快速理解复杂系统的能力。

2、它实际上是怎样的

把它想象成为你的代码建立一个智能图书管理员:

  1. 扫描器读取整个代码库,并将其分解为有意义的部分(不是随机文本块,而是实际的函数、类和模块)
  2. 嵌入引擎将每个部分转换为捕捉其含义的数学表示
  3. 向量数据库存储这些表示,以便高效地找到相似的代码
  4. 查询接口接收你的问题,找到最相关的代码,并让GPT解释它

关键的技术是语义搜索。我们不寻找精确的关键字匹配,而是寻找与你的问题概念相关的代码。询问“用户登录”时,它会找到身份验证中间件、会话处理和密码验证——即使这些确切的词从未出现过。

3、我尝试构建的内容

  • 接受任何代码库(我将以Ruby为例,但这些概念适用于任何语言)
  • 回答关于你的代码的自然语言问题
  • 完全本地运行(没有代码离开你的机器)
  • 高效处理超过100k行的代码库

3.1 智能代码分块(基础)

这是大多数尝试失败的地方:它们将代码视为普通文本并随意分割。但代码有结构。一个函数应该在一起。一个类定义需要它的方法。上下文很重要。

关键的见解是使用抽象语法树(ASTs)在有意义的边界上分割代码。

class CodebaseScanner  
  def initialize(root_dir)  
    @root_dir = Pathname.new(root_dir)  
    @results = []  
  end  

  def scan  
    scan_directory(@root_dir)  
    @results  
  end  

  private  

  def process_file(file)  
    content = File.read(file, encoding: "UTF-8")  

    if content.size <= 1000  
      # 小文件:保留整个内容  
      chunk = create_chunk(file.relative_path_from(@root_dir), content)  
      @results << chunk  
    else  
      # 大文件:按类/方法边界分割  
      chunks = split_by_ast_nodes(file, content)  
      @results.concat(chunks)  
    end  
  end  
end

这种方法解决了三个关键问题:

  1. 语义边界:函数保持在一起,类保持其方法,模块保持其结构
  2. 上下文保留:每个块知道它来自哪个文件,是什么类型的代码(控制器、模型、配置),以及它在系统中的作用
  3. 大小优化:小文件(小于1000字符)保持完整,大文件在自然边界上智能分割

元数据非常重要。当有人询问“数据库查询”时,我们可以优先考虑标记为“模型”或“迁移”的块,而不是通用的实用函数。

为什么AST很重要:传统的文本分割可能会将一个函数分成两半,或者将一个类与其方法分开。AST解析确保每个块是一个完整的、有意义的代码单元。

下面是递归解析的实际工作方式,使用实际的create\_chunk方法:

def create_chunk(file_path, content)  
  # 估计token数(粗略估算:1 token ≈ 4个字符)  
  estimated_tokens = content.length / 4  

  if estimated_tokens > 8000  
    # 解析内容并按块/方法分割  
    ast = Prism.parse(content)  
    return [fallback_chunk] unless ast  

    chunks = []  
    current_chunk = []  
    current_tokens = 0  

    # 递归函数提取有意义的节点  
    def extract_nodes(node)  
      nodes = []  
      case node  
      when Prism::DefNode, Prism::ClassNode, Prism::ModuleNode, Prism::CallNode  
        nodes << node  
      when Prism::StatementsNode  
        # 递归进入容器节点  
        node.body.each { |n| nodes.concat(extract_nodes(n)) }  
      end  
      nodes  
    end  

    # 首先提取所有相关节点  
    nodes = extract_nodes(ast.value.statements)  

    # 将节点处理为块,尊重token限制  
    nodes.each do |node|  
      source = node.location.slice  # 获取实际源代码  
      tokens = source.length / 4  

      if tokens > 8000  
        # 如果单个块太大,截断它  
        source = source[0..(8000 * 4)]  
      end  

      if current_tokens + tokens > 8000  
        # 当前块已满,开始新块  
        chunks << create_chunk(file_path, current_chunk.join("\n")) if current_chunk.any?  
        current_chunk = [source]  
        current_tokens = tokens  
      else  
        # 添加到当前块  
        current_chunk << source  
        current_tokens += tokens  
      end  
    end  

    # 不要忘记最后一个块  
    chunks << create_chunk(file_path, current_chunk.join("\n")) if current_chunk.any?  
    return chunks  
  end  

  # 对于小文件,返回单个块和元数据  
  {  
    id: generate_id(file_path),  
    context: content,  
    path: file_path,  
    metadata: {  
      type: determine_file_type(file_path),  
      gem: determine_gem(file_path),  
      layer: determine_layer(file_path)  
    }  
  }  
end

这种方法处理了现实世界的复杂性:

  1. 基于token的分块:保持块在8000个token以内(GPT的上下文窗口)
  2. 递归提取:自动遍历嵌套结构
  3. 智能批处理:将相关代码组合在一起,直到达到大小限制
  4. 元数据保留:每个块知道它的文件类型、gem和架构层
  5. 回退处理:优雅处理无法解析的代码

当处理一个大型Rails控制器时,这可能提取:

  • 类定义作为一个块
  • 每个方法作为单独的块
  • 复杂的方法在逻辑边界上分割
  • 所有都标记为“控制器”、“api层”、“business_logic gem”

结果:而不是任意的文本块,你得到了语义上有意义的代码单元,AI可以真正理解和解释。

3.2 将代码转化为数学(嵌入解释)

下一步是将代码转化为捕捉语义意义的数值向量。

把嵌入想象成这样:相似的代码获得相似的数字。一个登录函数和一个身份验证中间件会有数学上接近的向量,即使它们使用完全不同的变量名。

我构建了两种方法——一种用于速度和便利,另一种用于隐私和成本控制:

选项1:OpenAI嵌入(云)

def generate_embeddings(chunks)  
  client = OpenAI::Client.new(access_token: ENV['OPENAI_API_KEY'])  

  chunks.each do |chunk|  
    response = client.embeddings(  
      parameters: {  
        model: 'text-embedding-3-small',  
        input: chunk['context']  
      }  
    )  
    chunk['embedding'] = response.dig('data', 0, 'embedding')  
  end  
end
选项2:本地E5嵌入(离线)
from sentence_transformers import SentenceTransformer  

def generate_local_embeddings(chunks):  
    model = SentenceTransformer('intfloat/e5-large-v2')  

    for chunk in chunks:  
        # E5模型需要“query:”或“passage:”前缀  
        text = f"passage: {chunk['context']}"  
        embedding = model.encode(text, normalize_embeddings=True)  
        chunk['embedding'] = embedding.tolist()

权衡:OpenAI嵌入质量更高且更容易设置,但会花费金钱并将你的代码发送到他们的服务器。E5嵌入完全在你的机器上运行——没有API密钥,没有数据离开你的系统,没有持续成本。

对于大多数项目,质量差异可以忽略不计。实际上,我更喜欢本地方法,特别是对于包含专有代码的内容。

3.3 快速相似性搜索(进入向量数据库)

现在我们有数千个代码块,每个都表示为1,536个数字的向量。当有人提问时,我们需要在毫秒而不是分钟内找到最相似的向量。

向量数据库是为此用例设计的。我选择了Qdrant,因为它性能良好、可靠,并且可以在一个Docker容器中运行。

def store_embeddings(chunks, collection_name)  
  client = Qdrant::Client.new(url: "http://localhost:6333")  

  # 创建集合  
  client.collections.create(  
    collection_name: collection_name,  
    vectors: {  
      size: 1536, # OpenAI嵌入大小  
      distance: "Cosine"  
    }  
  )  

  # 批量上传以提高效率  
  chunks.each_slice(100) do |batch|  
    points = batch.map do |chunk|  
      {  
        id: generate_id(chunk["path"]),  
        vector: chunk["embedding"],  
        payload: {  
          context: chunk["context"],  
          path: chunk["path"],  
          type: chunk["metadata"]["type"]  
        }  
      }  
    end  

    client.points.upsert(  
      collection_name: collection_name,  
      points: points  
    )  
  end  
end

Qdrant高效处理向量相似性计算。你可以查询最相似的5个块,并在不到50毫秒内得到结果,即使有10,000多个代码块。

提示:批量上传至关重要。逐个上传块会耗费很长时间。每批100个可将上传时间从小时减少到分钟。

3.4 整合所有内容(查询界面)

这里是系统真正有用的时刻。用户提出一个自然语言问题,我们会返回一个包含相关代码示例的全面答案。

def answer_question(query)  
  # 生成问题的嵌入  
  query_embedding = generate_query_embedding(query)  

  # 搜索相似代码  
  results = qdrant_client.points.search(  
    collection_name: "code_chunks",  
    vector: query_embedding,  
    limit: 5,  
    with_payload: true  
  )  

  # 从搜索结果构建上下文  
  context = results["result"].map do |result|  
    "## 文件: #{result["payload"]["path"]}\n#{result["payload"]["context"]}"  
  end.join("\n\n")  

  # 让GPT解释  
  response = openai_client.chat(  
    parameters: {  
      model: "gpt-4-turbo",  
      messages: [  
        {  
          role: "system",  
          content: "你是一个代码分析助手。根据用户的问题解释以下代码。"  
        },  
        {  
          role: "user",   
          content: "问题: #{query}\n\n相关代码:\n#{context}"  
        }  
      ]  
    }  
  )  

  response.dig("choices", 0, "message", "content")  
end

该过程遵循以下步骤:

  1. 问题 → 向量:将用户的问题转换为与我们的代码相同的嵌入空间
  2. 向量 → 代码:找到最语义相似的代码块
  3. 代码 → 上下文:将相关代码与元数据和文件路径捆绑在一起
  4. 上下文 → 答案:让GPT分析代码并提供人类解释

关键的洞察是:我们不是要求GPT理解整个代码库。我们只是给它一些相关的部分,并让它解释这些特定的块。

4、原型结果

我在推动这个项目的20,000行Ruby代码库上测试了这个原型。以下是结果:

查询:“用户认证是如何工作的?”

  • 找到:身份验证中间件、会话处理、JWT令牌验证和密码哈希逻辑,分布在4个不同的文件中
  • 响应时间:2.1秒
  • 准确性:找到了所有相关组件,包括我忘记的边缘情况。

查询:“当API请求失败时会发生什么?”

  • 找到:服务层中的错误处理、重试逻辑、通知系统和数据库回滚程序
  • 响应时间:1.8秒
  • 准确性:全面覆盖整个故障恢复流程。

查询:“在哪里添加日志以跟踪用户行为?”

  • 找到:现有的日志模式、审计追踪实现和建议的集成点
  • 响应时间:1.6秒
  • 准确性:不仅是添加日志的位置,还有如何遵循现有模式。

5、我在构建这个过程中学到的经验

关键见解

  • 基于AST的分块至关重要。我最初尝试简单的文本分割,结果很差。函数被一分为二,类与它们的方法分离。AST解析确保每个块在语义上是完整的。
  • 本地嵌入的表现超出预期。虽然OpenAI嵌入的质量更高,但E5的表现几乎同样好,用于代码理解。隐私和成本优势使其对大多数用例来说值得考虑。
  • 元数据至关重要。仅仅拥有代码是不够的。知道一个块来自控制器、模型还是配置文件,可以极大地优先排序结果。

挑战和局限性

  • 跨文件关系难以捕捉。原型很好地理解了单个文件,但很难理解它们如何连接。如果一个控制器调用一个服务,它不会自动理解这种关系。
  • 上下文窗口很重要。GPT-4可以处理大量代码,但仍有限制。对于跨越许多文件的复杂查询,原型需要仔细的上下文总结和优先级排序。
  • 错误处理需要改进。当原型失败时,它经常无声地失败。用户问一个问题,得到没有结果,认为系统坏了。这一领域需要重大改进,才能用于生产系统。

6、实施指南

以下是实现类似原型的方法:

前提条件

  • Ruby 3+(用于扫描器)
  • Python 3.8+(用于本地嵌入)
  • Docker(用于Qdrant)
  • OpenAI API密钥(可选,用于云嵌入)

快速入门

启动Qdrant

docker run -p 6333:6333 qdrant/qdrant

扫描你的代码库

# 根据上面的例子实现扫描器  
./scan_codebase.rb /path/to/your/code

生成嵌入

# 本地方法(不需要API密钥)  
./embed_by_e5.py  
# 或者云方法(需要OPENAI_API_KEY)  
./embed.rb

存储到Qdrant

./store.rb

开始提问

./ask.rb "认证是如何工作的?"

成本分解

  • 本地设置:$0(只需你的计算能力)
  • 云嵌入:约5-10美元用于一个5万行的代码库
  • 持续查询:每次提问约0.01-0.05美元

整个原型运行在笔记本电脑上。不需要云基础设施。

未来方向

这个原型展示了AI辅助代码理解的潜力。几个领域需要进一步探索以用于生产系统:

立即扩展

  • 代码审查协助:“这个更改可能有什么问题?”
  • 架构文档:“为新团队成员生成系统概述”
  • 重构指导:“你会如何改进这段代码?”
  • 漏洞挖掘:“在请求处理流程中查找潜在的安全问题”

长期可能性

未来的开发环境可以支持:

  • 任何函数或模块的自然语言解释
  • 功能放置的架构指导
  • 旧代码分析和文档生成
  • 新团队成员的加速入职

原文链接:Building an AI Agent for Codebase Analysis and Understanding

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