schema.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
"""Recipe loop objects: DataRevision, Recipe, RunRecord, Checkpoint, EvalGate."""
 
from __future__ import annotations
 
from dataclasses import asdict, dataclass, field
from typing import Any
 
 
FAMILIES = ("tabular", "llm")
TABULAR_METHODS = ("linear", "boosting")
LLM_METHODS = ("lora",)
TABULAR_TASKS = ("classification", "regression")
 
 
class RecipeError(ValueError):
    """Invalid recipe or train inputs."""
 
 
@dataclass
class DataRevision:
    id: str
    path: str
    sha256: str
    format: str
    n_rows: int | None = None
    columns: list[str] | None = None
    snapshot_dir: str | None = None
 
    def to_dict(self) -> dict[str, Any]:
        return asdict(self)
 
    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> DataRevision:
        return cls(
            id=str(data["id"]),
            path=str(data["path"]),
            sha256=str(data["sha256"]),
            format=str(data.get("format") or "file"),
            n_rows=data.get("n_rows"),
            columns=list(data["columns"]) if data.get("columns") else None,
            snapshot_dir=data.get("snapshot_dir"),
        )
 
 
@dataclass
class EvalGate:
    passed: bool
    metric: str
    score: float | None
    min_score: float | None = None
    skipped: bool = False
    details: dict[str, Any] = field(default_factory=dict)
 
    def to_dict(self) -> dict[str, Any]:
        return asdict(self)
 
    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> EvalGate:
        return cls(
            passed=bool(data.get("passed")),
            metric=str(data.get("metric") or ""),
            score=data.get("score"),
            min_score=data.get("min_score"),
            skipped=bool(data.get("skipped")),
            details=dict(data.get("details") or {}),
        )
 
 
@dataclass
class Checkpoint:
    path: str
    kind: str
    step: int | None = None
    extra: dict[str, Any] = field(default_factory=dict)
 
    def to_dict(self) -> dict[str, Any]:
        return asdict(self)
 
    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> Checkpoint:
        return cls(
            path=str(data["path"]),
            kind=str(data.get("kind") or "file"),
            step=data.get("step"),
            extra=dict(data.get("extra") or {}),
        )
 
 
@dataclass
class Recipe:
    name: str
    family: str
    data: dict[str, Any]
    train: dict[str, Any]
    eval: dict[str, Any] = field(default_factory=dict)
    source: str | None = None
    extra: dict[str, Any] = field(default_factory=dict)
 
    def to_dict(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "family": self.family,
            "data": self.data,
            "train": self.train,
            "eval": self.eval,
            "source": self.source,
            "extra": self.extra,
        }
 
    @classmethod
    def from_dict(cls, data: dict[str, Any], *, source: str | None = None) -> Recipe:
        if not isinstance(data, dict):
            raise RecipeError("Recipe must be a mapping (YAML/JSON object).")
        name = str(data.get("name") or "").strip()
        if not name:
            raise RecipeError("Recipe needs a name.")
        family = str(data.get("family") or "").strip().lower()
        if family not in FAMILIES:
            raise RecipeError(f"family must be one of {list(FAMILIES)}, got {family!r}.")
        data_block = data.get("data")
        if not isinstance(data_block, dict) or not data_block.get("path"):
            raise RecipeError("Recipe data.path is required (CSV or jsonl).")
        train = data.get("train")
        if not isinstance(train, dict):
            raise RecipeError("Recipe train: block is required.")
        method = str(train.get("method") or "").strip().lower()
        if family == "tabular" and method not in TABULAR_METHODS:
            raise RecipeError(
                f"tabular train.method must be one of {list(TABULAR_METHODS)}, got {method!r}."
            )
        if family == "llm" and method not in LLM_METHODS:
            raise RecipeError(f"llm train.method must be one of {list(LLM_METHODS)}, got {method!r}.")
        if family == "tabular" and not data_block.get("target"):
            raise RecipeError("tabular recipes need data.target (label column).")
        if family == "llm" and not train.get("base"):
            raise RecipeError("llm recipes need train.base (HuggingFace model id).")
        task = str(train.get("task") or "classification").lower()
        if family == "tabular" and task not in TABULAR_TASKS:
            raise RecipeError(f"train.task must be one of {list(TABULAR_TASKS)}.")
        known = {"name", "family", "data", "train", "eval"}
        extra = {k: v for k, v in data.items() if k not in known}
        return cls(
            name=name,
            family=family,
            data=dict(data_block),
            train=dict(train),
            eval=dict(data.get("eval") or {}),
            source=source,
            extra=extra,
        )
 
 
@dataclass
class RunRecord:
    run_id: str
    name: str
    family: str
    status: str
    run_dir: str
    recipe: dict[str, Any]
    revision: dict[str, Any]
    checkpoint: dict[str, Any] | None = None
    gate: dict[str, Any] | None = None
    metrics: dict[str, Any] = field(default_factory=dict)
    inspect: dict[str, Any] = field(default_factory=dict)
    error: str | None = None
 
    def to_dict(self) -> dict[str, Any]:
        return asdict(self)