Skip to main content

自定义代码 RAG

虽然Continue自带@Codebase开箱即用，但您可能希望设置自己的向量数据库并构建自定义的检索增强生成（RAG）系统。这可以让您访问本地不可用的代码，在所有用户之间一次性索引代码，或包含自定义逻辑。在本指南中，我们将带您完成构建此系统的步骤。

步骤1：选择一个嵌入模型​

如果可能，我们推荐使用voyage-code-3，它将为代码提供现有嵌入模型中最准确的答案。您可以在此获取API密钥。因为他们的API是OpenAI兼容的，您可以通过替换URL来使用任何OpenAI客户端。

步骤 2: 选择一个向量数据库​

有许多可用的向量数据库，但由于大多数向量数据库都能够高效处理大型代码库，我们建议选择一个易于设置和实验的数据库。

LanceDB 是一个很好的选择，因为它可以与 Python 和 Node.js 的库一起在内存中运行。这意味着在开始时，你可以专注于编写代码，而不是设置基础设施。如果你已经选择了一个向量数据库，那么使用它而不是 LanceDB 也是一个不错的选择。

步骤3：选择一个“分块”策略​

大多数嵌入模型一次只能处理有限数量的文本。为了解决这个问题，我们将代码“分块”成更小的部分。

如果你使用voyage-code-3，它的最大上下文长度为16,000个标记，足以容纳大多数文件。这意味着在开始时，你可以采用一种更简单的策略，即截断超过限制的文件。从最简单到最全面的顺序，你可以使用的3种分块策略是：

  1. 当文件超过上下文长度时截断文件：在这种情况下，每个文件始终只有一个块。
  2. 将文件分割成固定长度的块：从文件顶部开始，将行添加到当前块中，直到达到限制，然后开始一个新的块。
  3. 使用递归的、基于抽象语法树（AST）的策略：这是最精确的，但也是最复杂的。在大多数情况下，您可以通过使用（1）或（2）获得高质量的结果，但如果您想尝试这种方法，您可以在我们的代码分块器或LlamaIndex中找到参考示例。

正如本指南中通常建议的那样，我们推荐从能够以20%的努力获得80%收益的策略开始。

步骤4：编写索引脚本​

索引，我们将以可检索的格式将您的代码插入向量数据库中，分为三个步骤：

  1. 分块
  2. 生成嵌入
  3. 插入到向量数据库中

使用LanceDB，我们可以同时执行步骤2和3，如在他们的文档中所示。例如，如果您使用Voyage AI，它将配置如下：

from lancedb.pydantic import LanceModel, Vector
from lancedb.embeddings import get_registry

db = lancedb.connect("/tmp/db")
func = get_registry().get("openai").create(
    name="voyage-code-3",
    base_url="https://api.voyageai.com/v1/",
    api_key=os.environ["VOYAGE_API_KEY"],
)

class CodeChunks(LanceModel):
    filename: str
    text: str = func.SourceField()
    # 1536 is the embedding dimension of the `voyage-code-3` model.
    vector: Vector(1536) = func.VectorField()

table = db.create_table("code_chunks", schema=CodeChunks, mode="overwrite")
table.add([
    {"text": "print('hello world!')", filename: "hello.py"},
    {"text": "print('goodbye world!')", filename: "goodbye.py"}
])

query = "greetings"
actual = table.search(query).limit(1).to_pydantic(CodeChunks)[0]
print(actual.text)
tip

如果你正在索引多个仓库，最好将这些仓库存储在单独的“表”（LanceDB使用的术语）或“集合”（其他一些向量数据库使用的术语）中。另一种方法是添加一个“仓库”字段，然后通过该字段进行过滤，但这种方法性能较差。

无论您选择了哪个数据库或模型，您的脚本都应遍历您希望索引的所有文件，将它们分块，为每个块生成嵌入，然后将所有块插入到您的向量数据库中。

步骤5：运行你的索引脚本​

tip

在一个完美的生产版本中，您可能希望构建“自动、增量索引”，这样每当文件发生变化时，只有该文件会自动重新索引，而其他文件则不会。这样做的好处是可以保持嵌入的最新状态，并且成本较低。

也就是说，我们强烈建议在尝试此操作之前先构建并测试管道。除非您的代码库经常被完全重写，否则完全刷新索引可能已经足够并且成本合理。

此时，您已经编写了索引脚本并测试了可以从向量数据库中进行查询。现在，您需要一个计划来确定何时运行索引脚本。

最初，您可能需要手动运行它。一旦您确信您的自定义RAG提供了价值并准备长期使用，那么您可以设置一个cron作业来定期运行它。由于代码库在短时间内基本不变，您不需要每天重新索引超过一次。每周或每月一次可能就足够了。

步骤6：设置您的服务器​

为了让Continue扩展访问您的自定义RAG系统，您需要设置一个服务器。该服务器负责接收来自扩展的查询，查询向量数据库，并以Continue预期的格式返回结果。

这里是一个使用FastAPI的参考实现，能够处理来自Continue的“http”上下文提供者的请求。

"""
This is an example of a server that can be used with the "http" context provider.
"""

from fastapi import FastAPI
from pydantic import BaseModel


class ContextProviderInput(BaseModel):
    query: str
    fullInput: str


app = FastAPI()


@app.post("/retrieve")
async def create_item(item: ContextProviderInput):
    results = [] # TODO: Query your vector database here.

    # Construct the "context item" format expected by Continue
    context_items = []
    for result in results:
        context_items.append({
            "name": result.filename,
            "description": result.filename,
            "content": result.text,
        })

    return context_items

在您设置好服务器后，您可以通过将“http”上下文提供者添加到config.json中的contextProviders数组来配置Continue以使用它：

config.json
{
  "name": "http",
  "params": {
    "url": "https://myserver.com/retrieve",
    "title": "http",
    "description": "Custom HTTP Context Provider",
    "displayTitle": "My Custom Context"
  }
}

步骤 7（额外）：设置重新排名​

如果您想提高结果的质量，一个很好的第一步是添加重新排序。这包括从向量数据库中检索更大的初始结果池，然后使用重新排序模型将它们从最相关到最不相关进行排序。这是因为重新排序模型可以对一小部分顶部结果进行稍微更昂贵的计算，因此可以提供比相似性搜索更准确的排序，而相似性搜索必须搜索数据库中的所有条目。

如果您希望每次查询返回10个结果，那么您需要：

  1. 使用相似性搜索从向量数据库中检索约50个结果
  2. 将所有这50个结果连同查询一起发送到reranker API，以获取每个结果的相关性分数
  3. 按相关性得分排序结果并返回前10个

我们推荐使用Voyage AI的rerank-2模型，这里有使用示例here。