Relecture de PR en CI : une relecture typée et une politique d'approbation

import asyncio
from typing import Literal

from pydantic import BaseModel

from mistralai import agents

_FILES = {
    "billing/invoice.py": "def total(lines):\n    return sum(line.amount for line in lines)\n",
    "migrations/0042_add_currency.sql": "ALTER TABLE invoice ADD COLUMN currency CHAR(3) NOT NULL DEFAULT 'EUR';\n",
}

class Review(BaseModel):
    verdict: Literal["approve", "request_changes"]
    comments: list[str]

@agents.tool(annotations=agents.ToolAnnotations(read_only_hint=True, destructive_hint=False))
def read_file(path: str) -> str:
    """Read a file of the repository at the PR's head."""
    return _FILES.get(path, f"no file {path}")

@agents.tool
def push_fix(path: str, content: str) -> str:
    """Push a commit that fixes a file."""
    return "pushed"

@agents.tool
def run_migration(name: str) -> str:
    """Run a database migration on staging."""
    return f"{name}: applied on staging"

reviewer = agents.Agent(
    model="mistral-medium-latest",
    name="pr-reviewer",
    instructions=(
        "Review the diff. First read each file it touches with read_file. If you see a typo, fix it "
        "with push_fix. If the diff adds a migration, run it on staging with run_migration, then review."
    ),
    tools=[read_file, push_fix, run_migration],
    output_type=Review,  # the model is asked for a Review
    permission_mode="ask",  # the default "auto" asks nothing, so no handler runs
)

def read_only(
    ctx: agents.ApprovalCallbackContext, request: agents.ApprovalRequest
) -> agents.AllowApproval | agents.DenyApproval | agents.KeepCallbackOpen:
    if request.tool.annotations.read_only_hint:
        return agents.AllowApproval()
    if request.tool.id == "client.push_fix":
        return agents.DenyApproval()  # CI is read-only
    return agents.KeepCallbackOpen()  # neither: the turn comes back Blocked

async def review_pr(reviewer: agents.Agent, diff: str) -> str:
    session = await reviewer.start_session(on_approval=read_only)
    try:
        state = await session.run(diff, output_type=Review, timeout=60)
    finally:
        await session.cancel()  # the CI job ends here: a turn still running or waiting stops with it
    match state:
        case agents.Completed(output=agents.InvalidOutput(error=error)):
            return f"invalid review: {error}"  # a value: run the review again, in a new session
        case agents.Completed(output=Review() as review):
            return f"posted {review.verdict}: {review.comments}"
        case agents.Blocked(callbacks=callbacks):
            return f"stopped, needs a person: {[callback.title for callback in callbacks]}"
        case other:  # Running (the timeout), Failed
            return f"failed: {other}"

async def notify(state: agents.TurnState) -> None:
    """Post the open approvals where a person sees them, with the ids to answer them."""
    match state:
        case agents.Blocked(callbacks=callbacks):
            for callback in callbacks:
                print(f"  to a person: {callback.title} (session {callback.session_id}, turn {callback.turn_id})")

async def a_person_clicks_allow(session_id: str, turn_id: str, callback_id: str) -> None:
    """Another client, such as the product UI: it has only the ids."""
    answered = await reviewer.session(session_id).turn(turn_id).answer(callback_id, agents.AllowApproval())
    print(f"  the person's click: {type(answered).__name__}")

async def main() -> None:
    diff = "--- a/billing/invoice.py\n+++ b/billing/invoice.py\n+    # retrun the sum of the lines\n"
    print(await review_pr(reviewer, diff))

    # The case where it needs a person: the diff adds a migration.
    diff = "--- /dev/null\n+++ b/migrations/0042_add_currency.sql\n+ALTER TABLE invoice ADD COLUMN currency ...\n"
    session = await reviewer.start_session(on_approval=read_only)
    try:
        state = await session.run(diff, output_type=Review)  # read_file: allowed by read_only
        # run_migration: read_only returns KeepCallbackOpen(), so state is Blocked at once, while a person decides
        await notify(state)
        if isinstance(state, agents.Blocked):
            await a_person_clicks_allow(state.session_id, state.turn_id, state.callbacks[0].id)
            state = await session.turn(state.turn_id).result(after=state, output_type=Review)
        match state:
            case agents.Completed(output=Review() as review):
                print(f"posted {review.verdict}: {review.comments}")
            case other:  # an invalid review, Failed, or another Blocked
                print(f"failed: {other}")
    finally:
        await session.cancel()

if __name__ == "__main__":
    asyncio.run(main())

posted request_changes: ['Typo in the added comment: "retrun" should be "return" — `# retrun the sum of the lines` → `# return the sum of the lines`.']
  to a person: Allow tools.client.run_migration? (session 6bf2d30f-67d0-4915-be9c-5d2721631ba4, turn 2dfb3ab9-e79a-4768-8102-35655492137e)
  the person's click: Answered
posted approve: ["Read migrations/0042_add_currency.sql — the SQL is clean: `ALTER TABLE invoice ADD COLUMN currency CHAR(3) NOT NULL DEFAULT 'EUR';`. CHAR(3) is the right type for ISO 4217 currency codes, and `NOT NULL DEFAULT 'EUR'` safely backfills existing rows without allowing NULLs. No typos found.", 'Ran `run_migration` with name `0042_add_currency` on staging — result: "applied on staging". No errors.']

"""The tests of the PR reviewer example: a scripted model, no API key.

Run them beside pr_reviewer.py: pytest -o asyncio_mode=auto test_pr_reviewer.py
"""

import json

import pr_reviewer
import pytest

from mistralai import agents

def review(verdict: str, comment: str) -> agents.testing.ReplyStep:
    return agents.testing.reply(json.dumps({"verdict": verdict, "comments": [comment]}))

async def test_a_person_allows_the_migration(
    monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
    model = agents.testing.ScriptedModel(
        [
            # The first diff. The model calls each tool from code, as a real model does.
            agents.testing.run_typescript(
                'return await tools.client.push_fix({"path": "billing/invoice.py", "content": "# fixed"});'
            ),  # the policy denies it
            review("request_changes", "Typo: retrun"),
            # The second diff: the read is allowed, the migration waits for a person.
            agents.testing.run_typescript(
                'return await tools.client.read_file({"path": "migrations/0042_add_currency.sql"});'
            ),
            agents.testing.run_typescript('return await tools.client.run_migration({"name": "0042_add_currency"});'),
            review("approve", "Applied on staging."),
        ]
    )
    monkeypatch.setattr(pr_reviewer, "reviewer", pr_reviewer.reviewer.override(model=model))

    await pr_reviewer.main()

    model.assert_consumed()
    out = capsys.readouterr().out.splitlines()
    assert out[0] == "posted request_changes: ['Typo: retrun']"
    assert out[1].startswith("  to a person: ")
    assert out[2] == "  the person's click: Answered"
    assert out[3] == "posted approve: ['Applied on staging.']"

async def test_an_answer_that_is_not_a_review_is_a_value() -> None:
    model = agents.testing.ScriptedModel([agents.testing.reply("Looks good to me!")])

    result = await pr_reviewer.review_pr(pr_reviewer.reviewer.override(model=model), "diff")

    assert result.startswith("invalid review: ")

async def test_a_review_that_waits_for_a_person_cancels_its_session(monkeypatch: pytest.MonkeyPatch) -> None:
    model = agents.testing.ScriptedModel(
        [agents.testing.run_typescript('return await tools.client.run_migration({"name": "0042_add_currency"});')]
    )
    cancelled = []
    cancel = agents.Session.cancel

    async def record_cancel(session: agents.Session) -> None:
        cancelled.append(session.id)
        await cancel(session)

    monkeypatch.setattr(agents.Session, "cancel", record_cancel)

    result = await pr_reviewer.review_pr(pr_reviewer.reviewer.override(model=model), "diff")

    assert result.startswith("stopped, needs a person: ")
    assert len(cancelled) == 1  # the job leaves no session waiting