local_cpu_test.py

# local_cpu_test.py - 로컬 CPU 환경용 RAG 테스트 스크립트
"""
로컬 CPU 환경에서 RAG 시스템을 테스트하는 스크립트
허깅페이스 스페이스 배포 전 로컬 검증용
"""

import os
import sys
import time
from datetime import datetime

# 로컬 환경임을 명시적으로 설정
os.environ.pop('SPACE_ID', None)  # 허깅페이스 환경 변수 제거

def test_imports():
    """필수 라이브러리 import 테스트"""
    print("=" * 50)
    print("라이브러리 import 테스트")
    print("=" * 50)
    
    try:
        import torch
        print(f"✓ PyTorch {torch.__version__}")
        print(f"  CUDA 사용 가능: {torch.cuda.is_available()}")
        print(f"  디바이스: {'cuda' if torch.cuda.is_available() else 'cpu'}")
    except ImportError as e:
        print(f"✗ PyTorch import 실패: {e}")
        return False
    
    try:
        import sentence_transformers
        print(f"✓ sentence-transformers {sentence_transformers.__version__}")
    except ImportError as e:
        print(f"✗ sentence-transformers import 실패: {e}")
        return False
    
    try:
        import faiss
        print(f"✓ FAISS 사용 가능")
    except ImportError as e:
        print(f"✗ FAISS import 실패: {e}")
        return False
    
    try:
        import sklearn
        print(f"✓ scikit-learn {sklearn.__version__}")
    except ImportError as e:
        print(f"✗ scikit-learn import 실패: {e}")
        return False
    
    print("✓ 모든 라이브러리 import 성공")
    return True

def test_config():
    """설정 테스트"""
    print("\n" + "=" * 50)
    print("설정 테스트")
    print("=" * 50)
    
    try:
        from config import RAG_CONFIG, IS_HUGGINGFACE_SPACE
        
        print(f"허깅페이스 환경: {IS_HUGGINGFACE_SPACE}")
        print(f"임베딩 모델: {RAG_CONFIG['embedding_models'][0]}")
        print(f"배치 크기: {RAG_CONFIG['batch_size']}")
        print(f"Top-K: {RAG_CONFIG['top_k']}")
        print(f"임계값: {RAG_CONFIG['similarity_threshold']}")
        print(f"하이브리드 가중치: {RAG_CONFIG['hybrid_weights']}")
        
        return True
        
    except Exception as e:
        print(f"설정 로드 실패: {e}")
        return False

def test_embedding_model():
    """임베딩 모델 로드 테스트"""
    print("\n" + "=" * 50)
    print("임베딩 모델 테스트")
    print("=" * 50)
    
    try:
        from sentence_transformers import SentenceTransformer
        from config import RAG_CONFIG
        
        # 가장 가벼운 모델로 테스트
        model_name = 'paraphrase-multilingual-MiniLM-L12-v2'
        
        print(f"모델 로딩: {model_name}")
        start_time = time.time()
        
        model = SentenceTransformer(
            model_name,
            device='cpu',  # 명시적으로 CPU 사용
            cache_folder='./model_cache'
        )
        
        load_time = time.time() - start_time
        print(f"✓ 모델 로드 완료 ({load_time:.2f}초)")
        
        # 간단한 임베딩 테스트
        test_texts = ["취득세율이 얼마인가요?", "주택 구입시 세금은?"]
        
        start_time = time.time()
        embeddings = model.encode(test_texts, convert_to_numpy=True)
        encode_time = time.time() - start_time
        
        print(f"✓ 임베딩 생성 완료 ({encode_time:.2f}초)")
        print(f"  임베딩 형태: {embeddings.shape}")
        
        return True
        
    except Exception as e:
        print(f"임베딩 모델 테스트 실패: {e}")
        return False

def test_law_fetcher():
    """법령 페처 테스트"""
    print("\n" + "=" * 50)
    print("법령 페처 테스트")
    print("=" * 50)
    
    try:
        from law_fetcher import HFLawAPIFetcher
        
        fetcher = HFLawAPIFetcher()
        print(f"✓ 법령 페처 초기화 완료")
        print(f"  캐시 디렉토리: {fetcher.cache_dir}")
        print(f"  캐시된 법령: {len(fetcher.cache_info)}개")
        
        # 캐시가 있는지 확인
        if fetcher.cache_info:
            print("  캐시 정보:")
            for law_name, info in fetcher.cache_info.items():
                cached_at = info.get('cached_at', 'Unknown')
                data_size = info.get('data_size', 0)
                print(f"    - {law_name}: {data_size/1024:.1f}KB ({cached_at[:10]})")
        
        return True
        
    except Exception as e:
        print(f"법령 페처 테스트 실패: {e}")
        return False

def test_rag_system_minimal():
    """RAG 시스템 최소 테스트"""
    print("\n" + "=" * 50)
    print("RAG 시스템 최소 테스트")
    print("=" * 50)
    
    try:
        from rag_system import HFSpacesTaxRAG
        
        print("RAG 시스템 초기화 중...")
        start_time = time.time()
        
        rag = HFSpacesTaxRAG()
        
        init_time = time.time() - start_time
        print(f"✓ RAG 시스템 초기화 완료 ({init_time:.2f}초)")
        
        print(f"  환경: {'허깅페이스' if rag.is_huggingface_space else '로컬'}")
        print(f"  디바이스: {rag.device}")
        print(f"  임베딩 모델: {rag.embedding_model.get_sentence_embedding_dimension()}차원")
        
        # 벡터 DB 상태 확인
        if rag.vector_db and rag.documents:
            print(f"  벡터 DB: {rag.vector_db.ntotal}개 벡터")
            print(f"  문서 수: {len(rag.documents)}개")
            
            # 간단한 검색 테스트
            test_query = "취득세 세율"
            print(f"\n검색 테스트: '{test_query}'")
            
            start_time = time.time()
            results = rag.search(test_query, top_k=2)
            search_time = time.time() - start_time
            
            print(f"✓ 검색 완료 ({search_time:.2f}초)")
            print(f"  결과: {len(results)}개 문서")
            
            for i, result in enumerate(results):
                score = result['hybrid_score']
                doc_preview = result['document'][:50]
                print(f"    {i+1}. 점수: {score:.3f} - {doc_preview}...")
        else:
            print("  벡터 DB 없음 - 시스템 구축 필요")
        
        return True
        
    except Exception as e:
        print(f"RAG 시스템 테스트 실패: {e}")
        import traceback
        traceback.print_exc()
        return False

def run_full_test():
    """전체 테스트 실행"""
    print("로컬 CPU 환경 RAG 시스템 테스트")
    print(f"실행 시간: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"Python 버전: {sys.version}")
    
    tests = [
        ("라이브러리 Import", test_imports),
        ("설정", test_config),
        ("임베딩 모델", test_embedding_model),
        ("법령 페처", test_law_fetcher),
        ("RAG 시스템", test_rag_system_minimal)
    ]
    
    success_count = 0
    total_start = time.time()
    
    for test_name, test_func in tests:
        try:
            print(f"\n{'='*60}")
            print(f"테스트: {test_name}")
            print(f"{'='*60}")
            
            if test_func():
                success_count += 1
                print(f"✓ {test_name} 테스트 성공")
            else:
                print(f"✗ {test_name} 테스트 실패")
                
        except KeyboardInterrupt:
            print(f"\n사용자에 의해 테스트 중단")
            break
        except Exception as e:
            print(f"✗ {test_name} 테스트 중 오류: {e}")
    
    total_time = time.time() - total_start
    
    print(f"\n{'='*60}")
    print(f"테스트 결과 요약")
    print(f"{'='*60}")
    print(f"성공: {success_count}/{len(tests)}")
    print(f"총 소요시간: {total_time:.2f}초")
    
    if success_count == len(tests):
        print("✓ 모든 테스트 통과 - 허깅페이스 배포 준비 완료")
        return True
    else:
        print("✗ 일부 테스트 실패 - 문제 해결 후 재시도")
        return False

if __name__ == "__main__":
    # 명령행 인수 처리
    if len(sys.argv) > 1:
        if sys.argv[1] == '--imports':
            success = test_imports()
        elif sys.argv[1] == '--config':
            success = test_config()
        elif sys.argv[1] == '--embedding':
            success = test_embedding_model()
        elif sys.argv[1] == '--rag':
            success = test_rag_system_minimal()
        else:
            print("사용법: python local_cpu_test.py [--imports|--config|--embedding|--rag]")
            sys.exit(1)
    else:
        success = run_full_test()
    
    sys.exit(0 if success else 1)