使用 AnalyticDB 作为 OpenAI 嵌入的向量数据库

本笔记本将指导您逐步了解如何将 AnalyticDB 用作 OpenAI 嵌入的向量数据库。

本笔记本将展示一个端到端的流程：

1. 使用 OpenAI API 创建的预计算嵌入。
2. 将嵌入存储在 AnalyticDB 的云实例中。
3. 使用 OpenAI API 将原始文本查询转换为嵌入。
4. 使用 AnalyticDB 在创建的集合中执行最近邻搜索。

什么是 AnalyticDB

AnalyticDB 是一个高性能的分布式向量数据库。它完全兼容 PostgreSQL 语法，您可以轻松使用它。AnalyticDB 是阿里云提供的云原生数据库，拥有高性能的向量计算引擎。开箱即用的体验允许处理数十亿数据向量，并提供丰富的特性，包括索引算法、结构化和非结构化数据特性、实时更新、距离度量、标量过滤、时间旅行搜索等。它还配备了完整的 OLAP 数据库功能和针对生产环境使用的 SLA 承诺；

部署选项

• 使用 AnalyticDB 云向量数据库。点击此处可快速部署。

先决条件

为了完成此练习，我们需要准备几项内容：

1. AnalyticDB 云服务器实例。
2. 用于与向量数据库交互的 'psycopg2' 库。任何其他 postgresql 客户端库也可以。
3. 一个 OpenAI API 密钥。

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

安装要求

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

! pip install openai psycopg2 pandas wget

准备您的 OpenAI API 密钥

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

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

获取密钥后，请将其添加为环境变量 OPENAI_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 已准备就绪

连接到 AnalyticDB

首先将其添加到您的环境变量中。或者，您可以直接更改下面的 "psycopg2.connect" 参数

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

import os
import psycopg2

# 注意：或者，您可以像这样设置一个临时环境变量：
# os.environ["PGHOST"] = "your_host"
# os.environ["PGPORT"] "5432"),
# os.environ["PGDATABASE"] "postgres"),
# os.environ["PGUSER"] "user"),
# os.environ["PGPASSWORD"] "password"),

connection = psycopg2.connect(
    host=os.environ.get("PGHOST", "localhost"),
    port=os.environ.get("PGPORT", "5432"),
    database=os.environ.get("PGDATABASE", "postgres"),
    user=os.environ.get("PGUSER", "user"),
    password=os.environ.get("PGPASSWORD", "password")
)

# 创建一个新的游标对象
cursor = connection.cursor()

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

# 执行一个简单的查询来测试连接
cursor.execute("SELECT 1;")
result = cursor.fetchone()

# 检查查询结果
if result == (1,):
    print("连接成功！")
else:
    print("连接失败。")

连接成功！

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.zip'

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

import zipfile
import os
import re
import tempfile

current_directory = os.getcwd()
zip_file_path = os.path.join(current_directory, "vector_database_wikipedia_articles_embedded.zip")
output_directory = os.path.join(current_directory, "../../data")

with zipfile.ZipFile(zip_file_path, "r") as zip_ref:
    zip_ref.extractall(output_directory)


# 检查 csv 文件是否存在
file_name = "vector_database_wikipedia_articles_embedded.csv"
data_directory = os.path.join(current_directory, "../../data")
file_path = os.path.join(data_directory, file_name)


if os.path.exists(file_path):
    print(f"文件 {file_name} 存在于数据目录中。")
else:
    print(f"文件 {file_name} 不存在于数据目录中。")

文件 vector_database_wikipedia_articles_embedded.csv 存在于数据目录中。

索引数据

AnalyticDB 将数据存储在 关系 中，其中每个对象至少由一个向量描述。我们的关系将称为 articles，每个对象将由 title 和 content 向量描述。

我们将首先创建一个关系并在 title 和 content 上创建向量索引，然后用我们预计算的嵌入填充它。

create_table_sql = '''
CREATE TABLE IF NOT EXISTS public.articles (
    id INTEGER NOT NULL,
    url TEXT,
    title TEXT,
    content TEXT,
    title_vector REAL[],
    content_vector REAL[],
    vector_id INTEGER
);

ALTER TABLE public.articles ADD PRIMARY KEY (id);
'''

# 创建索引的 SQL 语句
create_indexes_sql = '''
CREATE INDEX ON public.articles USING ann (content_vector) WITH (distancemeasure = l2, dim = '1536', pq_segments = '64', hnsw_m = '100', pq_centers = '2048');

CREATE INDEX ON public.articles USING ann (title_vector) WITH (distancemeasure = l2, dim = '1536', pq_segments = '64', hnsw_m = '100', pq_centers = '2048');
'''

# 执行 SQL 语句
cursor.execute(create_table_sql)
cursor.execute(create_indexes_sql)

# 提交更改
connection.commit()

加载数据

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

import io

# 本地 CSV 文件的路径
csv_file_path = '../../data/vector_database_wikipedia_articles_embedded.csv'

# 定义一个逐行处理文件的生成器函数
def process_file(file_path):
    with open(file_path, 'r') as file:
        for line in file:
            # 将 '[' 替换为 '{' 并将 ']' 替换为 '}'
            modified_line = line.replace('[', '{').replace(']', '}')
            yield modified_line

# 创建一个 StringIO 对象来存储修改后的行
modified_lines = io.StringIO(''.join(list(process_file(csv_file_path))))

# 为 copy_expert 方法创建 COPY 命令
copy_command = '''
COPY public.articles (id, url, title, content, title_vector, content_vector, vector_id)
FROM STDIN WITH (FORMAT CSV, HEADER true, DELIMITER ',');
'''

# 使用 copy_expert 方法执行 COPY 命令
cursor.copy_expert(copy_command, modified_lines)

# 提交更改
connection.commit()

# 检查集合大小以确保所有点都已存储
count_sql = """select count(*) from public.articles;"""
cursor.execute(count_sql)
result = cursor.fetchone()
print(f"计数:{result[0]}")

计数:25000

搜索数据

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

def query_analyticdb(query, collection_name, vector_name="title_vector", top_k=20):

    # 从用户查询创建嵌入向量
    embedded_query = openai.Embedding.create(
        input=query,
        model="text-embedding-3-small",
    )["data"][0]["embedding"]

    # 将 embedded_query 转换为 PostgreSQL 兼容格式
    embedded_query_pg = "{" + ",".join(map(str, embedded_query)) + "}"

    # 创建 SQL 查询
    query_sql = f"""
    SELECT id, url, title, l2_distance({vector_name},'{embedded_query_pg}'::real[]) AS similarity
    FROM {collection_name}
    ORDER BY {vector_name} <-> '{embedded_query_pg}'::real[]
    LIMIT {top_k};
    """
    # 执行查询
    cursor.execute(query_sql)
    results = cursor.fetchall()

    return results

import openai

query_results = query_analyticdb("modern art in Europe", "Articles")
for i, result in enumerate(query_results):
    print(f"{i + 1}. {result[2]} (Score: {round(1 - result[3], 3)})")

1. 现代艺术博物馆 (Score: 0.75)
2. 西欧 (Score: 0.735)
3. 文艺复兴艺术 (Score: 0.728)
4. 波普艺术 (Score: 0.721)
5. 北欧 (Score: 0.71)
6. 希腊化艺术 (Score: 0.706)
7. 现代主义文学 (Score: 0.694)
8. 艺术电影 (Score: 0.687)
9. 中欧 (Score: 0.685)
10. 欧洲人 (Score: 0.683)
11. 艺术 (Score: 0.683)
12. 拜占庭艺术 (Score: 0.682)
13. 后现代主义 (Score: 0.68)
14. 东欧 (Score: 0.679)
15. 欧洲 (Score: 0.678)
16. 立体主义 (Score: 0.678)
17. 印象派 (Score: 0.677)
18. 包豪斯 (Score: 0.676)
19. 超现实主义 (Score: 0.674)
20. 表现主义 (Score: 0.674)

# 这次我们将使用内容向量进行查询
query_results = query_analyticdb("Famous battles in Scottish history", "Articles", "content_vector")
for i, result in enumerate(query_results):
    print(f"{i + 1}. {result[2]} (Score: {round(1 - result[3], 3)})")

1. 班诺克本战役 (Score: 0.739)
2. 苏格兰独立战争 (Score: 0.723)
3. 1651 (Score: 0.705)
4. 苏格兰独立战争 (Score: 0.699)
5. 罗伯特一世 (苏格兰) (Score: 0.692)
6. 841 (Score: 0.688)
7. 1716 (Score: 0.688)
8. 1314 (Score: 0.674)
9. 1263 (Score: 0.673)
10. 威廉·华莱士 (Score: 0.671)
11. 斯特灵 (Score: 0.663)
12. 1306 (Score: 0.662)
13. 1746 (Score: 0.661)
14. 1040 年代 (Score: 0.656)
15. 1106 (Score: 0.654)
16. 1304 (Score: 0.653)
17. 大卫二世 (苏格兰) (Score: 0.65)
18. 勇敢的心 (Score: 0.649)
19. 1124 (Score: 0.648)
20. 7 月 27 日 (Score: 0.646)