Route pipeline failures to a fix or a person¶

A task failed overnight and the on-call engineer has to read the error, decide whether it was a blip worth rerunning, something a person has to fix now, or noise to ignore, and then act. This Dag hands the reading to a model that returns a pick and how sure it is. Airflow runs only the chosen branch, sends uncertain picks to a human with a deadline, and demands more confidence to page someone than to rerun a task.

What this demonstrates¶

Branch on an answer: LLMBranchOperator – LLMBranchOperator picks one downstream task id.
Classifier models – a classifier model, so the gate has a confidence to read.
Approval gates for LLM operators – on_uncertain="review" routes low-confidence picks to a human, with approval_timeout bounding the wait.

Run it¶

Install the provider with the classifier extra:

pip install "apache-airflow-providers-common-ai[typesafe]"

Set {"model": "typesafe:jev-1.13.0"} in the pydanticai_default extra. The second Dag below reads the same kind of connection as jev_default.

Trigger the Dag:

airflow dags test example_llm_branch_decision_policy

The triage_failure log shows the pick and its confidence. Exactly one of rerun, page_oncall and ignore runs. Below the confidence bar the run pauses for review instead (Airflow 3.1 or later); answer from Required Actions in the UI (see the note on What you can build).

The Dag¶

airflow/providers/common/ai/example_dags/example_llm_branch.py[source]

@dag(tags=["example"])
def example_llm_branch_decision_policy():
    # A classifier model reports how sure it is of each pick; a text model does not, and
    # with a min_confidence set every pick would count as uncertain and go to review.
    route = LLMBranchOperator(
        task_id="triage_failure",
        prompt=(
            "Task load_orders failed: psycopg2.OperationalError: could not connect to server: "
            "Connection timed out. Is the server running on host db.internal (10.0.4.12)?"
        ),
        llm_conn_id="pydanticai_default",
        model_id="typesafe:jev-1.13.0",
        system_prompt="Pick the remediation that addresses the cause of the failure.",
        branches={
            "rerun": "The failure looks transient: a timeout, a dropped connection, a rate limit.",
            # Paging someone on a wrong pick costs more than an extra rerun, so this branch needs more.
            "page_oncall": BranchOption(
                "Something a person has to fix now: data corruption, an outage, a security issue.",
                min_confidence=0.9,
            ),
            "ignore": "Expected or harmless: a known flaky check, a duplicate alert.",
        },
        decision_policy=DecisionPolicy(min_confidence=0.6, on_uncertain="review"),
        approval_timeout=timedelta(hours=4),
        allow_modifications=True,
    )

    @task
    def rerun():
        return "Clearing the failed task"

    @task
    def page_oncall():
        return "Paging on-call"

    @task
    def ignore():
        return "Leaving it"

    route >> [rerun(), page_oncall(), ignore()]

Classify-then-act variant¶

When the action depends on the score itself, classify in one task and act in the next (Classifier models explains reading the score). It uses the jev_default connection; run it as airflow dags test example_classifier_model_confidence:

airflow/providers/common/ai/example_dags/example_classifier_model.py[source]

@dag(tags=["example", "classifier"])
def example_classifier_model_confidence():
    """Classify a failure and escalate when the model says it does not know.

    The branch Dag above cannot do this: ``LLMBranchOperator`` takes the branch inside the
    operator, before any task can read the confidence.
    """

    @task
    def classify(log_line: str) -> dict:
        agent = PydanticAIHook(llm_conn_id="jev_default").create_agent(
            output_type=Literal["transient", "resource", "permanent"],
            instructions="Classify why this Airflow task failed.",
        )
        result = agent.run_sync(log_line)
        # Confidence is reported per output field; a bare output type lands under
        # "response". A bounded ``float`` output would report none at all -- there the
        # probability is the answer -- so this ``or 0.0`` would read as no confidence
        # rather than as a missing one. No operator surfaces this on XCom by default.
        details = result.response.provider_details or {}
        return {
            "category": result.output,
            "confidence": (details.get("confidence") or {}).get("response"),
        }

    @task
    def act(classification: dict) -> str:
        confidence = classification["confidence"] or 0.0
        if confidence >= ACT_ABOVE:
            return f"Remediating {classification['category']} automatically"
        if confidence >= REVIEW_ABOVE:
            return f"Filing {classification['category']} for review"
        return "Paging a human: the classification was a coin flip"

    act(classify(UNDERDETERMINED_INCIDENT))

Adapting it¶

Replace the hard-coded prompt with the failed task’s own error. An on_failure_callback on the production Dag can trigger this one with the exception text in conf, and the prompt reads {{ dag_run.conf["error"] }}.
Make the branch tasks do the work: rerun clears the failed task instance through the REST API, page_oncall posts to your alerting provider.
Tune min_confidence per branch. A wrong page costs more than an extra rerun, so page_oncall demands more certainty.
To decide retry versus fail inside Airflow’s retry loop, see Retry policies.