Aurora PostgreSQL에서 한국어 하이브리드 검색 구현하기: pg_bigm + pgvector로 만드는 한국어 특화 RAG

-- pgvector 확장 설치 (벡터 시맨틱 검색)
CREATE EXTENSION IF NOT EXISTS vector;

-- pg_bigm 확장 설치 (한국어 키워드 검색)
CREATE EXTENSION IF NOT EXISTS pg_bigm;

-- 설치 확인
SELECT extname, extversion FROM pg_extension
WHERE extname IN ('vector', 'pg_bigm');

CREATE TABLE kb_documents (
    id BIGSERIAL PRIMARY KEY,
    title TEXT NOT NULL,
    content TEXT NOT NULL,
    source VARCHAR(200),
    category VARCHAR(100),
    embedding vector(1024),  -- Amazon Titan Embeddings V2 (256/512/1024차원 선택 가능)
    created_at TIMESTAMPTZ DEFAULT NOW()
);

-- pg_bigm GIN 인덱스 (한국어 키워드 검색 가속)
CREATE INDEX idx_kb_content_bigm ON kb_documents USING gin (content gin_bigm_ops);
CREATE INDEX idx_kb_title_bigm ON kb_documents USING gin (title gin_bigm_ops);

-- pgvector HNSW 인덱스 (벡터 시맨틱 검색 가속)
-- normalize=True 임베딩이라면 vector_ip_ops(내적)로 변경하여 계산 효율을 높일 수 있습니다.
-- 여기서는 범용성을 위해 vector_cosine_ops를 사용합니다.
CREATE INDEX idx_kb_embedding ON kb_documents
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 128);

-- HNSW 검색 품질 튜닝: ef_search 값을 높이면 재현율이 올라가지만 속도가 느려집니다.
-- 기본값 40, RAG 환경에서는 100~200 권장
SET hnsw.ef_search = 100;

-- tsvector: 한국어 조사 처리 실패
SELECT to_tsvector('simple', '대한민국은 동아시아에 위치한 나라이다')
    @@ to_tsquery('simple', '대한민국');
-- 결과: false ❌

-- pg_bigm(LIKE): 조사가 붙어도 매칭 성공
SELECT '대한민국은 동아시아에 위치한 나라이다' LIKE '%대한민국%';
-- 결과: true ✅

-- pgvector 시맨틱 검색: 의미적으로 유사한 문서 찾기
-- query_embedding은 Python에서 Bedrock Titan으로 생성한 벡터
SELECT title, content, embedding <=> query_embedding AS distance
FROM kb_documents
ORDER BY embedding <=> query_embedding
LIMIT 5;

CREATE OR REPLACE FUNCTION hybrid_search_ko(
    query_text TEXT,
    query_embedding vector(1024),
    match_count INT DEFAULT 10,
    keyword_weight FLOAT DEFAULT 0.4,
    semantic_weight FLOAT DEFAULT 0.6,
    rrf_k INT DEFAULT 60
)
RETURNS TABLE (
    doc_id BIGINT, doc_title TEXT, doc_content TEXT,
    keyword_rank INT, semantic_rank INT, rrf_score FLOAT
)
LANGUAGE plpgsql
AS $$
BEGIN
    RETURN QUERY
    WITH
    -- pg_bigm 키워드 검색: title에 가중치 2배 (짧은 텍스트에서 similarity가 더 변별력 있음)
    keyword_candidates AS (
        SELECT d.id, d.title, d.content,
            (bigm_similarity(d.title, query_text) * 2.0 +
             bigm_similarity(d.content, query_text)) AS combined_sim
        FROM kb_documents d
        WHERE d.content LIKE '%' || query_text || '%'
           OR d.title LIKE '%' || query_text || '%'
    ),
    keyword_results AS (
        SELECT kc.id, kc.title, kc.content,
            ROW_NUMBER() OVER (ORDER BY kc.combined_sim DESC, kc.id)::INT AS rank_ix
        FROM keyword_candidates kc
        ORDER BY kc.combined_sim DESC
        LIMIT match_count * 5
    ),
    -- pgvector 시맨틱 검색
    semantic_results AS (
        SELECT d.id, d.title, d.content,
            ROW_NUMBER() OVER (ORDER BY d.embedding <=> query_embedding)::INT AS rank_ix
        FROM kb_documents d
        ORDER BY d.embedding <=> query_embedding
        LIMIT match_count * 5
    ),
    -- RRF 결합
    combined AS (
        SELECT COALESCE(k.id, s.id) AS id,
            COALESCE(k.title, s.title) AS title,
            COALESCE(k.content, s.content) AS content,
            k.rank_ix AS kw_rank, s.rank_ix AS sem_rank,
            COALESCE(keyword_weight / (rrf_k + k.rank_ix), 0.0) +
            COALESCE(semantic_weight / (rrf_k + s.rank_ix), 0.0) AS score
        FROM keyword_results k
        FULL OUTER JOIN semantic_results s ON k.id = s.id
    )
    SELECT c.id, c.title, c.content, c.kw_rank, c.sem_rank, c.score
    FROM combined c ORDER BY c.score DESC LIMIT match_count;
END;
$$;

SELECT * FROM hybrid_search_ko(
    'Aurora 인덱스 튜닝',           -- 키워드 (pg_bigm 검색용)
    '[0.023, -0.041, ...]'::vector, -- 임베딩 벡터 (pgvector 검색용)
    10,   -- 결과 수
    0.4,  -- 키워드 가중치
    0.6   -- 시맨틱 가중치
);

import boto3
import json
import os
import psycopg2
from pgvector.psycopg2 import register_vector

AWS_REGION = os.environ.get('AWS_REGION', 'us-west-2')
DB_SECRET_ID = os.environ['DB_SECRET_ID']  # 예: rds/aurora-pg/app 또는 시크릿 ARN

bedrock_runtime = boto3.client('bedrock-runtime', region_name=AWS_REGION)
secrets_client = boto3.client('secretsmanager', region_name=AWS_REGION)


def get_db_credentials(secret_id: str) -> dict:
    """AWS Secrets Manager에서 DB 접속 정보를 조회합니다.

    Aurora/RDS 관리형 시크릿은 다음 키를 JSON 형태로 포함합니다:
    username, password, host, port, dbname, engine
    """
    response = secrets_client.get_secret_value(SecretId=secret_id)
    return json.loads(response['SecretString'])


def connect_to_aurora() -> psycopg2.extensions.connection:
    """Secrets Manager에서 인증 정보를 가져와 Aurora PostgreSQL에 연결합니다."""
    creds = get_db_credentials(DB_SECRET_ID)
    conn = psycopg2.connect(
        host=creds['host'],
        port=creds.get('port', 5432),
        database=creds.get('dbname', 'mydb'),
        user=creds['username'],
        password=creds['password'],
        sslmode='require',       # Aurora는 SSL/TLS 연결 권장
        connect_timeout=5,
    )
    register_vector(conn)
    return conn


# 커넥션을 모듈 로딩 시 1회 생성 (Lambda 실행 환경 재사용 시 유리)
conn = connect_to_aurora()


def generate_embedding(text: str) -> list[float]:
    response = bedrock_runtime.invoke_model(
        modelId='amazon.titan-embed-text-v2:0',
        contentType='application/json', accept='application/json',
        body=json.dumps({'inputText': text, 'dimensions': 1024, 'normalize': True})
    )
    return json.loads(response['body'].read())['embedding']


# 참고: Bedrock 임베딩 호출은 매 검색마다 지연(~100ms)과 비용이 발생합니다.
# 프로덕션에서는 자주 사용되는 쿼리의 임베딩을 캐싱하는 것을 권장합니다.

def hybrid_search_korean(query: str, limit: int = 5,
                         keyword_weight: float = 0.4,
                         semantic_weight: float = 0.6):
    """한국어 하이브리드 검색 (pg_bigm + pgvector + RRF)"""
    query_embedding = generate_embedding(query)
    with conn.cursor() as cur:
        cur.execute(
            "SELECT * FROM hybrid_search_ko(%s, %s::vector, %s, %s, %s)",
            (query, query_embedding, limit, keyword_weight, semantic_weight)
        )
        return [
            {'id': r[0], 'title': r[1], 'content': r[2],
             'keyword_rank': r[3], 'semantic_rank': r[4], 'rrf_score': float(r[5])}
            for r in cur.fetchall()
        ]


def rag_answer_korean(question: str):
    """한국어 RAG: 하이브리드 검색 + LLM 응답 생성"""
    has_technical_term = any(
        term in question
        for term in ['pgvector', 'Aurora', 'HNSW', 'PostgreSQL', 'Lambda', 'RDS']
    )
    kw_weight = 0.6 if has_technical_term else 0.3
    sem_weight = 1.0 - kw_weight

    docs = hybrid_search_korean(question, limit=5,
                                keyword_weight=kw_weight, semantic_weight=sem_weight)
    context = "\n\n".join([
        f"[{d['title']}] (kw:{d['keyword_rank']}, sem:{d['semantic_rank']}, "
        f"RRF:{d['rrf_score']:.4f})\n{d['content']}"
        for d in docs
    ])

    response = bedrock_runtime.invoke_model(
        modelId='anthropic.claude-sonnet-4-5-20250929-v1:0',
        contentType='application/json', accept='application/json',
        body=json.dumps({
            'anthropic_version': 'bedrock-2023-05-31', 'max_tokens': 2048,
            'messages': [{'role': 'user', 'content':
                f"다음 컨텍스트를 기반으로 질문에 한국어로 답변해 주세요.\n\n"
                f"컨텍스트:\n{context}\n\n질문: {question}\n\n답변:"
            }]
        })
    )
    result = json.loads(response['body'].read())
    return {'answer': result['content'][0]['text'], 'sources': docs}


# 사용 예시
result = rag_answer_korean("Aurora PostgreSQL에서 HNSW 인덱스를 어떻게 튜닝하나요?")
print(f"답변: {result['answer']}")