thoxmesh_node/config.py

"""Typed configuration for the node agent, resolved from the environment.

Every field is settable by environment variable so the same module runs
unmodified in a Colab cell, a Docker Space and a local shell. ``.env.example``
documents the same names with placeholder values; no real secret is ever
committed.
"""

from __future__ import annotations

import os
from dataclasses import dataclass, field
from typing import Any, Mapping

from .errors import ConfigError

#: Supabase project hosting the MeshStack v2 control plane.
DEFAULT_SUPABASE_PROJECT = "rtaktyshnzbjktdyeemg"

#: Transport used to publish this node into the mesh.
TRANSPORT_DEVICE_V2 = "device_v2"
TRANSPORT_CONTROLLER = "controller"
TRANSPORT_AUTO = "auto"
VALID_TRANSPORTS = (TRANSPORT_AUTO, TRANSPORT_DEVICE_V2, TRANSPORT_CONTROLLER)

#: Protocol values allowed by the v2 schema CHECK constraint.
VALID_PROTOCOLS = ("openai", "ollama", "llama_cpp", "litert", "custom")


def _env_int(source: Mapping[str, str], key: str, default: int, *, minimum: int = 1) -> int:
    raw = (source.get(key) or "").strip()
    if not raw:
        return default
    try:
        value = int(raw)
    except ValueError as exc:
        raise ConfigError(f"{key} must be an integer, got {raw!r}") from exc
    if value < minimum:
        raise ConfigError(f"{key} must be >= {minimum}, got {value}")
    return value


def _env_bool(source: Mapping[str, str], key: str, default: bool) -> bool:
    raw = (source.get(key) or "").strip().lower()
    if not raw:
        return default
    return raw in ("1", "true", "yes", "on")


def _env_flags(source: Mapping[str, str], key: str, default: tuple[str, ...]) -> dict[str, bool]:
    """Parse a comma-separated capability list into the jsonb flag map.

    ThoxRoute and ``resolve_route`` both test capabilities by truthy key lookup,
    so the list is expanded into ``{"chat": True, ...}`` rather than kept as an
    array.
    """
    raw = (source.get(key) or "").strip()
    names = [part.strip() for part in raw.split(",") if part.strip()] if raw else list(default)
    return {name: True for name in names}


@dataclass(frozen=True)
class NodeConfig:
    """Fully resolved node configuration."""

    # Identity of what we serve
    model_id: str
    alias: str
    protocol: str
    context_window: int
    capabilities: dict[str, bool]

    # Where we serve it
    host: str
    port: int
    public_url: str

    # How we join the mesh
    transport: str
    supabase_project: str
    controller_url: str
    node_kind: str
    ephemeral: bool
    heartbeat_seconds: int
    ttl_seconds: int
    register_retry_seconds: int

    # Local backend
    model_repo: str
    model_file: str
    n_threads: int
    n_gpu_layers: str
    max_tokens_cap: int
    request_timeout_seconds: int

    extra: dict[str, Any] = field(default_factory=dict)

    @property
    def device_function_url(self) -> str:
        """Absolute URL of the ``meshstack-device-v2`` edge function."""
        return f"https://{self.supabase_project}.supabase.co/functions/v1/meshstack-device-v2"

    def validate(self) -> None:
        """Fail fast on contradictions the agent cannot resolve at runtime."""
        if self.protocol not in VALID_PROTOCOLS:
            raise ConfigError(
                f"THOX_PROTOCOL must be one of {VALID_PROTOCOLS}, got {self.protocol!r}; "
                "the v2 schema rejects anything else"
            )
        if self.transport not in VALID_TRANSPORTS:
            raise ConfigError(f"THOX_MESH_TRANSPORT must be one of {VALID_TRANSPORTS}")
        if self.ttl_seconds <= self.heartbeat_seconds:
            raise ConfigError(
                f"TTL ({self.ttl_seconds}s) must exceed the heartbeat interval "
                f"({self.heartbeat_seconds}s), otherwise the lease expires between beats"
            )
        if self.transport == TRANSPORT_CONTROLLER and not self.controller_url:
            raise ConfigError("controller transport selected but THOX_CONTROLLER_URL is empty")
        if not self.model_id:
            raise ConfigError("THOX_MODEL_ID must not be empty")
        if not self.alias:
            raise ConfigError("THOX_ALIAS must not be empty")


def load_config(env: Mapping[str, str] | None = None) -> NodeConfig:
    """Resolve configuration from ``env`` (defaults to ``os.environ``).

    Defaults target the free-tier case deliberately: a CPU-sized GGUF, a 30s
    heartbeat and a 120s TTL. The TTL is four heartbeats so a node survives a
    couple of transient failures but a dead Colab session drops out of routing
    inside two minutes.
    """
    source = env if env is not None else os.environ

    heartbeat = _env_int(source, "THOX_HEARTBEAT_SECONDS", 30)
    ttl = _env_int(source, "THOX_TTL_SECONDS", 120)

    config = NodeConfig(
        # Defaults name a model that exists. An earlier default was
        # "thox-micro-125m", which is not published on the Hub at all, so any
        # caller relying on defaults was pointed at nothing.
        model_id=(source.get("THOX_MODEL_ID") or "thoxmini-3b").strip(),
        alias=(source.get("THOX_ALIAS") or source.get("THOX_MODEL_ID") or "thoxmini-3b").strip(),
        protocol=(source.get("THOX_PROTOCOL") or "openai").strip(),
        context_window=_env_int(source, "THOX_N_CTX", 2048),
        capabilities=_env_flags(source, "THOX_CAPABILITIES", ("chat", "completion")),
        host=(source.get("THOX_HOST") or "0.0.0.0").strip(),
        port=_env_int(source, "THOX_PORT", 7860),
        public_url=(source.get("THOX_PUBLIC_URL") or "").strip().rstrip("/"),
        transport=(source.get("THOX_MESH_TRANSPORT") or TRANSPORT_AUTO).strip().lower(),
        supabase_project=(source.get("THOX_SUPABASE_PROJECT") or DEFAULT_SUPABASE_PROJECT).strip(),
        controller_url=(source.get("THOX_CONTROLLER_URL") or "").strip().rstrip("/"),
        node_kind=(source.get("THOX_NODE_KIND") or "unknown").strip(),
        ephemeral=_env_bool(source, "THOX_EPHEMERAL", True),
        heartbeat_seconds=heartbeat,
        ttl_seconds=ttl,
        register_retry_seconds=_env_int(source, "THOX_REGISTER_RETRY_SECONDS", 10),
        model_repo=(source.get("THOX_MODEL_REPO") or "Thox-ai/ThoxMini-3B").strip(),
        model_file=(source.get("THOX_MODEL_FILE") or "thoxmini-3b-Q4_K_M.gguf").strip(),
        n_threads=_env_int(source, "THOX_N_THREADS", 0, minimum=0),
        # "auto" offloads everything when a GPU is detected and nothing when it
        # is not, so one value is correct on both Colab runtime types.
        n_gpu_layers=(source.get("THOX_N_GPU_LAYERS") or "auto").strip(),
        max_tokens_cap=_env_int(source, "THOX_MAX_TOKENS_CAP", 512),
        request_timeout_seconds=_env_int(source, "THOX_REQUEST_TIMEOUT_SECONDS", 120),
    )
    config.validate()
    return config


def detect_node_kind(env: Mapping[str, str] | None = None) -> str:
    """Best-effort guess at where we are running.

    Used for telemetry labelling and to pick sensible tunnel behaviour. A Space
    is already publicly addressable and needs no tunnel; Colab always does.
    """
    source = env if env is not None else os.environ
    if source.get("SPACE_ID") or source.get("SPACE_HOST"):
        return "hf-space"
    if source.get("COLAB_RELEASE_TAG") or source.get("COLAB_GPU") is not None:
        return "colab"
    if "google.colab" in (source.get("PYTHONPATH") or ""):
        return "colab"
    return "local"


def space_public_url(env: Mapping[str, str] | None = None) -> str:
    """Public origin of the current Space, or ``""`` when not on a Space.

    Hugging Face injects ``SPACE_HOST`` (``owner-name.hf.space``) but not
    ``SPACE_URL``; this repo previously registered ``http://0.0.0.0:7860``
    because it read the variable that does not exist. Returning ``""`` when the
    origin is unknown is deliberate — advertising an unreachable address is
    worse than not registering.
    """
    source = env if env is not None else os.environ
    host = (source.get("SPACE_HOST") or "").strip()
    if host:
        return host if "://" in host else f"https://{host}"
    space_id = (source.get("SPACE_ID") or "").strip()
    if space_id and "/" in space_id:
        owner, name = space_id.split("/", 1)
        slug = f"{owner}-{name}".replace("_", "-").replace(".", "-").lower()
        return f"https://{slug}.hf.space"
    return ""