AML pack 1/5: schema, generator, licence gate, table snapshots

labagent/cli.py

@@ -108,8 +108,9 @@ def main(argv=None) -> int:
     pk.add_argument("--json", action="store_true", help="Machine-readable catalog with parameters, platforms and cost tier (used by the console)")
     from .ai.rag import cli as rag_cli
     rag_cli.add_parser(sub)
-    dp = sub.add_parser("dataset", help="Package a chunk file as a versioned, checksummed snapshot (maintainers)")
-    dp.add_argument("--chunks", required=True, help="JSONL of chunks: one object per line with id, title, content and meta")
+    dp = sub.add_parser("dataset", help="Package a chunk file or a folder of tables as a versioned, checksummed snapshot (maintainers)")
+    dp.add_argument("--chunks", help="JSONL of chunks: one object per line with id, title, content and meta")
+    dp.add_argument("--tables", help="Folder of CSV/Parquet files with a provenance.json (for example from python -m labagent.packs.aml.gen)")
     dp.add_argument("--name", required=True)
     dp.add_argument("--version", required=True)
     dp.add_argument("--out", required=True, help="Output directory")
@@ -128,13 +129,19 @@ def main(argv=None) -> int:
         print(STATE_BACKEND_DIR)
         return 0
     if args.cmd == "dataset":
-        from .datasets import DatasetError, build_snapshot
+        from .datasets import DatasetError, build_snapshot, build_table_snapshot
+        if bool(args.chunks) == bool(args.tables):
+            print("error: give exactly one of --chunks or --tables", file=sys.stderr)
+            return 2
         try:
-            m = build_snapshot(Path(args.chunks), args.name, args.version, Path(args.out))
+            if args.tables:
+                m = build_table_snapshot(Path(args.tables), args.name, args.version, Path(args.out))
+            else:
+                m = build_snapshot(Path(args.chunks), args.name, args.version, Path(args.out))
         except DatasetError as e:
             print(f"error: {e}", file=sys.stderr)
             return 2
-        print(json.dumps({k: m[k] for k in ("file", "chunks", "bytes", "sha256")}, indent=2))
+        print(json.dumps({k: m[k] for k in ("file", "chunks", "bytes", "sha256") if k in m}, indent=2))
         return 0
     if args.cmd == "rag":
         return rag_cli.run(args)

labagent/datasets.py

@@ -10,7 +10,9 @@
 import hashlib
 import json
 import os
+import re
 import urllib.request
+import zipfile
 from datetime import date
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional
@@ -79,6 +81,8 @@ def fetch_snapshot(name: str, cache_dir: Optional[Path] = None, download: Callab
     entry = entry or registry().get(name)
     if not entry:
         raise DatasetError(f"unknown dataset '{name}'. Known: {', '.join(sorted(registry()))}")
+    if entry.get("kind", "chunks") != "chunks":
+        raise DatasetError(f"dataset '{name}' holds tables, not chunks: use fetch_table_snapshot")
     root = cache_dir or Path(os.environ.get("LABAGENT_HOME") or Path.home() / ".labagent") / "cache" / "datasets"
     root.mkdir(parents=True, exist_ok=True)
     dest = root / f"{entry['sha256'][:16]}-{Path(entry['url']).name}"
@@ -99,3 +103,79 @@ def fetch_snapshot(name: str, cache_dir: Optional[Path] = None, download: Callab
             raise DatasetError(f"could not obtain dataset '{name}' with a matching checksum ({'; '.join(errors)})")
     with gzip.open(dest, "rt", encoding="utf-8") as f:
         return [json.loads(l) for l in f if l.strip()]
+
+
+# ---- table snapshots (relational data such as the AML pack's CSV files) ---------------------------------------------------------
+# A table snapshot is one zip of flat files (CSV, Parquet, provenance.json) plus a manifest and SHA256SUMS. The zip is reproducible
+# byte for byte. Which sources may go in is decided by labagent.packs.aml.licences, not here: this function only asks.
+
+SAFE_NAME = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]*$")
+ZIP_TIME = (1980, 1, 1, 0, 0, 0)
+
+
+def build_table_snapshot(src_dir: Path, name: str, version: str, out_dir: Path, built: Optional[str] = None) -> Dict[str, Any]:
+    from .packs.aml.licences import assert_redistributable
+    src_dir = Path(src_dir)
+    prov_path = src_dir / "provenance.json"
+    if not prov_path.exists():
+        raise DatasetError("provenance.json is missing: a snapshot must say where its data came from")
+    prov = json.loads(prov_path.read_text())
+    try:
+        assert_redistributable([s["id"] for s in prov.get("sources", [])])
+    except Exception as e:
+        raise DatasetError(str(e)) from e
+    files = sorted(p for p in src_dir.iterdir() if p.is_file())
+    if not files:
+        raise DatasetError("no files to package")
+    for p in files:
+        if not SAFE_NAME.match(p.name):
+            raise DatasetError(f"unsafe file name: {p.name}")
+    out_dir.mkdir(parents=True, exist_ok=True)
+    stem = f"{name}-{version}"
+    data = out_dir / f"{stem}.tables.zip"
+    with zipfile.ZipFile(data, "w", zipfile.ZIP_DEFLATED, compresslevel=9) as z:
+        for p in files:
+            info = zipfile.ZipInfo(p.name, ZIP_TIME)
+            info.compress_type, info.external_attr = zipfile.ZIP_DEFLATED, 0o644 << 16
+            z.writestr(info, p.read_bytes(), compresslevel=9)
+    manifest = {
+        "name": name, "version": version, "kind": "tables", "built": built or date.today().isoformat(), "labagent": __version__,
+        "file": data.name, "sha256": sha256_file(data), "bytes": data.stat().st_size,
+        "files": {p.name: sha256_file(p) for p in files}, "sources": prov["sources"], "info": prov.get("info", {}),
+    }
+    (out_dir / f"{stem}.manifest.json").write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
+    (out_dir / "SHA256SUMS").write_text(f"{manifest['sha256']}  {data.name}\n")
+    return manifest
+
+
+def fetch_table_snapshot(name: str, dest: Path, cache_dir: Optional[Path] = None, download: Callable[[str, Path], None] = default_http,
+                         entry: Optional[Dict[str, Any]] = None) -> Path:
+    """Download a pinned table snapshot once, check its SHA-256 every time, and unpack the flat files into `dest`."""
+    entry = entry or registry().get(name)
+    if not entry or entry.get("kind") != "tables":
+        raise DatasetError(f"unknown table dataset '{name}'")
+    root = cache_dir or Path(os.environ.get("LABAGENT_HOME") or Path.home() / ".labagent") / "cache" / "datasets"
+    root.mkdir(parents=True, exist_ok=True)
+    zpath = root / f"{entry['sha256'][:16]}-{Path(entry['url']).name}"
+    if not zpath.exists() or sha256_file(zpath) != entry["sha256"]:
+        errors = []
+        for url in [entry["url"], *entry.get("mirrors", [])]:
+            tmp = zpath.with_suffix(".part")
+            try:
+                download(url, tmp)
+                if sha256_file(tmp) != entry["sha256"]:
+                    raise DatasetError("checksum mismatch")
+                tmp.replace(zpath)
+                break
+            except Exception as e:
+                tmp.unlink(missing_ok=True)
+                errors.append(f"{url.split('/')[2]}: {type(e).__name__}")
+        else:
+            raise DatasetError(f"could not obtain dataset '{name}' with a matching checksum ({'; '.join(errors)})")
+    dest.mkdir(parents=True, exist_ok=True)
+    with zipfile.ZipFile(zpath) as z:
+        for info in z.infolist():
+            if not SAFE_NAME.match(info.filename):         # no paths, no '..': a snapshot is a flat folder
+                raise DatasetError(f"refusing unsafe name in snapshot: {info.filename!r}")
+            (dest / info.filename).write_bytes(z.read(info))
+    return dest

labagent/packs/aml/__init__.py

@@ -0,0 +1 @@
+"""AML lab pack: a synthetic anti-money-laundering dataset with ground truth, plus the documents learners search over."""

labagent/packs/aml/data/sdn_seed.json

@@ -0,0 +1,67 @@
+{
+  "source": "OFAC Specially Designated Nationals list, https://www.treasury.gov/ofac/downloads/sdn.csv",
+  "list_date": "2026-10-21",
+  "note": "Entity names only (no individuals, addresses or identifiers). US government work, used to seed screening practice. Not a current list.",
+  "entries": [
+    {
+      "uid": "19730",
+      "name": "MAZAKA GENERAL TRADING L.L.C.",
+      "program": "TCO"
+    },
+    {
+      "uid": "19731",
+      "name": "JETLINK TEXTILES TRADING",
+      "program": "TCO"
+    },
+    {
+      "uid": "19732",
+      "name": "SEVEN SEA GOLDEN GENERAL TRADING LLC",
+      "program": "TCO"
+    },
+    {
+      "uid": "19733",
+      "name": "AYDAH TRADING LLC",
+      "program": "TCO"
+    },
+    {
+      "uid": "19734",
+      "name": "WADI AL AFRAH TRADING LLC",
+      "program": "TCO"
+    },
+    {
+      "uid": "51354",
+      "name": "CERES SHIPPING LIMITED",
+      "program": "IRAN-EO13902"
+    },
+    {
+      "uid": "51355",
+      "name": "EUNOMIA LIMITED",
+      "program": "IRAN-EO13902"
+    },
+    {
+      "uid": "51351",
+      "name": "LIGHTSHIP MANAGEMENT LTD",
+      "program": "IRAN-EO13902"
+    },
+    {
+      "uid": "51320",
+      "name": "YURIMAGUAS LTD",
+      "program": "IRAN-EO13902"
+    },
+    {
+      "uid": "11318",
+      "name": "NAFTIRAN INTERTRADE CO. LIMITED",
+      "program": "IRAN"
+    },
+    {
+      "uid": "11621",
+      "name": "INTRA CHEM TRADING GMBH",
+      "program": "IRAN"
+    },
+    {
+      "uid": "11613",
+      "name": "ASCOTEC STEEL TRADING GMBH",
+      "program": "IRAN"
+    }
+  ]
+}

labagent/packs/aml/export.py

@@ -0,0 +1,30 @@
+"""Write a generated dataset to a folder: CSV always, Parquet when pyarrow is installed, and provenance.json."""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, Dict, List
+
+from .gen import Dataset
+from .licences import SOURCES
+from .schema import write_csv, write_parquet
+
+PROVENANCE = "provenance.json"
+
+
+def provenance(ds: Dataset) -> Dict[str, Any]:
+    return {"sources": [{"id": s, "licence": SOURCES[s]["licence"], "mode": SOURCES[s]["mode"]} for s in ds.sources], "info": ds.info}
+
+
+def write_dataset(ds: Dataset, out_dir: Path, parquet: str = "auto") -> List[str]:
+    """parquet: 'auto' (write if pyarrow is importable), 'yes' (fail if it is not) or 'no'. Returns the file names written."""
+    out_dir = Path(out_dir)
+    files = [p.name for p in write_csv(ds.tables, out_dir)]
+    if parquet != "no":
+        try:
+            files += [p.name for p in write_parquet(ds.tables, out_dir)]
+        except ImportError:
+            if parquet == "yes":
+                raise
+    (out_dir / PROVENANCE).write_text(json.dumps(provenance(ds), indent=2, sort_keys=True) + "\n")
+    return files + [PROVENANCE]

labagent/packs/aml/gen/__init__.py

@@ -0,0 +1,28 @@
+"""Generators. Each one returns a `Dataset`: the seven tables in `schema.TABLES` plus a record of which sources it used."""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List
+
+from ..schema import Tables
+
+
+@dataclass
+class Dataset:
+    tables: Tables
+    sources: List[str]                                   # ids from licences.SOURCES
+    info: Dict[str, Any] = field(default_factory=dict)   # generator name, seed, size, version: all non-secret
+
+
+class Generator:
+    id = ""
+
+    def generate(self, **params: Any) -> Dataset:
+        raise NotImplementedError
+
+
+def get(generator_id: str) -> Generator:
+    if generator_id == "builtin":
+        from .builtin import Builtin
+        return Builtin()
+    raise KeyError(f"unknown generator '{generator_id}'. Available now: builtin")

labagent/packs/aml/gen/__main__.py

@@ -0,0 +1,27 @@
+"""python -m labagent.packs.aml.gen --out DIR [--seed 42] [--size small] [--jurisdiction ca] [--parquet auto]"""
+import argparse
+import json
+from pathlib import Path
+
+from ..export import write_dataset
+from ..schema import table_counts
+from . import get
+
+
+def main() -> int:
+    p = argparse.ArgumentParser(description="Generate the synthetic AML tables")
+    p.add_argument("--out", required=True)
+    p.add_argument("--generator", default="builtin")
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--size", default="small", choices=["tiny", "small", "medium"])
+    p.add_argument("--jurisdiction", default="ca", choices=["ca", "us", "both"])
+    p.add_argument("--parquet", default="auto", choices=["auto", "yes", "no"])
+    a = p.parse_args()
+    ds = get(a.generator).generate(seed=a.seed, size=a.size, jurisdiction=a.jurisdiction)
+    files = write_dataset(ds, Path(a.out), a.parquet)
+    print(json.dumps({"files": files, "rows": table_counts(ds.tables)}, indent=2))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())