vector_store.py

import os
import pickle
import numpy as np
from typing import List, Dict, Tuple
from pathlib import Path
from sentence_transformers import SentenceTransformer
import faiss
from langchain_core.documents import Document
from config import Config

class VectorStore:
    """FAISS 기반 벡터 데이터베이스 클래스"""

    def __init__(self, embedding_model: str = None, cache_dir: str = None):
        self.embedding_model_name = embedding_model or Config.EMBEDDING_MODEL
        self.cache_dir = cache_dir or Config.VECTOR_DB_PATH
        self.model = None
        self.index = None
        self.documents = []
        self.doc_ids = []

        # 캐시 디렉토리 생성
        Path(self.cache_dir).mkdir(parents=True, exist_ok=True)

    def load_embedding_model(self):
        """임베딩 모델 로드"""
        if self.model is None:
            print(f"📥 임베딩 모델 로드: {self.embedding_model_name}")
            try:
                self.model = SentenceTransformer(self.embedding_model_name)
                print("✅ 임베딩 모델 로드 완료")
            except Exception as e:
                print(f"❌ 임베딩 모델 로드 실패: {str(e)}")
                print("🔄 다국어 모델로 대체 시도...")
                self.model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')

    def create_embeddings(self, texts: List[str]) -> np.ndarray:
        """텍스트 목록에 대한 임베딩 생성"""
        if self.model is None:
            self.load_embedding_model()

        print(f"🔄 {len(texts)}개 텍스트 임베딩 생성 중...")
        embeddings = self.model.encode(
            texts,
            batch_size=32,
            show_progress_bar=True,
            convert_to_numpy=True,
            normalize_embeddings=True
        )
        return embeddings

    def build_vector_index(self, documents: List[Document]) -> bool:
        """문서 목록으로부터 벡터 인덱스 구축"""
        if not documents:
            print("⚠️ 처리할 문서가 없습니다.")
            return False

        print(f"🏗️ {len(documents)}개 문서로 벡터 인덱스 구축 시작...")

        # 문서 저장
        self.documents = documents

        # 텍스트 추출
        texts = [doc.page_content for doc in documents]

        # 임베딩 생성
        embeddings = self.create_embeddings(texts)

        # FAISS 인덱스 생성
        dimension = embeddings.shape[1]
        self.index = faiss.IndexFlatIP(dimension)  # 내적 기반 유사도 검색

        # 임베딩 추가
        self.index.add(embeddings.astype('float32'))

        # 문서 ID 생성
        self.doc_ids = list(range(len(documents)))

        print(f"✅ 벡터 인덱스 구축 완료 (차원: {dimension}, 문서: {len(documents)})")

        # 인덱스 저장
        self.save_index()

        return True

    def search_similar(self, query: str, k: int = 5) -> List[Tuple[Document, float]]:
        """유사 문서 검색"""
        if self.index is None:
            print("⚠️ 벡터 인덱스가 생성되지 않았습니다.")
            return []

        if self.model is None:
            self.load_embedding_model()

        # 쿼리 임베딩 생성
        query_embedding = self.model.encode([query], normalize_embeddings=True)
        query_embedding = query_embedding.astype('float32')

        # 검색
        k = min(k, len(self.documents))
        similarities, indices = self.index.search(query_embedding, k)

        # 결과 변환
        results = []
        for i in range(k):
            idx = indices[0][i]
            similarity = similarities[0][i]

            if 0 <= idx < len(self.documents):
                doc = self.documents[idx]
                results.append((doc, float(similarity)))

        return results

    def save_index(self):
        """벡터 인덱스 및 문서 저장"""
        if self.index is None:
            return

        try:
            # FAISS 인덱스 저장
            index_path = os.path.join(self.cache_dir, "faiss_index.bin")
            faiss.write_index(self.index, index_path)

            # 문서 및 메타데이터 저장
            metadata_path = os.path.join(self.cache_dir, "metadata.pkl")
            metadata = {
                'documents': self.documents,
                'doc_ids': self.doc_ids,
                'embedding_model': self.embedding_model_name,
                'total_documents': len(self.documents)
            }

            with open(metadata_path, 'wb') as f:
                pickle.dump(metadata, f)

            print(f"💾 벡터 인덱스 저장 완료: {self.cache_dir}")

        except Exception as e:
            print(f"❌ 인덱스 저장 실패: {str(e)}")

    def load_index(self) -> bool:
        """저장된 벡터 인덱스 로드"""
        try:
            index_path = os.path.join(self.cache_dir, "faiss_index.bin")
            metadata_path = os.path.join(self.cache_dir, "metadata.pkl")

            if not os.path.exists(index_path) or not os.path.exists(metadata_path):
                return False

            # FAISS 인덱스 로드
            self.index = faiss.read_index(index_path)

            # 메타데이터 로드
            with open(metadata_path, 'rb') as f:
                metadata = pickle.load(f)

            self.documents = metadata['documents']
            self.doc_ids = metadata['doc_ids']
            self.embedding_model_name = metadata.get('embedding_model', Config.EMBEDDING_MODEL)

            # 임베딩 모델 로드
            self.load_embedding_model()

            print(f"📖 벡터 인덱스 로드 완료 (문서: {len(self.documents)}개)")
            return True

        except Exception as e:
            print(f"❌ 인덱스 로드 실패: {str(e)}")
            return False

    def get_stats(self) -> Dict:
        """벡터 데이터베이스 통계 정보"""
        if self.index is None:
            return {"status": "no_index"}

        return {
            "total_documents": len(self.documents),
            "embedding_model": self.embedding_model_name,
            "index_dimension": self.index.d,
            "cache_directory": self.cache_dir,
            "is_trained": self.index.is_trained if hasattr(self.index, 'is_trained') else True
        }

    def rebuild_if_needed(self, documents: List[Document], force_rebuild: bool = False) -> bool:
        """필요시 인덱스 재구축"""
        # 기존 인덱스가 있고 강제 재구축이 없는 경우
        if not force_rebuild and self.load_index():
            # 문서 개수가 크게 변경되지 않았으면 재사용
            if abs(len(self.documents) - len(documents)) / max(len(self.documents), 1) < 0.1:
                print("📦 기존 인덱스 재사용")
                return True

        print("🔄 벡터 인덱스 재구축")
        return self.build_vector_index(documents)

    def add_documents(self, new_documents: List[Document]) -> bool:
        """새 문서 추가 (동적 업데이트)"""
        if not new_documents:
            return False

        # 임베딩 생성
        new_texts = [doc.page_content for doc in new_documents]
        new_embeddings = self.create_embeddings(new_texts)

        if self.index is None:
            # 인덱스가 없으면 새로 생성
            return self.build_vector_index(new_documents)

        # 기존 인덱스에 추가
        self.index.add(new_embeddings.astype('float32'))

        # 문서 목록 업데이트
        start_id = len(self.documents)
        self.documents.extend(new_documents)
        self.doc_ids.extend(range(start_id, start_id + len(new_documents)))

        print(f"➕ {len(new_documents)}개 문서 추가 완료")

        # 저장
        self.save_index()
        return True

    def delete_document(self, doc_id: int) -> bool:
        """문서 삭제 (실제로는 인덱스 재구축 필요)"""
        if doc_id < 0 or doc_id >= len(self.documents):
            return False

        # 해당 문서 제외하고 재구축
        remaining_docs = [doc for i, doc in enumerate(self.documents) if i != doc_id]
        return self.build_vector_index(remaining_docs)

# 테스트용 함수
def test_vector_store():
    """벡터 데이터베이스 테스트"""
    from document_processor import DocumentProcessor

    # 문서 처리
    processor = DocumentProcessor()
    documents = processor.load_documents_from_folder("documents")

    if not documents:
        print("⚠️ 테스트할 문서가 없습니다.")
        return

    # 벡터 데이터베이스 생성
    vector_store = VectorStore()
    vector_store.build_vector_index(documents)

    # 검색 테스트
    test_queries = [
        "연차휴가 사용 방법",
        "근무시간은 어떻게 되나요?",
        "당직근무 절차"
    ]

    for query in test_queries:
        print(f"\n🔍 검색: {query}")
        results = vector_store.search_similar(query, k=3)

        for i, (doc, similarity) in enumerate(results):
            print(f"  {i+1}. 유사도: {similarity:.4f}")
            print(f"     내용: {doc.page_content[:100]}...")

if __name__ == "__main__":
    test_vector_store()