向量数据库使用轻松入门

Vector Database

在本文中,我们将学习如何使用Pinecone和OpenAI嵌入构建一个简单的向量数据库,你会发现这比想象中容易得多。

我们将通过使用Pinecone和OpenAI嵌入构建一个简单脚本,来聊聊向量数据库。这个脚本可以让你搜索一组我们姑且称之为 “食谱” 的数据。我们将探究向量数据库与你可能熟知且或爱或恨的传统SQL风格表有何不同,学习如何使用OpenAI生成嵌入以捕捉文本的 “含义”,然后将这些内容整合到一个简单的TypeScript项目中——因为我实在不想费神去研究Python中虚拟环境(virtualenvs)的工作原理。到最后,你将能够轻松、优雅地存储和语义查询各种食谱,或者任何你感兴趣的其他内容。

什么是向量数据库?

如果接下来我们要花些时间学习如何使用向量数据库,那么先花一两分钟快速回顾一下它们是什么以及为什么它们可能有用,这是很有必要的。

向量数据库是一种专门用于存储和搜索数据的高维表示(向量)的系统,而非传统的行和列形式。当人工智能模型(如大语言模型)将文本或图像转换为捕捉语义的数值嵌入时,这些嵌入可以高效地存储在向量数据库中。这使得能够进行 “相似性搜索”,即根据两个向量在高维空间中的接近程度来检索最相关的信息,而不是依赖于精确的关键词匹配。

在实际应用中,向量数据库在以下任务中表现出色:

  • 语义搜索:根据含义而非精确关键词,将用户查询与相似文档进行匹配。(剧透预警:这大致就是我们今天要实现的内容。)
  • 推荐系统:通过测量向量空间中的接近程度,找到具有相似特征的项目。
  • 问答系统:通过比较向量表示,从知识库中定位最相关的信息。

在处理大语言模型(LLMs,年轻人喜欢这么称呼)时,向量数据库非常有用,因为它们可以存储和比较由人工智能生成的、捕捉更深刻语义的向量 “嵌入”。例如,如果你要在大量法律文档上构建一个问答系统,数据库可以通过比较嵌入,快速返回相似的段落,即使它们没有完全相同的单词。

在推荐系统中,大语言模型可以嵌入产品描述,向量数据库则通过寻找具有重叠属性的项目,来推荐有意义的替代产品。此外,对于聊天机器人应用程序,你可以以向量形式存储知识库,并按需检索最相关的上下文,从而实现更准确、更具上下文感知的回复。

我们将使用Pinecone,据说它是向量数据库的热门选择。未来,我们也会探索其他一些替代方案。(我在看你呢,LanceDB。)

环境搭建

首先,我们定义一些要使用的简单类型。这个小示例将把一组食谱添加到我们的向量数据库中。稍后,我们将使用向量数据库搜索与给定查询相关的文档。由于食谱数据在recipes.json中是硬编码的,我就偷懒通过这个JSON对象的结构来推断类型。

import recipes from './recipes.json';
type Recipe = (typeof recipes)[number];

接下来,我们引入Pinecone和OpenAI的软件开发工具包(SDK)。此外,我们将使用dotenv.env文件中读取API密钥,并将其作为环境变量。你需要自己创建这个.env文件,因为它在.gitignore中被忽略了。毕竟,我可不想公开暴露我的API密钥,让大家替我产生一系列费用。

import 'dotenv/config';

import { Pinecone } from '@pinecone - database/pinecone';
import { OpenAI } from 'openai';

在继续之前,我们要确保环境变量配置正确。如果不正确,我们的代码会故意报错,但至少我们能确切知道原因。

const { PINECONE_API_KEY, OPEN_AI_API_KEY } = process.env;

if (!PINECONE_API_KEY) throw new Error('Pinecone API key is required');
if (!OPEN_AI_API_KEY) throw new Error('OpenAI API key is required');

我们将创建一个VectorDatabase类来封装相关逻辑。我故意把它设计得很简单,你可以把完善它当作课后作业。我们先从构造函数和一些属性开始。完成这些后,我会添加一些其他方法。为了确保一切正常,我们要保证这些API密钥已定义,如果没有则报错。

export class VectorDatabase {
    private pinecone: Pinecone;
    private openai: OpenAI;

    private readonly indexName ='recipes';
    private readonly dimension = 1536; // OpenAI的ada - 002嵌入维度
    private readonly metric = 'cosine'; // OpenAI的ada - 002嵌入度量

    constructor() {
        /** 实例化Pinecone SDK。 */
        this.pinecone = new Pinecone({ apiKey: PINECONE_API_KEY! });
        /** 实例化OpenAI SDK。 */
        this.openai = new OpenAI({ apiKey: OPEN_AI_API_KEY! });
    }

    //... 后续还有更多内容...
}

你可以把索引想象成传统数据库中的表。实际上,你也可以把它看作是数据库本身的一个实例。请别给我发消息争论这个。不管怎样,我们要确保 “recipes” 索引存在。如果存在,我们希望获取它的引用;如果不存在,请帮我们创建一个。

我要给自己一些实用的小工具函数:

  • this.#indexExists会给我Pinecone中当前所有索引的列表,并检查this.indexName是否在该列表中。
  • this.getIndex()将返回由this.indexName确定的Pinecone索引的引用。如果它还不存在,会为我创建一个。
/**
 * 验证索引是否存在于Pinecone中。
 */
get #indexExists() {
    return this.pinecone.listIndexes().then(({ indexes }) => {
        if (!indexes) return false;
        return indexes.some((index) => index.name === this.indexName);
    });
}

/**
 * 获取Pinecone中索引的引用。如果索引不存在,将创建它。
 */
async getIndex() {
    if (await this.#indexExists) return this.pinecone.Index(this.indexName);

    await this.pinecone.createIndex({
        name: this.indexName,
        dimension: this.dimension,
        metric: this.metric,
        spec: {
            serverless: {
                cloud: 'aws',
                region: 'us - east - 1',
            },
        },
    });

    return this.pinecone.Index(this.indexName);
}

如果你想使用不同的区域或云服务提供商,欢迎调整spec,但对于这个简单示例实现来说,这些是合理的默认值。

接下来,我们初始化与Pinecone的连接。

创建向量嵌入

我们将添加三个方法:

  • generateEmbedding:此方法将使用OpenAI根据文档的文本内容创建一个嵌入。
  • indexDocument:此方法将调用generateEmbedding,然后将文档添加到我们的向量数据库中。
  • semanticSearch:此方法将在我们的向量数据库中搜索与查询相似的内容。

第一个方法generateEmbedding将使用OpenAI把给定的一段文本转换为一个向量,也就是一个数字数组。这相当简单直接。我们指定要用于创建嵌入的模型,然后从响应中提取我们需要的数据。

export class VectorDatabase {
    //... 之前的代码...

    /**
     * 使用OpenAI的API生成嵌入。
     * @returns 文本的向量表示。
     */
    async generateEmbedding(text: string): Promise<number[]> {
        const response = await this.openai.embeddings.create({
            model: 'text - embedding - ada - 002',
            input: text,
        });

        return response.data[0].embedding;
    }
}

我们将使用this.generateEmbedding来存储文档,并且当有人尝试使用字符串查询数据库时,也用它来创建向量。接下来,看看this.insertDocument

export class VectorDatabase {
    //... 之前的代码...

    async indexDocument(document: Recipe) {
        const index = await this.getIndex();
        const embedding = await this.generateEmbedding(document.content);

        await index.upsert([
            {
                id: document.id,
                values: embedding,
                metadata: {
                    title: document.title,
                    content: document.content,
                },
            },
        ]);
    }
}

如你所见,我们既包含了文档(在这个例子中是食谱)的嵌入,也包含了关于该文档来源文件的一些元数据。这些元数据可以是任何能让你更容易将向量嵌入与原始源材料联系起来的信息。

最后,我们需要弄清楚如何搜索向量数据库,以获取与给定查询相关的内容。在这个例子中,由于我们的数据集不是特别大,我们默认返回前三个匹配项。

export class VectorDatabase {
    //... 之前的代码...

    /**
     * 在向量数据库中搜索与查询匹配的内容。
     */
    async semanticSearch(
        /**
         * 一个将被转换为嵌入并用于查询向量数据库的字符串。
         */
        query: string,
        /** 要返回的结果数量。最大值:10000。 */
        topK: number = 3,
    ) {
        // 为搜索查询生成嵌入。
        const vector = await this.generateEmbedding(query);
        const index = await this.getIndex();

        // 搜索相似向量
        const searchResults = await index.query({
            vector,
            topK,
            includeMetadata: true,
        });

        return searchResults.matches.map((match) => ({
            id: match.id,
            title: match.metadata?.title,
            content: match.metadata?.content.toString().slice(0, 50) + '…',
            score: match.score,
        }));
    }
}

尝试运行

现在一切都已设置好,我们可以将数据加载到Pinecone中并进行查询。注释掉你不需要的部分,并随意调整查询以查看结果的变化。

const database = new VectorDatabase();

// 如果你已经将食谱存储在数据库中,请注释掉这部分。
for (const recipe of recipes) {
    console.log(chalk.blue('Indexing recipe:'), recipe.title);
    await database.indexDocument(recipe);
}

const searchResults = await database.semanticSearch('recipes with ice cream');

console.table(searchResults);

就这样,我们成功地使用Pinecone和OpenAI嵌入拼凑出一个简单的向量数据库,用于存储和查询食谱数据。你不一定非要使用OpenAI来创建嵌入。唯一的规则是,用于为你的数据创建向量嵌入的模型,必须与用于为用户提供的查询创建向量嵌入的模型相同,不能混用。

作为练习,你可以使用更大的数据集(例如,你最喜欢的开源项目的文档、你过去五年的私人日记、你收藏并打算阅读的所有博客文章等)。

Publish on 2025-01-04,Update on 2025-02-10