"""
소방 복무관리 RAG 챗봇
깔끔하고 심플한 Gradio UI + 관리자 탭
"""
import gradio as gr
from rag import chat
from pdf_processor import process_pdf
from database import get_all_documents
import uuid
import os
# 세션별 대화 기록 저장
conversation_history = {}
# 관리자 비밀번호 (환경변수로 설정 권장)
ADMIN_PASSWORD = os.getenv("ADMIN_PASSWORD", "admin1234")
def respond(message: str, history: list, session_id: str) -> tuple:
"""채팅 응답 처리"""
if not message.strip():
return "", history
if session_id not in conversation_history:
conversation_history[session_id] = []
response = chat(message, conversation_history[session_id])
conversation_history[session_id].append({
"user": message,
"assistant": response
})
history.append((message, response))
return "", history
def clear_chat(session_id: str) -> tuple:
"""대화 초기화"""
if session_id in conversation_history:
conversation_history[session_id] = []
return [], ""
def verify_admin(password: str) -> tuple:
"""관리자 인증"""
if password == ADMIN_PASSWORD:
return gr.update(visible=False), gr.update(visible=True), "✅ 인증 성공"
return gr.update(visible=True), gr.update(visible=False), "❌ 비밀번호가 틀렸습니다"
def upload_pdf(file, title: str, category: str) -> str:
"""PDF 업로드 및 처리"""
if file is None:
return "❌ PDF 파일을 선택해주세요."
if not title.strip():
return "❌ 문서 제목을 입력해주세요."
result = process_pdf(
pdf_path=file.name,
title=title.strip(),
category=category
)
return result["message"]
def get_documents_list() -> list:
"""저장된 문서 목록 조회"""
docs = get_all_documents()
if not docs:
return [["데이터 없음", "-", "-"]]
return [[d.get("title", ""), d.get("category", ""), str(d.get("id", ""))] for d in docs]
# 커스텀 CSS
custom_css = """
.gradio-container {
max-width: 900px !important;
margin: auto !important;
font-family: 'Pretendard', -apple-system, BlinkMacSystemFont, sans-serif !important;
}
.header {
text-align: center;
padding: 20px 0;
border-bottom: 1px solid #e5e7eb;
margin-bottom: 20px;
}
.header h1 {
color: #dc2626;
font-size: 1.8rem;
font-weight: 700;
margin: 0;
}
.header p {
color: #6b7280;
font-size: 0.9rem;
margin-top: 8px;
}
.chatbot {
border-radius: 12px !important;
border: 1px solid #e5e7eb !important;
box-shadow: 0 1px 3px rgba(0,0,0,0.1) !important;
}
.input-area {
border-radius: 12px !important;
border: 1px solid #e5e7eb !important;
}
.primary-btn {
background-color: #dc2626 !important;
border: none !important;
border-radius: 8px !important;
color: white !important;
font-weight: 600 !important;
}
.primary-btn:hover {
background-color: #b91c1c !important;
}
.secondary-btn {
background-color: #f3f4f6 !important;
border: 1px solid #e5e7eb !important;
border-radius: 8px !important;
color: #374151 !important;
}
.admin-section {
background: #fef2f2;
border: 1px solid #fecaca;
border-radius: 12px;
padding: 20px;
margin: 10px 0;
}
.upload-section {
background: #f9fafb;
border: 1px solid #e5e7eb;
border-radius: 12px;
padding: 20px;
margin: 10px 0;
}
.footer {
text-align: center;
padding: 16px 0;
color: #9ca3af;
font-size: 0.8rem;
}
"""
# Gradio 앱 구성
with gr.Blocks(css=custom_css, title="소방 복무관리 챗봇", theme=gr.themes.Soft()) as app:
session_id = gr.State(value=str(uuid.uuid4()))
# 헤더
gr.HTML("""
""")
# 탭 구성
with gr.Tabs():
# ========== 채팅 탭 ==========
with gr.TabItem("💬 채팅"):
chatbot = gr.Chatbot(
label="",
height=450,
show_label=False,
container=True,
elem_classes=["chatbot"]
)
with gr.Row():
msg = gr.Textbox(
placeholder="메시지를 입력하세요...",
show_label=False,
container=False,
scale=9,
elem_classes=["input-area"]
)
submit_btn = gr.Button("전송", scale=1, elem_classes=["primary-btn"])
with gr.Row():
clear_btn = gr.Button("대화 초기화", elem_classes=["secondary-btn"])
gr.Examples(
examples=[
"연차 휴가 일수는 어떻게 계산하나요?",
"당직 근무 시간은 어떻게 되나요?",
"병가 신청 절차가 어떻게 되나요?",
"초과 근무 수당 기준이 어떻게 되나요?"
],
inputs=msg,
label="예시 질문"
)
# ========== 관리자 탭 ==========
with gr.TabItem("⚙️ 관리자"):
# 로그인 섹션
with gr.Column(visible=True) as login_section:
gr.Markdown("### 🔐 관리자 인증")
gr.Markdown("문서를 관리하려면 관리자 비밀번호를 입력하세요.")
with gr.Row():
password_input = gr.Textbox(
label="비밀번호",
type="password",
placeholder="비밀번호 입력...",
scale=3
)
login_btn = gr.Button("로그인", elem_classes=["primary-btn"], scale=1)
login_status = gr.Markdown("")
# 관리자 패널 (인증 후 표시)
with gr.Column(visible=False) as admin_panel:
gr.Markdown("### 📄 PDF 문서 업로드")
gr.Markdown("새로운 복무 관련 문서를 업로드하여 챗봇 지식을 확장하세요.")
with gr.Group(elem_classes=["upload-section"]):
pdf_file = gr.File(
label="PDF 파일 선택",
file_types=[".pdf"],
file_count="single"
)
with gr.Row():
doc_title = gr.Textbox(
label="문서 제목",
placeholder="예: 소방공무원 복무규정 2024",
scale=2
)
doc_category = gr.Dropdown(
label="카테고리",
choices=["일반", "휴가", "근무", "수당", "규정", "교육"],
value="일반",
scale=1
)
upload_btn = gr.Button("📤 업로드 및 처리", elem_classes=["primary-btn"])
upload_result = gr.Markdown("")
gr.Markdown("---")
gr.Markdown("### 📚 저장된 문서 목록")
refresh_btn = gr.Button("🔄 새로고침", elem_classes=["secondary-btn"])
documents_table = gr.Dataframe(
headers=["제목", "카테고리", "ID"],
label="",
interactive=False
)
# 푸터
gr.HTML("""
""")
# ========== 이벤트 핸들러 ==========
# 채팅
msg.submit(respond, [msg, chatbot, session_id], [msg, chatbot])
submit_btn.click(respond, [msg, chatbot, session_id], [msg, chatbot])
clear_btn.click(clear_chat, [session_id], [chatbot, msg])
# 관리자
login_btn.click(
verify_admin,
[password_input],
[login_section, admin_panel, login_status]
)
upload_btn.click(
upload_pdf,
[pdf_file, doc_title, doc_category],
[upload_result]
)
refresh_btn.click(
get_documents_list,
[],
[documents_table]
)
if __name__ == "__main__":
app.launch(server_name="0.0.0.0", server_port=7860)