Drizzle ORM - pgvector 확장을 활용한 벡터 유사성 검색

Drizzle | pgvector 확장을 활용한 벡터 유사성 검색

This guide assumes familiarity with:

PostgreSQL 시작하기
SELECT 문
인덱스
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 이상 버전이 필요합니다.

Drizzle ORM과 함께 PostgreSQL에서 벡터 유사성 검색을 구현하려면 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) => ({
    embeddingIndex: 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 => {
  const input = value.replaceAll('\n', ' ');

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

  return data[0].embedding;
};

임베딩을 기반으로 유사한 가이드를 검색하려면 gt와 sql 연산자를 사용하여 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`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: 'Turso와 함께하는 Drizzle',
    url: '/docs/tutorials/drizzle-with-turso',
    similarity: 0.8642314333984994
  },
  {
    name: 'Supabase 데이터베이스와 함께하는 Drizzle',
    url: '/docs/tutorials/drizzle-with-supabase',
    similarity: 0.8593631126014918
  },
  {
    name: 'Neon Postgres와 함께하는 Drizzle',
    url: '/docs/tutorials/drizzle-with-neon',
    similarity: 0.8541051184461372
  },
  {
    name: 'Vercel Edge Functions와 함께하는 Drizzle',
    url: '/docs/tutorials/drizzle-with-vercel-edge-functions',
    similarity: 0.8481551084241092
  }
]