scripts/reembed_with_bge.py

#!/usr/bin/env python
"""
bge-m3 재임베딩 스크립트
======================

기존 청크들을 bge-m3 모델로 재임베딩합니다.
로컬에서 실행하여 API 비용 없이 임베딩을 생성합니다.

사용법:
    python scripts/reembed_with_bge.py
    python scripts/reembed_with_bge.py --dry-run  # 테스트 모드
"""

import os
import sys
import json
from pathlib import Path
from typing import List, Dict, Any

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

from dotenv import load_dotenv
load_dotenv()

import click


@click.command()
@click.option('--dry-run', is_flag=True, help='실제 저장 없이 테스트')
@click.option('--limit', '-l', type=int, default=None, help='처리할 청크 수 제한')
@click.option('--batch-size', '-b', type=int, default=32, help='배치 크기')
def main(dry_run: bool, limit: int, batch_size: int):
    """기존 청크들을 bge-m3로 재임베딩"""
    
    print("🔄 bge-m3 재임베딩 시작")
    print("=" * 60)
    
    if dry_run:
        print("⚠️  DRY-RUN 모드: 실제 저장하지 않습니다.\n")
    
    # 1. bge-m3 모델 로드
    print("📦 bge-m3 모델 로딩 중...")
    from src.db.bge_embeddings import BGEEmbeddings
    embedder = BGEEmbeddings()
    print(f"✅ 모델 로드 완료 (차원: {embedder.EMBEDDING_DIMENSION})")
    
    # 2. Supabase 연결
    print("\n🔌 Supabase 연결 중...")
    from supabase import create_client
    
    url = os.getenv("SUPABASE_URL")
    key = os.getenv("SUPABASE_KEY")
    
    if not url or not key:
        print("❌ SUPABASE_URL 또는 SUPABASE_KEY가 설정되지 않았습니다.")
        return
    
    client = create_client(url, key)
    print("✅ Supabase 연결됨")
    
    # 3. 기존 청크 가져오기
    print("\n📥 기존 청크 로딩 중...")
    query = client.table("kb_chunks").select("chunk_id, content")
    
    if limit:
        query = query.limit(limit)
    
    result = query.execute()
    chunks = result.data or []
    print(f"   총 청크: {len(chunks)}개")
    
    if not chunks:
        print("❌ 처리할 청크가 없습니다.")
        return
    
    # 4. 배치 처리로 임베딩 생성
    print(f"\n📊 임베딩 생성 중... (배치 크기: {batch_size})")
    
    total_processed = 0
    total_errors = 0
    
    for i in range(0, len(chunks), batch_size):
        batch = chunks[i:i + batch_size]
        batch_contents = [c["content"] for c in batch]
        batch_ids = [c["chunk_id"] for c in batch]
        
        try:
            # 배치 임베딩 생성
            embeddings = embedder.embed_texts(batch_contents)
            
            # 각 청크 업데이트
            for j, (chunk_id, embedding) in enumerate(zip(batch_ids, embeddings)):
                if not dry_run:
                    try:
                        client.table("kb_chunks").update({
                            "embedding": embedding
                        }).eq("chunk_id", chunk_id).execute()
                    except Exception as e:
                        print(f"  ⚠️ {chunk_id} 업데이트 실패: {e}")
                        total_errors += 1
                        continue
                
                total_processed += 1
            
            # 진행 상황 출력
            progress = (i + len(batch)) / len(chunks) * 100
            print(f"   {i + len(batch):4d}/{len(chunks)} ({progress:5.1f}%) 완료")
            
        except Exception as e:
            print(f"  ⚠️ 배치 {i//batch_size + 1} 처리 실패: {e}")
            total_errors += len(batch)
    
    # 5. 결과 요약
    print("\n" + "=" * 60)
    print("📊 재임베딩 완료")
    print(f"   처리됨: {total_processed}개")
    print(f"   오류: {total_errors}개")
    
    if dry_run:
        print("\n⚠️  DRY-RUN 완료. 실제 저장하려면 --dry-run 옵션을 제거하세요.")
    else:
        print("\n✅ bge-m3 재임베딩 완료!")
        
    # 6. 검증 테스트
    if not dry_run and total_processed > 0:
        print("\n🧪 검증 테스트...")
        
        # 한국어 검색 테스트
        test_query = "웨스틴 파르나스"
        query_embedding = embedder.embed_query(test_query)
        
        try:
            search_result = client.rpc(
                "match_kb_chunks",
                {
                    "query_embedding": query_embedding,
                    "match_threshold": 0.3,
                    "match_count": 3,
                    "filter_chain": None
                }
            ).execute()
            
            print(f"   검색어: {test_query}")
            if search_result.data:
                for r in search_result.data[:3]:
                    sim = r.get("similarity", 0)
                    content = r.get("content", "")[:50].replace("\n", " ")
                    print(f"   - {sim:.4f}: {content}...")
            else:
                print("   ⚠️ 검색 결과 없음 (RPC 함수 업데이트 필요할 수 있음)")
        except Exception as e:
            print(f"   ⚠️ 검색 테스트 실패: {e}")
            print("   → Supabase에서 match_kb_chunks 함수를 1024차원으로 업데이트하세요.")


if __name__ == "__main__":
    main()