hgzero/rag/scripts/load_terms.py

"""
용어집 데이터 로딩 스크립트
"""
import sys
import json
import logging
from pathlib import Path
from typing import List

# 프로젝트 루트 디렉토리를 Python 경로에 추가
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from src.models.term import Term, DocumentSource
from src.db.postgres_vector import PostgresVectorDB
from src.utils.config import load_config, get_database_url
from src.utils.embedding import EmbeddingGenerator

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def load_terms_from_json(file_path: Path) -> List[Term]:
    """
    JSON 파일에서 용어 데이터 로드

    Args:
        file_path: JSON 파일 경로

    Returns:
        용어 리스트
    """
    logger.info(f"JSON 파일 로딩: {file_path}")

    with open(file_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    terms = []

    # 도메인별 데이터 처리
    for domain_data in data.get("terms", []):
        domain = domain_data.get("domain")
        source_type = domain_data.get("source_type")

        for term_data in domain_data.get("data", []):
            # DocumentSource 파싱
            doc_source = None
            if "document_source" in term_data:
                doc_source = DocumentSource(**term_data["document_source"])

            # Term 객체 생성
            term = Term(
                term_id=term_data["term_id"],
                term_name=term_data["term_name"],
                normalized_name=term_data["normalized_name"],
                category=term_data["category"],
                definition=term_data["definition"],
                context=term_data.get("context", ""),
                synonyms=term_data.get("synonyms", []),
                related_terms=term_data.get("related_terms", []),
                document_source=doc_source,
                confidence_score=term_data.get("confidence_score", 0.0),
                usage_count=term_data.get("usage_count", 0),
                last_updated=term_data.get("last_updated"),
                embedding=None  # 나중에 생성
            )

            terms.append(term)

    logger.info(f"  → {len(terms)}개 용어 로드 완료")
    return terms


def generate_embeddings(
    terms: List[Term],
    embedding_gen: EmbeddingGenerator
) -> List[Term]:
    """
    용어 임베딩 생성

    Args:
        terms: 용어 리스트
        embedding_gen: 임베딩 생성기

    Returns:
        임베딩이 추가된 용어 리스트
    """
    logger.info(f"임베딩 생성 시작: {len(terms)}개 용어")

    # 임베딩 텍스트 준비 (용어명 + 정의 + 맥락)
    texts = []
    for term in terms:
        text = f"{term.term_name}\n{term.definition}"
        if term.context:
            text += f"\n{term.context}"
        texts.append(text)

    # 배치 임베딩 생성
    embeddings = embedding_gen.generate_embeddings_batch(texts, batch_size=50)

    # 임베딩 추가
    for term, embedding in zip(terms, embeddings):
        term.embedding = embedding

    logger.info(f"  → 임베딩 생성 완료")
    return terms


def main():
    """메인 함수"""
    logger.info("=" * 60)
    logger.info("용어집 데이터 로딩 시작")
    logger.info("=" * 60)

    # 1. 설정 로드
    config = load_config(str(project_root / "config.yaml"))
    logger.info("✓ 설정 로드 완료")

    # 2. PostgreSQL 연결
    db_url = get_database_url(config)
    db = PostgresVectorDB(db_url)
    logger.info("✓ PostgreSQL 연결 완료")

    # 3. 데이터베이스 초기화
    db.init_database()
    logger.info("✓ 데이터베이스 초기화 완료")

    # 4. 임베딩 생성기 초기화
    azure_openai = config["azure_openai"]
    embedding_gen = EmbeddingGenerator(
        api_key=azure_openai["api_key"],
        endpoint=azure_openai["endpoint"],
        model=azure_openai["embedding_model"],
        dimension=azure_openai["embedding_dimension"],
        api_version=azure_openai["api_version"]
    )
    logger.info("✓ 임베딩 생성기 초기화 완료")

    # 5. 용어 데이터 로딩
    all_terms = []
    data_dir = project_root.parent / config["data"]["terms_dir"]

    for filename in config["data"]["terms_files"]:
        file_path = data_dir / filename
        if file_path.exists():
            terms = load_terms_from_json(file_path)
            all_terms.extend(terms)
        else:
            logger.warning(f"⚠ 파일 없음: {file_path}")

    logger.info(f"✓ 총 {len(all_terms)}개 용어 로드 완료")

    # 6. 임베딩 생성
    all_terms = generate_embeddings(all_terms, embedding_gen)

    # 7. 데이터베이스에 삽입
    logger.info(f"데이터베이스 삽입 시작: {len(all_terms)}개 용어")
    success_count = 0
    fail_count = 0

    for i, term in enumerate(all_terms, 1):
        if db.insert_term(term):
            success_count += 1
        else:
            fail_count += 1

        if i % 100 == 0:
            logger.info(f"  진행: {i}/{len(all_terms)} (성공: {success_count}, 실패: {fail_count})")

    logger.info(f"✓ 삽입 완료: 성공 {success_count}, 실패 {fail_count}")

    # 8. 통계 조회
    stats = db.get_stats()
    logger.info("=" * 60)
    logger.info("용어집 통계")
    logger.info("=" * 60)
    logger.info(f"전체 용어: {stats['total_terms']}개")
    logger.info(f"평균 신뢰도: {stats['avg_confidence']:.2f}")
    logger.info("\n카테고리별 통계:")
    for category, count in sorted(stats['by_category'].items(), key=lambda x: x[1], reverse=True):
        logger.info(f"  - {category}: {count}개")

    logger.info("=" * 60)
    logger.info("용어집 데이터 로딩 완료")
    logger.info("=" * 60)


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        logger.error(f"오류 발생: {str(e)}", exc_info=True)
        sys.exit(1)