PIG-Demo/envs/alert_triage/alert_triage/policies.py

"""The scripted analysts: three shipped references and the probe ladder's cheats.

Every policy is a function `(view) -> reply text`, where `view` is exactly what
the engine has rendered to the model (`Engine.view()`), and nothing else. No
policy sees the world, the label, the typology, the planted set or the
reference. `tests/test_policies.py` enforces the signature and scans the view
for hidden keys. They return TEXT so the parser and the rejection path are
exercised exactly as a model would exercise them.

The decision rules below are heuristics, written down so a reviewer can see
that the reference "analysts" are players, not oracles:

  fast      decides from the screen; pulls `documents` only where the screen
            says a document should explain the alert (a processor, a title
            company, an estate). One or two turns. Cannot see a hidden signal.
  targeted  `fast` on a screen that shows a typology; otherwise the three
            lookups that could carry a hidden signal — documents, prior
            alerts, each counterparty on screen — then decides. No history.
  thorough  always the full procedure: documents, prior alerts, each
            counterparty on screen, itemised history for the two most recent
            months, then decides; cites every constituting transaction it saw.

Hidden-signal rules, shared by targeted and thorough:
  * a counterparty shared with ≥ 3 other alerted accounts  → FUNNEL_ACCOUNT,
    cite the counterparty and its transactions;
  * ≥ 3 prior alerts, all closed within a day on one identical rationale
                                                           → the rule's typology,
    cite the prior alerts and the triggering transactions;
  * a document naming a party whose counterparty lookup carries a different
    name                                                   → the rule's typology,
    cite the document, the counterparty and its triggering credits.
"""

from __future__ import annotations

import json
from typing import Any, Callable

from .generator import INBOUND, NOMINAL_TYPOLOGY, RULE_IDS

Policy = Callable[[dict[str, Any]], str]

_FAMILY_OF = {rid: fam for table in RULE_IDS.values() for fam, rid in table.items()}
EXPLAINING_DOCS = ("processing_agreement", "closing_statement", "executor_letter")


def _emit(obj: dict[str, Any]) -> str:
    return json.dumps(obj)


def _family(view: dict[str, Any]) -> str:
    return _FAMILY_OF.get(view["screen"]["alert"]["rule"], "STR")


def _looked(view: dict[str, Any], what: str, key: str | None = None) -> dict[str, Any] | None:
    for l in view["lookups"]:
        if l["what"] == what and (key is None or l["key"] == key):
            return l["result"]
    return None


def _band(view: dict[str, Any], t: dict[str, Any]) -> bool:
    threshold = view["screen"]["threshold"]
    return t["channel"] == "cash_in" and threshold * 8 // 10 <= t["amount"] < threshold


def _close(reason: str, cites: list[str], note: str) -> dict[str, Any]:
    return {"action": "close", "reason": reason, "cites": cites, "note": note}


def _escalate(typology: str, cites: list[str], narrative: str) -> dict[str, Any]:
    return {"action": "escalate", "typology": typology, "cites": cites, "narrative": narrative}


def _dedupe(ids: list[str]) -> list[str]:
    out: list[str] = []
    for i in ids:
        if i not in out:
            out.append(i)
    return out


# ------------------------------------------------------------ decisions --


def screen_decision(view: dict[str, Any]) -> dict[str, Any] | None:
    """What the screen alone supports. None means "a document should explain this"."""
    screen = view["screen"]
    fam = _family(view)
    trig = screen["triggering_transactions"]
    trig_ids = [t["id"] for t in trig]
    kyc = screen["customer"]
    max_cash = max(r["cash_in"] for r in screen["summary"])

    if fam == "STR":
        band = [t for t in trig if _band(view, t)]
        branches = {t["branch"] for t in band}
        if len(band) >= 5 or (len(band) >= 4 and len(branches) >= 3):
            return _escalate("STRUCTURING", [t["id"] for t in band], "Repeated cash deposits just under the threshold across branches.")
        if kyc["expected_monthly_cash"] >= max_cash:
            return _close("CONSISTENT_WITH_PROFILE", ["kyc.expected_monthly_cash"], "Cash volume within the expectation on file; band deposits are ordinary takings.")
        return _escalate("STRUCTURING", [t["id"] for t in band] or trig_ids, "Band deposits with no expectation on file to cover them.")

    if fam == "CASH":
        if kyc["expected_monthly_cash"] >= max_cash:
            return _close("CONSISTENT_WITH_PROFILE", ["kyc.expected_activity", "kyc.expected_monthly_cash"], "Seasonal peak the file anticipated; within expected monthly cash.")
        return _escalate("STRUCTURING", trig_ids, "Cash far above the expectation on file and rising.")

    if fam == "VEL":
        senders = {t["cp"] for t in trig if t["channel"] == "p2p_in"}
        if len(senders) >= 3:
            return _escalate("MONEY_MULE", trig_ids, "Credits from unrelated individuals moved to an exchange within a day.")
        return None

    if fam == "WIRE":
        if kyc.get("remittance_declared", "none declared") != "none declared":
            return _close("CONSISTENT_WITH_PROFILE", ["kyc.remittance_declared"], "Destination declared at onboarding; amounts within income.")
        return _escalate("UNKNOWN", trig_ids, "Wires to a monitored jurisdiction with no declared purpose.")

    return None   # DORM: needs the document


def document_decision(view: dict[str, Any]) -> dict[str, Any]:
    """After `documents`: close on the explaining document, or escalate if there is none."""
    docs = _looked(view, "documents")
    trig_ids = [t["id"] for t in view["screen"]["triggering_transactions"]]
    if docs:
        for d in docs["documents"]:
            if d["kind"] in EXPLAINING_DOCS:
                return _close("DOCUMENTED_SOURCE_OF_FUNDS", [d["id"]], f"Source of funds documented: {d['kind']}.")
    return _escalate(NOMINAL_TYPOLOGY[_family(view)], trig_ids, "No document on file explains the movement.")


def lookup_decision(view: dict[str, Any]) -> dict[str, Any] | None:
    """A hidden signal in the lookups, or None."""
    screen = view["screen"]
    fam = _family(view)
    trig = screen["triggering_transactions"]
    trig_ids = [t["id"] for t in trig]

    for l in view["lookups"]:
        if l["what"] == "counterparty" and l["result"]["other_alerted_accounts"] >= 3:
            r = l["result"]
            return _escalate("FUNNEL_ACCOUNT", [r["id"]] + list(r["transactions"]),
                             f"{r['id']} is shared with {r['other_alerted_accounts']} other alerted accounts; funds converge on it.")

    priors = _looked(view, "prior_alerts")
    if priors and len(priors["alerts"]) >= 3:
        alerts = priors["alerts"]
        rationales = {p["rationale"] for p in alerts}
        if all(p["disposition"] == "closed" and p["elapsed_days"] <= 1 for p in alerts) and len(rationales) == 1:
            return _escalate(NOMINAL_TYPOLOGY[fam], [p["id"] for p in alerts] + trig_ids,
                             "Repeat alerts closed within a day on one boilerplate rationale; nobody has looked.")

    docs = _looked(view, "documents")
    if docs:
        for d in docs["documents"]:
            if not d["party"]:
                continue
            for cp_id in _dedupe([t["cp"] for t in trig if t["channel"] in INBOUND and t["cp"]]):
                cp = _looked(view, "counterparty", cp_id)
                if cp and cp["name"] != d["party"]:
                    credits = [t["id"] for t in trig if t["cp"] == cp_id and t["channel"] in INBOUND]
                    return _escalate(NOMINAL_TYPOLOGY[fam], [d["id"], cp_id] + credits,
                                     f"The document names {d['party']}; the funds came from {cp['name']}.")
    return None


def enriched_decision(view: dict[str, Any]) -> dict[str, Any]:
    """thorough's final call: the screen rules, with citations widened to everything itemised."""
    d = screen_decision(view)
    fam = _family(view)
    screen = view["screen"]
    months = screen["months"]
    seen: list[dict[str, Any]] = list(screen["triggering_transactions"])
    for l in view["lookups"]:
        if l["what"] == "history":
            seen += l["result"]["transactions"]
    if d is not None and d["action"] == "escalate":
        if fam == "STR":
            d["cites"] = _dedupe([t["id"] for t in seen if _band(view, t) and t["month"] == months[-1]] + d["cites"])
        elif fam == "CASH":
            last_two = set(months[-2:])
            d["cites"] = _dedupe([t["id"] for t in seen if t["channel"] == "cash_in" and t["month"] in last_two] + d["cites"])
        return d
    if d is None:
        return document_decision(view)
    return d


def _plan_next(view: dict[str, Any], plan: list[tuple[str, str | None]]) -> str | None:
    for what, key in plan:
        if _looked(view, what, key) is None:
            if what == "history":
                return _emit({"action": "lookup", "what": "history", "month": key})
            if what == "counterparty":
                return _emit({"action": "lookup", "what": "counterparty", "id": key})
            return _emit({"action": "lookup", "what": what})
    return None


# ------------------------------------------------------------- shipped --


def fast(view: dict[str, Any]) -> str:
    d = screen_decision(view)
    if d is not None:
        return _emit(d)
    if _looked(view, "documents") is None:
        return _emit({"action": "lookup", "what": "documents"})
    return _emit(document_decision(view))


def _investigation_plan(view: dict[str, Any]) -> list[tuple[str, str | None]]:
    plan: list[tuple[str, str | None]] = [("documents", None), ("prior_alerts", None)]
    plan += [("counterparty", cp["id"]) for cp in view["screen"]["screen_counterparties"]]
    return plan


def targeted(view: dict[str, Any]) -> str:
    d = screen_decision(view)
    if d is not None and d["action"] == "escalate":
        return _emit(d)
    nxt = _plan_next(view, _investigation_plan(view))
    if nxt:
        return nxt
    found = lookup_decision(view)
    if found:
        return _emit(found)
    return _emit(d if d is not None else document_decision(view))


def thorough(view: dict[str, Any]) -> str:
    months = view["screen"]["months"]
    plan = _investigation_plan(view) + [("history", months[-1]), ("history", months[-2])]
    nxt = _plan_next(view, plan)
    if nxt:
        return nxt
    found = lookup_decision(view)
    return _emit(found if found else enriched_decision(view))


def thorough_wasteful(view: dict[str, Any]) -> str:
    """thorough plus a month that carries nothing. Proves the hour meter binds."""
    months = view["screen"]["months"]
    plan = _investigation_plan(view) + [("history", months[-1]), ("history", months[-2]), ("history", months[-3])]
    nxt = _plan_next(view, plan)
    if nxt:
        return nxt
    found = lookup_decision(view)
    return _emit(found if found else enriched_decision(view))


# -------------------------------------------------------------- cheats --


def inaction(view: dict[str, Any]) -> str:
    return "Let me reason about this alert carefully before committing to a disposition."


def malformed(view: dict[str, Any]) -> str:
    return _emit({"action": "file_sar", "cites": []})


def close_blind(view: dict[str, Any]) -> str:
    return _emit(_close("CONSISTENT_WITH_PROFILE", ["kyc.customer_type"], "Looks fine."))


def close_cites_all(view: dict[str, Any]) -> str:
    screen = view["screen"]
    cites = [f"kyc.{k}" for k in screen["customer"] if k not in ("id", "name")]
    cites += [t["id"] for t in screen["triggering_transactions"]]
    return _emit(_close("CONSISTENT_WITH_PROFILE", cites, "See everything."))


def escalate_blind(view: dict[str, Any]) -> str:
    trig_ids = [t["id"] for t in view["screen"]["triggering_transactions"]]
    return _emit(_escalate(NOMINAL_TYPOLOGY[_family(view)], trig_ids, "Filing defensively."))


def fast_defensive(view: dict[str, Any]) -> str:
    """fast where the screen shows a typology; a blind filing everywhere else."""
    d = screen_decision(view)
    if d is not None and d["action"] == "escalate":
        return _emit(d)
    return escalate_blind(view)


def rule_table(view: dict[str, Any]) -> str:
    """Closes every wire-corridor and dormancy alert blind; files on everything else."""
    fam = _family(view)
    if fam == "WIRE":
        return _emit(_close("CONSISTENT_WITH_PROFILE", ["kyc.remittance_declared"], "Wires are always remittances."))
    if fam == "DORM":
        return _emit(_close("CONSISTENT_WITH_PROFILE", ["kyc.account_opened"], "Old accounts wake up."))
    return escalate_blind(view)


def cite_unseen(view: dict[str, Any]) -> str:
    """Cites a document it never opened; rejected, then falls back to closing blind."""
    if view["rejected"] == 0:
        return _emit(_close("DOCUMENTED_SOURCE_OF_FUNDS", ["doc.probate_letter"], "Per the letter on file."))
    return close_blind(view)


POLICIES: dict[str, Policy] = {
    "inaction": inaction,
    "malformed": malformed,
    "close_blind": close_blind,
    "close_cites_all": close_cites_all,
    "cite_unseen": cite_unseen,
    "escalate_blind": escalate_blind,
    "fast_defensive": fast_defensive,
    "rule_table": rule_table,
    "fast": fast,
    "targeted": targeted,
    "thorough": thorough,
    "thorough_wasteful": thorough_wasteful,
}

SHIPPED = ("fast", "targeted", "thorough")


def run_policy(policy: Policy, world: dict[str, Any]):
    """Play a world to the end under one policy. Returns the finished Engine."""
    from .engine import Engine

    engine = Engine(world)
    while not engine.done:
        engine.step(policy(engine.view()))
    return engine


def reference_for(world: dict[str, Any]) -> tuple[int | None, str | None]:
    """The cheapest shipped policy that reaches the correct disposition, in minutes.

    (None, None) if none does — which `tests/test_reference.py` asserts never
    happens for seeds 0–4095, so a null here is a CI failure, not a silent
    "not scored" on the page.
    """
    best: tuple[int, str] | None = None
    for name in SHIPPED:
        engine = run_policy(POLICIES[name], world)
        if engine.outcome == "solved" and (best is None or engine.minutes < best[0]):
            best = (engine.minutes, name)
    return best if best else (None, None)