scripts/check_vector_db.py

"""
Vector DB 상태 확인 스크립트

인덱싱이 완료된 후 벡터 DB의 내용을 확인합니다.
"""

import sys
from pathlib import Path

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

from dotenv import load_dotenv
from services.vector_store import VectorStore
from utils.config import settings


def main():
    """Vector DB 상태 확인"""
    load_dotenv()

    print("=" * 80)
    print("Vector DB 상태 확인")
    print("=" * 80)

    # Vector Store 초기화
    vector_store = VectorStore(
        persist_directory=settings.chroma_persist_directory,
        collection_name=settings.collection_name
    )

    # 통계 정보
    stats = vector_store.get_collection_stats()
    print(f"\n📊 기본 정보:")
    print(f"   컬렉션명: {stats['collection_name']}")
    print(f"   저장 경로: {stats['persist_directory']}")
    print(f"   전체 문서: {stats['total_documents']}개")
    print(f"   데이터 존재: {'✅ 예' if stats['has_data'] else '❌ 아니오'}")

    if not stats['has_data']:
        print("\n⚠️ Vector DB가 비어있습니다!")
        print("   python scripts/index_pdfs.py 를 먼저 실행하세요.")
        return

    # 샘플 데이터 확인
    print(f"\n📚 샘플 문서 확인:")
    sample = vector_store.collection.peek(limit=3)

    for i, (doc_id, doc, metadata) in enumerate(zip(
        sample['ids'],
        sample['documents'],
        sample['metadatas']
    ), 1):
        print(f"\n[{i}] {doc_id}")
        print(f"   출처: {metadata.get('source_filename', 'unknown')}")
        print(f"   제목: {metadata.get('title', 'N/A')}")
        print(f"   저자: {metadata.get('author', 'N/A')}")
        print(f"   내용: {doc[:150]}...")

    # 간단한 검색 테스트
    print(f"\n🔍 검색 테스트:")
    test_query = "financial crisis"
    print(f"   쿼리: '{test_query}'")

    results = vector_store.search_by_text(test_query, top_k=3)

    print(f"   결과: {len(results['documents'])}개 문서 발견")
    for i, (doc, metadata, distance) in enumerate(zip(
        results['documents'],
        results['metadatas'],
        results['distances']
    ), 1):
        similarity = 1 - distance
        print(f"\n   [{i}] {metadata.get('source_filename', 'unknown')}")
        print(f"       유사도: {similarity:.3f}")
        print(f"       내용: {doc[:100]}...")

    print("\n" + "=" * 80)
    print("✅ Vector DB 확인 완료!")
    print("=" * 80)


if __name__ == "__main__":
    main()