rag.py

"""
RAG (Retrieval-Augmented Generation) 모듈
- 문서 임베딩 생성
- 관련 문서 검색
- LLM 응답 생성 (OpenAI)
"""

from sentence_transformers import SentenceTransformer
from config import OPENAI_API_KEY, OPENAI_MODEL, EMBEDDING_MODEL
from database import search_documents
from typing import List, Dict

# 임베딩 모델 로드 (한국어 지원)
embedding_model = SentenceTransformer(EMBEDDING_MODEL)

# OpenAI 클라이언트 (lazy loading으로 gradio 충돌 방지)
_openai_client = None


def get_openai_client():
    """OpenAI 클라이언트 lazy loading"""
    global _openai_client
    if _openai_client is None and OPENAI_API_KEY:
        from openai import OpenAI
        _openai_client = OpenAI(api_key=OPENAI_API_KEY)
    return _openai_client


def get_embedding(text: str) -> List[float]:
    """텍스트를 벡터로 변환"""
    return embedding_model.encode(text).tolist()


def retrieve_context(query: str, top_k: int = 3) -> str:
    """질문과 관련된 문서 검색"""
    query_embedding = get_embedding(query)
    documents = search_documents(query_embedding, top_k)

    if not documents:
        return ""

    context_parts = []
    for doc in documents:
        context_parts.append(f"[{doc.get('title', '문서')}]\n{doc.get('content', '')}")

    return "\n\n".join(context_parts)


def generate_response(query: str, context: str = "", history: List[Dict] = None) -> str:
    """OpenAI를 사용하여 응답 생성"""
    client = get_openai_client()

    if not client:
        return "OpenAI API 연결이 필요합니다. API 키를 확인해주세요."

    # 시스템 프롬프트
    system_prompt = """당신은 소방 복무관리 전문 AI 어시스턴트입니다.
소방공무원의 복무, 근무, 휴가, 당직 등에 관한 질문에 정확하고 친절하게 답변합니다.
제공된 참고 자료를 바탕으로 답변하되, 자료에 없는 내용은 일반적인 지식으로 보완합니다.
답변은 간결하고 명확하게 작성합니다."""

    # 메시지 구성
    messages = [{"role": "system", "content": system_prompt}]

    # 대화 기록 추가
    if history:
        for h in history[-5:]:  # 최근 5개 대화만
            messages.append({"role": "user", "content": h.get("user", "")})
            messages.append({"role": "assistant", "content": h.get("assistant", "")})

    # 현재 질문 (컨텍스트 포함)
    if context:
        user_content = f"[참고 자료]\n{context}\n\n[질문]\n{query}"
    else:
        user_content = query

    messages.append({"role": "user", "content": user_content})

    try:
        response = client.chat.completions.create(
            model=OPENAI_MODEL,
            messages=messages,
            max_tokens=1024,
            temperature=0.7
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"응답 생성 중 오류가 발생했습니다: {str(e)}"


def chat(query: str, history: List[Dict] = None) -> str:
    """RAG 기반 채팅 함수"""
    # 1. 관련 문서 검색
    context = retrieve_context(query)

    # 2. LLM 응답 생성
    response = generate_response(query, context, history)

    return response