构建代码分析 AI 代理

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

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

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

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

1、代码理解的真正问题

我们都经历过这种情况。你盯着一个函数,它调用三个其他函数,这些函数从四个不同的模块导入,这些模块依赖两个外部服务。你的大脑开始构建心理地图,但它很脆弱。一次中断就会让你回到原点。

传统的代码搜索可以帮助回答"这个函数定义在哪里?"但无法回答"这个系统实际上是如何工作的?"我们需要一些能够理解意图而不仅仅是语法的东西。

如果你能问:

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

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

2、实际工作原理

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

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

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

3、我试图构建什么

  • 摄取任何代码库(我将使用Ruby,但这些概念适用于任何语言)
  • 回答关于代码的自然语言问题
  • 完全在本地运行(代码不会离开你的机器)
  • 高效处理多达100k+行的代码库

3.1 智能代码分块(基础)

大多数尝试失败的地方在于:他们像处理普通文本一样处理代码并任意分割。但代码有结构。一个函数应该在一起。一个类定义需要它的方法。上下文很重要。

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

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
      # Small files: keep entire content
      chunk = create_chunk(file.relative_path_from(@root_dir), content)
      @results << chunk
    else
      # Large files: split by class/method boundaries
      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)
  # Estimate tokens (rough approximation: 1 token ≈ 4 chars)
  estimated_tokens = content.length / 4

  if estimated_tokens > 8000
    # Parse the content and split by blocks/methods
    ast = Prism.parse(content)
    return [fallback_chunk] unless ast

    chunks = []
    current_chunk = []
    current_tokens = 0

    # Recursive function to extract meaningful nodes
    def extract_nodes(node)
      nodes = []
      case node
      when Prism::DefNode, Prism::ClassNode, Prism::ModuleNode, Prism::CallNode
        nodes << node
      when Prism::StatementsNode
        # Recurse into container nodes
        node.body.each { |n| nodes.concat(extract_nodes(n)) }
      end
      nodes
    end

    # Extract all relevant nodes first
    nodes = extract_nodes(ast.value.statements)

    # Process nodes into chunks, respecting token limits
    nodes.each do |node|
      source = node.location.slice  # Get the actual source code
      tokens = source.length / 4

      if tokens > 8000
        # If a single block is too large, truncate it
        source = source[0..(8000 * 4)]
      end

      if current_tokens + tokens > 8000
        # Current chunk is full, start a new one
        chunks << create_chunk(file_path, current_chunk.join("\n")) if current_chunk.any?
        current_chunk = [source]
        current_tokens = tokens
      else
        # Add to current chunk
        current_chunk << source
        current_tokens += tokens
      end
    end

    # Don't forget the last chunk
    chunks << create_chunk(file_path, current_chunk.join("\n")) if current_chunk.any?
    return chunks
  end

  # For small files, return as single chunk with metadata
  {
    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. 令牌感知分块:保持块在8000个令牌以下(GPT的上下文窗口)
  2. 递归提取:自动遍历嵌套结构
  3. 智能批处理:将相关代码分组在一起,直到达到大小限制
  4. 元数据保留:每个块都知道其文件类型、gem和架构层
  5. 回退处理:优雅地处理无法解析的代码

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

  • 类定义作为一个块
  • 每个方法作为单独的块
  • 复杂方法在逻辑边界处分割
  • 所有块都标记有元数据,如"controller"、"api layer"、"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 models expect "query:" or "passage:" prefixes
        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")

  # Create collection
  client.collections.create(
    collection_name: collection_name,
    vectors: {
      size: 1536, # OpenAI embedding size
      distance: "Cosine"
    }
  )

  # Batch upload for efficiency
  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)
  # Generate embedding for the question
  query_embedding = generate_query_embedding(query)

  # Search for similar code
  results = qdrant_client.points.search(
    collection_name: "code_chunks",
    vector: query_embedding,
    limit: 5,
    with_payload: true
  )

  # Build context from search results
  context = results["result"].map do |result|
    "## File: #{result["payload"]["path"]}\n#{result["payload"]["context"]}"
  end.join("\n\n")

  # Ask GPT to explain
  response = openai_client.chat(
    parameters: {
      model: "gpt-4-turbo",
      messages: [
        {
          role: "system",
          content: "You are a code analysis assistant. Explain the following code based on the user's question."
        },
        {
          role: "user",
          content: "Question: #{query}\n\nRelevant Code:\n#{context}"
        }
      ]
    }
  )

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

该过程遵循以下步骤:

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

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

4、原型结果

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

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

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

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

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

查询:"我应该在哪里为用户操作添加日志?"

  • 找到:现有的日志记录模式、审计跟踪实施和建议的集成点
  • 响应时间:1.6秒
  • 准确性:不仅在哪里添加日志,还如何遵循现有模式。

5、我构建这个的经验教训

关键见解

基于AST的分块至关重要。我最初尝试了简单的文本分割,效果很差。函数被分成两半,类与其方法分离。AST解析确保每个块在语义上是完整的。

本地嵌入的表现超出了预期。虽然OpenAI嵌入质量更高,但E5在代码理解方面表现几乎一样好。隐私和成本优势使其值得在大多数用例中考虑。

元数据至关重要。仅仅有代码是不够的。知道一个块是来自控制器、模型还是配置文件,可以显著优先考虑结果。

挑战和限制

跨文件关系难以捕获。该原型很好地理解单个文件,但在它们如何连接方面遇到困难。如果控制器调用服务,它不会自动理解这种关系。

上下文窗口很重要。GPT-4可以处理大量代码,但仍然有限制。对于跨越多个文件的复杂查询,该原型需要仔细的上下文摘要和优先级排序。

错误处理需要改进。当原型失败时,它通常会静默失败。用户提出问题,没有得到结果,并假设系统已损坏。对于生产系统,这个领域需要大量工作。

6、实施指南

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

先决条件

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

快速开始

1)启动Qdrant

docker run -p 6333:6333 qdrant/qdrant

2)扫描你的代码库

# Implement the scanner based on the examples above
./scan_codebase.rb /path/to/your/code

3)生成嵌入

# Local approach (no API key needed)
./embed_by_e5.py
# OR cloud approach (requires OPENAI_API_KEY)
./embed.rb

4)存储在Qdrant中

./store.rb

5)开始提问

./ask.rb "How does authentication work?"

成本细分

  • 本地设置:$0(仅你的计算资源)
  • 云嵌入:约$5–10,用于50k行代码库
  • 持续查询:约$0.01–0.05每个问题

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

7、未来方向

这个原型展示了AI辅助代码理解的潜力。对于生产系统,有几个领域值得进一步探索:

立即扩展

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

长期可能性

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

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

我们正在从代码搜索转向代码理解。我们不再寻找某个东西的定义,而是可以询问它的作用、它为什么存在以及它如何融入大局。

我们今天构建的工具将决定下一代开发者如何学习和使用代码。


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

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