_Geulbeot/04. design_agent/scripts/build_block_index.py

"""블록 라이브러리 FAISS 인덱스 빌드 스크립트.

catalog.yaml의 46개 블록을 임베딩하여 FAISS 인덱스를 생성한다.
블록 추가/수정 시 이 스크립트를 다시 실행하면 인덱스가 갱신된다.

사용법:
    python scripts/build_block_index.py

산출물:
    data/block_index.faiss      — FAISS 벡터 인덱스
    data/block_metadata.json    — 인덱스 순서 → 블록 매핑
"""
from __future__ import annotations

import json
import logging
import sys
from pathlib import Path

import faiss
import numpy as np
import yaml
from sentence_transformers import SentenceTransformer

logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
logger = logging.getLogger(__name__)

# Kei persona와 동일 모델 (1024차원, 한국어 최적화)
EMBEDDING_MODEL = "BAAI/bge-m3"

PROJECT_ROOT = Path(__file__).parent.parent
CATALOG_PATH = PROJECT_ROOT / "templates" / "catalog.yaml"
INDEX_PATH = PROJECT_ROOT / "data" / "block_index.faiss"
META_PATH = PROJECT_ROOT / "data" / "block_metadata.json"


def load_catalog() -> list[dict]:
    """catalog.yaml에서 블록 목록을 로드한다."""
    if not CATALOG_PATH.exists():
        logger.error(f"catalog.yaml 없음: {CATALOG_PATH}")
        sys.exit(1)

    with open(CATALOG_PATH, encoding="utf-8") as f:
        data = yaml.safe_load(f)

    blocks = data.get("blocks", [])
    if not blocks:
        logger.error("catalog.yaml에 블록이 없습니다.")
        sys.exit(1)

    logger.info(f"catalog 로드: {len(blocks)}개 블록")
    return blocks


def build_search_texts(blocks: list[dict]) -> list[str]:
    """각 블록의 검색용 텍스트를 생성한다.

    name + visual + when을 조합하여 검색 쿼리와 매칭되도록 한다.
    not_for는 네거티브이므로 검색 텍스트에 포함하지 않는다.
    """
    texts = []
    for block in blocks:
        parts = [
            block.get("name", ""),
            block.get("visual", ""),
            block.get("when", ""),
        ]
        text = ". ".join(p.strip() for p in parts if p.strip())
        texts.append(text)
    return texts


def build_index(texts: list[str]) -> tuple[faiss.IndexFlatIP, np.ndarray]:
    """텍스트를 임베딩하고 FAISS 인덱스를 생성한다."""
    logger.info(f"임베딩 모델 로딩: {EMBEDDING_MODEL}")
    model = SentenceTransformer(EMBEDDING_MODEL)

    logger.info(f"{len(texts)}개 텍스트 임베딩 중...")
    embeddings = model.encode(
        texts,
        normalize_embeddings=True,  # 코사인 유사도를 위해 정규화
        show_progress_bar=True,
    )
    embeddings = np.array(embeddings, dtype=np.float32)

    dim = embeddings.shape[1]
    logger.info(f"임베딩 완료: {embeddings.shape[0]}개 × {dim}차원")

    # Inner Product = 정규화된 벡터에서 코사인 유사도
    index = faiss.IndexFlatIP(dim)
    index.add(embeddings)

    logger.info(f"FAISS 인덱스 생성: {index.ntotal}개 벡터, {dim}차원")
    return index, embeddings


def save_metadata(blocks: list[dict]) -> None:
    """블록 메타데이터를 인덱스 순서대로 저장한다."""
    metadata = []
    for block in blocks:
        metadata.append({
            "id": block["id"],
            "name": block.get("name", ""),
            "template": block.get("template", ""),
            "category": block.get("template", "").split("/")[1] if "/" in block.get("template", "") else "",
            "height_cost": block.get("height_cost", "medium"),
            "visual": block.get("visual", ""),
            "when": block.get("when", ""),
            "not_for": block.get("not_for", ""),
        })

    META_PATH.parent.mkdir(parents=True, exist_ok=True)
    with open(META_PATH, "w", encoding="utf-8") as f:
        json.dump(metadata, f, ensure_ascii=False, indent=2)

    logger.info(f"메타데이터 저장: {META_PATH} ({len(metadata)}개)")


def main():
    # 1. catalog 로드
    blocks = load_catalog()

    # 2. 검색용 텍스트 생성
    texts = build_search_texts(blocks)

    # 3. 임베딩 + FAISS 인덱스
    index, embeddings = build_index(texts)

    # 4. 저장
    INDEX_PATH.parent.mkdir(parents=True, exist_ok=True)
    faiss.write_index(index, str(INDEX_PATH))
    logger.info(f"FAISS 인덱스 저장: {INDEX_PATH}")

    save_metadata(blocks)

    # 5. 검증
    logger.info("--- 검증 ---")
    test_index = faiss.read_index(str(INDEX_PATH))
    with open(META_PATH, encoding="utf-8") as f:
        test_meta = json.load(f)

    assert test_index.ntotal == len(blocks), f"벡터 수 불일치: {test_index.ntotal} vs {len(blocks)}"
    assert len(test_meta) == len(blocks), f"메타데이터 수 불일치: {len(test_meta)} vs {len(blocks)}"

    logger.info(f"✅ 검증 통과: {test_index.ntotal}개 벡터, {len(test_meta)}개 메타데이터")

    # 6. 테스트 검색
    model = SentenceTransformer(EMBEDDING_MODEL)
    test_queries = [
        "A vs B 두 개념 비교",
        "연도별 정책 로드맵",
        "핵심 수치 KPI 통계",
    ]
    for query in test_queries:
        q_emb = model.encode([query], normalize_embeddings=True)
        scores, indices = test_index.search(np.array(q_emb, dtype=np.float32), 3)
        results = [test_meta[i]["id"] for i in indices[0]]
        logger.info(f"  검색 '{query}' → {results}")


if __name__ == "__main__":
    main()