app.py

import gradio as gr
import logging
import uuid
from datetime import datetime
import pandas as pd
from i18n import get_text
from model_handler import ModelHandler
from tab_chat import create_chat_tab
from tab_chat import update_language as update_chat_language
from tab_code import create_code_tab
from tab_code import update_language as update_code_language
from tab_smart_writer import create_smart_writer_tab
from tab_smart_writer import update_language as update_writer_language
from tab_welcome import create_welcome_tab
from tab_welcome import update_language as update_welcome_language

# Configure logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


def get_history_df_from_app(history):
    if not history:
        return pd.DataFrame(
            {"ID": pd.Series(dtype="str"), "对话": pd.Series(dtype="str")}
        )
    df = pd.DataFrame(history)
    if "id" in df.columns and "title" in df.columns:
        return df[["id", "title"]].rename(columns={"id": "ID", "对话": "对话"})
    return pd.DataFrame({"ID": pd.Series(dtype="str"), "对话": pd.Series(dtype="str")})


def on_app_load(request: gr.Request, history, conv_id, current_lang_state):
    """
    Handles the application's initial state on load.
    - Determines language from URL parameter.
    - Loads conversation history or creates a new one.
    """
    # --- Language Detection ---
    query_params = dict(request.query_params)
    url_lang = query_params.get("lang")

    updated_lang = current_lang_state  # Start with the default
    if url_lang and url_lang in ["en", "zh"]:
        updated_lang = url_lang

    # --- History Loading Logic ---
    if not history:
        # First time ever, create a new conversation
        conv_id = str(uuid.uuid4())
        new_convo_title = get_text("chat_new_conversation_title", updated_lang)
        new_convo = {
            "id": conv_id,
            "title": new_convo_title,
            "messages": [],
            "timestamp": datetime.now().isoformat(),
        }
        history = [new_convo]
        return (
            conv_id,
            history,
            gr.update(value=get_history_df_from_app(history)),
            [],
            updated_lang,
        )

    if conv_id and any(c["id"] == conv_id for c in history):
        # Valid last session, load it
        for convo in history:
            if convo["id"] == conv_id:
                return (
                    conv_id,
                    history,
                    gr.update(value=get_history_df_from_app(history)),
                    convo["messages"],
                    updated_lang,
                )

    # Fallback to most recent conversation
    most_recent_convo = history[0]
    return (
        most_recent_convo["id"],
        history,
        gr.update(value=get_history_df_from_app(history)),
        most_recent_convo["messages"],
        updated_lang,
    )


CSS = """

#chatbot {
  height: calc(100vh - 21px - 16px);
  max-height: 1500px;
}

footer { 
    display: none !important; 
}

/* Disable transition and animation for no-transition class */
.no-transition, .no-transition * {
    transition: none !important;
    animation: none !important;
    animation-play-state: paused !important;
}
"""

if __name__ == "__main__":
    # Instantiate the model handler with the configuration
    logger.info("Starting Ling Space Application...")
    model_handler = ModelHandler()

    with gr.Blocks(analytics_enabled=False, fill_height=True, fill_width=True) as demo:
        # Language State
        current_lang_state = gr.State("en")

        # --- Collect all components that need language updates ---
        all_i18n_outputs = []

        with gr.Tabs(elem_id="indicator-space-app") as tabs:
            welcome_components = create_welcome_tab(current_lang_state.value)

            # The order of components MUST be consistent
            welcome_outputs = [
                welcome_components["tab"],
                welcome_components["header"],
                welcome_components["description"],
                welcome_components["chat_description"],
                welcome_components["code_description"],
                welcome_components["writer_description"],
                welcome_components["lang_select_header"],
                welcome_components["en_button"],
                welcome_components["zh_button"],
            ]
            all_i18n_outputs.extend(welcome_outputs)

            # --- Chat Tab ---
            with gr.TabItem(
                get_text("chat_tab_title", current_lang_state.value)
            ) as chat_tab:
                chat_components = create_chat_tab(
                    current_lang_state.value, current_lang_state
                )

            chat_outputs = [
                chat_tab,
                chat_components["new_chat_btn"],
                chat_components["history_df"],
                chat_components["chatbot"],
                chat_components["textbox"],
                chat_components["submit_btn"],
                chat_components["recommended_title"],
                chat_components["recommended_dataset"],
                chat_components["system_prompt_textbox"],
                chat_components["temperature_slider"],
            ]
            all_i18n_outputs.extend(chat_outputs)

            chat_tab.select(
                fn=None,
                js="() => {window.dispatchEvent(new CustomEvent('tabSelect.chat')); console.log('this'); return null;}",
            )

            # --- Code Tab ---
            with gr.TabItem(
                get_text("code_tab_title", current_lang_state.value)
            ) as code_tab:
                code_components = create_code_tab(
                    current_lang_state.value, current_lang_state
                )

            code_outputs = [
                code_tab,
                code_components["prompt_input"],
                code_components["overall_style_input"],
                code_components["decoration_input"],
                code_components["palette_display"],
                code_components["generate_style_btn"],
                code_components["examples_title"],
                code_components["examples_dataset"],
                code_components["generate_button"],
                code_components["preview_tab"],
                code_components["preview_header"],
                code_components["source_code_tab"],
                code_components["source_code_header"],
                code_components["code_output"],
                code_components["refresh_button"],
                code_components["log_chatbot"],
                code_components["js_error_channel"],
                # fullscreen_button is updated in its own handler, so it's excluded here
            ]
            all_i18n_outputs.extend(code_outputs)

            code_tab.select(
                fn=None,
                js="() => {window.dispatchEvent(new CustomEvent('tabSelect.code')); return null;}",
            )

            # --- Writer Tab ---
            with gr.TabItem(
                get_text("writer_tab_title", current_lang_state.value)
            ) as writer_tab:
                writer_components = create_smart_writer_tab(current_lang_state.value)

            writer_outputs = [
                writer_tab,
                writer_components["style_input"],
                writer_components["kb_accordion"],
                writer_components["kb_input"],
                writer_components["btn_suggest_kb"],
                writer_components["suggested_kb_dataframe"],
                writer_components["short_outline_accordion"],
                writer_components["short_outline_input"],
                writer_components["btn_sync_outline"],
                writer_components["long_outline_accordion"],
                writer_components["long_outline_input"],
                writer_components["flow_suggestion_display"],
                writer_components["btn_accept_flow"],
                writer_components["btn_change_flow"],
                writer_components["inspiration_prompt_input"],
                writer_components["prompt_suggestions_dataset"],
                writer_components["refresh_suggestions_btn"],
                writer_components["btn_generate_para"],
                writer_components["btn_change_para"],
                writer_components["btn_accept_para"],
                writer_components["para_suggestion_display"],
                writer_components["polish_title"],
                writer_components["polish_soon"],
                writer_components["stats_display"],
                writer_components["editor"],
            ]
            all_i18n_outputs.extend(writer_outputs)

            writer_tab.select(
                fn=None,
                js="() => {window.dispatchEvent(new CustomEvent('tabSelect.writing')); return null;}",
            )

        # --- Language Change Handler ---
        def on_language_change(lang):
            # Dispatch updates for each tab
            welcome_updates = update_welcome_language(lang, welcome_components)
            chat_updates = update_chat_language(lang, chat_components)
            chat_updates[chat_tab] = gr.update(label=get_text("chat_tab_title", lang))

            code_updates = update_code_language(lang, code_components)
            code_updates[code_tab] = gr.update(label=get_text("code_tab_title", lang))

            writer_updates = update_writer_language(lang, writer_components)
            writer_updates[writer_tab] = gr.update(
                label=get_text("writer_tab_title", lang)
            )

            all_updates = {
                **welcome_updates,
                **chat_updates,
                **code_updates,
                **writer_updates,
            }
            return tuple(all_updates.get(comp) for comp in all_i18n_outputs)

        current_lang_state.change(
            fn=on_language_change,
            inputs=[current_lang_state],
            outputs=all_i18n_outputs,
            show_progress="hidden",
        )

        # --- App Load Handler ---
        conversation_store = chat_components["conversation_store"]
        current_conversation_id = chat_components["current_conversation_id"]
        history_df = chat_components["history_df"]
        chatbot = chat_components["chatbot"]

        demo.load(
            on_app_load,
            inputs=[conversation_store, current_conversation_id, current_lang_state],
            outputs=[
                current_conversation_id,
                conversation_store,
                history_df,
                chatbot,
                current_lang_state,
            ],
            js="() => {window.dispatchEvent(new CustomEvent('appStart')); console.log('appStart'); return {};}",
        )

    demo.queue(default_concurrency_limit=32, max_size=128)
    # Launch the Gradio application
    demo.launch(
        theme=gr.themes.Default(),
        ssr_mode=False,
        max_threads=64,
        css=CSS,
        head="",
        head_paths=["./static/toastify.html", "./static/app.html"],
    )