openai_chatbot.py

import os
import time
from typing import List, Dict, Tuple
from dataclasses import dataclass
import openai
from supabase_vector_store import SupabaseVectorStore
from document_processor import DocumentProcessor
from config import Config

@dataclass
class ChatResponse:
    """챗봇 응답 결과 클래스"""
    answer: str
    sources: List[Dict]
    confidence: float
    response_time: float

class OpenAIRAGChatbot:
    """OpenAI 기반 RAG 챗봇"""

    def __init__(self):
        self.document_processor = DocumentProcessor(
            chunk_size=Config.CHUNK_SIZE,
            chunk_overlap=Config.CHUNK_OVERLAP
        )
        self.vector_store = None
        self.openai_client = None
        self.is_initialized = False

    def initialize(self, docs_folder: str = None, force_rebuild: bool = False) -> bool:
        """챗봇 초기화"""
        print("🤖 OpenAI 기반 소방 복무관리 RAG 챗봇 초기화 중...")

        # 1. OpenAI 클라이언트 초기화
        if not Config.OPENAI_API_KEY:
            print("❌ OpenAI API Key가 설정되지 않았습니다.")
            return False

        self.openai_client = openai.OpenAI(api_key=Config.OPENAI_API_KEY)

        # 2. 벡터 저장소 초기화
        try:
            self.vector_store = SupabaseVectorStore()
        except Exception as e:
            print(f"❌ 벡터 저장소 초기화 실패: {str(e)}")
            return False

        # 3. 문서 로드 및 처리
        docs_folder = docs_folder or Config.DOCS_FOLDER
        documents = self._load_documents(docs_folder)

        if not documents:
            print("❌ 처리할 문서가 없습니다. documents 폴더에 파일을 넣어주세요.")
            return False

        # 4. 벡터 데이터베이스 구축
        success = self.vector_store.rebuild_index(documents, force_rebuild, use_openai=True)
        if not success:
            print("❌ 벡터 데이터베이스 구축 실패")
            return False

        self.is_initialized = True
        print("✅ OpenAI RAG 챗봇 초기화 완료")
        return True

    def _load_documents(self, docs_folder: str) -> List:
        """문서 로드 및 처리"""
        if not os.path.exists(docs_folder):
            print(f"⚠️ 문서 폴더가 존재하지 않습니다: {docs_folder}")
            return []

        print(f"📂 문서 폴더: {docs_folder}")
        raw_documents = self.document_processor.load_documents_from_folder(docs_folder)
        processed_documents = self.document_processor.process_documents(raw_documents)

        print(f"✅ 총 {len(processed_documents)}개 문서 청크 생성 완료")
        return processed_documents

    def search_relevant_docs(self, query: str, k: int = 3) -> List[Tuple]:
        """관련 문서 검색"""
        if not self.is_initialized:
            print("⚠️ 챗봇이 초기화되지 않았습니다.")
            return []

        # 쿼리 전처리
        processed_query = self._preprocess_query(query)

        # 벡터 검색
        results = self.vector_store.search_similar(processed_query, k, use_openai=True)

        # 유사도 필터링
        filtered_results = [
            (doc, similarity) for doc, similarity in results
            if similarity > 0.3  # 최소 유사도 임계값
        ]

        return filtered_results

    def _preprocess_query(self, query: str) -> str:
        """쿼리 전처리"""
        import re

        # 불필요한 공백 제거
        query = re.sub(r'\s+', ' ', query.strip())

        # 복무관리 관련 키워드 강화
        keyword_mappings = {
            "연차": "연차휴가",
            "휴가": "휴가사용",
            "근무": "근무시간",
            "당직": "당직근무",
            "인사": "인사평가",
            "승진": "승진시험"
        }

        for keyword, enhanced in keyword_mappings.items():
            if keyword in query and enhanced not in query:
                query = query.replace(keyword, enhanced)

        return query

    def generate_answer(self, query: str) -> ChatResponse:
        """질문에 대한 답변 생성 (OpenAI 사용)"""
        start_time = time.time()

        if not self.is_initialized:
            return ChatResponse(
                answer="죄송합니다. 챗봇이 초기화되지 않았습니다. 관리자에게 문의해주세요.",
                sources=[],
                confidence=0.0,
                response_time=time.time() - start_time
            )

        # 1. 관련 문서 검색
        relevant_docs = self.search_relevant_docs(query, k=Config.MAX_RETRIEVE_DOCS)

        if not relevant_docs:
            return ChatResponse(
                answer="죄송합니다. 질문과 관련된 정보를 찾을 수 없습니다. 다른 방식으로 질문해주시거나 관련 부서에 문의해주시기 바랍니다.",
                sources=[],
                confidence=0.0,
                response_time=time.time() - start_time
            )

        # 2. OpenAI로 답변 생성
        answer = self._generate_openai_answer(query, relevant_docs)

        # 3. 출처 정보 준비
        sources = [
            {
                "source": doc.metadata.get("source", "알 수 없음"),
                "content": doc.page_content[:200] + "..." if len(doc.page_content) > 200 else doc.page_content,
                "similarity": f"{similarity:.4f}"
            }
            for doc, similarity in relevant_docs
        ]

        # 4. 신뢰도 계산
        confidence = min(sum(similarity for _, similarity in relevant_docs) / len(relevant_docs), 1.0)

        return ChatResponse(
            answer=answer,
            sources=sources,
            confidence=confidence,
            response_time=time.time() - start_time
        )

    def _generate_openai_answer(self, query: str, relevant_docs: List[Tuple]) -> str:
        """OpenAI로 답변 생성"""
        try:
            # 문맥 구성
            context = "\n\n".join([
                f"[출처 {i+1}] {doc.page_content}"
                for i, (doc, _) in enumerate(relevant_docs)
            ])

            # OpenAI API 호출
            messages = [
                {
                    "role": "system",
                    "content": f"""{Config.SYSTEM_PROMPT}

답변 시 다음 지침을 따르세요:
1. 반드시 아래 참고자료를 기반으로 답변하세요
2. 규정 조문이나 구체적인 절차를 명시하세요
3. 단계별 설명이 필요한 경우 번호로 구분해서 설명하세요
4. 필요한 서류나 양식을 구체적으로 안내하세요
5. 주의사항이나 중요 사항은 강조해주세요
6. 답변 마지막에 참고한 출처를 표시하세요"""
                },
                {
                    "role": "user",
                    "content": f"""[참고자료]
{context}

[질문]
{query}

위 참고자료를 바탕으로 질문에 답변해주세요. 정확하고 친절하게 설명해주세요."""
                }
            ]

            response = self.openai_client.chat.completions.create(
                model=Config.OPENAI_MODEL,
                messages=messages,
                max_tokens=2000,
                temperature=0.3,  # 더 일관된 답변을 위해 낮은 온도
                top_p=0.9
            )

            answer = response.choices[0].message.content.strip()
            return answer

        except Exception as e:
            print(f"⚠️ OpenAI 답변 생성 실패: {str(e)}")
            return self._generate_fallback_answer(query, relevant_docs)

    def _generate_fallback_answer(self, query: str, relevant_docs: List[Tuple]) -> str:
        """OpenAI 실패 시 대체 답변 생성"""
        top_doc, top_similarity = relevant_docs[0]

        answer = f"""📋 소방 복무관리 안내

질문: {query}

관련 정보:
{top_doc.page_content[:800]}...

📖 더 자세한 정보는 관련 규정 파일을 확인하시거나 담당 부서에 문의해주시기 바랍니다.

*참고자료 유사도: {top_similarity:.2%}*"""

        return answer

    def get_stats(self) -> Dict:
        """챗봇 통계 정보"""
        if not self.is_initialized:
            return {"status": "not_initialized"}

        vector_stats = self.vector_store.get_stats()

        return {
            "status": "initialized",
            "vector_store": vector_stats,
            "llm_provider": "openai",
            "llm_model": Config.OPENAI_MODEL,
            "embedding_model": Config.OPENAI_EMBEDDING_MODEL
        }

    def add_documents(self, documents: List) -> bool:
        """새 문서 추가"""
        if not self.is_initialized:
            print("⚠️ 챗봇이 초기화되지 않았습니다.")
            return False

        return self.vector_store.add_documents(documents, use_openai=True)

# 테스트용 함수
def test_openai_chatbot():
    """OpenAI RAG 챗봇 테스트"""
    # 환경 변수 확인
    if not Config.OPENAI_API_KEY:
        print("❌ OPENAI_API_KEY 환경 변수가 필요합니다.")
        return

    if not Config.SUPABASE_URL or not Config.SUPABASE_KEY:
        print("❌ SUPABASE_URL, SUPABASE_KEY 환경 변수가 필요합니다.")
        return

    # 챗봇 초기화
    chatbot = OpenAIRAGChatbot()
    success = chatbot.initialize()

    if not success:
        return

    # 테스트 질문
    test_questions = [
        "연차휴가는 어떻게 사용하나요?",
        "정규근무시간은 어떻게 되나요?",
        "당직근무가 무엇인가요?",
        "인사평가 절차가 궁금합니다."
    ]

    # 질문 테스트
    for question in test_questions:
        print(f"\n❓ 질문: {question}")
        response = chatbot.generate_answer(question)

        print(f"🤖 답변: {response.answer[:500]}...")
        print(f"📊 신뢰도: {response.confidence:.4f}")
        print(f"⏱️ 응답시간: {response.response_time:.4f}초")
        print(f"📚 출처: {len(response.sources)}개")

    # 통계 정보
    print(f"\n📈 챗봇 통계: {chatbot.get_stats()}")

if __name__ == "__main__":
    test_openai_chatbot()