vector_store.py

import json
from typing import Dict, List
from pathlib import Path
import numpy as np
from datetime import datetime
from sentence_transformers import SentenceTransformer
from huggingface_hub import HfApi
import os

class VectorStore:
    def __init__(self):
        self.documents = []
        self.metadata = []  # 문서 메타데이터 저장
        self.model = SentenceTransformer('all-MiniLM-L6-v2')
        self.hf_api = HfApi()
        self.dataset_name = "bluewhale2025/parseai_202506"  # Hugging Face dataset 이름
        
        # 데이터셋이 없으면 생성
        try:
            self.hf_api.create_repo(
                repo_id=self.dataset_name,
                repo_type="dataset",
                private=True  # 개인 데이터셋으로 설정
            )
            print(f"데이터셋 {self.dataset_name} 생성 완료")
        except Exception as e:
            print(f"데이터셋 생성 중 오류 발생: {str(e)}")
        
    def add_document(self, text: str, metadata: Dict) -> None:
        """문서를 저장"""
        try:
            # 문서 저장
            self.documents.append(text)
            
            # 메타데이터 저장
            metadata["timestamp"] = str(datetime.now())
            self.metadata.append(metadata)
            
            # 벡터 생성
            vector = self.model.encode(text)
            
            # 디렉토리 생성
            os.makedirs("vectors", exist_ok=True)
            os.makedirs("metadata", exist_ok=True)
            
            # 파일 경로 설정
            doc_id = len(self.documents)
            vector_path = f"vectors/{doc_id}.npy"
            metadata_path = f"metadata/{doc_id}.json"
            
            # 임시 파일로 저장
            np.save(vector_path, vector)
            with open(metadata_path, 'w', encoding='utf-8') as f:
                json.dump(metadata, f)
            
            # Hugging Face에 업로드
            self.hf_api.upload_file(
                path_or_fileobj=vector_path,
                path_in_repo=vector_path,
                repo_id=self.dataset_name,
                repo_type="dataset"
            )
            
            self.hf_api.upload_file(
                path_or_fileobj=metadata_path,
                path_in_repo=metadata_path,
                repo_id=self.dataset_name,
                repo_type="dataset"
            )
            
            # 임시 파일 삭제
            os.remove(vector_path)
            os.remove(metadata_path)
            
        except Exception as e:
            raise Exception(f"문서 저장 중 오류 발생: {str(e)}")
    
    def search(self, query: str, top_k: int = 5) -> List[Dict]:
        """키워드 검색"""
        try:
            # 쿼리 벡터 생성
            query_vector = self.model.encode(query)
            
            # Hugging Face에서 모든 벡터 로드
            vectors = []
            metadata = []
            
            # 모든 벡터 파일 로드
            files = self.hf_api.list_repo_files(
                repo_id=self.dataset_name,
                repo_type="dataset"
            )
            
            # 파일 정렬 (1부터 시작)
            vector_files = sorted([f for f in files if f.startswith("vectors/")])
            metadata_files = sorted([f for f in files if f.startswith("metadata/")])
            
            if not vector_files or not metadata_files:
                return []
            
            # 파일 로드
            for vector_file, metadata_file in zip(vector_files, metadata_files):
                vector = np.load(self.hf_api.download_file(
                    repo_id=self.dataset_name,
                    filename=vector_file,
                    repo_type="dataset"
                ))
                vectors.append(vector)
                
                meta = json.load(self.hf_api.download_file(
                    repo_id=self.dataset_name,
                    filename=metadata_file,
                    repo_type="dataset"
                ))
                metadata.append(meta)
            
            # 유사도 계산
            similarities = cosine_similarity(vectors, [query_vector]).flatten()
            
            # 유사도 기반 정렬
            sorted_idx = np.argsort(similarities)[::-1][:top_k]
            
            # 결과 생성
            results = []
            for idx in sorted_idx:
                results.append({
                    "filename": metadata[idx]["filename"],
                    "total_pages": metadata[idx]["total_pages"],
                    "summary": metadata[idx]["summary"],
                    "timestamp": metadata[idx]["timestamp"],
                    "similarity": float(similarities[idx])
                })
            
            return results
            
        except Exception as e:
            raise Exception(f"검색 중 오류 발생: {str(e)}")
    
    def _save_metadata(self) -> None:
        """메타데이터 저장"""
        try:
            Path(self.metadata_path).parent.mkdir(parents=True, exist_ok=True)
            with open(self.metadata_path, 'w', encoding='utf-8') as f:
                json.dump({
                    "documents": self.documents,
                    "metadata": self.metadata
                }, f, ensure_ascii=False, indent=2)
        except Exception as e:
            raise Exception(f"메타데이터 저장 중 오류 발생: {str(e)}")
    
    def _load_metadata(self):
        """메타데이터 로드"""
        try:
            if Path(self.metadata_path).exists():
                with open(self.metadata_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    self.documents = data["documents"]
                    self.metadata = data["metadata"]
        except Exception as e:
            raise Exception(f"메타데이터 로드 중 오류 발생: {str(e)}")

    def load(self) -> None:
        """저장된 메타데이터 불러오기"""
        self._load_metadata()

# 싱글톤 인스턴스 생성
vector_store = VectorStore()