AML pack 3/5: rag mode

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"]]

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

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

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."""

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)

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