youlbot/app.py

"""Gradio Web UI — 율봇 Phase 4 + Phase 9/10 + Phase 14(음성)."""
import os
import subprocess
import tempfile
import gradio as gr
from dotenv import load_dotenv
load_dotenv()

from container import Container
from services.agent.agent_service import AgentService

container = Container()

db = container.db_service()
db.connect()
db.init_schema()

ingestion = container.ingestion_service()
retriever = container.retriever_service()

_cfg = container.config()
_agent_cache: dict[str, AgentService] = {}

USER_LABELS = ["아록", "근혜", "도율", "하율"]
DEFAULT_USER = "아록"

_whisper_model = None


def _get_whisper():
    global _whisper_model
    if _whisper_model is None:
        import whisper
        _whisper_model = whisper.load_model(_cfg.whisper_model_size)
    return _whisper_model


def transcribe_audio(filepath: str) -> str:
    if not filepath:
        return ""
    model = _get_whisper()
    result = model.transcribe(filepath, language="ko")
    return result["text"].strip()


def tts_speak(text: str, voice: str) -> str | None:
    """텍스트를 macOS say 명령어로 음성 변환, 재생용 wav 파일 경로 반환."""
    if not text:
        return None
    try:
        tmp = tempfile.NamedTemporaryFile(suffix=".aiff", delete=False)
        tmp.close()
        subprocess.run(
            ["say", "-v", voice, "-o", tmp.name, text],
            check=True,
            capture_output=True,
        )
        return tmp.name
    except Exception:
        return None


def _get_agent(user_id: str) -> AgentService:
    if user_id not in _agent_cache:
        _agent_cache[user_id] = AgentService(
            chat_model=container.chat_model(),
            retriever_service=retriever,
            system_prompt=_cfg.system_prompt,
            rag_verbose=_cfg.rag_verbose,
            rag_show_sources=_cfg.rag_show_sources,
            langgraph_verbose=_cfg.langgraph_verbose,
            think_verbose=_cfg.think_verbose,
            user_profile_repository=container.user_profile_repository(),
            conversation_repository=container.conversation_repository(),
            user_id=user_id,
        )
    return _agent_cache[user_id]


async def respond(message, history, show_thinking, user_id, use_tts):
    if not message.strip():
        yield history, "", None
        return

    agent = _get_agent(user_id)
    history = list(history)
    history.append({"role": "user", "content": message})
    history.append({"role": "assistant", "content": ""})
    yield history, "", None

    async for token in agent.stream_response(message, show_thinking=show_thinking):
        history[-1]["content"] += token
        yield history, "", None

    if use_tts:
        response_text = history[-1]["content"]
        audio_path = tts_speak(response_text, _cfg.tts_voice)
        yield history, "", audio_path


def switch_user(user_id):
    """사용자 전환 시 채팅 화면만 초기화 (대화 이력은 유지)."""
    return []


def reset_chat(user_id):
    agent = _get_agent(user_id)
    agent.reset()
    return []


def ingest_files(files):
    if not files:
        return "파일을 선택해주세요."
    paths = [f if isinstance(f, str) else f.name for f in files]
    try:
        count = ingestion.ingest(paths)
        names = ", ".join(p.split("/")[-1] for p in paths)
        return f"완료: {names} → {count}개 청크 저장됨"
    except Exception as e:
        return f"오류: {e}"


def list_docs():
    try:
        sources = retriever.list_documents()
        return [[os.path.basename(s), s] for s in sources]
    except Exception as e:
        return [[f"오류: {e}", ""]]


def delete_doc(source):
    if not source.strip():
        return "삭제할 파일 경로를 입력하세요.", list_docs()
    try:
        retriever.delete_document(source.strip())
        return f"삭제 완료: {os.path.basename(source.strip())}", list_docs()
    except Exception as e:
        return f"오류: {e}", list_docs()


with gr.Blocks(title="율봇") as demo:
    gr.Markdown("# 율봇\n육아·금융 전문 AI 상담 도우미")

    user_state = gr.State(DEFAULT_USER)

    with gr.Tab("대화"):
        with gr.Row():
            user_selector = gr.Dropdown(
                choices=USER_LABELS,
                value=DEFAULT_USER,
                label="사용자",
                scale=1,
            )

        chatbot = gr.Chatbot(label="율봇", height=500)
        with gr.Row():
            msg_box = gr.Textbox(
                placeholder="질문을 입력하세요... (Enter로 전송)",
                label="",
                scale=5,
                autofocus=True,
            )
            send_btn = gr.Button("전송", variant="primary", scale=1)

        # 음성 입력 (STT)
        with gr.Row():
            audio_input = gr.Audio(
                sources=["microphone"],
                type="filepath",
                label="음성으로 질문하기",
                scale=4,
            )
            transcribe_btn = gr.Button("음성 → 텍스트 변환", scale=1)

        with gr.Row():
            show_thinking = gr.Checkbox(label="사고 과정 표시", value=False)
            use_tts = gr.Checkbox(label="음성으로 답변 읽기 (TTS)", value=False)
            reset_btn = gr.Button("대화 초기화", size="sm")

        # TTS 출력
        tts_output = gr.Audio(label="음성 답변", autoplay=True, visible=False)
        use_tts.change(lambda v: gr.Audio(visible=v), inputs=[use_tts], outputs=[tts_output])

        user_selector.change(
            switch_user,
            inputs=[user_selector],
            outputs=[chatbot],
        ).then(
            lambda u: u, inputs=[user_selector], outputs=[user_state]
        )

        transcribe_btn.click(
            transcribe_audio,
            inputs=[audio_input],
            outputs=[msg_box],
        )

        send_btn.click(
            respond,
            inputs=[msg_box, chatbot, show_thinking, user_state, use_tts],
            outputs=[chatbot, msg_box, tts_output],
        )
        msg_box.submit(
            respond,
            inputs=[msg_box, chatbot, show_thinking, user_state, use_tts],
            outputs=[chatbot, msg_box, tts_output],
        )
        reset_btn.click(reset_chat, inputs=[user_state], outputs=[chatbot])

    with gr.Tab("문서 등록"):
        gr.Markdown("PDF 또는 TXT 파일을 업로드하면 율봇이 내용을 참고해 답변합니다.")
        file_input = gr.File(
            file_types=[".pdf", ".txt"],
            file_count="multiple",
            label="파일 선택",
        )
        ingest_btn = gr.Button("문서 수집", variant="primary")
        ingest_status = gr.Textbox(label="결과", interactive=False)
        ingest_btn.click(ingest_files, inputs=[file_input], outputs=[ingest_status])

    with gr.Tab("문서 관리"):
        gr.Markdown("Qdrant에 등록된 문서 목록입니다. 불필요한 문서를 삭제할 수 있습니다.")
        doc_table = gr.Dataframe(
            headers=["파일명", "전체 경로"],
            label="등록된 문서",
            interactive=False,
        )
        refresh_btn = gr.Button("새로고침")
        gr.Markdown("---")
        with gr.Row():
            delete_source = gr.Textbox(
                label="삭제할 파일 경로",
                placeholder="위 표에서 전체 경로를 복사해 붙여넣으세요",
                scale=4,
            )
            delete_btn = gr.Button("삭제", variant="stop", scale=1)
        delete_status = gr.Textbox(label="결과", interactive=False)

        refresh_btn.click(list_docs, outputs=[doc_table])
        delete_btn.click(
            delete_doc,
            inputs=[delete_source],
            outputs=[delete_status, doc_table],
        )
        demo.load(list_docs, outputs=[doc_table])


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, theme=gr.themes.Soft())