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())