Neon Postgres 向量相似性搜索

本 Notebook 将指导您使用 Neon Serverless Postgres 作为 OpenAI 嵌入的向量数据库。它演示了如何：

1. 使用 OpenAI API 创建的嵌入。
2. 将嵌入存储在 Neon Serverless Postgres 数据库中。
3. 使用 OpenAI API 将原始文本查询转换为嵌入。
4. 使用 Neon 和 pgvector 扩展执行向量相似性搜索。

先决条件

在开始之前，请确保您拥有以下条件：

1. Neon Postgres 数据库。您可以在几个简单的步骤中创建一个帐户并设置一个具有即用型 neondb 数据库的项目。有关说明，请参阅 注册 和 创建您的第一个项目。
2. 您的 Neon 数据库连接字符串。您可以从 Neon Dashboard 上的“连接详细信息”小部件复制它。请参阅 从任何应用程序连接。
3. pgvector 扩展。通过运行 CREATE EXTENSION vector; 在 Neon 中安装扩展。有关说明，请参阅 启用 pgvector 扩展。
4. 您的 OpenAI API 密钥。
5. Python 和 pip。

安装所需的模块

此 Notebook 需要 openai、psycopg2、pandas、wget 和 python-dotenv 包。您可以使用 pip 安装它们：

! pip install openai psycopg2 pandas wget python-dotenv

准备您的 OpenAI API 密钥

需要 OpenAI API 密钥来为文档和查询生成向量。

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

将 OpenAI API 密钥添加为操作系统环境变量，或在提示时为其提供会话。如果您定义了环境变量，请将变量命名为 OPENAI_API_KEY。

有关将 OpenAI API 密钥配置为环境变量的信息，请参阅 API 密钥安全最佳实践。

测试您的 OpenAPI 密钥

# 测试以确保您的 OpenAI API 密钥已设置为环境变量或在提示时提供
# 如果您在本地运行此 Notebook，则可能需要重新加载终端和 Notebook 才能使环境可用

import os
from getpass import getpass

# 检查 OPENAI_API_KEY 是否已设置为环境变量
if os.getenv("OPENAI_API_KEY") is not None:
    print("您的 OPENAI_API_KEY 已准备就绪")
else:
    # 如果没有，则提示输入
    api_key = getpass("请输入您的 OPENAI_API_KEY：")
    if api_key:
        print("您的 OPENAI_API_KEY 现在可用于此会话")
        # 可选地，您可以将其设置为当前会话的环境变量
        os.environ["OPENAI_API_KEY"] = api_key
    else:
        print("您未输入您的 OPENAI_API_KEY")

您的 OPENAI_API_KEY 已准备就绪

连接到您的 Neon 数据库

在此处直接提供您的 Neon 数据库连接字符串，或使用 DATABASE_URL 变量在 .env 文件中定义它。有关获取 Neon 连接字符串的信息，请参阅 从任何应用程序连接。

import os
import psycopg2
from dotenv import load_dotenv

# 从 .env 文件加载环境变量
load_dotenv()

# 连接字符串可以直接在此处提供。
# 用您的 Neon 连接字符串替换下一行。
connection_string = "postgres://<user>:<password>@<hostname>/<dbname>"

# 如果连接字符串未直接在上面提供，
# 则检查 DATABASE_URL 是否已在环境变量或 .env 中设置。
if not connection_string:
    connection_string = os.environ.get("DATABASE_URL")

    # 如果两种方法都未提供连接字符串，则引发错误。
    if not connection_string:
        raise ValueError("请在代码中或 .env 文件中以 DATABASE_URL 的形式提供有效的连接字符串。")

# 使用连接字符串进行连接
connection = psycopg2.connect(connection_string)

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

测试到数据库的连接：

# 执行此查询以测试数据库连接
cursor.execute("SELECT 1;")
result = cursor.fetchone()

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

您的数据库连接成功！

本指南使用 OpenAI Cookbook examples 目录中提供的预计算维基百科文章嵌入，因此您无需使用自己的 OpenAI 积分来计算嵌入。

导入预计算的嵌入 zip 文件：

import wget

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

# 该文件约为 700 MB。导入它将需要几分钟时间。
wget.download(embeddings_url)

'vector_database_wikipedia_articles_embedded.zip'

解压下载的 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"csv 文件 {file_name} 存在于数据目录中。")
else:
    print(f"csv 文件 {file_name} 不存在于数据目录中。")

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

创建表并为您的向量嵌入添加索引

在数据库中创建的向量表名为 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 = '''
CREATE INDEX ON public.articles USING ivfflat (content_vector) WITH (lists = 1000);

CREATE INDEX ON public.articles USING ivfflat (title_vector) WITH (lists = 1000);
'''

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

# 提交更改
connection.commit()

加载数据

将预计算的向量数据从 .csv 文件加载到您的 articles 表中。有 25000 条记录，因此该操作预计需要几分钟时间。

import io

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

# 定义一个处理 csv 文件的生成器函数
def process_file(file_path):
    with open(file_path, 'r', encoding='utf-8') 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 = '''
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()

检查记录数以确保数据已加载。应有 25000 条记录。

# 检查数据大小
count_sql = """select count(*) from public.articles;"""
cursor.execute(count_sql)
result = cursor.fetchone()
print(f"计数:{result[0]}")

计数:25000

搜索您的数据

将数据存储在 Neon 数据库中后，您可以查询数据以查找最近邻。

首先定义 query_neon 函数，该函数在运行向量相似性搜索时执行。该函数根据用户查询创建嵌入，准备 SQL 查询，并使用嵌入运行 SQL 查询。您加载到数据库中的预计算嵌入是使用 text-embedding-3-small OpenAI 模型创建的，因此您必须使用相同的模型来创建用于相似性搜索的嵌入。

提供了一个 vector_name 参数，允许您按“title”或“content”进行搜索。

def query_neon(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}'::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

基于 title_vector 嵌入运行相似性搜索：

# 基于 `title_vector` 嵌入进行查询
import openai

query_results = query_neon("Greek mythology", "Articles")
for i, result in enumerate(query_results):
    print(f"{i + 1}. {result[2]} (Score: {round(1 - result[3], 3)})")

1. Greek mythology (Score: 0.998)
2. Roman mythology (Score: 0.7)
3. Greek underworld (Score: 0.637)
4. Mythology (Score: 0.635)
5. Classical mythology (Score: 0.629)
6. Japanese mythology (Score: 0.615)
7. Norse mythology (Score: 0.569)
8. Greek language (Score: 0.566)
9. Zeus (Score: 0.534)
10. List of mythologies (Score: 0.531)
11. Jupiter (mythology) (Score: 0.53)
12. Greek (Score: 0.53)
13. Gaia (mythology) (Score: 0.526)
14. Titan (mythology) (Score: 0.522)
15. Mercury (mythology) (Score: 0.521)
16. Ancient Greece (Score: 0.52)
17. Greek alphabet (Score: 0.52)
18. Venus (mythology) (Score: 0.515)
19. Pluto (mythology) (Score: 0.515)
20. Athena (Score: 0.514)

基于 content_vector 嵌入运行相似性搜索：

# 基于 `content_vector` 嵌入进行查询
query_results = query_neon("Famous battles in Greek history", "Articles", "content_vector")
for i, result in enumerate(query_results):
    print(f"{i + 1}. {result[2]} (Score: {round(1 - result[3], 3)})")

1. 222 BC (Score: 0.489)
2. Trojan War (Score: 0.458)
3. Peloponnesian War (Score: 0.456)
4. History of the Peloponnesian War (Score: 0.449)
5. 430 BC (Score: 0.441)
6. 168 BC (Score: 0.436)
7. Ancient Greece (Score: 0.429)
8. Classical Athens (Score: 0.428)
9. 499 BC (Score: 0.427)
10. Leonidas I (Score: 0.426)
11. Battle (Score: 0.421)
12. Greek War of Independence (Score: 0.421)
13. Menelaus (Score: 0.419)
14. Thebes, Greece (Score: 0.417)
15. Patroclus (Score: 0.417)
16. 427 BC (Score: 0.416)
17. 429 BC (Score: 0.413)
18. August 2 (Score: 0.412)
19. Ionia (Score: 0.411)
20. 323 (Score: 0.409)