Develop Apache Airflow® Dags with AI

AI coding agents can draft an Apache Airflow® Dag for you. Whether that draft is useful depends on the agent’s context, the quality of your spec, and whether your agent was able to iteratively improve its work by testing against given success criteria. Astronomer’s data engineering agent, Otto, has the best practices described in this guide built in, plus additional advanced capabilities for local data engineering. See Working with Otto. This guide covers:

How to use Otto to write best practice Dags
Writing a spec with defined success criteria
Letting an agent run in a loop to iteratively write Dags

Other ways to learnThere are multiple resources for learning about this topic. See also:

Blog post: Best practices for writing Airflow Dags with AI.

Assumed knowledge

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

A local Airflow environment. See Run Airflow locally.
An agent that knows about your Airflow environment. See AI context for data engineering.
Basic familiarity with an AI coding agent harness, such as Claude Code or Cursor.

Working with Otto

LabsThis feature is in Labs.

The easiest way to develop best practice Dags with AI is by using Otto, Astronomer’s data engineering agent. You can interact with Otto directly or delegate tasks to Otto from another coding agent like Claude Code, Codex, or Gemini CLI.

Work with Otto directly

You can work with Otto in the Astro UI, in the Astro IDE, through an API endpoint, or using the Astro CLI. This guide focuses on how to use the Astro CLI to interact with Otto. For other interfaces, see the Otto documentation. To interact with Otto from your terminal:

Install the Astro CLI. Make sure you are on the latest version.
Sign in to your Astro account with astro login (a free trial is available). This is necessary to work with Otto because Otto uses an LLM gateway hosted by Astronomer.
Make sure you are in your Astro project directory, or create a new one by running astro dev init.
Start your Airflow project by running astro dev start.

If your Dag connects to other systems, you’ll need to provide the credentials in Airflow connections. See Manage connections in Apache Airflow for general information on Airflow connections, and Manage Airflow connections, variables, and environment variables for how to add Airflow connections to Astro Deployments.

Run astro otto to start your interactive Otto session in the terminal. Type /model to choose your preferred model.
Ask Otto to write you a spec and describe the pipeline you’d like to create in natural language, for example, “Write me a spec to have an AI agent answer support tickets”.
Answer any clarifying questions Otto asks.
Read the spec Otto writes for the planned Dag and suggest any necessary changes. The following is a section of a detailed Otto-written spec for an AI orchestration pipeline.
Once you are happy with the spec, tell Otto to implement it. Otto automatically develops in an agentic loop, testing its own changes with astro dev parse, while being aware of your current Airflow environment, following the conventions in your project, and having access to all relevant Airflow-related skills.
Go to localhost:8080 to inspect and run your Dag. Ask Otto to make any necessary changes.

Delegate to Otto

If you prefer to work from within your existing AI agent setup, such as Claude Code, you can still use Otto to develop your Dags by delegating Dag-writing related tasks to Otto.

Add the Delegating to Otto skill to your AI agent.
Now, whenever you mention “use Otto” or a similar phrase, your AI agent delegates the work to Otto, running as a sub-agent.

Using other coding agents

If you cannot use Otto, you can manually follow this section’s steps to improve the Dag-writing capabilities of your generic AI agent.

Give your agent context

Since generic AI agents don’t have Otto’s skills assessing your Airflow environment, you need to give them this context explicitly. Relevant information includes your Airflow version, your provider versions, and the conventions your team has already settled on. See AI context for data engineering. With the context in place, use your harness’s planning mode, where the agent writes a plan without permission to edit any files. After you review and approve that plan, have it build a spec as described in the following section.

Write the spec

Once you are happy with your agent’s plan, tell it to turn it into a structured spec. You can think of a spec like a skill or workflow to follow for a specific goal, in this case to write an Airflow pipeline.

There are several tools that help your agent with writing specs. A commonly used one is spec-kit.

Some important considerations when writing a spec (or reviewing an AI-written spec) are:

Structure the spec with Markdown headings so the model works through requirements in a fixed order: goal, context, data contract (between tasks), and constraints.
Instruct the agent to write a skeleton Dag first using EmptyOperator placeholders for the task graph, then optionally stop to ask you to review the Dag structure before filling in the task logic.
Specify the format of the input and the output as well as the action you want performed. For example: input is a DataFrame with columns x and y, output is a dictionary with fields a and b.
State idempotency requirements explicitly.
Keep raw data separate from instructions. If the prompt includes an example payload, a schema, or sample output, wrap it in a fenced code block or an XML-style tag so the model doesn’t confuse the data with the instructions around it. Explicitly state that the example is a sample of the general structure.

The example at the end of this guide shows a full spec written this way: a goal, the environment and its constraints, the data contracts for what goes in and out of each task, the dependencies between tasks, and the acceptance criteria the finished Dag has to fulfill.

Define success criteria

A spec only helps if you can check whether the agent fulfilled it. Before you ask an agent to write a Dag, decide what “done” means, ideally in the form of programmatic tests. Tests give you and the agent something concrete to check against, and they let an agent iterate independently.

To learn more about different options for testing Airflow Dags, see the Best practices for testing Apache Airflow® Dags eBook.

When using the Astro CLI, you can define tests in the tests directory and run them with the astro dev pytest command (independently of whether they use the pytest package or another testing framework such as unittest). There are five main types of tests in Airflow:

Parsing tests check whether the Dag parses correctly or results in an import error. Your agent can run one with astro dev parse.
Unit tests check the logic in the Python functions an agent writes, such as a custom hook, a custom operator, or functions used in an @task or PythonOperator task. Write these the same way you would for any Python code. See Unit testing for examples.
Dag validation tests define rules for your Dags, for example requiring tags to be defined, or only allowing specific operators or Dag schedules. See Write Dag validation tests.
Integration tests check that individual tasks work against the systems they talk to, rather than against mocks. Often the easiest way to run one is with a helper Dag that interacts with an external system, for example querying a few rows from a specific table in a database and checking the schema against what your Airflow tasks expect.
End-to-end (E2E) tests run a whole Dag or a set of connected Dags. If your agent has access to the Astro CLI, it can trigger these itself to check its work. These Dag runs change data and incur cost in external systems; therefore, Astronomer recommends pointing your development environment at a development replica rather than at production.

Testing Airflow Dag code is a separate concept from testing the quality of the data orchestrated with Airflow Dags. For the latter, see Data quality and Airflow.

Safety

Before you let an agent iterate on its own, decide what it can do without you.

Decide which tools and commands your agent can run without asking. Allowlist safe commands such as parsing Dags or running tests.
Decide which actions need your approval every time. Anything you can’t easily undo belongs here, along with anything that writes outside your project folder.

Use version control and commit often when working with AI agents. Review diffs closely: agents sometimes refactor code you didn’t ask them to change. For more information, see Safety.

Let agents write in a loop

Once your spec, success criteria, and guardrails are in place, it is time to let the agent write your Dags while you grab a nice cup of tea. With tests defined, an agent can check its own output and only hand you a draft after its code passes the success criteria. The loop is the same one as in regular test-driven development: the agent writes code, runs a check, reads the result, and tries again until the check passes or until it hits a pre-defined timeout in terms of time or tokens spent.

Agentic hooks

For Dags, the fastest check in that loop is astro dev parse, which catches import errors. Beyond the command permissions covered in Restrict agent commands, Claude Code and similar harnesses let you enforce this check automatically with a hook that runs after every file edit. A PostToolUse hook, for example, can run astro dev parse after every Edit or Write call and feed the result back to the agent for the next iteration.

Claude Code settings.json example

The following settings.json example shows this in practice for Claude Code: an allow list for a small set of low-risk commands, a PreToolUse hook that blocks destructive commands like rm -rf or DROP TABLE outright, and a PostToolUse hook that runs astro dev parse after every Edit or Write call and logs the result, feeding a pass or fail message straight back to the agent.The env and enabledPlugins entries enable a Python language server. This lets the agent see unresolved imports and type errors from the provider versions installed in your project.

{
  "permissions": {
    "allow": [
      "Bash(astro dev parse)",
      "Bash(astro dev pytest *)",
      "Bash(astro dev logs *)",
      "Bash(astro dev run dags list)",
      "Bash(git diff *)",
      "Bash(git log *)"
    ]
  },
  "hooks": {
    "PostToolUse": [
      {
        "matcher": "Edit|Write",
        "hooks": [
          {
            "type": "command",
            "command": "ts=$(date '+%Y-%m-%d %H:%M:%S'); file=$(jq -r '.tool_input.file_path'); echo \"[$ts] PostToolUse Edit|Write on $file\" >> .claude/hook-activity.log"
          },
          {
            "type": "command",
            "command": "ts=$(date '+%Y-%m-%d %H:%M:%S'); file=$(jq -r '.tool_input.file_path'); case \"$file\" in *dags/*.py) output=$(astro dev parse 2>&1); if [ $? -eq 0 ]; then echo \"[$ts] PASS parse $file\" >> .claude/hook-activity.log; jq -n --arg msg \"Dag parse passed: $file\" '{decision: \"block\", reason: $msg}'; else echo \"[$ts] FAIL parse $file\" >> .claude/hook-activity.log; echo \"$output\" >> .claude/hook-activity.log; jq -n --arg msg \"Dag parse failed: $file\\n\\n$output\" '{decision: \"block\", reason: $msg}'; fi ;; esac"
          },
          {
            "type": "command",
            "command": "file=$(jq -r '.tool_input.file_path'); case \"$file\" in *dags/*.py) echo \"Auto-tested Dag: $file\" >> .claude/dag-changelog.md ;; esac"
          }
        ]
      }
    ],
    "PreToolUse": [
      {
        "matcher": "Bash",
        "hooks": [
          {
            "type": "command",
            "command": "cmd=$(jq -r '.tool_input.command'); if echo \"$cmd\" | grep -qE '(rm -rf|DROP TABLE|DROP DATABASE)'; then echo 'BLOCKED: Dangerous command detected' >&2; exit 2; fi"
          }
        ]
      }
    ]
  },
  "env": {
    "ENABLE_LSP_TOOL": "1"
  },
  "enabledPlugins": {
    "pyright-lsp@claude-plugins-official": true
  }
}

This example is for Claude Code. Other harnesses such as OpenAI Codex or Google Gemini CLI support similar automation, but with different configuration formats. See the relevant harness documentation for more information.

Example: Write an AI orchestration Dag

The following example shows a spec written with the structure from Write the spec, and the Dag an agent (Claude Code using Sonnet 4.6 with the astronomer/agent skills available) produced from that spec in a single pass.

Spec

# Spec: AI-Drafted Support Ticket Responder

## Goal
An Airflow Dag that drafts an LLM reply to a support ticket, routes it through a human reviewer, and sends the approved/edited reply. Teaching example for **agentic LLM tasks + human-in-the-loop (HITL) gates**.

## Stack / constraints
- Airflow 3.2, TaskFlow SDK (`@dag`, `@task`, `chain`, `Param`)
- LLM via `@task.agent` (`llm_conn_id="pydanticai_default"`)
- HITL via `HITLOperator` (standard provider)
- Structured LLM output validated by Pydantic
- Input: one ticket JSON file on disk

## Data contracts
- **Ticket (in):** `{ ticket_id, customer, customer_id, subject, body }`
- **DraftAnswer (LLM out):** `{ answer: str, urgency: "P0".."P4" }`

## Flow
1. **fetch_support_ticket**: read ticket JSON → dict. `TICKET_PATH = Path(__file__).resolve().parents[1] / "include" / "support_ticket.json"`
2. **draft_answer** (`@task.agent`): ticket → `DraftAnswer`. Persona: friendly Starfleet support, short replies for busy captains.
3. **human_in_the_loop** (`HITLOperator`): render ticket + draft. Options: `Approve AI answer` (default), `Request Changes`, `Manual Response`. Free-text param `override_llm_response` carries edit instructions or the manual reply.
4. **route** (`@task.branch`): reviewer choice → dispatch: approve→send; changes→rewrite; manual→process_manual_response.
5. **rewrite** (`@task.agent`): apply ONLY requested changes; preserve the rest verbatim.
6. **process_manual_response**: wrap reviewer's text as the reply (no LLM).
7. **send_email** (×3, one per branch): emit the final reply.

## Requirements
- Exactly one send task fires per run (mutually exclusive branch).
- LLM output MUST conform to `DraftAnswer`; retry on violation.
- Rewrite preserves unrequested content unchanged.
- HITL defaults to approve so an unattended run still completes.

## Non-goals
- No real email send (print/log is fine).
- No live ticket queue or batching (single static file).
- No reply persistence.

## Acceptance criteria
- Parses and runs end-to-end on Airflow 3.x.
- Each reviewer choice reaches its correct send task; the other two skip.
- Draft and rewrite both produce schema-valid output.

Result

import json
from pathlib import Path
from typing import Literal

from airflow.providers.standard.operators.hitl import HITLOperator
from airflow.sdk import Param, chain, dag, task
from pydantic import BaseModel, Field

TICKET_PATH = Path(__file__).resolve().parents[1] / "include" / "support_ticket.json"


class DraftAnswer(BaseModel):
    answer: str = Field(description="The reply to send to the customer")
    urgency: Literal["P0", "P1", "P2", "P3", "P4"] = Field(
        description="Ticket priority: P0 is critical, P4 is low"
    )


@dag
def answer_support_tickets_sonnet_4_6_with_skills():
    @task
    def fetch_support_ticket() -> dict:
        return json.loads(TICKET_PATH.read_text())

    @task.agent(
        llm_conn_id="pydanticai_default",
        system_prompt=(
            "You are a friendly Starfleet support officer. "
            "Keep replies short and practical — the captain is busy."
        ),
        output_type=DraftAnswer,
    )
    def draft_answer(ticket: dict) -> str:
        return (
            f"Draft a support reply for ticket #{ticket['ticket_id']} "
            f"from {ticket['customer']}.\n\n"
            f"Subject: {ticket['subject']}\n\n"
            f"{ticket['body']}"
        )

    review = HITLOperator(
        task_id="human_in_the_loop",
        subject="Review AI-drafted support reply",
        body=(
            "**Ticket**\n\n"
            "{{ ti.xcom_pull(task_ids='fetch_support_ticket') }}\n\n"
            "**AI Draft**\n\n"
            "{{ ti.xcom_pull(task_ids='draft_answer') }}"
        ),
        options=["Approve AI answer", "Request Changes", "Manual Response"],
        defaults="Approve AI answer",
        params={
            "override_llm_response": Param(
                "",
                type=["string", "null"],
                title="Edit instructions (for changes) or full reply (for manual)",
            )
        },
    )

    @task.branch
    def route(decision: dict) -> str:
        choice = decision["chosen_options"][0]
        if choice == "Request Changes":
            return "rewrite"
        if choice == "Manual Response":
            return "process_manual_response"
        return "send_email_approved"

    @task.agent(
        llm_conn_id="pydanticai_default",
        system_prompt=(
            "You are a friendly Starfleet support officer. "
            "Apply ONLY the requested changes to the draft. "
            "Preserve all other content verbatim."
        ),
        output_type=DraftAnswer,
    )
    def rewrite(draft: dict, decision: dict) -> str:
        instructions = decision["params_input"].get("override_llm_response", "")
        return (
            f"Rewrite the following support reply applying ONLY these changes: {instructions}\n\n"
            f"Original answer:\n{draft['answer']}\n\n"
            f"Original urgency: {draft['urgency']}"
        )

    @task
    def process_manual_response(decision: dict) -> dict:
        reply = decision["params_input"].get("override_llm_response", "")
        return {"answer": reply, "urgency": "P2"}

    @task
    def send_email(reply: dict):
        print(f"[SEND] urgency={reply['urgency']}\n\n{reply['answer']}")

    ticket = fetch_support_ticket()
    draft = draft_answer(ticket)

    chain(draft, review)

    router = route(review.output)

    rewritten = rewrite(draft=draft, decision=review.output)
    manual_reply = process_manual_response(decision=review.output)

    send_approved = send_email.override(task_id="send_email_approved")(reply=draft)
    send_rewritten = send_email.override(task_id="send_email_rewritten")(reply=rewritten)
    send_manual = send_email.override(task_id="send_email_manual")(reply=manual_reply)

    chain(router, [send_approved, rewritten, manual_reply])


answer_support_tickets_sonnet_4_6_with_skills()

Graph view of the Dag showing the task dependencies.

This is exactly what the spec produced, with no edits. Two things a human reviewer would likely change: the human-in-the-loop step renders the raw JSON payload instead of formatted Markdown, and process_manual_response hardcodes the urgency to P2 for every manually written reply, regardless of what the reply says.

Overview

Get started

Airflow concepts

Tutorials

Integrations & connections

Reference Architectures

Observability concepts

Glossary

Develop Apache Airflow® Dags with AI

Assumed knowledge

Working with Otto

Work with Otto directly

Delegate to Otto

Using other coding agents

Give your agent context

Write the spec

Define success criteria

Safety

Let agents write in a loop

Agentic hooks

Example: Write an AI orchestration Dag

​Assumed knowledge

​Working with Otto

​Work with Otto directly

​Delegate to Otto

​Using other coding agents

​Give your agent context

​Write the spec

​Define success criteria

​Safety

​Let agents write in a loop

​Agentic hooks

​Example: Write an AI orchestration Dag

Assumed knowledge

Working with Otto

Work with Otto directly

Delegate to Otto

Using other coding agents

Give your agent context

Write the spec

Define success criteria

Safety

Let agents write in a loop

Agentic hooks

Example: Write an AI orchestration Dag