runtime.py

from __future__ import annotations

import os
import tempfile
import time
from collections.abc import Iterator
from typing import Any, Callable

DEFAULT_PROVIDER_FALLBACK = (
    "auto",
    "hf-inference",
    "novita",
    "together",
    "fireworks-ai",
)
PRO_PROVIDER_FALLBACK = (
    "hf-inference",
    "auto",
    "novita",
    "together",
    "fireworks-ai",
)
MAX_INFERENCE_ATTEMPTS = 3
MAX_PRO_INFERENCE_ATTEMPTS = 2
MAX_PRO_PROVIDER_HOPS = 3
RETRY_BACKOFF_SECONDS = (0.75, 1.5, 3.0)
RETRYABLE_ERROR_MARKERS = (
    "429",
    "rate limit",
    "502",
    "503",
    "504",
    "timeout",
    "timed out",
    "temporarily",
    "overloaded",
    "unavailable",
    "connection reset",
    "connection aborted",
)

_pro_status_by_token: dict[str, bool] = {}


def detect_hosting() -> str:
    if os.getenv("SPACE_ID"):
        return "huggingface"
    if os.getenv("CODESPACES") == "true":
        return "github_codespaces"
    if os.getenv("GITHUB_ACTIONS") == "true":
        return "github_actions"
    if os.getenv("PORT") or os.getenv("WEBSITES_PORT"):
        return "container"
    return "local"


def resolve_hf_token() -> str | None:
    try:
        from huggingface_hub import get_token

        token = get_token()
        if token:
            return token
    except Exception:
        pass
    return os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")


def reset_hf_pro_cache() -> None:
    _pro_status_by_token.clear()


def resolve_hf_pro_status(*, token: str | None = None) -> bool:
    token = token or resolve_hf_token()
    if not token:
        return False
    if token in _pro_status_by_token:
        return _pro_status_by_token[token]
    try:
        from huggingface_hub import whoami

        info = whoami(token=token, cache=True)
        is_pro = bool(info.get("isPro"))
    except Exception:
        return False
    _pro_status_by_token[token] = is_pro
    return is_pro


def resolve_server_port() -> int | None:
    for key in ("PORT", "WEBSITES_PORT", "GRADIO_SERVER_PORT", "SPACE_PORT"):
        raw = os.getenv(key)
        if not raw:
            continue
        try:
            return int(raw)
        except ValueError:
            continue
    return None


def gradio_launch_kwargs(*, css: str | None = None) -> dict[str, Any]:
    hosting = detect_hosting()
    kwargs: dict[str, Any] = {
        "css": css,
        "ssr_mode": False,
        "show_error": True,
        "allowed_paths": [tempfile.gettempdir()],
    }

    if hosting in {"huggingface", "github_codespaces", "container"}:
        kwargs["server_name"] = "0.0.0.0"

    port = resolve_server_port()
    if port is not None:
        kwargs["server_port"] = port

    root_path = os.getenv("GRADIO_ROOT_PATH")
    if root_path:
        kwargs["root_path"] = root_path

    return kwargs


def inference_provider_chain(preferred: str | None, *, is_pro: bool = False) -> list[str]:
    preferred_value = (preferred or "auto").strip() or "auto"
    fallback = PRO_PROVIDER_FALLBACK if is_pro else DEFAULT_PROVIDER_FALLBACK

    if is_pro and preferred_value == "auto":
        chain = list(PRO_PROVIDER_FALLBACK)
    else:
        chain = [preferred_value]
        for provider in fallback:
            if provider not in chain:
                chain.append(provider)

    if is_pro:
        return chain[:MAX_PRO_PROVIDER_HOPS]
    return chain


def inference_attempts_per_provider(*, is_pro: bool = False) -> int:
    return MAX_PRO_INFERENCE_ATTEMPTS if is_pro else MAX_INFERENCE_ATTEMPTS


def inference_routing_hint() -> str:
    token = resolve_hf_token()
    if not token:
        return ""
    if resolve_hf_pro_status(token=token):
        return (
            "<span style='color:#69ff53;font-size:0.92em'>"
            "HF Pro routing active — hf-inference preferred.</span>"
        )
    return (
        "<span style='color:#9fdfff;font-size:0.92em'>"
        "Standard Hugging Face inference routing.</span>"
    )


def is_retryable_inference_error(exc: Exception) -> bool:
    message = str(exc).lower()
    return any(marker in message for marker in RETRYABLE_ERROR_MARKERS)


def stream_chat_completion(
    client_factory: Callable[[str], Any],
    *,
    providers: list[str],
    model: str,
    messages: list[dict[str, str]],
    max_tokens: int,
    temperature: float,
    max_attempts_per_provider: int = MAX_INFERENCE_ATTEMPTS,
) -> Iterator[str]:
    last_error: Exception | None = None

    for provider in providers:
        for attempt in range(max_attempts_per_provider):
            try:
                client = client_factory(provider)
                stream = client.chat_completion(
                    model=model,
                    messages=messages,
                    max_tokens=max_tokens,
                    temperature=temperature,
                    stream=True,
                )
                for chunk in stream:
                    delta = chunk.choices[0].delta.content or ""
                    if delta:
                        yield delta
                return
            except Exception as exc:
                last_error = exc
                if (
                    attempt + 1 >= max_attempts_per_provider
                    or not is_retryable_inference_error(exc)
                ):
                    break
                time.sleep(RETRY_BACKOFF_SECONDS[min(attempt, len(RETRY_BACKOFF_SECONDS) - 1)])

    if last_error is not None:
        raise last_error
    raise RuntimeError("Inference failed without a provider response.")


def format_inference_failure(
    exc: Exception,
    providers: list[str],
    *,
    is_pro: bool = False,
) -> str:
    hosting = detect_hosting()
    routing = "HF Pro routing (hf-inference preferred)" if is_pro else "standard routing"
    lines = [
        f"Inference failed after trying providers: {', '.join(providers)} ({routing}).",
        f"Last error: {exc}",
    ]
    if hosting == "huggingface":
        lines.append(
            "On Hugging Face Spaces, add an HF token with Inference Providers permission "
            "as the HF_TOKEN secret."
        )
    elif hosting.startswith("github"):
        lines.append(
            "On GitHub-hosted runtimes, set HF_TOKEN (or HUGGING_FACE_HUB_TOKEN) in repository "
            "or Codespace secrets."
        )
    else:
        lines.append(
            "Set HF_TOKEN locally or choose a different inference provider/model."
        )
    if is_pro:
        lines.append(
            "HF Pro is active; if errors persist, try a specific provider or model."
        )
    lines.append(
        "Transient provider outages are retried automatically; persistent errors need "
        "a new token, model, or provider."
    )
    return "\n\n".join(lines)