app.py

#!/usr/bin/env python3
"""
허깅페이스 Spaces 배포용 메인 파일
소방 복무관리 RAG 챗봇
"""

import os
import sys
import argparse
import gradio as gr
import pandas as pd
from pathlib import Path
from typing import List, Dict, Tuple

# 현재 디렉토리를 Python 경로에 추가
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(current_dir)

# 모듈 임포트
from config import Config
from rag_chatbot import RAGChatbot
from openai_chatbot import OpenAIRAGChatbot
from document_processor import DocumentProcessor

class HuggingFaceApp:
    """허깅페이스 Spaces 배포용 앱 클래스"""

    def __init__(self):
        # 벡터 DB 타입에 따라 챗봇 선택
        if Config.VECTOR_DB_TYPE == "supabase":
            self.chatbot = OpenAIRAGChatbot()
        else:
            self.chatbot = RAGChatbot()
        self.is_initialized = False

        # 예시 질문 (허깅페이스 환경 최적화)
        self.example_questions = [
            "연차휴가 사용 방법을 알려주세요",
            "정규근무시간은 어떻게 되나요?",
            "당직근무 절차가 궁금합니다",
            "인사평가는 언제 진행되나요?",
            "파견근무 신청 방법",
            "복무규정 위반 시 처리"
        ]

        # 앱 초기화
        self._initialize_app()

    def _initialize_app(self):
        """앱 초기화"""
        try:
            print("Starting Fire Service Management RAG Chatbot...")

            # 문서 폴더 확인 및 샘플 데이터 생성
            self._ensure_documents()

            # 챗봇 초기화
            success = self.chatbot.initialize()
            if success:
                self.is_initialized = True
                print("Chatbot initialization completed")
            else:
                print("Chatbot initialization failed - running in template mode")

        except Exception as e:
            print(f"Initialization error: {str(e)}")
            # 오류가 있어도 앱은 계속 실행

    def _ensure_documents(self):
        """문서 폴더 및 샘플 데이터 확인"""
        docs_folder = Path("documents")
        docs_folder.mkdir(exist_ok=True)

        # 샘플 문서가 없으면 생성
        sample_files = list(docs_folder.glob("*.txt"))
        if not sample_files:
            print("Creating sample documents...")
            self._create_sample_documents()

    def _create_sample_documents(self):
        """샘플 복무관리 문서 생성"""
        sample_docs = [
            {
                "filename": "복무관리규정.txt",
                "content": """소방공무원 복무관리 규정

제1장 총칙
제1조 (목적)
이 규정은 소방공무원의 복무에 관한 기본사항을 규정하여 직무수행의 효율성을 높이고 조직의 발전에 기여함을 목적으로 한다.

제2조 (근무시간)
1. 정규근무시간은 09:00부터 18:00까지로 한다.
2. 점심시간은 12:00부터 13:00까지로 한다.
3. 토요일, 일요일 및 법정공휴일은 휴무일로 한다.

제3조 (연차휴가)
1. 연차휴가는 1년간 정상 근무한 자에게 15일을 부여한다.
2. 연차휴가 사용 시 3일 전까지 신청서를 제출해야 한다.
3. 부서장의 승인을 받아 사용하며, 긴급한 경우에는 사후 승인도 가능하다.

제4조 (당직근무)
1. 당직근무는 정규근무시간 외에 수행하는 근무를 말한다.
2. 당직자는 비상상황에 대비한 통신장비를 항시 점검해야 한다.
3. 당직 중에는 음주를 엄금하며, 직무 수행에 지장이 없는 행위만 가능하다.
"""
            },
            {
                "filename": "인사평가규정.txt",
                "content": """소방공무원 인사평가 규정

제1장 평가의 기본원칙
제1조 (평가목적)
소방공무원의 직무수행 능력과 성과를 객관적으로 평가하여 능력위주의 인사관리를 정립하고 공정한 보상 및 승진의 기초 자료로 활용한다.

제2조 (평가주기)
1. 정기평가는 연 1회 실시하며, 평가기간은 매년 1월 1일부터 12월 31일까지로 한다.
2. 수시평가는 특별한 사유가 있을 경우 실시할 수 있다.

제3조 (평가항목)
1. 직무수행 능력 (40점)
2. 업무 성과 (30점)
3. 근무 태도 (20점)
4. 협업 능력 (10점)

제4조 (평가등급)
- 수 (90점 이상)
- 우 (80점 이상 90점 미만)
- 양 (70점 이상 80점 미만)
- 가 (60점 이상 70점 미만)
- 미 (60점 미만)
"""
            },
            {
                "filename": "교육훈련.txt",
                "content": """소방공무원 교육훈련 안내

제1조 (교육목적)
소방공무원의 전문성 향상과 직무능력 개발을 위한 체계적인 교육훈련을 실시한다.

제2조 (필수교육)
1. 신임교육: 신규 임용자 대상 2주간 집체교육
2. 직무연수: 매년 1회, 직무별 전문교육
3. 안전교육: 분기별 1회, 안전사고 예방 교육

제3조 (선택교육)
1. 외국어 교육
2. 정보통신 기술 교육
3. 리더십 교육
4. 전문 자격증 취득 지원 교육

제4조 (교육신청)
1. 교육 희망자는 소속 기관을 통해 신청한다.
2. 신청 시기는 교육 시작일 1개월 전까지이다.
3. 업무에 지장이 없는 경우 우선 선발한다.
"""
            }
        ]

        for doc in sample_docs:
            file_path = Path("documents") / doc["filename"]
            try:
                with open(file_path, 'w', encoding='utf-8') as f:
                    f.write(doc["content"])
                print(f"✅ {doc['filename']} 생성 완료")
            except Exception as e:
                print(f"❌ {doc['filename']} 생성 실패: {str(e)}")

    def format_response(self, response) -> str:
        """응답을 Gradio 형식으로 변환"""
        try:
            # 메시지 형식으로 변환
            answer = response.answer

            # 신뢰도 표시
            confidence = getattr(response, 'confidence', 0.0)
            if confidence >= 0.8:
                confidence_emoji = "🟢"
            elif confidence >= 0.5:
                confidence_emoji = "🟡"
            else:
                confidence_emoji = "🔴"

            # 응답 시간
            response_time = getattr(response, 'response_time', 0.0)

            # 출처 정보
            sources = getattr(response, 'sources', [])
            source_text = ""
            if sources:
                source_text = "\n\n📚 **참고자료:**\n"
                for i, source in enumerate(sources[:3], 1):  # 최대 3개만 표시
                    source_name = source.get('source', '알 수 없음')
                    source_text += f"{i}. {source_name}\n"

            # 전체 응답
            full_response = f"""{answer}

---
{confidence_emoji} 신뢰도: {confidence:.1%}
⏱️ 응답시간: {response_time:.2f}초
📄 참고문서: {len(sources)}개{source_text}"""

            return full_response

        except Exception as e:
            return f"응답 형식 변환 중 오류 발생: {str(e)}"

    def chat_function(self, message: str, history: List[List[str]]) -> List[List[str]]:
        """채팅 함수"""
        if not message.strip():
            return history

        try:
            # 챗봇 응답 생성
            if self.is_initialized:
                response = self.chatbot.generate_answer(message, use_llm=False)
                answer = self.format_response(response)
            else:
                answer = "죄송합니다. 챗봇이 초기화되지 않았습니다. 페이지를 새로고침해주세요."

            # 히스토리에 추가
            history.append([message, answer])

        except Exception as e:
            error_msg = f"답변 생성 중 오류 발생: {str(e)}\n\n관리자에게 문의해주세요."
            history.append([message, error_msg])

        return history

    def create_demo(self):
        """Gradio 데모 생성"""
        # 커스텀 CSS
        custom_css = """
        .gradio-container {
            max-width: 900px !important;
            margin: auto !important;
        }
        .message.user {
            background-color: #e3f2fd;
            border-radius: 15px 15px 0 15px;
        }
        .message.assistant {
            background-color: #f1f8e9;
            border-radius: 15px 15px 15px 0;
        }
        """

        with gr.Blocks(
            title="소방 복무관리 RAG 챗봇",
            theme=gr.themes.Soft(),
            css=custom_css
        ) as demo:

            gr.Markdown("""
            # 🚒 소방 복무관리 RAG 챗봇

            소방공무원 복무관리 규정, 인사평가, 교육훈련 등 업무 관련 질문에 답변해 드립니다.

            💡 **사용 방법**: 아래에 복무관리 관련 질문을 입력하고 Enter 키를 누르세요.
            """)

            # 상태 표시
            with gr.Row():
                status_text = gr.HTML(
                    "✅ **챗봇 준비 완료** - 복무관리 관련 질문을 입력해주세요"
                )

            # 채팅 인터페이스
            chatbot = gr.Chatbot(
                height=500,
                placeholder="안녕하세요! 소방 복무관리에 대해 무엇이 궁금하신가요?",
                avatar_images=["👤", "🤖"]
            )

            # 입력 영역
            with gr.Row():
                msg = gr.Textbox(
                    placeholder="복무관리 관련 질문을 입력해주세요 (예: 연차휴가 사용 방법)",
                    container=False,
                    scale=7
                )
                submit_btn = gr.Button("전송", scale=1, variant="primary")
                clear_btn = gr.Button("초기화", scale=1)

            # 예시 질문
            gr.Markdown("### 💡 예시 질문 (클릭하면 자동 입력됩니다)")

            with gr.Row():
                with gr.Column():
                    gr.Button("연차휴가 사용 방법", size="sm").click(
                        fn=lambda: "연차휴가 사용 방법을 알려주세요",
                        outputs=msg
                    )
                    gr.Button("정규근무시간", size="sm").click(
                        fn=lambda: "정규근무시간은 어떻게 되나요?",
                        outputs=msg
                    )
                    gr.Button("당직근무 절차", size="sm").click(
                        fn=lambda: "당직근무 절차가 궁금합니다",
                        outputs=msg
                    )
                with gr.Column():
                    gr.Button("인사평가 기준", size="sm").click(
                        fn=lambda: "인사평가는 어떤 기준으로 이루어지나요?",
                        outputs=msg
                    )
                    gr.Button("교육훈련 안내", size="sm").click(
                        fn=lambda: "교육훈련 종류와 신청 방법을 알려주세요",
                        outputs=msg
                    )
                    gr.Button("파견근무 신청", size="sm").click(
                        fn=lambda: "파견근무 신청 방법을 알려주세요",
                        outputs=msg
                    )

            # 정보 섹션
            with gr.Accordion("📊 시스템 정보", open=False):
                if self.is_initialized:
                    stats = self.chatbot.get_stats()
                    vector_stats = stats.get("vector_store", {})
                    gr.Markdown(f"""
                    - **문서 수**: {vector_stats.get('total_documents', 0)}개
                    - **임베딩 모델**: {vector_stats.get('embedding_model', 'N/A')}
                    - **응답 모드**: 템플릿 기반
                    - **최대 검색 문서**: {Config.MAX_RETRIEVE_DOCS}개
                    """)
                else:
                    gr.Markdown("⚠️ 챗봇이 초기화되지 않았습니다.")

            # 이벤트 핸들러
            def user_input(user_message, history):
                """사용자 입력 처리"""
                return "", history + [[user_message, None]]

            def bot_response(history):
                """봇 응답 처리"""
                if history and history[-1][1] is None:
                    user_message = history[-1][0]
                    bot_message = self.chat_function(user_message, history[:-1])
                    if bot_message:
                        history[-1][1] = bot_message[-1][1]  # 마지막 응답만 가져오기
                    else:
                        history[-1][1] = "죄송합니다. 응답을 생성할 수 없습니다."
                return history

            # 메시지 전송 이벤트
            msg.submit(
                user_input,
                [msg, chatbot],
                [msg, chatbot],
                queue=False
            ).then(
                bot_response,
                chatbot,
                chatbot
            )

            submit_btn.click(
                user_input,
                [msg, chatbot],
                [msg, chatbot],
                queue=False
            ).then(
                bot_response,
                chatbot,
                chatbot
            )

            clear_btn.click(
                lambda: ([], ""),
                outputs=[chatbot, msg]
            )

        return demo

def main():
    """메인 실행 함수"""
    parser = argparse.ArgumentParser(description="소방 복무관리 RAG 챗봇")
    parser.add_argument("--share", action="store_true", help="공유 링크 생성")
    parser.add_argument("--port", type=int, default=7860, help="서버 포트")
    args = parser.parse_args()

    # 앱 생성
    app = HuggingFaceApp()
    demo = app.create_demo()

    # 실행
    print("🚀 허깅페이스 Spaces 앱 시작 중...")
    demo.launch(
        share=args.share,
        server_port=args.port,
        server_name="0.0.0.0",
        show_error=True,
        show_tips=False
    )

if __name__ == "__main__":
    main()