> ## Documentation Index
> Fetch the complete documentation index at: https://astronomer.io/docs/llms.txt
> Use this file to discover all available pages before exploring further.

# Task loops in Apache Airflow®

> Learn how to use task loops in Airflow 3.4 to repeat a task group until a condition is met or a maximum number of iterations has run.

Task loops, which were added in Airflow 3.4, allow you to repeatedly run the tasks in a [task group](https://astronomer.io/docs/airflow/task-groups.md) within a Dag. The loop continues `until` a condition is met or it has run `max_iterations` times. You can use the result of a previous iteration in the next iteration.

In this guide, you'll learn:

* When to use task loops, and how they compare to dynamic task mapping, retries, and dynamic Dag generation.
* How to loop a task group with `.loop()`, `max_iterations`, and `until`.
* How to use information about the current and previous iteration in your tasks.
* How to pass loop results to tasks outside of the loop.
* How trigger rules and reruns work in a loop.
* Which combinations with other Airflow features are supported.
* Two example Dags: one that processes a hierarchy one level per iteration, and one that improves AI-generated text with a human-in-the-loop review.

## Assumed knowledge

To get the most out of this guide, you should have an understanding of:

* Airflow Dags. See [Introduction to Apache Airflow® Dags](https://astronomer.io/docs/airflow/dags.md).
* Airflow task groups. See [Airflow task groups](https://astronomer.io/docs/airflow/task-groups.md).
* Passing data between tasks. See [Pass data between tasks](https://astronomer.io/docs/airflow/airflow-passing-data-between-tasks.md).
* Airflow trigger rules. See [Apache Airflow® trigger rules](https://astronomer.io/docs/airflow/airflow-trigger-rules.md).

## When to use task loops

You can use loops to repeat any number and combination of sequential and parallel tasks in a task group during a Dag run. After all tasks in the task group are executed, a gate task determines whether another iteration of the loop runs. The loop continues until either a condition is fulfilled or the maximum number of iterations have been performed.

> [!NOTE]
>
> Task loops can't deadlock a Dag or run forever. Airflow creates the task instances of the next iteration only after the gate task of the current iteration has finished, so the Dag graph has no circular dependencies. Every loop also has an upper limit on the number of iterations, because `max_iterations` is a required parameter.

Common use cases include:

* Improving an LLM output over several iterations: one task generates a draft, another evaluates it, and the loop repeats until the result passes the evaluation. The evaluation can be structured using [human-in-the-loop tasks](https://astronomer.io/docs/airflow/airflow-human-in-the-loop.md), [AI-as-a-judge](https://astronomer.io/docs/airflow/ai-orchestration-overview#review-ai-output), or [deterministic evaluations](https://astronomer.io/docs/airflow/ai-orchestration-model-evals.md). See [Example: Improve AI output with HITL](#example-improve-ai-output-with-hitl).
* Processing nested data one level at a time, for example a folder tree or parent-child records, where each iteration uses the IDs returned by the previous one. See [Example: Deterministically process a hierarchy](#example-deterministically-process-a-hierarchy).
* Processing a large dataset in chunks, using the offset or cursor returned by the previous iteration, until the entire dataset is processed.
* Training a machine learning model, where each iteration trains the model for one or more epochs and evaluates it, until the validation loss stops improving or a target metric is met.
* Repeating a calculation until the result is within a target threshold or converges.

### Task loops vs. other features

| | Task loops | Dynamic task mapping | Retries |
| - | - | - | - |
| Syntax | `.loop()` | `.expand()` | `retries=n` |
| Use it to | Repeat a group of tasks where each iteration can use the result of the previous one, until a condition is met or the loop has run `max_iterations` times | Run the same task or task group once for each item in a collection, in parallel | Try a failed task again, for example after a temporary API error |
| How it works | After each iteration, Airflow creates the task instances for the next iteration if the loop has run fewer than `max_iterations` times and the optional `until` function returns `False`. | Airflow creates one mapped task instance or mapped task group instance per item in a list that is passed at runtime. The mapped instances run in parallel. | Airflow runs the same task instance again. Retries don't create new task instances and don't count as loop iterations. You can use retry policies to determine whether a task is retried based on the type of failure. |
| Airflow version | 3.4+ | Tasks: 2.3+, task groups: 2.5+ | All versions, retry policies: 3.3+ |
| Learn more | This guide | [Dynamic tasks](https://astronomer.io/docs/airflow/dynamic-tasks.md), [Map over task groups](https://astronomer.io/docs/airflow/dynamic-tasks.md#map-over-task-groups) | [Automatically retry tasks](https://astronomer.io/docs/airflow/rerunning-dags.md#automatically-retry-tasks) |

You can combine some of these features. For example, you can use dynamically mapped tasks inside a looped task group, and you can set retries on tasks inside a loop and on dynamically mapped tasks. Other combinations, such as a looped task group inside a dynamically mapped task group, aren't possible. See [Limitations and compatibility](#limitations-and-compatibility).

> [!NOTE]
>
> Creating Dags or tasks programmatically is sometimes called dynamic Dag generation. This is a different concept from task loops and dynamic task mapping. With dynamic Dag generation, you use Python code, for example a `for` loop in the Dag file, to create tasks or Dags based on information that is available when the Dag file is parsed. The Dag has the same tasks in every Dag run. See [Dynamically generate Dags](https://astronomer.io/docs/airflow/dynamically-generating-dags.md).

## Task loop basics

This section covers how to define a task loop, how tasks access information about the current iteration, and a basic example.

### Basic task loop syntax

To loop a task group, add the `.loop` method to your task group function call.

```python theme={null}
from airflow.sdk import task_group, task

def my_func(*, loop): ...

@task_group
def my_task_group(): ...

my_task_group.loop(max_iterations=10, until=my_func)
```

The `.loop` method has two parameters:

* `max_iterations` (mandatory): requires a positive integer that limits how often the loop can repeat.
* `until` (optional): takes a callable that determines whether the loop should continue. Airflow adds this function as the final gate task of the loop, with the function name as the task ID, or `__loop_gate` if the function has no name. If the value returned by the `until` function isn't a boolean, the task fails.

If no `until` function is provided, the loop runs for the number of iterations defined by `max_iterations` and then ends, with the gate task marked as `success`.

If an `until` function is given, the loop continues as long as the function returns `False`. If the function returns `True`, the loop ends with the gate task marked as `success`. If the loop has run `max_iterations` times and the `until` function still returns `False`, the loop ends with the gate task marked as `failed`.

> [!NOTE]
>
> If you have tasks downstream of your loop that should run even if the loop fails or ends at `max_iterations` without the `until` function returning `True`, adjust the [trigger rules](https://astronomer.io/docs/airflow/airflow-trigger-rules.md) of those downstream tasks.

To pass regular arguments into your task group function, use the `.partial` method. To override task group parameters, you can use `.override`. These values are passed for all iterations of the loop.

```python theme={null}
from airflow.sdk import task_group, task

def my_func(*, loop): ...

@task_group(ui_color="#ba5252")
def my_task_group(my_arg): ...

my_task_group.override(ui_color="#57BA52").partial(my_arg=10).loop(max_iterations=10, until=my_func)
```

Within the looped task group, you can add any number of sequential and parallel tasks (with some [limitations](#limitations-and-compatibility) for specific features). The only requirement is that there is one final task in the dependency structure. The value pushed to XCom by this final task under the `return_value` key is the iteration's result (`loop.result`).

### Access loop information

Inside the task group loop, you can access information about the loop contained in the [Airflow context](https://astronomer.io/docs/airflow/airflow-context.md) `context["loop"]` from within any task in the loop, including the `until` function.

```python theme={null}
from airflow.sdk import task

@task
def my_task(*, loop):
    print(loop.index)
```

There are four loop properties you can access and use in your looping logic:

* `loop.index`: The iteration that the task instance belongs to. Starts at 1.
* `loop.max_iterations`: The `max_iterations` value set in `.loop()` for this task group.
* `loop.previous`: The XCom pushed under the `return_value` key of the task group's terminal task in the *previous* iteration. This value is `None` on iteration `1`.
* `loop.result`: The XCom pushed under the `return_value` key of the task group's terminal task in the *current* iteration, which the `until` callable typically uses to determine whether to stop. This value is `None` if the terminal task in this iteration hasn't run yet.

### Simple task loop example

In the following example, a task group containing three tasks is repeated until a number that is doubled in each iteration exceeds 100 or the loop has run five times.

```python expandable theme={null}
from airflow.sdk import dag, task, task_group, chain


# function for the until condition
def over_100(*, loop):
    loop_result = loop.result["doubled_number"]
    print(f"Loop result: {loop_result}")
    return loop_result > 100  # Loop stops if True


@dag
def basic_task_loop_doubling():

    @task
    def my_upstream_task():
        import random

        return random.randint(1, 10)

    @task_group
    def my_task_group(my_number):
        @task
        def my_task_1(my_number, *, loop):
            return my_number if loop.index == 1 else loop.previous["doubled_number"]

        @task
        def my_task_2(my_number):
            print(f"Still looping, the number is: {my_number}")

        @task
        def my_task_3(my_number):
            return {"doubled_number": my_number * 2}

        _my_task_1 = my_task_1(my_number)
        _my_task_2 = my_task_2(my_number=_my_task_1)
        _my_task_3 = my_task_3(my_number=_my_task_1)

        # this dependency is needed because looped task groups need a singular final task
        chain(_my_task_2, _my_task_3)

    @task
    def my_downstream_task(my_results):
        print(my_results[-1])

    _my_upstream_task = my_upstream_task()

    _my_task_group = my_task_group.partial(my_number=_my_upstream_task).loop(
        max_iterations=5, until=over_100
    )

    _my_downstream_task = my_downstream_task(_my_task_group["my_task_3"].output)


basic_task_loop_doubling()
```

The Dag created by this code has one looped task group, `my_task_group`, which contains the three tasks defined in the Dag code and the `over_100` gate task that Airflow creates from the `until` function. Outside of the looped task group, there is one upstream task and one downstream task.

![Airflow UI Graph view of the basic\_task\_loop\_doubling Dag showing my\_upstream\_task, the looped my\_task\_group with a loop icon and the count 5 containing my\_task\_1, my\_task\_2, my\_task\_3, and the over\_100 LoopGateOperator task, followed by my\_downstream\_task](https://mintcdn.com/astronomer/s42Zt9FGU2QJB03d/images/img/guides/3-4_task_loops_graph_view.png?fit=max\&auto=format\&n=s42Zt9FGU2QJB03d\&q=85\&s=d3c2ff85f256a2eef6eda1ceedacfcbb)

In this Dag, the `my_upstream_task` task picks a random number between 1 and 10 that is doubled until the result is over 100.

For example, if the random number picked is 8:

* Iteration `1` of the loop uses `my_number=8` and produces a result of `16`; the `over_100` function returns `False` and the next iteration starts.
* Iteration `2` of the loop uses `my_number=16` and produces a result of `32`; the `over_100` function returns `False` and the next iteration starts.
* Iteration `3` of the loop uses `my_number=32` and produces a result of `64`; the `over_100` function returns `False` and the next iteration starts.
* Iteration `4` of the loop uses `my_number=64` and produces a result of `128`; the `over_100` function returns `True` and the loop exits.

The `my_downstream_task` prints the value pushed to XCom by the `my_task_3` task in the last task loop iteration, which is `128`.

In the Gantt view, you can see the tasks in `my_task_group` repeat sequentially. Each task appears four times, once for each of iterations `1` to `4`. `my_downstream_task` starts after the gate task of iteration `4` returns `True`.

![Airflow UI Gantt view of the basic\_task\_loop\_doubling Dag with the looped my\_task\_group selected, showing my\_task\_1, my\_task\_2, my\_task\_3, and over\_100 each running four times in sequence, followed by my\_downstream\_task, and the outcome Stopped early at iteration 4 with 4/5 iterations](https://mintcdn.com/astronomer/s42Zt9FGU2QJB03d/images/img/guides/3-4_task_loops_gantt_view.png?fit=max\&auto=format\&n=s42Zt9FGU2QJB03d\&q=85\&s=8d68eb606c16a243b68f165d0fcd1610)

The loop has a `max_iterations` value of `5`, which limits the loop to five iterations. For example, if `my_upstream_task` returns `1`, iteration `5` returns `32` after doubling five times, which is still below 100. The `over_100` function returns `False` in the last allowed iteration, so the `over_100` gate task fails with a `LoopMaxIterationsExceeded` error.

![Airflow UI Gantt view of the basic\_task\_loop\_doubling Dag with the looped my\_task\_group in the failed state, showing my\_task\_1, my\_task\_2, and my\_task\_3 succeeding in all five iterations, over\_100 succeeding in iterations 1 to 4 and failing in iteration 5, my\_downstream\_task in the upstream\_failed state, and the outcome Iteration 5 failed in my\_task\_group.over\_100 with 5/5 iterations](https://mintcdn.com/astronomer/s42Zt9FGU2QJB03d/images/img/guides/3-4_task_loops_max_iterations_failed_gantt.png?fit=max\&auto=format\&n=s42Zt9FGU2QJB03d\&q=85\&s=4b50c2e30705b71cd2af51ec10dcf886)

## Advanced concepts

### `until` function

If you add a docstring to the `until` function, Airflow shows it as the **Behaviour** of the looped task group in the Airflow UI and as the documentation of the gate task. Use it to describe the stop condition, as shown in the [deterministic example](#example-deterministically-process-a-hierarchy).

```python theme={null}
def hierarchy_ends(*, loop):
    """Stop once the level just populated has no child level."""
    return not loop.result["has_more"]
```

Besides `loop`, the `until` function can accept any other key of the [Airflow context](https://astronomer.io/docs/airflow/airflow-context.md) as a keyword argument, for example `params` or `ti`. The following Dag stops the loop once the result is at least the `threshold` value set in the Dag's [params](https://astronomer.io/docs/airflow/airflow-params.md), which means you can change the stop condition when triggering a Dag run.

```python expandable theme={null}
from airflow.sdk import dag, task, task_group


def at_least_threshold(*, loop, params):
    return loop.result >= params["threshold"]


@dag(params={"threshold": 10})
def basic_task_loop_until_params():
    @task_group
    def my_task_group():
        @task
        def my_count_task(*, loop):
            return (0 if loop.index == 1 else loop.previous) + 1

        my_count_task()

    my_task_group.loop(max_iterations=20, until=at_least_threshold)


basic_task_loop_until_params()
```

To use a value that another task determines at runtime, pull it from XCom inside the `until` function. In the following Dag, `my_fetch_threshold_task` returns the threshold, and the `until` function uses `ti.xcom_pull` to compare each iteration's result against it.

```python expandable theme={null}
from airflow.sdk import chain, dag, task, task_group


def at_least_threshold(*, loop, ti):
    threshold = ti.xcom_pull(task_ids="my_fetch_threshold_task")
    print(f"Iteration {loop.index}: result {loop.result}, threshold {threshold}")
    return loop.result >= threshold


@dag
def basic_task_loop_until_upstream_threshold():
    @task
    def my_fetch_threshold_task():
        # for example, read the threshold from a file in object storage
        return 4

    @task_group
    def my_task_group():
        @task
        def my_count_task(*, loop):
            return (0 if loop.index == 1 else loop.previous) + 1

        my_count_task()

    _my_fetch_threshold_task = my_fetch_threshold_task()
    _my_task_group = my_task_group.loop(max_iterations=20, until=at_least_threshold)

    chain(_my_fetch_threshold_task, _my_task_group)


basic_task_loop_until_upstream_threshold()
```

### Use loop results outside the loop

To pass the results of a looped task group to a downstream task, index the task group object with the task ID of its terminal task and use `.output`. The downstream task receives one value per iteration, in iteration order. Use `[-1]` to get the result of the last iteration and `len()` to get the number of iterations that ran.

```python theme={null}
@task
def my_downstream_task(my_results):
    print(f"Iterations: {len(my_results)}, last result: {my_results[-1]}")

_my_task_group = my_task_group.partial(my_number=5).loop(max_iterations=10, until=over_100)

my_downstream_task(_my_task_group["my_task_3"].output)
```

If the terminal task is dynamically mapped, the downstream task receives a flat sequence with one value per mapped task instance per iteration. In this case `[-1]` returns the result of the last mapped task instance in the last iteration, and `len()` is the number of iterations multiplied by the number of mapped task instances.

For more information, see the [Airflow documentation](https://airflow.staged.apache.org/docs/apache-airflow/3.4.0/authoring-and-scheduling/loops.html#use-a-loop-s-result-outside-the-loop).

### Trigger rules and reruns

Tasks inside a looped task group behave like any other Airflow task. Keep the following points in mind:

* Trigger rules work as usual, including for the gate task. You can't configure the gate task directly. It takes its `trigger_rule` and `retries` from the `default_args` of the Dag or of the looped task group, for example `@task_group(default_args={"retries": 2})`.
* If the terminal task of an iteration is skipped or fails, the gate task doesn't run and no further iteration is created. Tasks downstream of the looped task group are skipped or set to `upstream_failed` unless you adjust their [trigger rules](https://astronomer.io/docs/airflow/airflow-trigger-rules.md), for example to `none_failed` or `all_done`.
* A retry happens within the same iteration and doesn't count toward `max_iterations`.
* If the `until` function raises an exception, the gate task fails and is retried according to its `retries`. If `until` returns a value that isn't a boolean, or the loop ends at `max_iterations` without `until` returning `True`, the gate task fails without retrying.
* Manually marking the gate task as successful ends the loop without evaluating the `until` function. Downstream tasks run as if `until` had returned `True`.

When you clear a task inside a loop, the **Later iterations** option, selected by default, determines what happens to the iterations after it when the cleared tasks include the gate task, for example because **Downstream** is selected. If selected, later iterations are cleared and the gate task decides again whether the loop continues, so the loop can end earlier or run longer than before, up to `max_iterations`. If you deselect it, later iterations are kept as they are.

![Airflow UI Clear Task Instance dialog for my\_task\_group.my\_task\_2 in iteration 3 with the Downstream and Later iterations options selected, listing 8 affected tasks from iterations 3 and 4, including the over\_100 gate tasks and my\_downstream\_task](https://mintcdn.com/astronomer/s42Zt9FGU2QJB03d/images/img/guides/3-4_task_loops_clear_later_iterations.png?fit=max\&auto=format\&n=s42Zt9FGU2QJB03d\&q=85\&s=89953bf52981304fb93e2c8d960fbeb2)

For more information, see the [Airflow documentation](https://airflow.staged.apache.org/docs/apache-airflow/3.4.0/authoring-and-scheduling/loops.html#failures-skips-and-retries).

## Limitations and compatibility

The following table shows which combinations of task loops with dynamic task mapping and task groups are supported.

| Combination | Supported |
| - | - |
| Dynamically mapped task inside a looped task group ([example](https://airflow.staged.apache.org/docs/apache-airflow/3.4.0/authoring-and-scheduling/loops.html#mapped-tasks-inside-a-loop)) | Yes |
| Dynamically mapped task group inside a looped task group | Yes |
| Regular task group inside a looped task group | Yes |
| Looped task group inside a regular task group | Yes |
| Looped task group inside a dynamically mapped task group | No |
| Looped task group inside another looped task group | No |
| `.expand()` and `.loop()` on the same task group | No |
| `.loop()` on a task group created with the `TaskGroup` class | No |

You can't pull XCom from previous Dag runs for tasks inside a loop. Calling [`xcom_pull`](https://astronomer.io/docs/airflow/airflow-passing-data-between-tasks#xcom) with `include_prior_dates=True` for a task inside a loop raises an error.

If you use a [custom XCom backend](https://astronomer.io/docs/airflow/custom-xcom-backend-strategies.md), it must store values with the `store_value()` method instead of `serialize_value()` to support task loops. See [Use a custom XCom backend class](https://astronomer.io/docs/airflow/custom-xcom-backend-strategies.md#use-a-custom-xcom-backend-class). If you use the Object Storage XCom Backend, you need a version of the Common IO provider later than 1.10.1.

### Common errors

The following import errors might occur when you try to define a task loop:

| Error | Cause |
| - | - |
| `ValueError: A loop body must have exactly one terminal task definition` | The looped task group has more than one task without a downstream task inside the group. Use `chain()` to set a single final task. |
| `TypeError: missing a required keyword-only argument: 'loop'` | A task accesses `loop` from the context, but its task group isn't looped. This happens when you call the task group function without `.loop()` while a task inside it accepts the `loop` argument, or when you call `.expand()` on the task group before `.loop()`. A task group can't be both dynamically mapped and looped. |
| `AttributeError: 'LoopTaskGroup' object has no attribute 'expand'` | `.expand()` was called after `.loop()`. A looped task group can't be dynamically mapped. |
| `TypeError: Dag '<dag_id>' uses a loop, but XCom backend <backend> stores values through serialize_value rather than store_value` | The configured custom XCom backend doesn't implement `store_value()`, or you use the Object Storage XCom Backend with Common IO provider version 1.10.1 or earlier. Update the backend or upgrade the provider. |

The `LoopMaxIterationsExceeded` error is raised when the `until` function still returned `False` in the last allowed iteration according to the `max_iterations` limit. The gate task fails and tasks downstream of the looped task group are set to `upstream_failed`. This is intended behavior. To run downstream tasks regardless of the outcome of the looped task group, change the [trigger rule](https://astronomer.io/docs/airflow/airflow-trigger-rules.md) of the downstream tasks, for example to `all_done`.

## Example: Deterministically process a hierarchy

The following Dag processes a nested hierarchy one level per iteration, from a company down to individual assets. This pattern is common when you load parent-child records, for example to build a knowledge graph, and each level needs the results of the level above it.

```python expandable theme={null}
import json

from airflow.sdk import dag, task, task_group

MAX_LEVELS = 8

LEVEL_NAMES = [
    "company",
    "state",
    "location",
    "building",
    "floor",
    "unit",
    "asset",
    "subasset",
]


def hierarchy_ends(*, loop):
    """Stop once the level just populated has no child level."""
    return not loop.result["has_more"]


@dag(params={"depth": 4})
def hierarchy_walk():

    @task
    def seed(params):
        depth = int(params["depth"])
        print(f"This run walks {depth} level(s) (cap: {MAX_LEVELS})")
        return {"depth": depth, "path": []}

    @task_group
    def walk_level(start):
        @task
        def populate(start, *, loop):
            level = loop.index - 1
            previous = start if loop.index == 1 else loop.previous
            name = LEVEL_NAMES[level]
            path = previous["path"] + [name]
            node = {
                "level": level,
                "name": name,
                "risk_score": round(0.1 * loop.index, 3),
            }
            print(f"Populating level {level} ({name}), path: {' > '.join(path)}")
            return {"depth": previous["depth"], "path": path, "node": node}

        @task
        def decide(populated, *, loop):
            has_more = loop.index < populated["depth"]
            print(f"Level {loop.index - 1}: has_more={has_more}")
            return {**populated, "has_more": has_more}

        decide(populate(start))

    @task
    def merge(levels):
        merged = {
            "levels_walked": len(levels),
            "hierarchy": [level["node"] for level in levels],
        }
        print("Final knowledge graph:\n" + json.dumps(merged, indent=2))
        return merged

    _seed = seed()

    _walk_level = walk_level.partial(start=_seed).loop(
        max_iterations=MAX_LEVELS, until=hierarchy_ends
    )

    _merge = merge(_walk_level["decide"].output)


hierarchy_walk()
```

This Dag contains five tasks, three of which are in the looped task group `walk_level`, including the `hierarchy_ends` gate task:

* The `seed` task reads the `depth` [param](https://astronomer.io/docs/airflow/airflow-params.md), which sets how many levels the hierarchy has in this Dag run.
* In each iteration, the `populate` task uses `loop.previous` to extend the path of the level above and creates a node for the current level. Typically, this task would query the child records of the previous level's IDs.
* The `decide` task is the terminal task. It returns the populated level and whether another level exists below it.
* The `hierarchy_ends` gate task stops the loop when there is no further level. `max_iterations=MAX_LEVELS` limits the loop to eight levels.
* The `merge` task receives the result of every iteration through `_walk_level["decide"].output` and combines the nodes into one hierarchy.

The following screenshot shows a Dag run with the default `depth` of `4`. The loop stopped after iteration `4`, after 4 of the 8 allowed iterations. The **Behaviour** field shows the docstring of the `hierarchy_ends` function, and the **Iteration** filter shows the task instances of iteration `4`.

![Airflow UI Gantt view of the hierarchy\_walk Dag with the walk\_level task group selected, showing seed, four iterations of populate, decide, and hierarchy\_ends, and merge, next to the walk\_level details with the outcome Stopped early at iteration 4, 4/8 iterations, the Behaviour field Stop once the level just populated has no child level, and the task instances of iteration 4](https://mintcdn.com/astronomer/s42Zt9FGU2QJB03d/images/img/guides/3-4_task_loops_hierarchy_walk.png?fit=max\&auto=format\&n=s42Zt9FGU2QJB03d\&q=85\&s=92b900afcf0838af2f04e3f8a06db4ca)

## Example: Improve AI output with HITL

The following Dag drafts a reply to a customer support ticket with an AI agent and loops until a human approves the reply or escalates the ticket. It uses the `@task.agent` decorator from the [Common AI provider](https://astronomer.io/docs/airflow/airflow-common-ai-provider.md) and a [human-in-the-loop](https://astronomer.io/docs/airflow/airflow-human-in-the-loop.md) task. To run it, you need version 1.0 or later of the Common AI provider and an LLM connection with the ID `pydanticai_default`.

````python expandable theme={null}
from datetime import timedelta
from typing import Literal

from pydantic import BaseModel

from airflow.providers.standard.operators.hitl import HITLOperator
from airflow.sdk import Param, chain, dag, task, task_group

MAX_DRAFTS = 5

TICKET = {
    "ticket_id": "T-1042",
    "customer": "Benjamin Sisko",
    "subject": "Charged twice for my annual plan",
    "message": (
        "Hi, I upgraded to the annual plan yesterday and my card was charged twice ($240 each). "
        "I only want one plan. Can you refund the extra charge? This is pretty frustrating."
    ),
}


class TicketResponse(BaseModel):
    summary: str
    response: str
    priority: Literal["low", "medium", "high"]
    confidence_score: float
    suggested_tags: list[str]


def review_finished(*, loop):
    return loop.result["decision"] != "Request changes"


@dag
def agent_support_reply():

    @task
    def fetch_ticket():
        return TICKET

    @task_group
    def draft_support_reply(ticket):
        @task
        def gather_feedback(ticket, *, loop):
            if loop.index == 1:
                return {"ticket": ticket, "feedback": "", "previous_response": ""}
            return {
                "ticket": ticket,
                "feedback": loop.previous["feedback"],
                "previous_response": loop.previous["draft"].response,
            }

        @task.agent(
            llm_conn_id="pydanticai_default",
            system_prompt=(
                "You are a friendly and helpful support agent for a software company. "
                "Write a reply to the customer's ticket and address the customer by name."
            ),
            output_type=TicketResponse,
        )
        def draft_reply(context):
            ticket = context["ticket"]
            prompt = (
                f"Ticket {ticket['ticket_id']} from {ticket['customer']}\n"
                f"Subject: {ticket['subject']}\n\n{ticket['message']}"
            )
            if context["feedback"]:
                prompt += (
                    f"\n\nYour previous reply was:\n{context['previous_response']}"
                    f"\n\nThe reviewer asked for these changes:\n{context['feedback']}"
                )
            return prompt

        _gather_feedback = gather_feedback(ticket)
        _draft_reply = draft_reply(_gather_feedback)

        _review_reply = HITLOperator(
            task_id="review_reply",
            subject="AI support reply ready for review",
            body="""{% set draft = ti.xcom_pull(task_ids='draft_support_reply.draft_reply') %}
**Summary:** {{ draft.summary }}

**Suggested priority:** {{ draft.priority }}
**Confidence:** {{ "%.0f" | format(draft.confidence_score * 100) }}%
**Suggested tags:** {{ draft.suggested_tags | join(', ') }}

**Reply:**
```
{{ draft.response }}
```

- **Approve**: send this reply to the customer.
- **Request changes**: write your feedback below, the agent drafts a new reply.
- **Escalate to CRE**: hand the ticket to a human engineer.""",
            options=["Approve", "Request changes", "Escalate to CRE"],
            defaults="Escalate to CRE",
            params={"feedback": Param(None, type=["null", "string"], description="What should change?")},
            response_timeout=timedelta(hours=4),
        )

        @task
        def record_review(draft, review):
            decision = review["chosen_options"][0]
            feedback = review["params_input"].get("feedback") or ""
            print(f"Decision: {decision}\nFeedback: {feedback or 'none'}")
            return {"draft": draft, "decision": decision, "feedback": feedback}

        chain(_draft_reply, _review_reply)
        record_review(_draft_reply, _review_reply.output)

    @task
    def handle_ticket(ticket, reviews):
        final = reviews[-1]
        if final["decision"] == "Approve":
            print(f"Sending reply to {ticket['customer']} after {len(reviews)} draft(s):\n{final['draft'].response}")
        else:
            print(f"Escalating ticket {ticket['ticket_id']} to CRE after {len(reviews)} draft(s).")

    _fetch_ticket = fetch_ticket()

    _draft_support_reply = draft_support_reply.partial(ticket=_fetch_ticket).loop(
        max_iterations=MAX_DRAFTS, until=review_finished
    )

    _handle_ticket = handle_ticket(_fetch_ticket, _draft_support_reply["record_review"].output)


agent_support_reply()
````

This Dag contains seven tasks, five of which are in the looped task group `draft_support_reply`, including the `review_finished` gate task:

* The `fetch_ticket` task returns the support ticket.
* In each iteration, the `gather_feedback` task collects the reviewer feedback and the previous draft from `loop.previous`. In iteration `1`, there is no feedback yet.
* The `draft_reply` task uses an AI agent to write the reply. The `output_type` makes the agent return a `TicketResponse` object with a summary, the reply, a priority, a confidence score, and suggested tags.
* The `review_reply` task waits until a human chooses **Approve**, **Request changes**, or **Escalate to CRE** in the Airflow UI. The optional `feedback` param lets the reviewer describe the changes they want. Without a response within 4 hours, the task selects **Escalate to CRE**.
* The `record_review` task is the terminal task. It returns the draft, the decision, and the feedback.
* The `review_finished` gate task stops the loop unless the reviewer chose **Request changes**. If the reviewer requests changes in all five iterations, the gate task fails with a `LoopMaxIterationsExceeded` error.
* The `handle_ticket` task uses the last iteration's result to either send the reply or escalate the ticket.

The following screenshot shows the `review_reply` task in the first iteration of the loop. The first draft says that fixing the issue takes 5 to 7 days. In the `feedback` field, the reviewer asks the agent to change this to 2 business days and to sign the reply with the name Jani, then chooses **Request changes**.

![Airflow UI Required Action tab of the draft\_support\_reply.review\_reply task in the agent\_support\_reply Dag, showing the drafted reply to Benjamin, the Approve, Request changes, and Escalate to CRE options, and the feedback field filled in with a request to add the name Jani and to change the timeline to 2 business days](https://mintcdn.com/astronomer/s42Zt9FGU2QJB03d/images/img/guides/3-4_task_loops_ai_review_reply.png?fit=max\&auto=format\&n=s42Zt9FGU2QJB03d\&q=85\&s=22aab94e1097c2e7f679110198b44309)

In the second iteration, the `gather_feedback` task passes the feedback and the previous draft to the agent. The new draft is signed with the name Jani and states that the issue will be resolved in 2 business days. The reviewer chooses **Approve**, which ends the loop.

![Airflow UI Required Action tab of the draft\_support\_reply.review\_reply task in the second iteration of the agent\_support\_reply Dag, showing the revised reply signed by Jani, and the Gantt view with two iterations of the draft\_support\_reply loop](https://mintcdn.com/astronomer/s42Zt9FGU2QJB03d/images/img/guides/3-4_task_loops_ai_review_reply_2.png?fit=max\&auto=format\&n=s42Zt9FGU2QJB03d\&q=85\&s=c745f1e9b331864e6fbbd0e551c900a2)