如何使用Neon Serverless Postgres构建OpenAI矢量数据库

随着人工智能和深度学习技术的迅猛发展，矢量数据在数据科学和应用领域中扮演着越来越重要的角色。从自然语言处理到图像识别，矢量数据是我们理解和处理各种信息的关键。然而，有效地存储和查询矢量数据一直是一个挑战，特别是在云计算环境中。在本教程中，我们将介绍如何使用Neon Serverless Postgres构建OpenAI矢量数据库，以便存储和查询OpenAI生成的嵌入向量。

准备工作

在开始之前，请确保您具备以下条件：

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

安装所需模块

这个笔记本需要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密钥

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")

连接到您的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>"

# 如果上面没有直接提供connection_string，
# 则检查环境或.env中是否设置了DATABASE_URL。
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("连接失败。")

加载数据

将预先计算的矢量数据加载到名为articles的数据库表中，从.csv文件加载数据。由于有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]}")

搜索您的数据

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

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

提供了一个vector_name参数，允许基于“标题”或“内容”搜索。

def query_neon(query, collection

_name, vector_name="title_vector", top_k=20):

    # 从用户查询创建嵌入向量
    embedded_query = openai.Embedding.create(
        input=query,
        model="text-embedding-ada-002",
    )["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

# 基于title_vector嵌入运行相似性搜索
query_results = query_neon("Greek mythology", "Articles")
for i, result in enumerate(query_results):
    print(f"{i + 1}. {result[2]} (分数: {round(1 - result[3], 3)})")

# 基于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]} (分数: {round(1 - result[3], 3)})")

结论

通过使用Neon Serverless Postgres构建OpenAI矢量数据库，您可以有效地存储和查询OpenAI生成的嵌入向量。这为您提供了强大的工具来处理和分析文本数据，从而使您能够更好地理解和利用矢量数据的潜力。