PostgreSQL과 pgvector를 사용하여 RAG 기반 데이터베이스 어시스턴트 구축하기
태그: ai, database, postgres, tutorial

문제점
모든 기업은 데이터베이스와 동일한 대화를 나눕니다:

"지난 30일 동안 놓친 주문이 몇 건인가요?" "세 번 이상 주문했지만 리뷰를 남기지 않은 고객은 누구인가요?" "지난 분기 지역별 매출 트렌드를 보여주세요."

이러한 질문들은 인간 분석가에게는 간단한 질문이지만, SQL 쿼리로 신속하게 변환되지 않습니다. 비즈니스 팀은 엔지니어링 팀을 기다립니다. 엔지니어링 팀은 바쁩니다. 쿼리는 결국 3일 후에 작성되며, 그 답변은 이미 너무 늦은 정보가 됩니다.

만약 데이터베이스 자체가 자연어 질문에 몇 초 만에 답할 수 있다면 어떨까요?

이것이 바로 이 튜토리얼에서 구축하는 것입니다: 관련 스키마 컨텍스트 (schema context)를 검색하고, LLM (대규모 언어 모델)으로 SQL을 생성하며, 답변을 직접 반환하는 RAG (검색 증강 생성) 기반 데이터베이스 어시스턴트입니다.

"그냥 ChatGPT에게 물어보는 것" 대신 RAG를 사용하는 이유?

스키마 인식 (Schema-aware): LLM이 데이터베이스의 실제 테이블, 컬럼 및 관계를 확인합니다.
최신 상태 유지 (Up-to-date): 재학습 없이 메타데이터 변경 사항이 즉시 반영됩니다.
안전함 (Safe): 생성된 SQL은 읽기 전용 (read-only) 연결에서 실행되므로 데이터를 변경할 수 없습니다.
설명 가능함 (Explainable): 사용자는 질문에 답하기 위해 어떤 스키마 컨텍스트가 사용되었는지 확인할 수 있습니다.

왜 PostgreSQL + pgvector인가
PostgreSQL은 오늘날 사용되는 가장 다재다능한 프로덕션 데이터베이스입니다. 이미 JSON, 전문 검색 (full-text search), 윈도우 함수 (window functions) 및 풍부한 인덱스를 지원합니다. 2023년부터 pgvector 확장을 통해 스택에 다른 데이터베이스를 추가하지 않고도 프로덕션급 벡터 검색 (vector search) 기능을 제공합니다.

운영 데이터와 검색 모두에 동일한 데이터베이스를 사용하는 것은 다음을 의미합니다:

단일 백업 전략, 단일 모니터링 스택, 단일 자격 증명 세트.
메타데이터와 실제 테이블 간의 SQL 조인 (joins).
별도의 벡터 저장소 (vector store)로 이동하기 위한 추가적인 네트워크 홉 (network hop) 없음.

아키텍처 (Architecture)
사용자 질문 (User question)
|
v
[1] 질문 임베딩 (Embed question) (text-embedding-3-small)
|
v
[2] schema_metadata에 대한 pgvector 유사도 검색 (similarity search)
|
v
[3] 프롬프트 구축 (Build prompt): 질문 + 관련 테이블/컬럼/예시
|
v
[4] LLM이 SQL 생성 (읽기 전용)
|
v
[5] SQL 실행 -> 답변 + 설명 반환

핵심 아이디어는 문서를 저장하는 것이 아니라, 데이터베이스 스키마 메타데이터 (database schema metadata)를 검색 코퍼스 (retrieval corpus)로 저장한다는 것입니다. 이것이 일반적인 챗봇 (chatbot)과 데이터베이스 어시스턴트 (database assistant)의 차이점입니다.

프로젝트 설정 (Project Setup)

사전 요구 사항 (Prerequisites)

• pgvector가 설치된 PostgreSQL 15 이상
• Python 3.11 이상
• OpenAI 호환 API 키

pip install "psycopg[binary]" openai

import os
import json
import psycopg
from openai import OpenAI

client = OpenAI() # 환경 변수에서 OPENAI_API_KEY를 읽어옵니다

DB_URL = os.getenv(
"DATABASE_URL",
"postgresql://postgres:postgres@localhost:5432/rag_db",
)

단계 1: pgvector 활성화 및 메타데이터 테이블 생성

with psycopg.connect(DB_URL, autocommit=True) as conn:
    conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
    conn.execute("""
    CREATE TABLE IF NOT EXISTS schema_metadata (
        id BIGSERIAL PRIMARY KEY,
        object_name TEXT NOT NULL,
        object_type TEXT NOT NULL, -- table | column | example
        definition TEXT NOT NULL, -- DDL 또는 컬럼 설명
        description TEXT NOT NULL, -- 자연어 의미론 (natural-language semantics)
        embedding vector(1536) -- text-embedding-3-small
    );
    """)
    conn.execute("""
    CREATE INDEX IF NOT EXISTS schema_metadata_embedding_idx
    ON schema_metadata
    USING hnsw (embedding vector_cosine_ops);
    """)

HNSW 인덱스는 적당한 규모의 데이터셋에서 10ms 미만의 지연 시간으로 근사 최근접 이웃 (approximate nearest-neighbor) 검색을 제공합니다. 수천 개의 메타데이터 행의 경우, 이는 사실상 즉각적입니다.

2단계: 데이터베이스에서 스키마 메타데이터 (Schema Metadata) 추출하기
테이블 설명을 수동으로 작성하는 대신, information_schema에서 실제 스키마를 추출합니다:

def fetch_schema_metadata(conn) -> list[dict]:  
    rows = conn.execute("""  
    SELECT  
    c.table_name,  
    c.column_name,  
    c.data_type,  
    pg_catalog.col_description(  
    ('"' || c.table_schema || '"."' || c.table_name || '"')::regclass,  
    c.ordinal_position  
    ) AS column_comment  
    FROM information_schema.columns c  
    WHERE c.table_schema = 'public'  
    ORDER BY c.table_name, c.ordinal_position  
    """).fetchall()

metadata = []
for table_name, column_name, data_type, comment in rows:
    metadata.append({
...

또한 쿼리 예시 (query examples)를 추가합니다. 이를 통해 LLM(Large Language Model)에게 팀에서 사용하는 쿼리 스타일을 학습시킬 수 있습니다:

QUERY_EXAMPLES = [  
{
"object_name": "example:monthly_revenue",  
"object_type": "example",  
"definition": (  
"Question: What is the monthly revenue?\n"  
"SQL: SELECT DATE_TRUNC('month', created_at) AS month, "  
"SUM(total_amount) FROM orders GROUP BY 1 ORDER BY 1;"
),
"description": "Monthly revenue aggregation on orders",  
},
{
"object_name": "example:top_customers",  
"object_type": "example",  
"definition": (  
"Question: Who are the top 10 customers by spend?\n"  
"SQL: SELECT c.name, SUM(o.total_amount) AS spend "  
"FROM customers c JOIN orders o ON o.customer_id = c.id "  
"GROUP BY c.id ORDER BY spend DESC LIMIT 10;"
),
"description": "Top customers by total order amount",  
},
]

3단계: 메타데이터 임베딩 (Embed) 및 저장하기

def embed_texts(texts: list[str]) -> list[list[float]]:  
    resp = client.embeddings.create(  
        model="text-embedding-3-small",  
        input=texts,  
    )  
    return [item.embedding for item in resp.data]

def ingest_metadata():  
    with psycopg.connect(DB_URL) as conn:  
        schema_metadata = fetch_schema_metadata(conn)

docs = []
for item in schema_metadata + QUERY_EXAMPLES:
    docs.append(f"{item['object_name']}\n{item['definition']}\n{item['description']}")
...

참고: vector 타입은 psycopg를 통해 Python 리스트를 직접 전달받을 수 있으며, 확장 기능(extension)이 이를 자체 네이티브 형식으로 직렬화(serialize)합니다.

4단계: 관련 스키마 컨텍스트 검색 (Retrieve Relevant Schema Context)
def retrieve_context(question: str, top_k: int = 5) -> str:
q_vector = embed_texts([question])[0]

with psycopg.connect(DB_URL) as conn:
    rows = conn.execute("""
        SELECT object_name, object_type, definition, description,
...

코사인 연산자 <=>가 핵심적인 부분입니다. 이 연산자는 사용자의 질문과 의미론적 관련성(semantic relevance)에 따라 스키마 조각들의 순위를 매깁니다.

5단계: SQL 생성 및 답변 반환 (Generate SQL and Return an Answer)
SYSTEM_PROMPT = """
당신은 시니어 데이터베이스 분석가입니다. 당신의 임무는
자연어 질문을 안전하고 정확한 PostgreSQL로 변환하는 것입니다.

규칙:

1. 제공된 컨텍스트(context)에 있는 테이블과 컬럼만 사용하세요.
2. 필요한 경우 식별자(identifier)를 항상 큰따옴표(")로 감싸세요.
3. DELETE, UPDATE, INSERT, DROP, TRUNCATE 또는 DDL을 절대 사용하지 마세요.
4. 사용자가 명시적으로 모든 행을 요청하지 않는 한 LIMIT을 추가하세요.
5. sql 코드 펜스(code fence) 안에 SQL만 반환하고, 그 뒤에 일반 텍스트로 짧은 설명을 하나 덧붙이세요.
   """

def generate_sql(question: str, context: str) -> str:
response = client.chat.completions.create(
model="gpt-4o-mini",
messages=[
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": (
f"Database schema context:\n{context}\n\n"
f"Question: {question}\n\n"
"Generate PostgreSQL."
)},
],
temperature=0,
)
content = response.choices[0].message.content
return extract_sql(content)

def extract_sql(content: str) -> str:
if "sql" in content: return content.split("sql")[1].split("```")[0].strip()
return content.strip()
...