Vector DB 마스터 클래스 2편: HNSW 인덱스와 하이브리드 검색 최적화

import hnswlib
import numpy as np

def build_hnsw_index(vectors: np.ndarray, dim: int) -> hnswlib.Index:
    index = hnswlib.Index(space="cosine", dim=dim)

    index.init_index(
        max_elements=len(vectors),
        ef_construction=200,  # 높을수록 정확, 느린 인덱싱 (권장: 100~400)
        M=16,                 # 연결 수. 높을수록 정확, 많은 메모리 (권장: 8~48)
    )
    index.add_items(vectors)

    # 검색 시 ef: 높을수록 정확, 느린 쿼리 (권장: ef_construction 이상)
    index.set_ef(50)
    return index

# 정확도 vs 속도 트레이드오프 측정
def benchmark(index, queries, ground_truth, ef_values):
    for ef in ef_values:
        index.set_ef(ef)
        labels, _ = index.knn_query(queries, k=10)
        recall = compute_recall(labels, ground_truth)
        print(f"ef={ef:3d} → recall={recall:.3f}")

from rank_bm25 import BM25Okapi
import numpy as np

class HybridSearcher:
    def __init__(self, docs: list[str], embeddings: np.ndarray):
        tokenized = [d.split() for d in docs]
        self.bm25 = BM25Okapi(tokenized)
        self.embeddings = embeddings
        self.docs = docs

    def search(self, query: str, query_emb: np.ndarray,
               alpha: float = 0.5, top_k: int = 10) -> list:
        # BM25 점수 (키워드 매칭)
        bm25_scores = self.bm25.get_scores(query.split())
        bm25_norm = bm25_scores / (bm25_scores.max() + 1e-9)

        # 코사인 유사도 (의미 매칭)
        cosine_scores = self.embeddings @ query_emb
        cosine_norm = (cosine_scores + 1) / 2  # [-1,1] → [0,1]

        # 가중 합산 (RRF 방식도 가능)
        final = alpha * cosine_norm + (1 - alpha) * bm25_norm
        top_idx = np.argsort(final)[::-1][:top_k]
        return [(self.docs[i], float(final[i])) for i in top_idx]

# pgvector 예시 — 필터 후 벡터 검색
async def filtered_search(
    query_emb: list[float],
    department: str,
    date_from: str,
    limit: int = 10
) -> list:
    return await db.fetch("""
        SELECT id, title, content,
               1 - (embedding <=> $1::vector) AS similarity
        FROM documents
        WHERE department = $2
          AND created_at >= $3
          AND status = 'published'
        ORDER BY embedding <=> $1::vector
        LIMIT $4
    """, query_emb, department, date_from, limit)

import hashlib
import redis

class EmbeddingCache:
    def __init__(self, ttl: int = 3600):
        self.redis = redis.Redis()
        self.ttl = ttl

    def _key(self, text: str, model: str) -> str:
        h = hashlib.sha256(f"{model}:{text}".encode()).hexdigest()[:16]
        return f"emb:{h}"

    async def get_or_embed(self, text: str, model: str, embed_fn) -> list[float]:
        key = self._key(text, model)
        cached = self.redis.get(key)
        if cached:
            return json.loads(cached)
        emb = await embed_fn(text)
        self.redis.setex(key, self.ttl, json.dumps(emb))
        return emb