使用 Redis 和 OpenAI 作为向量数据库

本笔记本介绍了如何将 Redis 与 OpenAI 嵌入结合使用作为向量数据库。Redis 是一个可扩展的实时数据库，在使用 RediSearch 模块 时可用作向量数据库。RediSearch 模块允许您在 Redis 中索引和搜索向量。本笔记本将向您展示如何使用 RediSearch 模块来索引和搜索使用 OpenAI API 创建并存储在 Redis 中的向量。

什么是 Redis？

大多数来自 Web 服务背景的开发人员可能都熟悉 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）都在其 marketplace 中提供了 Redis Enterprise。

先决条件

在我们开始这个项目之前，我们需要进行以下设置：

• 使用 RediSearch (redis-stack) 启动 Redis 数据库
• 安装库
  • Redis-py
• 获取您的 OpenAI API 密钥

===========================================================

启动 Redis

为了使此示例保持简单，我们将使用 Redis Stack docker 容器，我们可以如下启动它：

$ docker-compose up -d

这还包括用于管理 Redis 数据库的 RedisInsight GUI，您可以在启动 docker 容器后在 http://localhost:8001 上查看它。

您已全部设置完毕，可以开始使用了！接下来，我们将导入并创建与我们刚刚创建的 Redis 数据库通信的客户端。

安装要求

Redis-Py 是用于与 Redis 通信的 Python 客户端。我们将使用它与我们的 Redis-stack 数据库进行通信。

! 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

连接到 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 中的搜索索引。我们将：

1. 设置一些用于定义索引的常量，例如距离度量和索引名称
2. 使用 RediSearch 字段定义索引架构
3. 创建索引

# 常量
VECTOR_DIM = len(data['title_vector'][0]) # 向量的长度
VECTOR_NUMBER = len(data)                 # 初始向量数量
INDEX_NAME = "embeddings-index"           # 搜索索引的名称
PREFIX = "doc"                            # 文档键的前缀
DISTANCE_METRIC = "COSINE"                # 向量的距离度量（例如 COSINE、IP、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]:

    # 使用 OpenAI 为用户查询创建嵌入向量
    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}"'

# 搜索关于苏格兰历史著名战役的内容向量，并仅包含标题中包含苏格兰的结果
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)

# 对标题向量中关于艺术的文章运行混合查询，并仅包含文本中包含“莱昂纳多·达·芬奇”一词的结果
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 在大多数情况下比 FLAT 构建时间更长，消耗的内存也更多，但查询速度更快，尤其是在大型数据集上。

以下单元格将展示如何创建 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 为现有文档在后台创建索引，我们将等待索引完成后再运行查询。
# 虽然第一次查询不需要这样做，但某些查询可能需要更长时间才能运行，如果索引未完全构建。
# 总的来说，将新文档添加到现有索引比在现有文档上创建新索引效果更好。
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
 ------------------------