chat.py

"""Gradio adapter layer — wires :mod:`core.chat` into Gradio events.

All pure conversation logic lives in :mod:`core.chat`. This module only deals
with Gradio types (``gr.update``, ``gr.Warning``) and yields the tuples that
the Blocks event handlers expect.

Architecture
------------
The chat surface is a custom HTML container (``<div id="hy-chat">``) owned
and mutated by ``static/chat.js``. The server ships **deltas**, never full
snapshots, via two hidden bridge components:

* ``chat_delta`` (``gr.HTML``, ``elem_id="hy-chat-delta"``): a single
  ``<div>`` whose textContent is the latest JSON payload of the form
  ``{"seq": N, "ops": [...]}``. The browser observes mutations on this
  element, parses, and applies each op to the chat container.
* ``busy_marker`` (``gr.HTML``, ``elem_id="hy-busy-marker"``): drives the
  Send button's enabled state. ``HY_BUSY_HTML`` while the model is
  generating, ``HY_IDLE_HTML`` when it's done.

Delta wire payload is just the new characters, and the client only
re-renders the trailing in-progress paragraph; everything above it is
frozen DOM that never changes again.

Send-button busy-state ownership
--------------------------------
The Send button's busy state is **owned entirely by the browser** — there
is NO server-side ``gr.update(interactive=...)`` toggling. See
``static/app.js`` for the MutationObserver-based gate.

The terminal IDLE-marker frame is shipped *separately* from any
chat-content frame so that even when WebSocket is back-pressured the
tiny marker frame still slips through and the UI unblocks promptly.
"""

from __future__ import annotations

import itertools
import json
import logging
from typing import Iterator

import gradio as gr

from core.chat import (
    ChatState,
    build_api_kwargs,
    finalize_response,
    flush_tool_results,
    init_state,
    record_tool_result,
    reset_state_in_place,
    stream_response,
)
from display import (
    format_tool_call_prompt,
    format_tool_calls_for_display,
    preprocess_latex,
)
from i18n import t

logger = logging.getLogger(__name__)

# Sentinel values driven into the dedicated #hy-busy-marker gr.HTML
# component on the server side, observed by the browser to drive the Send
# button's enabled state.
HY_BUSY_HTML = '<span data-hy-busy="1" data-hy-streaming="1" hidden aria-hidden="true"></span>'
HY_IDLE_HTML = '<span data-hy-busy="0" hidden aria-hidden="true"></span>'

# Initial value for the chat container (the inner <div> that the renderer
# owns). Wrapped in another <div> so Gradio's gr.HTML wrapper has a stable
# child to host the chat surface.
HY_CHAT_INITIAL_HTML = '<div id="hy-chat" class="hy-chat" role="log" aria-live="polite"></div>'

# Thread-safe monotonically-increasing sequence numbers for the delta
# channel. Each delta payload carries one; the client uses it to dedup
# (Gradio occasionally re-fires the same value on reconnect, and we never
# want to re-apply the same ops twice).
_seq_counter = itertools.count(1)


def _next_seq() -> int:
    return next(_seq_counter)


# Bubble IDs are pre-allocated per assistant turn so all ops for the same
# bubble (assistant_begin → reasoning_delta → content_delta → assistant_end →
# tool_call) can be routed to the right DOM node on the client.
_bubble_counter = itertools.count(1)


def _next_bubble_id() -> str:
    return f"a-{next(_bubble_counter)}"


def _delta(ops: list[dict], epoch: int | None = None) -> str:
    """Encode a list of ops as a delta payload string.

    The payload is shipped through ``gr.HTML``'s value, which the front
    end renders into the DOM via ``innerHTML``. Plain ``json.dumps``
    output is valid HTML text 99 % of the time, but any user- or
    model-supplied ``<`` would be parsed as the opening of an HTML tag
    and silently truncate the payload (also a stored-XSS sink). We
    pre-escape the three HTML-meta characters into their JSON
    ``\\uNNNN`` form — the browser's ``JSON.parse`` reverses the escapes
    so semantics are preserved end-to-end.

    ``epoch`` tags the payload with the session's chat-epoch id so the
    client can drop stale deltas that arrive after a "+" / new-chat
    reset. The reset payload itself carries the NEW epoch — the client
    adopts it, then ignores any later payload that still carries the
    old one.
    """
    payload: dict = {"seq": _next_seq(), "ops": ops}
    if epoch is not None:
        payload["epoch"] = epoch
    raw = json.dumps(payload, ensure_ascii=False)
    return (
        raw.replace("<", "\\u003c")
           .replace(">", "\\u003e")
           .replace("&", "\\u0026")
    )


# Re-export so ``from chat import init_state`` keeps working in app.py.
__all__ = [
    "init_state",
    "send_message",
    "new_chat",
    "submit_tool_result",
    "api_chat",
    "HY_BUSY_HTML",
    "HY_IDLE_HTML",
    "HY_CHAT_INITIAL_HTML",
]


def _ensure_state(state) -> ChatState:
    """Normalize ``state`` into a fresh ``ChatState`` dict.

    ``gr.State(init_state)`` treats ``init_state`` as a per-session factory in
    the browser, but stateless ``gradio_client`` calls hand the factory itself
    to the handler. Accept callable/None and fall back to a fresh state so the
    API works without breaking the in-browser per-session semantics.
    """
    if state is None or callable(state):
        return init_state()
    return state


def _validate_user_message(message: str, functions_json_str: str) -> bool:
    if not message or not message.strip():
        gr.Warning(t("warn.empty_msg"))
        return False
    if functions_json_str and functions_json_str.strip():
        try:
            json.loads(functions_json_str)
        except json.JSONDecodeError as e:
            gr.Warning(t("warn.invalid_fn_json", err=str(e)))
            return False
    return True


def _run_streaming_turn(
    state: ChatState,
    bubble_id: str,
    epoch: int,
    system_prompt: str,
    think_level: str,
    temperature: float,
    temperature_use_default: bool,
    max_tokens: int,
    top_p: float,
    functions_json_str: str,
) -> Iterator[tuple[list[dict], bool, str]]:
    """Stream a single turn. Yields ``(ops, show_tool, tool_md)``.

    ``ops`` is the list of bubble-scoped ops (each carrying ``id=bubble_id``)
    to ship on this frame. ``show_tool`` / ``tool_md`` carry the tool-area
    UI state for the caller to translate into ``gr.update(...)`` outputs.

    ``epoch`` is the chat-epoch captured by the caller at turn start.
    We compare it against the live ``state["epoch"]`` between yields;
    if they diverge the user clicked "+", so we stop yielding and let
    the caller exit without finalising the assistant turn.
    """
    # Collapse the (slider value, "Use model default" checkbox) pair into a
    # single tristate that ``build_api_kwargs`` understands: ``None`` means
    # "omit the field, let the server apply its own default"; any float
    # (including 0, i.e. greedy decoding) is sent literally.
    resolved_temperature: float | None = (
        None if temperature_use_default else float(temperature)
    )
    kwargs = build_api_kwargs(
        state, system_prompt, functions_json_str,
        think_level, resolved_temperature, max_tokens, top_p,
    )
    if logger.isEnabledFor(logging.DEBUG):
        # Dump the actual request body (minus the bulky messages list) so we
        # can verify which sampling knobs reached the wire.
        loggable = {k: v for k, v in kwargs.items() if k != "messages"}
        logger.debug("streaming turn kwargs: %r", loggable)

    def _cancelled() -> bool:
        return state.get("epoch") != epoch

    last_ac, last_rc, last_tca = "", "", []
    seen_first_content = False
    for ops, ac, rc, tca, _rid in stream_response(kwargs, is_cancelled=_cancelled):
        if _cancelled():
            return
        last_ac, last_rc, last_tca = ac, rc, tca
        if not ops:
            # Heartbeat — no new content. We still surface an empty-ops
            # frame so the wire ticks (Gradio dedups equal values; the
            # heartbeat carries the new seq via the delta payload below).
            yield [], False, ""
            continue
        # Apply per-op LaTeX preprocessing on content_delta. The model may
        # emit bare \begin{equation}...\end{equation} blocks etc.; KaTeX
        # only renders math inside $$ delimiters.
        out_ops: list[dict] = []
        for op in ops:
            if op["type"] == "content_delta":
                # First content delta implicitly closes the thinking block on
                # the client. We tag the FIRST content op with a flag so the
                # client can switch the summary text and auto-collapse.
                delta_text = preprocess_latex(op["delta"])
                wrapped = {
                    "type": "content_delta",
                    "id": bubble_id,
                    "delta": delta_text,
                }
                if not seen_first_content:
                    wrapped["thinking_done"] = True
                    seen_first_content = True
                out_ops.append(wrapped)
            elif op["type"] == "reasoning_delta":
                out_ops.append({
                    "type": "reasoning_delta",
                    "id": bubble_id,
                    "delta": op["delta"],
                })
            elif op["type"] == "tool_calls":
                # Carry the snapshot through; we materialise it as a UI
                # block only at end-of-turn (see the finalize section
                # below). Skipping here keeps mid-stream wire payloads
                # tiny — the same partial tool-call gets re-shipped on
                # every frame otherwise.
                continue
        if out_ops:
            yield out_ops, False, ""

    # ── Terminal phase ─────────────────────────────────────────────────
    if _cancelled():
        return
    has_tools, pending = finalize_response(
        state, last_ac, last_rc, last_tca,
    )
    end_ops: list[dict] = [{"type": "assistant_end", "id": bubble_id}]
    if has_tools:
        # Render the tool-call display block as Markdown; the client will
        # run marked over it. Append AFTER the assistant_end op so the DOM
        # ordering stays correct (content → tool calls).
        tc_markdown = format_tool_calls_for_display(pending)
        end_ops.append({
            "type": "tool_call",
            "id": bubble_id,
            "markdown": tc_markdown,
        })
        tool_md = format_tool_call_prompt(pending[0], 1, len(pending))
        yield end_ops, True, tool_md
    else:
        yield end_ops, False, ""


# ─── send_message ────────────────────────────────────────────────────────────
# Output slot order (must match app.py's ``send_outputs``):
#   0 chat_delta (gr.HTML, elem_id="hy-chat-delta")
#   1 state
#   2 msg_input
#   3 tool_area
#   4 tool_call_info
#   5 tool_result_input
#   6 busy_marker (gr.HTML, elem_id="hy-busy-marker")
def send_message(
    message, state,
    system_prompt, think_level,
    temperature, temperature_use_default,
    max_tokens, top_p,
    functions_json_str,
):
    state = _ensure_state(state)

    if not _validate_user_message(message, functions_json_str):
        # Validation failure — never started streaming, marker stays idle,
        # no chat ops emitted.
        yield (
            gr.update(), state, gr.update(),
            gr.update(), gr.update(), gr.update(),
            HY_IDLE_HTML,
        )
        return

    message = message.strip()
    state["messages"].append({"role": "user", "content": message})
    bubble_id = _next_bubble_id()
    # Capture the chat epoch at the start of the turn. ``new_chat`` bumps
    # ``state["epoch"]`` in place; once that happens we stop yielding so
    # the new fresh chat surface stays clean.
    epoch = state["epoch"]

    def _cancelled() -> bool:
        return state.get("epoch") != epoch

    # First frame: append the user bubble, open a new assistant bubble
    # (with typing indicator), set busy marker, clear input.
    yield (
        _delta([
            {"type": "user", "text": message},
            {"type": "assistant_begin", "id": bubble_id},
        ], epoch=epoch),
        state, "",
        gr.update(visible=False), gr.update(), gr.update(),
        HY_BUSY_HTML,
    )

    try:
        for ops, show_tool, tool_md in _run_streaming_turn(
            state, bubble_id, epoch,
            system_prompt, think_level,
            temperature, temperature_use_default,
            max_tokens, top_p,
            functions_json_str,
        ):
            if _cancelled():
                return
            # Heartbeat / no-op frame: still ship an empty-ops payload so
            # the wire ticks and the proxy doesn't drop the connection.
            payload = _delta(ops, epoch=epoch) if ops else _delta([], epoch=epoch)
            if show_tool:
                yield (
                    payload, state, gr.update(),
                    gr.update(visible=True), gr.update(value=tool_md), gr.update(value=""),
                    gr.update(),
                )
            else:
                yield (
                    payload, state, gr.update(),
                    gr.update(visible=False), gr.update(), gr.update(),
                    gr.update(),
                )

    except Exception as e:
        if _cancelled():
            # Reset already cleaned the UI; swallow late-arriving errors
            # from the abandoned upstream so we don't bounce a stale
            # warning at the user.
            return
        logger.exception("send_message failed: %s", e)
        if state["messages"] and state["messages"][-1].get("role") == "user":
            state["messages"].pop()
        gr.Warning(t("warn.request_failed"))
        # Drop the typing-indicator placeholder bubble client-side, restore
        # the user's message into the input box, release the busy marker.
        yield (
            _delta([{"type": "remove_bubble", "id": bubble_id}], epoch=epoch),
            state, gr.update(value=message),
            gr.update(visible=False), gr.update(), gr.update(),
            HY_IDLE_HTML,
        )
        return

    if _cancelled():
        return

    # ── Terminal IDLE-marker frame ──────────────────────────────────────
    # Only changes the busy marker; every other slot is gr.update(). Tens
    # of bytes on the wire — slips past any WebSocket back-pressure so the
    # browser MutationObserver fires the moment the stream genuinely ends,
    # releasing the Send button.
    yield (
        gr.update(), state, gr.update(),
        gr.update(), gr.update(), gr.update(),
        HY_IDLE_HTML,
    )


# ─── Programmatic API endpoint ───────────────────────────────────────────────
# The UI handlers above (``send_message`` / ``submit_tool_result``) ship chat
# content as DOM-mutation deltas through a hidden ``chat_delta`` HTML
# component, which is unusable from ``gradio_client``: a remote caller would
# only ever see the final ``gr.update()`` placeholders and the busy marker.
#
# ``api_chat`` is the headless counterpart: a stateless, generator function
# whose YIELDED tuple is the entire reply so far. ``gradio_client.predict``
# returns the last yield, so callers get the full ``(content, reasoning,
# tool_calls, history)`` snapshot once the stream finishes; users that want
# token-by-token streaming can iterate via ``client.submit(...)``.
def api_chat(
    message: str,
    system_prompt: str = "",
    history: list | None = None,
    think_level: str = "high",
    temperature: float | None = None,
    max_tokens: int = 0,
    top_p: float = 0,
    functions_json_str: str = "",
) -> Iterator[tuple[str, str, list, list]]:
    """Stateless chat endpoint for ``gradio_client`` callers.

    Args:
        message: The new user turn.
        system_prompt: Optional system prompt prepended to every request.
        history: Prior conversation as a list of OpenAI-style messages
            (``[{"role": "user"|"assistant"|"tool", "content": ...}, ...]``).
            Pass ``None`` or ``[]`` to start a fresh conversation. The caller
            owns the history; pass the returned ``updated_history`` back on
            the next call to continue multi-turn.
        think_level: One of ``"no_think" | "low" | "high"``.
        temperature: Sampling temperature. Pass ``None`` (the default) to omit
            the field from the API request and let the server apply its own
            default; pass any float (including ``0`` for greedy decoding) to
            send it literally.
        max_tokens, top_p: Sampling knobs. Pass ``0`` (the default) to omit
            the field from the API request and let the server apply its own
            default.
        functions_json_str: JSON string of tool definitions (OpenAI tools
            schema). Empty string disables function calling.

    Yields:
        ``(content, reasoning_content, tool_calls, updated_history)``
        cumulative on every yield. The final yield is the complete reply.

        * ``content``: assistant visible text.
        * ``reasoning_content``: chain-of-thought (when ``think_level``
          requests it). May be empty for ``no_think``.
        * ``tool_calls``: list of OpenAI-style tool-call dicts the model
          requested. Empty when no function was called. The caller is
          responsible for executing each call and appending the
          corresponding ``{"role": "tool", "tool_call_id": ..., "content":
          ...}`` messages to ``history`` on the next ``api_chat`` call.
        * ``updated_history``: full message list including the new user
          turn and the assistant reply (with ``tool_calls`` attached when
          present). Suitable for round-tripping into the next call.

    Raises:
        ValueError: when ``message`` is empty or ``functions_json_str`` is
            not valid JSON.
    """
    if not message or not message.strip():
        raise ValueError("message must be a non-empty string")
    if functions_json_str and functions_json_str.strip():
        try:
            json.loads(functions_json_str)
        except json.JSONDecodeError as e:
            raise ValueError(f"functions_json_str is not valid JSON: {e}") from e

    state = init_state()
    if history:
        # Defensive copy — never mutate the caller's list.
        state["messages"] = [dict(m) for m in history]
    state["messages"].append({"role": "user", "content": message.strip()})

    kwargs = build_api_kwargs(
        state, system_prompt, functions_json_str,
        think_level, temperature, max_tokens, top_p,
    )

    last_ac, last_rc = "", ""
    last_tca: list[dict] = []
    yielded = False
    for _ops, ac, rc, tca, _rid in stream_response(kwargs):
        last_ac, last_rc, last_tca = ac, rc, tca
        # Snapshot the in-flight assistant turn into a transient history
        # view so streaming consumers see the assistant text growing.
        in_flight = list(state["messages"])
        in_flight.append({
            "role": "assistant",
            "content": ac or "",
            **({"reasoning_content": rc} if rc else {}),
            **({"tool_calls": list(tca)} if tca else {}),
        })
        yield ac or "", rc or "", list(tca), in_flight
        yielded = True

    finalize_response(state, last_ac, last_rc, last_tca)
    final_history = list(state["messages"])

    # Always emit a terminal frame so callers consuming only the last yield
    # observe the persisted history (the in-flight snapshots above attach a
    # provisional assistant message; this one is the canonical record).
    if not yielded:
        yield "", "", [], final_history
    else:
        yield last_ac or "", last_rc or "", list(last_tca), final_history


def new_chat(state):
    """Reset the conversation.

    Critical: when there's an in-flight stream we MUTATE the existing
    state dict in place (bumping its epoch) instead of returning a brand
    new dict. The streaming generator holds a reference to this same
    dict and polls ``state["epoch"]`` between yields — the in-place
    bump is what tells it to abandon the rest of its turn. Returning a
    fresh dict would leave the generator pointed at a stale dict it
    would happily keep streaming into.

    We also push ``HY_IDLE_HTML`` to the busy marker so the Send button
    unlocks immediately. Without this the cancelled generator never
    yields its terminal IDLE frame, so the button would stay stuck
    until the soft watchdog in static/app.js force-clears it (~4s).
    """
    if not state or not isinstance(state, dict) or not state.get("messages"):
        gr.Info(t("info.new_chat"))
        return gr.update(), state, gr.update()
    new_epoch = reset_state_in_place(state)
    return (
        _delta([{"type": "reset"}], epoch=new_epoch),
        state,
        HY_IDLE_HTML,
    )


# ─── submit_tool_result ──────────────────────────────────────────────────────
# Output slot order (must match app.py's ``tool_submit_outputs``):
#   0 chat_delta
#   1 state
#   2 tool_area
#   3 tool_call_info
#   4 tool_result_input
#   5 busy_marker
def submit_tool_result(
    result_text, state,
    system_prompt, think_level,
    temperature, temperature_use_default,
    max_tokens, top_p,
    functions_json_str,
):
    state = _ensure_state(state)
    pending = state.get("pending_tool_calls", [])
    if not pending:
        yield (
            gr.update(), state,
            gr.update(visible=False), gr.update(), gr.update(),
            HY_IDLE_HTML,
        )
        return

    record_tool_result(state, pending[0], result_text)
    state["pending_tool_calls"] = pending[1:]

    if state["pending_tool_calls"]:
        # More tool results still pending — keep the tool area visible,
        # marker stays idle (we're not streaming yet).
        remaining = state["pending_tool_calls"]
        submitted = state.get("submitted_tool_results", [])
        total = len(remaining) + len(submitted)
        current_idx = len(submitted) + 1
        tool_info_md = format_tool_call_prompt(remaining[0], current_idx, total)
        yield (
            gr.update(), state,
            gr.update(visible=True), gr.update(value=tool_info_md), gr.update(value=""),
            HY_IDLE_HTML,
        )
        return

    flush_tool_results(state)
    bubble_id = _next_bubble_id()
    epoch = state["epoch"]

    def _cancelled() -> bool:
        return state.get("epoch") != epoch

    # First streaming frame: open a new assistant bubble for the tool
    # follow-up response, hide the tool area, set the busy marker.
    yield (
        _delta([{"type": "assistant_begin", "id": bubble_id}], epoch=epoch),
        state,
        gr.update(visible=False), gr.update(), gr.update(),
        HY_BUSY_HTML,
    )

    try:
        for ops, show_tool, tool_md in _run_streaming_turn(
            state, bubble_id, epoch,
            system_prompt, think_level,
            temperature, temperature_use_default,
            max_tokens, top_p,
            functions_json_str,
        ):
            if _cancelled():
                return
            payload = _delta(ops, epoch=epoch) if ops else _delta([], epoch=epoch)
            if show_tool:
                yield (
                    payload, state,
                    gr.update(visible=True), gr.update(value=tool_md), gr.update(value=""),
                    gr.update(),
                )
            else:
                yield (
                    payload, state,
                    gr.update(visible=False), gr.update(), gr.update(),
                    gr.update(),
                )

    except Exception as e:
        if _cancelled():
            return
        logger.exception("submit_tool_result failed: %s", e)
        gr.Warning(t("warn.request_failed"))
        yield (
            _delta([{"type": "remove_bubble", "id": bubble_id}], epoch=epoch),
            state,
            gr.update(visible=False), gr.update(), gr.update(),
            HY_IDLE_HTML,
        )
        return

    if _cancelled():
        return

    yield (
        gr.update(), state,
        gr.update(), gr.update(), gr.update(),
        HY_IDLE_HTML,
    )