diff --git a/evomerge/archival.py b/evomerge/archival.py new file mode 100644 index 0000000..d7958cb --- /dev/null +++ b/evomerge/archival.py @@ -0,0 +1,215 @@ +"""Trace archival — retention policies and tiered cold-storage migration. + +This module is the **retention / tiering** half of the Milestone-5 trace +archival and retrieval feature (issue #53). It builds on +:mod:`evomerge.storage` and provides two data-side primitives: + + 1. ``RetentionPolicy`` — declares how long traces stay in the *hot* tier + before ageing out to *cold*, and the absolute age cap after which they + are deleted. ``plan()`` inspects a ``TraceStorage`` and returns a + ``RetentionReport`` of planned migrate/delete actions; ``apply()`` + executes them. + 2. ``ArchivalStore`` — pairs a hot ``TraceStorage`` with a cold one (a + different backend, e.g. S3 Glacier-class, and/or a cheaper codec, e.g. + Parquet with zstd) and moves aged partitions hot → cold, deleting from + hot once safely persisted. + +Both primitives are backend-agnostic: the actual object-storage lifecycle +(S3 lifecycle rules, GCS object archival) can additionally enforce retention +at the bucket level; this module provides the explicit, testable in-process +equivalent so retention is reproducible and auditable. +""" +from __future__ import annotations + +import re +from dataclasses import dataclass, field +from datetime import datetime, timezone + +from evomerge.storage import _GRANULARITY, TraceStorage + +#: Regex extracting the ``dt=`` bucket from a Hive-style partition directory. +_DT_RE = re.compile(r"dt=([^/]+)/?") + + +@dataclass(frozen=True) +class RetentionPolicy: + """Declarative retention schedule for trace partitions. + + Attributes: + hot_days: partitions older than this many days are migrated hot → cold. + delete_after_days: partitions older than this are deleted outright + (must be ``> hot_days`` when set; ``None`` = retain forever in cold). + """ + + hot_days: int = 30 + delete_after_days: int | None = None + + def __post_init__(self) -> None: + if self.hot_days < 0: + raise ValueError("hot_days must be non-negative") + if self.delete_after_days is not None: + if self.delete_after_days <= self.hot_days: + raise ValueError("delete_after_days must be greater than hot_days") + if self.delete_after_days < 0: + raise ValueError("delete_after_days must be non-negative") + + def action_for(self, age_days: float) -> tuple[str, str]: + """Return ``(action, reason)`` for a partition of the given age. + + ``action`` is ``"delete"``, ``"migrate"``, or ``"keep"``. + """ + if self.delete_after_days is not None and age_days >= self.delete_after_days: + return "delete", f"age {age_days:.1f}d >= delete_after_days {self.delete_after_days}d" + if age_days >= self.hot_days: + return "migrate", f"age {age_days:.1f}d >= hot_days {self.hot_days}d" + return "keep", f"age {age_days:.1f}d < hot_days {self.hot_days}d" + + +@dataclass +class RetentionAction: + """A single planned retention action on one partition.""" + + partition: str + action: str # "migrate" | "delete" + reason: str + age_days: float + tier: str = "hot" # which tier the partition currently lives in + + +@dataclass +class RetentionReport: + """Outcome of a retention pass — planned or executed actions.""" + + actions: list[RetentionAction] = field(default_factory=list) + + @property + def migrated(self) -> list[RetentionAction]: + return [a for a in self.actions if a.action == "migrate"] + + @property + def deleted(self) -> list[RetentionAction]: + return [a for a in self.actions if a.action == "delete"] + + @property + def kept(self) -> list[RetentionAction]: + return [a for a in self.actions if a.action == "keep"] + + def __len__(self) -> int: + return len(self.actions) + + +def partition_age_days(partition_dir: str, granularity: str, *, now: datetime | None = None) -> float: + """Age (in days) of a partition derived from its ``dt=`` bucket. + + For ``day`` granularity the bucket is ``YYYY-MM-DD``; for ``month`` it is + ``YYYY-MM`` (aged from the first of the month); for ``hour`` it is + ``YYYY-MM-DDTHH``. Returns ``+inf`` if no ``dt=`` segment is present. + """ + now = now if now is not None else datetime.now(tz=timezone.utc) + match = _DT_RE.search(partition_dir) + if not match: + return float("inf") + bucket = match.group(1) + fmt = _GRANULARITY[granularity] + try: + part_dt = datetime.strptime(bucket, fmt).replace(tzinfo=timezone.utc) + except ValueError: + return float("inf") + return (now - part_dt).total_seconds() / 86400.0 + + +class ArchivalStore: + """Tiered trace store — migrates aged partitions hot → cold, then deletes. + + The *hot* store is the low-latency retrieval tier (default JSONL on local + disk or S3 standard); the *cold* store is the archival tier (Parquet with + a stronger codec, a Glacier-class bucket, or simply a different prefix). + Records are re-encoded on migration so the two tiers may use different + codecs — a columnar cold tier is cheaper to scan than row-based hot. + """ + + def __init__(self, hot: TraceStorage, cold: TraceStorage) -> None: + if hot.granularity != cold.granularity: + raise ValueError( + "hot and cold stores must share time granularity " + f"({hot.granularity!r} vs {cold.granularity!r})" + ) + self.hot = hot + self.cold = cold + + # -- planning ------------------------------------------------------------ + + def plan(self, policy: RetentionPolicy, *, now: datetime | None = None) -> RetentionReport: + """Plan retention actions for the hot tier (no side effects).""" + now = now if now is not None else datetime.now(tz=timezone.utc) + actions: list[RetentionAction] = [] + for pdir in self.hot.partitions(): + age = partition_age_days(pdir, self.hot.granularity, now=now) + action, reason = policy.action_for(age) + actions.append(RetentionAction(partition=pdir, action=action, reason=reason, age_days=age)) + # Also consider deletion of already-archived cold partitions. + for pdir in self.cold.partitions(): + age = partition_age_days(pdir, self.cold.granularity, now=now) + action, reason = policy.action_for(age) + if action == "delete": + actions.append( + RetentionAction(partition=pdir, action="delete", reason=reason, age_days=age, tier="cold") + ) + return RetentionReport(actions=actions) + + # -- execution ----------------------------------------------------------- + + def migrate_partition(self, partition_dir: str) -> int: + """Move one hot partition to cold; return the record count migrated. + + Records are decoded from the hot codec and re-encoded into the cold + codec, then removed from hot. Idempotent: re-migrating an emptied + partition moves zero records. + """ + records: list[dict] = [] + for key in self.hot.backend.list_keys(partition_dir): + records.extend(self.hot.codec.loads(self.hot.backend.read_bytes(key))) + if records: + self.cold.write(records) + for key in self.hot.backend.list_keys(partition_dir): + self.hot.backend.delete(key) + return len(records) + + def delete_partition(self, partition_dir: str, *, tier: str = "hot") -> int: + """Delete every object in a partition from the given tier.""" + store = self.cold if tier == "cold" else self.hot + keys = store.backend.list_keys(partition_dir) + for key in keys: + store.backend.delete(key) + return len(keys) + + def apply( + self, + policy: RetentionPolicy, + *, + now: datetime | None = None, + dry_run: bool = False, + ) -> RetentionReport: + """Execute the retention pass: migrate aged hot partitions, delete expired. + + With ``dry_run=True`` this is identical to :meth:`plan` (no side effects). + Returns a report of the actions taken (or planned). + """ + report = self.plan(policy, now=now) + if dry_run: + return report + for action in report.actions: + if action.action == "migrate": + self.migrate_partition(action.partition) + elif action.action == "delete": + self.delete_partition(action.partition, tier=action.tier) + return report + + +__all__ = [ + "ArchivalStore", + "RetentionAction", + "RetentionPolicy", + "RetentionReport", + "partition_age_days", +] diff --git a/evomerge/storage.py b/evomerge/storage.py new file mode 100644 index 0000000..5e53593 --- /dev/null +++ b/evomerge/storage.py @@ -0,0 +1,700 @@ +"""Trace archival storage — time-partitioned record store with indexed retrieval. + +This module is the **storage layer** of the Milestone-5 trace archival and +retrieval feature (issue #53). It implements three data-side primitives: + + 1. ``StorageBackend`` — a byte-store Protocol with ``LocalBackend`` (disk, + the default) plus thin ``S3Backend`` / ``GCSBackend`` adapters that lazily + pull in ``boto3`` / ``google-cloud-storage``. Real bucket I/O is wired + through these; the cluster/bucket provisioning lives outside this repo, + mirroring how ``training_pipeline.py`` treats Ray/Lightning. + 2. ``TraceCodec`` — row ⇄ bytes serialisation. ``JsonLinesCodec`` is the + zero-dependency default; ``ParquetCodec`` writes genuine columnar Parquet + via ``pyarrow`` (optional ``[arch]`` extra). Both are interchangeable. + 3. ``TraceStorage`` — time-partitioned record store with an indexed query + API. Records are laid out Hive-style under partition directories + (``subject_id=/dt=YYYY-MM-DD/.``) so that audit-trail, + reproduction, and historical-trust queries prune partitions by + ``subject_id`` / ``user_id`` / date-range without scanning the whole store. + +The Parquet + S3/GCS pieces are *opt-in*: ``TraceStorage`` defaults to the +local backend and the JSON-Lines codec so ``python -m pytest`` passes with no +extra dependencies. Install the ``[arch]`` extra to enable the Parquet codec. +""" +from __future__ import annotations + +import io +import json +import os +import time +from collections.abc import Iterable, Sequence +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Protocol, runtime_checkable + +# =========================================================================== +# Record helpers +# =========================================================================== + +#: Supported time-partition granularities → (strftime bucket, timedelta unit). +_GRANULARITY: dict[str, str] = { + "hour": "%Y-%m-%dT%H", + "day": "%Y-%m-%d", + "month": "%Y-%m", +} + + +def _record_to_dict(rec: Any) -> dict[str, Any]: + """Normalise a trace record (pydantic model or plain dict) to a dict.""" + if isinstance(rec, dict): + return rec + if hasattr(rec, "model_dump"): # pydantic v2 + return rec.model_dump(mode="json") + if hasattr(rec, "dict"): # pydantic v1 fallback + return rec.dict() + return dict(rec) + + +def _extract_field(record: dict[str, Any], name: str) -> Any: + """Read ``name`` from a record, falling back to AEP v0.3 ``run_context``. + + AEP v0.3 nests ``user_id`` / ``subject_id`` under ``run_context``; rollout / + training records carry them at the top level. This helper covers both. + """ + if name in record: + return record[name] + ctx = record.get("run_context") + if isinstance(ctx, dict) and name in ctx: + return ctx[name] + return None + + +def _parse_time(value: Any) -> datetime | None: + """Coerce a timestamp value to an aware ``datetime``. + + Accepts ISO-8601 strings (with or without trailing ``Z``), epoch seconds, + and epoch milliseconds. Returns ``None`` if the value is missing/blank. + """ + if value is None or value == "": + return None + if isinstance(value, (int, float)): + # Heuristic: values > 1e12 are milliseconds. + secs = value / 1000.0 if abs(value) > 1e12 else float(value) + return datetime.fromtimestamp(secs, tz=timezone.utc) + if isinstance(value, str): + s = value.strip() + if not s: + return None + if s.endswith("Z"): + s = s[:-1] + "+00:00" + try: + dt = datetime.fromisoformat(s) + except ValueError: + return None + if dt.tzinfo is None: + dt = dt.replace(tzinfo=timezone.utc) + return dt + return None + + +def _record_time(record: dict[str, Any], time_field: str) -> datetime: + """Extract the record timestamp, trying ``time_field`` then common aliases.""" + for candidate in (time_field, "timestamp", "ts", "created_at", "recorded_at", "time"): + value = _extract_field(record, candidate) + dt = _parse_time(value) + if dt is not None: + return dt + # Last resort: ingestion time, so a partition is always derivable. + return datetime.now(tz=timezone.utc) + + +def _segment(value: Any) -> str: + """Render a partition value as a path-safe segment (Hive ``key=value``).""" + return str(value).replace("/", "_").replace(os.sep, "_") + + +def _join(prefix: str, child: str) -> str: + """Join a parent key prefix (may lack a trailing '/') and a child segment. + + ``child`` is a relative path returned by ``StorageBackend.list_dirs`` (it + carries its own trailing '/' for directories). This helper guarantees + exactly one '/' separates the two, so a namespace like ``"traces"`` does + not fuse with its first partition segment. + """ + if not prefix: + return child + return prefix.rstrip("/") + "/" + child + + +def _bucket_start(dt: datetime, granularity: str) -> datetime: + """Floor ``dt`` to the start of its time bucket.""" + if granularity == "hour": + return dt.replace(minute=0, second=0, microsecond=0) + if granularity == "day": + return dt.replace(hour=0, minute=0, second=0, microsecond=0) + return dt.replace(day=1, hour=0, minute=0, second=0, microsecond=0) + + +def _advance(dt: datetime, granularity: str) -> datetime: + """Advance to the start of the next bucket (boundary, not time-of-day).""" + start = _bucket_start(dt, granularity) + if granularity == "hour": + return datetime.fromtimestamp(start.timestamp() + 3600, tz=start.tzinfo) + if granularity == "day": + return datetime.fromtimestamp(start.timestamp() + 86400, tz=start.tzinfo) + # month: anchor on day-1 to avoid skipped-day overflow. + if start.month == 12: + return start.replace(year=start.year + 1, month=1) + return start.replace(month=start.month + 1) + + +def _date_buckets(start: datetime | None, end: datetime | None, granularity: str) -> list[str] | None: + """Enumerate the distinct time buckets spanning ``[start, end]``. + + Returns ``None`` when neither bound is given (meaning "all partitions"); + otherwise returns the inclusive list of bucket strings for the granularity, + enumerated by bucket *boundary* so a range that lands mid-bucket still + covers every bucket it touches. + """ + fmt = _GRANULARITY[granularity] + if start is None and end is None: + return None + if start is None: + start = end # type: ignore[assignment] + if end is None: + end = start # type: ignore[assignment] + if end < start: + start, end = end, start + + buckets: list[str] = [] + cursor = _bucket_start(start, granularity) + end_bucket = _bucket_start(end, granularity) + # Cap to avoid pathological enumeration of unbounded ranges. + max_buckets = 100_000 + while cursor <= end_bucket and len(buckets) < max_buckets: + buckets.append(cursor.strftime(fmt)) + cursor = _advance(cursor, granularity) + return buckets + + +# =========================================================================== +# Storage backends +# =========================================================================== + + +@runtime_checkable +class StorageBackend(Protocol): + """Byte-store Protocol implemented by local disk, S3, and GCS backends. + + Keys are POSIX-style relative paths (``"subject_id=s/dt=2026-07-28/b.parquet"``). + Implementations are responsible for translating keys to their physical + location (filesystem path, S3 object key, GCS blob name). + """ + + def write_bytes(self, key: str, data: bytes) -> None: ... + def read_bytes(self, key: str) -> bytes: ... + def exists(self, key: str) -> bool: ... + def delete(self, key: str) -> None: ... + def list_keys(self, prefix: str = "") -> list[str]: ... + def list_dirs(self, prefix: str = "") -> list[str]: ... + + +class LocalBackend: + """Filesystem-backed storage (the default; used by tests and on-disk prod).""" + + def __init__(self, root: str | Path) -> None: + self.root = Path(root) + + def _path(self, key: str) -> Path: + return self.root / key + + def write_bytes(self, key: str, data: bytes) -> None: + path = self._path(key) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(data) + + def read_bytes(self, key: str) -> bytes: + return self._path(key).read_bytes() + + def exists(self, key: str) -> bool: + return self._path(key).exists() + + def delete(self, key: str) -> None: + self._path(key).unlink(missing_ok=True) + + def list_keys(self, prefix: str = "") -> list[str]: + base = self._path(prefix) + if not base.exists(): + return [] + out: list[str] = [] + for p in sorted(base.rglob("*")): + if p.is_file(): + out.append(str(p.relative_to(self.root)).replace(os.sep, "/")) + return out + + def list_dirs(self, prefix: str = "") -> list[str]: + base = self._path(prefix) + if not base.exists(): + return [] + return [ + child.name + "/" + for child in sorted(base.iterdir()) + if child.is_dir() + ] + + +class S3Backend: + """Amazon S3 storage adapter — lazily imports ``boto3``. + + Real bucket I/O (get/put/list/delete) is implemented; credentials, endpoint + configuration, and bucket provisioning live with the deployment, not here. + """ + + def __init__( + self, + bucket: str, + prefix: str = "", + *, + client: Any = None, + endpoint_url: str | None = None, + region: str | None = None, + not_found_exc: type[BaseException] | tuple[type[BaseException], ...] | None = None, + ) -> None: + self.bucket = bucket + self.prefix = prefix.strip("/") + # Exception(s) ``head_object`` raises for a missing key. Defaults to + # ``botocore.exceptions.ClientError`` (lazy); injectable so the adapter + # is testable without the AWS SDK installed. + self._not_found_exc = not_found_exc + if client is not None: + self._client = client + else: # pragma: no cover - exercised only when boto3 is installed + import boto3 # type: ignore[import-not-found] + + kwargs: dict[str, Any] = {} + if endpoint_url: + kwargs["endpoint_url"] = endpoint_url + if region: + kwargs["region_name"] = region + self._client = boto3.client("s3", **kwargs) + + def _key(self, key: str) -> str: + return f"{self.prefix}/{key}" if self.prefix else key + + def write_bytes(self, key: str, data: bytes) -> None: + self._client.put_object(Bucket=self.bucket, Key=self._key(key), Body=data) + + def read_bytes(self, key: str) -> bytes: + return self._client.get_object(Bucket=self.bucket, Key=self._key(key))["Body"].read() + + def exists(self, key: str) -> bool: + if self._not_found_exc is None: + from botocore.exceptions import ClientError # pragma: no cover - needs AWS SDK + + not_found = ClientError + else: + not_found = self._not_found_exc + try: + self._client.head_object(Bucket=self.bucket, Key=self._key(key)) + return True + except not_found: + return False + + def delete(self, key: str) -> None: + self._client.delete_object(Bucket=self.bucket, Key=self._key(key)) + + def list_keys(self, prefix: str = "") -> list[str]: + full = self._key(prefix) + paginator = self._client.get_paginator("list_objects_v2") + out: list[str] = [] + for page in paginator.paginate(Bucket=self.bucket, Prefix=full): + for obj in page.get("Contents", []): + if self.prefix: + out.append(obj["Key"][len(self.prefix) + 1 :]) + else: + out.append(obj["Key"]) + return out + + def list_dirs(self, prefix: str = "") -> list[str]: + full = self._key(prefix) + resp = self._client.list_objects_v2( + Bucket=self.bucket, Prefix=full, Delimiter="/" + ) + return [p["Prefix"].rstrip("/").split("/")[-1] + "/" for p in resp.get("CommonPrefixes", [])] + + +class GCSBackend: + """Google Cloud Storage adapter — lazily imports ``google-cloud-storage``.""" + + def __init__(self, bucket: str, prefix: str = "", *, client: Any = None) -> None: + self.bucket_name = bucket + self.prefix = prefix.strip("/") + if client is not None: + self._client = client + else: # pragma: no cover + from google.cloud import storage # type: ignore[import-not-found] + + self._client = storage.Client() + self._bucket = self._client.bucket(bucket) + + def _name(self, key: str) -> str: + return f"{self.prefix}/{key}" if self.prefix else key + + def write_bytes(self, key: str, data: bytes) -> None: # pragma: no cover + self._bucket.blob(self._name(key)).upload_from_string(data) + + def read_bytes(self, key: str) -> bytes: # pragma: no cover + return self._bucket.blob(self._name(key)).download_as_bytes() + + def exists(self, key: str) -> bool: # pragma: no cover + return self._bucket.blob(self._name(key)).exists() + + def delete(self, key: str) -> None: # pragma: no cover + self._bucket.blob(self._name(key)).delete() + + def list_keys(self, prefix: str = "") -> list[str]: # pragma: no cover + full = self._name(prefix) + out: list[str] = [] + for blob in self._client.list_blobs(self.bucket_name, prefix=full): + if self.prefix: + out.append(blob.name[len(self.prefix) + 1 :]) + else: + out.append(blob.name) + return out + + def list_dirs(self, prefix: str = "") -> list[str]: # pragma: no cover + full = self._name(prefix) + # list_blobs with delimiter yields "prefixes" for immediate sub-dirs. + kwargs: dict[str, Any] = {"prefix": full, "delimiter": "/"} + _iter = self._client.list_blobs(self.bucket_name, **kwargs) + list(_iter) # consume to populate prefixes + return [ + p.rstrip("/").split("/")[-1] + "/" + for p in _iter.prefixes + ] + + +# =========================================================================== +# Codecs +# =========================================================================== + + +@runtime_checkable +class TraceCodec(Protocol): + """Row ⇄ bytes serialisation contract for stored trace batches.""" + + ext: str + + def dumps(self, rows: list[dict[str, Any]]) -> bytes: ... + def loads(self, data: bytes) -> list[dict[str, Any]]: ... + + +class JsonLinesCodec: + """Zero-dependency newline-delimited JSON codec (default).""" + + ext = "jsonl" + + def dumps(self, rows: list[dict[str, Any]]) -> bytes: + return ("\n".join(json.dumps(r, ensure_ascii=False) for r in rows)).encode("utf-8") + + def loads(self, data: bytes) -> list[dict[str, Any]]: + if not data: + return [] + return [json.loads(line) for line in data.decode("utf-8").splitlines() if line.strip()] + + +class ParquetCodec: + """Columnar Parquet codec backed by ``pyarrow`` (``[arch]`` extra). + + Each top-level record field becomes a Parquet column; heterogeneous batches + (records with differing keys) are handled by ``pyarrow.Table.from_pylist``, + which fills absent fields with nulls — exactly the Hive/Iceberg behaviour + trace archival relies on. + """ + + ext = "parquet" + + def __init__(self, compression: str = "snappy") -> None: + self.compression = compression + + def _import(self): # pragma: no cover - imported lazily + try: + import pyarrow as pa # type: ignore[import-not-found] + import pyarrow.parquet as pq # type: ignore[import-not-found] + except ImportError as exc: # pragma: no cover + raise ImportError( + "ParquetCodec requires pyarrow. Install with: pip install 'evomerge[arch]'" + ) from exc + return pa, pq + + def dumps(self, rows: list[dict[str, Any]]) -> bytes: + pa, pq = self._import() # pragma: no cover + if not rows: + rows = [{"_empty": True}] # parquet needs at least one column + table = pa.Table.from_pylist(rows) + sink = io.BytesIO() + pq.write_table(table, sink, compression=self.compression) + return sink.getvalue() + + def loads(self, data: bytes) -> list[dict[str, Any]]: # pragma: no cover + pa, pq = self._import() + table = pq.read_table(io.BytesIO(data)) + rows = table.to_pylist() + return [r for r in rows if not r.get("_empty")] + + +# =========================================================================== +# Trace storage +# =========================================================================== + + +@dataclass +class StorageManifest: + """Summary of a ``TraceStorage`` write batch.""" + + written: int # records persisted + partitions: list[str] # partition dirs touched + objects: list[str] # object keys written + + @property + def n_partitions(self) -> int: + return len(self.partitions) + + +class TraceStorage: + """Time-partitioned trace store with indexed query APIs. + + Records are persisted under Hive-style partition directories whose segments + correspond to ``partition_by`` (default ``subject_id``) and a time bucket + ``dt`` (granularity-controlled). The directory layout *is* the index: an + audit-trail / reproduction / historical-trust query supplies a + ``subject_id`` / ``user_id`` / date-range and only the matching partitions + are read — the rest are pruned at the backend-listing level. + """ + + def __init__( + self, + backend: StorageBackend, + *, + codec: TraceCodec | None = None, + partition_by: Sequence[str] = ("subject_id",), + time_field: str = "timestamp", + granularity: str = "day", + namespace: str = "traces", + ) -> None: + if granularity not in _GRANULARITY: + raise ValueError( + f"granularity must be one of {sorted(_GRANULARITY)}, got {granularity!r}" + ) + for dim in partition_by: + if dim not in ("subject_id", "user_id", "tenant"): + raise ValueError( + f"partition_by dimension {dim!r} not supported " + "(allowed: subject_id, user_id, tenant)" + ) + self.backend = backend + self.codec = codec if codec is not None else JsonLinesCodec() + self.partition_dims = tuple(partition_by) + self.time_field = time_field + self.granularity = granularity + self.namespace = namespace.strip("/") + + # -- partition layout --------------------------------------------------- + + def _partition_dir(self, record: dict[str, Any]) -> str: + """Compute the Hive-style partition directory for one record.""" + dt = _record_time(record, self.time_field) + parts = [self.namespace] if self.namespace else [] + for dim in self.partition_dims: + value = _extract_field(record, dim) + # Unknown shard key → "_unknown" bucket so the record is still + # retrievable (and visibly un-partitioned) rather than dropped. + parts.append(f"{dim}={_segment(value) if value is not None else '_unknown'}") + parts.append(f"dt={dt.strftime(_GRANULARITY[self.granularity])}") + return "/".join(parts) + + def _object_key(self, partition_dir: str, batch_id: str) -> str: + return f"{partition_dir}/{batch_id}.{self.codec.ext}" + + # -- write -------------------------------------------------------------- + + def write(self, records: Iterable[Any]) -> StorageManifest: + """Persist records, grouped by partition, one object per (batch, partition). + + ``write`` is append-safe: each call emits new objects keyed by a + monotonic batch id, so concurrent writers do not clobber one another. + """ + buckets: dict[str, list[dict[str, Any]]] = {} + for rec in records: + d = _record_to_dict(rec) + buckets.setdefault(self._partition_dir(d), []).append(d) + + partition_dirs: list[str] = [] + object_keys: list[str] = [] + total = 0 + for pdir, rows in sorted(buckets.items()): + partition_dirs.append(pdir + "/") + batch_id = f"batch-{int(time.time() * 1_000_000)}" + key = self._object_key(pdir, batch_id) + self.backend.write_bytes(key, self.codec.dumps(rows)) + object_keys.append(key) + total += len(rows) + return StorageManifest(written=total, partitions=partition_dirs, objects=object_keys) + + # -- indexed query ------------------------------------------------------ + + def _candidate_dirs( + self, + subject_id: str | None, + user_id: str | None, + start: datetime | None, + end: datetime | None, + ) -> list[str]: + """Return the partition directories a query must descend into. + + Implements partition pruning: exact-match dimensions narrow to one + child, unspecified dimensions expand to all children (via + ``backend.list_dirs``), and the time range narrows to the specific + ``dt=`` buckets it spans. + """ + bounds = {"subject_id": subject_id, "user_id": user_id} + candidates: list[str] = [self.namespace] if self.namespace else [""] + for dim in self.partition_dims: + value = bounds.get(dim) + if value is not None: + candidates = [_join(c, f"{dim}={_segment(value)}/") for c in candidates] + else: + expanded: list[str] = [] + for c in candidates: + expanded.extend(_join(c, d) for d in self.backend.list_dirs(c)) + candidates = expanded + # Time bucket level. + buckets = _date_buckets(start, end, self.granularity) + if buckets is not None: + candidates = [_join(c, f"dt={b}/") for c in candidates for b in buckets] + else: + expanded = [] + for c in candidates: + expanded.extend( + _join(c, d) for d in self.backend.list_dirs(c) if d.startswith("dt=") + ) + candidates = expanded + return candidates + + def query( + self, + *, + subject_id: str | None = None, + user_id: str | None = None, + start: datetime | str | None = None, + end: datetime | str | None = None, + rollout_id: str | None = None, + limit: int | None = None, + ) -> list[dict[str, Any]]: + """Indexed retrieval of stored records. + + Partitions are pruned by ``subject_id`` / ``user_id`` / ``[start, end]`` + before any object is read; remaining records are filtered in-memory by + the full predicate and ordered by timestamp (ascending). ``rollout_id`` + supports deterministic single-trace reproduction. + """ + start_dt = _parse_time(start) if not isinstance(start, datetime) else start + end_dt = _parse_time(end) if not isinstance(end, datetime) else end + + results: list[dict[str, Any]] = [] + for cdir in self._candidate_dirs(subject_id, user_id, start_dt, end_dt): + for key in self.backend.list_keys(cdir): + rows = self.codec.loads(self.backend.read_bytes(key)) + for row in rows: + if not self._matches(row, subject_id, user_id, start_dt, end_dt, rollout_id): + continue + results.append(row) + + results.sort(key=lambda r: _record_time(r, self.time_field)) + if limit is not None: + results = results[:limit] + return results + + @staticmethod + def _matches( + row: dict[str, Any], + subject_id: str | None, + user_id: str | None, + start: datetime | None, + end: datetime | None, + rollout_id: str | None, + ) -> bool: + if subject_id is not None and str(_extract_field(row, "subject_id")) != str(subject_id): + return False + if user_id is not None and str(_extract_field(row, "user_id")) != str(user_id): + return False + if rollout_id is not None and str(row.get("rollout_id")) != str(rollout_id): + return False + if start is not None or end is not None: + dt = _record_time(row, "timestamp") + if start is not None and dt < start: + return False + if end is not None and dt > end: + return False + return True + + # -- introspection ------------------------------------------------------ + + def partitions(self) -> list[str]: + """List all non-empty partition directories currently in the store. + + A directory whose objects were migrated/deleted (and is now empty) is + not reported as a live partition. + """ + candidates: list[str] = [self.namespace] if self.namespace else [""] + # Walk all partition levels (shard dims + dt). + levels = len(self.partition_dims) + 1 + for _ in range(levels): + expanded: list[str] = [] + for c in candidates: + expanded.extend(_join(c, d) for d in self.backend.list_dirs(c)) + candidates = expanded + return [c for c in candidates if self.backend.list_keys(c)] + + def count(self) -> int: + """Total record count across all partitions.""" + total = 0 + for cdir in self.partitions(): + for key in self.backend.list_keys(cdir): + total += len(self.codec.loads(self.backend.read_bytes(key))) + return total + + def get(self, rollout_id: str) -> list[dict[str, Any]]: + """Deterministic reproduction fetch — all records for one rollout id.""" + return self.query(rollout_id=rollout_id) + + +def open_trace_storage( + root: str | Path, + *, + codec: str | TraceCodec = "jsonl", + **kwargs: Any, +) -> TraceStorage: + """Convenience factory: a ``TraceStorage`` over a local directory. + + ``codec`` may be ``"jsonl"`` (default), ``"parquet"``, or a ``TraceCodec`` + instance. + """ + if isinstance(codec, str): + codec = ParquetCodec() if codec == "parquet" else JsonLinesCodec() + return TraceStorage(LocalBackend(root), codec=codec, **kwargs) + + +__all__ = [ + "GCSBackend", + "JsonLinesCodec", + "LocalBackend", + "ParquetCodec", + "S3Backend", + "StorageBackend", + "StorageManifest", + "TraceCodec", + "TraceStorage", + "open_trace_storage", +] diff --git a/evomerge/trust_score.py b/evomerge/trust_score.py index 51cead2..a9f0039 100644 --- a/evomerge/trust_score.py +++ b/evomerge/trust_score.py @@ -130,6 +130,50 @@ def _geometric_mean(values: list[float]) -> float: return math.exp(log_sum / len(values)) +def _trace_to_dict(trace: Any) -> dict[str, Any]: + """Normalise a historical trace (pydantic model or dict) to a dict.""" + if isinstance(trace, dict): + return trace + if hasattr(trace, "model_dump"): + return trace.model_dump(mode="json") + if hasattr(trace, "dict"): + return trace.dict() + return dict(trace) + + +def historical_consistency(traces: Any) -> float | None: + """Reduce a set of retrieved historical traces to a [0, 1] consistency score. + + Each trace contributes a per-run trustworthiness value: + - ``objective_status == "pass"`` → 1.0, ``"fail"`` → 0.0, otherwise + - ``objective_score`` (a number) clamped to [0, 1]. + + Returns the mean over all scored traces, or ``None`` when no trace carried + a usable signal (so the dimension is recorded as *unknown* rather than a + misleading 0.0 that would collapse the geometric mean). + """ + traces = list(traces) + if not traces: + return None + scores: list[float] = [] + for trace in traces: + rec = _trace_to_dict(trace) + status = rec.get("objective_status") + if status == "pass": + scores.append(1.0) + elif status == "fail": + scores.append(0.0) + else: + score = rec.get("objective_score") + if isinstance(score, bool): + scores.append(1.0 if score else 0.0) + elif isinstance(score, (int, float)): + scores.append(max(0.0, min(1.0, float(score)))) + if not scores: + return None + return sum(scores) / len(scores) + + def _verify_receipt_digest(receipt_path: Path) -> bool: """Re-compute the receipt's own canonical SHA-256 and compare to receipt_digest field. @@ -378,6 +422,39 @@ def add_training_telemetry(self, telemetry: Any) -> AgentTrustScoreBuilder: self._dims["training_health"] = max(0.0, min(1.0, health)) return self + def add_historical_traces(self, traces: Any) -> AgentTrustScoreBuilder: + """Fold retrieved historical traces into the score as ``historical_consistency``. + + Closes the retrieval → trust-score loop (Milestone 5, issue #53): traces + materialised from the archival/retrieval layer + (``evomerge.storage.TraceStorage.query``) are reduced to a [0, 1] + consistency dimension and enter the geometric mean, so a subject whose + historical runs mostly succeeded lifts the score and one with a poor + track record does not. + + Accepts: + - a plain ``int``/``float`` in [0, 1] (used directly as the consistency), or + - an iterable of trace records (dicts or pydantic models), e.g. the + output of ``TraceStorage.query(...)``, scored via + :func:`historical_consistency`. An empty / signal-less iterable + records the dimension as *None* (unknown). + """ + if isinstance(traces, bool): + # bool is an int subclass — treat as an explicit 0/1 consistency. + consistency: float | None = 1.0 if traces else 0.0 + elif isinstance(traces, (int, float)): + consistency = float(traces) + else: + consistency = historical_consistency(traces) + if consistency is None: + self._dims["historical_consistency"] = None + self._notes.append( + "historical_consistency: no scored historical traces — dimension unknown" + ) + else: + self._dims["historical_consistency"] = max(0.0, min(1.0, consistency)) + return self + def build(self) -> AgentTrustScore: # Only include non-None dimensions in the geometric mean known_values = [v for v in self._dims.values() if v is not None] diff --git a/pyproject.toml b/pyproject.toml index c816246..312e601 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,10 @@ dependencies = [ dev = ["pytest>=8", "ruff>=0.5", "scipy>=1.11", "sympy>=1.12"] ml = ["scikit-learn>=1.3", "joblib>=1.3"] train = ["scikit-learn>=1.3", "joblib>=1.3", "peft>=0.10", "trl>=0.8", "transformers>=4.40", "bitsandbytes>=0.43"] +# Archival / retrieval layer (Milestone 5, issue #53): enables the genuine +# columnar Parquet codec (evomerge.storage.ParquetCodec). JSONL is the default +# codec and needs no extra deps, so the base/dev install stays dependency-light. +arch = ["pyarrow>=14.0"] [project.scripts] evomerge = "evomerge.__main__:main" diff --git a/tests/test_archival.py b/tests/test_archival.py new file mode 100644 index 0000000..3cbbb72 --- /dev/null +++ b/tests/test_archival.py @@ -0,0 +1,205 @@ +"""Tests for evomerge.archival — Milestone 5 retention & tiered cold-storage (issue #53). + +Covers the two archival sub-requirements, each in its own test class: + + - TestRetentionPolicy — retention schedule thresholds + validation + - TestPartitionAge — partition age derivation from ``dt=`` buckets + - TestArchivalStore — hot → cold migration, deletion, dry-run planning +""" +from __future__ import annotations + +from datetime import datetime, timezone + +import pytest + +from evomerge.archival import ( + ArchivalStore, + RetentionAction, + RetentionPolicy, + RetentionReport, + partition_age_days, +) +from evomerge.storage import JsonLinesCodec, LocalBackend, ParquetCodec, TraceStorage + +NOW = datetime(2026, 7, 28, tzinfo=timezone.utc) + + +def _trace(rollout_id, subject_id="s1", ts="2026-07-01T00:00:00Z", status="pass", score=1): + return { + "rollout_id": rollout_id, + "subject_id": subject_id, + "timestamp": ts, + "objective_status": status, + "objective_score": score, + } + + +@pytest.fixture +def hot_cold(tmp_path): + hot = TraceStorage(LocalBackend(tmp_path / "hot"), granularity="day") + cold = TraceStorage(LocalBackend(tmp_path / "cold"), granularity="day") + return ArchivalStore(hot, cold) + + +# --------------------------------------------------------------------------- +# Retention policy +# --------------------------------------------------------------------------- + +class TestRetentionPolicy: + def test_keep_below_hot_days(self): + policy = RetentionPolicy(hot_days=30, delete_after_days=365) + action, reason = policy.action_for(5) + assert action == "keep" + assert "5.0d" in reason + + def test_migrate_at_hot_days(self): + policy = RetentionPolicy(hot_days=30, delete_after_days=365) + action, _ = policy.action_for(40) + assert action == "migrate" + + def test_delete_at_delete_after_days(self): + policy = RetentionPolicy(hot_days=30, delete_after_days=365) + action, _ = policy.action_for(400) + assert action == "delete" + + def test_delete_takes_precedence_over_migrate(self): + # A partition older than delete_after is deleted, not migrated. + policy = RetentionPolicy(hot_days=30, delete_after_days=60) + action, _ = policy.action_for(90) + assert action == "delete" + + def test_retain_forever_when_delete_after_none(self): + policy = RetentionPolicy(hot_days=30, delete_after_days=None) + action, _ = policy.action_for(10_000) + assert action == "migrate" + + def test_reject_negative_hot_days(self): + with pytest.raises(ValueError, match="hot_days"): + RetentionPolicy(hot_days=-1) + + def test_reject_delete_after_not_greater_than_hot(self): + with pytest.raises(ValueError, match="greater than hot_days"): + RetentionPolicy(hot_days=30, delete_after_days=30) + with pytest.raises(ValueError, match="greater than hot_days"): + RetentionPolicy(hot_days=30, delete_after_days=10) + + def test_zero_hot_days_migrates_immediately(self): + policy = RetentionPolicy(hot_days=0) + action, _ = policy.action_for(0.0) + assert action == "migrate" + + +# --------------------------------------------------------------------------- +# Partition age +# --------------------------------------------------------------------------- + +class TestPartitionAge: + def test_day_granularity_age(self): + age = partition_age_days("traces/subject_id=s1/dt=2026-07-01/", "day", now=NOW) + assert age == 27.0 # 2026-07-01 → 2026-07-28 + + def test_month_granularity_age_from_first_of_month(self): + age = partition_age_days("traces/dt=2026-06/", "month", now=NOW) + # 2026-06-01 → 2026-07-28 = 57 days + assert age == 57.0 + + def test_missing_dt_returns_inf(self): + assert partition_age_days("traces/subject_id=s1/", "day", now=NOW) == float("inf") + + def test_malformed_dt_returns_inf(self): + assert partition_age_days("traces/dt=not-a-date/", "day", now=NOW) == float("inf") + + +# --------------------------------------------------------------------------- +# Archival store +# --------------------------------------------------------------------------- + +class TestArchivalStore: + def test_plan_identifies_aged_partitions(self, hot_cold): + hot_cold.hot.write([ + _trace("old1", ts="2026-01-01T00:00:00Z"), + _trace("old2", ts="2026-01-02T00:00:00Z"), + _trace("new1", ts="2026-07-27T00:00:00Z"), + ]) + policy = RetentionPolicy(hot_days=30, delete_after_days=365) + report = hot_cold.plan(policy, now=NOW) + actions = {a.partition: a.action for a in report.actions} + assert actions["traces/subject_id=s1/dt=2026-01-01/"] == "migrate" + assert actions["traces/subject_id=s1/dt=2026-01-02/"] == "migrate" + assert actions["traces/subject_id=s1/dt=2026-07-27/"] == "keep" + + def test_apply_migrates_aged_to_cold(self, hot_cold): + hot_cold.hot.write([ + _trace("old1", ts="2026-01-01T00:00:00Z"), + _trace("new1", ts="2026-07-27T00:00:00Z"), + ]) + policy = RetentionPolicy(hot_days=30, delete_after_days=365) + report = hot_cold.apply(policy, now=NOW) + assert len(report.migrated) == 1 + # Hot retains only the recent partition. + assert hot_cold.hot.count() == 1 + assert [r["rollout_id"] for r in hot_cold.hot.query()] == ["new1"] + # Cold holds the migrated record. + assert hot_cold.cold.count() == 1 + assert [r["rollout_id"] for r in hot_cold.cold.query()] == ["old1"] + + def test_apply_deletes_expired_from_cold(self, hot_cold): + # 2026-04-01 at NOW (2026-07-28) = 118 days: migrates (>=30) but is not + # yet expired (<365). At a later date it ages past delete_after → deleted. + hot_cold.hot.write([_trace("aging", ts="2026-04-01T00:00:00Z")]) + policy = RetentionPolicy(hot_days=30, delete_after_days=365) + # First pass: migrates to cold (age 118d is between hot and delete thresholds). + hot_cold.apply(policy, now=NOW) + assert hot_cold.cold.count() == 1 + # Second pass at a later date: now exceeds delete_after → deleted. + later = datetime(2027, 7, 28, tzinfo=timezone.utc) + report = hot_cold.apply(policy, now=later) + assert len(report.deleted) == 1 + assert hot_cold.cold.count() == 0 + + def test_dry_run_has_no_side_effects(self, hot_cold): + hot_cold.hot.write([_trace("old1", ts="2026-01-01T00:00:00Z")]) + policy = RetentionPolicy(hot_days=30) + report = hot_cold.apply(policy, now=NOW, dry_run=True) + assert len(report.migrated) == 1 + # Nothing actually moved. + assert hot_cold.hot.count() == 1 + assert hot_cold.cold.count() == 0 + + def test_migrate_partition_re_encodes_across_codec(self, tmp_path): + # Hot uses JSONL; cold uses Parquet — migration re-encodes. + pytest.importorskip("pyarrow") + hot = TraceStorage(LocalBackend(tmp_path / "hot"), codec=JsonLinesCodec(), granularity="day") + cold = TraceStorage(LocalBackend(tmp_path / "cold"), codec=ParquetCodec(), granularity="day") + arch = ArchivalStore(hot, cold) + hot.write([_trace("old1", ts="2026-01-01T00:00:00Z")]) + moved = arch.migrate_partition("traces/subject_id=s1/dt=2026-01-01/") + assert moved == 1 + assert hot.count() == 0 + assert cold.count() == 1 + # Cold objects are genuine Parquet. + cold_keys = [k for p in cold.partitions() for k in cold.backend.list_keys(p)] + assert all(k.endswith(".parquet") for k in cold_keys) + + def test_migrate_partition_idempotent(self, hot_cold): + hot_cold.hot.write([_trace("old1", ts="2026-01-01T00:00:00Z")]) + first = hot_cold.migrate_partition("traces/subject_id=s1/dt=2026-01-01/") + second = hot_cold.migrate_partition("traces/subject_id=s1/dt=2026-01-01/") + assert first == 1 + assert second == 0 # already migrated, nothing to move + + def test_granularity_mismatch_rejected(self, tmp_path): + hot = TraceStorage(LocalBackend(tmp_path / "h"), granularity="day") + cold = TraceStorage(LocalBackend(tmp_path / "c"), granularity="hour") + with pytest.raises(ValueError, match="granularity"): + ArchivalStore(hot, cold) + + def test_report_action_types(self, hot_cold): + # Structural check: RetentionAction / RetentionReport expose the API. + hot_cold.hot.write([_trace("new1", ts="2026-07-27T00:00:00Z")]) + report = hot_cold.plan(RetentionPolicy(hot_days=30), now=NOW) + assert isinstance(report, RetentionReport) + assert all(isinstance(a, RetentionAction) for a in report.actions) + assert report.kept == [a for a in report.actions if a.action == "keep"] + assert report.migrated == [] + assert report.deleted == [] diff --git a/tests/test_storage.py b/tests/test_storage.py new file mode 100644 index 0000000..f870dc2 --- /dev/null +++ b/tests/test_storage.py @@ -0,0 +1,334 @@ +"""Tests for evomerge.storage — Milestone 5 trace archival & retrieval (issue #53). + +Covers the storage-layer sub-requirements, each in its own test class: + + - TestPartitioning — time-partitioned layout (Hive-style subject_id=/dt= dirs) + - TestCodecs — JsonLinesCodec round-trip + ParquetCodec (pyarrow opt-in) + - TestIndexedQuery — audit-trail / reproduction / historical-trust retrieval + - TestBackendAdapter — S3Backend routing via an injected fake client + - TestStorageTrustLoop — end-to-end: query() feeds add_historical_traces() +""" +from __future__ import annotations + +import builtins + +import pytest + +from evomerge.storage import ( + GCSBackend, + JsonLinesCodec, + LocalBackend, + ParquetCodec, + S3Backend, + StorageBackend, + StorageManifest, + TraceStorage, + open_trace_storage, +) + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +def _trace(rollout_id, subject_id="s1", user_id="u1", ts="2026-07-01T10:00:00Z", + status="pass", score=1): + return { + "rollout_id": rollout_id, + "subject_id": subject_id, + "user_id": user_id, + "timestamp": ts, + "objective_status": status, + "objective_score": score, + } + + +@pytest.fixture +def store(tmp_path): + return TraceStorage(LocalBackend(tmp_path), granularity="day") + + +# --------------------------------------------------------------------------- +# Partitioning +# --------------------------------------------------------------------------- + +class TestPartitioning: + def test_records_partitioned_by_subject_and_day(self, store): + recs = [ + _trace("r1", "s1", ts="2026-07-01T10:00:00Z"), + _trace("r2", "s1", ts="2026-07-02T10:00:00Z"), + _trace("r3", "s2", ts="2026-07-01T11:00:00Z"), + ] + manifest = store.write(recs) + assert isinstance(manifest, StorageManifest) + assert manifest.written == 3 + assert manifest.n_partitions == 3 + parts = set(store.partitions()) + assert "traces/subject_id=s1/dt=2026-07-01/" in parts + assert "traces/subject_id=s1/dt=2026-07-02/" in parts + assert "traces/subject_id=s2/dt=2026-07-01/" in parts + + def test_unknown_shard_key_bucketed_not_dropped(self, store): + # A record missing subject_id must still be retrievable. + rec = {"rollout_id": "rx", "timestamp": "2026-07-01T00:00:00Z"} + store.write([rec]) + out = store.query() + assert [r["rollout_id"] for r in out] == ["rx"] + assert any("subject_id=_unknown" in p for p in store.partitions()) + + def test_granularity_variants(self, tmp_path): + for gran, expected in [("hour", "2026-07-01T10"), ("month", "2026-07")]: + s = TraceStorage(LocalBackend(tmp_path / gran), granularity=gran) + s.write([_trace("r1", ts="2026-07-01T10:30:00Z")]) + assert any(f"dt={expected}" in p for p in s.partitions()) + + def test_namespace_layout(self, tmp_path): + s = TraceStorage(LocalBackend(tmp_path), namespace="archive", granularity="day") + s.write([_trace("r1", ts="2026-07-01T00:00:00Z")]) + assert all(p.startswith("archive/") for p in s.partitions()) + + def test_invalid_granularity_rejected(self, tmp_path): + with pytest.raises(ValueError, match="granularity"): + TraceStorage(LocalBackend(tmp_path), granularity="week") + + def test_invalid_partition_dim_rejected(self, tmp_path): + with pytest.raises(ValueError, match="partition_by"): + TraceStorage(LocalBackend(tmp_path), partition_by=("session_id",)) + + +# --------------------------------------------------------------------------- +# Codecs +# --------------------------------------------------------------------------- + +class TestCodecs: + def test_jsonl_codec_round_trip(self, tmp_path): + codec = JsonLinesCodec() + rows = [{"a": 1, "b": "x"}, {"a": 2, "b": "y", "c": [1, 2]}] + assert codec.loads(codec.dumps(rows)) == rows + assert codec.loads(b"") == [] + assert isinstance(LocalBackend(tmp_path), StorageBackend) + + def test_parquet_codec_round_trip(self): + pytest.importorskip("pyarrow") + codec = ParquetCodec() + rows = [ + {"rollout_id": "r1", "subject_id": "s1", "score": 1}, + {"rollout_id": "r2", "subject_id": "s1", "score": 0}, + ] + out = codec.loads(codec.dumps(rows)) + assert out == rows + + def test_parquet_codec_missing_dependency_message(self, monkeypatch): + real_import = builtins.__import__ + + def fake_import(name, *args, **kwargs): + if name in ("pyarrow", "pyarrow.parquet"): + raise ImportError("simulated missing") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", fake_import) + with pytest.raises(ImportError, match=r"evomerge\[arch\]"): + ParquetCodec().dumps([{"a": 1}]) + + def test_storage_with_parquet_codec(self, tmp_path): + pytest.importorskip("pyarrow") + s = TraceStorage(LocalBackend(tmp_path), codec=ParquetCodec(), granularity="day") + manifest = s.write([_trace("r1", ts="2026-07-01T00:00:00Z")]) + assert all(key.endswith(".parquet") for key in manifest.objects) + out = s.query(subject_id="s1") + assert [r["rollout_id"] for r in out] == ["r1"] + + +# --------------------------------------------------------------------------- +# Indexed query +# --------------------------------------------------------------------------- + +class TestIndexedQuery: + @pytest.fixture + def populated(self, store): + store.write([ + _trace("r1", "s1", user_id="u1", ts="2026-07-01T10:00:00Z"), + _trace("r2", "s1", user_id="u1", ts="2026-07-02T10:00:00Z", status="fail", score=0), + _trace("r3", "s2", user_id="u2", ts="2026-07-01T11:00:00Z"), + ]) + return store + + def test_audit_trail_by_subject(self, populated): + ids = [r["rollout_id"] for r in populated.query(subject_id="s1")] + assert ids == ["r1", "r2"] + + def test_audit_trail_by_subject_and_date_range(self, populated): + out = populated.query( + subject_id="s1", + start="2026-07-02T00:00:00Z", + end="2026-07-02T23:59:59Z", + ) + assert [r["rollout_id"] for r in out] == ["r2"] + + def test_date_range_spanning_multiple_buckets(self, populated): + out = populated.query(start="2026-07-01T00:00:00Z", end="2026-07-02T23:59:59Z") + assert {r["rollout_id"] for r in out} == {"r1", "r2", "r3"} + + def test_reproduction_by_rollout_id(self, populated): + out = populated.get("r3") + assert [r["rollout_id"] for r in out] == ["r3"] + + def test_partition_pruning_skips_other_subjects(self, populated): + # Querying s1 must never read s2's partition: r3 absent. + out = populated.query(subject_id="s1") + assert "r3" not in {r["rollout_id"] for r in out} + + def test_user_id_query(self, populated): + out = populated.query(user_id="u1") + assert {r["rollout_id"] for r in out} == {"r1", "r2"} + + def test_results_ordered_by_timestamp(self, populated): + out = populated.query() + ts = [r["timestamp"] for r in out] + assert ts == sorted(ts) + + def test_limit_applied_after_order(self, populated): + out = populated.query(limit=2) + assert len(out) == 2 + + def test_count(self, populated): + assert populated.count() == 3 + + def test_empty_query_returns_empty(self, store): + assert store.query(subject_id="nobody") == [] + + +# --------------------------------------------------------------------------- +# Backend adapter (S3 via injected fake client — no boto3 required) +# --------------------------------------------------------------------------- + +class TestBackendAdapter: + def _fake_s3(self): + class FakeS3: + def __init__(self): + self.objects = {} # key -> bytes + + def put_object(self, *, Bucket, Key, Body): + self.objects[(Bucket, Key)] = Body + + def get_object(self, *, Bucket, Key): + return {"Body": _Reader(self.objects[(Bucket, Key)])} + + def head_object(self, *, Bucket, Key): + if (Bucket, Key) not in self.objects: + raise KeyError(Key) + return {} + + def delete_object(self, *, Bucket, Key): + self.objects.pop((Bucket, Key), None) + + def get_paginator(self, name): + assert name == "list_objects_v2" + return _Paginator(self) + + def list_objects_v2(self, *, Bucket, Prefix, Delimiter=None): + # Only the delimiter form (CommonPrefixes) is used by list_dirs. + prefix = Prefix + commons = set() + for (b, k) in self.objects: + if b != Bucket or not k.startswith(prefix): + continue + rest = k[len(prefix):] + if Delimiter == "/" and "/" in rest: + commons.add(prefix + rest.split("/", 1)[0] + "/") + return {"CommonPrefixes": [{"Prefix": c} for c in sorted(commons)]} + + class _Reader: + def __init__(self, data): + self._data = data + + def read(self): + return self._data + + class _Paginator: + def __init__(self, s3): + self.s3 = s3 + + def paginate(self, *, Bucket, Prefix): + keys = [k for (b, k) in self.s3.objects if b == Bucket and k.startswith(Prefix)] + yield {"Contents": [{"Key": k} for k in sorted(keys)]} + + return FakeS3() + + def test_s3_backend_routes_through_client(self): + client = self._fake_s3() + backend = S3Backend("mybucket", prefix="traces", client=client, not_found_exc=KeyError) + backend.write_bytes("a/b/c.jsonl", b"hello") + assert backend.read_bytes("a/b/c.jsonl") == b"hello" + assert backend.exists("a/b/c.jsonl") is True + assert backend.exists("a/b/missing.jsonl") is False + backend.delete("a/b/c.jsonl") + assert backend.list_keys("a/b/") == [] + + def test_s3_backend_list_keys_and_dirs(self): + client = self._fake_s3() + backend = S3Backend("mybucket", prefix="traces", client=client) + backend.write_bytes("subject_id=s1/dt=2026-07-01/b.jsonl", b"x") + backend.write_bytes("subject_id=s1/dt=2026-07-02/b.jsonl", b"x") + backend.write_bytes("subject_id=s2/dt=2026-07-01/b.jsonl", b"x") + # list_keys returns keys relative to the prefix. + all_keys = backend.list_keys("") + assert len(all_keys) == 3 + assert all(k.startswith("subject_id=") for k in all_keys) + # list_dirs at root returns the two subject partitions. + assert set(backend.list_dirs("")) == {"subject_id=s1/", "subject_id=s2/"} + # list_dirs under s1 returns the two date partitions. + assert set(backend.list_dirs("subject_id=s1/")) == { + "dt=2026-07-01/", "dt=2026-07-02/" + } + + def test_storage_over_s3_fake_backend(self): + client = self._fake_s3() + # The S3 prefix is the sole top-level here (namespace="") so backend + # prefix and storage namespace do not double up. + backend = S3Backend("mybucket", prefix="traces", client=client) + store = TraceStorage(backend, namespace="", granularity="day") + store.write([_trace("r1", "s1", ts="2026-07-01T00:00:00Z")]) + out = store.query(subject_id="s1") + assert [r["rollout_id"] for r in out] == ["r1"] + + def test_local_backend_is_storage_backend(self, tmp_path): + assert isinstance(LocalBackend(tmp_path), StorageBackend) + + def test_gcs_backend_construct(self): + # GCSBackend should be constructable with an injected client (no SDK). + class FakeBucket: + pass + + class FakeClient: + def bucket(self, name): + return FakeBucket() + + backend = GCSBackend("bucket", prefix="traces", client=FakeClient()) + assert backend.bucket_name == "bucket" + assert isinstance(backend, object) + + +# --------------------------------------------------------------------------- +# End-to-end: retrieval → historical trust +# --------------------------------------------------------------------------- + +class TestStorageTrustLoop: + def test_query_feeds_historical_trust(self, store): + store.write([ + _trace("r1", "s1", ts="2026-06-01T00:00:00Z", status="pass", score=1), + _trace("r2", "s1", ts="2026-06-02T00:00:00Z", status="pass", score=1), + _trace("r3", "s1", ts="2026-06-03T00:00:00Z", status="fail", score=0), + ]) + from evomerge.trust_score import AgentTrustScoreBuilder + + traces = store.query(subject_id="s1", start="2026-06-01T00:00:00Z", end="2026-06-30T00:00:00Z") + score = AgentTrustScoreBuilder().add_historical_traces(traces).build() + # 2 pass + 1 fail → consistency 2/3. + assert abs(score.breakdown["historical_consistency"] - 2 / 3) < 1e-9 + + +def test_open_trace_storage_factory(tmp_path): + s = open_trace_storage(tmp_path, codec="jsonl", granularity="day") + assert isinstance(s, TraceStorage) + s.write([_trace("r1", ts="2026-07-01T00:00:00Z")]) + assert s.count() == 1 diff --git a/tests/test_trust_score.py b/tests/test_trust_score.py index d897a06..1780e47 100644 --- a/tests/test_trust_score.py +++ b/tests/test_trust_score.py @@ -6,6 +6,7 @@ AgentTrustScoreBuilder, _geometric_mean, compute_trust_score, + historical_consistency, ) @@ -162,3 +163,64 @@ def test_replay_determinism_zero_collapses(): builder.add_replay_determinism(0.0) score = builder.build() assert score.overall == 0.0 + + +# --------------------------------------------------------------------------- +# Historical-traces hook (Milestone 5, issue #53 — retrieval → trust loop) +# --------------------------------------------------------------------------- + +def test_historical_consistency_mixed_signals(): + traces = [ + {"objective_status": "pass"}, + {"objective_status": "fail"}, + {"objective_score": 0.5}, + ] + # (1.0 + 0.0 + 0.5) / 3 + assert historical_consistency(traces) == pytest.approx(0.5) + + +def test_historical_consistency_empty_is_none(): + assert historical_consistency([]) is None + # No usable signal → None, not 0.0. + assert historical_consistency([{"objective_status": "unknown"}]) is None + + +def test_add_historical_traces_from_float(): + score = AgentTrustScoreBuilder().add_historical_traces(0.8).build() + assert score.breakdown["historical_consistency"] == pytest.approx(0.8) + + +def test_add_historical_traces_from_records(): + traces = [{"objective_status": "pass"}, {"objective_status": "pass"}, {"objective_status": "fail"}] + score = ( + AgentTrustScoreBuilder() + .add_task_success(True) + .add_historical_traces(traces) + .build() + ) + assert score.breakdown["historical_consistency"] == pytest.approx(2 / 3) + + +def test_add_historical_traces_empty_records_dimension_unknown(): + builder = AgentTrustScoreBuilder().add_historical_traces([]) + score = builder.build() + # Empty history → dimension recorded as None (unknown), not 0.0. + assert score.breakdown["historical_consistency"] is None + assert any("historical_consistency" in n for n in builder._notes) + + +def test_add_historical_traces_from_storage_query(tmp_path): + # End-to-end: archival retrieval output feeds the trust-score dimension. + from evomerge.storage import LocalBackend, TraceStorage + + store = TraceStorage(LocalBackend(tmp_path), granularity="day") + store.write([ + {"rollout_id": "r1", "subject_id": "s1", "timestamp": "2026-06-01T00:00:00Z", + "objective_status": "pass"}, + {"rollout_id": "r2", "subject_id": "s1", "timestamp": "2026-06-02T00:00:00Z", + "objective_status": "fail"}, + ]) + traces = store.query(subject_id="s1") + score = AgentTrustScoreBuilder().add_historical_traces(traces).build() + assert score.breakdown["historical_consistency"] == pytest.approx(0.5) +