rag_faiss.py

"""
rag_faiss.py - FAISS 기반 RAG 엔진
SentenceTransformer + FAISS를 사용한 알레르기 정보 검색

기존 Pinecone 기반 rag_tool.py를 대체합니다.
"""

import os
import numpy as np
import pandas as pd
import faiss
import json
from typing import List, Dict
from sentence_transformers import SentenceTransformer
from langchain.tools import tool


# ============================================
# 설정
# ============================================

# CSV 파일 경로 (프로젝트 루트 기준)
CSV_PATH = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), 
    "data", 
    "allergens_small_extended.csv"
)

# 임베딩 모델 (로컬에서 빠르고 가벼운 모델)
EMBEDDING_MODEL = "all-MiniLM-L6-v2"

# OCR 오타 정규화 맵
NORM_MAP = {
    "mieko": "milk",
    "soya": "soy",
    "wheatt": "wheat",
    "wheat flour": "wheat",
    "milk powder": "milk",
    "almond bits": "almond",
    "pork meat": "pork",
    # 한글 추가
    "우유": "milk",
    "밀가루": "wheat",
    "소맥분": "wheat",
    "대두": "soy",
    "콩": "soy",
    "땅콩": "peanut",
    "계란": "egg",
    "난백": "egg",
    "난황": "egg"
}


# ============================================
# 헬퍼 함수
# ============================================

def clean_and_tokenize(ocr_text: str) -> List[str]:
    """
    OCR 텍스트를 정규화하고 토큰화
    
    Args:
        ocr_text: OCR 결과 텍스트
    
    Returns:
        정규화된 성분 토큰 리스트
    """
    import re
    
    s = str(ocr_text).lower()
    s = re.sub(r"[^a-z0-9가-힣, ]", " ", s)  # 한글 지원
    
    # "ingredients" 또는 "원재료명" 이후만 추출
    if "ingredients" in s:
        s = s.split("ingredients")[-1]
    elif "원재료명" in s or "원재료" in s:
        s = s.split("원재료")[-1]
    
    # 쉼표로 분리
    tokens = [t.strip() for t in s.split(",") if t.strip()]
    
    # 정규화 맵 적용
    normalized = []
    for t in tokens:
        # 정규화 맵 체크
        matched = False
        for k, v in NORM_MAP.items():
            if k in t:
                t = v
                matched = True
                break
        normalized.append(t)
    
    return normalized


def load_allergen_db(csv_path: str = CSV_PATH) -> pd.DataFrame:
    """
    알레르기 CSV 데이터 로드
    
    Args:
        csv_path: CSV 파일 경로
    
    Returns:
        DataFrame
    """
    if not os.path.exists(csv_path):
        raise FileNotFoundError(
            f"Allergen CSV not found at {csv_path}\n"
            f"Please ensure the CSV file is in the correct location."
        )
    
    df = pd.read_csv(csv_path)
    
    # doc_text 컬럼 생성 (검색용)
    if "doc_text" not in df.columns:
        df["doc_text"] = (
            df["ingredient"].astype(str) + " - " + 
            df["allergy_type"].astype(str) + " - " + 
            df["notes"].astype(str)
        )
    
    return df


# ============================================
# RAG 엔진 클래스
# ============================================

class RagEngine:
    """FAISS 기반 RAG 엔진"""
    
    def __init__(self, df: pd.DataFrame = None, model_name: str = EMBEDDING_MODEL):
        """
        RAG 엔진 초기화
        
        Args:
            df: 알레르기 데이터프레임
            model_name: SentenceTransformer 모델명
        """
        self.df = df if df is not None else load_allergen_db()
        
        print("📦 임베딩 모델 로딩 중... (첫 실행시 모델 다운로드)")
        self.model = SentenceTransformer(model_name)
        
        # 문서 임베딩 생성
        print("🔄 알레르기 데이터베이스 임베딩 생성 중...")
        self.doc_texts = self.df["doc_text"].tolist()
        self.doc_embeddings = self.model.encode(
            self.doc_texts, 
            normalize_embeddings=True,
            show_progress_bar=False
        )
        self.doc_embeddings = np.array(self.doc_embeddings).astype("float32")
        
        # FAISS 인덱스 구축
        self.dimension = self.doc_embeddings.shape[1]
        self.index = faiss.IndexFlatL2(self.dimension)
        self.index.add(self.doc_embeddings)
        
        print(f"✅ FAISS 인덱스 구축 완료 ({self.index.ntotal}개 문서)")
    
    def search(self, token: str, top_k: int = 1) -> List[Dict]:
        """
        단일 토큰 검색
        
        Args:
            token: 검색할 성분명
            top_k: 반환할 결과 개수
        
        Returns:
            매칭 결과 리스트
        """
        # 토큰 임베딩
        emb = self.model.encode([token], normalize_embeddings=True).astype("float32")
        
        # FAISS 검색
        distances, indices = self.index.search(emb, top_k)
        
        results = []
        for dist, idx in zip(distances[0], indices[0]):
            # 거리가 너무 크면 매칭 안 됨
            if dist > 1.5:  # 임계값 (조정 가능)
                continue
            
            results.append({
                "token": token,
                "ingredient": self.df.loc[idx, "ingredient"],
                "allergy_type": self.df.loc[idx, "allergy_type"],
                "notes": self.df.loc[idx, "notes"],
                "distance": float(dist),
                "confidence": 1.0 - min(float(dist) / 2.0, 1.0)  # 거리를 신뢰도로 변환
            })
        
        return results
    
    def search_multiple(self, tokens: List[str], top_k: int = 1) -> List[Dict]:
        """
        여러 토큰 검색
        
        Args:
            tokens: 검색할 성분 리스트
            top_k: 각 토큰당 반환할 결과 개수
        
        Returns:
            모든 매칭 결과 리스트
        """
        all_results = []
        for token in tokens:
            matches = self.search(token, top_k=top_k)
            if matches:
                all_results.extend(matches)
        return all_results


# ============================================
# 싱글톤 엔진
# ============================================

_rag_engine = None

def get_rag_engine() -> RagEngine:
    """RAG 엔진 싱글톤"""
    global _rag_engine
    if _rag_engine is None:
        df = load_allergen_db()
        _rag_engine = RagEngine(df=df)
    return _rag_engine


# ============================================
# LangChain Tool 정의
# ============================================

@tool
def search_allergen_info(ingredients: List[str]) -> str:
    """
    식품 성분 리스트를 입력받아 FAISS에서 알레르기 정보를 검색합니다.
    
    Args:
        ingredients: 검색할 성분 리스트. 예: ["milk", "wheat", "soy", "sugar"]
    
    Returns:
        각 성분에 대한 알레르기 정보를 담은 JSON 문자열
    
    Example:
        >>> search_allergen_info(["milk", "wheat"])
        '[{"ingredient": "milk", "matched": true, ...}, ...]'
    """
    engine = get_rag_engine()
    
    # 정규화
    normalized_ingredients = []
    for ing in ingredients:
        ing_lower = ing.lower().strip()
        # 정규화 맵 적용
        for k, v in NORM_MAP.items():
            if k in ing_lower:
                ing_lower = v
                break
        normalized_ingredients.append(ing_lower)
    
    # 검색
    all_matches = engine.search_multiple(normalized_ingredients, top_k=1)
    
    # 결과 포맷팅 (기존 Pinecone 형식과 호환)
    results = []
    
    for ingredient in ingredients:
        # 해당 성분의 매칭 찾기
        matched = [m for m in all_matches if m["token"].lower() in ingredient.lower()]
        
        if matched:
            best_match = matched[0]
            results.append({
                "ingredient": ingredient,
                "matched": True,
                "matched_allergen": best_match["ingredient"],
                "allergen_name_ko": best_match["ingredient"],  # 한글명 (CSV에 추가 필요시)
                "allergen_type": best_match["allergy_type"],
                "severity": "HIGH" if best_match["allergy_type"] in ["dairy", "nut", "seafood"] else "MEDIUM",
                "cross_reactive": [],  # CSV에 추가 가능
                "symptoms": [],  # CSV에 추가 가능
                "alternative_names": [],
                "confidence": best_match["confidence"],
                "safe": False,
                "notes": best_match["notes"]
            })
        else:
            # 매칭 안 됨 (안전)
            results.append({
                "ingredient": ingredient,
                "matched": False,
                "matched_allergen": None,
                "safe": True,
                "confidence": 0.0,
                "note": "알려진 알레르기 성분 아님"
            })
    
    return json.dumps(results, ensure_ascii=False, indent=2)


# ============================================
# 통합 함수
# ============================================

def run_allergy_check(ocr_text: str) -> Dict:
    """
    AllerLens 통합 함수 (main.py용)
    
    Args:
        ocr_text: OCR 추출 텍스트
    
    Returns:
        {
            "ingredients_detected": [...],
            "allergy_types": [...],
            "notes": [...],
            "explanation_basic": "..."
        }
    """
    # 토큰화 및 정규화
    tokens = clean_and_tokenize(ocr_text)
    
    # RAG 검색
    result_json = search_allergen_info.invoke({"ingredients": tokens})
    matches = json.loads(result_json)
    
    # 결과 추출
    detected = [m["matched_allergen"] for m in matches if m.get("matched")]
    allergy_types = list({m["allergen_type"] for m in matches if m.get("matched")})
    notes = [m.get("notes", "") for m in matches if m.get("matched")]
    
    # 기본 설명 생성
    if detected:
        explanation = f"이 제품에는 {', '.join(detected)}이(가) 포함되어 있습니다. 알레르기 유형: {', '.join(allergy_types)}"
    else:
        explanation = "주요 알레르기 성분이 감지되지 않았습니다."
    
    return {
        "ingredients_detected": detected,
        "allergy_types": allergy_types,
        "notes": notes,
        "explanation_basic": explanation,
        "raw_matches": matches
    }


# ============================================
# 테스트 코드
# ============================================

if __name__ == "__main__":
    print("=== FAISS RAG 엔진 테스트 ===\n")
    
    # 테스트 케이스
    tests = [
        "Ingredients: sugar, mieko, soya, wheatt",
        "원재료명: 우유, 밀가루, 설탕",
        "Ingredients: wheat flour, MILK powder, alcohol, salt",
        "Ingredients: pea protein, almond bits, sugar, pork"
    ]
    
    for test_text in tests:
        print(f"\n📄 입력: {test_text}")
        result = run_allergy_check(test_text)
        print(f"✅ 결과:")
        print(f"   - 감지된 성분: {result['ingredients_detected']}")
        print(f"   - 알레르기 유형: {result['allergy_types']}")
        print(f"   - 설명: {result['explanation_basic']}")
        print()