tsagentkit 1.0.2__py3-none-any.whl

tsagentkit/contracts/task_spec.py

@@ -0,0 +1,300 @@
+"""Pydantic specs for tsagentkit contracts and configuration.
+
+These models are the JSON-serializable configuration and artifact contracts
+used by agents and orchestration layers. They mirror docs/PRD.md Appendix B.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+# ---------------------------
+# Common
+# ---------------------------
+
+class BaseSpec(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+
+CovariateRole = Literal["static", "past", "future_known"]
+AggregationMode = Literal["reject", "sum", "mean", "median", "last"]
+MissingPolicy = Literal["error", "ffill", "bfill", "zero", "mean"]
+IntervalMode = Literal["level", "quantiles"]
+AnomalyMethod = Literal["interval_breach", "conformal", "mad_residual"]
+SeasonalityMethod = Literal["acf", "stl", "periodogram"]
+CovariatePolicy = Literal["ignore", "known", "observed", "auto", "spec"]
+
+
+# ---------------------------
+# Data contracts (column-level)
+# ---------------------------
+
+class PanelContract(BaseSpec):
+    unique_id_col: str = "unique_id"
+    ds_col: str = "ds"
+    y_col: str = "y"
+    aggregation: AggregationMode = "reject"
+
+
+class ForecastContract(BaseSpec):
+    long_format: bool = True
+    model_col: str = "model"
+    yhat_col: str = "yhat"
+    cutoff_col: str = "cutoff"  # required for CV output
+    interval_mode: IntervalMode = "level"
+    levels: list[int] = Field(default_factory=lambda: [80, 95])
+    quantiles: list[float] = Field(default_factory=lambda: [0.1, 0.5, 0.9])
+
+
+class CovariateSpec(BaseSpec):
+    # Explicit typing strongly preferred for agent safety.
+    roles: dict[str, CovariateRole] = Field(default_factory=dict)
+    missing_policy: MissingPolicy = "error"
+
+
+# ---------------------------
+# Task / execution specs
+# ---------------------------
+
+class BacktestSpec(BaseSpec):
+    h: int | None = Field(None, gt=0)
+    n_windows: int = Field(5, gt=0)
+    step: int = Field(1, gt=0)
+    min_train_size: int = Field(56, gt=1)
+    regularize_grid: bool = True
+
+
+class TaskSpec(BaseSpec):
+    # Forecast horizon
+    h: int = Field(..., gt=0)
+
+    # Frequency handling
+    freq: str | None = None
+    infer_freq: bool = True
+
+    # Contracts
+    panel_contract: PanelContract = Field(default_factory=PanelContract)
+    forecast_contract: ForecastContract = Field(default_factory=ForecastContract)
+
+    # Covariates
+    covariates: CovariateSpec | None = None
+    covariate_policy: CovariatePolicy = "auto"
+
+    # Backtest defaults (can be overridden by the caller)
+    backtest: BacktestSpec = Field(default_factory=BacktestSpec)
+
+    @model_validator(mode="before")
+    @classmethod
+    def _normalize_inputs(cls, data: Any) -> Any:
+        if not isinstance(data, dict):
+            return data
+
+        payload = dict(data)
+
+        # Backward-compat aliases
+        if "horizon" in payload and "h" not in payload:
+            payload["h"] = payload.pop("horizon")
+        if "rolling_step" in payload:
+            backtest = payload.get("backtest", {})
+            if isinstance(backtest, BacktestSpec):
+                backtest = backtest.model_dump()
+            if isinstance(backtest, dict):
+                backtest = dict(backtest)
+                if "step" not in backtest:
+                    backtest["step"] = payload.pop("rolling_step")
+                payload["backtest"] = backtest
+
+        # Legacy quantiles/levels mapping to forecast_contract
+        if "quantiles" in payload or "levels" in payload:
+            fc = payload.get("forecast_contract", {})
+            if isinstance(fc, ForecastContract):
+                fc = fc.model_dump()
+            if isinstance(fc, dict):
+                fc = dict(fc)
+                if "quantiles" in payload:
+                    fc["quantiles"] = payload.pop("quantiles")
+                if "levels" in payload:
+                    fc["levels"] = payload.pop("levels")
+                payload["forecast_contract"] = fc
+
+        return payload
+
+    @model_validator(mode="after")
+    def _apply_backtest_defaults(self) -> TaskSpec:
+        if self.backtest.h is None:
+            object.__setattr__(self, "backtest", self.backtest.model_copy(update={"h": self.h}))
+        return self
+
+    @property
+    def horizon(self) -> int:
+        return self.h
+
+    @property
+    def quantiles(self) -> list[float]:
+        return self.forecast_contract.quantiles
+
+    @property
+    def levels(self) -> list[int]:
+        return self.forecast_contract.levels
+
+    @property
+    def season_length(self) -> int | None:
+        return self._infer_season_length(self.freq)
+
+    @staticmethod
+    def _infer_season_length(freq: str | None) -> int | None:
+        if not freq:
+            return None
+        freq_map: dict[str, int] = {
+            "D": 7,
+            "B": 5,
+            "H": 24,
+            "T": 60,
+            "min": 60,
+            "M": 12,
+            "MS": 12,
+            "Q": 4,
+            "QS": 4,
+            "W": 52,
+        }
+        base_freq = freq.lstrip("0123456789")
+        return freq_map.get(base_freq)
+
+    def model_hash(self) -> str:
+        import hashlib
+        import json
+
+        data = self.model_dump(exclude_none=True)
+        json_str = json.dumps(data, sort_keys=True, separators=(",", ":"))
+        return hashlib.sha256(json_str.encode()).hexdigest()[:16]
+
+
+# ---------------------------
+# Router / planning
+# ---------------------------
+
+class RouterThresholds(BaseSpec):
+    min_train_size: int = Field(56, gt=1)
+    max_missing_ratio: float = Field(0.15, ge=0.0, le=1.0)
+
+    # Intermittency classification (heuristic, deterministic)
+    max_intermittency_adi: float = Field(1.32, gt=0.0)
+    max_intermittency_cv2: float = Field(0.49, ge=0.0)
+
+    # Seasonality
+    seasonality_method: SeasonalityMethod = "acf"
+    min_seasonality_conf: float = Field(0.70, ge=0.0, le=1.0)
+
+    # Practical routing guardrails
+    max_series_count_for_tsfm: int = Field(20000, gt=0)
+    max_points_per_series_for_tsfm: int = Field(5000, gt=0)
+
+
+class PlanSpec(BaseSpec):
+    plan_name: str
+    candidate_models: list[str] = Field(..., min_length=1)
+
+    # Covariate usage rules
+    use_static: bool = True
+    use_past: bool = True
+    use_future_known: bool = True
+
+    # Training policy
+    min_train_size: int = Field(56, gt=1)
+    max_train_size: int | None = None  # if set, truncate oldest points deterministically
+
+    # Output policy
+    interval_mode: IntervalMode = "level"
+    levels: list[int] = Field(default_factory=lambda: [80, 95])
+    quantiles: list[float] = Field(default_factory=lambda: [0.1, 0.5, 0.9])
+
+    # Fallback policy
+    allow_drop_covariates: bool = True
+    allow_baseline: bool = True
+
+
+class RouteDecision(BaseSpec):
+    # Series statistics used in routing (computed deterministically)
+    stats: dict[str, Any] = Field(default_factory=dict)
+
+    # Bucket tags
+    buckets: list[str] = Field(default_factory=list)
+
+    # Which plan template was selected
+    selected_plan: PlanSpec
+
+    # Human-readable deterministic reasons (safe for logs)
+    reasons: list[str] = Field(default_factory=list)
+
+
+class RouterConfig(BaseSpec):
+    thresholds: RouterThresholds = Field(default_factory=RouterThresholds)
+
+    # Mapping bucket -> plan template name, resolved by registry
+    bucket_to_plan: dict[str, str] = Field(default_factory=dict)
+
+    # Default plan when no bucket matches
+    default_plan: str = "default"
+
+
+# ---------------------------
+# Calibration + anomaly
+# ---------------------------
+
+class CalibratorSpec(BaseSpec):
+    method: Literal["none", "conformal"] = "conformal"
+    level: int = Field(99, ge=50, le=99)
+    by: Literal["unique_id", "global"] = "unique_id"
+
+
+class AnomalySpec(BaseSpec):
+    method: AnomalyMethod = "conformal"
+    level: int = Field(99, ge=50, le=99)
+    score: Literal["margin", "normalized_margin", "zscore"] = "normalized_margin"
+
+
+# ---------------------------
+# Provenance artifacts (config-level, serializable)
+# ---------------------------
+
+class RunArtifactSpec(BaseSpec):
+    run_id: str
+    created_at: datetime
+
+    task_spec: TaskSpec
+    router_config: RouterConfig | None = None
+    route_decision: RouteDecision | None = None
+
+    # Identifiers / hashes for reproducibility (implementation-defined)
+    data_signature: str | None = None
+    code_signature: str | None = None
+
+    # Output references (implementation-defined; typically file paths or object-store keys)
+    outputs: dict[str, str] = Field(default_factory=dict)
+
+
+__all__ = [
+    "AggregationMode",
+    "AnomalyMethod",
+    "AnomalySpec",
+    "BacktestSpec",
+    "BaseSpec",
+    "CalibratorSpec",
+    "CovariatePolicy",
+    "CovariateRole",
+    "CovariateSpec",
+    "ForecastContract",
+    "IntervalMode",
+    "MissingPolicy",
+    "PanelContract",
+    "PlanSpec",
+    "RouteDecision",
+    "RouterConfig",
+    "RouterThresholds",
+    "RunArtifactSpec",
+    "SeasonalityMethod",
+    "TaskSpec",
+]

tsagentkit/covariates/__init__.py

@@ -0,0 +1,340 @@
+"""Covariate typing, alignment, and guardrails."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import pandas as pd
+
+from tsagentkit.contracts import (
+    CovariateSpec,
+    ECovariateIncompleteKnown,
+    ECovariateLeakage,
+    ECovariateStaticInvalid,
+    ETaskSpecInvalid,
+    TaskSpec,
+)
+from tsagentkit.time import make_future_index
+
+
+@dataclass(frozen=True)
+class CovariateBundle:
+    static_x: pd.DataFrame | None = None
+    past_x: pd.DataFrame | None = None
+    future_x: pd.DataFrame | None = None
+
+
+@dataclass(frozen=True)
+class AlignedDataset:
+    panel: pd.DataFrame
+    static_x: pd.DataFrame | None
+    past_x: pd.DataFrame | None
+    future_x: pd.DataFrame | None
+    covariate_spec: CovariateSpec | None
+    future_index: pd.DataFrame | None
+
+
+def align_covariates(
+    panel: pd.DataFrame,
+    task_spec: TaskSpec,
+    covariates: CovariateBundle | None = None,
+) -> AlignedDataset:
+    """Align covariates and enforce coverage/leakage guardrails."""
+    contract = task_spec.panel_contract
+    uid_col = contract.unique_id_col
+    ds_col = contract.ds_col
+    y_col = contract.y_col
+
+    panel = panel.copy()
+    future_index = make_future_index(panel, task_spec.h, task_spec.freq, uid_col, ds_col, y_col)
+
+    if covariates is not None:
+        static_x = _validate_static_covariates(covariates.static_x, uid_col)
+        past_x = _validate_past_covariates(covariates.past_x, future_index, uid_col, ds_col)
+        future_x = _validate_future_covariates(
+            covariates.future_x, future_index, uid_col, ds_col
+        )
+        return AlignedDataset(
+            panel=_panel_base(panel, uid_col, ds_col, y_col),
+            static_x=static_x,
+            past_x=past_x,
+            future_x=future_x,
+            covariate_spec=task_spec.covariates,
+            future_index=future_index,
+        )
+
+    covariate_cols = [c for c in panel.columns if c not in {uid_col, ds_col, y_col}]
+    if not covariate_cols or task_spec.covariate_policy == "ignore":
+        return AlignedDataset(
+            panel=_panel_base(panel, uid_col, ds_col, y_col),
+            static_x=None,
+            past_x=None,
+            future_x=None,
+            covariate_spec=task_spec.covariates,
+            future_index=future_index,
+        )
+
+    static_cols: list[str] = []
+    past_cols: list[str] = []
+    future_cols: list[str] = []
+
+    policy = task_spec.covariate_policy
+    spec = task_spec.covariates
+
+    if policy == "spec" and spec is None and covariate_cols:
+        raise ETaskSpecInvalid(
+            "covariate_policy='spec' requires task_spec.covariates.",
+            context={"missing": "covariates"},
+        )
+
+    if policy == "spec" and spec is not None:
+        _validate_spec_roles(spec, covariate_cols)
+        for col, role in spec.roles.items():
+            if role == "static":
+                static_cols.append(col)
+            elif role == "past":
+                past_cols.append(col)
+            elif role == "future_known":
+                future_cols.append(col)
+    elif policy == "known":
+        future_cols = covariate_cols
+    elif policy == "observed":
+        past_cols = covariate_cols
+    elif policy == "auto":
+        for col in covariate_cols:
+            if _has_future_values(panel, future_index, col, uid_col, ds_col):
+                # Candidate future-known: require full coverage
+                future_cols.append(col)
+                _enforce_future_coverage(panel, future_index, col, uid_col, ds_col)
+            else:
+                past_cols.append(col)
+    else:
+        past_cols = covariate_cols
+
+    for col in future_cols:
+        _enforce_future_coverage(panel, future_index, col, uid_col, ds_col)
+
+    static_x = _extract_static(panel, uid_col, static_cols)
+    past_x = _extract_time_covariates(panel, uid_col, ds_col, past_cols)
+    future_x = _extract_future_covariates(panel, future_index, uid_col, ds_col, future_cols)
+
+    _enforce_past_leakage(panel, future_index, uid_col, ds_col, past_cols)
+
+    return AlignedDataset(
+        panel=_panel_base(panel, uid_col, ds_col, y_col),
+        static_x=static_x,
+        past_x=past_x,
+        future_x=future_x,
+        covariate_spec=spec,
+        future_index=future_index,
+    )
+
+
+def _panel_base(panel: pd.DataFrame, uid_col: str, ds_col: str, y_col: str) -> pd.DataFrame:
+    cols = [c for c in [uid_col, ds_col, y_col] if c in panel.columns]
+    return panel[cols].copy()
+
+
+def _validate_spec_roles(spec: CovariateSpec, covariate_cols: list[str]) -> None:
+    if not spec.roles:
+        if covariate_cols:
+            raise ETaskSpecInvalid(
+                "covariate_policy='spec' requires explicit roles for all covariate columns.",
+                context={"missing_roles_for": sorted(covariate_cols)},
+            )
+        return
+
+    missing_in_panel = [col for col in spec.roles if col not in covariate_cols]
+    if missing_in_panel:
+        raise ETaskSpecInvalid(
+            "CovariateSpec roles include columns not present in panel data.",
+            context={"missing_in_panel": sorted(missing_in_panel)},
+        )
+
+    extra_in_panel = [col for col in covariate_cols if col not in spec.roles]
+    if extra_in_panel:
+        raise ETaskSpecInvalid(
+            "covariate_policy='spec' requires roles for all panel covariates.",
+            context={"missing_roles_for": sorted(extra_in_panel)},
+        )
+
+
+def _has_future_values(
+    panel: pd.DataFrame,
+    future_index: pd.DataFrame,
+    col: str,
+    uid_col: str,
+    ds_col: str,
+) -> bool:
+    merged = future_index.merge(
+        panel[[uid_col, ds_col, col]],
+        on=[uid_col, ds_col],
+        how="left",
+    )
+    return merged[col].notna().any()
+
+
+def _enforce_future_coverage(
+    panel: pd.DataFrame,
+    future_index: pd.DataFrame,
+    col: str,
+    uid_col: str,
+    ds_col: str,
+) -> None:
+    merged = future_index.merge(
+        panel[[uid_col, ds_col, col]],
+        on=[uid_col, ds_col],
+        how="left",
+    )
+    if merged[col].isna().any():
+        missing = int(merged[col].isna().sum())
+        raise ECovariateIncompleteKnown(
+            f"Future-known covariate '{col}' missing {missing} values in horizon.",
+            context={"covariate": col, "missing": missing},
+        )
+
+
+def _enforce_past_leakage(
+    panel: pd.DataFrame,
+    future_index: pd.DataFrame,
+    uid_col: str,
+    ds_col: str,
+    past_cols: list[str],
+) -> None:
+    if not past_cols:
+        return
+
+    merged = future_index.merge(
+        panel[[uid_col, ds_col] + past_cols],
+        on=[uid_col, ds_col],
+        how="left",
+    )
+    for col in past_cols:
+        if merged[col].notna().any():
+            count = int(merged[col].notna().sum())
+            raise ECovariateLeakage(
+                f"Past covariate '{col}' has {count} values in forecast horizon.",
+                context={"covariate": col, "future_values_count": count},
+            )
+
+
+def _extract_static(panel: pd.DataFrame, uid_col: str, cols: list[str]) -> pd.DataFrame | None:
+    if not cols:
+        return None
+    df = panel[[uid_col] + cols].dropna(subset=[uid_col]).copy()
+    # Validate constant per unique_id
+    for col in cols:
+        counts = df.groupby(uid_col)[col].nunique(dropna=True)
+        if (counts > 1).any():
+            bad = counts[counts > 1].index.tolist()[:5]
+            raise ECovariateStaticInvalid(
+                f"Static covariate '{col}' varies within series.",
+                context={"covariate": col, "unique_id_examples": bad},
+            )
+    return df.groupby(uid_col, as_index=False).first()
+
+
+def _extract_time_covariates(
+    panel: pd.DataFrame,
+    uid_col: str,
+    ds_col: str,
+    cols: list[str],
+) -> pd.DataFrame | None:
+    if not cols:
+        return None
+    return panel[[uid_col, ds_col] + cols].copy()
+
+
+def _extract_future_covariates(
+    panel: pd.DataFrame,
+    future_index: pd.DataFrame,
+    uid_col: str,
+    ds_col: str,
+    cols: list[str],
+) -> pd.DataFrame | None:
+    if not cols:
+        return None
+    merged = future_index.merge(
+        panel[[uid_col, ds_col] + cols],
+        on=[uid_col, ds_col],
+        how="left",
+    )
+    for col in cols:
+        if merged[col].isna().any():
+            missing = int(merged[col].isna().sum())
+            raise ECovariateIncompleteKnown(
+                f"Future-known covariate '{col}' missing {missing} values in horizon.",
+                context={"covariate": col, "missing": missing},
+            )
+    return merged
+
+
+def _validate_static_covariates(
+    static_x: pd.DataFrame | None,
+    uid_col: str,
+) -> pd.DataFrame | None:
+    if static_x is None or static_x.empty:
+        return None
+    counts = static_x.groupby(uid_col).size()
+    if (counts != 1).any():
+        bad = counts[counts != 1].index.tolist()[:5]
+        raise ECovariateStaticInvalid(
+            "Static covariates must have exactly one row per unique_id.",
+            context={"unique_id_examples": bad},
+        )
+    return static_x.copy()
+
+
+def _validate_past_covariates(
+    past_x: pd.DataFrame | None,
+    future_index: pd.DataFrame,
+    uid_col: str,
+    ds_col: str,
+) -> pd.DataFrame | None:
+    if past_x is None or past_x.empty:
+        return None
+    merged = future_index.merge(
+        past_x,
+        on=[uid_col, ds_col],
+        how="left",
+    )
+    covariate_cols = [c for c in past_x.columns if c not in {uid_col, ds_col}]
+    for col in covariate_cols:
+        if merged[col].notna().any():
+            count = int(merged[col].notna().sum())
+            raise ECovariateLeakage(
+                f"Past covariate '{col}' has {count} values in forecast horizon.",
+                context={"covariate": col, "future_values_count": count},
+            )
+    return past_x.copy()
+
+
+def _validate_future_covariates(
+    future_x: pd.DataFrame | None,
+    future_index: pd.DataFrame,
+    uid_col: str,
+    ds_col: str,
+) -> pd.DataFrame | None:
+    if future_x is None or future_x.empty:
+        return None
+    merged = future_index.merge(
+        future_x,
+        on=[uid_col, ds_col],
+        how="left",
+    )
+    covariate_cols = [c for c in future_x.columns if c not in {uid_col, ds_col}]
+    for col in covariate_cols:
+        if merged[col].isna().any():
+            missing = int(merged[col].isna().sum())
+            raise ECovariateIncompleteKnown(
+                f"Future-known covariate '{col}' missing {missing} values in horizon.",
+                context={"covariate": col, "missing": missing},
+            )
+    return future_x.copy()
+
+
+__all__ = [
+    "AlignedDataset",
+    "CovariateBundle",
+    "align_covariates",
+]