使用 Qdrant 作为 OpenAI embedding 的向量数据库

本笔记本将指导您逐步使用 Qdrant 作为 OpenAI embedding 的向量数据库。Qdrant 是一个用 Rust 编写的高性能向量搜索数据库。它提供 RESTful 和 gRPC API 来管理您的 embedding。有一个官方的 Python qdrant-client，可以简化与应用程序的集成。

本笔记本介绍了端到端流程：

1. 使用 OpenAI API 创建的预计算 embedding。
2. 将 embedding 存储在本地 Qdrant 实例中。
3. 使用 OpenAI API 将原始文本查询转换为 embedding。
4. 使用 Qdrant 在创建的 collection 中执行最近邻搜索。

什么是 Qdrant

Qdrant 是一个开源向量数据库，允许存储神经 embedding 以及元数据，也称为 payload。Payload 不仅可以用于保留特定点的其他属性，还可以用于过滤。 Qdrant 提供了一种独特的过滤机制，该机制内置于向量搜索阶段，这使其非常高效。

部署选项

Qdrant 可以根据应用程序的目标负载以多种方式启动：

• 使用 Docker 容器在本地或本地部署
• 使用 Helm chart 在 Kubernetes 集群上部署
• 使用 Qdrant Cloud

集成

Qdrant 同时提供 RESTful 和 gRPC API，无论您使用哪种编程语言，都可以轻松集成。但是，有一些官方的客户端可用于最流行的语言，如果您使用 Python，那么 Python Qdrant 客户端库 可能是最佳选择。

先决条件

为了进行此练习，我们需要准备一些东西：

1. Qdrant 服务器实例。在我们的例子中是本地 Docker 容器。
2. qdrant-client 库，用于与向量数据库进行交互。
3. OpenAI API 密钥。

启动 Qdrant 服务器

我们将使用在 Docker 容器中运行的本地 Qdrant 实例。启动它的最简单方法是使用附加的 [docker-compose.yaml] 文件并运行以下命令：

! docker compose up -d

 [1A [1B [0G [?25l[+] Running 1/0
  [32m✔ [0m Container qdrant-qdrant-1   [32mRunning [0m                                       [34m0.0s  [0m
 [?25h

我们可以通过运行一个简单的 curl 命令来验证服务器是否已成功启动：

! curl http://localhost:6333

{"title":"qdrant - vector search engine","version":"1.3.0"}

安装依赖

本笔记本显然需要 openai 和 qdrant-client 包，但我们还将使用一些其他附加库。以下命令将它们全部安装：

! pip install openai qdrant-client pandas wget

准备您的 OpenAI API 密钥

OpenAI API 密钥用于文档和查询的向量化。

如果您没有 OpenAI API 密钥，可以从 https://beta.openai.com/account/api-keys 获取。

获得密钥后，请通过运行以下命令将其添加为环境变量 OPENAI_API_KEY：

! export OPENAI_API_KEY="your API key"

# 测试您的 OpenAI API 密钥是否已正确设置为环境变量
# 注意：如果您在本地运行此笔记本，则需要重新加载终端和笔记本才能使环境变量生效。
import os

# 注意：或者，您可以像这样设置一个临时环境变量：
# os.environ["OPENAI_API_KEY"] = "sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"

if os.getenv("OPENAI_API_KEY") is not None:
    print("OPENAI_API_KEY 已准备就绪")
else:
    print("未找到 OPENAI_API_KEY 环境变量")

OPENAI_API_KEY 已准备就绪

连接到 Qdrant

使用官方 Python 库可以轻松连接到正在运行的 Qdrant 服务器实例：

import qdrant_client

client = qdrant_client.QdrantClient(
    host="localhost",
    prefer_grpc=True,
)

我们可以通过运行任何可用方法来测试连接：

client.get_collections()

CollectionsResponse(collections=[])

加载数据

在本节中，我们将加载在此会话之前准备好的数据，这样您就不必花费自己的积分重新计算维基百科文章的 embedding 了。

import wget

embeddings_url = "https://cdn.openai.com/API/examples/data/vector_database_wikipedia_articles_embedded.zip"

# 该文件约为 700 MB，因此需要一些时间
wget.download(embeddings_url)

100% [......................................................................] 698933052 / 698933052




'vector_database_wikipedia_articles_embedded (9).zip'

下载的文件必须然后解压：

import zipfile

with zipfile.ZipFile("vector_database_wikipedia_articles_embedded.zip","r") as zip_ref:
    zip_ref.extractall("../data")

最后，我们可以从提供的 CSV 文件加载它：

import pandas as pd

from ast import literal_eval

article_df = pd.read_csv('../data/vector_database_wikipedia_articles_embedded.csv')
# 将字符串中的向量读回列表
article_df["title_vector"] = article_df.title_vector.apply(literal_eval)
article_df["content_vector"] = article_df.content_vector.apply(literal_eval)
article_df.head()

索引数据

Qdrant 将数据存储在__collections__中，其中每个对象至少由一个向量描述，并可能包含称为__payload__的附加元数据。我们的 collection 将被称为 Articles，每个对象将由 title 和 content 向量描述。Qdrant 不需要您预先设置任何类型的模式，因此您可以通过简单的设置自由地将点放入 collection 中。

我们将从创建一个 collection 开始，然后用我们预先计算的 embedding 填充它。

from qdrant_client.http import models as rest

vector_size = len(article_df["content_vector"][0])

client.create_collection(
    collection_name="Articles",
    vectors_config={
        "title": rest.VectorParams(
            distance=rest.Distance.COSINE,
            size=vector_size,
        ),
        "content": rest.VectorParams(
            distance=rest.Distance.COSINE,
            size=vector_size,
        ),
    }
)

True

client.upsert(
    collection_name="Articles",
    points=[
        rest.PointStruct(
            id=k,
            vector={
                "title": v["title_vector"],
                "content": v["content_vector"],
            },
            payload=v.to_dict(),
        )
        for k, v in article_df.iterrows()
    ],
)

UpdateResult(operation_id=0, status=<UpdateStatus.COMPLETED: 'completed'>)

# 检查 collection 大小以确保所有点都已存储
client.count(collection_name="Articles")

CountResult(count=25000)

搜索数据

将数据放入 Qdrant 后，我们将开始查询 collection 以查找最接近的向量。我们可以提供一个附加参数 vector_name 来从基于标题的搜索切换到基于内容的搜索。由于预计算的 embedding 是使用 text-embedding-ada-002 OpenAI 模型创建的，因此我们在搜索时也必须使用它。

from openai import OpenAI

openai_client = OpenAI()

def query_qdrant(query, collection_name, vector_name="title", top_k=20):
    # 从用户查询创建 embedding 向量
    embedded_query = openai_client.embeddings.create(
        input=query,
        model="text-embedding-ada-002",
    ).data[0].embedding

    query_results = client.search(
        collection_name=collection_name,
        query_vector=(
            vector_name, embedded_query
        ),
        limit=top_k,
    )

    return query_results

query_results = query_qdrant("modern art in Europe", "Articles")
for i, article in enumerate(query_results):
    print(f"{i + 1}. {article.payload['title']} (Score: {round(article.score, 3)})")

1. Museum of Modern Art (Score: 0.875)
2. Western Europe (Score: 0.867)
3. Renaissance art (Score: 0.864)
4. Pop art (Score: 0.86)
5. Northern Europe (Score: 0.855)
6. Hellenistic art (Score: 0.853)
7. Modernist literature (Score: 0.847)
8. Art film (Score: 0.843)
9. Central Europe (Score: 0.843)
10. European (Score: 0.841)
11. Art (Score: 0.841)
12. Byzantine art (Score: 0.841)
13. Postmodernism (Score: 0.84)
14. Eastern Europe (Score: 0.839)
15. Cubism (Score: 0.839)
16. Europe (Score: 0.839)
17. Impressionism (Score: 0.838)
18. Bauhaus (Score: 0.838)
19. Surrealism (Score: 0.837)
20. Expressionism (Score: 0.837)

# 这次我们将使用内容向量进行查询
query_results = query_qdrant("Famous battles in Scottish history", "Articles", "content")
for i, article in enumerate(query_results):
    print(f"{i + 1}. {article.payload['title']} (Score: {round(article.score, 3)})")

1. Battle of Bannockburn (Score: 0.869)
2. Wars of Scottish Independence (Score: 0.861)
3. 1651 (Score: 0.852)
4. First War of Scottish Independence (Score: 0.85)
5. Robert I of Scotland (Score: 0.846)
6. 841 (Score: 0.844)
7. 1716 (Score: 0.844)
8. 1314 (Score: 0.837)
9. 1263 (Score: 0.836)
10. William Wallace (Score: 0.835)
11. Stirling (Score: 0.831)
12. 1306 (Score: 0.831)
13. 1746 (Score: 0.83)
14. 1040s (Score: 0.828)
15. 1106 (Score: 0.827)
16. 1304 (Score: 0.826)
17. David II of Scotland (Score: 0.825)
18. Braveheart (Score: 0.824)
19. 1124 (Score: 0.824)
20. Second War of Scottish Independence (Score: 0.823)