terrine 0.1.0.dev0__py3-none-any.whl

terrine/agent/definitions.py

@@ -0,0 +1,143 @@
+"""Typed discovery records for the existing DatasetDefinition codec."""
+
+import json
+from datetime import datetime
+from typing import Annotated, Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, field_validator
+
+from terrine._time import utc_instant
+from terrine.agent.design_schema import design_schema
+from terrine.dataset.splits import policy_from_record
+from terrine.errors import thaw
+from terrine.model import Model
+from terrine.spec.anchors import FixedDate
+from terrine.spec.dataset import INVALID_DEFINITION, DatasetDefinition
+from terrine.spec.io import from_json
+
+
+class ViewRecord(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
+    features: tuple[str, ...]
+
+    @field_validator("features")
+    @classmethod
+    def _features(cls, value: tuple[str, ...]) -> tuple[str, ...]:
+        if len(set(value)) != len(value) or any(not item for item in value):
+            raise ValueError("features must contain unique nonempty names")
+        return value
+
+
+class GroupRecord(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
+    kind: Literal["group"]
+    group: str
+    train_fraction: str
+    validation_fraction: str
+    seed: int
+    codec: Literal["terrine.group.v1"]
+
+
+class TimeRecord(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
+    kind: Literal["time"]
+    train_until: str
+    validation_until: str | None
+    embargo_microseconds: Annotated[int, Field(ge=0)]
+    groups: GroupRecord | None
+
+
+class TrainingView(ViewRecord):
+    kind: Literal["training"]
+    labels: Annotated[tuple[str, ...], Field(min_length=1)]
+    incomplete_labels: Literal["drop", "keep"]
+    split: Annotated[TimeRecord | GroupRecord, Field(discriminator="kind")] | None
+
+    @field_validator("labels")
+    @classmethod
+    def _labels(cls, value: tuple[str, ...]) -> tuple[str, ...]:
+        return cls._features(value)
+
+
+class ScoringView(ViewRecord):
+    kind: Literal["scoring"]
+    as_of: str
+    eligibility: dict[str, Any]
+    parent_design_hash: Annotated[str, Field(pattern=r"^[0-9a-f]{64}$")]
+
+    @field_validator("as_of")
+    @classmethod
+    def _instant(cls, value: str) -> str:
+        utc_instant(value)
+        return value
+
+
+_VIEW = TypeAdapter(Annotated[TrainingView | ScoringView, Field(discriminator="kind")])
+
+
+def parse_definition(serialized: str, model: Model) -> DatasetDefinition:
+    definition = DatasetDefinition.from_json(serialized)
+    if definition.view is not None:
+        try:
+            _VIEW.validate_json(json.dumps(thaw(definition.view)))
+            if definition.view["kind"] == "training" and definition.view["split"] is not None:
+                policy_from_record(thaw(definition.view["split"]))
+        except (ValueError, OverflowError) as error:
+            raise INVALID_DEFINITION.error(
+                reason="use the discovered typed training or scoring view"
+            ) from error
+        if definition.view["kind"] == "scoring":
+            design = definition.base_design
+            if (
+                design.shape != "cross_section"
+                or not isinstance(design.anchor, FixedDate)
+                or design.outcomes
+                or set(definition.view["features"]) != set(design.features)
+            ):
+                raise INVALID_DEFINITION.error(
+                    reason="scoring view must match its fixed feature-only base Design"
+                )
+            anchor = design.anchor.date
+            resolved = (
+                utc_instant(anchor) if isinstance(anchor, datetime) else model.day_boundary(anchor)
+            )
+            if utc_instant(definition.view["as_of"]) != resolved:
+                raise INVALID_DEFINITION.error(
+                    reason="scoring instant must match its base Design anchor"
+                )
+            document = json.loads(definition.canonical_json())["base_design"]
+            document["eligibility"] = thaw(definition.view["eligibility"])
+            try:
+                eligibility_design = from_json(
+                    json.dumps({"schema_version": 1, "design": document})
+                )
+            except (KeyError, TypeError, ValueError, OverflowError) as error:
+                raise INVALID_DEFINITION.error(
+                    reason="scoring eligibility must use typed predicates"
+                ) from error
+            if eligibility_design.eligibility != design.eligibility:
+                raise INVALID_DEFINITION.error(
+                    reason="scoring eligibility must match its base Design"
+                )
+    return definition
+
+
+def definition_schema() -> dict[str, Any]:
+    design = design_schema()
+    view = _VIEW.json_schema()
+    definitions = {**design["$defs"], **view.pop("$defs", {})}
+    definitions["ScoringView"]["properties"]["eligibility"] = definitions["Design"]["properties"][
+        "eligibility"
+    ]
+    return {
+        "$schema": design["$schema"],
+        "$defs": definitions,
+        "type": "object",
+        "additionalProperties": False,
+        "required": ["schema_version", "base_design", "view"],
+        "properties": {
+            "schema_version": {"const": 1},
+            "base_design": {"$ref": "#/$defs/Design"},
+            "view": {"anyOf": [view, {"type": "null"}]},
+        },
+    }

terrine/agent/design_schema.py

@@ -0,0 +1,220 @@
+"""Discovery schema for the existing versioned Design codec."""
+
+from typing import Any
+
+
+def design_schema() -> dict[str, Any]:
+    def ref(name):
+        return {"$ref": f"#/$defs/{name}"}
+
+    def union(*names):
+        return {"anyOf": [ref(name) for name in names]}
+
+    def nullable(schema):
+        return {"anyOf": [schema, {"type": "null"}]}
+
+    def array(schema):
+        return {"type": "array", "items": schema}
+
+    def mapping(schema):
+        return {"type": "object", "additionalProperties": schema}
+
+    def obj(properties, required=None):
+        return {
+            "type": "object",
+            "properties": properties,
+            "required": list(properties) if required is None else required,
+            "additionalProperties": False,
+        }
+
+    text = {"type": "string"}
+    number = {"type": "number"}
+    integer = {"type": "integer"}
+    boolean = {"type": "boolean"}
+    scalar = {"type": ["string", "number", "boolean", "null"]}
+    duration = {"type": ["integer", "string"]}
+    definitions = {"Temporal": obj({"type": {"enum": ["date", "datetime"]}, "value": text})}
+    definitions["Cutpoint"] = {
+        "anyOf": [
+            obj({"type": {"enum": ["int", "float64"]}, "value": text}),
+            obj(
+                {
+                    "type": {"const": "decimal"},
+                    "sign": {"enum": [0, 1]},
+                    "digits": array({"type": "integer", "minimum": 0, "maximum": 9}),
+                    "exponent": integer,
+                }
+            ),
+        ]
+    }
+    expression = union("MetricAgg", "AttributeAsOf", "EventTime")
+    predicate = union("Gt", "Eq", "IsNull", "And")
+    named = mapping(expression)
+    outputs = mapping({"anyOf": [expression, predicate]})
+    nodes = {
+        "Window": {"start": number, "end": number},
+        "FixedDate": {"date": ref("Temporal")},
+        "EventAnchor": {
+            "fact": text,
+            "which": {"enum": ["first", "nth", "each"]},
+            "n": nullable(integer),
+        },
+        "Schedule": {"every": duration, "start": ref("Temporal"), "end": ref("Temporal")},
+        "AnchorRange": {"start": nullable(ref("Temporal")), "end": nullable(ref("Temporal"))},
+        "MetricAgg": {
+            "metric": text,
+            "window": ref("Window"),
+            "relative_to": nullable(ref("EventTime")),
+        },
+        "AttributeAsOf": {
+            "attribute": text,
+            "before": number,
+            "relative_to": nullable(ref("EventTime")),
+        },
+        "EventTime": {
+            "fact": text,
+            "window": ref("Window"),
+            "label_property": nullable(text),
+            "equals": scalar,
+            "occurrence": {"enum": ["first", "last", "nth"]},
+            "n": nullable(integer),
+            "relative_to": nullable(ref("EventTime")),
+        },
+        "Gt": {"expr": expression, "value": scalar},
+        "Eq": {"left": expression, "value": scalar},
+        "IsNull": {"expr": expression, "negate": boolean},
+        "And": {"terms": array(predicate)},
+        "Treatment": {"attribute": text, "at": number, "known_at_t0": boolean},
+        "StateObservations": {
+            "property": text,
+            "levels": array(scalar),
+            "ordered": boolean,
+            "absorbing": array(scalar),
+        },
+        "BinnedMetricState": {
+            "metric": text,
+            "cutpoints": array(ref("Cutpoint")),
+            "levels": array(scalar),
+        },
+        "StatePanel": {
+            "state": union("StateObservations", "BinnedMetricState"),
+            "max_periods": integer,
+            "baseline_features": named,
+        },
+        "Trajectory": {
+            "metric": text,
+            "period": duration,
+            "periods": integer,
+            "cumulative": boolean,
+        },
+        "CohortSpec": {"origin": text, "trajectory": ref("Trajectory")},
+        "TimeToEvent": {
+            "event": text,
+            "exit": nullable(text),
+            "entry": nullable(text),
+            "delayed_entry": boolean,
+        },
+        "LifeTable": {
+            "time_to_event": ref("TimeToEvent"),
+            "interval": duration,
+            "group_by": array(text),
+        },
+        "SessionSpec": {"gap_seconds": number},
+        "FunnelSpec": {
+            "event_property": text,
+            "steps": array(text),
+            "within_seconds": nullable(duration),
+        },
+        "EventSequence": {
+            "fact": text,
+            "window": ref("Window"),
+            "max_events": integer,
+            "representation": text,
+            "properties": array(text),
+            "label_property": nullable(text),
+            "sessions": nullable(ref("SessionSpec")),
+            "funnel": nullable(ref("FunnelSpec")),
+        },
+        "Initiation": {"source": text, "property": text, "history_start": text, "knowledge": text},
+        "PersistentAdoption": {
+            "source": text,
+            "property": text,
+            "adopted_value": scalar,
+            "history_start": text,
+            "persistence": text,
+            "knowledge": text,
+        },
+        "MatchedT0Sample": {
+            "strata": array(text),
+            "controls_per_treated": integer,
+            "seed": integer,
+        },
+        "TrialWeight": {"kind": text, "values": nullable(mapping(number))},
+        "SequentialSpec": {
+            "treatment": union("Initiation", "PersistentAdoption"),
+            "trial_weight": nullable(ref("TrialWeight")),
+            "anticipation_seconds": number,
+            "reference_period_seconds": number,
+        },
+        "EventPanelSpec": {
+            "periods": array(integer),
+            "period_seconds": number,
+            "reference_period": integer,
+        },
+        "CalendarComparisonSpec": {
+            "cohort": ref("Initiation"),
+            "control": text,
+            "reference_period": integer,
+            "comparison_periods": array(integer),
+            "anticipation_periods": array(integer),
+        },
+    }
+    optional = {
+        "EventAnchor": {"n"},
+        "MetricAgg": {"relative_to"},
+        "AttributeAsOf": {"relative_to"},
+        "EventTime": {"relative_to"},
+        "IsNull": {"negate"},
+        "Treatment": {"at"},
+        "StateObservations": {"ordered", "absorbing"},
+    }
+    for name, fields in nodes.items():
+        properties = {"type": {"const": name}, **fields}
+        definitions[name] = obj(
+            properties, [key for key in properties if key not in optional.get(name, set())]
+        )
+    definitions["Design"] = obj(
+        {
+            "unit": text,
+            "anchor": union("FixedDate", "EventAnchor", "Schedule"),
+            "features": outputs,
+            "outcomes": outputs,
+            "eligibility": mapping(predicate),
+            "treatment": nullable(ref("Treatment")),
+            "gap": number,
+            "shape": {
+                "anyOf": [
+                    text,
+                    union(
+                        "StatePanel",
+                        "CohortSpec",
+                        "TimeToEvent",
+                        "LifeTable",
+                        "EventSequence",
+                        "SequentialSpec",
+                        "EventPanelSpec",
+                        "CalendarComparisonSpec",
+                    ),
+                ]
+            },
+            "anchor_range": nullable(ref("AnchorRange")),
+            "sample": nullable(ref("MatchedT0Sample")),
+            "administrative_cutoff": nullable(ref("Temporal")),
+        },
+        ["unit", "anchor"],
+    )
+    return {
+        "$schema": "https://json-schema.org/draft/2020-12/schema",
+        "$defs": definitions,
+        **obj({"schema_version": {"const": 1}, "design": ref("Design")}),
+    }

terrine/agent/policy.py

@@ -0,0 +1,144 @@
+"""Owner-only grants and resource limits for dataset operations."""
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Literal
+
+from terrine.data import Budget
+from terrine.errors import RefusalSpec
+from terrine.runtime.api import Executor
+from terrine.runtime.saved_plan import DestinationRef, ModelReference, ModelResolver
+from terrine.runtime.snapshot import SnapshotAuthority
+
+__all__ = ["AgentPolicy"]
+
+type Permission = Literal[
+    "inspect", "validate", "prepare", "write", "export", "preview", "read", "cancel"
+]
+
+_DENIED = RefusalSpec("terrine.agent.permission_denied", "Owner policy denies this operation.")
+_SCOPE = RefusalSpec("terrine.agent.scope_denied", "Requested resources are outside owner policy.")
+_LIMIT = RefusalSpec("terrine.agent.request_limit", "Owner tool request limit has been reached.")
+
+
+def effective_budget(owner: Budget, requested: Budget | None) -> Budget:
+    """Intersect optional nonnegative limits; None denotes no numeric bound."""
+    if not isinstance(owner, Budget) or (
+        requested is not None and not isinstance(requested, Budget)
+    ):
+        raise TypeError("budgets must be Budget values")
+    if requested is None:
+        return owner
+    values = {}
+    for name in ("max_rows", "max_output_bytes", "max_scan_bytes"):
+        limits = [
+            value for value in (getattr(owner, name), getattr(requested, name)) if value is not None
+        ]
+        values[name] = min(limits) if limits else None
+    return Budget(**values, strict_scan=owner.strict_scan or requested.strict_scan)
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class AgentPolicy:
+    """Trusted configuration, never deserialized from a tool request.
+
+    Source identities are Catalog source names, destinations are exact owner
+    references, and catalogs use the owner ModelReference's catalog hash.
+    """
+
+    inspect: bool = True
+    validate: bool = True
+    prepare: bool = False
+    write: bool = False
+    export: bool = False
+    preview: bool = False
+    read: bool = False
+    cancel: bool = False
+    allowed_sources: frozenset[str] = frozenset()
+    allowed_catalogs: frozenset[str] = frozenset()
+    allowed_destinations: frozenset[DestinationRef] = frozenset()
+    allowed_exports: frozenset[Path] = frozenset()
+    design_root: Path | None = None
+    budget: Budget = field(default_factory=Budget)
+    max_run_seconds: int = 300
+    max_concurrency: int = 1
+    max_requests: int = 100
+    max_preview_rows: int = 20
+    max_bounded_count: int = 1000
+    batch_cap: int = 100
+    preview_columns: frozenset[str] = frozenset()
+    redacted_columns: frozenset[str] = frozenset()
+    model_reference: ModelReference | None = None
+    resolver: ModelResolver | None = field(default=None, repr=False, compare=False)
+    snapshot_authority: SnapshotAuthority | None = field(default=None, repr=False, compare=False)
+    executor: Executor | None = field(default=None, repr=False, compare=False)
+
+    def __post_init__(self) -> None:
+        for name in (
+            "inspect",
+            "validate",
+            "prepare",
+            "write",
+            "export",
+            "preview",
+            "read",
+            "cancel",
+        ):
+            if not isinstance(getattr(self, name), bool):
+                raise TypeError(f"{name} must be bool")
+        for name in (
+            "max_run_seconds",
+            "max_concurrency",
+            "max_requests",
+            "max_preview_rows",
+            "max_bounded_count",
+            "batch_cap",
+        ):
+            value = getattr(self, name)
+            if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
+                raise ValueError(f"{name} must be a positive integer")
+        for name in ("allowed_sources", "allowed_catalogs", "preview_columns", "redacted_columns"):
+            value = frozenset(getattr(self, name))
+            if any(not isinstance(item, str) or not item for item in value):
+                raise TypeError(f"{name} must contain nonempty strings")
+            object.__setattr__(self, name, value)
+        destinations = frozenset(self.allowed_destinations)
+        if any(not isinstance(item, DestinationRef) for item in destinations):
+            raise TypeError("allowed_destinations must contain DestinationRef values")
+        object.__setattr__(self, "allowed_destinations", destinations)
+        exports = frozenset(Path(path).expanduser().resolve() for path in self.allowed_exports)
+        object.__setattr__(self, "allowed_exports", exports)
+        if not isinstance(self.budget, Budget):
+            raise TypeError("budget must be Budget")
+        if self.design_root is not None:
+            object.__setattr__(self, "design_root", Path(self.design_root).resolve())
+
+    def authorize(
+        self,
+        permission: Permission,
+        *,
+        sources: tuple[str, ...] = (),
+        catalog: str | None = None,
+        destination: DestinationRef | None = None,
+    ) -> None:
+        if permission not in {
+            "inspect",
+            "validate",
+            "prepare",
+            "write",
+            "export",
+            "preview",
+            "read",
+            "cancel",
+        } or not getattr(self, permission):
+            raise _DENIED.error(
+                operation=permission, route="Ask the owner to grant this operation."
+            )
+        if (
+            not set(sources) <= self.allowed_sources
+            or (catalog is not None and catalog not in self.allowed_catalogs)
+            or (destination is not None and destination not in self.allowed_destinations)
+        ):
+            raise _SCOPE.error(
+                operation=permission, route="Use owner-approved sources and destinations."
+            )