gradio_interface.py

import gradio as gr
import time
from typing import List, Dict, Tuple
import os
from rag_chatbot import RAGChatbot, ChatResponse

class GradioInterface:
    """Gradio 웹 인터페이스 클래스"""

    def __init__(self):
        self.chatbot = RAGChatbot()
        self.is_initialized = False

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

    def initialize_chatbot(self, docs_folder: str = None, force_rebuild: bool = False) -> str:
        """챗봇 초기화"""
        try:
            success = self.chatbot.initialize(docs_folder, force_rebuild)
            if success:
                self.is_initialized = True
                return "✅ RAG 챗봇이 성공적으로 초기화되었습니다!"
            else:
                return "❌ 챗봇 초기화에 실패했습니다. documents 폴더에 파일이 있는지 확인해주세요."
        except Exception as e:
            return f"❌ 초기화 중 오류 발생: {str(e)}"

    def format_chat_response(self, response: ChatResponse) -> Tuple[str, str]:
        """챗봇 응답을 채팅 형식으로 변환"""
        # 메인 답변
        answer_html = response.answer.replace('\n', '<br>')

        # 신뢰도 색상
        if response.confidence >= 0.8:
            confidence_color = "green"
            confidence_text = "높음"
        elif response.confidence >= 0.5:
            confidence_color = "orange"
            confidence_text = "보통"
        else:
            confidence_color = "red"
            confidence_text = "낮음"

        # 정보 메시지
        info_html = f"""
        <div style="background-color: #f0f0f0; padding: 10px; border-radius: 5px; margin-top: 10px;">
            <strong>📊 답변 정보</strong><br>
            • 신뢰도: <span style="color: {confidence_color}; font-weight: bold;">{confidence_text} ({response.confidence:.2%})</span><br>
            • 응답시간: {response.response_time:.2f}초<br>
            • 참고문서: {len(response.sources)}개
        </div>
        """

        # 출처 정보
        if response.sources:
            sources_html = "<br><strong>📚 참고자료:</strong><ul>"
            for i, source in enumerate(response.sources, 1):
                sources_html += f"<li><strong>{source['source']}</strong> (유사도: {source['similarity']})<br><small>{source['content'][:150]}...</small></li>"
            sources_html += "</ul>"
            info_html += sources_html

        return answer_html, info_html

    def chat_interface(self, message: str, history: List[List[str]]) -> List[List[str]]:
        """채팅 인터페이스 핸들러"""
        if not self.is_initialized:
            history.append([message, "⚠️ 챗봇이 초기화되지 않았습니다. 먼저 초기화를 눌러주세요."])
            return history

        if not message.strip():
            return history

        try:
            # 답변 생성
            response = self.chatbot.generate_answer(message, use_llm=False)  # 템플릿 모드로 가볍게
            answer, info = self.format_chat_response(response)

            # 전체 응답
            full_response = f"{answer}<br><br>{info}"

            history.append([message, full_response])

        except Exception as e:
            error_response = f"❌ 답변 생성 중 오류 발생: {str(e)}"
            history.append([message, error_response])

        return history

    def test_question(self, question: str) -> str:
        """예시 질문 테스트"""
        if not self.is_initialized:
            return "⚠️ 먼저 챗봇 초기화를 진행해주세요."

        try:
            response = self.chatbot.generate_answer(question, use_llm=False)
            answer, info = self.format_chat_response(response)
            return f"❓ 질문: {question}\n\n🤖 답변:\n{answer}\n\n{info}"
        except Exception as e:
            return f"❌ 테스트 중 오류 발생: {str(e)}"

    def get_chatbot_stats(self) -> str:
        """챗봇 통계 정보"""
        if not self.is_initialized:
            return "챗봇이 초기화되지 않았습니다."

        try:
            stats = self.chatbot.get_stats()
            stats_html = "<h3>🤖 챗봇 상태 정보</h3>"

            if stats.get("status") == "initialized":
                vector_stats = stats.get("vector_store", {})
                stats_html += f"""
                <ul>
                    <li>상태: <span style="color: green;">✅ 정상 작동 중</span></li>
                    <li>총 문서 수: {vector_stats.get('total_documents', 0)}개</li>
                    <li>임베딩 모델: {vector_stats.get('embedding_model', 'N/A')}</li>
                    <li>벡터 차원: {vector_stats.get('index_dimension', 'N/A')}</li>
                    <li>LLM 사용: {'✅' if stats.get('llm_available') else '❌ (템플릿 모드)'}</li>
                </ul>
                """
            else:
                stats_html += "<p>⚠️ 챗봇이 초기화되지 않았습니다.</p>"

            return stats_html

        except Exception as e:
            return f"❌ 통계 정보 조회 실패: {str(e)}"

    def create_interface(self):
        """Gradio 인터페이스 생성"""
        # 커스텀 CSS
        custom_css = """
        .chat-message {
            padding: 15px;
            border-radius: 10px;
            margin: 10px 0;
        }
        .user-message {
            background-color: #e3f2fd;
            border-left: 4px solid #2196f3;
        }
        .assistant-message {
            background-color: #f1f8e9;
            border-left: 4px solid #4caf50;
        }
        .info-box {
            background-color: #fff3e0;
            border: 1px solid #ffb74d;
            border-radius: 8px;
            padding: 12px;
            margin: 10px 0;
        }
        .stats-box {
            background-color: #f5f5f5;
            border-radius: 8px;
            padding: 15px;
            margin: 10px 0;
        }
        """

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

            gr.Markdown("# 🚒 소방 복무관리 RAG 챗봇")
            gr.Markdown("소방업무 복무관리 규정 및 절차에 대한 질문에 답변해 드립니다.")

            with gr.Tab("💬 채팅"):
                with gr.Row():
                    with gr.Column(scale=4):
                        chatbot = gr.Chatbot(
                            height=500,
                            avatar_images=["👤", "🤖"]
                        )
                        msg = gr.Textbox(
                            placeholder="복무관리 관련 질문을 입력해주세요 (예: 연차휴가 사용 방법)",
                            label="질문 입력",
                            submit_btn="전송"
                        )
                        with gr.Row():
                            submit_btn = gr.Button("💬 전송", variant="primary")
                            clear_btn = gr.Button("🗑️ 대화 초기화")

                    with gr.Column(scale=1):
                        gr.Markdown("### 🚀 빠른 시작")
                        init_btn = gr.Button("🔧 챗봇 초기화", variant="secondary")
                        init_status = gr.HTML("⏳ 초기화를 눌러주세요.")

                        gr.Markdown("### 💡 예시 질문")
                        for question in self.example_questions:
                            gr.Button(question, size="sm").click(
                                fn=lambda q=question: q,
                                outputs=msg
                            )

                        gr.Markdown("### 📊 상태 정보")
                        stats_info = gr.HTML("상태 정보를 확인해주세요.")

            with gr.Tab("🔍 테스트"):
                gr.Markdown("### 예시 질문 테스트")
                with gr.Row():
                    with gr.Column():
                        test_question = gr.Textbox(
                            placeholder="테스트할 질문을 입력하세요",
                            label="질문"
                        )
                        test_btn = gr.Button("🧪 테스트 실행", variant="primary")
                    with gr.Column():
                        test_result = gr.HTML("테스트 결과가 여기에 표시됩니다.")

                gr.Markdown("### 통계 정보")
                stats_display = gr.HTML(self.get_chatbot_stats())

            with gr.Tab("⚙️ 설정"):
                gr.Markdown("### 시스템 설정")
                with gr.Row():
                    with gr.Column():
                        docs_folder = gr.Textbox(
                            value="documents",
                            label="문서 폴더 경로"
                        )
                        force_rebuild = gr.Checkbox(
                            label="강제 재구축",
                            info="체크 시 기존 인덱스를 새로构建"
                        )
                    with gr.Column():
                        rebuild_btn = gr.Button("🔄 인덱스 재구축", variant="primary")

                gr.Markdown("### 시스템 정보")
                system_info = gr.HTML("""
                <div class="stats-box">
                    <h4>🤖 시스템 사양</h4>
                    <ul>
                        <li>임베딩 모델: jhgan/ko-sbert-nli</li>
                        <li>LLM 모델: beomi/Llama-3-Open-Ko-8B</li>
                        <li>청크 크기: 500</li>
                        <li>검색 문서 수: 3개</li>
                    </ul>
                </div>
                """)

            # 이벤트 핸들러
            def init_chatbot_wrapper():
                status = self.initialize_chatbot()
                return status, self.get_chatbot_stats()

            init_btn.click(
                fn=init_chatbot_wrapper,
                outputs=[init_status, stats_info]
            )

            def rebuild_wrapper(docs_path, force):
                return self.initialize_chatbot(docs_path, force)

            rebuild_btn.click(
                fn=rebuild_wrapper,
                inputs=[docs_folder, force_rebuild],
                outputs=stats_display
            )

            # 채팅 이벤트
            def chat_wrapper(message, history):
                return self.chat_interface(message, history), ""

            msg.submit(
                fn=chat_wrapper,
                inputs=[msg, chatbot],
                outputs=[chatbot, msg]
            )

            submit_btn.click(
                fn=chat_wrapper,
                inputs=[msg, chatbot],
                outputs=[chatbot, msg]
            )

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

            # 테스트 이벤트
            test_btn.click(
                fn=self.test_question,
                inputs=test_question,
                outputs=test_result
            )

            # 자동 새로고침 (통계 정보)
            interface.load(
                fn=self.get_chatbot_stats,
                outputs=stats_info,
                every=30  # 30초마다 새로고침
            )

        return interface

    def launch(self, share: bool = False, server_port: int = 7860):
        """인터페이스 실행"""
        interface = self.create_interface()

        print("🚀 Gradio 웹 인터페이스 시작 중...")
        print(f"📍 접속 주소: http://localhost:{server_port}")

        interface.launch(
            share=share,
            server_port=server_port,
            show_error=True,
            show_tips=True,
            inbrowser=True
        )

# 메인 실행
if __name__ == "__main__":
    # 웹 인터페이스 실행
    gradio_app = GradioInterface()
    gradio_app.launch(share=False, server_port=7860)