使用 pgvector 扩展进行向量相似性搜索

Drizzle | 使用 pgvector 扩展进行向量相似性搜索

This guide assumes familiarity with:

从 PostgreSQL 开始入门。
选择语句
索引
sql 操作符
pgvector 扩展
Drizzle kit
您应该安装 openai 包来生成嵌入。

npm

yarn

pnpm

bun

npm i openai

yarn add openai

pnpm add openai

bun add openai

您应该安装 drizzle-orm@0.31.0 和 drizzle-kit@0.22.0 或更高版本。

要在 PostgreSQL 中使用 Drizzle ORM 实现向量相似性搜索，您可以使用 pgvector 扩展。这个扩展提供了一系列函数来处理向量并进行相似性搜索。

目前，Drizzle 并不会自动创建扩展，因此您需要手动创建。创建一个空的迁移文件并添加 SQL 查询：

npx drizzle-kit generate --custom

CREATE EXTENSION vector;

为了进行相似性搜索，您需要创建一个包含向量列的表，并在该列上创建一个 HNSW 或 IVFFlat 索引以提高性能：

schema.ts

migration.sql

import { index, pgTable, serial, text, vector } from 'drizzle-orm/pg-core';

export const guides = pgTable(
  'guides',
  {
    id: serial('id').primaryKey(),
    title: text('title').notNull(),
    description: text('description').notNull(),
    url: text('url').notNull(),
    embedding: vector('embedding', { dimensions: 1536 }),
  },
  (table) => [
    index('embeddingIndex').using('hnsw', table.embedding.op('vector_cosine_ops')),
  ]
);

CREATE TABLE IF NOT EXISTS "guides" (
  "id" serial PRIMARY KEY NOT NULL,
  "title" text NOT NULL,
  "description" text NOT NULL,
  "url" text NOT NULL,
  "embedding" vector(1536)
);
--> statement-breakpoint
CREATE INDEX IF NOT EXISTS "embeddingIndex" ON "guides" USING hnsw (embedding vector_cosine_ops);

embedding 列用于存储指南描述的向量嵌入。向量嵌入是一种对某些数据的表示。它将不同类型的数据转换为语言模型可以处理的通用格式（向量）。这使我们能够执行数学运算，例如测量两个向量之间的距离，以确定两个数据项是多么相似或不同。

在此示例中，我们将使用 OpenAI 模型生成描述的嵌入：

import OpenAI from 'openai';

const openai = new OpenAI({
  apiKey: process.env['OPENAI_API_KEY'],
});

export const generateEmbedding = async (value: string): Promise<number[]> => {
  const input = value.replaceAll('\n', ' ');

  const { data } = await openai.embeddings.create({
    model: 'text-embedding-ada-002',
    input,
  });

  return data[0].embedding;
};

要通过嵌入搜索相似的指南，您可以使用 gt 和 sql 操作符与 cosineDistance 函数来计算 embedding 列与生成的嵌入之间的相似性：

import { cosineDistance, desc, gt, sql } from 'drizzle-orm';
import { generateEmbedding } from './embedding';
import { guides } from './schema';

const db = drizzle(...);

const findSimilarGuides = async (description: string) => {
  const embedding = await generateEmbedding(description);

  const similarity = sql<number>`1 - (${cosineDistance(guides.embedding, embedding)})`;

  const similarGuides = await db
    .select({ name: guides.title, url: guides.url, similarity })
    .from(guides)
    .where(gt(similarity, 0.5))
    .orderBy((t) => desc(t.similarity))
    .limit(4);

  return similarGuides;
};

const description = '使用 Drizzle ORM 的各种平台指南';

const similarGuides = await findSimilarGuides(description);

[
  {
    name: 'Drizzle with Turso',
    url: '/docs/tutorials/drizzle-with-turso',
    similarity: 0.8642314333984994
  },
  {
    name: 'Drizzle with Supabase Database',
    url: '/docs/tutorials/drizzle-with-supabase',
    similarity: 0.8593631126014918
  },
  {
    name: 'Drizzle with Neon Postgres',
    url: '/docs/tutorials/drizzle-with-neon',
    similarity: 0.8541051184461372
  },
  {
    name: 'Drizzle with Vercel Edge Functions',
    url: '/docs/tutorials/drizzle-with-vercel-edge-functions',
    similarity: 0.8481551084241092
  }
]