Workflows with Prefect¶

"""The flow that summarizes documents: thin orchestration over the domain code."""

from __future__ import annotations

import asyncio

from prefect import flow, get_run_logger, task
from prefect.artifacts import acreate_markdown_artifact

from app.core.documents import Document, Summary, summarize

# Lowercase letters, numbers and dashes only: Prefect rejects any other artifact key.
ARTIFACT_KEY = "document-summaries"


@task(retries=2, retry_delay_seconds=5)
async def summarize_document(document: Document) -> Summary:
    """Summarize one document, retried when it fails."""
    return summarize(document)


def _table(summaries: list[Summary]) -> str:
    rows = "".join(f"| {s.document_id} | {s.words} | {s.first_sentence} |\n" for s in summaries)
    return f"| document | words | first sentence |\n|---|---|---|\n{rows}"


@flow
async def summarize_documents(documents: list[Document]) -> list[Summary]:
    """Summarize the documents at the same time and publish the summaries as an artifact."""
    logger = get_run_logger()
    logger.info("summarizing %d documents", len(documents))
    async with asyncio.TaskGroup() as group:
        runs = [group.create_task(summarize_document(document)) for document in documents]
    summaries = [run.result() for run in runs]
    await acreate_markdown_artifact(
        markdown=_table(summaries),
        key=ARTIFACT_KEY,
        description="The summary of each document of the run",
    )
    return summaries

"""The summary of a document: domain code, which knows nothing of Prefect."""

from __future__ import annotations

from pydantic import BaseModel

from app.core.logger import get_logger

log = get_logger(__name__)


class Document(BaseModel):
    """A document to summarize."""

    id: str
    text: str


class Summary(BaseModel):
    """What the summary of a document says about it."""

    document_id: str
    words: int
    first_sentence: str


def summarize(document: Document) -> Summary:
    """Return the summary of a document."""
    words = len(document.text.split())
    first_sentence = document.text.split(".", 1)[0].strip()
    log.info("document_summarized", document_id=document.id, words=words)
    return Summary(document_id=document.id, words=words, first_sentence=first_sentence)

"""The process that serves the deployments of the package and runs their flow runs.

It is not a worker: a worker polls a work pool and starts infrastructure for each run, and this
process runs each flow run itself, as a subprocess.
"""

from __future__ import annotations

import asyncio
from typing import TYPE_CHECKING

from prefect import aserve

from app.flows.documents import summarize_documents

if TYPE_CHECKING:
    from prefect.deployments.runner import RunnerDeployment


async def deployments() -> list[RunnerDeployment]:
    """Return the deployments that this process serves."""
    return [await summarize_documents.ato_deployment(name="summarize-documents")]


async def main() -> None:
    """Register the deployments and run their flow runs until the process stops."""
    await aserve(*await deployments())


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

"""The flow runs of a deployment, started and awaited from another process, such as the API.

The module names the deployment and imports no flow, so the API does not load the code of the
flows nor their dependencies.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

from prefect.client.orchestration import get_client
from prefect.deployments import arun_deployment
from prefect.exceptions import FlowRunWaitTimeout
from prefect.flow_runs import wait_for_flow_run
from prefect.states import Cancelling

if TYPE_CHECKING:
    from uuid import UUID

    from prefect.client.schemas.objects import StateType

    from app.core.documents import Document

# `<flow name>/<deployment name>`, as `app/deploy.py` serves it.
SUMMARIZE_DOCUMENTS = "summarize-documents/summarize-documents"


async def start_summaries(documents: list[Document]) -> UUID:
    """Create a flow run of the deployment and return its id without waiting for it."""
    flow_run = await arun_deployment(
        name=SUMMARIZE_DOCUMENTS,
        parameters={"documents": [document.model_dump() for document in documents]},
        timeout=0,
    )
    return flow_run.id


async def wait_or_cancel(flow_run_id: UUID, timeout: int) -> StateType | None:
    """Wait for a flow run up to `timeout` seconds; past it, cancel the run."""
    try:
        flow_run = await wait_for_flow_run(flow_run_id, timeout=timeout)
    except FlowRunWaitTimeout:
        async with get_client() as client:
            await client.set_flow_run_state(flow_run_id, state=Cancelling(), force=True)
            flow_run = await client.read_flow_run(flow_run_id)
    return flow_run.state_type

"""The endpoints of documents: the API starts flow runs and imports no flow."""

from fastapi import APIRouter, status

from app.api.v1.schemas.requests.documents import SummarizeRequest
from app.api.v1.schemas.responses.documents import SummarizeResponse
from app.core.documents import Document
from app.runs import start_summaries

router = APIRouter(prefix="/documents", tags=["documents"])


@router.post("/summaries", status_code=status.HTTP_202_ACCEPTED)
async def summarize(payload: SummarizeRequest) -> SummarizeResponse:
    """Start a flow run that summarizes the documents, and return its id."""
    documents = [Document(id=item.id, text=item.text) for item in payload.documents]
    return SummarizeResponse(flow_run_id=await start_summaries(documents))

"""The automation that calls a webhook when a flow run of the summaries fails or crashes.

The server runs it, so it fires even when the process of the flow run dies, which a state hook of
the flow cannot promise.
"""

from __future__ import annotations

from prefect.automations import Automation
from prefect.blocks.webhook import Webhook
from prefect.events.actions import CallWebhook
from prefect.events.schemas.automations import EventTrigger, Posture
from pydantic import SecretStr

PAYLOAD = '{"flow_run": "{{ flow_run.name }}", "state": "{{ flow_run.state.name }}"}'


async def notify_failed_summaries(webhook_url: str) -> Automation:
    """Create the automation, and the webhook block it calls, on the server of Prefect."""
    webhook = Webhook(method="POST", url=SecretStr(webhook_url))
    block_id = await webhook.asave("summaries-failed", overwrite=True)
    automation = Automation(
        name="notify-failed-summaries",
        trigger=EventTrigger(
            expect={"prefect.flow-run.Failed", "prefect.flow-run.Crashed"},
            match={"prefect.resource.id": "prefect.flow-run.*"},
            match_related={
                "prefect.resource.role": "flow",
                "prefect.resource.name": "summarize-documents",
            },
            posture=Posture.Reactive,
            threshold=1,
        ),
        actions=[CallWebhook(block_document_id=block_id, payload=PAYLOAD)],
    )
    return await automation.acreate()