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) => [
    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 = 'Guides on using Drizzle ORM with different platforms';

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
  }
]