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

本笔记本将指导您逐步使用Hologres作为OpenAI嵌入的向量数据库。

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

1. 使用OpenAI API创建的预计算嵌入。
2. 将嵌入存储在Hologres的云实例中。
3. 使用OpenAI API将原始文本查询转换为嵌入。
4. 使用Hologres在创建的集合中执行最近邻搜索。
5. 在提示工程中，将搜索结果作为上下文提供给大型语言模型。

什么是Hologres

Hologres是阿里云开发的一体化实时数据仓库服务。您可以使用Hologres实时写入、更新、处理和分析大量数据。Hologres支持标准SQL语法，兼容PostgreSQL，并支持大多数PostgreSQL函数。Hologres支持高达PB级别数据的在线分析处理（OLAP）和即席分析，并为高达PB级别的数据提供高并发、低延迟的在线数据服务。Hologres支持多工作负载的细粒度隔离和企业级安全能力。Hologres与MaxCompute、实时计算（Apache Flink）和DataWorks深度集成，为企业提供全栈的在线和离线数据仓库解决方案。

Hologres通过采用Proxima提供向量数据库功能。

Proxima是阿里云达摩院开发的高性能软件库。它允许您搜索向量的最近邻。与Facebook AI Similarity Search（Faiss）等类似的开源软件相比，Proxima提供了更高的稳定性和性能。Proxima提供的基础模块在行业内具有领先的性能和效果，并允许您搜索相似的图像、视频或人脸。Hologres与Proxima深度集成，提供高性能的向量搜索服务。

部署选项

• 点击此处快速部署Hologres数据仓库。

先决条件

为了完成本次练习，我们需要准备几件事：

1. Hologres云服务器实例。
2. 'psycopg2-binary'库，用于与向量数据库进行交互。任何其他postgresql客户端库都可以。
3. 一个OpenAI API密钥。

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

安装要求

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

! pip install openai psycopg2-binary 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 已准备就绪

连接到Hologres

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

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

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")
)
connection.set_session(autocommit=True)

# 创建一个新的游标对象
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"

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

下载的文件随后需要解压：

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 存在于数据目录中。

加载数据

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

!unzip -n vector_database_wikipedia_articles_embedded.zip
!ls -lh vector_database_wikipedia_articles_embedded.csv

Archive:  vector_database_wikipedia_articles_embedded.zip
-rw-r--r--@ 1 geng  staff   1.7G Jan 31 01:19 vector_database_wikipedia_articles_embedded.csv

看一下数据。

import pandas, json
data = pandas.read_csv('../../data/vector_database_wikipedia_articles_embedded.csv')
data

title_vector_length = len(json.loads(data['title_vector'].iloc[0]))
content_vector_length = len(json.loads(data['content_vector'].iloc[0]))

print(title_vector_length, content_vector_length)

1536 1536

创建表和proxima向量索引

Hologres将数据存储在__表中，其中每个对象至少由一个向量描述。我们的表将命名为articles，每个对象将由title和content向量描述。

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

cursor.execute('CREATE EXTENSION IF NOT EXISTS proxima;')
create_proxima_table_sql = '''
BEGIN;
DROP TABLE IF EXISTS articles;
CREATE TABLE articles (
    id INT PRIMARY KEY NOT NULL,
    url TEXT,
    title TEXT,
    content TEXT,
    title_vector float4[] check(
        array_ndims(title_vector) = 1 and 
        array_length(title_vector, 1) = 1536
    ), -- 定义向量
    content_vector float4[] check(
        array_ndims(content_vector) = 1 and 
        array_length(content_vector, 1) = 1536
    ),
    vector_id INT
);

-- 为向量字段创建索引。
call set_table_property(
    'articles',
    'proxima_vectors', 
    '{
        "title_vector":{"algorithm":"Graph","distance_method":"Euclidean","builder_params":{"min_flush_proxima_row_count" : 10}},
        "content_vector":{"algorithm":"Graph","distance_method":"Euclidean","builder_params":{"min_flush_proxima_row_count" : 10}}
    }'
);  

COMMIT;
'''

# 执行SQL语句（将自动提交）
cursor.execute(create_proxima_table_sql)

上传数据

现在，让我们使用COPY语句将数据上传到Hologres云实例。根据网络带宽，这可能需要5-10分钟。

import io

# 解压后的CSV文件路径
csv_file_path = '../../data/vector_database_wikipedia_articles_embedded.csv'

# 在SQL中，数组用{}包围，而不是[]
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)

proxima索引将在后台构建。在此期间我们可以进行搜索，但如果没有向量索引，查询速度会很慢。使用此命令等待索引构建完成。

cursor.execute('vacuum articles;')

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

计数:25000

搜索数据

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

import openai
def query_knn(query, table_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, pm_approx_euclidean_distance({vector_name},'{embedded_query_pg}'::float4[]) AS distance
    FROM {table_name}
    ORDER BY distance
    LIMIT {top_k};
    """
    # 执行查询
    cursor.execute(query_sql)
    results = cursor.fetchall()

    return results

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

1. 现代艺术博物馆 (分数: 0.501)
2. 西欧 (分数: 0.485)
3. 文艺复兴艺术 (分数: 0.479)
4. 波普艺术 (分数: 0.472)
5. 北欧 (分数: 0.461)
6. 希腊化艺术 (分数: 0.458)
7. 现代主义文学 (分数: 0.447)
8. 艺术电影 (分数: 0.44)
9. 中欧 (分数: 0.439)
10. 艺术 (分数: 0.437)
11. 欧洲 (分数: 0.437)
12. 拜占庭艺术 (分数: 0.436)
13. 后现代主义 (分数: 0.435)
14. 东欧 (分数: 0.433)
15. 立体主义 (分数: 0.433)
16. 欧洲 (分数: 0.432)
17. 印象派 (分数: 0.432)
18. 包豪斯 (分数: 0.431)
19. 超现实主义 (分数: 0.429)
20. 表现主义 (分数: 0.429)

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

1. 班诺克本战役 (分数: 0.489)
2. 苏格兰独立战争 (分数: 0.474)
3. 1651年 (分数: 0.457)
4. 苏格兰独立战争 (分数: 0.452)
5. 罗伯特一世 (苏格兰) (分数: 0.445)
6. 841年 (分数: 0.441)
7. 1716年 (分数: 0.441)
8. 1314年 (分数: 0.429)
9. 1263年 (分数: 0.428)
10. 威廉·华莱士 (分数: 0.426)
11. 斯特灵 (分数: 0.419)
12. 1306年 (分数: 0.419)
13. 1746年 (分数: 0.418)
14. 1040年代 (分数: 0.414)
15. 1106年 (分数: 0.412)
16. 1304年 (分数: 0.411)
17. 大卫二世 (苏格兰) (分数: 0.408)
18. 勇敢的心 (分数: 0.407)
19. 1124年 (分数: 0.406)
20. 7月27日 (分数: 0.405)