backend/app/rag_core.py

# ./backend/app/rag_core.py

from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
from typing import List, Dict, Tuple

# 전역 변수로 모델 로드 (앱 시작 시 한 번만 로드되도록)
# 네이버 클로바의 한국어 SentenceBERT 모델을 로드합니다.
try:
    # 네이버 클로바 HyperCLOVAX-SEED-Text-Instruct-0.5B 모델 로드
    model = SentenceTransformer('jhgan/ko-sroberta-multitask', device='cpu')
    print("INFO: 임베딩 모델 'jhgan/ko-sroberta-multitask' 로드 완료.")
except Exception as e:
    print(f"ERROR: 임베딩 모델 'jhgan/ko-sroberta-multitask' 로드 실패: {e}. 다국어 모델로 시도합니다.")
    # 대체 모델 로직 (필요하다면 유지하거나 제거할 수 있습니다)
    try:
        model = SentenceTransformer('sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2', device='cpu')
        print("INFO: 임베딩 모델 'sentence-transformers/paraphrase-multilingual-L12-v2' 로드 완료.")
    except Exception as e:
        print(f"ERROR: 대체 임베딩 모델 로드 실패: {e}. RAG 기능을 사용할 수 없습니다.")
        raise

async def perform_rag_query(chunks_with_timestamps: List[Dict], query: str, top_k: int = 5) -> List[Dict]:
    """
    제공된 텍스트 청크들과 쿼리를 사용하여 RAG(Retrieval-Augmented Generation) 검색을 수행합니다.
    현재는 임베딩 기반 유사도 검색만 수행하며, LLM 호출은 추후 추가됩니다.

    Args:
        chunks_with_timestamps: [{"text": "...", "timestamp": "...", "start_seconds": ...}] 형태의 리스트.
        query: 사용자 쿼리 문자열.
        top_k: 쿼리와 가장 유사한 상위 N개의 청크를 반환.

    Returns:
        쿼리와 가장 관련성 높은 상위 N개의 청크 (Dict) 리스트.
    """
    if not chunks_with_timestamps:
        print("WARNING: RAG 검색을 위한 텍스트 청크가 없습니다.")
        return []

    # 1. 텍스트 임베딩 생성
    # 모든 청크의 텍스트만 추출
    texts = [chunk["text"] for chunk in chunks_with_timestamps]
    print(f"INFO: 총 {len(texts)}개의 텍스트 청크 임베딩 시작.")
    
    # 모델의 encode 메서드는 비동기가 아니므로, 직접 호출.
    # 만약 시간이 오래 걸린다면 FastAPI의 `run_in_threadpool` 등을 고려.
    try:
        chunk_embeddings = model.encode(texts, convert_to_numpy=True)
        print("INFO: 텍스트 청크 임베딩 완료.")
    except Exception as e:
        print(f"ERROR: 텍스트 청크 임베딩 중 오류 발생: {e}")
        return []

    # 2. FAISS 인덱스 생성 및 청크 임베딩 추가
    dimension = chunk_embeddings.shape[1] # 임베딩 벡터의 차원
    index = faiss.IndexFlatL2(dimension) # L2 유클리드 거리를 사용하는 간단한 인덱스
    index.add(chunk_embeddings)
    print("INFO: FAISS 인덱스 생성 및 임베딩 추가 완료.")

    # 3. 쿼리 임베딩
    query_embedding = model.encode([query], convert_to_numpy=True)
    print("INFO: 쿼리 임베딩 완료.")

    # 4. 유사도 검색 (FAISS)
    # D: 거리 (Distance), I: 인덱스 (Index)
    distances, indices = index.search(query_embedding, top_k)
    print(f"INFO: FAISS 유사도 검색 완료. 상위 {top_k}개 결과.")

    retrieved_chunks = []

    # ✨✨✨ 유사도 임계값 설정 및 필터링 추가 ✨✨✨
    # 이 값은 실험을 통해 최적의 값을 찾아야 합니다.
    # 거리가 낮을수록 유사하므로, 이 값보다 '거리(score)'가 낮아야만 결과를 포함합니다.
    # 예를 들어, 0.5는 '거리가 0.5 미만인 경우에만 결과에 포함하라'는 의미입니다.
    # 거리가 0이면 완벽히 일치합니다.
    MIN_DISTANCE_THRESHOLD = 150 # 예시 값: 이 값보다 거리가 작아야 합니다 (더 유사해야 함)

    for i in range(len(indices[0])):
        idx = indices[0][i]
        original_chunk = chunks_with_timestamps[idx]
        score = float(distances[0][i]) # FAISS에서 반환된 거리 값 (낮을수록 유사)

        # 설정된 임계값보다 거리가 작을 때만 (즉, 유사도가 높을 때만) 결과에 포함
        if score < MIN_DISTANCE_THRESHOLD:
            retrieved_chunks.append({
                "text": original_chunk["text"],
                "timestamp": original_chunk["timestamp"],
                "score": score
            })
        else:
            # 디버깅용: 임계값 때문에 제외된 청크를 로그로 확인
            print(f"DEBUG: 유사도 임계값({MIN_DISTANCE_THRESHOLD:.4f}) 초과로 제외된 청크 (거리: {score:.4f}): {original_chunk['text'][:50]}...")
    
    # 거리가 작은 순서(유사도가 높은 순서)로 정렬하여 반환
    retrieved_chunks.sort(key=lambda x: x['timestamp'])

    print(f"DEBUG: 최종 검색된 청크 수: {len(retrieved_chunks)}")
    return retrieved_chunks