diff --git a/labagent/cli.py b/labagent/cli.py index 57d0e99..483f9b8 100644 --- a/labagent/cli.py +++ b/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) diff --git a/labagent/datasets.py b/labagent/datasets.py index 9e6d97d..da7eb99 100644 --- a/labagent/datasets.py +++ b/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 diff --git a/labagent/packs/aml/__init__.py b/labagent/packs/aml/__init__.py new file mode 100644 index 0000000..398e59c --- /dev/null +++ b/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.""" diff --git a/labagent/packs/aml/data/sdn_seed.json b/labagent/packs/aml/data/sdn_seed.json new file mode 100644 index 0000000..8e3a9ba --- /dev/null +++ b/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" + } + ] +} diff --git a/labagent/packs/aml/export.py b/labagent/packs/aml/export.py new file mode 100644 index 0000000..d880e0a --- /dev/null +++ b/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] diff --git a/labagent/packs/aml/gen/__init__.py b/labagent/packs/aml/gen/__init__.py new file mode 100644 index 0000000..f825fb2 --- /dev/null +++ b/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") diff --git a/labagent/packs/aml/gen/__main__.py b/labagent/packs/aml/gen/__main__.py new file mode 100644 index 0000000..80e7ee2 --- /dev/null +++ b/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()) diff --git a/labagent/packs/aml/gen/builtin.py b/labagent/packs/aml/gen/builtin.py new file mode 100644 index 0000000..3b44831 --- /dev/null +++ b/labagent/packs/aml/gen/builtin.py @@ -0,0 +1,476 @@ +"""Pure-Python synthetic AML data. No Java, no network, deterministic for a given seed. + +Normal activity (payroll, rent, card spending, business sales and suppliers) is generated first. Then labelled cases are laid +over it: the 8 laundering patterns from Altman et al. (2023), two ways of disguising laundering as ordinary payments, name +matches against a fixed set of OFAC SDN entity names, look-alikes that deserve a closer look (monitor) and ordinary events +that trigger a rule for a harmless reason (dismiss). + +Labels are set by construction, not by an investigator. The dataset card says so. + +Only these random calls are used: random(), randint(), choice(), gauss(), shuffle(), sample(). They give the same numbers for +the same seed on every supported Python. +""" +from __future__ import annotations + +import difflib +import json +import random +from datetime import date, datetime, timedelta +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +from .. import names as N +from ..schema import EXTERNAL_PREFIX, LIST_SEP, PATTERNS, RULES, TABLES, Tables, validate +from . import Dataset, Generator + +SIZES = { # customers, laundering cases per kind, monitor cases, dismiss cases, sanctions seeds (exact, fuzzy) + "tiny": dict(customers=60, per_kind=1, monitor=4, dismiss=8, sanctions=(1, 1)), # for tests + "small": dict(customers=300, per_kind=2, monitor=16, dismiss=30, sanctions=(2, 4)), + "medium": dict(customers=2000, per_kind=12, monitor=120, dismiss=300, sanctions=(4, 8)), +} +PROSE = {"fan_out": "fan-out", "fan_in": "fan-in", "gather_scatter": "gather-scatter", "scatter_gather": "scatter-gather", "random": "random-walk"} +THRESHOLD = 10000 # reporting threshold, in the account currency, for cash and wires +DATA = Path(__file__).resolve().parent.parent / "data" +VERSION = "1" + + +def _iso(dt: datetime) -> str: + return dt.strftime("%Y-%m-%dT%H:%M:%S") + + +class _Build: + def __init__(self, seed: int, size: str, jurisdiction: str, base: date): + self.rng = random.Random(seed) + self.cfg = SIZES[size] + self.jur = jurisdiction + self.base = base + self.customers: List[Dict[str, Any]] = [] + self.accounts: List[Dict[str, Any]] = [] + self.tx: List[Dict[str, Any]] = [] # drafts: carry "_k" until final ids are assigned in time order + self.cases: List[Dict[str, Any]] = [] + self.sanctions: List[Dict[str, Any]] = [] + self.kyc: Dict[str, str] = {} + self.cur: Dict[str, str] = {} # account id -> currency + self.primary: Dict[str, str] = {} # customer id -> main account id + self.owner: Dict[str, str] = {} # account id -> customer id + self.amt: Dict[int, float] = {} + self._k = 0 + self._ext = 0 + self._illicit = 0 + + # ---- small helpers ------------------------------------------------------------------------------------------------------- + def when(self, day: float, hour: Optional[int] = None) -> datetime: + d = self.base + timedelta(days=int(day)) + h = self.rng.randint(8, 18) if hour is None else hour + return datetime(d.year, d.month, d.day, h, self.rng.randint(0, 59), self.rng.randint(0, 59)) + + def ext(self) -> str: + self._ext += 1 + return f"{EXTERNAL_PREFIX}{(self._ext % 400) + 1:05d}" # a shared pool of ordinary outside parties + + def illicit_ext(self) -> str: + self._illicit += 1 + return f"{EXTERNAL_PREFIX}9{self._illicit:04d}" + + def amount(self, mean: float, spread: float = 0.25) -> float: + return round(max(5.0, self.rng.gauss(mean, mean * spread)), 2) + + def add(self, ts: datetime, src: str, dst: str, amt: float, channel: str, desc: str, cur: Optional[str] = None) -> int: + self._k += 1 + self.amt[self._k] = round(amt, 2) + self.tx.append({"_k": self._k, "timestamp": _iso(ts), "from_account": src, "to_account": dst, "amount": round(amt, 2), + "currency": cur or self.cur.get(src) or self.cur.get(dst, "CAD"), "channel": channel, "description": desc}) + return self._k + + # ---- customers and accounts ---------------------------------------------------------------------------------------------- + def make_customers(self) -> None: + r = self.rng + for i in range(1, self.cfg["customers"] + 1): + cid = f"C{i:05d}" + if r.random() < 0.15: + stem, (industry, suffix, rev, cash) = r.choice(N.BUSINESS_STEMS), r.choice(N.BUSINESS_KINDS) + ctype, name, occ, monthly = "business", f"{stem} {suffix}", industry, rev * r.uniform(0.6, 1.4) + else: + (occ, income) = r.choice(N.OCCUPATIONS) + ctype, name, monthly, cash = "person", f"{r.choice(N.FIRST)} {r.choice(N.LAST)}", income * r.uniform(0.7, 1.3), False + country = r.choice(["CA", "US"]) if self.jur == "both" else self.jur.upper() + if r.random() < 0.06: + country = r.choice(N.FOREIGN) + risk = "low" + if ctype == "business" and cash: + risk = "medium" + if country in N.HIGH_RISK_COUNTRIES or r.random() < 0.05: + risk = "high" if r.random() < 0.5 else "medium" + years = r.randint(0, 9) * 365 + r.randint(30, 360) if r.random() > 0.08 else -r.randint(5, 300) # a few are new this year + onboard = self.base - timedelta(days=years) + self.customers.append({"id": cid, "name": name, "type": ctype, "occupation": occ, "country": country, "risk_rating": risk, + "expected_monthly_activity": int(round(monthly, -2)), "onboarding_date": onboard.isoformat(), + "_cash": cash}) + home = "CAD" if country == "CA" else "USD" + kinds = ["business chequing"] if ctype == "business" else ["chequing"] + (["savings"] if r.random() < 0.5 else []) + for k in kinds: + aid = f"ACC{len(self.accounts) + 1:06d}" + self.accounts.append({"id": aid, "customer_id": cid, "type": k, "open_date": onboard.isoformat(), "branch": r.choice(N.BRANCHES)}) + self.cur[aid] = home + self.owner[aid] = cid + self.primary.setdefault(cid, aid) + + # ---- normal activity ----------------------------------------------------------------------------------------------------- + def normal_activity(self) -> None: + r = self.rng + for c in self.customers: + acct, monthly = self.primary[c["id"]], float(c["expected_monthly_activity"]) + start = max(0, (date.fromisoformat(c["onboarding_date"]) - self.base).days) + employer = r.choice(N.EMPLOYERS) + for month in range(12): + m0 = month * 30 + if m0 + 29 < start: + continue + if c["type"] == "person": + for half in (0, 15): + self.add(self.when(m0 + half, 9), self.ext(), acct, self.amount(monthly / 2, 0.03), "e-transfer", f"PAYROLL {employer.upper()}") + self.add(self.when(m0 + 1, 10), acct, self.ext(), self.amount(monthly * 0.3, 0.02), "e-transfer", "RENT") + self.add(self.when(m0 + r.randint(3, 25)), acct, self.ext(), self.amount(monthly * 0.04, 0.4), "card", "UTILITY BILL") + for _ in range(r.randint(4, 9)): + what = r.choice(["GROCERY", "FUEL", "RESTAURANT", "PHARMACY", "ONLINE SHOP"]) + self.add(self.when(m0 + r.randint(0, 29)), acct, self.ext(), self.amount(monthly * 0.015, 0.6), "card", f"POS PURCHASE {what}") + if r.random() < 0.6: + self.add(self.when(m0 + r.randint(0, 29)), acct, self.ext(), self.amount(monthly * 0.03, 0.5), "cash", "ATM WITHDRAWAL") + if r.random() < 0.5: + self.add(self.when(m0 + r.randint(0, 29)), acct, self.ext(), self.amount(monthly * 0.02, 0.6), "e-transfer", "E-TRANSFER SENT") + else: + for _ in range(r.randint(6, 12)): + self.add(self.when(m0 + r.randint(0, 29)), self.ext(), acct, self.amount(monthly / 9, 0.5), "card", "CARD SETTLEMENT") + if c["_cash"]: + for _ in range(r.randint(2, 4)): + self.add(self.when(m0 + r.randint(0, 29)), self.ext(), acct, self.amount(min(monthly / 12, 6000), 0.4), "cash", "BRANCH CASH DEPOSIT") + for _ in range(r.randint(2, 5)): + self.add(self.when(m0 + r.randint(0, 29)), acct, self.ext(), self.amount(monthly / 7, 0.5), "wire" if r.random() < 0.3 else "e-transfer", + f"SUPPLIER INVOICE {r.randint(1000, 9999)}") + self.add(self.when(m0 + 27, 11), acct, self.ext(), self.amount(monthly * 0.25, 0.05), "e-transfer", "PAYROLL RUN") + + # ---- cases --------------------------------------------------------------------------------------------------------------- + def established(self, kind: Optional[str] = None) -> List[Dict[str, Any]]: + """Customers whose account was open before the data starts, so a case never happens before the account existed.""" + cut = self.base.isoformat() + return [c for c in self.customers if c["onboarding_date"] < cut and (kind is None or c["type"] == kind)] + + def mule_pool(self) -> List[str]: + """Accounts whose owners expect little activity, so a large flow does not match their profile.""" + low = sorted(self.customers, key=lambda c: (c["expected_monthly_activity"], c["id"])) + low = [c for c in low if c["type"] == "person" and c["onboarding_date"] < self.base.isoformat()][: max(12, int(len(self.customers) * 0.4))] + pool = [self.primary[c["id"]] for c in low] + self.rng.shuffle(pool) + return pool + + def take(self, pool: List[str], k: int) -> List[str]: + out = [] + for _ in range(k): + out.append(pool.pop(0)) + pool.append(out[-1]) # reuse across cases only after the whole pool has been used + return out + + def mix(self, pool: List[str], k: int, internal_share: float = 0.6) -> List[str]: + n_in = max(1, int(round(k * internal_share))) + accts = self.take(pool, n_in) + [self.illicit_ext() for _ in range(k - n_in)] + self.rng.shuffle(accts) + return accts + + def case(self, label: str, pattern: str, rule: str, key_account: str, involved: List[str], keys: List[int], reason: str, + amount: float, day: float, score: Tuple[int, int]) -> None: + self.cases.append({"label": label, "pattern": pattern, "rule": rule, "account": key_account, "involved": involved, "keys": keys, + "reason": reason, "amount": round(amount, 2), "day": int(day), "score": self.rng.randint(*score)}) + + def expected_of(self, acct: str) -> float: + c = next(x for x in self.customers if x["id"] == self.owner[acct]) + return float(c["expected_monthly_activity"]) + + def launder(self, pattern: str, pool: List[str]) -> None: + r = self.rng + total = float(r.randint(18, 80) * 1000) + t0 = r.randint(20, 300) + keys: List[int] = [] + ch = lambda: r.choice(["wire", "e-transfer", "e-transfer"]) + legs: List[Tuple[int, str, str, float]] = [] # (day offset, from, to, amount) + rule, key_account = "R02", "" + + if pattern == "fan_out": + k = r.randint(4, 8) + src, hub = self.illicit_ext(), self.take(pool, 1)[0] + outs = self.mix(pool, k) + legs.append((0, src, hub, total)) + legs += [(1 + (i % 2), hub, o, total / k * r.uniform(0.9, 0.99)) for i, o in enumerate(outs)] + rule, key_account = "R04", hub + elif pattern == "fan_in": + k = r.randint(4, 8) + ins, hub = self.mix(pool, k), self.take(pool, 1)[0] + legs += [(i % 3, s, hub, total / k * r.uniform(0.9, 1.05)) for i, s in enumerate(ins)] + rule, key_account = "R03", hub + elif pattern == "gather_scatter": + k = r.randint(3, 5) + ins, hub, outs = self.mix(pool, k), self.take(pool, 1)[0], self.mix(pool, k) + legs += [(i % 2, s, hub, total / k) for i, s in enumerate(ins)] + legs += [(3 + i % 2, hub, o, total / k * 0.97) for i, o in enumerate(outs)] + rule, key_account = "R02", hub + elif pattern == "scatter_gather": + k = r.randint(3, 5) + origin, mids, sink = self.take(pool, 1)[0], self.mix(pool, k), self.illicit_ext() + legs += [(i % 2, origin, m, total / k) for i, m in enumerate(mids)] + legs += [(2 + i % 2, m, sink, total / k * 0.97) for i, m in enumerate(mids)] + rule, key_account = "R02", origin + elif pattern == "cycle": + ring = self.take(pool, r.randint(3, 5)) + amt = total + for i, a in enumerate(ring): + b = ring[(i + 1) % len(ring)] + legs.append((i, a, b, amt)) + amt *= r.uniform(0.97, 0.995) + rule, key_account = "R05", ring[0] + elif pattern == "random": + walk = self.mix(pool, r.randint(5, 8), 0.7) + amt = total + for i in range(len(walk) - 1): + legs.append((i, walk[i], walk[i + 1], amt)) + amt *= r.uniform(0.95, 0.99) + rule = "R02" + key_account = next((a for a in walk if a in self.owner), self.take(pool, 1)[0]) + elif pattern == "bipartite": + ins, outs = self.mix(pool, r.randint(2, 4), 0.7), self.mix(pool, r.randint(2, 4), 0.5) + legs += [(j % 3, s, o, total / (len(ins) * len(outs))) for j, (s, o) in enumerate((s, o) for s in ins for o in outs)] + rule = "R02" + key_account = next((a for a in ins + outs if a in self.owner), self.take(pool, 1)[0]) + elif pattern == "stack": + l1, l2, l3 = self.mix(pool, 2, 0.7), self.mix(pool, 3, 0.6), self.mix(pool, 3, 0.4) + legs += [(j % 2, s, o, total / 6) for j, (s, o) in enumerate((s, o) for s in l1 for o in l2)] + legs += [(3 + j % 2, s, o, total / 9 * 0.97) for j, (s, o) in enumerate((s, o) for s in l2 for o in l3)] + rule = "R02" + key_account = next((a for a in l1 + l2 + l3 if a in self.owner), self.take(pool, 1)[0]) + else: + raise ValueError(pattern) + + involved: List[str] = [] + for off, s, d, a in legs: + if s == d: + continue + channel = "cash" if (s.startswith(EXTERNAL_PREFIX) and pattern == "fan_out" and off == 0 and total < 60000) else ch() + keys.append(self.add(self.when(t0 + off), s, d, a, channel, self.cover_story(channel))) + for x in (s, d): + if x not in involved: + involved.append(x) + span = max(o for o, *_ in legs) + 1 + moved = sum(a for *_, a in legs) + ratio = moved / max(self.expected_of(key_account), 1.0) + self.case("escalate", pattern, rule, key_account, involved, keys, + f"{len(keys)} transfers moved about {int(moved):,} across {len(involved)} accounts in {span} days, in a {PROSE.get(pattern, pattern)} shape. " + f"That is {ratio:.1f} times the monthly activity this customer's profile expects, and most counterparties are new.", + moved, t0 + span, (72, 96)) + + def cover_story(self, channel: str) -> str: + r = self.rng + if channel == "cash": + return "BRANCH CASH DEPOSIT" + return r.choice(["E-TRANSFER SENT", "PERSONAL TRANSFER", "FAMILY SUPPORT", "WIRE TRANSFER", "INVOICE", "LOAN REPAYMENT"]) + + def disguised(self, kind: str, pool: List[str]) -> None: + """A front business collects 'sales' from mule accounts, then pays them out as payroll or supplies.""" + r = self.rng + fronts = self.established("business") + front_c = r.choice(fronts) + front, t0 = self.primary[front_c["id"]], r.randint(30, 280) + payers, keys = self.take(pool, r.randint(4, 6) if kind == "disguised_supplies" else r.randint(3, 5)), [] + involved = [front] + payers + collected = 0.0 + for i, p in enumerate(payers): + a = round(r.uniform(6, 14) * 1000, 2) + collected += a + keys.append(self.add(self.when(t0 + i), p, front, a, "e-transfer", f"INVOICE PAYMENT {r.randint(1000, 9999)}")) + outs, paid = [], 0.0 + n_out = r.randint(3, 5) + if kind == "disguised_supplies": + n_out = max(2, min(n_out, int(collected // 9600))) + for i in range(n_out): + if kind == "disguised_payroll": + dst = self.take(pool, 1)[0] if i < n_out - 1 else self.illicit_ext() + a, ch, desc = round(collected / n_out * r.uniform(0.95, 1.0), 2), "e-transfer", "PAYROLL" + else: + dst = self.illicit_ext() + a, ch, desc = round(r.uniform(9300, max(9301, min(9950, collected / n_out))), 2), "wire", f"SUPPLIES INVOICE {r.randint(1000, 9999)}" + outs.append(dst) + paid += a + keys.append(self.add(self.when(t0 + 3 + i * r.randint(1, 3)), front, dst, a, ch, desc)) + involved += outs + rule = "R06" if kind == "disguised_payroll" else "R08" + why = ("pays 'payroll' to people who are not on its payroll, at no regular schedule, right after receiving invoice payments from individuals" + if kind == "disguised_payroll" else + "pays 'supplies' invoices to new payees in amounts just under the 10,000 reporting threshold, funded by invoice payments from individuals") + self.case("escalate", kind, rule, front, involved, keys, + f"The business received {int(collected):,} and sent out {int(paid):,} within about a week. It {why}. " + f"The activity is {collected / max(float(front_c['expected_monthly_activity']), 1):.1f} times what the business profile expects.", + collected, t0 + 8, (68, 92)) + + def lookalike(self, kind: str, pool: List[str]) -> None: + r = self.rng + t0, keys = r.randint(20, 320), [] + c = r.choice(self.established("person")) + acct = self.primary[c["id"]] + if kind == "possible_fan_in": + friends = [self.ext() for _ in range(r.randint(4, 6))] + keys = [self.add(self.when(t0 + i % 2), f, acct, round(r.uniform(40, 150), 2), "e-transfer", "BIRTHDAY GIFT") for i, f in enumerate(friends)] + rule, involved = "R03", [acct] + friends + reason = (f"{len(friends)} small e-transfers arrived in two days, which looks like a fan-in. The amounts are small, the notes say gift, and the funds stayed in the account. " + "Keep watching: no evidence of onward movement.") + elif kind == "possible_cash_structuring": + acct = self.primary[r.choice([x for x in self.established("business") if x["_cash"]] or [c])["id"]] + keys = [self.add(self.when(t0 + i * 2), self.ext(), acct, round(r.uniform(8200, 9800), 2), "cash", "BRANCH CASH DEPOSIT") for i in range(3)] + rule, involved = "R01", [acct] + reason = ("Three cash deposits just under 10,000 in one week. The customer runs a cash business and makes regular cash deposits, " + "but these are much larger than usual and sit close to the threshold. Keep watching and ask for the sales records.") + elif kind == "possible_cycle": + other = self.mix(pool, 1, 1.0)[0] + while other == acct: + other = self.mix(pool, 1, 1.0)[0] + a = round(r.uniform(900, 2400), 2) + keys = [self.add(self.when(t0), acct, other, a, "e-transfer", "LOAN TO FRIEND"), self.add(self.when(t0 + 20), other, acct, a, "e-transfer", "LOAN REPAYMENT")] + rule, involved = "R05", [acct, other] + reason = "Money went to another customer and came back in full 20 days later. The notes say loan and repayment, and the amounts are small. Keep watching for repeats." + else: + supplier = self.ext() + a = round(r.uniform(4000, 9000), 2) + keys = [self.add(self.when(t0, 10), self.ext(), acct, a, "wire", "CLIENT PAYMENT"), self.add(self.when(t0, 15), acct, supplier, round(a * 0.97, 2), "wire", "SUPPLIER INVOICE 4471")] + rule, involved = "R02", [acct, supplier] + reason = "A payment came in and 97 percent left the same day to a supplier payee. Businesses do pay suppliers by wire, but the invoice is not on file. Keep watching and ask for it." + self.case("monitor", kind, rule, acct, involved, keys, reason, sum(self.amt[k] for k in keys), t0 + 3, (45, 65)) + + def benign(self, kind: int) -> None: + r = self.rng + t0, c = r.randint(20, 330), r.choice(self.established()) + acct, keys = self.primary[c["id"]], [] + if kind == 0: + a = round(r.uniform(60, 240) * 1000, 2) + keys = [self.add(self.when(t0, 11), self.ext(), acct, a, "wire", "REAL ESTATE CLOSING PROCEEDS")] + rule, why = "R06", "A single large wire from a law firm for the sale of a home. The customer told the branch about the sale in advance and the deed is on file." + elif kind == 1: + a = round(c["expected_monthly_activity"] * r.uniform(2.5, 4.0), 2) + keys = [self.add(self.when(t0, 9), self.ext(), acct, a, "e-transfer", "PAYROLL BONUS")] + rule, why = "R06", "A one-time bonus from the customer's regular employer. The sender matches the payroll deposits of the last twelve months." + elif kind == 2: + payee, each = self.ext(), round(r.uniform(900, 1100), 2) + t0 = max(t0, 100) + for back in (90, 60): # the payee is not new: the same amount went to it earlier (ordinary activity, not part of the case) + self.add(self.when(t0 - back), acct, payee, each, "e-transfer", "FAMILY SUPPORT") + keys = [self.add(self.when(t0 + i * 7), acct, payee, each, "e-transfer", "FAMILY SUPPORT") for i in range(4)] + a = sum(self.amt[k] for k in keys) + rule, why = "R08", "Four equal weekly e-transfers with the note family support. The customer paid the same payee the same amount twice in the previous three months, so the payee is not new." + else: + a = round(r.uniform(18, 32) * 1000, 2) + keys = [self.add(self.when(t0, 10), acct, self.ext(), a, "wire", "TUITION PAYMENT")] + rule, why = "R07", "A wire to a foreign university for tuition. The invoice is on file and the customer is enrolled." + self.case("dismiss", "none", rule, acct, [acct], keys, why, a, t0 + 1, (15, 35)) + + def sanctions_cases(self) -> None: + entries = json.loads((DATA / "sdn_seed.json").read_text()) + meta = {"list_date": entries["list_date"]} + n_exact, n_fuzzy = self.cfg["sanctions"] + biz = self.established("business") + picks = self.rng.sample(biz, min(len(biz), n_exact + n_fuzzy)) + for i, c in enumerate(picks): + e = entries["entries"][i % len(entries["entries"])] + exact = i < n_exact + name = e["name"] if exact else self.vary(e["name"]) + c["name"], c["occupation"], c["_cash"] = name, "import and export", False + score = 1.0 if exact else round(difflib.SequenceMatcher(None, name, e["name"]).ratio(), 2) + self.sanctions.append({"customer_id": c["id"], "customer_name": name, "list": "OFAC SDN (seeded test entry)", "sdn_uid": e["uid"], + "sdn_name": e["name"], "match_type": "exact" if exact else "fuzzy", "score": f"{score:.2f}", "list_date": meta["list_date"], + "note": "TEST DATA: this customer is invented and was named to match a list entry for screening practice"}) + acct = self.primary[c["id"]] + t0 = self.rng.randint(30, 330) + keys = [self.add(self.when(t0 + j * 3), acct, self.illicit_ext(), round(self.rng.uniform(12, 40) * 1000, 2), "wire", f"WIRE INVOICE {self.rng.randint(1000, 9999)}") for j in range(2)] + reason = ("The customer name is identical to a sanctioned entity on the list and the account sends wires abroad. Treat as a true match until cleared." + if exact else + f"The customer name is close to a listed entity (score {score:.2f}) but not identical. Compare registration details and ownership before deciding.") + self.case("escalate" if exact else "monitor", "sanctions_match", "R09", acct, [acct], keys, reason, + sum(self.amt[k] for k in keys), t0 + 4, (80, 98) if exact else (50, 70)) + + def vary(self, name: str) -> str: + """A near-miss spelling of a list name: two adjacent letters swapped, and sometimes LIMITED shortened.""" + r = self.rng + for _ in range(8): + words = name.replace(".", "").split() + i = r.randint(0, len(words) - 1) + w = words[i] + if len(w) > 4: + j = r.randint(1, len(w) - 3) + words[i] = w[:j] + w[j + 1] + w[j] + w[j + 2:] + out = " ".join(words) + if r.random() < 0.5: + out = out.replace("LIMITED", "LTD") + if out != name: + return out + return name + " INTL" + + # ---- assemble ------------------------------------------------------------------------------------------------------------ + def tables(self) -> Tables: + ordered = sorted(self.tx, key=lambda t: (t["timestamp"], t["_k"])) + final: Dict[int, str] = {} + txs = [] + for n, t in enumerate(ordered, 1): + final[t["_k"]] = f"T{n:07d}" + txs.append({"id": final[t["_k"]], **{c: t[c] for c in TABLES["transactions"][1:]}, }) + for t in txs: + t["amount"] = f"{t['amount']:.2f}" + cases = sorted(self.cases, key=lambda c: (c["day"], c["account"], c["pattern"])) + alerts, truth = [], [] + for n, c in enumerate(cases, 1): + aid = f"AL{n:05d}" + alerts.append({"id": aid, "rule": c["rule"], "account": c["account"], "date": (self.base + timedelta(days=c["day"])).isoformat(), + "amount": f"{c['amount']:.2f}", "risk_score": c["score"]}) + truth.append({"alert_id": aid, "label": c["label"], "pattern": c["pattern"], "involved_accounts": LIST_SEP.join(c["involved"]), + "key_transactions": LIST_SEP.join(sorted(final[k] for k in c["keys"])), "reason": c["reason"]}) + sanctions = [{"id": f"S{i:04d}", **s} for i, s in enumerate(self.sanctions, 1)] + customers = [{k: v for k, v in c.items() if not k.startswith("_")} for c in self.customers] + for c in customers: + self.kyc[c["id"]] = self.profile(c) + return {"customers": customers, "accounts": self.accounts, "transactions": txs, "alerts": alerts, "ground_truth": truth, + "kyc_profiles": [{"customer_id": cid, "profile": text} for cid, text in sorted(self.kyc.items())], "sanctions_hits": sanctions} + + def profile(self, c: Dict[str, Any]) -> str: + cur = self.cur[self.primary[c["id"]]] + if c["type"] == "business": + who = f"{c['name']} is a {c['occupation']} business in {N.COUNTRY_NAMES.get(c['country'], c['country'])}." + else: + job = {"student": "is a student", "retired": "is retired"}.get(c["occupation"], f"works as a {c['occupation']}") + who = f"{c['name']} lives in {N.COUNTRY_NAMES.get(c['country'], c['country'])} and {job}." + return (f"{who} Customer since {c['onboarding_date']}. Expects about {c['expected_monthly_activity']:,} {cur} of deposits each month. " + f"Risk rating: {c['risk_rating']}.") + + +class Builtin(Generator): + id = "builtin" + + def generate(self, seed: Any = 42, size: str = "small", jurisdiction: str = "ca", base_date: str = "2025-01-01", **_: Any) -> Dataset: + if size not in SIZES: + raise ValueError(f"size must be one of {', '.join(SIZES)}") + if jurisdiction not in ("ca", "us", "both"): + raise ValueError("jurisdiction must be ca, us or both") + b = _Build(int(seed), size, jurisdiction, date.fromisoformat(str(base_date))) + b.make_customers() + b.normal_activity() + pool = b.mule_pool() + for pattern in PATTERNS: + for _ in range(b.cfg["per_kind"]): + b.launder(pattern, pool) + for kind in ("disguised_payroll", "disguised_supplies"): + for _ in range(b.cfg["per_kind"]): + b.disguised(kind, pool) + for i in range(b.cfg["monitor"]): + b.lookalike(("possible_fan_in", "possible_cash_structuring", "possible_cycle", "possible_rapid_in_out")[i % 4], pool) + for i in range(b.cfg["dismiss"]): + b.benign(i % 4) + b.sanctions_cases() + tables = b.tables() + problems = validate(tables) + if problems: + raise ValueError("generated data failed validation: " + "; ".join(problems[:5])) + return Dataset(tables, ["synthetic", "ofac-sdn"], {"generator": "builtin", "version": VERSION, "seed": int(seed), "size": size, + "jurisdiction": jurisdiction, "base_date": str(base_date), + "institution": "Maple Ridge Credit Union"}) diff --git a/labagent/packs/aml/licences.py b/labagent/packs/aml/licences.py new file mode 100644 index 0000000..fe3409a --- /dev/null +++ b/labagent/packs/aml/licences.py @@ -0,0 +1,42 @@ +"""Where every AML input comes from, under what licence, and whether it may be written into a published snapshot. + +Rule: a snapshot only ever contains sources whose `mode` is "bundled". Everything else is fetched on the user's machine at +launch, or imported only when the user asks and has accepted the licence. `assert_redistributable` is the gate that +`build_table_snapshot` calls, so a fetch-only source cannot reach a release by accident. +""" +from __future__ import annotations + +from typing import Dict, Iterable + + +class LicenceError(Exception): + pass + + +SOURCES: Dict[str, Dict[str, str]] = { + "synthetic": {"licence": "CC-BY-4.0", "mode": "bundled", "note": "Generated by labagent. Names, numbers and documents are invented."}, + "ofac-sdn": {"licence": "US government work (public domain)", "mode": "bundled", + "note": "A fixed set of 12 entity names with their SDN uid. Seeds screening practice. Not a current list."}, + "ecfr-31-cfr-x": {"licence": "US government work (public domain)", "mode": "bundled", + "note": "31 CFR Chapter X from govinfo bulk XML. The eCFR is an unofficial compilation."}, + "justice-laws-ca": {"licence": "Open Government Licence - Canada", "mode": "bundled", + "note": "PCMLTFA and regulations from Justice Laws XML. Needs attribution and a notice that this is not the official version."}, + "ffiec-bsa-aml": {"licence": "not verified", "mode": "fetched", "note": "Fetched on the user's machine at launch."}, + "fintrac": {"licence": "Government of Canada terms (not reviewed)", "mode": "fetched", "note": "Indicators and STR guidance. Fetched at launch."}, + "opensanctions": {"licence": "CC-BY-NC-4.0", "mode": "fetched", "note": "Non-commercial only. Opt-in fetch with an explicit acknowledgement."}, + "ibm-aml-kaggle": {"licence": "CDLA-Sharing-1.0 (per search result, not verified against the Kaggle page)", "mode": "import", + "note": "User supplies Kaggle credentials. Never bundled."}, + "saml-d": {"licence": "not verified", "mode": "import", "note": "User downloads it. Never bundled."}, +} + + +def assert_redistributable(source_ids: Iterable[str]) -> None: + bad = [] + for s in source_ids: + info = SOURCES.get(s) + if info is None: + bad.append(f"{s} (unknown source)") + elif info["mode"] != "bundled": + bad.append(f"{s} ({info['mode']} only, {info['licence']})") + if bad: + raise LicenceError("these sources may not be written into a snapshot: " + "; ".join(bad)) diff --git a/labagent/packs/aml/names.py b/labagent/packs/aml/names.py new file mode 100644 index 0000000..5b88f26 --- /dev/null +++ b/labagent/packs/aml/names.py @@ -0,0 +1,56 @@ +"""Invented names and the pools the builtin generator draws from. No real institution appears here: DENYLIST is checked by tests +against every file the pack ships and every table it generates.""" +from __future__ import annotations + +import re + +FIRST = ["Aiden", "Amara", "Anika", "Arjun", "Beatrice", "Bram", "Camila", "Chloe", "Dara", "Declan", "Elena", "Emeka", "Farah", "Felix", + "Gwen", "Hamid", "Hana", "Hugo", "Imani", "Isaac", "Jonah", "Julia", "Kai", "Keira", "Lars", "Leila", "Liam", "Lucia", "Mateo", + "Maya", "Nadia", "Nolan", "Olga", "Omar", "Priya", "Quinn", "Rafael", "Rhea", "Sam", "Selma", "Tariq", "Thea", "Uma", "Viktor", + "Wren", "Xavier", "Yara", "Zane", "Ines", "Joaquin", "Katya", "Malik", "Noor", "Owen", "Petra", "Ravi", "Sofia", "Tomas", + "Una", "Vera", "Wes", "Yusuf", "Zoe", "Abel", "Bianca", "Cyrus", "Dalia", "Eli", "Freya", "Gideon", "Hallie", "Idris", "Jade"] +LAST = ["Abara", "Alder", "Banerjee", "Brightwater", "Caldwell", "Chen", "Dalgleish", "Delacroix", "Eriksen", "Fairbairn", "Garnier", + "Haddad", "Ishikawa", "Jovanovic", "Kowalczyk", "Lindqvist", "Moreau", "Nakamura", "Okafor", "Petrov", "Quill", "Rosario", + "Sandhu", "Thorne", "Ueda", "Valdez", "Whitlock", "Yilmaz", "Zielinski", "Ashworth", "Boyle", "Cordero", "Dube", "Everhart", + "Fontaine", "Grewal", "Holloway", "Ibarra", "Jansen", "Kapoor", "Larkin", "Marchetti", "Novak", "Ocampo", "Pereira", "Raines", + "Sorensen", "Tremblay", "Underhill", "Vasquez", "Wexler", "Yamada", "Zhou", "Armitage", "Bastien", "Costa", "Doherty", "Elwood"] +BUSINESS_STEMS = ["Birchwood", "Cedar Hollow", "Driftline", "Eastgate", "Fernbrook", "Granite Peak", "Harbourview", "Ironbridge", "Juniper", + "Kestrel", "Lakeshore", "Millstone", "Northfield", "Oakridge", "Pinecrest", "Quarry Lane", "Redfern", "Stonehaven", + "Timberline", "Upland", "Valleyfield", "Westmere", "Yarrow", "Alder Creek", "Brightside", "Copperfield", "Dunmore"] +BUSINESS_KINDS = [ + # (industry, name suffix, mean monthly revenue, cash intensive) + ("restaurant", "Kitchen Ltd.", 42000, True), ("landscaping", "Landscaping Inc.", 26000, True), ("auto repair", "Auto Service Ltd.", 38000, True), + ("retail clothing", "Apparel Inc.", 34000, False), ("construction", "Builders Ltd.", 90000, False), ("dental clinic", "Dental Group Inc.", 70000, False), + ("convenience store", "Market Ltd.", 52000, True), ("trucking", "Freight Ltd.", 80000, False), ("consulting", "Advisory Inc.", 45000, False), + ("bakery", "Bakery Ltd.", 21000, True), ("software services", "Digital Inc.", 60000, False), ("import and export", "Trading Ltd.", 75000, False), +] +OCCUPATIONS = [ + # (occupation, mean monthly income) + ("teacher", 5200), ("nurse", 6100), ("software developer", 8200), ("electrician", 6400), ("student", 1400), ("retired", 3100), + ("truck driver", 5600), ("accountant", 7000), ("server", 2600), ("warehouse worker", 3600), ("civil engineer", 8800), + ("pharmacist", 9200), ("graphic designer", 5000), ("paramedic", 5900), ("sales associate", 3200), ("plumber", 6000), +] +EMPLOYERS = ["Northbridge School Board", "Lakeview Health Region", "Crestline Logistics", "Ardent Software", "Brightpath Staffing", + "Summit Retail Group", "Harlow Construction", "Clearwater Pharmacy", "Ostrander Engineering", "Pembroke Foods", + "Tidewater Transit", "Greenleaf Grocers", "Ridgeway Manufacturing", "Silverline Dental", "Fairhaven Care Homes"] +CITIES = {"CA": ["Maple Ridge", "Kelowna", "Red Deer", "Guelph", "Moncton", "Brandon", "Kamloops", "Sudbury"], + "US": ["Spokane", "Boise", "Duluth", "Dayton", "Asheville", "Fresno", "Tulsa", "Madison"]} +FOREIGN = ["GB", "IN", "PH", "MX", "AE", "TR", "NG", "CN", "PA", "CY"] +COUNTRY_NAMES = {"CA": "Canada", "US": "the United States", "GB": "the United Kingdom", "IN": "India", "PH": "the Philippines", "MX": "Mexico", + "AE": "the United Arab Emirates", "TR": "Turkey", "NG": "Nigeria", "CN": "China", "PA": "Panama", "CY": "Cyprus"} +HIGH_RISK_COUNTRIES = {"AE", "TR", "NG", "PA", "CY"} +BRANCHES = ["B01", "B02", "B03", "B04", "B05"] + +# Real institutions and payment brands that must never appear in generated data or shipped documents. +DENYLIST = [ + r"royal bank", r"\brbc\b", r"\btd bank\b", r"toronto[- ]dominion", r"scotiabank", r"bank of nova scotia", r"\bbmo\b", r"bank of montreal", + r"\bcibc\b", r"national bank", r"desjardins", r"tangerine", r"\bhsbc\b", r"jp ?morgan", r"\bchase\b", r"wells fargo", r"bank of america", + r"citibank", r"citigroup", r"goldman", r"morgan stanley", r"u\.?s\.? bancorp", r"u\.?s\.? bank", r"\bpnc\b", r"truist", r"capital one", + r"barclays", r"santander", r"deutsche bank", r"\bubs\b", r"credit suisse", r"western union", r"paypal", r"\bvisa\b", r"mastercard", + r"\binterac\b", r"\bzelle\b", r"coast capital", r"vancity", r"servus", r"meridian credit", +] +DENY_RE = re.compile("|".join(DENYLIST), re.I) + + +def real_bank_hits(text: str): + return sorted({m.group(0).lower() for m in DENY_RE.finditer(text)}) diff --git a/labagent/packs/aml/schema.py b/labagent/packs/aml/schema.py new file mode 100644 index 0000000..23ae4c9 --- /dev/null +++ b/labagent/packs/aml/schema.py @@ -0,0 +1,182 @@ +"""The AML table schema. Every generator (builtin, AMLSim, snapshot, importers) produces exactly these tables and columns.""" +from __future__ import annotations + +import csv +import re +from pathlib import Path +from typing import Any, Dict, List, Optional + +Rows = List[Dict[str, Any]] +Tables = Dict[str, Rows] + +INSTITUTION = "Maple Ridge Credit Union" + +TABLES: Dict[str, List[str]] = { + "customers": ["id", "name", "type", "occupation", "country", "risk_rating", "expected_monthly_activity", "onboarding_date"], + "accounts": ["id", "customer_id", "type", "open_date", "branch"], + "transactions": ["id", "timestamp", "from_account", "to_account", "amount", "currency", "channel", "description"], + "alerts": ["id", "rule", "account", "date", "amount", "risk_score"], + "ground_truth": ["alert_id", "label", "pattern", "involved_accounts", "key_transactions", "reason"], + "kyc_profiles": ["customer_id", "profile"], + "sanctions_hits": ["id", "customer_id", "customer_name", "list", "sdn_uid", "sdn_name", "match_type", "score", "list_date", "note"], +} + +CHANNELS = ("wire", "cash", "e-transfer", "card", "crypto off-ramp") +LABELS = ("escalate", "monitor", "dismiss") +CUSTOMER_TYPES = ("person", "business") +RISK_RATINGS = ("low", "medium", "high") + +# The 8 patterns from Altman et al. (2023), arXiv 2306.16424, Figure 2. +PATTERNS = ("fan_out", "fan_in", "gather_scatter", "scatter_gather", "cycle", "random", "bipartite", "stack") +# Laundering that looks like ordinary payments, and a name match. Escalate cases only. +DISGUISED = ("disguised_payroll", "disguised_supplies") +OTHER_TYPOLOGIES = ("sanctions_match",) +# Monitor cases carry the look-alike shape with a "possible_" prefix. Dismiss cases carry "none". +LOOKALIKES = ("possible_fan_in", "possible_cash_structuring", "possible_cycle", "possible_rapid_in_out") +ALL_PATTERNS = PATTERNS + DISGUISED + OTHER_TYPOLOGIES + LOOKALIKES + ("none",) + +RULES = { + "R01": "Cash deposits just under the reporting threshold", + "R02": "Rapid movement of funds in and out", + "R03": "Many senders to one account", + "R04": "One account sends to many recipients", + "R05": "Funds return to the sender through a loop", + "R06": "Activity far above expected", + "R07": "Wire to or from a high-risk country", + "R08": "Round or just-under-threshold payments to new payees", + "R09": "Sanctions name match", +} + +EXTERNAL_PREFIX = "X" # accounts at other institutions: they appear in transactions but not in `accounts` +TS_RE = re.compile(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}$") +DATE_RE = re.compile(r"^\d{4}-\d{2}-\d{2}$") +LIST_SEP = ";" + + +def write_csv(tables: Tables, out_dir: Path) -> List[Path]: + """One CSV per table. Fixed column order, LF line ends, UTF-8: the same input always gives the same bytes.""" + out_dir.mkdir(parents=True, exist_ok=True) + paths = [] + for name, cols in TABLES.items(): + path = out_dir / f"{name}.csv" + with open(path, "w", newline="", encoding="utf-8") as f: + w = csv.DictWriter(f, fieldnames=cols, lineterminator="\n") + w.writeheader() + for row in tables[name]: + w.writerow({c: row[c] for c in cols}) + paths.append(path) + return paths + + +def write_parquet(tables: Tables, out_dir: Path) -> List[Path]: + """Optional: needs pyarrow (`pip install 'labagent[aml]'`). Raises ImportError if it is missing.""" + import pyarrow as pa + import pyarrow.parquet as pq + out_dir.mkdir(parents=True, exist_ok=True) + paths = [] + for name, cols in TABLES.items(): + path = out_dir / f"{name}.parquet" + pq.write_table(pa.table({c: [str(r[c]) if not isinstance(r[c], (int, float)) else r[c] for r in tables[name]] for c in cols}), path) + paths.append(path) + return paths + + +def read_csv(in_dir: Path) -> Tables: + out: Tables = {} + for name in TABLES: + with open(in_dir / f"{name}.csv", newline="", encoding="utf-8") as f: + out[name] = list(csv.DictReader(f)) + return out + + +def validate(tables: Tables) -> List[str]: + """Return a list of problems (empty means the dataset is valid). Works on rows from a generator or from read_csv.""" + problems: List[str] = [] + + def bad(msg: str) -> None: + if len(problems) < 50: + problems.append(msg) + + for name, cols in TABLES.items(): + if name not in tables: + bad(f"missing table {name}") + continue + for i, row in enumerate(tables[name]): + if list(row.keys()) != cols: + bad(f"{name} row {i}: columns differ from the schema") + break + if problems: + return problems + + ids = lambda t, k="id": [str(r[k]) for r in tables[t]] + for t in ("customers", "accounts", "transactions", "alerts", "sanctions_hits"): + if len(set(ids(t))) != len(tables[t]): + bad(f"{t}: duplicate ids") + cust, acct = set(ids("customers")), set(ids("accounts")) + tx_ids = set(ids("transactions")) + alert_ids = set(ids("alerts")) + + for r in tables["customers"]: + if r["type"] not in CUSTOMER_TYPES: + bad(f"customer {r['id']}: bad type {r['type']}") + if r["risk_rating"] not in RISK_RATINGS: + bad(f"customer {r['id']}: bad risk rating") + if not DATE_RE.match(str(r["onboarding_date"])): + bad(f"customer {r['id']}: bad onboarding date") + for r in tables["accounts"]: + if str(r["customer_id"]) not in cust: + bad(f"account {r['id']}: unknown customer") + for r in tables["transactions"]: + if not TS_RE.match(str(r["timestamp"])): + bad(f"transaction {r['id']}: bad timestamp") + if r["channel"] not in CHANNELS: + bad(f"transaction {r['id']}: bad channel {r['channel']}") + if float(r["amount"]) <= 0: + bad(f"transaction {r['id']}: amount must be positive") + for side in ("from_account", "to_account"): + a = str(r[side]) + if a not in acct and not a.startswith(EXTERNAL_PREFIX): + bad(f"transaction {r['id']}: unknown {side} {a}") + if r["from_account"] == r["to_account"]: + bad(f"transaction {r['id']}: sends to itself") + for r in tables["alerts"]: + if str(r["account"]) not in acct: + bad(f"alert {r['id']}: unknown account") + if str(r["rule"]) not in RULES: + bad(f"alert {r['id']}: unknown rule") + gt_ids = [str(r["alert_id"]) for r in tables["ground_truth"]] + if set(gt_ids) != alert_ids or len(gt_ids) != len(alert_ids): + bad("ground_truth must have exactly one row per alert") + for r in tables["ground_truth"]: + if r["label"] not in LABELS: + bad(f"ground_truth {r['alert_id']}: bad label") + if r["pattern"] not in ALL_PATTERNS: + bad(f"ground_truth {r['alert_id']}: unknown pattern {r['pattern']}") + for t in filter(None, str(r["key_transactions"]).split(LIST_SEP)): + if t not in tx_ids: + bad(f"ground_truth {r['alert_id']}: unknown transaction {t}") + for a in filter(None, str(r["involved_accounts"]).split(LIST_SEP)): + if a not in acct and not a.startswith(EXTERNAL_PREFIX): + bad(f"ground_truth {r['alert_id']}: unknown account {a}") + if not str(r["reason"]).strip(): + bad(f"ground_truth {r['alert_id']}: empty reason") + kyc = [str(r["customer_id"]) for r in tables["kyc_profiles"]] + if set(kyc) != cust or len(kyc) != len(cust): + bad("kyc_profiles must have exactly one row per customer") + for r in tables["sanctions_hits"]: + if str(r["customer_id"]) not in cust: + bad(f"sanctions hit {r['id']}: unknown customer") + if "TEST DATA" not in str(r["note"]): + bad(f"sanctions hit {r['id']}: must be tagged TEST DATA") + return problems + + +def table_counts(tables: Tables) -> Dict[str, int]: + return {name: len(tables[name]) for name in TABLES} + + +def find_pattern_labels(tables: Tables) -> Dict[str, Optional[str]]: + """Which of the 8 patterns have at least one escalate case. Used by tests and the dataset card.""" + seen = {p for p in PATTERNS} + have = {r["pattern"] for r in tables["ground_truth"] if r["label"] == "escalate"} + return {p: p if p in have else None for p in sorted(seen)} diff --git a/pyproject.toml b/pyproject.toml index 8ec573e..822df15 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,7 @@ dependencies = ["boto3>=1.34", "cryptography>=42"] rag = ["psycopg[binary]>=3.1"] rag-local = ["pydantic>=2", "chromadb>=0.5", "sentence-transformers>=3", "streamlit>=1.35", "pypdf>=4"] modal = ["modal>=1.0"] +aml = ["openpyxl>=3.1", "python-docx>=1.1", "pyarrow>=15"] rag-lab = ["pydantic>=2", "chromadb>=0.5", "sentence-transformers>=3", "streamlit>=1.35", "pypdf>=4", "groq>=0.9", "langchain>=0.3", "qdrant-client>=1.9", "jupyterlab>=4", "ipykernel>=6"] dev = ["pytest>=8", "pydantic>=2", "moto[all]>=5", "build>=1.2", "twine>=5"] @@ -61,6 +62,7 @@ labagent = [ "terraform/*/.terraform.lock.hcl", "terraform/*/tests/*.tftest.hcl", "packs/*/sources/*.json", + "packs/aml/data/*.json", "packs/*/app/*", "packs/*/app/static/*", "packs/*/tutorial/notebooks/*.py", diff --git a/tests/test_aml_generator.py b/tests/test_aml_generator.py new file mode 100644 index 0000000..bd771b2 --- /dev/null +++ b/tests/test_aml_generator.py @@ -0,0 +1,299 @@ +"""AML pack, part 1: the builtin generator, the schema, the licence gate and table snapshots.""" +import json +import re +import subprocess +import sys +from pathlib import Path + +import pytest + +from labagent.datasets import DatasetError, build_table_snapshot, fetch_table_snapshot, sha256_file +from labagent.packs.aml import names as N +from labagent.packs.aml.export import write_dataset +from labagent.packs.aml.gen import get +from labagent.packs.aml.licences import SOURCES, LicenceError, assert_redistributable +from labagent.packs.aml.schema import (ALL_PATTERNS, DISGUISED, EXTERNAL_PREFIX, INSTITUTION, LABELS, PATTERNS, TABLES, read_csv, + table_counts, validate, write_csv) + +ROOT = Path(__file__).resolve().parents[1] + + +def make(**kw): + kw.setdefault("seed", 42) + kw.setdefault("size", "tiny") + return get("builtin").generate(**kw) + + +@pytest.fixture(scope="module") +def small(): + return make(size="small") + + +# ---- determinism and schema --------------------------------------------------------------------------------------------------- + +def test_same_seed_gives_identical_files_and_a_different_seed_does_not(tmp_path): + for d, seed in (("a", 7), ("b", 7), ("c", 8)): + write_dataset(make(seed=seed), tmp_path / d, parquet="no") + same = lambda x, y: all(sha256_file(tmp_path / x / f) == sha256_file(tmp_path / y / f) for f in sorted(p.name for p in (tmp_path / x).iterdir())) + assert same("a", "b") and not same("a", "c") + + +def test_every_table_has_exactly_the_schema_columns_and_validates(small): + assert validate(small.tables) == [] + for name, cols in TABLES.items(): + assert small.tables[name] and list(small.tables[name][0].keys()) == cols + + +def test_csv_round_trip_still_validates(tmp_path, small): + write_csv(small.tables, tmp_path) + assert validate(read_csv(tmp_path)) == [] + + +def test_validate_catches_broken_data(small): + broken = {k: [dict(r) for r in v] for k, v in small.tables.items()} + broken["transactions"][0]["to_account"] = "ACC999999" + broken["ground_truth"][0]["label"] = "maybe" + broken["sanctions_hits"][0]["note"] = "real" + msgs = " ".join(validate(broken)) + assert "unknown to_account" in msgs and "bad label" in msgs and "TEST DATA" in msgs + + +def test_sizes_and_unknown_options_are_checked(): + assert table_counts(make(size="tiny").tables)["customers"] == 60 + with pytest.raises(ValueError): + make(size="huge") + with pytest.raises(ValueError): + make(jurisdiction="fr") + + +def test_jurisdiction_sets_the_home_currency(): + ca = {t["currency"] for t in make(jurisdiction="ca").tables["transactions"]} + us = {t["currency"] for t in make(jurisdiction="us").tables["transactions"]} + assert "CAD" in ca and "USD" in us + assert {c["country"] for c in make(jurisdiction="both", size="small").tables["customers"]} >= {"CA", "US"} + + +# ---- labels -------------------------------------------------------------------------------------------------------------------- + +def test_each_of_the_eight_patterns_appears_labelled_escalate(small): + escalate = {g["pattern"] for g in small.tables["ground_truth"] if g["label"] == "escalate"} + assert set(PATTERNS) <= escalate and set(DISGUISED) <= escalate and "sanctions_match" in escalate + + +def test_all_three_labels_exist_and_follow_the_rules_of_construction(small): + gt = small.tables["ground_truth"] + assert {g["label"] for g in gt} == set(LABELS) + for g in gt: + if g["label"] == "dismiss": + assert g["pattern"] == "none" + if g["label"] == "monitor": + assert g["pattern"].startswith("possible_") or g["pattern"] == "sanctions_match" + assert g["pattern"] in ALL_PATTERNS and g["reason"].strip() and g["key_transactions"] + + +def test_graph_patterns_have_the_right_shape(small): + tx = {t["id"]: t for t in small.tables["transactions"]} + gt = [g for g in small.tables["ground_truth"] if g["label"] == "escalate"] + for g in gt: + legs = [tx[k] for k in g["key_transactions"].split(";")] + senders, receivers = {t["from_account"] for t in legs}, {t["to_account"] for t in legs} + if g["pattern"] == "fan_out": + hub = max(senders, key=lambda s: sum(1 for t in legs if t["from_account"] == s)) + assert len({t["to_account"] for t in legs if t["from_account"] == hub}) >= 4 + if g["pattern"] == "fan_in": + assert len(senders) >= 4 and len(receivers) == 1 + if g["pattern"] == "cycle": + assert senders == receivers and len(legs) == len(senders) # every account both sends and receives, once + if g["pattern"] == "bipartite": + assert not (senders & receivers) + if g["pattern"] == "scatter_gather": + assert len(senders - receivers) == 1 and len(receivers - senders) == 1 + + +def test_disguised_cases_use_ordinary_looking_descriptions(small): + tx = {t["id"]: t for t in small.tables["transactions"]} + for g in small.tables["ground_truth"]: + if g["pattern"] == "disguised_payroll": + assert any(tx[k]["description"] == "PAYROLL" for k in g["key_transactions"].split(";")) + if g["pattern"] == "disguised_supplies": + assert any(tx[k]["description"].startswith("SUPPLIES INVOICE") and 9000 < float(tx[k]["amount"]) < 10000 + for k in g["key_transactions"].split(";")) + + +def test_laundering_does_not_match_the_customer_profile(small): + """The point of the lab: the alert account moved far more than its KYC profile says to expect.""" + expected = {c["id"]: int(c["expected_monthly_activity"]) for c in small.tables["customers"]} + owner = {a["id"]: a["customer_id"] for a in small.tables["accounts"]} + alert = {a["id"]: a for a in small.tables["alerts"]} + ratios = [float(alert[g["alert_id"]]["amount"]) / expected[owner[alert[g["alert_id"]]["account"]]] + for g in small.tables["ground_truth"] if g["label"] == "escalate" and g["pattern"] in PATTERNS] + assert sum(r > 2 for r in ratios) / len(ratios) > 0.8 + + +def test_no_case_happens_before_its_account_was_opened(small): + opened = {a["id"]: a["open_date"] for a in small.tables["accounts"]} + tx = {t["id"]: t for t in small.tables["transactions"]} + for g in small.tables["ground_truth"]: + for k in g["key_transactions"].split(";"): + for side in ("from_account", "to_account"): + a = tx[k][side] + if a in opened: + assert tx[k]["timestamp"][:10] >= opened[a] + + +def test_case_reasons_match_the_data(): + """Reasons are written for learners to check against the tables, so the checkable claims must be true.""" + for seed in (42, 1, 2): + ds = make(seed=seed, size="small") + tx = {t["id"]: t for t in ds.tables["transactions"]} + allx = ds.tables["transactions"] + for g in ds.tables["ground_truth"]: + legs = [tx[k] for k in g["key_transactions"].split(";")] + if g["pattern"] == "disguised_supplies": + inn = sum(float(t["amount"]) for t in legs if t["description"].startswith("INVOICE")) + out = sum(float(t["amount"]) for t in legs if t["description"].startswith("SUPPLIES")) + assert out <= inn + 0.01 + if g["label"] == "dismiss" and legs[0]["description"] == "FAMILY SUPPORT": + payee, amt = legs[0]["to_account"], legs[0]["amount"] + earlier = [t for t in allx if t["to_account"] == payee and t["amount"] == amt and t["timestamp"] < legs[0]["timestamp"]] + assert len(earlier) >= 2 + if g["pattern"] == "possible_cash_structuring": + assert all(float(t["amount"]) > 8000 for t in legs) + + +# ---- sanctions seeds ----------------------------------------------------------------------------------------------------------- + +def test_sanctions_hits_are_tagged_test_data_and_use_entity_names_only(small): + seed = json.loads((ROOT / "labagent/packs/aml/data/sdn_seed.json").read_text()) + assert len(seed["entries"]) == 12 and all(set(e) == {"uid", "name", "program"} and e["uid"].isdigit() for e in seed["entries"]) # no addresses, no people + hits = small.tables["sanctions_hits"] + assert {h["match_type"] for h in hits} == {"exact", "fuzzy"} + for h in hits: + assert "TEST DATA" in h["note"] and h["list_date"] == seed["list_date"] + assert (h["customer_name"] == h["sdn_name"]) == (h["match_type"] == "exact") + assert h["sdn_name"] in {e["name"] for e in seed["entries"]} + + +# ---- names and hygiene --------------------------------------------------------------------------------------------------------- + +def test_no_real_bank_or_payment_brand_in_generated_tables_or_shipped_code(small): + for rows in small.tables.values(): + assert N.real_bank_hits(json.dumps(rows)) == [] + assert INSTITUTION == "Maple Ridge Credit Union" + for f in (ROOT / "labagent/packs/aml").rglob("*.py"): + if f.name != "names.py": + assert N.real_bank_hits(f.read_text()) == [], f + + +def test_denylist_catches_what_it_should(): + assert N.real_bank_hits("sent via Interac to a TD Bank account at Scotiabank") == ["interac", "scotiabank", "td bank"] + assert N.real_bank_hits("Maple Ridge Credit Union") == [] + + +def test_generated_files_pass_the_repo_pii_scan(tmp_path): + """scripts/check_no_pii.py only reads tracked files, so run the same patterns over a generated folder.""" + write_dataset(make(size="small"), tmp_path, parquet="no") + sys.path.insert(0, str(ROOT / "scripts")) + try: + import check_no_pii as c + finally: + sys.path.pop(0) + bad = [] + for f in tmp_path.iterdir(): + for n, line in enumerate(f.read_text().splitlines(), 1): + if any(not c.ALLOWED_EMAIL.search(m.group(0)) for m in c.EMAIL.finditer(line)): + bad.append(f"{f.name}:{n} email") + if c.ACCOUNT.search(c.UUID.sub("", line)): + bad.append(f"{f.name}:{n} 12-digit number") + if c.KEYS.search(line): + bad.append(f"{f.name}:{n} key") + assert bad == [] + + +def test_ids_are_short_and_external_parties_are_marked(small): + assert all(len(a["id"]) == 9 for a in small.tables["accounts"]) + ext = {t[s] for t in small.tables["transactions"] for s in ("from_account", "to_account") if t[s].startswith(EXTERNAL_PREFIX)} + assert ext and not ext & {a["id"] for a in small.tables["accounts"]} + + +# ---- licences and snapshots ---------------------------------------------------------------------------------------------------- + +def test_only_bundled_sources_may_enter_a_snapshot(): + assert_redistributable(["synthetic", "ofac-sdn", "ecfr-31-cfr-x", "justice-laws-ca"]) + for s in ("opensanctions", "fintrac", "ffiec-bsa-aml", "ibm-aml-kaggle", "saml-d"): + assert SOURCES[s]["mode"] != "bundled" + with pytest.raises(LicenceError): + assert_redistributable(["synthetic", s]) + with pytest.raises(LicenceError): + assert_redistributable(["something-new"]) + + +def test_builtin_data_uses_only_bundled_sources(): + ds = make() + assert ds.sources == ["synthetic", "ofac-sdn"] + assert "opensanctions" not in json.dumps(ds.info).lower() + + +def test_snapshot_is_reproducible_and_lists_its_sources(tmp_path): + for d in "ab": + write_dataset(make(), tmp_path / f"src-{d}", parquet="no") + a = build_table_snapshot(tmp_path / "src-a", "aml-synthetic", "t1", tmp_path / "a", built="2026-01-01") + b = build_table_snapshot(tmp_path / "src-b", "aml-synthetic", "t1", tmp_path / "b", built="2026-01-01") + assert a["sha256"] == b["sha256"] and a["kind"] == "tables" + assert [s["id"] for s in a["sources"]] == ["synthetic", "ofac-sdn"] and set(a["files"]) >= {"transactions.csv", "provenance.json"} + assert (tmp_path / "a" / "SHA256SUMS").read_text().split()[0] == a["sha256"] + + +def test_snapshot_refuses_a_fetch_only_source_and_missing_provenance(tmp_path): + write_dataset(make(), tmp_path / "src", parquet="no") + prov = json.loads((tmp_path / "src" / "provenance.json").read_text()) + prov["sources"].append({"id": "opensanctions", "licence": "CC-BY-NC-4.0", "mode": "fetched"}) + (tmp_path / "src" / "provenance.json").write_text(json.dumps(prov)) + with pytest.raises(DatasetError, match="opensanctions"): + build_table_snapshot(tmp_path / "src", "x", "1", tmp_path / "out") + (tmp_path / "src" / "provenance.json").unlink() + with pytest.raises(DatasetError, match="provenance"): + build_table_snapshot(tmp_path / "src", "x", "1", tmp_path / "out") + + +def test_fetch_checks_the_hash_and_unpacks_only_flat_files(tmp_path): + write_dataset(make(), tmp_path / "src", parquet="no") + m = build_table_snapshot(tmp_path / "src", "demo", "1", tmp_path / "out") + entry = {"kind": "tables", "url": "https://x.example/demo.zip", "sha256": m["sha256"]} + zbytes = (tmp_path / "out" / m["file"]).read_bytes() + dl = lambda u, d: d.write_bytes(zbytes) + dest = fetch_table_snapshot("demo", tmp_path / "dest", cache_dir=tmp_path / "cache", download=dl, entry=entry) + assert validate(read_csv(dest)) == [] + with pytest.raises(DatasetError, match="matching checksum"): + fetch_table_snapshot("demo", tmp_path / "d2", cache_dir=tmp_path / "c2", download=lambda u, d: d.write_bytes(b"tampered"), entry=entry) + + +def test_zip_with_a_path_in_it_is_refused(tmp_path): + import zipfile + z = tmp_path / "evil.zip" + with zipfile.ZipFile(z, "w") as f: + f.writestr("../escape.csv", "x") + entry = {"kind": "tables", "url": "https://x.example/evil.zip", "sha256": sha256_file(z)} + with pytest.raises(DatasetError, match="unsafe"): + fetch_table_snapshot("evil", tmp_path / "d", cache_dir=tmp_path / "c", download=lambda u, d: d.write_bytes(z.read_bytes()), entry=entry) + assert not (tmp_path / "escape.csv").exists() + + +def test_table_and_chunk_snapshots_are_not_mixed_up(tmp_path): + from labagent.datasets import fetch_snapshot + with pytest.raises(DatasetError, match="tables"): + fetch_snapshot("t", cache_dir=tmp_path, entry={"kind": "tables", "url": "https://x.example/a.zip", "sha256": "0" * 64}) + + +def test_dataset_cli_packages_tables(tmp_path, capsys): + from labagent.cli import main + write_dataset(make(), tmp_path / "src", parquet="no") + assert main(["dataset", "--tables", str(tmp_path / "src"), "--name", "aml-synthetic", "--version", "t1", "--out", str(tmp_path / "o")]) == 0 + assert json.loads(capsys.readouterr().out)["file"] == "aml-synthetic-t1.tables.zip" + assert main(["dataset", "--name", "x", "--version", "1", "--out", str(tmp_path / "o")]) == 2 + + +def test_module_entry_point_writes_the_tables(tmp_path): + p = subprocess.run([sys.executable, "-m", "labagent.packs.aml.gen", "--out", str(tmp_path), "--size", "tiny", "--parquet", "no"], + capture_output=True, text=True, cwd=ROOT) + assert p.returncode == 0 and json.loads(p.stdout)["rows"]["customers"] == 60