scripts/quick_setup_test_db.py

"""
빠른 테스트용 벡터 DB 생성 스크립트
간단한 금융 관련 샘플 문서로 벡터 DB 생성 (1분 이내)
"""

import sys
import os
from pathlib import Path

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

from services.vector_store import VectorStore
from services.embedder import Embedder
from loguru import logger

# 샘플 금융/경제 문서 (한국어)
SAMPLE_DOCUMENTS = [
    {
        "text": "포트폴리오 다각화는 투자 위험을 분산시키는 중요한 전략입니다. 서로 다른 자산 클래스에 투자함으로써 특정 자산의 하락이 전체 포트폴리오에 미치는 영향을 최소화할 수 있습니다.",
        "metadata": {"source": "portfolio_theory.pdf", "page": 1, "title": "포트폴리오 이론"}
    },
    {
        "text": "현대 포트폴리오 이론(MPT)에 따르면, 효율적 투자선은 주어진 위험 수준에서 최대 수익을 제공하는 포트폴리오들의 집합입니다. 마코위츠는 이 이론으로 노벨 경제학상을 수상했습니다.",
        "metadata": {"source": "portfolio_theory.pdf", "page": 2, "title": "포트폴리오 이론"}
    },
    {
        "text": "금융위기는 대체로 과도한 레버리지, 자산 가격 거품, 시스템적 리스크의 축적으로 인해 발생합니다. 2008년 글로벌 금융위기는 서브프라임 모기지 시장의 붕괴에서 시작되었습니다.",
        "metadata": {"source": "financial_crisis.pdf", "page": 1, "title": "금융위기의 원인"}
    },
    {
        "text": "금융위기 예방을 위해서는 건전한 규제, 적절한 자본 요구사항, 그리고 시스템적 리스크 모니터링이 필수적입니다. 바젤III는 이러한 규제 강화의 대표적인 예입니다.",
        "metadata": {"source": "financial_crisis.pdf", "page": 2, "title": "금융위기의 원인"}
    },
    {
        "text": "효율적 시장 가설(EMH)은 주가가 모든 가용 정보를 반영한다고 주장합니다. 약형, 준강형, 강형 효율성의 세 가지 형태로 구분됩니다.",
        "metadata": {"source": "market_efficiency.pdf", "page": 1, "title": "시장 효율성"}
    },
    {
        "text": "행동재무학은 투자자들의 비합리적 행동이 시장에 미치는 영향을 연구합니다. 과잉반응, 과소반응, 군집행동 등이 대표적인 현상입니다.",
        "metadata": {"source": "behavioral_finance.pdf", "page": 1, "title": "행동재무학"}
    },
    {
        "text": "자산가격결정모형(CAPM)은 기대수익률이 시장 위험 프리미엄과 베타의 곱에 무위험 수익률을 더한 값이라고 설명합니다. 베타는 시장 대비 자산의 변동성을 나타냅니다.",
        "metadata": {"source": "asset_pricing.pdf", "page": 1, "title": "자산가격결정"}
    },
    {
        "text": "옵션 가격결정에는 블랙-숄즈 모형이 널리 사용됩니다. 이 모형은 주가의 로그정규분포 가정 하에 유럽형 옵션의 이론 가격을 계산합니다.",
        "metadata": {"source": "derivatives.pdf", "page": 1, "title": "파생상품"}
    },
    {
        "text": "위험 관리에서 VaR(Value at Risk)는 특정 신뢰수준과 기간에서 예상되는 최대 손실을 측정합니다. 금융기관의 리스크 관리에 널리 사용됩니다.",
        "metadata": {"source": "risk_management.pdf", "page": 1, "title": "위험 관리"}
    },
    {
        "text": "중앙은행의 통화정책은 금리 조정, 공개시장운영, 지급준비율 변경 등을 통해 경제 안정화를 추구합니다. 인플레이션 타겟팅이 주요 정책 프레임워크입니다.",
        "metadata": {"source": "monetary_policy.pdf", "page": 1, "title": "통화정책"}
    },
    {
        "text": "ESG 투자는 환경(Environmental), 사회(Social), 지배구조(Governance) 요소를 고려하는 투자 전략입니다. 지속가능한 투자의 중요성이 증가하고 있습니다.",
        "metadata": {"source": "esg_investing.pdf", "page": 1, "title": "ESG 투자"}
    },
    {
        "text": "암호화폐는 블록체인 기술을 기반으로 한 디지털 자산입니다. 비트코인, 이더리움 등이 대표적이며, 탈중앙화 금융(DeFi)의 기반이 되고 있습니다.",
        "metadata": {"source": "cryptocurrency.pdf", "page": 1, "title": "암호화폐"}
    },
    {
        "text": "신용평가는 채무자의 채무 이행 능력을 평가합니다. 신용등급은 투자 의사결정과 금리 결정에 중요한 영향을 미칩니다.",
        "metadata": {"source": "credit_rating.pdf", "page": 1, "title": "신용평가"}
    },
    {
        "text": "주식시장의 이상현상(anomaly)에는 규모효과, 가치효과, 모멘텀 효과 등이 있습니다. 이는 효율적 시장 가설에 대한 반례로 제시됩니다.",
        "metadata": {"source": "market_anomalies.pdf", "page": 1, "title": "시장 이상현상"}
    },
    {
        "text": "구조화 금융은 자산유동화증권(ABS), 담보부채권(CDO) 등 복잡한 금융상품을 포함합니다. 2008년 금융위기에서 중요한 역할을 했습니다.",
        "metadata": {"source": "structured_finance.pdf", "page": 1, "title": "구조화 금융"}
    },
]


def main():
    """테스트용 벡터 DB 빠르게 생성"""
    logger.info("=" * 80)
    logger.info("테스트용 벡터 DB 생성 시작...")
    logger.info("=" * 80)

    try:
        # 1. Embedder 초기화 (무료 sentence-transformers 사용)
        logger.info("1️⃣ Embedder 초기화 중...")
        embedder = Embedder(
            model_type="sentence-transformers",
            model_name="all-MiniLM-L6-v2"
        )
        logger.info(f"✅ Embedder 준비 완료 ({embedder.get_embedding_dimension()}차원)")

        # 2. Vector Store 초기화
        logger.info("2️⃣ Vector Store 초기화 중...")
        persist_dir = project_root / "data" / "chroma_db"
        persist_dir.mkdir(parents=True, exist_ok=True)

        vector_store = VectorStore(
            persist_directory=str(persist_dir),
            collection_name="financial_papers"
        )
        logger.info("✅ Vector Store 준비 완료")

        # 3. 샘플 문서 임베딩 생성
        logger.info(f"3️⃣ {len(SAMPLE_DOCUMENTS)}개 샘플 문서 임베딩 중...")
        texts = [doc["text"] for doc in SAMPLE_DOCUMENTS]

        embeddings = embedder.embed_batch(texts)
        logger.info(f"✅ 임베딩 생성 완료 ({len(embeddings)}개)")

        # 4. Vector Store에 추가 (chunks 형식으로 변환)
        logger.info("4️⃣ Vector Store에 문서 추가 중...")
        chunks = []
        for i, doc in enumerate(SAMPLE_DOCUMENTS):
            chunks.append({
                'text': doc['text'],
                'source_filename': doc['metadata']['source'],
                'source_filepath': f"test_data/{doc['metadata']['source']}",
                'chunk_id': i,  # 고유 ID를 위해 인덱스 사용
                'total_chunks': 1,
                'metadata': doc['metadata'],
                'page_count': doc['metadata'].get('page', 1)
            })

        vector_store.add_documents(
            chunks=chunks,
            embeddings=embeddings
        )
        logger.info("✅ 문서 추가 완료")

        # 5. 검증
        logger.info("5️⃣ 생성된 DB 검증 중...")
        count = vector_store.collection.count()
        logger.info("=" * 80)
        logger.info(f"✨ 테스트 벡터 DB 생성 완료!")
        logger.info(f"📊 총 문서: {count}개")
        logger.info(f"📁 위치: {persist_dir}")
        logger.info(f"🔍 컬렉션: financial_papers")
        logger.info("=" * 80)

        # 6. 간단한 테스트 검색 (선택사항)
        try:
            logger.info("\n🧪 테스트 검색 실행 중...")
            test_query = "포트폴리오 다각화는 무엇인가요?"
            query_embedding = embedder.embed_batch([test_query])[0]
            results = vector_store.search(
                query_embedding=query_embedding,
                top_k=3
            )

            logger.info(f"질문: {test_query}")
            logger.info(f"검색 결과: {len(results['documents'])}개 문서 발견")
            for i, (doc, distance) in enumerate(zip(results['documents'], results['distances']), 1):
                metadata = results['metadatas'][i-1] if i-1 < len(results['metadatas']) else {}
                title = metadata.get('title', 'Unknown')
                logger.info(f"  [{i}] {title} (거리: {distance:.4f})")
        except Exception as e:
            logger.warning(f"테스트 검색 실패 (무시 가능): {str(e)}")

        logger.info("\n✅ DB 생성 완료! 백엔드 서버를 실행할 준비가 되었습니다.")

    except Exception as e:
        logger.error(f"❌ 오류 발생: {str(e)}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()