ui/chat_tab.py

"""Single Chat 탭 UI"""

import gradio as gr
from typing import Any

from models.model_registry import get_all_models, get_model_info
from characters import get_character_loader, build_system_prompt
from utils import parse_thinking_response, format_thinking_for_display


def create_chat_tab(
    model_manager: Any = None,
    use_mock: bool = False,
):
    """Single Chat 탭 생성"""

    # 데이터 로더
    char_loader = get_character_loader()

    # 모델 목록
    all_models = get_all_models()
    model_choices = [(f"{get_model_info(m).get('description', m)}", m) for m in all_models]

    # 캐릭터 목록
    characters = char_loader.get_character_names()

    # ============================================================
    # UI 구성
    # ============================================================

    gr.Markdown("## 단일 모델 채팅")
    gr.Markdown("선택한 모델과 캐릭터로 대화를 나눠보세요.")

    with gr.Row():
        with gr.Column(scale=1):
            model_dropdown = gr.Dropdown(
                choices=model_choices,
                value=all_models[0] if all_models else None,
                label="모델 선택",
            )
        with gr.Column(scale=1):
            character_dropdown = gr.Dropdown(
                choices=characters,
                value=characters[0] if characters else None,
                label="캐릭터 선택",
            )

    # 채팅 영역
    chatbot = gr.Chatbot(
        label="대화",
        height=400,
        type="messages",
    )

    with gr.Accordion("Thinking Process (마지막 응답)", open=False):
        thinking_display = gr.Markdown("*(응답 생성 후 표시됩니다)*")

    with gr.Row():
        user_input = gr.Textbox(
            label="메시지 입력",
            placeholder="메시지를 입력하세요...",
            lines=2,
            scale=4,
        )
        send_btn = gr.Button("전송", variant="primary", scale=1)

    with gr.Row():
        clear_btn = gr.Button("대화 초기화")

    metadata_display = gr.Markdown("")

    # ============================================================
    # 이벤트 핸들러
    # ============================================================

    def respond(
        model_id: str,
        character: str,
        message: str,
        history: list,
    ):
        """응답 생성"""
        if not message.strip():
            return history, "", "*(메시지를 입력해주세요)*", ""

        # 대화 히스토리 구성
        messages = []
        for msg in history:
            if msg["role"] == "user":
                messages.append({"role": "user", "content": msg["content"]})
            elif msg["role"] == "assistant":
                # Thinking 제거한 클린 응답만 히스토리에
                _, clean = parse_thinking_response(msg["content"])
                messages.append({"role": "assistant", "content": clean})

        messages.append({"role": "user", "content": message})

        system_prompt = build_system_prompt(character)

        # Mock 또는 실제 추론
        if use_mock or model_manager is None:
            response_full = f"<think>\n{character}로서 생각해보면...\n</think>\n\n안녕~ 반가워! (Mock Response)"
            meta = {"latency_s": 0.5, "output_tokens": 30}
        else:
            try:
                response_full, meta = model_manager.generate_response(
                    model_id, messages, system_prompt
                )
            except Exception as e:
                response_full = f"*Error: {str(e)}*"
                meta = {"latency_s": 0, "output_tokens": 0}

        # Thinking 파싱
        thinking, clean_response = parse_thinking_response(response_full)

        # 히스토리 업데이트
        history.append({"role": "user", "content": message})
        history.append({"role": "assistant", "content": response_full})

        # 메타데이터
        meta_str = f"⏱️ {meta.get('latency_s', 0):.2f}s | {meta.get('output_tokens', 0)} tokens"

        return (
            history,
            "",  # 입력 초기화
            format_thinking_for_display(thinking) if thinking else "*No thinking*",
            meta_str,
        )

    def clear_chat():
        """대화 초기화"""
        return [], "", "*(응답 생성 후 표시됩니다)*", ""

    # ============================================================
    # 이벤트 바인딩
    # ============================================================

    send_btn.click(
        fn=respond,
        inputs=[model_dropdown, character_dropdown, user_input, chatbot],
        outputs=[chatbot, user_input, thinking_display, metadata_display],
    )

    user_input.submit(
        fn=respond,
        inputs=[model_dropdown, character_dropdown, user_input, chatbot],
        outputs=[chatbot, user_input, thinking_display, metadata_display],
    )

    clear_btn.click(
        fn=clear_chat,
        outputs=[chatbot, user_input, thinking_display, metadata_display],
    )