跳到主内容

自定义代码 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 个 token，这足以容纳大多数文件。这意味着一开始您可以采用一种更简单的方法，即截断超出限制的文件。按照从易到难的顺序，您可以使用以下 3 种分块策略：

  1. 当文件超出上下文长度时截断文件：在这种情况下，每个文件始终只有 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()
    # 1024 is the default dimension for `voyage-code-3`: https://docs.voyageai.com/docs/embeddings#model-choices
    vector: Vector(1024) = 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)
提示

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

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

第 5 步：运行索引脚本​

提示

在完美的生产版本中，您会希望构建“自动增量索引”，以便每当文件更改时，该文件（仅该文件）会自动重新索引。这样做的好处是嵌入信息完全最新且成本较低。

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

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

一开始，您可能应该手动运行它。一旦您确信您的自定义 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”上下文提供商添加到配置中的 contextProviders 数组来配置 Continue 使用它。

config.yaml
context:
  - provider: http
    params:
      url: https://myserver.com/retrieve
      name: http
      description: Custom HTTP Context Provider
      displayTitle: My Custom Context

第 7 步（附加）：设置重新排序​

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

例如，如果您希望为每个查询返回总共 10 个结果，那么您可以：

  1. 使用相似性搜索从向量数据库中检索约 50 个结果
  2. 将所有这 50 个结果以及查询发送到重新排序 API，以便获取每个结果的相关性分数
  3. 按相关性分数对结果进行排序并返回前 10 个

我们建议使用 Voyage AI 的 rerank-2 模型，其用法示例可在此处找到。