Review a suspicious payment with an agent and an analyst

READ_ONLY = agents.ToolAnnotations(read_only_hint=True, destructive_hint=False)

@agents.tool(annotations=READ_ONLY)
def get_alert(alert_id: str) -> dict[str, object]:
    """Transaction monitoring: the alert and the payment that raised it."""
    if alert_id not in _ALERTS:
        raise agents.ToolError(f"No alert {alert_id}.")
    return _ALERTS[alert_id]

@agents.tool
def message_customer(account: str, text: str) -> str:
    """Customer messaging: send a message to the account holder."""
    MESSAGES.append(text)
    return "sent"

@agents.tool
async def open_case(alert_id: str, summary: str, evidence: list[str]) -> str:
    """Open a case file on a suspicious payment, with the evidence. Returns the case id."""
    case_id = f"CASE-{alert_id}"
    CASES[case_id] = {"alert_id": alert_id, "summary": summary, "evidence": evidence}
    return case_id


def review_policy(
    ctx: agents.ApprovalCallbackContext, request: agents.ApprovalRequest
) -> agents.AllowApproval | agents.DenyApproval | agents.KeepCallbackOpen:
    if request.tool.annotations.read_only_hint:  # request.tool: every approval has one
        return agents.AllowApproval()
    if request.tool.id == "client.open_case":
        return agents.KeepCallbackOpen()  # not mine: the approval stays open for an analyst
    return agents.DenyApproval()

payment_reviewer = agents.Agent(
    model="mistral-large-4",
    name="payment-reviewer",
    instructions=(
        "Review the alert. Collect the context: the alert, the account history, the KYC profile, "
        "sanctions screening and adverse media for the beneficiary. Do not contact the customer. "
        "If the payment is suspicious, open a case file with the evidence. An analyst approves each "
        "case file: if one is denied, do not open it again; end with your review."
    ),
    tools=[
        get_alert,
        get_account_history,
        get_kyc_profile,
        screen_sanctions,
        search_adverse_media,
        message_customer,
        open_case,
    ],
    permission_mode="ask",  # the default "auto" asks nothing, so no handler runs
)

async def on_alert(alert_id: str) -> agents.TurnState:
    """The transaction monitoring process: one review per alert."""
    session = await payment_reviewer.start_session(on_approval=review_policy, idempotency_key=f"alert:{alert_id}")
    return await session.run(
        f"Review transaction monitoring alert {alert_id}.", idempotency_key=f"alert:{alert_id}:review"
    )
    # Blocked when open_case waits: the case management app lists state.callbacks


async def on_analyst_button(session_id: str, turn_id: str, callback_id: str, approve: bool) -> agents.TurnState:
    """The case management app, another service: an analyst allows or denies the case file."""
    turn = payment_reviewer.session(session_id, on_approval=review_policy).turn(turn_id)
    answer = agents.AllowApproval() if approve else agents.DenyApproval()
    await turn.answer(callback_id, answer)  # Answered | AlreadySettled | NotFound
    return await turn.result()  # drives the rest; review_policy answers the reads

async def main(alert_id: str, approve: bool) -> agents.TurnState:
    async with agents.local_runtime():  # at exit: every local session of the event loop stops
        state = await on_alert(alert_id)
        try:
            while isinstance(state, agents.Blocked):
                # The case management app shows the open approval to an analyst.
                callback = state.callbacks[0]
                print(f"Waiting for an analyst: {callback.title}")
                if isinstance(callback, agents.OpenApproval) and isinstance(callback.request.input, dict):
                    print(f"  summary: {callback.request.input['summary']}")

                # The analyst clicks a button; the app passes on the ids it stored.
                state = await on_analyst_button(state.session_id, state.turn_id, callback.id, approve)
            print(f"Review: {type(state).__name__}: {state.text}")
            print(f"Case files: {list(CASES)}")
            return state
        finally:
            await payment_reviewer.session(state.session_id).cancel()  # the review is over

import asyncio

from payment_review import payment_reviewer, review_policy

from mistralai import agents


def show(event: agents.Event) -> None:
    match event:
        case agents.CallbackOpened(callback=callback):
            print(f"  [approval: {callback.title}]")
        case agents.TurnEnded(state=state):
            print(f"  [turn ended: {type(state).__name__}]")


async def ask_about(alert_id: str) -> None:
    session = await payment_reviewer.start_session(on_approval=review_policy)
    try:
        state = await session.run(
            f"Summarise alert {alert_id}: who paid whom, and why it was flagged. Don't open a case yet."
        )
        print(state.text)
        state = await session.run(
            "Compare it with the customer's payments of the last 12 months.",
            on_event=show,  # each event as it comes
        )
        print(state.text)
    finally:
        await session.cancel()  # returning does not end the session


if __name__ == "__main__":
    asyncio.run(ask_about("4711"))

"""The tests of the suspicious payment review: a scripted model, no API key.

Run them beside payment_review.py: pytest -o asyncio_mode=auto test_payment_review.py
"""

import ask_about
import payment_review
import pytest
import run_review

from mistralai import agents


def use_model(monkeypatch: pytest.MonkeyPatch, model: agents.testing.ScriptedModel) -> None:
    """Give the agent the scripted model, in each module that imported it."""
    agent = payment_review.payment_reviewer.override(model=model)
    for module in (payment_review, run_review, ask_about):
        monkeypatch.setattr(module, "payment_reviewer", agent)


@pytest.mark.parametrize("approve", [True, False])
async def test_the_review_waits_for_the_analyst(
    monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str], approve: bool
) -> None:
    payment_review.CASES.clear()
    model = agents.testing.ScriptedModel(
        [
            # The model calls each tool from code, as a real model does.
            agents.testing.run_typescript('return await tools.client.get_alert({"alert_id": "4711"});'),
            agents.testing.run_typescript(
                'return await tools.client.screen_sanctions({"name": "Northwind Trading FZE"});'
            ),
            agents.testing.run_typescript(
                'return await tools.client.open_case({"alert_id": "4711", '
                '"summary": "9,800 EUR to a beneficiary with a 0.82 sanctions match.", '
                '"evidence": ["sanctions match 0.82"]});'
            ),
            agents.testing.reply("Case CASE-4711 opened." if approve else "The analyst denied the case file."),
        ]
    )
    use_model(monkeypatch, model)

    state = await run_review.main("4711", approve=approve)

    model.assert_consumed()
    assert isinstance(state, agents.Completed)
    assert list(payment_review.CASES) == (["CASE-4711"] if approve else [])  # a denied case is never written
    out = capsys.readouterr().out
    assert "Waiting for an analyst: " in out
    assert "summary: 9,800 EUR to a beneficiary" in out


async def test_the_policy_allows_reads_and_denies_the_customer_message(monkeypatch: pytest.MonkeyPatch) -> None:
    payment_review.MESSAGES.clear()
    model = agents.testing.ScriptedModel(
        [
            agents.testing.run_typescript('return await tools.client.get_kyc_profile({"account": "FR76"});'),
            agents.testing.run_typescript(
                'return await tools.client.message_customer({"account": "FR76", "text": "Please explain."});'
            ),
            agents.testing.reply("Reviewed."),
        ]
    )
    use_model(monkeypatch, model)

    state = await payment_review.on_alert("4711")

    assert isinstance(state, agents.Completed)
    assert [c.name for c in state.tool_calls] == ["client.get_kyc_profile", "client.message_customer"]
    assert payment_review.MESSAGES == []  # the policy denied it, so the tool never ran


async def test_the_analysts_conversation_takes_two_turns(
    monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
    model = agents.testing.ScriptedModel(
        [
            agents.testing.reply("Mr Le Chonk paid 9,800 EUR to Northwind Trading FZE."),
            agents.testing.run_typescript('return await tools.client.get_account_history({"account": "FR76"});'),
            agents.testing.reply("Its largest payment in 12 months was 1,150 EUR."),
        ]
    )
    use_model(monkeypatch, model)

    await ask_about.ask_about("4711")

    model.assert_consumed()
    out = capsys.readouterr().out
    assert "Mr Le Chonk paid 9,800 EUR" in out
    assert "[approval: " in out  # the read asked, and the policy answered
    assert "[turn ended: Completed]" in out

pytest -o asyncio_mode=auto test_payment_review.py