使用Redis作为OpenAI的向量数据库
本笔记本介绍了如何使用Redis作为OpenAI嵌入向量数据库的入门知识。Redis是一个可扩展的实时数据库，当使用RediSearch模块时，可以用作向量数据库。RediSearch模块允许您在Redis中索引和搜索向量。本笔记本将向您展示如何使用RediSearch模块来索引和搜索通过使用OpenAI API创建并存储在Redis中的向量。
什么是Redis？
大多数来自Web服务背景的开发人员可能对Redis很熟悉。在其核心，Redis是一个开源的键值存储，可用作缓存、消息代理和数据库。开发人员选择Redis是因为它快速，拥有大量的客户端库生态系统，并且多年来已被主要企业部署使用。
除了Redis的传统用途外，Redis还提供了Redis模块，这是一种通过新的数据类型和命令扩展Redis的方式。示例模块包括RedisJSON、RedisTimeSeries、RedisBloom和RediSearch。
什么是RediSearch？
RediSearch是一个Redis模块，为Redis提供了查询、二级索引、全文搜索和向量搜索功能。要使用RediSearch，首先在Redis数据上声明索引。然后可以使用RediSearch客户端查询该数据。有关RediSearch功能集的更多信息，请参阅README或RediSearch文档。
部署选项
有许多部署Redis的方式。对于本地开发，最快速的方法是使用Redis Stack docker容器，我们将在这里使用。Redis Stack包含许多Redis模块，可以一起使用以创建快速的多模型数据存储和查询引擎。
对于生产用例，最简单的方法是使用Redis Cloud服务开始。Redis Cloud是一个完全托管的Redis服务。您还可以使用Redis Enterprise在自己的基础设施上部署Redis。Redis Enterprise是一个完全托管的Redis服务，可以在kubernetes、本地或云中部署。
此外，每个主要的云提供商（AWS Marketplace、Google Marketplace或Azure Marketplace）都提供Redis Enterprise的市场提供。
先决条件
在开始这个项目之前，我们需要设置以下内容：
- 启动带有 RediSearch 的 Redis 数据库（redis-stack）
- 安装库
- 获取您的 OpenAI API 密钥
===========================================================
启动 Redis
为了保持这个示例简单，我们将使用 Redis Stack docker 容器，可以按照以下方式启动：
$ docker-compose up -d
这还包括用于管理您的 Redis 数据库的 RedisInsight GUI，您可以在启动 docker 容器后通过 http://localhost:8001 查看它。
您已经设置好并且准备就绪！接下来，我们将导入并创建用于与我们刚刚创建的 Redis 数据库通信的客户端。
安装要求
Redis-Py是用于与Redis通信的Python客户端。我们将使用它来与我们的Redis堆栈数据库进行通信。
! pip install redis wget pandas openai
=========================================================== ## 准备你的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
import openai
# 注意：或者，您也可以像这样设置一个临时的环境变量：
# os.environ["OPENAI_API_KEY"] = 'sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx'
if os.getenv("OPENAI_API_KEY") is not None:
openai.api_key = os.getenv("OPENAI_API_KEY")
print ("OPENAI_API_KEY is ready")
else:
print ("OPENAI_API_KEY environment variable not found")
OPENAI_API_KEY is ready
加载数据
在这一部分，我们将加载已经转换为向量的嵌入数据。我们将使用这些数据在Redis中创建一个索引，然后搜索相似的向量。
import sys
import numpy as np
import pandas as pd
from typing import List
# 使用 `nbutils.py` 中的辅助函数下载并读取数据
# 这大约需要5到10分钟来运行。
if os.getcwd() not in sys.path:
sys.path.append(os.getcwd())
import nbutils
nbutils.download_wikipedia_data()
data = nbutils.read_wikipedia_data()
data.head()
File Downloaded
| id | url | title | text | title_vector | content_vector | vector_id | |
|---|---|---|---|---|---|---|---|
| 0 | 1 | https://simple.wikipedia.org/wiki/April | April | April is the fourth month of the year in the J... | [0.001009464613161981, -0.020700545981526375, ... | [-0.011253940872848034, -0.013491976074874401,... | 0 |
| 1 | 2 | https://simple.wikipedia.org/wiki/August | August | August (Aug.) is the eighth month of the year ... | [0.0009286514250561595, 0.000820168002974242, ... | [0.0003609954728744924, 0.007262262050062418, ... | 1 |
| 2 | 6 | https://simple.wikipedia.org/wiki/Art | Art | Art is a creative activity that expresses imag... | [0.003393713850528002, 0.0061537534929811954, ... | [-0.004959689453244209, 0.015772193670272827, ... | 2 |
| 3 | 8 | https://simple.wikipedia.org/wiki/A | A | A or a is the first letter of the English alph... | [0.0153952119871974, -0.013759135268628597, 0.... | [0.024894846603274345, -0.022186409682035446, ... | 3 |
| 4 | 9 | https://simple.wikipedia.org/wiki/Air | Air | Air refers to the Earth's atmosphere. Air is a... | [0.02224554680287838, -0.02044147066771984, -0... | [0.021524671465158463, 0.018522677943110466, -... | 4 |
连接到Redis
现在我们的Redis数据库正在运行，我们可以使用Redis-py客户端连接到它。我们将使用Redis数据库的默认主机和端口，即localhost:6379。
import redis
from redis.commands.search.indexDefinition import (
IndexDefinition,
IndexType
)
from redis.commands.search.query import Query
from redis.commands.search.field import (
TextField,
VectorField
)
REDIS_HOST = "localhost"
REDIS_PORT = 6379
REDIS_PASSWORD = "" # 无密码默认设置的Redis
# 连接到 Redis
redis_client = redis.Redis(
host=REDIS_HOST,
port=REDIS_PORT,
password=REDIS_PASSWORD
)
redis_client.ping()
True
在Redis中创建搜索索引
下面的单元格将展示如何在Redis中指定和创建一个搜索索引。我们将：
- 设置一些常量来定义我们的索引，比如距离度量和索引名称
- 使用RediSearch字段定义索引模式
- 创建索引
# 常量
VECTOR_DIM = len(data['title_vector'][0]) # 向量的长度
VECTOR_NUMBER = len(data) # 初始向量数量
INDEX_NAME = "embeddings-index" # 搜索索引的名称
PREFIX = "doc" # 文档键的前缀
DISTANCE_METRIC = "COSINE" # 向量间的距离度量方法（例如：余弦距离、内积、欧氏距离L2）
# 为数据集中的每一列定义RediSearch字段
title = TextField(name="title")
url = TextField(name="url")
text = TextField(name="text")
title_embedding = VectorField("title_vector",
"FLAT", {
"TYPE": "FLOAT32",
"DIM": VECTOR_DIM,
"DISTANCE_METRIC": DISTANCE_METRIC,
"INITIAL_CAP": VECTOR_NUMBER,
}
)
text_embedding = VectorField("content_vector",
"FLAT", {
"TYPE": "FLOAT32",
"DIM": VECTOR_DIM,
"DISTANCE_METRIC": DISTANCE_METRIC,
"INITIAL_CAP": VECTOR_NUMBER,
}
)
fields = [title, url, text, title_embedding, text_embedding]
# 检查索引是否存在
try:
redis_client.ft(INDEX_NAME).info()
print("Index already exists")
except:
# 创建RediSearch索引
redis_client.ft(INDEX_NAME).create_index(
fields = fields,
definition = IndexDefinition(prefix=[PREFIX], index_type=IndexType.HASH)
)
将文档加载到索引中
现在我们有了一个搜索索引，我们可以将文档加载到其中。我们将使用之前示例中使用的相同文档。在Redis中，可以使用HASH或JSON（如果除了RediSearch还使用RedisJSON）数据类型来存储文档。在本示例中，我们将 使用HASH数据类型。下面的单元格将展示如何将文档加载到索引中。
def index_documents(client: redis.Redis, prefix: str, documents: pd.DataFrame):
records = documents.to_dict("records")
for doc in records:
key = f"{prefix}:{str(doc['id'])}"
# 为标题和内容创建字节向量
title_embedding = np.array(doc["title_vector"], dtype=np.float32).tobytes()
content_embedding = np.array(doc["content_vector"], dtype=np.float32).tobytes()
# 将浮点数列表替换为字节向量
doc["title_vector"] = title_embedding
doc["content_vector"] = content_embedding
client.hset(key, mapping = doc)
index_documents(redis_client, PREFIX, data)
print(f"Loaded {redis_client.info()['db0']['keys']} documents in Redis search index with name: {INDEX_NAME}")
Loaded 25000 documents in Redis search index with name: embeddings-index
使用OpenAI查询嵌入进行简单的向量搜索查询
现在我们已经有了一个搜索索引并加载了文档，我们可以运行搜索查询。下面我们将提供一个函数，该函数将运行一个搜索查询并返回结果。使用这个函数，我们运行一些查询，展示如何利用Redis作为向量数据库。
def search_redis(
redis_client: redis.Redis,
user_query: str,
index_name: str = "embeddings-index",
vector_field: str = "title_vector",
return_fields: list = ["title", "url", "text", "vector_score"],
hybrid_fields = "*",
k: int = 20,
print_results: bool = True,
) -> List[dict]:
# 从用户查询生成嵌入向量
embedded_query = openai.Embedding.create(input=user_query,
model="text-embedding-3-small",
)["data"][0]['embedding']
# 准备查询
base_query = f'{hybrid_fields}=>[KNN {k} @{vector_field} $vector AS vector_score]'
query = (
Query(base_query)
.return_fields(*return_fields)
.sort_by("vector_score")
.paging(0, k)
.dialect(2)
)
params_dict = {"vector": np.array(embedded_query).astype(dtype=np.float32).tobytes()}
# 执行向量搜索
results = redis_client.ft(index_name).search(query, params_dict)
if print_results:
for i, article in enumerate(results.docs):
score = 1 - float(article.vector_score)
print(f"{i}. {article.title} (Score: {round(score ,3) })")
return results.docs
# 使用OpenAI生成查询嵌入
results = search_redis(redis_client, 'modern art in Europe', k=10)
0. Museum of Modern Art (Score: 0.875)
1. Western Europe (Score: 0.868)
2. Renaissance art (Score: 0.864)
3. Pop art (Score: 0.86)
4. Northern Europe (Score: 0.855)
5. Hellenistic art (Score: 0.853)
6. Modernist literature (Score: 0.847)
7. Art film (Score: 0.843)
8. Central Europe (Score: 0.843)
9. European (Score: 0.841)
results = search_redis(redis_client, 'Famous battles in Scottish history', vector_field='content_vector', k=10)
0. Battle of Bannockburn (Score: 0.869)
1. Wars of Scottish Independence (Score: 0.861)
2. 1651 (Score: 0.853)
3. First War of Scottish Independence (Score: 0.85)
4. Robert I of Scotland (Score: 0.846)
5. 841 (Score: 0.844)
6. 1716 (Score: 0.844)
7. 1314 (Score: 0.837)
8. 1263 (Score: 0.836)
9. William Wallace (Score: 0.835)
使用Redis进行混合查询
前面的示例展示了如何在RediSearch中运行向量搜索查询。在本节中，我们将展示如何将向量搜索与其他RediSearch字段结合起来进行混合搜索。在下面的示例中，我们将结合向量搜索和全文搜索。
def create_hybrid_field(field_name: str, value: str) -> str:
return f'@{field_name}:"{value}"'
# 在内容向量 中搜索关于苏格兰历史上著名战役的文章，并仅包含标题中带有“Scottish”的结果。
results = search_redis(redis_client,
"Famous battles in Scottish history",
vector_field="title_vector",
k=5,
hybrid_fields=create_hybrid_field("title", "Scottish")
)
0. First War of Scottish Independence (Score: 0.892)
1. Wars of Scottish Independence (Score: 0.889)
2. Second War of Scottish Independence (Score: 0.879)
3. List of Scottish monarchs (Score: 0.873)
4. Scottish Borders (Score: 0.863)
# run a hybrid query for articles about Art in the title vector and only include results with the phrase "Leonardo da Vinci" in the text
results = search_redis(redis_client,
"Art",
vector_field="title_vector",
k=5,
hybrid_fields=create_hybrid_field("text", "Leonardo da Vinci")
)
# 在我们的全文搜索查询返回的文本中，找到了关于莱昂纳多·达·芬奇的具体提及。
mention = [sentence for sentence in results[0].text.split("\n") if "Leonardo da Vinci" in sentence][0]
mention
0. Art (Score: 1.0)
1. Paint (Score: 0.896)
2. Renaissance art (Score: 0.88)
3. Painting (Score: 0.874)
4. Renaissance (Score: 0.846)
'In Europe, after the Middle Ages, there was a "Renaissance" which means "rebirth". People rediscovered science and artists were allowed to paint subjects other than religious subjects. People like Michelangelo and Leonardo da Vinci still painted religious pictures, but they also now could paint mythological pictures too. These artists also invented perspective where things in the distance look smaller in the picture. This was new because in the Middle Ages people would paint all the figures close up and just overlapping each other. These artists used nudity regularly in their art.'
HNSW索引
到目前为止，我们一直在使用FLAT或“暴力”索引来运行我们的查询。Redis还支持HNSW索引，这是一种快速的、近似的索引。HNSW索引是一种基于图的索引，它使用分层可导航小世界图来存储 向量。HNSW索引是在大型数据集中运行近似查询时的一个很好的选择。
对于大多数情况，HNSW构建时间更长，消耗的内存更多，但在运行查询时会更快，特别是对于大型数据集。
接下来的单元格将展示如何使用与之前相同的数据创建HNSW索引并运行查询。
# 重新定义 RediSearch 向量字段以使用 HNSW 索引
title_embedding = VectorField("title_vector",
"HNSW", {
"TYPE": "FLOAT32",
"DIM": VECTOR_DIM,
"DISTANCE_METRIC": DISTANCE_METRIC,
"INITIAL_CAP": VECTOR_NUMBER
}
)
text_embedding = VectorField("content_vector",
"HNSW", {
"TYPE": "FLOAT32",
"DIM": VECTOR_DIM,
"DISTANCE_METRIC": DISTANCE_METRIC,
"INITIAL_CAP": VECTOR_NUMBER
}
)
fields = [title, url, text, title_embedding, text_embedding]
import time
# 检查索引是否存在
HNSW_INDEX_NAME = INDEX_NAME+ "_HNSW"
try:
redis_client.ft(HNSW_INDEX_NAME).info()
print("Index already exists")
except:
# 创建RediSearch索引
redis_client.ft(HNSW_INDEX_NAME).create_index(
fields = fields,
definition = IndexDefinition(prefix=[PREFIX], index_type=IndexType.HASH)
)
# 由于RediSearch会为现有文档在后台创建索引，我们将等待直至
# 在执行我们的查询之前，索引工作已经完成。尽管这对于第一个查询来说并非必要，
# 如果索引尚未完全构建，某些查询可能需要更长时间才能运行。通常情况下，Redis会执行
# 在现有索引中添加新文档时效果最佳，而不是在现有文档上创建新索引。
while redis_client.ft(HNSW_INDEX_NAME).info()["indexing"] == "1":
time.sleep(5)
results = search_redis(redis_client, 'modern art in Europe', index_name=HNSW_INDEX_NAME, k=10)
0. Western Europe (Score: 0.868)
1. Northern Europe (Score: 0.855)
2. Central Europe (Score: 0.843)
3. European (Score: 0.841)
4. Eastern Europe (Score: 0.839)
5. Europe (Score: 0.839)
6. Western European Union (Score: 0.837)
7. Southern Europe (Score: 0.831)
8. Western civilization (Score: 0.83)
9. Council of Europe (Score: 0.827)
# 将HNSW索引的结果与FLAT索引的结果进行比较，并记录两者的查询时间。
def time_queries(iterations: int = 10):
print(" ----- Flat Index ----- ")
t0 = time.time()
for i in range(iterations):
results_flat = search_redis(redis_client, 'modern art in Europe', k=10, print_results=False)
t0 = (time.time() - t0) / iterations
results_flat = search_redis(redis_client, 'modern art in Europe', k=10, print_results=True)
print(f"Flat index query time: {round(t0, 3)} seconds\n")
time.sleep(1)
print(" ----- HNSW Index ------ ")
t1 = time.time()
for i in range(iterations):
results_hnsw = search_redis(redis_client, 'modern art in Europe', index_name=HNSW_INDEX_NAME, k=10, print_results=False)
t1 = (time.time() - t1) / iterations
results_hnsw = search_redis(redis_client, 'modern art in Europe', index_name=HNSW_INDEX_NAME, k=10, print_results=True)
print(f"HNSW index query time: {round(t1, 3)} seconds")
print(" ------------------------ ")
time_queries()
----- Flat Index -----
0. Museum of Modern Art (Score: 0.875)
1. Western Europe (Score: 0.867)
2. Renaissance art (Score: 0.864)
3. Pop art (Score: 0.861)
4. Northern Europe (Score: 0.855)
5. Hellenistic art (Score: 0.853)
6. Modernist literature (Score: 0.847)
7. Art film (Score: 0.843)
8. Central Europe (Score: 0.843)
9. Art (Score: 0.842)
Flat index query time: 0.263 seconds
----- HNSW Index ------
0. Western Europe (Score: 0.867)
1. Northern Europe (Score: 0.855)
2. Central Europe (Score: 0.843)
3. European (Score: 0.841)
4. Eastern Europe (Score: 0.839)
5. Europe (Score: 0.839)
6. Western European Union (Score: 0.837)
7. Southern Europe (Score: 0.831)
8. Western civilization (Score: 0.83)
9. Council of Europe (Score: 0.827)
HNSW index query time: 0.129 seconds
------------------------