database.py

"""
Supabase 데이터베이스 연동 모듈
- 문서 저장 및 검색
- 벡터 유사도 검색 (pgvector)
"""

from supabase import create_client
from config import SUPABASE_URL, SUPABASE_KEY
from typing import List, Dict, Optional

# Supabase 클라이언트 초기화
supabase = create_client(SUPABASE_URL, SUPABASE_KEY) if SUPABASE_URL and SUPABASE_KEY else None


def save_document(title: str, content: str, embedding: List[float], category: str = "일반") -> bool:
    """문서를 데이터베이스에 저장"""
    if not supabase:
        return False

    try:
        supabase.table("documents").insert({
            "title": title,
            "content": content,
            "embedding": embedding,
            "category": category
        }).execute()
        return True
    except Exception as e:
        print(f"문서 저장 오류: {e}")
        return False


def search_documents(query_embedding: List[float], top_k: int = 3) -> List[Dict]:
    """벡터 유사도 기반 문서 검색"""
    if not supabase:
        return []

    try:
        # Supabase RPC 함수 호출 (pgvector 유사도 검색)
        result = supabase.rpc(
            "match_documents",
            {
                "query_embedding": query_embedding,
                "match_count": top_k
            }
        ).execute()
        return result.data if result.data else []
    except Exception as e:
        print(f"문서 검색 오류: {e}")
        return []


def get_all_documents(category: Optional[str] = None) -> List[Dict]:
    """모든 문서 조회"""
    if not supabase:
        return []

    try:
        query = supabase.table("documents").select("id, title, content, category")
        if category:
            query = query.eq("category", category)
        result = query.execute()
        return result.data if result.data else []
    except Exception as e:
        print(f"문서 조회 오류: {e}")
        return []


def save_chat_history(user_message: str, bot_response: str, session_id: str) -> bool:
    """채팅 기록 저장"""
    if not supabase:
        return False

    try:
        supabase.table("chat_history").insert({
            "session_id": session_id,
            "user_message": user_message,
            "bot_response": bot_response
        }).execute()
        return True
    except Exception as e:
        print(f"채팅 기록 저장 오류: {e}")
        return False