使用PolarDB-PG作为OpenAI嵌入向量数据库

本笔记将逐步指导您如何将PolarDB-PG用作OpenAI嵌入向量数据库。

本笔记介绍了以下端到端的过程： 1. 使用OpenAI API创建的预先计算的嵌入向量。 2. 将嵌入向量存储在PolarDB-PG的云实例中。 3. 将原始文本查询转换为嵌入向量，使用OpenAI API。 4. 使用PolarDB-PG在创建的集合中执行最近邻搜索。

什么是PolarDB-PG

PolarDB-PG 是一种高性能向量数据库，采用读写分离架构。它是由阿里云管理的云原生数据库，与PostgreSQL 100%兼容，与Oracle语法高度兼容。它支持处理大规模向量数据存储和查询，并通过优化底层执行算法大大提高向量计算的效率，为用户提供快速、弹性、高性能、大规模存储、安全可靠的向量数据库服务。此外，PolarDB-PG还支持多维和多模态时空信息引擎和地理信息引擎。同时，PolarDB-PG配备完整的OLAP功能和服务级别协议，得到许多用户的认可和使用；

部署选项

• 使用PolarDB-PG云向量数据库。点击这里快速部署。

先决条件

为了完成这个练习，我们需要准备一些东西：

1. PolarDB-PG 云服务器实例。
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。

如果您对通过环境变量设置API密钥有任何疑问，请参考API密钥安全最佳实践。

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

if os.getenv("OPENAI_API_KEY") is not None:
    print("OPENAI_API_KEY is ready")
else:
    print("OPENAI_API_KEY environment variable not found")

OPENAI_API_KEY is ready

连接到PolarDB

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

使用官方的Python库连接到正在运行的PolarDB服务器实例非常简单：

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("Connection successful!")
else:
    print("Connection failed.")

Connection successful!

import wget

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

# 文件大小约为700MB，因此需要一些时间来完成。
wget.download(embeddings_url)

'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"The file {file_name} exists in the data directory.")
else:
    print(f"The file {file_name} does not exist in the data directory.")

The file vector_database_wikipedia_articles_embedded.csv exists in the data directory.

索引数据

PolarDB将数据存储在__关系__中，其中每个对象至少由一个向量描述。我们的关系将被称为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 vector(1536),
    content_vector vector(1536),
    vector_id INTEGER
);

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

# 创建索引的SQL语句
create_indexes_sql = '''
在public.articles表上使用ivfflat方法为content_vector字段创建索引，设置lists参数为1000；

在public.articles表上使用ivfflat方法为title_vector字段创建索引，设置lists参数为1000。
'''

# 执行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:
            yield line

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

# 创建用于 copy_expert 方法的 COPY 命令
copy_command = '''
复制 public.articles 表中的数据（包括 id、url、title、content、title_vector、content_vector、vector_id 字段）
从标准输入读取，使用 CSV 格式，包含表头，分隔符为逗号。
'''

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

# 提交更改
connection.commit()

# 检查集合大小，确保所有点都已存储。
count_sql = """从 public.articles 表中选择计数(*)；"""
cursor.execute(count_sql)
result = cursor.fetchone()
print(f"Count:{result[0]}")

Count:25000

搜索数据

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

def query_polardb(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"]

    # 将嵌入式查询转换为与PostgreSQL兼容的格式
    embedded_query_pg = "[" + ",".join(map(str, embedded_query)) + "]"

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

    return results

import openai

query_results = query_polardb("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. Museum of Modern Art (Score: 0.5)
2. Western Europe (Score: 0.485)
3. Renaissance art (Score: 0.479)
4. Pop art (Score: 0.472)
5. Northern Europe (Score: 0.461)
6. Hellenistic art (Score: 0.457)
7. Modernist literature (Score: 0.447)
8. Art film (Score: 0.44)
9. Central Europe (Score: 0.439)
10. European (Score: 0.437)
11. Art (Score: 0.437)
12. Byzantine art (Score: 0.436)
13. Postmodernism (Score: 0.434)
14. Eastern Europe (Score: 0.433)
15. Europe (Score: 0.433)
16. Cubism (Score: 0.432)
17. Impressionism (Score: 0.432)
18. Bauhaus (Score: 0.431)
19. Surrealism (Score: 0.429)
20. Expressionism (Score: 0.429)

# 这次我们将使用内容向量进行查询。
query_results = query_polardb("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. Battle of Bannockburn (Score: 0.489)
2. Wars of Scottish Independence (Score: 0.474)
3. 1651 (Score: 0.457)
4. First War of Scottish Independence (Score: 0.452)
5. Robert I of Scotland (Score: 0.445)
6. 841 (Score: 0.441)
7. 1716 (Score: 0.441)
8. 1314 (Score: 0.429)
9. 1263 (Score: 0.428)
10. William Wallace (Score: 0.426)
11. Stirling (Score: 0.419)
12. 1306 (Score: 0.419)
13. 1746 (Score: 0.418)
14. 1040s (Score: 0.414)
15. 1106 (Score: 0.412)
16. 1304 (Score: 0.411)
17. David II of Scotland (Score: 0.408)
18. Braveheart (Score: 0.407)
19. 1124 (Score: 0.406)
20. July 27 (Score: 0.405)