src/adapters/__init__.py

"""The adapter allow-list and the factory that reaches it.

Enrolling a model means naming an adapter *family* and a Hub model id. The
family list is fixed in this file. Nothing here ever imports, downloads or
executes code chosen by a user: a user-supplied id is loaded only through an
already-vetted family's loader, and `trust_remote_code` is never set anywhere
in this package.

That is the whole security model for the "enroll a model" flow, and it is why
an unsupported family is a clear rejection rather than a best-effort attempt.
"""

from __future__ import annotations

import re

from .base import (HARDWARE_CPU, HARDWARE_GPU, OHLCV_COLUMNS,
                   OUTPUT_OHLCV_PATHS, OUTPUT_QUANTILE_LINE, AdapterError,
                   Capabilities, ContextError, ForecastAdapter,
                   ForecastResult, LookaheadError, ModelNotAllowed,
                   check_context, default_device, seed_everything)
from .baseline import BASELINE_MODELS, BaselineAdapter
from .chronos import CHRONOS_MODELS, ChronosAdapter
from .kronos import KRONOS_MODELS, KronosAdapter
from .timesfm import TIMESFM_MODELS, TimesFMAdapter

__all__ = [
    "AdapterError", "Capabilities", "ContextError", "ForecastAdapter",
    "ForecastResult", "LookaheadError", "ModelNotAllowed", "check_context",
    "default_device", "seed_everything", "get_adapter", "validate_model_id",
    "ALLOWED_ADAPTER_FAMILIES", "KNOWN_MODELS", "family_for",
    "OHLCV_COLUMNS", "OUTPUT_OHLCV_PATHS", "OUTPUT_QUANTILE_LINE",
    "HARDWARE_CPU", "HARDWARE_GPU",
]

# The allow-list. Adding an entry here is the deliberate act of accepting a new
# code path; there is no dynamic discovery.
ADAPTERS: dict[str, type[ForecastAdapter]] = {
    "kronos": KronosAdapter,
    "chronos": ChronosAdapter,
    "timesfm": TimesFMAdapter,
    "baseline": BaselineAdapter,
}

ALLOWED_ADAPTER_FAMILIES = tuple(ADAPTERS)

# Every model this app ships knowing about, and the family that loads it. A
# user may enroll an id that is not on this list, but only under a family that
# is -- see `runtime.enroll`.
KNOWN_MODELS: dict[str, str] = {
    **{mid: "kronos" for mid in KRONOS_MODELS},
    **{mid: "chronos" for mid in CHRONOS_MODELS},
    **{mid: "timesfm" for mid in TIMESFM_MODELS},
    **{mid: "baseline" for mid in BASELINE_MODELS},
}

# `namespace/name`, the Hub's own shape. Anything else is not a model id and is
# rejected before it reaches a loader.
_MODEL_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,95}/[A-Za-z0-9][A-Za-z0-9._-]{0,95}$")


def validate_model_id(model_id: str) -> str:
    """Accept a Hub model id, or raise.

    Model ids arrive from the enrollment form and are user input. Validating
    the shape here keeps path traversal and injection out of every downstream
    consumer -- the store writes files under a slug derived from this.
    """
    if not isinstance(model_id, str):
        raise AdapterError("model id must be a string")
    model_id = model_id.strip()
    if not _MODEL_ID.match(model_id):
        raise AdapterError(
            f"{model_id!r} is not a valid Hugging Face model id "
            f"(expected 'namespace/name')"
        )
    if ".." in model_id:
        raise AdapterError("model id may not contain '..'")
    return model_id


def family_for(model_id: str) -> str | None:
    """The family that ships support for this id, if any."""
    return KNOWN_MODELS.get(model_id)


def get_adapter(family: str, model_id: str, revision: str | None = None,
                **kwargs) -> ForecastAdapter:
    """Build an adapter, or refuse.

    Refusing is the important half: an unknown family is `ModelNotAllowed`
    with the supported list in the message, never a silent fallback to a
    generic loader.
    """
    fam = (family or "").strip().lower()
    if fam not in ADAPTERS:
        raise ModelNotAllowed(
            f"adapter family {family!r} is not supported. "
            f"Supported families: {', '.join(ALLOWED_ADAPTER_FAMILIES)}"
        )
    return ADAPTERS[fam](validate_model_id(model_id), revision=revision, **kwargs)


def model_slug(model_id: str) -> str:
    """Store path segment for a model id.

    `amazon/chronos-bolt-small` -> `chronos-bolt-small`. The namespace is
    dropped because it is redundant with the registry entry, and the result is
    checked against the same character class the id was, so it can never climb
    out of the store's directory tree.
    """
    validate_model_id(model_id)
    slug = model_id.split("/", 1)[1].lower()
    slug = re.sub(r"[^a-z0-9._-]", "-", slug)
    if not slug or slug in (".", ".."):
        raise AdapterError(f"cannot derive a slug from {model_id!r}")
    return slug