app.py

"""Multi-model comparison playground (Gradio).

Broadcasts the same system prompt + intro + each user message to several backends.
Each backend keeps its OWN conversation thread and streams its reply into its own
panel, live and concurrently. You pick the best response per turn; every request,
response, and the preferred pick are saved under DATA_DIR. Entrypoint for Hugging
Face Spaces (must be app.py).

Prompts/intros may contain {VARIABLE} tokens — those are surfaced as fill-in fields
in the sidebar and substituted into the effective prompt when a session starts.
"""

import json
import queue
import re
import threading
from datetime import datetime

import gradio as gr

import config
from src.clients import BackendError, stream_backend, transcribe_audio
from src.display import format_for_display
from src.storage import (
    append_request_log,
    append_to_dataset,
    new_session,
    save_session,
)

BACKENDS = config.BACKENDS
KEYS = [b["key"] for b in BACKENDS]
BY_KEY = {b["key"]: b for b in BACKENDS}
TIE = "Tie"

# "Preferred response" only makes sense when comparing 2+ backends. With a single
# backend (the default — see config.ENABLE_BACKEND_B) the picker is hidden.
SHOW_COMPARE = len(KEYS) > 1

# Max number of fill-in fields surfaced in the sidebar (per-preset variables).
MAX_VARS = 10

# Anonymous display labels so the panels reveal no model identity (avoids bias).
# The mapping back to the real backend is kept here and saved with each pick.
ANON = {k: f"Response {chr(65 + i)}" for i, k in enumerate(KEYS)}  # A, B, ...
ANON_TO_KEY = {v: k for k, v in ANON.items()}


# ---------------------------------------------------------------------------
# Prompt variables — {TOKEN} placeholders become fill-in fields
# ---------------------------------------------------------------------------

_VAR_RE = re.compile(r"\{([A-Za-z0-9_]+)\}")


def fill_vars(text: str, names: list, values) -> str:
    """Substitute filled {TOKEN}s into `text`; leave unfilled tokens untouched."""
    mapping = {n: v for n, v in zip(names, values) if v}
    return _VAR_RE.sub(lambda m: mapping.get(m.group(1), m.group(0)), text or "")


def load_preset_fields(name):
    """Load a preset's prompt + intro and its fixed list of fill-in variables.

    Returns updates for the prompt + intro textboxes, the var-name state, and
    each variable box: one visible, labelled, pre-filled box per token declared
    in the preset's `vars`, the rest hidden. Variables are fixed per preset —
    users don't add their own tokens.
    """
    preset = config.get_preset(name)
    pvars = preset.get("vars", {})
    tokens = list(pvars.keys())[:MAX_VARS]
    box_updates = []
    for i in range(MAX_VARS):
        if i < len(tokens):
            token = tokens[i]
            box_updates.append(gr.update(visible=True, label=token, value=pvars[token]))
        else:
            box_updates.append(gr.update(visible=False, value=""))
    return (preset["system_prompt"], preset["intro"], tokens, *box_updates)


# ---------------------------------------------------------------------------
# Rendering helpers
# ---------------------------------------------------------------------------


def render_history(history: list) -> list:
    """Convert a backend's raw conversation into chatbot-safe display messages."""
    out = []
    for h in history:
        content = h["content"]
        if h["role"] == "assistant":
            content = format_for_display(content)
        out.append({"role": h["role"], "content": content})
    return out


def format_metrics(m: dict) -> str:
    """Render a compact metrics line for one backend panel (tokens only)."""
    if not m:
        return ""
    if m.get("error"):
        return f"⚠️ {m['error']}"
    parts = []
    if m.get("tool_called"):
        parts.append(f"📞 tool: `{m['tool_called']}`")
    if m.get("prompt_tokens") is not None:
        parts.append(f"In: {m['prompt_tokens']}")
    if m.get("completion_tokens") is not None:
        parts.append(f"Out: {m['completion_tokens']}")
    if m.get("cached_tokens") is not None:
        parts.append(f"Cached: {m['cached_tokens']}")
    # if m.get("latency_s") is not None:
    #     parts.append(f"TTFB: {m['latency_s']}")
    return " &nbsp;·&nbsp; ".join(parts)


def format_debug(acc: dict, mstate: dict, rstate: dict) -> str:
    """Render the exact request + model response(s) for the current turn as JSON.

    Shows, per backend: the exact request body sent to the model, the
    live-streamed text, and — once the turn completes — the raw assistant
    message the model produced (`content` + any `tool_calls` with verbatim
    arguments), so both sides of the exchange are visible exactly as sent.
    """
    payload = {}
    for k in KEYS:
        m = mstate.get(k) or {}
        payload[ANON[k]] = {
            "label": BY_KEY[k]["label"],
            "request": rstate.get(k),
            "streamed_text": acc.get(k, ""),
            "raw_response": m.get("raw_response"),
            "tool_called": m.get("tool_called"),
            "error": m.get("error"),
        }
    return json.dumps(payload, ensure_ascii=False, indent=2)


def backend_messages(system_prompt: str, conversation: list) -> list:
    """Build the request messages for one backend: system + its own thread."""
    msgs = []
    if system_prompt and system_prompt.strip():
        msgs.append({"role": "system", "content": system_prompt.strip()})
    msgs.extend({"role": t["role"], "content": t["content"]} for t in conversation)
    return msgs


# ---------------------------------------------------------------------------
# Session lifecycle
# ---------------------------------------------------------------------------


def init_state(
    system_prompt: str, intro: str, temperature: float, max_tokens: int
) -> dict:
    """Fresh per-backend histories seeded with the intro + a new session log."""
    seed = (
        [{"role": "assistant", "content": intro.strip()}]
        if intro and intro.strip()
        else []
    )
    histories = {k: [dict(m) for m in seed] for k in KEYS}
    session = new_session(system_prompt, intro, temperature, max_tokens)
    return {"histories": histories, "session": session}


def start_session(
    system_prompt, intro, temperature, max_tokens, var_names, *var_values
):
    """New session / Apply: substitute filled variables, rebuild threads, reset panels."""
    sp = fill_vars(system_prompt, var_names, var_values)
    intro_s = fill_vars(intro, var_names, var_values)
    state = init_state(sp, intro_s, temperature, max_tokens)
    chatbots = [render_history(state["histories"][k]) for k in KEYS]
    metrics = ["" for _ in KEYS]
    status = f"🆕 New session `{state['session']['session_id']}` — settings applied."
    return (state, *chatbots, *metrics, status, None)


def transcribe(filepath, current_msg):
    """Transcribe a recording (STT) and append it to the message box."""
    if not filepath:
        return current_msg, "🎤 Record something first."
    try:
        text = transcribe_audio(filepath)
    except BackendError as e:
        return current_msg, f"⚠️ {e}"
    except Exception as e:  # noqa: BLE001
        return current_msg, f"⚠️ Transcription failed: {e}"
    if not text:
        return current_msg, "🎤 Didn't catch that — try again."
    combined = f"{current_msg.strip()} {text}".strip() if current_msg else text
    return (
        combined,
        f"✅ Heard: “{combined}” — press **Send** (switch to ⌨️ Text to edit).",
    )


# ---------------------------------------------------------------------------
# Concurrent streaming for all backends
# ---------------------------------------------------------------------------


def respond_all(user_msg, state, temperature, max_tokens):
    """Broadcast the user message to all backends and stream replies concurrently."""
    histories = state["histories"]
    session = state["session"]

    if not user_msg or not user_msg.strip():
        chatbots = [render_history(histories[k]) for k in KEYS]
        metrics = [gr.update() for _ in KEYS]
        yield (state, *chatbots, *metrics, "", None, gr.update())
        return

    user_msg = user_msg.strip()
    req_msgs = {}
    for k in KEYS:
        histories[k].append({"role": "user", "content": user_msg})
        # Capture the request BEFORE adding the empty assistant placeholder.
        req_msgs[k] = backend_messages(session["system_prompt"], histories[k])
        histories[k].append({"role": "assistant", "content": ""})

    acc = {k: "" for k in KEYS}
    mstate = {k: {} for k in KEYS}
    rstate = {k: None for k in KEYS}  # exact request payload sent per backend
    q: queue.Queue = queue.Queue()

    def worker(key):
        backend = BY_KEY[key]
        try:
            for item in stream_backend(backend, req_msgs[key], temperature, max_tokens):
                if isinstance(item, dict) and item.get("__metrics__"):
                    q.put((key, "metrics", item))
                elif isinstance(item, dict) and item.get("__request__"):
                    q.put((key, "request", item.get("payload")))
                else:
                    q.put((key, "delta", item))
        finally:
            q.put((key, "done", None))

    threads = [threading.Thread(target=worker, args=(k,), daemon=True) for k in KEYS]
    for t in threads:
        t.start()

    def snapshot():
        chatbots = []
        for k in KEYS:
            histories[k][-1]["content"] = acc[k]
            chatbots.append(render_history(histories[k]))
        metrics = [format_metrics(mstate[k]) for k in KEYS]
        debug = format_debug(acc, mstate, rstate)
        return (state, *chatbots, *metrics, "", None, debug)

    remaining = set(KEYS)
    yield snapshot()
    while remaining:
        try:
            key, kind, payload = q.get(timeout=0.1)
        except queue.Empty:
            yield snapshot()
            continue
        if kind == "delta":
            acc[key] += payload
        elif kind == "metrics":
            mstate[key] = payload
        elif kind == "request":
            rstate[key] = payload
        elif kind == "done":
            remaining.discard(key)
        yield snapshot()

    # Finalize: write the raw replies into each thread and log the turn.
    for k in KEYS:
        histories[k][-1]["content"] = acc[k]
    turn = {
        "turn_index": len(session["turns"]),
        "timestamp": datetime.now().isoformat(),
        "temperature": temperature,
        "max_tokens": max_tokens,
        "user_message": user_msg,
        "responses": {
            k: {
                "display_label": ANON[k],
                "label": BY_KEY[k]["label"],
                "model": BY_KEY[k].get("model") or BY_KEY[k].get("deployment"),
                "endpoint_or_deployment": BY_KEY[k].get("endpoint_id")
                or BY_KEY[k].get("deployment"),
                "request_messages": req_msgs[k],
                "content": acc[k],
                "metrics": {
                    mk: mstate[k].get(mk)
                    for mk in (
                        "prompt_tokens",
                        "completion_tokens",
                        "total_tokens",
                        "cached_tokens",
                        "latency_s",
                        "tool_called",
                    )
                },
                "error": mstate[k].get("error"),
            }
            for k in KEYS
        },
        "preferred": None,
    }
    session["turns"].append(turn)
    # Persist the full session AND append every request to the durable log.
    try:
        save_session(session)
        append_request_log(session, turn)
    except Exception as e:  # noqa: BLE001
        print(f"request log failed: {e}")
    yield snapshot()


def save_preferred(state, preferred):
    """Persist the last turn's preferred pick to the comparison dataset."""
    session = state["session"]
    if not session["turns"]:
        return "Nothing to save yet — send a message first."
    if not preferred:
        return "Pick a preferred response (or 'Tie') first."
    turn = session["turns"][-1]
    # The user picked an anonymous label (e.g. "Response A"); resolve it back to
    # the real backend so the saved data records which model was actually preferred.
    pref_key = ANON_TO_KEY.get(preferred)
    turn["preferred"] = preferred  # what the user saw (anonymous)
    turn["preferred_key"] = pref_key  # real backend key, or None for Tie/None
    if pref_key:
        turn["preferred_label"] = BY_KEY[pref_key]["label"]
        turn["preferred_model"] = BY_KEY[pref_key].get("model") or BY_KEY[pref_key].get(
            "deployment"
        )
    else:
        turn["preferred_label"] = None
        turn["preferred_model"] = None
    try:
        save_session(session)
        append_to_dataset(session, turn)
    except Exception as e:  # noqa: BLE001
        return f"⚠️ Save failed: {e}"
    detail = turn["preferred_label"] or preferred
    return (
        f"💾 Saved turn {turn['turn_index']} — preferred **{preferred}** "
        f"(_{detail}_) to comparisons.jsonl"
    )


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

# Default the UI to the first preset, with its declared variables pre-filled
# from their default values so the opening panels read cleanly.
_PRESET0 = config.PRESETS[0]
_DEFAULT_SP = _PRESET0["system_prompt"]
_DEFAULT_INTRO = _PRESET0["intro"]
_DEFAULT_VARS = list(_PRESET0["vars"].keys())
_DEFAULT_VALS = list(_PRESET0["vars"].values())

_INIT = init_state(
    fill_vars(_DEFAULT_SP, _DEFAULT_VARS, _DEFAULT_VALS),
    fill_vars(_DEFAULT_INTRO, _DEFAULT_VARS, _DEFAULT_VALS),
    config.DEFAULT_TEMPERATURE,
    config.DEFAULT_MAX_TOKENS,
)

with gr.Blocks(
    theme=gr.Theme.from_hub("Maani/MonoNeo"), title="LLM Comparison Playground"
) as demo:
    _intro_line = (
        "Same prompt & conversation across multiple backends — pick the best reply, "
        "and the request/responses are saved for analysis."
        if SHOW_COMPARE
        else "Chat with the model, and every request/response is saved for analysis."
    )

    gr.Markdown(f"# 🆚 LLM Comparison Playground\n{_intro_line}")

    state = gr.State(_INIT)
    var_names = gr.State([])  # token names currently shown as fill-in fields

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### ⚙️ Settings")
            preset = gr.Dropdown(
                choices=[p["name"] for p in config.PRESETS],
                value=config.PRESETS[0]["name"],
                label="Preset (loads a prompt + intro to try)",
            )
            system_prompt = gr.Textbox(
                label="System Prompt (only enter task and FAQs)",
                value=_DEFAULT_SP,
                lines=6,
            )
            intro = gr.Textbox(
                label="Intro Message (first assistant turn)",
                value=_DEFAULT_INTRO,
                lines=2,
            )

            with gr.Group():
                gr.Markdown("### 📝 Call details")
                gr.Markdown(
                    "Fill in the values below, then click **New session / Apply** "
                    "to use them in the prompt."
                )
                var_boxes = [
                    gr.Textbox(label=f"var{i}", visible=False, interactive=True)
                    for i in range(MAX_VARS)
                ]

            temperature = gr.Slider(
                label="Temperature",
                minimum=0.0,
                maximum=2.0,
                step=0.1,
                value=config.DEFAULT_TEMPERATURE,
            )
            max_tokens = gr.Slider(
                label="Max Tokens",
                minimum=16,
                maximum=4096,
                step=16,
                value=config.DEFAULT_MAX_TOKENS,
            )
            new_btn = gr.Button("🆕 New session / Apply", variant="secondary")
            gr.Markdown(
                "_Editing the System Prompt / Intro or filling in variables takes "
                "effect after **New session / Apply** (it rebuilds all threads)._"
            )

        with gr.Column(scale=4):
            chatbots = []
            metrics = []
            with gr.Row():
                for b in BACKENDS:
                    with gr.Column():
                        gr.Markdown(f"### {ANON[b['key']]}")
                        cb = gr.Chatbot(
                            label=None,
                            height=460,
                            value=render_history(_INIT["histories"][b["key"]]),
                            show_label=False,
                        )
                        md = gr.Markdown("")
                        chatbots.append(cb)
                        metrics.append(md)

            input_mode = gr.Radio(
                choices=["🎤 Voice", "⌨️ Text"],
                value="🎤 Voice",
                label="Input mode",
            )
            with gr.Group():
                # Voice is the default; the text bar shows only when the user
                # switches Input mode to "⌨️ Text". Toggling the *column*
                # visibility (not the field's) renders reliably in Gradio.
                with gr.Column(visible=False) as text_input_col:
                    msg = gr.Textbox(
                        label="Your message",
                        placeholder="Type a message — sent to all backends…",
                        lines=2,
                    )
                with gr.Column(visible=True) as voice_input_col:
                    audio_in = gr.Audio(
                        sources=["microphone"],
                        type="filepath",
                        label="🎤 Voice input (STT)",
                    )
                stt_status = gr.Markdown("")
            with gr.Row():
                send_btn = gr.Button("Send", variant="primary")

            # Debug: the exact request body sent and the raw response the model
            # produced this turn (text + any tool calls with verbatim
            # arguments), per backend. Collapsed by default.
            with gr.Accordion("🐞 Debug — exact request & response", open=False):
                debug_box = gr.Code(
                    label="Exact request & raw response (per backend)",
                    language="json",
                    value="",
                    interactive=False,
                )

            # Pick the best response for this turn and save it to the dataset.
            with gr.Group(visible=SHOW_COMPARE):
                gr.Markdown("### 💾 Save preferred response")
                preferred = gr.Radio(
                    choices=[ANON[k] for k in KEYS] + [TIE],
                    label="Which response was best?",
                    visible=SHOW_COMPARE,
                )
                save_btn = gr.Button(
                    "Save preferred", variant="primary", visible=SHOW_COMPARE
                )
            status = gr.Markdown("")

    # Wiring -----------------------------------------------------------------
    respond_inputs = [msg, state, temperature, max_tokens]
    respond_outputs = [state, *chatbots, *metrics, msg, preferred, debug_box]

    preset_outputs = [system_prompt, intro, var_names, *var_boxes]

    apply_inputs = [
        system_prompt,
        intro,
        temperature,
        max_tokens,
        var_names,
        *var_boxes,
    ]
    apply_outputs = [state, *chatbots, *metrics, status, preferred]

    def _clear_voice():
        return None, ""

    def set_input_mode(mode):
        """Voice mode shows the mic; Text mode shows the text bar."""
        voice = mode == "🎤 Voice"
        return gr.update(visible=not voice), gr.update(visible=voice)

    input_mode.change(
        set_input_mode, inputs=[input_mode], outputs=[text_input_col, voice_input_col]
    )

    msg.submit(respond_all, respond_inputs, respond_outputs).then(
        _clear_voice, outputs=[audio_in, stt_status]
    )
    send_btn.click(respond_all, respond_inputs, respond_outputs).then(
        _clear_voice, outputs=[audio_in, stt_status]
    )

    new_btn.click(start_session, inputs=apply_inputs, outputs=apply_outputs)

    # Picking a preset loads its prompt + intro and its fixed fill-in variables,
    # then starts a fresh session.
    preset.change(
        load_preset_fields, inputs=[preset], outputs=preset_outputs
    ).then(start_session, inputs=apply_inputs, outputs=apply_outputs)

    save_btn.click(save_preferred, inputs=[state, preferred], outputs=[status])

    # Voice input: transcribe via STT into the message box when recording stops.
    audio_in.stop_recording(
        transcribe, inputs=[audio_in, msg], outputs=[msg, stt_status]
    )

    # On page load, populate the default preset's prompt/intro + variable fields.
    demo.load(load_preset_fields, inputs=[preset], outputs=preset_outputs)


if __name__ == "__main__":
    launch_kwargs = {}
    if config.AUTH_ENABLED:
        if config.AUTH_IS_DEFAULT:
            print(
                "⚠️  No APP_PASSWORD/APP_AUTH set — using the default dev login "
                "(user / slm-demo). Set APP_USERNAME + APP_PASSWORD (or APP_AUTH) "
                "in your environment / Space Secrets to secure the app."
            )
        launch_kwargs["auth"] = config.get_auth()
        launch_kwargs["auth_message"] = (
            "Please log in to use the comparison playground."
        )
    else:
        print(
            "ℹ️  ENABLE_AUTH=false — login is disabled; the app is open to "
            "anyone who can reach it."
        )
    demo.launch(**launch_kwargs)