server.py

"""Tencent-branded Hy3 chat demo, served from a :class:`gradio.Server`.

This module replaces the previous ``gr.Blocks`` entry point (``app.py``)
with a :class:`gradio.Server` backend that pairs a fully custom, Tencent-
branded frontend (``static/index.html``) with Gradio's queue / concurrency
engine. See the HF article "gradio.Server: Any Custom Frontend with
Gradio's Backend" for the pattern.

Why a Server instead of Blocks?
-------------------------------
The chat surface is a bespoke web app (drag-friendly message bubbles,
collapsible reasoning, in-place markdown + KaTeX + highlight.js, tool-call
panels, an empty-state examples grid). None of that maps cleanly onto
Gradio components. ``gradio.Server`` lets us ship our own HTML/JS frontend
while keeping:

  * Gradio's queue + concurrency control (``@app.api`` goes through the
    queue; multiple users are serialized, GPU requests don't collide),
  * SSE streaming of the generator's per-yield snapshots,
  * ``gradio_client`` compatibility (the endpoints are callable from
    Python too),
  * Hugging Face Spaces hosting + ZeroGPU plumbing.

Endpoints
---------
``/``                       GET   — serves ``static/index.html``.
``/static/<path>``          GET   — static asset passthrough (logo, fonts).
``/chat``                   api  — streaming chat (generator). The browser
                                   consumes it via the ``@gradio/client`` JS
                                   module's async-iterator.
``/new_chat``               api  — reset the conversation server-side and
                                   return a fresh epoch token.
``/submit_tool_result``     api  — submit a function return value, then
                                   stream the model's follow-up reply.
``/validate_functions``     api  — validate + pretty-print tool JSON.

The chat *logic* is unchanged: we reuse ``core.chat`` and the headless
``api_chat`` generator from ``chat.py``. ``api_chat`` already yields the
cumulative ``(content, reasoning, tool_calls, updated_history)`` snapshot
per frame — exactly the shape the browser renders incrementally. We wrap
it with session state (a per-session ``ChatState`` + epoch) so the
multi-turn + tool-calling UX works without the user round-tripping the
full history each turn.
"""

from __future__ import annotations

import logging
import os
import threading
from pathlib import Path
from typing import Any, Iterator

from fastapi.responses import HTMLResponse, FileResponse
from fastapi.staticfiles import StaticFiles

from gradio import Server

from chat import api_chat
from core.chat import (
    ChatState,
    finalize_response,
    flush_tool_results,
    init_state,
    record_tool_result,
    reset_state_in_place,
)
from tools import parse_functions_json

logger = logging.getLogger(__name__)

_STATIC_DIR = Path(__file__).parent / "static"
_INDEX_HTML = _STATIC_DIR / "index.html"

# ---------------------------------------------------------------------------
# Per-session state
# ---------------------------------------------------------------------------
# ``gradio.Server`` endpoints are stateless at the transport layer (each
# call is an independent HTTP request), so multi-turn conversation state
# has to live somewhere we can re-attach on the next call. We keep an
# in-process registry keyed by a session id the browser holds (a random
# token minted on first load and persisted in ``sessionStorage``). Each
# entry is a ``ChatState`` plus a lock so a fast double-send on the same
# session can't race two streams into the same message log.
#
# This is intentionally simple and lives in-process. On HF Spaces (single
# worker, sticky-ish routing) it's adequate for a demo; a horizontally-
# scaled deployment would externalise this into a shared store.

_SESSIONS: dict[str, dict] = {}
_SESSIONS_LOCK = threading.Lock()


def _get_session(session_id: str) -> dict:
    """Return (creating if needed) the session record for ``session_id``."""
    with _SESSIONS_LOCK:
        sess = _SESSIONS.get(session_id)
        if sess is None:
            sess = {"state": init_state(), "lock": threading.Lock()}
            _SESSIONS[session_id] = sess
        return sess


# ---------------------------------------------------------------------------
# Server + routes
# ---------------------------------------------------------------------------
app = Server(title="Hy3 · Tencent")


@app.get("/", response_class=HTMLResponse)
async def homepage() -> HTMLResponse:
    """Serve the custom Tencent-branded frontend."""
    return HTMLResponse(_INDEX_HTML.read_text(encoding="utf-8"))


# Static asset passthrough (logo, vendored libs the frontend pulls from the
# same origin instead of a CDN, etc.). Mounted under /static so it never
# collides with Gradio's own routes.
app.mount(
    "/static",
    StaticFiles(directory=str(_STATIC_DIR)),
    name="hy-static",
)


# ---------------------------------------------------------------------------
# API endpoints
# ---------------------------------------------------------------------------
# All endpoints accept a ``session_id`` so we can re-attach the right
# ``ChatState`` across the stateless HTTP boundary. The frontend mints one
# on first load and reuses it for the lifetime of the tab.
#
# ``concurrency_limit`` mirrors the Blocks app's chat pool: many sessions
# can stream concurrently. ``stream_every`` is small so the browser sees
# near-token-rate updates.

CHAT_CONCURRENCY = 64


def _resolve_session(session_id: Any) -> tuple[dict, ChatState]:
    sess = _get_session(session_id or "")
    return sess, sess["state"]


@app.api(
    name="chat",
    concurrency_limit=CHAT_CONCURRENCY,
    concurrency_id="chat",
    stream_every=0.06,
)
def chat(
    message: str,
    session_id: str = "",
    system_prompt: str = "",
    think_level: str = "high",
    temperature: float | None = None,
    max_tokens: int = 0,
    top_p: float = 0,
    preserved_thinking: bool | None = None,
    functions_json_str: str = "",
) -> Iterator[tuple[str, str, list, list, str]]:
    """Stream a chat turn for ``session_id``.

    Yields cumulative ``(content, reasoning, tool_calls, history, epoch)``
    snapshots as the model streams. The browser renders each frame
    incrementally; the final frame is the complete reply with the
    canonical, persisted ``history``.

    ``epoch`` is returned so the frontend can drop frames that arrived
    after the user clicked "new chat" (the same staleness guard the
    Blocks app uses, just over the stateless boundary now).
    """
    sess, state = _resolve_session(session_id)
    with sess["lock"]:
        epoch = state["epoch"]
        # ``api_chat`` is stateless (it builds its own internal state from
        # the supplied ``history``), so hand it the live session history
        # and then write the finalised history back into the session.
        history = list(state["messages"])

        for content, reasoning, tool_calls, updated_history in api_chat(
            message=message,
            system_prompt=system_prompt,
            history=history,
            think_level=think_level,
            temperature=temperature,
            max_tokens=max_tokens,
            top_p=top_p,
            preserved_thinking=preserved_thinking,
            functions_json_str=functions_json_str,
        ):
            yield content, reasoning, tool_calls, updated_history, str(epoch)

        # Persist the final history the headless generator produced.
        # ``api_chat`` appends the assistant turn (and tool_calls, when
        # the model called a function) into ``updated_history``; we adopt
        # it verbatim so the next turn sees the full conversation.
        state["messages"] = [dict(m) for m in updated_history]

        # If the model requested tool calls, mirror the pending queue the
        # Blocks app would have built so ``submit_tool_result`` can pick
        # up from here.
        _sync_pending_tool_calls(state, updated_history)


def _sync_pending_tool_calls(state: ChatState, updated_history: list[dict]) -> None:
    """Re-derive ``pending_tool_calls`` from the latest assistant message.

    ``api_chat`` doesn't touch ``state`` at all (it works off a defensive
    copy), so after it finishes we reconstruct the tool-call pending queue
    from the assistant message it appended. This keeps the
    ``submit_tool_result`` flow working: the next call sees the same
    ``pending_tool_calls`` list the Blocks UI would have.
    """
    for msg in reversed(updated_history):
        if msg.get("role") == "assistant":
            tcs = msg.get("tool_calls")
            if tcs:
                state["pending_tool_calls"] = [dict(t) for t in tcs]
                state["submitted_tool_results"] = []
                state["pending_assistant_msg"] = dict(msg)
            else:
                state["pending_tool_calls"] = []
                state["pending_assistant_msg"] = None
            return


@app.api(name="new_chat", concurrency_limit=CHAT_CONCURRENCY)
def new_chat(session_id: str = "") -> dict:
    """Reset the session's conversation and return the new epoch.

    Returns ``{"epoch": "<new epoch>", "ok": true}``. The frontend adopts
    the new epoch and drops any in-flight frames from the abandoned turn.
    """
    sess, state = _resolve_session(session_id)
    with sess["lock"]:
        new_epoch = reset_state_in_place(state)
    return {"epoch": str(new_epoch), "ok": True}


@app.api(
    name="submit_tool_result",
    concurrency_limit=CHAT_CONCURRENCY,
    concurrency_id="chat",
    stream_every=0.06,
)
def submit_tool_result(
    result_text: str,
    session_id: str = "",
    system_prompt: str = "",
    think_level: str = "high",
    temperature: float | None = None,
    max_tokens: int = 0,
    top_p: float = 0,
    preserved_thinking: bool | None = None,
    functions_json_str: str = "",
) -> Iterator[tuple[str, str, list, list, str, bool]]:
    """Submit one tool-call result, then stream the model's follow-up.

    If multiple tool calls were requested, the FIRST pending result is
    recorded; if more remain, we yield a single empty frame with
    ``needs_more_results=True`` and stop. Once all results are in, we
    flush them into the message log and stream the follow-up reply with
    ``needs_more_results=False`` (same shape as ``chat`` for the extra
    trailing bool).
    """
    sess, state = _resolve_session(session_id)
    with sess["lock"]:
        pending = state.get("pending_tool_calls", [])
        if not pending:
            yield "", "", [], list(state["messages"]), str(state["epoch"]), False
            return

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

        if state["pending_tool_calls"]:
            # Still waiting on more results — yield the REMAINING pending
            # tool calls in the tool_calls slot so the frontend re-opens
            # the tool panel for the next one. (Empty content/reasoning;
            # the trailing ``True`` flags "more results pending".)
            yield (
                "", "", list(state["pending_tool_calls"]),
                list(state["messages"]), str(state["epoch"]), True,
            )
            return

        # All results in — flush them into the message log, then stream the
        # model's follow-up turn over the tool-result-augmented history.
        # We drive ``core.chat.stream_response`` directly (rather than
        # ``api_chat``, which requires a non-empty user message) so the
        # tool results are the last thing the model sees — exactly what
        # the Blocks app does in ``chat.submit_tool_result``.
        flush_tool_results(state)
        epoch = state["epoch"]

        from core.chat import (
            build_api_kwargs as _build_kwargs,
            stream_response as _stream,
        )
        kwargs = _build_kwargs(
            state, system_prompt, functions_json_str,
            think_level, temperature, max_tokens, top_p,
            preserved_thinking,
        )

        last_ac, last_rc, last_tca = "", "", []
        for ops, ac, rc, tca, _rid in _stream(kwargs):
            last_ac, last_rc, last_tca = ac, rc, tca
            # Surface a cumulative snapshot on every frame so the frontend
            # can render incrementally — same shape as /chat's yield.
            in_flight = list(state["messages"]) + [{
                "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, str(epoch), False

        # Persist the finalised assistant turn.
        finalize_response(state, last_ac, last_rc, last_tca)
        _sync_pending_tool_calls(state, state["messages"])


@app.api(name="config", concurrency_limit=32)
def config() -> dict:
    """Return the showcase example prompts + i18n strings for the frontend.

    Lazily imports the curated prompts from :mod:`app` (which also builds
    the legacy Blocks, but only at this first call) and the translation
    table from :mod:`i18n`. Kept in its own endpoint so the static HTML
    can stay a plain file and the curated content lives in one place.
    """
    from app import EXAMPLES  # noqa: WPS433 — lazy to keep server import light
    from i18n import LANG, t, TRANSLATIONS

    table = TRANSLATIONS.get(LANG, TRANSLATIONS["en"])
    # Ship the keys the frontend actually reads; everything else stays
    # server-side. Mirrors the HY_I18N object the Blocks app injected.
    i18n = {
        "title": t("title"),
        "title.notice_html": t("title.notice_html"),
        "examples_heading": t("examples_heading"),
        "msg_placeholder": t("msg_placeholder"),
        "sidebar.think_level": t("sidebar.think_level"),
        "sidebar.think_level.info": t("sidebar.think_level.info"),
        "sidebar.system_prompt": t("sidebar.system_prompt"),
        "sidebar.system_prompt.placeholder": t("sidebar.system_prompt.placeholder"),
        "sidebar.temperature": t("sidebar.temperature"),
        "sidebar.temperature.info": t("sidebar.temperature.info"),
        "sidebar.temperature.use_default": t("sidebar.temperature.use_default"),
        "sidebar.max_tokens": t("sidebar.max_tokens"),
        "sidebar.max_tokens.info": t("sidebar.max_tokens.info"),
        "sidebar.top_p": t("sidebar.top_p"),
        "sidebar.top_p.info": t("sidebar.top_p.info"),
        "sidebar.preserved_thinking": t("sidebar.preserved_thinking"),
        "sidebar.preserved_thinking.info": t("sidebar.preserved_thinking.info"),
        "sidebar.preserved_thinking.use_default": t("sidebar.preserved_thinking.use_default"),
        "sidebar.preserved_thinking.use_default.info": t("sidebar.preserved_thinking.use_default.info"),
        "sidebar.functions": t("sidebar.functions"),
        "sidebar.functions.label": t("sidebar.functions.label"),
        "sidebar.validate_btn": t("sidebar.validate_btn"),
        "tool.result_label": t("tool.result_label"),
        "tool.result_placeholder": t("tool.result_placeholder"),
        "tool.submit": t("tool.submit"),
        "tool.call_header": t("tool.call_header"),
        "tool.call_label": t("tool.call_label"),
        "display.thinking": t("display.thinking"),
        "display.thinking_done": t("display.thinking_done"),
        "display.code_copy": t("display.code_copy"),
        "display.code_copied": t("display.code_copied"),
        "warn.empty_msg": t("warn.empty_msg"),
        "warn.model_busy": t("warn.model_busy"),
        "warn.request_failed": t("warn.request_failed"),
        "info.new_chat": t("info.new_chat"),
    }
    # The frontend needs the tool-call display label, which the Blocks app
    # rendered server-side via display.format_tool_calls_for_display. We
    # surface the raw label here and let the frontend build the markdown.
    return {"examples": list(EXAMPLES), "i18n": i18n, "lang": LANG, "table": table}


@app.api(name="validate_functions", concurrency_limit=8)
def validate_functions(functions_json_str: str) -> dict:
    """Validate + pretty-print a tools-JSON string.

    Returns ``{"ok": bool, "message": str, "formatted": str}``. Mirrors
    the Blocks app's "Validate & Format" button but without ``gr.Warning``
    (no Gradio chrome here to surface a toast through), so we return the
    human-readable result for the frontend to display.
    """
    if not functions_json_str or not functions_json_str.strip():
        return {"ok": False, "message": "Please enter function definition JSON", "formatted": ""}
    import json
    try:
        data = json.loads(functions_json_str)
    except json.JSONDecodeError as e:
        return {"ok": False, "message": f"Invalid JSON format: {e}", "formatted": ""}
    if isinstance(data, dict):
        data = [data]
    if not isinstance(data, list):
        return {"ok": False, "message": "JSON must be an array [...] or a single object {...}", "formatted": ""}
    names: list[str] = []
    for i, item in enumerate(data):
        if not isinstance(item, dict):
            return {"ok": False, "message": f"Item {i + 1} is not a JSON object", "formatted": ""}
        fn = _normalize(item)
        if not fn:
            return {
                "ok": False,
                "message": (
                    f"Item {i + 1} has invalid format, expected "
                    "{{type, function}} or {{name, parameters}}"
                ),
                "formatted": "",
            }
        if fn["name"] in names:
            return {"ok": False, "message": f"Duplicate function name '{fn['name']}'", "formatted": ""}
        names.append(fn["name"])
    formatted = json.dumps(data, indent=2, ensure_ascii=False)
    return {
        "ok": True,
        "message": f"Validation passed, {len(names)} function(s): {', '.join(names)}",
        "formatted": formatted,
    }


def _normalize(item: dict) -> dict | None:
    """Local copy of tools._normalize_tool_item (avoids a private import)."""
    if item.get("type") == "function" and "function" in item:
        fn = item["function"]
        if isinstance(fn, dict) and fn.get("name"):
            return fn
        return None
    if item.get("name"):
        return item
    return None


if __name__ == "__main__":
    # HuggingFace Spaces invokes ``python server.py``. Guarded so importing
    # the module (e.g. from tests) doesn't start the server as a side effect.
    app.launch(show_error=True)