backend/forecasting/config.py

from __future__ import annotations

from typing import Dict, Mapping, Tuple


FORECAST_MODEL_ORDER: Tuple[str, ...] = ("kronos", "timesfm", "chronos")

FORECAST_MODEL_LABELS: Dict[str, str] = {
    "kronos": "Kronos",
    "timesfm": "TimesFM",
    "chronos": "Chronos",
}

DEFAULT_FORECAST_MODEL_SELECTION: Dict[str, bool] = {
    model_key: True
    for model_key in FORECAST_MODEL_ORDER
}


def normalize_forecast_model_selection(
    selection: Mapping[str, bool] | None = None,
) -> Dict[str, bool]:
    normalized = {
        model_key: bool(
            DEFAULT_FORECAST_MODEL_SELECTION[model_key]
            if selection is None
            else selection.get(model_key, DEFAULT_FORECAST_MODEL_SELECTION[model_key])
        )
        for model_key in FORECAST_MODEL_ORDER
    }
    if not any(normalized.values()):
        normalized[FORECAST_MODEL_ORDER[0]] = True
    return normalized


def build_query_model_selection(
    *,
    use_kronos: bool,
    use_timesfm: bool,
    use_chronos: bool,
) -> Dict[str, bool]:
    return normalize_forecast_model_selection(
        {
            "kronos": use_kronos,
            "timesfm": use_timesfm,
            "chronos": use_chronos,
        }
    )


def forecast_model_signature(model_selection: Mapping[str, bool]) -> str:
    normalized = normalize_forecast_model_selection(model_selection)
    return "".join(
        f"{model_key[0]}{int(normalized[model_key])}"
        for model_key in FORECAST_MODEL_ORDER
    )


def enabled_model_keys(model_selection: Mapping[str, bool]) -> list[str]:
    normalized = normalize_forecast_model_selection(model_selection)
    return [
        model_key
        for model_key in FORECAST_MODEL_ORDER
        if normalized[model_key]
    ]