PostgreSQL 教程: 使用 pgvector 实现语义搜索

二月 20, 2025

摘要：在本教程中，您将学习如何在 PostgreSQL 中使用 pgvector 实现语义搜索。

初始设置

近年来，由于生成式 AI 和大型语言模型生态系统的所有进步，语义搜索变得越来越流行。

在本教程中，我们将会使用 Kaggle 上的电影数据集，其中包括了截至 2017 年的 45000 部电影，这样的数据量对我们来说已经足够了。

我们将会使用 pgvector，这是一个 PostgreSQL 的开源扩展，它增加了对向量运算和相似性搜索的支持。它允许您直接在 PostgreSQL 数据库中存储、索引和查询向量数据。首先，让我们创建一个存放电影数据集的表：

CREATE TABLE IF NOT EXISTS movies (
    id integer,
    imdb_id text,
    title text,
    tagline text,
    overview text,
    release_date date,
    vote_average float,
    embedding vector(1536)
);

现在，让我们来使用 pgvector 和 OpenAI 嵌入模块，构建一个简单的语义搜索引擎！

import openai
import psycopg2
import csv

# Set up OpenAI API (replace with your actual API key)
openai.api_key = "your_openai_api_key"

# Function to get embeddings from OpenAI
def get_embedding(text):
    response = openai.embeddings.create(input=text, model="text-embedding-ada-002")
    return response['data'][0]['embedding']

# Function to add a document
def add_document(cur, row):
    content = """
                 Title: {0}
                 Tagline: {1}
                 Overview: {2}
                 Release date: {3}
                 Genres: {4}""".format(
                 row['title'],
                 row['tagline'],
                 row['overview'],
                 row['release_date'],
                 row['genres'])
    embedding = get_embedding(content)
    cur.execute("""INSERT INTO movies
                   (id, imdb_id, title, tagline, overview, release_date, vote_average, embedding)
                   VALUES (%s, %s, %s, %s, %s, %s, %s, %s)""",
                (row['id'], row['imdb_id'],
                 row['title'], row['tagline'], row['overview'], row['release_date'],
                 row['vote_average'], embedding))

# Connect to the database
conn = psycopg2.connect("dbname=your_database user=your_username")
cur = conn.cursor()

with open("./movies_metadata.csv", "r") as csvfile:
    reader = csv.DictReader(csvfile)
    for row in reader:
        print(row['first_name'], row['last_name'])
        add_document(cur, row)
        conn.commit()

# Clean up
cur.close()
conn.close()

它使用了 OpenAI 的文本嵌入模型来嵌入文档，并将它们存储到带有 pgvector 的 PostgreSQL 数据库中。

演示

如果你搜索 “old movies”，你会得到一些老电影，因为我们在嵌入内容中包含了发布日期。

还有有趣的地方；如果你搜索 “virtual reality”，你会在前 20 个结果中得到 Oculus，这是一部由 Mike Flanagan 创作的恐怖电影。这部电影与 VR 没有任何关系，但标题刚好是一种流行的 VR 耳机的名称，这表明该模型有某种世界性的知识。

# Function to search for similar documents
def search_documents(query, limit=5):
    query_embedding = get_embedding(query)
    cur.execute("""
        SELECT title, overview, embedding <-> %s AS distance
        FROM movies
        ORDER BY distance
        LIMIT %s
    """, (query_embedding, limit))
    return cur.fetchall()

# Connect to the database
conn = psycopg2.connect("dbname=your_database user=your_username")
cur = conn.cursor()

# Perform a search
search_query = "wizard school poor boy"
results = search_documents(search_query)
print(f"Search results for: '{search_query}'")
for i, (title, overview, distance) in enumerate(results, 1):
    print(f"{i}. {title} {overview} (Distance: {distance:.4f})")

# Clean up
cur.close()
conn.close()