qodo如何构建大型代码仓库的RAG？

qodo如何构建大型代码仓库的RAG？

拥有数千个仓库和数百万行代码的代码库，大多是遗留代码，这使得上下文感知成为企业开发者采用AI时的关键因素。而检索增强型生成（Retrieval Augmented Generation，简称RAG）技术正是解决这一问题的有效手段。本文将详细介绍如何将RAG应用于大规模代码库，以及qodo（前身为Codium）在构建生成式AI编码平台时所采取的策略。

在大的代码仓库中使用RAG

RAG大致可以分为两个部分：索引知识库（这里是指代码库）和检索。对于不断变化的生产代码库，索引并非一次性或定期的任务，而是需要一个强大的管道来持续维护最新的索引。

下图展示了我们的数据摄取管道，文件被路由到适当的分块器进行分块，分块后会添加自然语言描述，并为每个分块生成向量嵌入，然后存储在向量数据库中。

分块（Chunking）

对于自然语言文本，分块相对简单——段落（和句子）提供了创建语义上有意义的段落的明显边界点。然而，简单的分块方法在准确划分代码的有意义段落时却会遇到问题，导致边界定义不准确以及包含无关或不完整信息的问题。我们发现，向大型语言模型（LLM）提供无效或不完整的代码段实际上会损害性能并增加幻觉现象，而不是提供帮助。

Sweep AI团队去年发表了一篇很棒的博客文章，详细介绍了他们的代码分块策略。他们开源了使用具体语法树（CST）解析器创建内聚块的方法，该方法已被LlamaIndex采用。

这是我们出发点，但我们遇到了一些问题：

1. 尽管有所改进，但块仍然不总是完整的，有时会缺少关键上下文，如导入语句或类定义。
2. 对可嵌入块大小的硬性限制并不总是允许捕获较大代码结构的完整上下文。
3. 该方法没有考虑到企业级代码库的独特挑战。

为了解决这些问题，我们开发了几种策略：

智能分块策略（Intelligent Chunking Strategies）

Sweep AI使用静态分析实现分块，这比以前的方法有了很大的改进。但在当前节点超出token limit(令牌限制)并开始在不考虑上下文的情况下将其子节点分割成块时，这种方法并不理想。这可能导致在方法或if语句中间中断块（例如，“if”在一个块中，“else”在另一个块中）。

为缓解此问题，qodo使用特定于语言的静态分析递归地将节点分割成更小的块，并执行追溯处理以重新添加任何被移除的关键上下文。这使我们能够创建保持代码结构的块，将相关元素保持在一起。

from utilities import format_complex
class ComplexNumber:

        def __init__(self, real, imag):
                 self.real = real
                 self.imag = imag

        def modulus(self):
                 return math.sqrt(self.real**2 + self.imag**2)

        def add(self, other):
                 return ComplexNumber(self.real + other.real, self.imag + other.imag)

        def multiply(self, other):
                 new_real = self.real * other.real - self.imag * other.imag 
                 new_imag = self.real * other.imag + self.imag * other.real
                 return ComplexNumber(new_real, new_imag)

        def __str__(self): 
                return format_complex(self.real, self.imag)

原来的分块：

        def __str__(self): 
                return format_complex(self.real, self.imag)

qodo的分块：

from utilities import format_complex

class ComplexNumber:

       def __init__(self, real, imag):
                 self.real = real
                 self.imag = imag
      # … 

      def __str__(self): 
                return format_complex(self.real, self.imag)

qodo的分块器将关键上下文与类方法保持在一起，包括任何相关的导入以及类定义和初始化方法，确保AI模型拥有理解并处理此代码所需的所有信息。

在分块中维护上下文（Maintaining Context in Chunks）

（embedding）嵌入较小的块通常可以获得更好的性能。理想情况下，您希望拥有尽可能小的块，同时包含相关上下文——包含的任何无关内容都会稀释嵌入的语义含义。

qodo的目标是使块尽可能小，并将大小限制在大约500个字符左右。但是，较大的类或复杂的代码结构通常会超出此限制，导致代码表示不完整或支离破碎。

因此，qodo开发了一个系统，允许灵活的块大小，并确保将关键上下文（如类定义和导入语句）包含在相关的块中。

对于大型类，qodo可以创建一个嵌入(embedding)，并分别索引各个方法，但每个方法块中都包含类定义和相关导入（import）项。这样一来，当检索到特定方法时，AI模型就能获得理解和使用该方法所需的完整上下文。

针对不同文件类型的特殊处理

不同的文件类型（例如代码文件、配置文件、文档）需要不同的分块策略以保持其语义结构。

我们为各种文件类型实现了专门的分块策略，特别关注像OpenAPI/Swagger规范这样的具有复杂、相互关联结构的文件。

对于OpenAPI文件，我们不是按行或字符进行分块，而是按端点进行分块，确保每个块包含特定API端点的所有信息，包括其参数、响应和安全定义。

使用功能描述增强嵌入

代码嵌入通常无法捕捉代码的语义含义，尤其是对于自然语言查询。

我们使用LLM为每个代码块生成自然语言描述。然后，这些描述与代码一起嵌入，增强了我们检索自然语言查询相关代码的能力。

对于前面提到的map_finish_reason函数：

def map_finish_reason(finish_reason: str,):
    # openai supports 5 stop sequences - 'stop', 'length', 'function_call', 'content_filter', 'null'
    # anthropic mapping
    if finish_reason == "stop_sequence":
        return "stop"
    # cohere mapping - https://docs.cohere.com/reference/generate
    elif finish_reason == "COMPLETE":
        return "stop"
    elif finish_reason == "MAX_TOKENS":  # cohere + vertex ai
        return "length"
    elif finish_reason == "ERROR_TOXIC":
        return "content_filter"
    elif (
        finish_reason == "ERROR"
    ):  # openai currently doesn't support an 'error' finish reason
        return "stop"
    # huggingface mapping https://huggingface.github.io/text-generation-inference/#/Text%20Generation%20Inference/generate_stream
    elif finish_reason == "eos_token" or finish_reason == "stop_sequence":
        return "stop"
    elif (
        finish_reason == "FINISH_REASON_UNSPECIFIED" or finish_reason == "STOP"
    ):  # vertex ai - got from running `print(dir(response_obj.candidates[0].finish_reason))`: ['FINISH_REASON_UNSPECIFIED', 'MAX_TOKENS', 'OTHER', 'RECITATION', 'SAFETY', 'STOP',]
        return "stop"
    elif finish_reason == "SAFETY" or finish_reason == "RECITATION":  # vertex ai
        return "content_filter"
    elif finish_reason == "STOP":  # vertex ai
        return "stop"
    elif finish_reason == "end_turn" or finish_reason == "stop_sequence":  # anthropic
        return "stop"
    elif finish_reason == "max_tokens":  # anthropic
        return "length"
    elif finish_reason == "tool_use":  # anthropic
        return "tool_calls"
    elif finish_reason == "content_filtered":
        return "content_filter"
    return finish_reason

我们可能会生成如下描述：

“Python函数，用于标准化来自不同AI平台的完成原因，将特定于平台的原因映射到诸如‘stop’、‘length’和‘content_filter’等通用术语。”

然后，此描述与代码一起嵌入，提高了检索类似“如何在不同平台之间规范化AI完成状态”这类查询的能力。这种方法旨在解决当前嵌入模型的差距，这些模型并非面向代码，且在自然语言和代码之间缺乏有效的转换。

高级检索和排名

简单的向量相似性搜索通常会检索到无关或上下文不符的代码片段，特别是在拥有数百万个索引块的大型多样化代码库中。

qodo实现了一个两阶段检索过程。首先，我们从向量存储中执行初始检索。然后，我们使用LLM根据其与特定任务或查询的相关性来过滤和排名结果。

如果开发人员查询“如何处理API速率限制”，我们的系统可能会首先检索与API调用和错误处理相关的几个代码片段。然后，LLM会根据查询的上下文分析这些片段，将那些专门处理速率限制逻辑的片段排名更高，并丢弃不相关的片段。

为企业仓库扩展RAG

当仓库数量达到数千个时，如果每个查询都跨所有仓库搜索，检索将变得嘈杂且效率低下。

qodo正在开发仓库级别的过滤策略，以在深入到各个代码块之前缩小搜索范围。这包括“黄金仓库”的概念——允许组织指定符合最佳实践且包含良好组织代码的特定仓库。

对于有关特定微服务架构模式的查询，qodo系统可能会首先根据元数据和高级内容分析确定最有可能包含相关信息的前5到10个仓库。然后，它在这些仓库内执行详细的代码搜索，显著减少了噪声并提高了相关性。

RAG测试和评价

由于缺乏标准化的基准，评估代码RAG系统的性能是一项挑战。

qodo开发了一种多方面的评估方法，结合了自动化指标和来自企业客户的实际使用数据。

使用相关性评分（开发人员实际使用检索到的代码片段的频率）、准确性指标（对于代码补全任务）和效率测量（响应时间、资源使用）的组合。另外，还与企业客户密切合作，收集反馈和实际性能数据。

总结

为大规模企业代码库实施RAG带来了独特的挑战，这些挑战超出了典型的RAG应用范围。通过专注于智能分块、增强嵌入、高级检索技术和可扩展架构，qodo开发了一个能够有效导航并利用企业级代码库中丰富知识的系统。

参考资料

• https://blog.kuzudb.com/post/enhancing-graph-rag-with-vector-search/

• https://www.qodo.ai/blog/what-is-rag-retrieval-augmented-generation/