四、从 0 开始构建一个代码库-向量数据库的选择与集成

四、从 0 开始构建一个代码库-向量数据库的选择与集成

主流向量数据库对比

常用的向量数据库有很多，如 Faiss、Chroma、LanceDB、Milvus 等。它们各有特点和适用场景，需要根据具体需求进行选择。以下是一些常见的对比点：

• faiss
• chroma
• lanceDB
• milvus

以下是 Faiss、Chroma、LanceDB、Milvus 四个数据库的特点和适用场景的对比：

数据库	类型	**核心特点	适用场景	数据规模	实时性	生态集成
Faiss	向量检索库（非数据库）	- 专注高效向量相似度搜索与聚类 - 支持多种索引算法（L2、HNSW、IVF等） - C++高性能实现，需结合外部存储扩展	- 机器学习推理（推荐系统、图像检索） - 离线批量检索 - 科研与算法验证	百万级-十亿级（需扩展）	离线为主	PyTorch、TensorFlow
Chroma	嵌入式向量数据库	- 轻量级、易集成（Python/JS客户端） - 支持向量与结构化数据混合存储 - 内置SQLite，可对接PostgreSQL	- 中小型应用（聊天机器人、知识管理） - 快速原型开发 - 边缘计算/本地部署	百万级以内	支持实时	LangChain、QA工具
LanceDB	数据湖向量数据库	- 基于Parquet格式，兼容大数据生态（Spark、Pandas） - 支持向量+SQL联合查询 - 数据版本管理与时间序列支持	- 数据湖分析（日志、传感器数据） - 多模态检索（图像+文本标签） - 时间序列场景	千万级-百亿级（依赖数据湖）	近实时	Spark、Pandas、Delta Lake
Milvus	分布式向量数据库	- 原生分布式架构（Kubernetes部署） - 支持十亿级向量，毫秒级实时查询 - 企业级功能（持久化、故障恢复） - 集成Pinecone、LangChain	- 大规模生产环境（推荐系统、人脸识别） - 高并发场景（社交平台实时推荐） - 金融/医疗等可靠性要求高的领域	十亿级+（原生支持）	实时（毫秒级）	Kubernetes、Pinecone、LangChain

在我们的这个教程中，使用的是 LanceDB，因为它是一个轻量级的向量数据库，支持多种索引算法，并且支持多种语言的客户端。

向量数据库的初始化与配置

执行下面的命令安装 lancedb 库

npm install @lancedb/lancedb

以 LanceDB 为例，介绍数据库的安装和启动

创建文件在 src/core/embed.js, 并将下面的代码复制到文件中

下面的代码主要实现了以下功能：

• 初始化 LanceDB 数据库
• 配置向量数据库的索引类型
• 设置数据库的存储路径和内存使用策略
• 向量模型（embed model）的使用

import * as lancedb from "@lancedb/lancedb";
import { env, pipeline } from "@xenova/transformers";

const databaseDir = "./db/lancedb";
const tableName = "codebase";

/**
 * 初始化embed模型
 * @returns
 */
async function initEmbedModel() {
  env.allowLocalModels = true;
  env.allowRemoteModels = false;
  env.localModelPath = "models";

  // 创建一个来自管道的嵌入函数，该函数从批次中返回一个向量列表。
  // sourceColumn 是要嵌入的数据中的列的名称。
  // 管道的输出是一个张量 { data: Float32Array(384) }，因此要过滤出向量。
  let pipe = await pipeline("feature-extraction", "Xenova/all-MiniLM-L6-v2");
  return {
    sourceColumn: "", // 嵌入的列名
    embed: async function (batch) {
      let result = [];
      for (let text of batch) {
        const res = await pipe(text, { pooling: "mean", normalize: true });
        result.push(res.tolist());
      }
      return result;
    },
  };
}

/**
 * 将文本内容转换为向量并保存到 LanceDB
 * @param {*} data []Object
 */
async function saveLancedb(data) {
  const embedModel = await initEmbedModel();
  for (const row of data) {
    // 计算向量
    row.vector = (await embedModel.embed([row.text]))[0];
  }
  const db = await lancedb.connect(databaseDir);
  const table = await db.createTable(tableName, data, { mode: "overwrite" });
}

/**
 * 检索LanceDB中的内容
 * @param query
 * @param topK
 */
async function searchLancedb(query, topK = 5) {
  const embedModel = await initEmbedModel();
  // 查询相似内容
  let queryEmbed = (await embedModel.embed(query))[0];
  const db = await lancedb.connect(databaseDir);
  const _tb = await db.openTable(tableName);
  const results = await _tb
    .vectorSearch(queryEmbed)
    // .where("type == 'fruit'") // 可添加过滤条件
    .distanceType("cosine")
    .limit(topK)
    .toArray();

  return results.map((r) => r.text);
}

export { saveLancedb, searchLancedb };

选择合适的向量模型

Xenova/all-MiniLM-L6-v2

介绍：

huggingface.co/Xenova/all-...

all-MiniLM-L6-v2 是一种基于 Transformer 的向量模型，具有强大的语义表示能力，可以将内容映射到384 维的密集向量空间。它可以处理最多256个词的输入文本，并将文本转换为向量形式。其基于 MiniLM 架构，是一种轻量级、低延迟的语言模型，专门设计用于高效的文本嵌入生成，为众多轻量级应用提供了高效且实用的离线解决方案。

特点与优势：

轻量级：参数量约为 38M，模型文件大小仅约 70MB。运行时资源消耗少，在 CPU 上推理速度可达 780 字 / 秒，GPU 显存需求仅 2GB，适合在边缘设备、集成显卡或资源受限的环境中运行。

性能出色：在句子相似度、信息检索等任务中表现优异，在相关的 MTEB 榜单准确率接近大型模型，尤其擅长处理短文本。并且对多语言有较好的兼容性，支持 30 多种语言场景。

开发友好：借助 sentence-transformers 库，仅需几行代码即可加载模型并生成句子嵌入，开发成本低，便于集成到各种应用中。同时，它还支持与 Faiss 等向量数据库结合，实现高效的文本聚类与检索。

如何将文本数据转换为向量表示

我们这里使用 all-MiniLM-L6-v2 作为向量模型，将文本数据转换为向量表示。

import { env, pipeline } from "@xenova/transformers";
import path from "path";

// env.localModelPath = path.join(__dirname, "models");
env.localModelPath = "models";
env.allowLocalModels = true;
// 禁用从 Hugging Face Hub 加载远程模型：
env.allowRemoteModels = false;

// env.backends.onnx.wasm.wasmPaths = "/path/to/files/";
export const sentences = ["This is an example sentence", "Each sentence is converted"];

/**
 * 将文本内容转换为向量
 * @param {string[]} contents - 要转换的文本内容数组
 * @returns {Promise<number[][]>} 返回向量数组的Promise
 */
export const embedder = async (contents) => {
    const extractor = await pipeline(
        "feature-extraction",
        "Xenova/all-MiniLM-L6-v2"
    );

    // 计算向量
    const output = await extractor(contents, {
      pooling: "mean",
      normalize: true,
    });
    return output.tolist();
};

codebase 的嵌入与存储

代码块的分块策略

在将代码块嵌入到向量数据库之前，需要对代码进行分块。常见的代码块分块策略包括：

• 按行分块：将代码按行分割，每一行作为一个代码块。
• 按函数分块：将代码按函数分割，每个函数作为一个代码块。
• 按语句分块：将代码按语句分割，每个语句作为一个代码块。
• 按文件分块：将代码按文件分割，每个文件作为一个代码块。

将分块后的代码数据存储到向量数据库中

上一篇文章中，我们已经实现了按行分块的功能。这里我们将内容分块的内容保存到数据库中

在 src/core/codeSnippetsIndex.js 中添加下面的代码

import sqlite3 from "sqlite3";
import { getDB } from ".";

export class CodeSnippetsCodebaseIndex {
  relativeExpectedTime = 1;
  artifactId = "codeSnippets";

  static async _createTables(db) {
    await db.exec(`CREATE TABLE IF NOT EXISTS code_snippets (
        id INTEGER PRIMARY KEY,
        path TEXT NOT NULL,
        content TEXT NOT NULL,
        signature TEXT,
        startLine INTEGER NOT NULL,
        endLine INTEGER NOT NULL
    )`);
  }
  
  static async add(datas) {
    const db = await getDB();
    await CodeSnippetsCodebaseIndex._createTables(db);
    for (const data of datas) {
      const { path, signature, content, startLine, endLine } = data;
      db.run(
      `INSERT INTO code_snippets (path, content, signature, startLine, endLine)
      VALUES (?, ?, ?, ?, ?)`,
      [path, signature, content, startLine, endLine],
    )}
  }

  static async getAll(signature) {
    const db = await getDB();
    await CodeSnippetsCodebaseIndex._createTables(db);
    try {
      const rows = await db.all(
        `SELECT cs.id, cs.path, cs.title
        FROM code_snippets cs
        WHERE cs.signature = ?;
        `,
        signature
      );

      // 检查 rows 是否为数组，避免调用 map 方法时出错
      if (Array.isArray(rows)) {
        return rows.map((row) => ({
          title: row.title,
          path: row.path,
          id: row.id,
        }));
      }
    } catch (error) {
      console.error("Error fetching code snippets:", error);
    }
    return [];
  }
}

示例代码展示数据嵌入和存储的过程：

将整个流程整合到一个函数中，实现代码块的嵌入和存储。

import CodeSnapped from "./treesitter.js";
import { CodeSnippetsCodebaseIndex } from "./CodeSnippetsIndex.js";

const fileDir = "./src";
const CodeSnappeder = new CodeSnapped();
const result = await CodeSnappeder.chunkCode(fileDir);

// 添加代码片段到数据库中，并包向量的结果保存到 LanceDB
await CodeSnippetsCodebaseIndex.add(result);

向量数据库的查询与检索

// 从数据库中检索相似代码片段
const query = "const";
const topK = 5;
const results = await CodeSnippetsCodebaseIndex.getAll(query, topK);
console.log(results);

相关文章

• 一、从 0 开始构建一个代码库-codebase 介绍
• 二、从 0 开始构建一个代码库-技术选型
• 三、从 0 开始构建一个代码库-代码分块技术实践
• 四、从 0 开始构建一个代码库-向量数据库的选择与集成
• 五、从 0 开始构建一个代码库-四：RAG 与 codebase 的深度融合
• 六、从 0 开始构建一个代码库-codebase 的测试与验证
• 七、从 0 开始构建一个代码库-Trae 编辑器中 Builder 模式的具体实现
• 八、从 0 开始构建一个代码库-mcp 实践