path: root/app/backends/probe.py

"""Identify which TTS backend answers at a URL (remote-server probing).

The hub keeps locally-managed backends distinct from externally-run ones: a
server this tool started is tagged "[local]", and a server found by probing a
configured remote URL (``*_REMOTE_URL`` in ``app/converter/config.py``) is
tagged "[remote]". To know that a remote URL really hosts the backend we
think it does (and not some other HTTP service), each backend exposes a small
identity check over plain HTTP:

  * audio.cpp   ``GET /health`` -> ``{"status": "ok"}`` and ``GET /v1/models``
                -> ``{"data": [{"id": ...}, ...]}``.
  * faster      ``GET /health`` -> a JSON object with a ``model_loaded`` key.
  * qwen-tts    a Gradio app: ``GET /info`` -> ``named_endpoints`` containing
                the endpoint names the converter calls (``/run_instruct`` /
                ``/run_custom_voice`` / ``/generate_custom_voice`` for the
                CustomVoice demo; ``/run_voice_clone`` / ``/generate_voice_clone``
                for the Base demo; ``/run_voice_design`` for the VoiceDesign
                demo).

``identify_server`` returns one of the IDENTITY_* constants, or None when the
URL does not answer or answers as something unrecognized. It is stdlib-only
(urllib) and deliberately imports nothing from the other backend modules, so
it stays cheap to import alongside ``backends.common``.
"""

import json
import urllib.parse
import urllib.request
from typing import List, Optional

from backends import common

IDENTITY_AUDIOCPP = "audiocpp"
IDENTITY_FASTER = "faster"
IDENTITY_QWEN_CUSTOM = "qwen-custom"
IDENTITY_QWEN_CLONE = "qwen-clone"
IDENTITY_QWEN_DESIGN = "qwen-design"
IDENTITY_SGLOMNI = "sglomni"

# Endpoint names the converter resolves for each qwen demo server (see
# converter.clients QwenTTSClient). Mirror them here so identification matches
# exactly what the converter would call.
_QWEN_CUSTOM_ENDPOINTS = (
    "/run_instruct", "/run_custom_voice", "/generate_custom_voice")
_QWEN_CLONE_ENDPOINTS = ("/run_voice_clone", "/generate_voice_clone")
_QWEN_DESIGN_ENDPOINTS = ("/run_voice_design",)

DEFAULT_TIMEOUT = 3.0


def identify_server(url: str, timeout: float = DEFAULT_TIMEOUT) -> Optional[str]:
    """Return the backend identity answering at URL, or None.

    A cheap TCP-connect gate runs first (``common.server_running``) so a dead
    or unrouteable host returns quickly; the HTTP probes only run when
    something is listening. Returns None when the URL is empty/unparsable,
    unreachable, or answers as none of the known backends.
    """
    if not url:
        return None
    base = url.rstrip("/")
    if not common.server_running(url):
        return None
    identity = _identify_health(base, timeout)
    if identity is not None:
        return identity
    return _identify_gradio(base, timeout)


def _get_json(url: str, timeout: float) -> Optional[dict]:
    """GET URL and parse a JSON object, or None on any error."""
    try:
        with urllib.request.urlopen(url, timeout=timeout) as response:
            payload = json.loads(response.read().decode("utf-8"))
    except (OSError, ValueError):
        return None
    return payload if isinstance(payload, dict) else None


def _identify_health(base: str, timeout: float) -> Optional[str]:
    """Identify audio.cpp / faster / sglang-omni from their /health responses."""
    payload = _get_json(f"{base}/health", timeout)
    if payload is None:
        return None
    # faster's /health reports model load state under "model_loaded".
    if "model_loaded" in payload:
        return IDENTITY_FASTER
    # audio.cpp's /health reports {"status": "ok"}; confirm it also serves
    # the /v1/models catalog (id-bearing entries) to avoid mistaking some
    # other service that happens to return {"status": "ok"}.
    if payload.get("status") == "ok":
        models = _get_json(f"{base}/v1/models", timeout)
        entries = models.get("data") if models is not None else None
        if isinstance(entries, list) and entries \
                and any(isinstance(e, dict) and e.get("id") for e in entries):
            return IDENTITY_AUDIOCPP
    # sglang-omni's /health reports {"status": "healthy", "stages": [...]}
    # (200 when serving, 503 with "unhealthy" while booting). A 503 body
    # still parses as JSON here, so require the healthy word explicitly —
    # an "unhealthy" sgl-omni must not count as usable. The pipeline
    # "stages" list is confirmed as a secondary mark (present on every
    # sgl-omni 0.1.x server) before trusting the generic-sounding status.
    if payload.get("status") == "healthy" \
            and isinstance(payload.get("stages"), list):
        return IDENTITY_SGLOMNI
    return None


def _identify_gradio(base: str, timeout: float) -> Optional[str]:
    """Identify a qwen-tts Gradio demo from its ``/info`` named endpoints."""
    payload = _get_json(f"{base}/info", timeout)
    if payload is None:
        return None
    endpoints = payload.get("named_endpoints")
    if not isinstance(endpoints, dict):
        return None
    if any(name in endpoints for name in _QWEN_CUSTOM_ENDPOINTS):
        return IDENTITY_QWEN_CUSTOM
    if any(name in endpoints for name in _QWEN_CLONE_ENDPOINTS):
        return IDENTITY_QWEN_CLONE
    if any(name in endpoints for name in _QWEN_DESIGN_ENDPOINTS):
        return IDENTITY_QWEN_DESIGN
    return None


def _canonical_host(host: str) -> str:
    """Fold the loopback aliases so "localhost" and "127.0.0.1" compare equal."""
    return "127.0.0.1" if host in ("localhost", "::1", "[::1]") else host


def same_endpoint(url_a: str, url_b: str) -> bool:
    """True when URL_A and URL_B address the same host and port.

    Scheme and path are ignored (127.0.0.1:8080 and http://127.0.0.1:8080/
    are the same server), and the loopback names are folded together
    ("localhost:8080" equals "127.0.0.1:8080") — the config's remote-URL
    defaults point at the managed servers, so a user writing either form
    must not get their own server double-counted as "[remote]".
    Returns False when either URL is empty/unparsable.
    """
    if not url_a or not url_b:
        return False
    try:
        a = urllib.parse.urlsplit(url_a)
        b = urllib.parse.urlsplit(url_b)
    except ValueError:
        return False
    host_a = _canonical_host(a.hostname or "127.0.0.1")
    host_b = _canonical_host(b.hostname or "127.0.0.1")
    port_a = a.port or (443 if (a.scheme or "http") == "https" else 80)
    port_b = b.port or (443 if (b.scheme or "http") == "https" else 80)
    return host_a == host_b and port_a == port_b


def detect_remote_url(remote_url: str, local_url: str, identity: str,
                      name: str, managed: bool = False) -> tuple:
    """The shared backend "is something answering at the remote URL?" check.

    Returns ``(running, {name: url})``. An empty URL disables the check; a
    remote URL equal to the configured local endpoint is ignored while
    MANAGED (that server was started by this tool and is already reported
    as "[local]"); otherwise the URL must answer HTTP as IDENTITY to count.
    The faster and audio.cpp backends use it verbatim; qwen's is
    model-aware and keeps its own variant.
    """
    url = (remote_url or "").strip()
    if not url:
        return False, {}
    if managed and same_endpoint(url, local_url):
        return False, {}
    if identify_server(url) == identity:
        return True, {name: url}
    return False, {}


def health_payload(url: str, timeout: float = DEFAULT_TIMEOUT) -> Optional[dict]:
    """Return the server's ``/health`` JSON document, or None.

    A cheaper, raw check than ``identify_server``: used by the run view's
    background poll to tell "process alive" from "server answering" without
    probing every identity endpoint.
    """
    if not url:
        return None
    return _get_json(f"{url.rstrip('/')}/health", timeout)


def faster_model_loaded(url: str, timeout: float = DEFAULT_TIMEOUT) -> bool:
    """True when a faster server at URL reports its model loaded.

    faster's ``/health`` answers with a ``model_loaded`` flag only after the
    weights are resident, so this is the "truly ready" signal used while
    waiting for a started server to become usable.
    """
    payload = health_payload(url, timeout)
    return bool(payload and payload.get("model_loaded"))


def sglomni_served_model(url: str,
                         timeout: float = DEFAULT_TIMEOUT) -> Optional[str]:
    """The HuggingFace repo id a sglang-omni server at URL hosts, or None.

    ``GET /v1/models`` answers ``{"data": [{"id": <served repo>}, ...]}``
    with exactly one entry (one model per server process) — the same
    model-identity role qwen's probe plays for its three demos. Used to
    name the running model in statuses and to decide when a managed
    server must be restarted to host the model a run selected.
    """
    if not url:
        return None
    models = _get_json(f"{url.rstrip('/')}/v1/models", timeout)
    if models is None:
        return None
    entries = models.get("data")
    if isinstance(entries, list) and entries \
            and isinstance(entries[0], dict):
        model_id = entries[0].get("id")
        if isinstance(model_id, str) and model_id:
            return model_id
    return None


def sglomni_voice_names(url: str,
                        timeout: float = DEFAULT_TIMEOUT) -> Optional[List[str]]:
    """The uploaded voice names registered on a sglang-omni server, or None.

    ``GET /v1/audio/voices?names_only=true`` answers
    ``{"uploaded_voice_names": [...]}`` — the server-side voices a remote
    Convert form can offer in its voice picker (uploaded clips persist
    across server restarts). None when the URL does not answer.
    """
    if not url:
        return None
    payload = _get_json(f"{url.rstrip('/')}/v1/audio/voices?names_only=true",
                        timeout)
    if payload is None:
        return None
    names = payload.get("uploaded_voice_names")
    if isinstance(names, list):
        return [str(name) for name in names if isinstance(name, str) and name]
    return None