From 4575f925278845847b2e50e2c692f97dd79cd3e5 Mon Sep 17 00:00:00 2001 From: Chandan Kumar Date: Sat, 24 Oct 2026 12:30:51 -0400 Subject: [PATCH] AML pack, part 3: rag mode on the local platform Registers the 'aml' pack (alias aml-rag). mode=rag generates the data, writes the corpus, indexes it with the existing RAG client, runs three quality checks (retrieval hit rate, citation validity, triage agreement with the answer key) and leaves a lab folder with a guide and a Python helper. Adds the evidence bundle after arXiv 2604.19755 and a generator=snapshot reader. The answer key is kept out of tables/. Co-Authored-By: Claude Sonnet 5.5 --- labagent/packs/aml/corpus/documents.py | 15 ++ labagent/packs/aml/gen/__init__.py | 5 +- labagent/packs/aml/gen/snapshot.py | 28 +++ labagent/packs/aml/pack.py | 293 +++++++++++++++++++++++++ labagent/packs/aml/rag/__init__.py | 1 + labagent/packs/aml/rag/bundle.py | 74 +++++++ labagent/packs/aml/rag/qa.py | 37 ++++ labagent/packs/aml/rag/quality.py | 201 +++++++++++++++++ labagent/packs/aml/rag/tools.py | 80 +++++++ labagent/packs/aml/schema.py | 10 +- labagent/packs/registry.py | 3 +- tests/test_aml_rag.py | 245 +++++++++++++++++++++ tests/test_packs.py | 4 +- 13 files changed, 989 insertions(+), 7 deletions(-) create mode 100644 labagent/packs/aml/gen/snapshot.py create mode 100644 labagent/packs/aml/pack.py create mode 100644 labagent/packs/aml/rag/__init__.py create mode 100644 labagent/packs/aml/rag/bundle.py create mode 100644 labagent/packs/aml/rag/qa.py create mode 100644 labagent/packs/aml/rag/quality.py create mode 100644 labagent/packs/aml/rag/tools.py create mode 100644 tests/test_aml_rag.py diff --git a/labagent/packs/aml/corpus/documents.py b/labagent/packs/aml/corpus/documents.py index 6266efd..3419544 100644 --- a/labagent/packs/aml/corpus/documents.py +++ b/labagent/packs/aml/corpus/documents.py @@ -207,6 +207,21 @@ def __init__(self, t: Tables, window_days: int = 30, depth: int = 8, history_per self.alerts_by_account: Dict[str, List[Dict[str, Any]]] = {} for a in t["alerts"]: self.alerts_by_account.setdefault(a["account"], []).append(a) + self._degree: Dict[str, int] = {} + + def uncommon(self, party: str) -> bool: + """True for another customer of the credit union, or an outside party that few accounts deal with. Ordinary payees (employers, + shops, landlords) are dealt with by many accounts. Transfers between customers do not happen in ordinary activity.""" + if party in self.acct: + return True + if not self._degree: + seen: Dict[str, set] = {} + for x in self.t["transactions"]: + for me, other in ((x["from_account"], x["to_account"]), (x["to_account"], x["from_account"])): + if me not in self.acct: + seen.setdefault(me, set()).add(other) + self._degree = {k: len(v) for k, v in seen.items()} + return self._degree.get(party, 0) < 8 def customer(self, alert: Dict[str, Any]) -> Dict[str, Any]: return self.cust[self.acct[alert["account"]]["customer_id"]] diff --git a/labagent/packs/aml/gen/__init__.py b/labagent/packs/aml/gen/__init__.py index f825fb2..0f9974a 100644 --- a/labagent/packs/aml/gen/__init__.py +++ b/labagent/packs/aml/gen/__init__.py @@ -25,4 +25,7 @@ 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") + if generator_id == "snapshot": + from .snapshot import Snapshot + return Snapshot() + raise KeyError(f"unknown generator '{generator_id}'. Available now: builtin, snapshot") diff --git a/labagent/packs/aml/gen/snapshot.py b/labagent/packs/aml/gen/snapshot.py new file mode 100644 index 0000000..216660f --- /dev/null +++ b/labagent/packs/aml/gen/snapshot.py @@ -0,0 +1,28 @@ +"""generator=snapshot: read a pinned, checksummed table snapshot instead of generating one.""" +from __future__ import annotations + +import tempfile +from pathlib import Path +from typing import Any, Callable, Dict, Optional + +from ....datasets import DatasetError, default_http, fetch_table_snapshot +from ..schema import read_csv, validate +from . import Dataset, Generator + + +class Snapshot(Generator): + id = "snapshot" + + def __init__(self, download: Callable[[str, Path], None] = default_http, entry: Optional[Dict[str, Any]] = None): + self.download, self.entry = download, entry + + def generate(self, dataset: str = "aml-synthetic", **_: Any) -> Dataset: + import json + with tempfile.TemporaryDirectory() as tmp: + folder = fetch_table_snapshot(dataset, Path(tmp), download=self.download, entry=self.entry) + tables = read_csv(folder) + prov = json.loads((folder / "provenance.json").read_text()) + problems = validate(tables) + if problems: + raise DatasetError("the snapshot failed validation: " + "; ".join(problems[:3])) + return Dataset(tables, [s["id"] for s in prov["sources"]], {**prov.get("info", {}), "generator": "snapshot", "dataset": dataset}) diff --git a/labagent/packs/aml/pack.py b/labagent/packs/aml/pack.py new file mode 100644 index 0000000..3bf81ed --- /dev/null +++ b/labagent/packs/aml/pack.py @@ -0,0 +1,293 @@ +"""aml: a synthetic anti-money-laundering lab with ground truth, for two audiences. + +mode=rag Developers build a RAG app in Python. The pack generates the data, builds the corpus, indexes it, runs the quality checks + and leaves a lab folder with a guide and a small Python helper. +mode=copilot Non-technical staff learn Microsoft 365 Copilot on ordinary Excel, Word and PDF files. (Added in a later change.) + +Everything the pack creates lives in one folder (`out`, default ./aml-) marked with .labagent-aml. `destroy` removes only what +the pack wrote there. The pack touches no cloud account in the local platform. +""" +from __future__ import annotations + +import json +import shutil +from pathlib import Path +from typing import Any, Dict, List, Optional + +from ... import datasets +from ..base import Pack, PackContext, PackError, Param, Step +from . import gen as generators +from .corpus import build as corpus_build +from .corpus import documents as D +from .corpus import regs as R +from .schema import TABLES, table_counts, validate, write_csv +from .rag import quality as Q +from .rag import tools + +MARKER = ".labagent-aml" +GUIDE = "LAB_GUIDE.md" +SHIM = "aml_tools.py" +OWN_PATHS = ("tables", "answer_key", "data", "corpus", "index", tools.CONFIG, GUIDE, SHIM, MARKER) + +SHIM_TEXT = '"""Lab helper. Run Python in this folder: from aml_tools import AmlLab; lab = AmlLab()"""\nfrom labagent.packs.aml.rag.tools import AmlLab # noqa: F401\n' + + +def guide_text(cfg: Dict[str, Any], quality: Dict[str, Any]) -> str: + j = {"ca": "Canada", "us": "the United States", "both": "Canada and the United States"}[cfg["jurisdiction"]] + q = quality + note = ("\n**Your embedding model is the keyword-only fallback.** Search matches shared words, not meaning, so some answers will miss. " + "Install the full lab with `pip install 'labagent[rag-local]'` and launch again.\n") if q["keyword_only_embedder"] else "" + return f"""# AML RAG lab: {cfg['institution']} ({j}) + +You build a research assistant for an anti-money laundering team. It answers policy questions with citations, pulls the evidence for an +alert, and helps draft a report narrative. All names and transactions are invented. {cfg['institution']} is not a real institution. +{note} +## What is in this folder + +| Path | What it holds | +|---|---| +| `data/` | Policy manual, alert procedure, report guide and one case file per alert (Markdown). Start `labagent rag serve --dir .` to search them. | +| `tables/` | `customers`, `accounts`, `transactions`, `alerts`, `kyc_profiles`, `sanctions_hits` as CSV | +| `answer_key/` | The decisions and reasons for each alert. Open it after you decide. | +| `aml_tools.py` | Helpers for the tasks below | + +## Starting numbers + +- Retrieval hit rate on {q['retrieval']['questions']} questions with known answers: **{q['retrieval']['hit_rate']:.2f}** +- Citation validity of the reference narratives: **{q['citations']['rate']:.2f}** +- A fixed reference rule agrees with the answer key on **{q['triage']['reference_rule']['agreement']:.2f}** of {q['triage']['alerts']} sample alerts. This shows the evidence is enough to decide. It is not a measure of a model. + +## Task 1: ask policy questions with citations + +```python +from aml_tools import AmlLab +lab = AmlLab() +r = lab.ask("When must a suspicious transaction report be filed?") +print(r.answer) +for c in r.chunks: print(c.rank, c.metadata.get("citation") or c.chunk_id) +``` + +Try a filter: `lab.ask("How long are records kept?", doc_type="policy")`. Find one question the search gets wrong and say why. + +## Task 2: pull the evidence for an alert + +```python +lab.show("{cfg['sample_alert']}") +``` + +The bundle has the customer profile, the alert trigger, the transactions around it, the matching policy and rule text, and similar earlier cases. +Read it and decide: **escalate**, **monitor** or **dismiss**. Write down the transaction IDs that support your decision. + +## Task 3: draft a narrative and check it + +Write a short narrative (who, what, when, where, why, how) for an alert you would escalate. Cite transaction IDs and the policy section. + +```python +print(lab.check(my_narrative)) +``` + +Every cited ID and section must exist. Fix any that do not. Then compare: `lab.score({{"{cfg['sample_alert']}": "escalate"}})`. + +## Task 4 (optional): let a chat model decide + +Set a provider key (for example `GROQ_API_KEY`) and run `lab.triage("{cfg['sample_alert']}")` on ten alerts. Score the answers with `lab.score`. Where does it disagree with you, and why? + +## Cleaning up + +`labagent destroy aml --name {cfg['name']}` removes this folder's lab files. +""" + + +class AmlLab(Pack): + id = "aml" + title = "Synthetic AML lab (RAG and Copilot)" + summary = "Synthetic anti-money laundering data with ground truth, policies and regulations: a RAG lab for developers, a Copilot lab for staff" + aliases = ("aml-rag", "aml-lab", "anti-money-laundering", "aml-copilot") + platforms = ("local",) + cost_tier = "free" + default_name = "demo" + output_keys = ("out_dir", "mode", "platform", "jurisdiction", "counts", "quality", "guide", "embedding_model") + params = [ + Param("mode", "rag", "rag (developers: indexed corpus and Python helpers) | copilot (staff: Excel, Word and PDF files)"), + Param("jurisdiction", "ca", "ca | us | both"), + Param("size", "small", "small (about 300 customers, 40k transfers) | medium (about 2,000 customers, 280k transfers)"), + Param("generator", "builtin", "builtin (pure Python, deterministic) | snapshot (pinned release from labagent-datasets)"), + Param("seed", 42, "Same seed, same data"), + Param("platform", "local", "local"), + Param("out", None, "Folder for the lab files (default ./aml-)"), + Param("regulations", "auto", "auto | fetch (download the regulation text now) | skip (policies and cases only)"), + Param("embedder", "auto", "auto (full model if installed, else keyword-only fallback) | hashing (force the fast keyword-only embedder)"), + Param("llm", "none", "none (mock answers) | env (use the chat model configured in the environment, and score it on triage)"), + Param("dataset", "aml-synthetic", "generator=snapshot: the pinned dataset name"), + Param("min_hit_rate", Q.HIT_GATE, "Retrieval hit rate the launch must reach (not applied to the keyword-only fallback)"), + ] + + # ---- helpers --------------------------------------------------------------------------------------------------------------- + def services(self, profile: Optional[str]) -> Dict[str, Any]: + return {} + + def out_dir(self, ctx: PackContext) -> Path: + return Path(ctx.params["out"] or f"aml-{ctx.name}").expanduser().resolve() + + def _check_params(self, ctx: PackContext) -> None: + p = ctx.params + if p["mode"] not in ("rag", "copilot"): + raise PackError("mode must be rag or copilot") + if p["mode"] == "copilot": + raise PackError("mode=copilot is not available in this build yet. Use mode=rag.") + if p["jurisdiction"] not in ("ca", "us", "both"): + raise PackError("jurisdiction must be ca, us or both") + if p["size"] not in ("small", "medium") and not ctx.services.get("allow_tiny"): + raise PackError("size must be small or medium") + if p["generator"] not in ("builtin", "snapshot"): + raise PackError("generator must be builtin or snapshot") + if p["platform"] != "local": + raise PackError("platform must be local") + if p["regulations"] not in ("auto", "fetch", "skip"): + raise PackError("regulations must be auto, fetch or skip") + if p["embedder"] not in ("auto", "hashing") or p["llm"] not in ("none", "env"): + raise PackError("embedder must be auto or hashing, and llm must be none or env") + try: + int(p["seed"]) + except (TypeError, ValueError): + raise PackError("seed must be a whole number") from None + + def _client(self, ctx: PackContext, out: Path): + make = ctx.services.get("client") + if make: + return make(out) + return tools.open_client(out) + + def _write_config(self, ctx: PackContext, out: Path) -> Dict[str, Any]: + p = ctx.params + cfg = {"institution": "Maple Ridge Credit Union", "name": ctx.name, "jurisdiction": p["jurisdiction"], "mode": p["mode"], + "regulations": p["regulations"] != "skip", "embedder": p["embedder"], "seed": int(p["seed"]), "size": p["size"], + "sample_alert": ctx.state.get("sample_alert", "AL00001")} + (out / tools.CONFIG).write_text(json.dumps(cfg, indent=2, sort_keys=True) + "\n") + return cfg + + # ---- steps ----------------------------------------------------------------------------------------------------------------- + def deploy_steps(self, ctx: PackContext) -> List[Step]: + out = self.out_dir(ctx) + + def preflight(c: PackContext): + self._check_params(c) + if out.exists() and not (out / MARKER).exists() and any(out.iterdir()): + raise PackError(f"{out} already exists and was not created by this pack. Choose another folder with --param out=...") + out.mkdir(parents=True, exist_ok=True) + (out / MARKER).write_text("created by labagent pack aml\n") + return {"out_dir": str(out), "mode": c.params["mode"], "platform": "local", "jurisdiction": c.params["jurisdiction"]} + + def data(c: PackContext): + try: + ds = generators.get(c.params["generator"]).generate( + seed=c.params["seed"], size=c.params["size"], jurisdiction=c.params["jurisdiction"], dataset=c.params["dataset"]) + except (datasets.DatasetError, KeyError) as e: + raise PackError(str(e)) from e + problems = validate(ds.tables) + if problems: + raise PackError("generated data failed validation: " + "; ".join(problems[:3])) + write_csv(ds.tables, out / "tables", only=[n for n in TABLES if n != "ground_truth"]) + write_csv(ds.tables, out / "answer_key", only=["ground_truth"]) + (out / "tables" / "provenance.json").write_text(json.dumps({"sources": ds.sources, "info": ds.info}, indent=2, sort_keys=True) + "\n") + first = next((g["alert_id"] for g in ds.tables["ground_truth"] if g["label"] == "escalate"), "AL00001") + return {"counts": table_counts(ds.tables), "sample_alert": first, "data_sources": ds.sources} + + def corpus(c: PackContext): + from .schema import read_csv + tables = read_csv(out / "tables", only=[n for n in TABLES if n != "ground_truth"]) + tables.update(read_csv(out / "answer_key", only=["ground_truth"])) + want_regs = c.params["regulations"] != "skip" + http = c.services.get("regs_http") or R.default_http + try: + chunks = corpus_build.build_chunks(tables, c.params["jurisdiction"], with_regulations=want_regs, http=http, cache_dir=c.services.get("regs_cache")) + except (OSError, ValueError) as e: + raise PackError(f"could not download the regulation text ({type(e).__name__}). Check your network, or launch with --param regulations=skip") from e + (out / "corpus").mkdir(exist_ok=True) + (out / "corpus" / "chunks.jsonl").write_text("\n".join(json.dumps(x, ensure_ascii=False, sort_keys=True) for x in chunks) + "\n") + docs = D.policy_documents(c.params["jurisdiction"]) + D.case_documents(tables, c.params["jurisdiction"]) + corpus_build.write_markdown(docs, out / "data") + if want_regs: + reg_by = {} + for x in chunks: + if x["meta"]["doc_type"] == "regulation": + reg_by.setdefault(x["meta"]["instrument"], []).append(x) + for inst, rows in reg_by.items(): + slug = inst.lower().replace(" ", "-") + body = [f"# {inst}", "", rows[0]["meta"]["notice"], ""] + for x in rows: + body += [f"## {x['title']}", "", x["content"], ""] + (out / "data" / f"regulation-{slug}.md").write_text("\n".join(body)) + return {"chunks": len(chunks), "documents": len(docs)} + + def index(c: PackContext): + self._write_config(c, out) + client = self._client(c, out) + rows = [json.loads(l) for l in (out / "corpus" / "chunks.jsonl").read_text().splitlines() if l.strip()] + res = client.ingest_chunks(rows) + return {"indexed_chunks": res.chunks, "embedding_model": getattr(client.embedder, "name", "unknown")} + + def quality(c: PackContext): + from .schema import read_csv + tables = read_csv(out / "tables", only=[n for n in TABLES if n != "ground_truth"]) + tables.update(read_csv(out / "answer_key", only=["ground_truth"])) + client = self._client(c, out) + q = Q.run_quality(client, tables, c.params["jurisdiction"], c.params["regulations"] != "skip", + min_hit_rate=float(c.params["min_hit_rate"]), use_chat_model=c.params["llm"] == "env", + policy_sections={s.id for d in D.policy_documents(c.params["jurisdiction"]) for s in d.sections}) + c.state["quality_full"] = q + if not q["passed"]: + raise PackError("quality check failed: " + "; ".join(q["failures"])) + return {"quality_full": q, "quality": summarize(q)} + + def guide(c: PackContext): + cfg = self._write_config(c, out) + (out / GUIDE).write_text(guide_text(cfg, c.state["quality_full"])) + (out / SHIM).write_text(SHIM_TEXT) + return {"guide": str(out / GUIDE)} + + return [ + Step("preflight", "Check the options and the output folder", preflight, check=lambda c: (out / MARKER).exists() and "out_dir" in c.state), + Step("data", "Generate the synthetic transactions, customers, alerts and answer key", data, + check=lambda c: "counts" in c.state and (out / "tables" / "transactions.csv").exists()), + Step("corpus", "Write the policies, case files and regulation text", corpus, + check=lambda c: "chunks" in c.state and (out / "corpus" / "chunks.jsonl").exists()), + Step("index", "Index the corpus for search", index, check=lambda c: "indexed_chunks" in c.state), + Step("quality", "Run the quality checks", quality, check=lambda c: bool(c.state.get("quality", {}).get("passed"))), + Step("guide", "Write the lab guide", guide, check=lambda c: (out / GUIDE).exists() and "guide" in c.state), + ] + + def destroy_steps(self, ctx: PackContext) -> List[Step]: + out = self.out_dir(ctx) + + def remove(c: PackContext): + if not out.exists(): + return {} + if not (out / MARKER).exists(): + raise PackError(f"{out} was not created by this pack, so it is left alone") + for name in OWN_PATHS: + path = out / name + if path.is_dir(): + shutil.rmtree(path) + elif path.exists(): + path.unlink() + try: + out.rmdir() # only if empty: anything the learner added stays + except OSError: + pass + return {} + + return [Step("remove", "Delete the lab files this pack wrote", remove)] + + def outputs(self, state: Dict[str, Any]) -> Dict[str, Any]: + o = {k: v for k, v in state.items() if k in self.output_keys} + o["embedding_model"] = state.get("embedding_model") + return o + + +def summarize(q: Dict[str, Any]) -> Dict[str, Any]: + chat = q["triage"].get("chat_model") + return {"passed": q["passed"], "retrieval_hit_rate": q["retrieval"]["hit_rate"], "retrieval_questions": q["retrieval"]["questions"], + "retrieval_gate": q["retrieval"]["gate"], "citation_validity": q["citations"]["rate"], + "triage_reference_agreement": q["triage"]["reference_rule"]["agreement"], "triage_alerts": q["triage"]["alerts"], + "triage_chat_model_agreement": chat["agreement"] if chat else None, "keyword_only_embedder": q["keyword_only_embedder"]} diff --git a/labagent/packs/aml/rag/__init__.py b/labagent/packs/aml/rag/__init__.py new file mode 100644 index 0000000..82e77e9 --- /dev/null +++ b/labagent/packs/aml/rag/__init__.py @@ -0,0 +1 @@ +"""RAG mode of the AML lab: the evidence bundle, the quality checks and the helpers learners call.""" diff --git a/labagent/packs/aml/rag/bundle.py b/labagent/packs/aml/rag/bundle.py new file mode 100644 index 0000000..6311bd3 --- /dev/null +++ b/labagent/packs/aml/rag/bundle.py @@ -0,0 +1,74 @@ +"""The evidence bundle for one alert: policy, customer profile, alert trigger, transactions and similar earlier cases. + +Design follows the triage paper (arXiv 2604.19755): everything a reviewer needs to decide, with a source for each part. Nothing here reads +the label, pattern or reason of the alert itself. "Similar cases" are earlier alerts with the same rule: their decisions were already made. +""" +from __future__ import annotations + +from typing import Any, Callable, Dict, List, Optional + +from ..corpus.documents import CaseIndex +from ..schema import RULES, Tables + +# retrieve(question, k, where) -> [{"chunk_id", "text", "metadata"}] +Retrieve = Callable[[str, int, Optional[Dict[str, Any]]], List[Dict[str, Any]]] + + +def client_retriever(client: Any) -> Retrieve: + """Adapt RAGLabClient.retrieve() to the shape the bundle uses.""" + def retrieve(question: str, k: int = 3, where: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]: + return [{"chunk_id": c.chunk_id, "text": c.text, "metadata": c.metadata or {}} for c in client.retrieve(question, top_k=k, where=where)] + return retrieve + + +def jurisdiction_of(ix: CaseIndex, alert: Dict[str, Any], jurisdiction: str) -> str: + if jurisdiction != "both": + return jurisdiction + cur = next((x["currency"] for x in ix.by_account.get(alert["account"], []) + if x["from_account"] == alert["account"] or x["to_account"] == alert["account"]), "CAD") + return "us" if cur == "USD" else "ca" + + +def evidence_bundle(t: Tables, alert_id: str, retrieve: Retrieve, jurisdiction: str, index: Optional[CaseIndex] = None, k: int = 3, + with_regulations: bool = True) -> Dict[str, Any]: + ix = index or CaseIndex(t) + alert = next((a for a in t["alerts"] if a["id"] == alert_id), None) + if alert is None: + raise KeyError(f"no alert '{alert_id}'") + j = jurisdiction_of(ix, alert, jurisdiction) + cust = ix.customer(alert) + rows, trimmed = ix.rows(alert) + rule_text = RULES[alert["rule"]] + policy = retrieve(f"{rule_text}: what does the policy say?", k, {"doc_type": "policy", "jurisdiction": j}) + regs = retrieve(f"{rule_text} reporting requirement", k, {"doc_type": "regulation", "jurisdiction": j}) if with_regulations else [] + truth = {g["alert_id"]: g for g in t.get("ground_truth", [])} + earlier = sorted((a for a in t["alerts"] if a["rule"] == alert["rule"] and a["date"] < alert["date"]), key=lambda a: a["date"])[-3:] + return { + "alert": alert, "rule": rule_text, "customer": cust, "kyc": ix.kyc[cust["id"]], "jurisdiction": j, + "sanctions": [h for h in t["sanctions_hits"] if h["customer_id"] == cust["id"]], + "transactions": rows, "trimmed": trimmed, + "uncommon_parties": sorted({p for x in rows for p in (x["from_account"], x["to_account"]) if ix.uncommon(p)}), + "policy": policy, "regulations": regs, + "similar_cases": [{"alert_id": a["id"], "date": a["date"], "decision": truth[a["id"]]["label"], "reason": truth[a["id"]]["reason"]} for a in earlier if a["id"] in truth], + } + + +def render(b: Dict[str, Any], max_rows: int = 40) -> str: + """Plain text for a person or a prompt.""" + a, c = b["alert"], b["customer"] + lines = [f"ALERT {a['id']} on {a['date']}: {b['rule']} (rule {a['rule']}), amount {a['amount']}, score {a['risk_score']}", + f"CUSTOMER {c['id']}: {b['kyc']}"] + if b["sanctions"]: + lines += [f"SANCTIONS SCREENING: {h['match_type']} match to {h['sdn_name']} (score {h['score']})" for h in b["sanctions"]] + lines.append("POLICY:") + lines += [f" [{p['chunk_id']}] {p['text'][:300]}" for p in b["policy"]] + if b["regulations"]: + lines.append("REGULATIONS:") + lines += [f" [{p['metadata'].get('citation', p['chunk_id'])}] {p['text'][:240]}" for p in b["regulations"]] + lines.append("TRANSACTIONS:") + lines += [f" {x['id']} {x['timestamp']} {x['from_account']} -> {x['to_account']} {x['amount']} {x['currency']} {x['channel']} {x['description']}" + for x in b["transactions"][-max_rows:]] + if b["similar_cases"]: + lines.append("SIMILAR EARLIER CASES:") + lines += [f" {s['alert_id']} ({s['date']}): {s['decision']}. {s['reason']}" for s in b["similar_cases"]] + return "\n".join(lines) diff --git a/labagent/packs/aml/rag/qa.py b/labagent/packs/aml/rag/qa.py new file mode 100644 index 0000000..f0c58d4 --- /dev/null +++ b/labagent/packs/aml/rag/qa.py @@ -0,0 +1,37 @@ +"""Questions with known answers, for the retrieval check. Each one names the metadata a retrieved chunk must carry to count as a hit.""" +from __future__ import annotations + +from typing import Any, Dict, List + +POLICY = [ + ("Who owns the AML program and who do they report to?", "policy", "2"), + ("What must staff do before the first transaction for a new customer?", "policy", "3"), + ("How often is a customer risk rating reviewed?", "policy", "4"), + ("Which rules raise a transaction monitoring alert?", "policy", "5"), + ("When must a cash transaction report be filed?", "policy", "6"), + ("What happens when a customer name is a close match to a sanctions list entry?", "policy", "7"), + ("When should an investigator escalate instead of monitor?", "policy", "8"), + ("How long must records be kept?", "policy", "9"), + ("Can staff tell a customer that a report was filed?", "policy", "10"), + ("How soon must an investigator open a new alert?", "procedure", "2"), + ("What are the six questions a report narrative must answer?", "guide", "2"), +] + +REGULATIONS = { + "ca": [("Which section of the Act requires reporting a transaction when there are reasonable grounds to suspect money laundering?", "PCMLTFA s. 7"), + ("When must a suspicious transaction report be sent to the Centre?", "PCMLTF STR Regulations s. 9"), + ("What must a financial entity keep when it receives $10,000 or more in cash?", "PCMLTF Regulations s. 10")], + "us": [("When must a bank file a suspicious activity report?", "31 CFR 1020.320"), + ("What is the currency transaction reporting requirement for a financial institution?", "31 CFR 1010.311"), + ("Which rule forbids structuring transactions to evade reporting?", "31 CFR 1010.314")], +} + + +def policy_qa(jurisdiction: str, with_regulations: bool) -> List[Dict[str, Any]]: + """Cases in the shape RAGLabClient.evaluate() takes: {question, expect}. With jurisdiction 'both' every question is asked per country.""" + out: List[Dict[str, Any]] = [] + for j in (["ca", "us"] if jurisdiction == "both" else [jurisdiction]): + out += [{"question": q, "expect": {"doc_type": kind, "section": sec, "jurisdiction": j}} for q, kind, sec in POLICY] + if with_regulations: + out += [{"question": q, "expect": {"doc_type": "regulation", "citation": cite, "jurisdiction": j}} for q, cite in REGULATIONS[j]] + return out diff --git a/labagent/packs/aml/rag/quality.py b/labagent/packs/aml/rag/quality.py new file mode 100644 index 0000000..94771be --- /dev/null +++ b/labagent/packs/aml/rag/quality.py @@ -0,0 +1,201 @@ +"""Quality checks for the AML RAG lab. Three checks, reported in the result envelope: + +1. Retrieval hit rate on a policy and regulation Q&A set. A hard gate (default 0.6, like the SEC tutorial) when the embedding model + understands meaning. With the keyword-only fallback embedder it is reported but does not fail the launch. +2. Citation validity: every transaction ID, regulation section and policy section a text cites must exist. A hard gate, run on a model + narrative built from the evidence bundle, so it proves the corpus supports the citations a good answer needs. +3. Triage agreement: escalate, monitor or dismiss compared with ground_truth on a small, balanced set of alerts. Reported only. + A deterministic reference rule shows the evidence is enough to decide; a chat model is scored the same way when one is configured. + The labels come from how the data was built, so this is a check that the lab is solvable, not a benchmark of real-world accuracy. +""" +from __future__ import annotations + +import re +from datetime import date +from typing import Any, Callable, Dict, List, Optional + +from ..corpus import regs as R +from ..corpus.documents import CaseIndex +from ..schema import LABELS, Tables +from .bundle import Retrieve, evidence_bundle, render +from .qa import policy_qa + +HIT_GATE = 0.6 +TX_RE = re.compile(r"\bT\d{7}\b") +POLICY_SECTION_RE = re.compile(r"\b(?:policy(?: manual)?|procedure|guide) (?:section|§) ?(\d+)\b", re.I) +LABEL_RE = re.compile(r"\b(escalate|monitor|dismiss)\b", re.I) +BENIGN = {"REAL ESTATE CLOSING PROCEEDS", "PAYROLL BONUS", "TUITION PAYMENT"} + + +# ---- 1. retrieval -------------------------------------------------------------------------------------------------------------- + +def retrieval_check(client: Any, jurisdiction: str, with_regulations: bool, top_k: int = 5) -> Dict[str, Any]: + cases = policy_qa(jurisdiction, with_regulations) + res = client.evaluate(cases, top_k=top_k) + return {"questions": res.cases, "hits": res.hits, "hit_rate": res.hit_rate, "misses": [m["question"] for m in res.misses][:10]} + + +# ---- 2. citations -------------------------------------------------------------------------------------------------------------- + +def citation_validity(text: str, t: Tables, index: Optional[Dict[str, Any]] = None, policy_sections: Optional[set] = None) -> Dict[str, Any]: + """Check every transaction ID, regulation section and 'policy section N' reference in `text`. rate is valid / cited (1.0 when nothing is cited).""" + index = index or R.load_index() + tx_ids = {x["id"] for x in t["transactions"]} + policy_sections = policy_sections if policy_sections is not None else {str(n) for n in range(1, 12)} + cited_tx = sorted(set(TX_RE.findall(text))) + bad_tx = [x for x in cited_tx if x not in tx_ids] + cited_reg = [f"{i} {n}" for i, nums in R.cited_sections(text).items() for n in nums] + bad_reg = R.missing_citations(text, index) + cited_pol = sorted(set(POLICY_SECTION_RE.findall(text))) + bad_pol = [s for s in cited_pol if s not in policy_sections] + total = len(cited_tx) + len(cited_reg) + len(cited_pol) + invalid = bad_tx + bad_reg + [f"policy section {s}" for s in bad_pol] + return {"cited": total, "invalid": invalid, "rate": 1.0 if total == 0 else round((total - len(invalid)) / total, 3), + "transactions": len(cited_tx), "regulations": len(cited_reg), "policy_sections": len(cited_pol)} + + +def model_narrative(b: Dict[str, Any]) -> str: + """A reference narrative built only from the bundle. Every ID and section it cites comes from retrieved or listed evidence.""" + a, c = b["alert"], b["customer"] + acct = a["account"] + mine = [x for x in b["transactions"] if acct in (x["from_account"], x["to_account"])] + top = sorted(mine, key=lambda x: -float(x["amount"]))[:5] + top.sort(key=lambda x: x["timestamp"]) + ids = ", ".join(x["id"] for x in top) + total = sum(float(x["amount"]) for x in top) + section = next((p["metadata"].get("section") for p in b["policy"] if p["metadata"].get("doc_type") == "policy"), None) + reg = next((p["metadata"].get("citation") for p in b["regulations"]), None) + out = [f"Alert {a['id']} on account {acct} for {b['rule'].lower()}.", + f"Between {top[0]['timestamp'][:10]} and {top[-1]['timestamp'][:10]} the account had large transfers ({ids}), about {total:,.0f} in total." if top else + "The account had no transfers in the review window.", + f"The profile expects about {int(c['expected_monthly_activity']):,} a month."] + if section: + out.append(f"Policy section {section} applies.") + if reg: + out.append(f"See {reg}.") + return " ".join(out) + + +# ---- 3. triage ----------------------------------------------------------------------------------------------------------------- + +def _loop_of_three_or_more(acct: str, edges: Dict[str, set], limit: int = 7) -> bool: + def walk(node: str, depth: int, path: List[str]) -> bool: + for nxt in sorted(edges.get(node, ())): + if nxt == acct and depth >= 3: + return True + if nxt not in path and depth < limit and walk(nxt, depth + 1, path + [nxt]): + return True + return False + return walk(acct, 1, [acct]) + + +def triage_baseline(b: Dict[str, Any]) -> str: + """A fixed reference rule that uses only the bundle. It shows the evidence is enough to decide; it is not a model.""" + a, c = b["alert"], b["customer"] + rule, acct = a["rule"], a["account"] + if rule == "R09": + return "escalate" if any(h["match_type"] == "exact" for h in b["sanctions"]) else "monitor" + ratio = float(a["amount"]) / max(float(c["expected_monthly_activity"]), 1.0) + unusual = set(b["uncommon_parties"]) - {acct} + end = date.fromisoformat(a["date"]) + rows = b["transactions"] + days = lambda x: (end - date.fromisoformat(x["timestamp"][:10])).days + recent = [x for x in rows if acct in (x["from_account"], x["to_account"]) and 0 <= days(x) <= 21] + ins = [x for x in recent if x["to_account"] == acct] + outs = [x for x in recent if x["from_account"] == acct] + if rule in ("R06", "R07") and {x["description"] for x in recent} & BENIGN: + return "dismiss" + if rule == "R08": + for x in outs: + same = [y for y in rows if y["to_account"] == x["to_account"] and y["from_account"] == acct and y["amount"] == x["amount"] and y["timestamp"] < x["timestamp"]] + if len(same) >= 2: + return "dismiss" # the same amount went to the same payee before, so the payee is not new + if any(x["description"].startswith("SUPPLIES INVOICE") and 9000 <= float(x["amount"]) < 10000 for x in outs): + return "escalate" + return "monitor" + if rule == "R03": + return "escalate" if len({x["from_account"] for x in ins} & unusual) >= 3 and ratio >= 1 else "monitor" + if rule == "R04": + return "escalate" if len({x["to_account"] for x in outs} & unusual) >= 3 and ratio >= 1 else "monitor" + if rule == "R05": + edges: Dict[str, set] = {} + for x in rows: + if x["from_account"] in unusual | {acct} and x["to_account"] in unusual | {acct}: + edges.setdefault(x["from_account"], set()).add(x["to_account"]) + return "escalate" if ratio >= 1 and _loop_of_three_or_more(acct, edges) else "monitor" + if rule == "R06": + return "escalate" if len({x["from_account"] for x in ins} & unusual) >= 3 else "monitor" + if rule == "R02": + legs = [x for x in rows if x["from_account"] in unusual | {acct} and x["to_account"] in unusual | {acct}] + return "escalate" if len(legs) >= 3 and ratio >= 1 else "monitor" + return "monitor" + + +def triage_llm(client: Any, b: Dict[str, Any]) -> Optional[str]: + prompt = ("You review an anti-money laundering alert. Read the evidence and answer with one word: escalate, monitor or dismiss. " + "Then give the transaction IDs and policy section you relied on.") + res = client.llm.complete(prompt, render(b)) + m = LABEL_RE.search(res.text or "") + return m.group(1).lower() if m else None + + +def sample_alerts(t: Tables, per_label: int = 4) -> List[str]: + """A small balanced set: the first `per_label` alerts of each label, by alert id.""" + chosen: Dict[str, List[str]] = {l: [] for l in LABELS} + for g in sorted(t["ground_truth"], key=lambda g: g["alert_id"]): + if len(chosen[g["label"]]) < per_label: + chosen[g["label"]].append(g["alert_id"]) + return sorted(i for ids in chosen.values() for i in ids) + + +def triage_check(t: Tables, retrieve: Retrieve, jurisdiction: str, with_regulations: bool, client: Any = None, per_label: int = 4, + index: Optional[CaseIndex] = None) -> Dict[str, Any]: + ix = index or CaseIndex(t) + truth = {g["alert_id"]: g["label"] for g in t["ground_truth"]} + ids = sample_alerts(t, per_label) + base_ok = llm_ok = llm_n = 0 + wrong: List[str] = [] + for aid in ids: + b = evidence_bundle(t, aid, retrieve, jurisdiction, ix, with_regulations=with_regulations) + if triage_baseline(b) == truth[aid]: + base_ok += 1 + else: + wrong.append(aid) + if client is not None: + guess = triage_llm(client, b) + llm_n += 1 + llm_ok += guess == truth[aid] + out = {"alerts": len(ids), "reference_rule": {"agreement": round(base_ok / len(ids), 3), "disagreed": wrong[:10]}, + "chat_model": None} + if client is not None and llm_n: + out["chat_model"] = {"agreement": round(llm_ok / llm_n, 3), "alerts": llm_n} + return out + + +# ---- all three ----------------------------------------------------------------------------------------------------------------- + +def run_quality(client: Any, t: Tables, jurisdiction: str, with_regulations: bool, min_hit_rate: float = HIT_GATE, use_chat_model: bool = False, + policy_sections: Optional[set] = None) -> Dict[str, Any]: + """Returns {retrieval, citations, triage, passed, failures}. `passed` covers the two hard gates only.""" + from .bundle import client_retriever + retrieve = client_retriever(client) + fallback = str(getattr(client.embedder, "name", "")).startswith("hashing") + retrieval = retrieval_check(client, jurisdiction, with_regulations) + retrieval["gate"] = None if fallback else min_hit_rate + ix = CaseIndex(t) + narratives = [] + for aid in sample_alerts(t, 3): + b = evidence_bundle(t, aid, retrieve, jurisdiction, ix, with_regulations=with_regulations) + narratives.append(citation_validity(model_narrative(b), t, policy_sections=policy_sections)) + cited = sum(n["cited"] for n in narratives) + invalid = [x for n in narratives for x in n["invalid"]] + citations = {"narratives": len(narratives), "cited": cited, "invalid": invalid[:10], + "rate": 1.0 if not cited else round((cited - len(invalid)) / cited, 3)} + triage = triage_check(t, retrieve, jurisdiction, with_regulations, client if use_chat_model else None, index=ix) + failures = [] + if not fallback and retrieval["hit_rate"] < min_hit_rate: + failures.append(f"retrieval hit rate {retrieval['hit_rate']:.2f} is below the {min_hit_rate} gate") + if citations["rate"] < 1.0: + failures.append(f"citation validity {citations['rate']:.2f}: invalid citations {citations['invalid']}") + return {"retrieval": retrieval, "citations": citations, "triage": triage, "embedding_model": getattr(client.embedder, "name", "unknown"), + "keyword_only_embedder": fallback, "passed": not failures, "failures": failures} diff --git a/labagent/packs/aml/rag/tools.py b/labagent/packs/aml/rag/tools.py new file mode 100644 index 0000000..6628e36 --- /dev/null +++ b/labagent/packs/aml/rag/tools.py @@ -0,0 +1,80 @@ +"""What a learner calls from Python or a notebook after `labagent launch aml`. + + from aml_tools import AmlLab + lab = AmlLab() # run it in the lab folder + lab.ask("When must a suspicious transaction report be filed?") + lab.show("AL00007") # the evidence bundle for one alert + lab.check(my_narrative) # are the transaction IDs and sections I cited real? + lab.score({"AL00007": "escalate"}) # compare my decisions with the answer key +""" +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Dict, List, Optional, Union + +from ..corpus.documents import CaseIndex +from ..schema import TABLES, read_csv +from . import quality as Q +from .bundle import client_retriever, evidence_bundle, render + +CONFIG = "config.json" +CHUNKS = "corpus/chunks.jsonl" + + +def open_client(out_dir: Path, settings_update: Optional[Dict[str, Any]] = None): + """Open the lab's index. If the index is empty (a fresh checkout, or a memory store), rebuild it from corpus/chunks.jsonl.""" + from ...ai.rag import RAGLabClient + from ...ai.rag.config import RAGSettings + cfg = json.loads((out_dir / CONFIG).read_text()) + s = RAGSettings.from_env().model_copy(update={"chroma_path": str(out_dir / "index"), "collection": "aml_lab", + **({"embedding_model_name": "hashing"} if cfg.get("embedder") == "hashing" else {}), + **(settings_update or {})}) + client = RAGLabClient(settings=s) + if client.store.count() == 0: + rows = [json.loads(l) for l in (out_dir / CHUNKS).read_text().splitlines() if l.strip()] + client.ingest_chunks(rows) + return client + + +class AmlLab: + def __init__(self, out_dir: Union[str, Path] = ".", client: Any = None): + self.dir = Path(out_dir) + self.config = json.loads((self.dir / CONFIG).read_text()) + self.tables = read_csv(self.dir / "tables", only=[n for n in TABLES if n != "ground_truth"]) + key = self.dir / "answer_key" + if (key / "ground_truth.csv").exists(): # the answer key is kept apart: open it after you have decided + self.tables.update(read_csv(key, only=["ground_truth"])) + self.client = client or open_client(self.dir) + self._ix = CaseIndex(self.tables) + + def alerts(self) -> List[Dict[str, Any]]: + """Every alert, without the answers.""" + return list(self.tables["alerts"]) + + def ask(self, question: str, top_k: int = 5, **filters: Any): + """Ask the corpus. Filters narrow the search, for example doc_type='policy' or jurisdiction='ca'.""" + return self.client.query(question, top_k=top_k, where=filters or None) + + def evidence(self, alert_id: str) -> Dict[str, Any]: + return evidence_bundle(self.tables, alert_id, client_retriever(self.client), self.config["jurisdiction"], self._ix, + with_regulations=self.config["regulations"]) + + def show(self, alert_id: str) -> None: + print(render(self.evidence(alert_id))) + + def check(self, text: str) -> Dict[str, Any]: + """Citation check: every transaction ID, regulation section and 'policy section N' you cite must exist.""" + return Q.citation_validity(text, self.tables) + + def triage(self, alert_id: str) -> Optional[str]: + """Ask the configured chat model for escalate, monitor or dismiss. Returns None when the model gives no clear answer.""" + return Q.triage_llm(self.client, self.evidence(alert_id)) + + def score(self, decisions: Dict[str, str]) -> Dict[str, Any]: + """Compare your decisions with the answer key. Look at this after you have decided, not before.""" + truth = {g["alert_id"]: g["label"] for g in self.tables["ground_truth"]} if "ground_truth" in self.tables else {} + known = {k: v for k, v in decisions.items() if k in truth} + right = sorted(k for k, v in known.items() if truth[k] == v) + return {"decided": len(known), "agree": len(right), "agreement": round(len(right) / len(known), 3) if known else None, + "disagree": {k: {"you": v, "key": truth[k]} for k, v in known.items() if truth[k] != v}} diff --git a/labagent/packs/aml/schema.py b/labagent/packs/aml/schema.py index 23ae4c9..cdb64a9 100644 --- a/labagent/packs/aml/schema.py +++ b/labagent/packs/aml/schema.py @@ -53,11 +53,13 @@ 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.""" +def write_csv(tables: Tables, out_dir: Path, only: Optional[List[str]] = None) -> List[Path]: + """One CSV per table (or just the `only` tables). 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(): + if only is not None and name not in only: + continue path = out_dir / f"{name}.csv" with open(path, "w", newline="", encoding="utf-8") as f: w = csv.DictWriter(f, fieldnames=cols, lineterminator="\n") @@ -81,9 +83,11 @@ def write_parquet(tables: Tables, out_dir: Path) -> List[Path]: return paths -def read_csv(in_dir: Path) -> Tables: +def read_csv(in_dir: Path, only: Optional[List[str]] = None) -> Tables: out: Tables = {} for name in TABLES: + if only is not None and name not in only: + continue with open(in_dir / f"{name}.csv", newline="", encoding="utf-8") as f: out[name] = list(csv.DictReader(f)) return out diff --git a/labagent/packs/registry.py b/labagent/packs/registry.py index 124e6f1..0f310a3 100644 --- a/labagent/packs/registry.py +++ b/labagent/packs/registry.py @@ -10,7 +10,8 @@ def builtin() -> Dict[str, Pack]: from .rag.tutorial_pack import RagSecTutorial from .db_agent.pack import DbAgent from .rag_ide.pack import RagIde - return {p.id: p for p in (RagSecTutorial(), DbAgent(), RagIde())} + from .aml.pack import AmlLab + return {p.id: p for p in (RagSecTutorial(), DbAgent(), RagIde(), AmlLab())} def all_packs() -> Dict[str, Pack]: diff --git a/tests/test_aml_rag.py b/tests/test_aml_rag.py new file mode 100644 index 0000000..310545c --- /dev/null +++ b/tests/test_aml_rag.py @@ -0,0 +1,245 @@ +"""AML pack, part 3: the RAG mode. Uses the keyword-only embedder and an in-memory store, so no model download and no network.""" +import json +from pathlib import Path + +import pytest + +from labagent.ai.rag import RAGLabClient +from labagent.ai.rag.config import RAGSettings +from labagent.cli import main +from labagent.packs import registry +from labagent.packs.aml.corpus import documents as D +from labagent.packs.aml.corpus.documents import CaseIndex +from labagent.packs.aml.gen import get +from labagent.packs.aml.pack import MARKER, AmlLab +from labagent.packs.aml.rag import quality as Q +from labagent.packs.aml.rag.bundle import client_retriever, evidence_bundle, render +from labagent.packs.aml.rag.qa import policy_qa +from labagent.packs.aml.rag.tools import AmlLab as Lab +from labagent.packs.base import PackContext, load_state, run_pack, state_path +from tests.test_aml_corpus import ECFR, JUSTICE + + +@pytest.fixture +def home(tmp_path, monkeypatch): + monkeypatch.setenv("LABAGENT_HOME", str(tmp_path / "home")) + monkeypatch.setenv("EMBEDDING_MODEL_NAME", "hashing") + monkeypatch.delenv("CHROMA_PATH", raising=False) + return tmp_path + + +def memory_client(out=None): + return RAGLabClient(vector_db="memory", llm_provider="mock", settings=RAGSettings.from_env()) + + +def services(tmp_path): + return {"allow_tiny": True, "client": memory_client, "regs_http": lambda url: ECFR if "ECFR" in url else JUSTICE, "regs_cache": tmp_path / "regs"} + + +def launch(tmp_path, **params): + base = {"size": "tiny", "out": str(tmp_path / "lab"), "embedder": "hashing"} + base.update(params) + return run_pack(AmlLab(), "launch", "t", base, services=services(tmp_path)) + + +@pytest.fixture(scope="module") +def tables(): + return get("builtin").generate(seed=42, size="small", jurisdiction="ca").tables + + +def built_client(tables, with_regs=False, j="ca"): + from labagent.packs.aml.corpus.build import build_chunks + c = memory_client() + c.ingest_chunks(build_chunks(tables, j, with_regulations=False)) + return c + + +# ---- the pack ------------------------------------------------------------------------------------------------------------------ + +def test_pack_is_registered_with_its_catalog_entry(capsys, home): + assert registry.resolve("aml-rag").id == "aml" and registry.resolve("anti-money-laundering").id == "aml" + assert main(["packs", "--json"]) == 0 + entry = next(p for p in json.loads(capsys.readouterr().out)["packs"] if p["id"] == "aml") + assert entry["platforms"] == ["local"] and entry["cost_tier"] == "free" + assert {"mode", "jurisdiction", "size", "generator", "seed"} <= {q["name"] for q in entry["params"]} + + +def test_launch_builds_the_lab_folder_and_reports_quality(home): + env = launch(home) + assert env.status == "ok", env.error + out = home / "lab" + assert (out / MARKER).exists() and (out / "tables" / "transactions.csv").exists() and (out / "data" / "policy-ca.md").exists() + assert (out / "answer_key" / "ground_truth.csv").exists() and not (out / "tables" / "ground_truth.csv").exists() + assert (out / "LAB_GUIDE.md").exists() and (out / "aml_tools.py").exists() and (out / "corpus" / "chunks.jsonl").exists() + o = env.summary["outputs"] + assert o["mode"] == "rag" and o["counts"]["customers"] == 60 and Path(o["out_dir"]) == out.resolve() + q = o["quality"] + assert q["passed"] and q["citation_validity"] == 1.0 and 0 <= q["retrieval_hit_rate"] <= 1 and q["triage_reference_agreement"] >= 0.8 + assert q["keyword_only_embedder"] is True and q["triage_chat_model_agreement"] is None + + +def test_launch_is_resumable_and_a_second_run_skips_everything(home): + launch(home) + again = launch(home) + assert again.status == "ok" and again.summary["skipped"] == 6 and again.summary["done"] == 0 + + +def test_regulation_text_is_in_the_lab_unless_skipped(home): + launch(home) + assert (home / "lab" / "data" / "regulation-pcmltfa.md").exists() + assert "Open Government Licence - Canada" in (home / "lab" / "data" / "regulation-pcmltfa.md").read_text() + launch(home, out=str(home / "lab2"), regulations="skip") + assert not list((home / "lab2" / "data").glob("regulation-*")) + assert not any(json.loads(l)["meta"]["doc_type"] == "regulation" for l in (home / "lab2" / "corpus" / "chunks.jsonl").read_text().splitlines()) + + +def test_a_network_failure_is_a_clear_step_failure(home): + def down(url): + raise OSError("no route") + env = run_pack(AmlLab(), "launch", "t", {"size": "tiny", "out": str(home / "lab")}, services={**services(home), "regs_http": down}) + assert env.status == "failed" and "regulations=skip" in env.error and "no route" not in env.error + + +def test_bad_options_fail_early_with_a_plain_message(home): + for bad, msg in (({"mode": "x"}, "mode"), ({"jurisdiction": "fr"}, "jurisdiction"), ({"seed": "abc"}, "seed"), ({"generator": "magic"}, "generator"), + ({"mode": "copilot"}, "not available")): + env = launch(home, out=str(home / "x"), **bad) + assert env.status == "failed" and msg in env.error, (bad, env.error) + env = run_pack(AmlLab(), "launch", "t", {"size": "tiny", "out": str(home / "x")}, services={**services(home), "allow_tiny": False}) + assert "small or medium" in env.error + + +def test_it_refuses_a_folder_it_did_not_create(home): + (home / "mine").mkdir() + (home / "mine" / "notes.txt").write_text("keep") + env = launch(home, out=str(home / "mine")) + assert env.status == "failed" and "was not created by this pack" in env.error and (home / "mine" / "notes.txt").read_text() == "keep" + + +def test_a_failed_quality_gate_stops_the_launch_and_names_the_reason(home, monkeypatch): + monkeypatch.setattr(Q, "retrieval_check", lambda *a, **k: {"questions": 5, "hits": 0, "hit_rate": 0.0, "misses": []}) + + def client(out=None): + c = memory_client() # pretend the embedder understands meaning, so the gate applies + c.embedder.name = "all-MiniLM-L6-v2" + return c + env = run_pack(AmlLab(), "launch", "t", {"size": "tiny", "out": str(home / "lab")}, services={**services(home), "client": client}) + assert env.status == "failed" and "retrieval hit rate 0.00 is below the 0.6 gate" in env.error + assert "guide" not in load_state("aml", "t") + + +def test_destroy_removes_only_what_the_pack_wrote(home): + launch(home) + (home / "lab" / "my_notes.md").write_text("mine") + env = run_pack(AmlLab(), "destroy", "t", {"out": str(home / "lab")}, services={}) + assert env.status == "ok" + assert [p.name for p in (home / "lab").iterdir()] == ["my_notes.md"] and not state_path("aml", "t").exists() + launch(home, out=str(home / "lab3")) + run_pack(AmlLab(), "destroy", "t", {"out": str(home / "lab3")}, services={}) + assert not (home / "lab3").exists() # nothing of the learner's was there, so the folder goes too + + +def test_destroy_leaves_alone_a_folder_without_the_marker(home): + (home / "other").mkdir() + (home / "other" / "tables").mkdir() + env = run_pack(AmlLab(), "destroy", "t", {"out": str(home / "other")}, services={}) + assert env.status == "failed" and (home / "other" / "tables").exists() + + +def test_envelope_carries_no_secret_or_answer(home): + env = launch(home) + text = json.dumps(env.to_dict()) + state = json.dumps(load_state("aml", "t")) + for g in get("builtin").generate(seed=42, size="tiny").tables["ground_truth"][:5]: + assert g["reason"] not in text and g["reason"] not in state + + +# ---- the quality library -------------------------------------------------------------------------------------------------------- + +def test_evidence_bundle_has_every_part_and_never_reads_the_alerts_own_answer(tables): + blind = {k: v for k, v in tables.items() if k != "ground_truth"} + c = built_client(tables) + aid = tables["alerts"][20]["id"] + b = evidence_bundle(blind, aid, client_retriever(c), "ca", with_regulations=False) + assert set(b) >= {"alert", "customer", "kyc", "sanctions", "transactions", "policy", "regulations", "similar_cases", "uncommon_parties"} + assert b["policy"] and all(p["metadata"]["doc_type"] == "policy" and p["metadata"]["jurisdiction"] == "ca" for p in b["policy"]) + assert b["similar_cases"] == [] # without an answer key there is nothing to recall + full = evidence_bundle(tables, aid, client_retriever(c), "ca", with_regulations=False) + assert all(s["date"] < b["alert"]["date"] for s in full["similar_cases"]) and aid not in {s["alert_id"] for s in full["similar_cases"]} + assert aid in render(full) and "ALERT" in render(full) + + +def test_citation_check_catches_invented_ids_and_sections(tables): + real = tables["transactions"][100]["id"] + ok = Q.citation_validity(f"See {real}, PCMLTFA s. 7 and policy section 8.", tables) + assert ok["rate"] == 1.0 and ok["cited"] == 3 + bad = Q.citation_validity(f"See {real}, T9999999, PCMLTFA s. 777 and policy section 40.", tables) + assert bad["rate"] == 0.25 and set(bad["invalid"]) == {"T9999999", "PCMLTFA 777", "policy section 40"} + assert Q.citation_validity("Nothing cited here.", tables)["rate"] == 1.0 + + +def test_reference_rule_agrees_with_the_answer_key_on_many_datasets(): + right = total = 0 + for seed in (42, 1, 2, 3): + for j in ("ca", "both"): + t = get("builtin").generate(seed=seed, size="small", jurisdiction=j).tables + ix = CaseIndex(t) + for g in t["ground_truth"]: + b = evidence_bundle(t, g["alert_id"], lambda q, k, w: [], j, ix, with_regulations=False) + right += Q.triage_baseline(b) == g["label"] + total += 1 + assert right / total >= 0.97 + + +def test_sample_alerts_are_balanced(tables): + ids = Q.sample_alerts(tables, 4) + truth = {g["alert_id"]: g["label"] for g in tables["ground_truth"]} + assert len(ids) == 12 and sorted(truth[i] for i in ids).count("monitor") == 4 + + +def test_chat_model_triage_is_scored_the_same_way(tables): + class Res: + text = "I would escalate this one. T0000001, policy section 8." + + class Llm: + def complete(self, prompt, context): + return Res() + + c = built_client(tables) + c.llm = Llm() + r = Q.triage_check(tables, client_retriever(c), "ca", False, client=c, per_label=2) + assert r["chat_model"]["alerts"] == 6 and r["chat_model"]["agreement"] == pytest.approx(2 / 6, abs=0.01) # says escalate every time + + +def test_qa_set_covers_both_countries_and_skips_regulations_when_asked(): + both = policy_qa("both", True) + assert {c["expect"]["jurisdiction"] for c in both} == {"ca", "us"} and any("citation" in c["expect"] for c in both) + assert not any("citation" in c["expect"] for c in policy_qa("ca", False)) + + +# ---- the learner helper --------------------------------------------------------------------------------------------------------- + +def test_learner_helper_works_from_the_lab_folder(home): + launch(home) + lab = Lab(home / "lab", client=memory_client()) + lab.client.ingest_chunks([json.loads(l) for l in (home / "lab" / "corpus" / "chunks.jsonl").read_text().splitlines() if l.strip()]) + assert len(lab.alerts()) > 10 and "label" not in lab.alerts()[0] + r = lab.ask("When must a suspicious transaction report be filed?", doc_type="policy") + assert r.chunks and all(c.metadata["doc_type"] == "policy" for c in r.chunks) + b = lab.evidence(lab.alerts()[3]["id"]) + assert b["transactions"] and b["policy"] + first = lab.alerts()[0]["id"] + truth = {g["alert_id"]: g["label"] for g in lab.tables["ground_truth"]} + assert lab.score({first: truth[first]})["agreement"] == 1.0 + wrong = lab.score({first: "dismiss" if truth[first] != "dismiss" else "escalate"}) + assert wrong["agreement"] == 0.0 and wrong["disagree"][first]["key"] == truth[first] + assert lab.check(f"See {lab.tables['transactions'][0]['id']} and T9999999.")["invalid"] == ["T9999999"] + + +def test_the_guide_names_real_commands_and_follows_the_writing_rules(home): + launch(home) + text = (home / "lab" / "LAB_GUIDE.md").read_text() + for needle in ("lab.ask(", "lab.show(", "lab.check(", "lab.score(", "labagent destroy aml --name t", "keyword-only fallback"): + assert needle in text + assert "—" not in text and "–" not in text + assert "rag-ide" not in text diff --git a/tests/test_packs.py b/tests/test_packs.py index f5faf33..3308ce8 100644 --- a/tests/test_packs.py +++ b/tests/test_packs.py @@ -75,7 +75,7 @@ def go(c): def test_registry_lists_the_current_packs_and_the_cli_refuses_unknown_ones(capsys, home): - assert sorted(registry.all_packs()) == ["db-agent", "rag-ide", "rag-sec-tutorial"] + assert sorted(registry.all_packs()) == ["aml", "db-agent", "rag-ide", "rag-sec-tutorial"] assert main(["packs"]) == 0 and "rag-sec-tutorial" in capsys.readouterr().out assert main(["launch", "nope", "--name", "x"]) == 2 @@ -149,7 +149,7 @@ def test_an_explicit_profile_is_never_overridden(home): def test_packs_json_describes_every_pack_with_platforms_and_params(capsys, home): assert main(["packs", "--json"]) == 0 cat = {p["id"]: p for p in json.loads(capsys.readouterr().out)["packs"]} - assert sorted(cat) == ["db-agent", "rag-ide", "rag-sec-tutorial"] + assert sorted(cat) == ["aml", "db-agent", "rag-ide", "rag-sec-tutorial"] assert cat["rag-ide"]["platforms"] == ["local", "ec2", "modal"] and cat["rag-ide"]["cost_tier"] == "metered" assert {"name": "hours", "default": 4, "help": "Hard cap per session, 1 to 24", "required": False} in cat["rag-ide"]["params"] assert all(p["title"] and p["summary"] and p["platforms"] and p["cost_tier"] for p in cat.values())