Show live progress from a hook

[project]
name = "live-progress"
version = "0.1.0"
description = "An invoice checking agent whose hook pushes each step to a progress bar, in your process or when deployed."
requires-python = "==3.12.*"
dependencies = [
    "mistralai-agents==1.1.0",
    "httpx>=0.28",
    "starlette>=0.46",
    "uvicorn>=0.34",
]

[tool.mistralai.agents]
entrypoint = "invoice_checker:agent"

[tool.uv]
package = false

"""Your app's side: the progress board the hook pushes to, and the endpoint a deployed agent posts to."""

import hmac
from collections.abc import Callable
from typing import Literal

from pydantic import BaseModel, ValidationError
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import Response
from starlette.routing import Route


class StepUpdate(BaseModel):
    """One finished step of a session, such as one invoice checked."""

    session_id: str
    step: str  # what the step worked on: the same step sent twice counts once
    status: Literal["done", "failed"]


def print_bar(session_id: str, finished: int, total: int, update: StepUpdate) -> None:
    width = 20
    filled = width * finished // total
    print(f"progress: [{'#' * filled}{'.' * (width - filled)}] {finished}/{total} {update.step} {update.status}")


class ProgressBoard:
    """The progress of each session your app follows, keyed by session id."""

    def __init__(self, on_change: Callable[[str, int, int, StepUpdate], None] = print_bar) -> None:
        self._on_change = on_change
        self._totals: dict[str, int] = {}
        self._steps: dict[str, dict[str, StepUpdate]] = {}

    def expect(self, session_id: str, total: int) -> None:
        """Follow a session that has `total` steps. Updates of any other session are ignored."""
        self._totals[session_id] = total
        self._steps[session_id] = {}

    async def report(self, update: StepUpdate) -> None:
        steps = self._steps.get(update.session_id)
        if steps is None or steps.get(update.step) == update:
            return  # a session the board does not follow, or a step it already has
        steps[update.step] = update
        self._on_change(update.session_id, len(steps), self._totals[update.session_id], update)


def receiver(board: ProgressBoard, token: str) -> Starlette:
    """The endpoint a deployed agent's hook posts each step to."""

    async def receive(request: Request) -> Response:
        sent = request.headers.get("authorization", "")
        if not hmac.compare_digest(sent.encode(), f"Bearer {token}".encode()):
            return Response(status_code=401)
        try:
            update = StepUpdate.model_validate_json(await request.body())
        except ValidationError:
            return Response(status_code=422)
        await board.report(update)
        return Response(status_code=204)

    return Starlette(routes=[Route("/progress", receive, methods=["POST"])])

"""An invoice checking agent whose hook pushes each step to a progress board.

The hook holds a reporter in a PrivateAttr. In your process, the reporter is the board itself. When the
agent is deployed (`agents serve` or `agents deploy`), it is an HttpReporter that posts each step to
your app. The accounting system is faked, so the example runs with only a Mistral API key.

`agents serve` and `agents deploy` load `agent`, at the bottom of this file.
"""

import asyncio
import logging
import os
from typing import Protocol

import httpx
from progress_board import StepUpdate
from pydantic import PrivateAttr

from mistralai import agents

log = logging.getLogger(__name__)

# --- The accounting system, faked ---------------------------------------------------------------
# In production, check_invoice reads the invoice and its purchase order. Here the amounts are fixed.

_INVOICES = {
    "INV-1001": (1200.00, 1200.00),
    "INV-1002": (845.50, 845.50),
    "INV-1003": (3100.00, 2900.00),
    "INV-1004": (99.90, 99.90),
    "INV-1005": (12500.00, 12500.00),
    "INV-1006": (430.00, 340.00),
}
BATCH = list(_INVOICES)
DELAY_SECONDS = 0.5  # the accounting system is slow


@agents.tool
async def check_invoice(invoice_id: str) -> dict[str, object]:
    """Compare an invoice with its purchase order. Returns both amounts and whether they match."""
    await asyncio.sleep(DELAY_SECONDS)
    invoiced, ordered = _INVOICES[invoice_id]
    return {"invoice_id": invoice_id, "invoiced": invoiced, "ordered": ordered, "matches": invoiced == ordered}


# --- Where the hook pushes ----------------------------------------------------------------------


class Reporter(Protocol):
    async def report(self, update: StepUpdate) -> None: ...


class HttpReporter:
    """When deployed: post each step to your app, at PROGRESS_URL, with the PROGRESS_TOKEN bearer token."""

    async def report(self, update: StepUpdate) -> None:
        async with httpx.AsyncClient(timeout=2) as client:
            response = await client.post(
                os.environ["PROGRESS_URL"],
                content=update.model_dump_json(),
                headers={"Authorization": f"Bearer {os.environ['PROGRESS_TOKEN']}"},
            )
            response.raise_for_status()


# --- The hook -----------------------------------------------------------------------------------


class ReportProgress(agents.Hook):
    """Push a step to the reporter each time one of the `steps` tools finishes."""

    steps: list[str]  # the tool ids that count as a step
    _reporter: Reporter = PrivateAttr()

    def __init__(self, reporter: Reporter, **settings: object) -> None:
        super().__init__(**settings)
        self._reporter = reporter

    async def post_tool_call(self, ctx, result):
        if result.tool_id not in self.steps:
            return
        update = StepUpdate(
            session_id=ctx.session.id,
            step=" ".join(str(value) for value in result.arguments.values()),  # here, the invoice id
            status="failed" if result.tool_result.is_error else "done",
        )
        try:
            await asyncio.wait_for(self._reporter.report(update), timeout=3)
        except Exception as error:  # a progress bar never stops the agent
            log.warning("progress update lost: %r", error)


# --- The agent ----------------------------------------------------------------------------------


def build_agent(reporter: Reporter, **options) -> agents.Agent:
    return agents.Agent(
        model="mistral-large-4",
        name="invoice-checker",
        instructions=(
            "You check invoices with check_invoice, one invoice at a time, each call after the previous "
            "one returned. Then list the invoices whose amount does not match its purchase order, in one sentence."
        ),
        tools=[check_invoice],
        hooks=[ReportProgress(reporter, steps=["client.check_invoice"])],
        **options,
    )


# What `agents serve` and `agents deploy` load: the hook posts each step to your app.
agent = build_agent(
    HttpReporter(),
    secrets={
        "PROGRESS_URL": agents.MistralSecret(worker_env=True),
        "PROGRESS_TOKEN": agents.MistralSecret(worker_env=True),
    },
)

PROMPT = f"Check these invoices: {', '.join(BATCH)}."

"""Run the invoice checker in your process: the hook pushes each step straight to the board.

Set MISTRAL_API_KEY, then run in this folder:
    uv run python run_local.py
"""

import asyncio

from invoice_checker import BATCH, PROMPT, build_agent
from progress_board import ProgressBoard

from mistralai import agents


async def main() -> None:
    board = ProgressBoard()
    checker = build_agent(board)  # the hook holds the board itself
    async with agents.local_runtime():
        session = await checker.start_session()
        board.expect(session.id, total=len(BATCH))
        state = await session.run(PROMPT)
        print(f"agent: {state.text}")
        await session.cancel()


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

"""Drive the served or deployed invoice checker, and show the steps its hook posts to this app.

Serve or deploy the agent first (see the cookbook), then run in this folder, with the same
PROGRESS_TOKEN as the deployed agent:
    uv run python run_remote.py
"""

import asyncio
import os

import uvicorn
from invoice_checker import BATCH, PROMPT
from progress_board import ProgressBoard, receiver

from mistralai import agents


async def main() -> None:
    board = ProgressBoard()
    app = receiver(board, token=os.environ["PROGRESS_TOKEN"])
    server = uvicorn.Server(uvicorn.Config(app, host="127.0.0.1", port=8765, log_level="warning"))
    serving = asyncio.create_task(server.serve())
    try:
        checker = agents.Agent.load_from(agent_name="invoice-checker")
        session = await checker.start_session()
        board.expect(session.id, total=len(BATCH))  # before the first step can arrive
        state = await session.run(PROMPT)
        print(f"agent: {state.text}")
        await session.cancel()
    finally:
        server.should_exit = True
        await serving


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