Agent orchestration with Apache Airflow®

Agent orchestration means running a model in a model-and-tool-calling loop as an Apache Airflow® task. The task submits the system prompt and the user prompt to the model alongside the agent configuration and a set of available tools, and the AI harness orchestrates the model API calls and tool calls. Unlike a single model call, neither the number of model calls nor the sequence of actions is known before the task runs. This guide builds on AI orchestration and LLM orchestration with Apache Airflow®. Everything in the LLM guide still applies: an agent task receives context from upstream tasks and pushes its result to XCom. What’s added is the ability to perform actions in other systems through tool calls, and to make more than one model API call per task. This guide covers:

Why use Airflow for agent orchestration
The structure of an agent orchestration pipeline
How to give an agent tools and restrict tool actions
How to use tools you already have with an agent task
Common agent orchestration patterns
Single- and multi-agent patterns
How to make an agent task production-ready

Assumed knowledge

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

Airflow basics. See Introduction to Apache Airflow®.
Airflow decorators. See Introduction to the TaskFlow API and Airflow decorators.
Airflow connections. See Manage connections in Apache Airflow®.
Airflow hooks. See Airflow hooks.
Single model call as a task. See LLM orchestration with Apache Airflow®.

Why Airflow for agent orchestration

Everything in Why Airflow for AI orchestration applies to agent tasks. Three features are specific to a model running in a loop with tools:

Durable execution: AI agents often perform many steps in their loop within a task. The Airflow task state store feature allows you to save information at any point inside a task, which is available to the task if it needs to retry, for example after hitting a rate limit. The Common AI provider’s @task.agent can be set to cache all model call and tool call outputs by setting durable=True. See Recoverability.
Tool credentials: The Common AI provider contains several pre-built toolsets that can use Airflow connections to give your agent access to external systems. This means you can store and govern all credentials in a central location. See Restrict what an agent can do.
Dag-as-a-tool: A multi-step process that an agent would otherwise re-implement every time it runs can be defined as a Dag that the agent starts through the Airflow REST API. Turning multi-step workflows into Dags makes them more predictable and reliable, and saves token cost. See Dag-as-a-tool.

Anatomy of an agent orchestration pipeline

Single-agent orchestration pipelines often follow patterns similar to an LLM orchestration pipeline. There are two main differences:

The number of model calls isn’t fixed. The model calls a tool, reads the result, and decides what to do next. The loop ends when the model returns a final answer that satisfies the output_type, when a usage limit is exceeded, or when a tool keeps failing.
Upstream tasks assemble only the initial context. The agent can retrieve additional information through tool calls. If you know in advance which information your agent needs and can fetch it deterministically in an upstream task, doing so is usually more cost efficient than letting the agent gather it in its loop.

Multi-agent orchestration pipelines come in many different structures. See Multi-agent orchestration.

Agent tasks aren’t idempotent, and they vary more than single model calls do. Two runs on the same input can differ in the actions the agent takes, not only in the wording of the output. See Non-determinism and idempotency.

Run an agent with any harness

If you are using a Python-based agent harness today such as LangChain, LangGraph, CrewAI, or Temporal, you can use your existing harness inside of an @task task, as shown in Start with the harness you already have.

Run an agent with the Common AI provider

@task.agent runs a PydanticAI agent as a task, analogously to @task.llm but with additional agent-specific parameters.

support_agent.py

from typing import Literal

from airflow.providers.common.ai.toolsets.mcp import MCPToolset
from airflow.providers.common.ai.toolsets.sql import SQLToolset
from airflow.sdk import task
from pydantic import BaseModel, Field
from pydantic_ai.usage import UsageLimits


class TicketResponse(BaseModel):
    summary: str = Field(description="Two sentences at most, in plain language.")
    reply: str = Field(description="The drafted reply to the customer.")
    priority: Literal["P0", "P1", "P2", "P3", "P4"] = Field(
        description="Incident-style severity, P0 most urgent, P4 least."
    )
    escalate: bool = Field(
        description="True if the ticket needs a human owner rather than a reply."
    )


@task.agent(
    llm_conn_id="pydanticai_default",
    system_prompt=(
        "You draft replies to customer support tickets. Look up the customer's "
        "orders and shipments before you answer, and cite the product "
        "documentation you used. Never promise a delivery date that the "
        "shipment record does not confirm."
    ),
    output_type=TicketResponse,
    toolsets=[
        SQLToolset(
            db_conn_id="support_readonly",
            allowed_tables=["orders", "shipments"],
            max_rows=20,
        ),
        MCPToolset(mcp_conn_id="mcp_product_docs"),
    ],
    usage_limits=UsageLimits(
        request_limit=10,
        tool_calls_limit=15,
        total_tokens_limit=60_000,
    ),
)
def draft_reply(ticket: dict) -> str:
    return f"Draft a reply to this ticket:\n\n{ticket['body']}"  # The string the decorated function returns is the user prompt.

The following parameters are available in @task.agent:

Parameter	Description
`llm_conn_id`	The Airflow connection for the model provider. Required
`model_id`	The model to use, in the format `<provider>:<model>`. Default: `None`, which uses the model set in the connection
`system_prompt`	Instructions loaded into the context before the user prompt. Default: `""`
`output_type`	The type the agent has to return, enforced by the harness. Use a Pydantic model for structured output. Default: `str`, see Output type.
`toolsets`	The toolsets the agent can call, as a list of PydanticAI toolset instances. Each one builds its tools from an Airflow connection. Default: `None`. See Toolsets in the Common AI provider.
`agent_params`	Additional keyword arguments passed to the PydanticAI `Agent` constructor, for example `model_settings` or `output_retries`. Default: `None`
`usage_limits`	A `UsageLimits` object capping requests, tokens, and tool calls per run. Exceeding any cap fails the task with `UsageLimitExceeded`. Default: `None`, no limits
`enable_tool_logging`	Logs each tool call with its timing at INFO level and its arguments at DEBUG level. Default: `True`. See Observability
`durable`	Caches completed model calls and tool calls and uses the cached information during an Airflow retry instead of running them again. Can’t be combined with `enable_hitl_review` or `code_mode`. Default: `False`. See Recoverability
`enable_hitl_review`	Adds an iterative review loop controlled by `max_hitl_iterations` (default `5`), `hitl_timeout` (default no limit), and `hitl_poll_interval` (default `10.0` seconds). Needs Airflow 3.1 or later. Default: `False`. See Downstream review of agent outputs
`message_history`	Seeds the run with a prior conversation and pushes the full transcript to XCom under the key `message_history`, so a later Dag run can continue the session. Default: `None`, a single-turn run
`code_mode`	Collapses the tools into a single `run_code` tool that the model drives by writing Python, so one model turn can make many tool calls. Needs the `code-mode` extra. Default: `False`
`serialize_output`	Converts a Pydantic `output_type` to a dictionary before pushing it to XCom. Default: `False`

See @task.agent for more information.

Define the class you pass to output_type at module scope. A class nested inside the Dag function can’t be deserialized from XCom.

To try a single @task.agent task without starting a scheduler, add dag.test() to the Dag file and run it as a plain Python script. It executes every task in one process against your real connections, which makes it the fastest way to check a prompt change or a new output_type. Note that tasks that wait for human input can’t complete this way, because there is no way for you to respond. See Debug interactively with dag.test().

Restrict what an agent can do

An agent that can call a tool can perform any action that tool allows.

Prompt instructions are not a guardrail. There are many ways in which an agent stops following rules given in its context, from prompt injection to instructions lost during context window compaction. What an agent can actually do is limited only by the permissions on the tool, based on the credentials the toolset uses.

Toolsets in the Common AI provider

The Common AI provider contains toolsets that connect to external systems using an Airflow connection. By restricting the permissions on the credentials in the Airflow connection you can add a guardrail to what the agent can do using the toolset. Some toolsets have additional restriction options as listed in the table below.

Toolset	Tools	Additional restriction options
`SQLToolset`	`list_tables`, `get_schema`, `query`, and `check_query` against any database supported by a `DbApiHook`	`allowed_tables` (enforced on queries with `sqlglot`, see Pre-built-toolsets), `allowed_functions`, `schema`, `allow_writes` (`False` by default), and `max_rows` (50 by default).
`MCPToolset`	Whatever the Model Context Protocol (MCP) server exposes, over streamable HTTP or stdio	The MCP server’s configuration
`HookToolset`	The methods you name on any Airflow hook	`allowed_methods` (auto-discovery is disabled)
`DataFusionToolset`	`list_tables`, `get_schema`, and `query` against Parquet, CSV, Avro, and Iceberg files on object storage, through Apache DataFusion	`datasource_configs` (one per table), `allow_writes` (`False` by default), and `max_rows` (50 by default)
`AgentSkillsToolset`	The skill tools from `pydantic-ai-skills`, among them `read_skill_resource` and `run_skill_script`, over Agent Skills loaded from folders or Git repositories	`exclude_tools`, for example to disable `run_skill_script`, and `exclude_resources` glob patterns

See Toolsets for more information and examples for the above toolsets. You don’t list the tools in the system prompt. The harness sends the name, description, and parameter schema of every tool in toolsets to the model with each request. For a HookToolset, the descriptions are from the hook’s docstrings: the first paragraph of a method docstring becomes the tool description, and :param entries become the parameter descriptions.

Tool descriptions and schemas occupy context window tokens on every model call in the loop, so a toolset with many tools costs tokens even when the agent calls none of them. This is one reason to give an agent only the tools its task needs.

In addition to the Common AI toolsets, any PydanticAI toolset works with @task.agent.

Move tools between harnesses

Any agent harness can run inside a regular @task, the same way a single model call does (see Start with the harness you already have). What is specific to agents is the tools, but when moving between LangChain or LangGraph and the Common AI provider you don’t have to rewrite tools, because they bridge in both directions. First, install the langchain and sql extras:

apache-airflow-providers-common-ai[langchain, sql]==<version>

With the extras installed, existing LangChain tools work inside a Common AI agent through LangChainToolset:

langchain_tools_in_agent.py

from airflow.sdk import task
from pydantic_ai.ext.langchain import LangChainToolset


@task.agent(
    llm_conn_id="pydanticai_default",
    system_prompt=SYSTEM_PROMPT,
    output_type=TicketResponse,
    toolsets=[LangChainToolset([my_langchain_tool])],
)
def draft_reply(ticket: dict) -> str:
    return build_user_prompt(ticket)

Going the other way, airflow_toolset_to_langchain_tools converts an Airflow toolset into a list of LangChain StructuredTool objects, so a LangChain or LangGraph agent can call Common AI toolsets that use Airflow connections.

airflow_tools_in_langchain.py

from airflow.providers.common.ai.toolsets.langchain_bridge import (
    airflow_toolset_to_langchain_tools,
)
from airflow.providers.common.ai.toolsets.sql import SQLToolset
from airflow.sdk import task


@task
def generate_ai_response(ticket: dict) -> dict:
    from langchain_openai import ChatOpenAI
    from langgraph.prebuilt import create_react_agent

    toolset = SQLToolset(
        db_conn_id="support_readonly",
        allowed_tables=["orders", "shipments"],
        max_rows=20,
    )

    agent = create_react_agent(
        model=ChatOpenAI(model="gpt-5-mini"),
        tools=airflow_toolset_to_langchain_tools(toolset),
        prompt=SYSTEM_PROMPT,
        response_format=TicketResponse,
    )
    result = agent.invoke(
        {"messages": [{"role": "user", "content": build_user_prompt(ticket)}]}
    )
    return result["structured_response"].model_dump()

Bridged tools are called outside a PydanticAI agent run, and each call gets its own connection. An MCP server is reconnected on every tool call, so a server using the stdio transport starts a fresh process each time and loses any state it kept.

Downstream review of agent outputs

Agent output goes through the same two review steps as any other AI output, an AI-as-a-judge task and a human-in-the-loop task. See Review AI output for both, and Human-in-the-loop workflows with Apache Airflow® for the operators and example code. One review option only exists for agent tasks: enable_hitl_review adds an iterative loop where a reviewer reads the output, sends feedback in natural language, and the agent regenerates, under a HITL Review tab on the task instance.

enable_hitl_review can’t be combined with durable=True, which fails at parse time, or with message_history, which fails when the task runs. max_hitl_iterations counts the outputs shown to the reviewer, so the default of five allows feedback rounds at iterations one through four. Requesting changes at the limit fails the task with HITLMaxIterationsError without calling the model again.

Make agent tasks production-ready

An agent runs a variable number of model and tool calls per task, so the three properties in From prototype to production take more configuration here than they do for a single model call. Make LLM tasks production-ready covers what all AI tasks share, including output typing and retry policies. This section covers what the loop and the tools add.

Control

Scope the connections and toolsets. See Restrict what an agent can do.
Cap tool calls as well as tokens. Set tool_calls_limit in usage_limits alongside request_limit and total_tokens_limit. While tool calls are often cheap in terms of AI tokens, they can be expensive in the target system, for example when running a costly SQL query on a large table.
Add a human-in-the-loop step for high-stakes output. See Downstream review of agent outputs.

Observability

An agent run produces a tree of calls rather than one call, and that tree differs between runs. PydanticAI emits OpenTelemetry GenAI spans for the agent run, each model call, and each tool call, so the trace shows which tools the agent chose, in what order, how long each one took, and what each model call cost. The spans nest under the task span and inherit the Dag ID, run ID, task ID, and try number. An automatic retry reuses the task instance’s trace context, so every attempt appears on the same trace, distinguished by try number. Tracing is enabled the same way as for a single model call. See Observability in the LLM orchestration guide for the configuration.

Recoverability

A plain Airflow retry restarts the loop from the beginning, which means every model call and tool call runs again, incurring cost. Setting durable=True caches each completed model response and tool result, and an Airflow retry uses the cached information without re-executing model and tool calls that ran in the previous try of the task. The task logs show whether cached results were used:

[2026-04-26 20:37:33] INFO - Durable: replayed 3 cached steps (2 model, 1 tool), executed 4 new steps (2 model, 2 tool)

In Airflow 3.3 and later, the cache uses the task state store and needs no configuration (see Durable execution in the Common AI guide for earlier versions). Cache entries are written with NEVER_EXPIRE so they stay available however long a retry is delayed. The keys a run used are deleted when the task succeeds. A task that fails permanently leaves its entries behind until the Dag run is deleted. Durable execution has two limits:

Each cached step is verified against the current request before it’s used. When the prompt, model, model settings, tools, or message history changed since the failed attempt, the affected steps re-run and the task logs a warning.
Tool results are only cached for tools you pass through toolsets=. Tools that are made available to the agent another way, such as a capability in agent_params, re-run on every retry.

As with single model calls, which task failures to retry at all can be configured using Airflow’s retry policies, including the Common AI provider’s LLMRetryPolicy. See Recoverability in the LLM orchestration guide.

Agent orchestration patterns

The following patterns cover some of the most common agentic use cases.

Data exploration agent

A stakeholder asks a data question, which starts a Dag run. The agent queries the warehouse, reasons over the results, decides what to look at next, and writes a report that a downstream task sends back to Slack. Use an agent rather than text-to-SQL when one query isn’t enough to answer the question. A single model call gets one attempt with whatever schema context you gave it, while an agent inspects the schema, runs a query, and bases the next query on the result. This pattern needs two guardrails:

The agent can’t be allowed to modify or delete data. Give the SQLToolset a connection whose role has SELECT grants only, and leave allow_writes at False.
The agent can only see data the requesting stakeholder is allowed to see. A product manager asking about feature adoption shouldn’t get back sensitive personally identifiable information (PII). Map stakeholders to connections with matching grants, and select the connection ID from the request rather than hardcoding one.

Because the output goes to an internal stakeholder who can ask again or escalate to the data team, this pattern works without a human-in-the-loop review step in the pipeline. Astronomer runs a similar architecture internally, described in Building Kepler, Astronomer’s internal data assistant.

Support ticket agent

An incoming support ticket starts the Dag, the agent looks up account, order, and product documentation through its tools, and it drafts a reply. This is an AI data product for external stakeholders, which changes two things compared to the data exploration agent. First, tool scoping affects correctness as well as security: a connection that can read every customer’s records can leak customer A’s purchase history into customer B’s reply. Second, the pipeline needs quality control before the output leaves your organization, see Review AI output. For more information about this pattern, see the AI-powered education operations reference architecture, which runs the pipeline behind support for the Astronomer Academy.

Self-improving agents

The support ticket agent can’t learn from its mistakes. This changes when you capture feedback and make it available to future agent runs in a decision tracing context graph pattern:

A decision trace is everything relevant to one decision instance: the inputs, every decision by an agent and by a human, the reasoning behind each one, and the outcome.
A context graph is the accumulation of those traces for the same or similar business processes, which an agent can search for precedent on its next run.

To capture decision traces and context graphs you add additional Dags around the agent Dag: one that retrieves relevant past traces into the initial context, and one that assembles the trace after the fact from the AI draft, the AI review, the human decision and its stated reasoning, the final output, and the stakeholder’s response. For more information, see Context graphs for self-improving AI agents with Apache Airflow®.

Dag-as-a-tool

The patterns so far run agents inside an Airflow Dag. You can also invert the relationship and let an agent run a Dag, by giving it a tool that calls the Airflow REST API. This works for agents inside a task and for agents in a local harness such as Claude Code. Use this pattern when an agent skill describes a multi-step process that the agent rebuilds on every invocation. For example, an employee onboarding skill that creates accounts across identity, chat, source control, and payroll systems can be unreliable, especially if some steps need to wait on other steps to complete. It is also not cost efficient, because the code to perform the steps is regenerated on every skill use. When you rewrite the skill as a Dag-as-a-tool, Airflow handles the dependencies, the waiting, and the retries. The skill now only has two steps:

onboard-employee.md

1. Run the `onboard_employee` Airflow Dag using your `run_airflow_dag` tool providing the ID of the new employee as a Dag param.
2. Use the account details that are returned to send the employee a personalized welcome email with their login details.

The pattern is useful for work that a specialized model does better than a language model. For example, a sales assistant agent that needs a revenue estimate for an opportunity can run a Dag that performs inference with a trained regression model. A Dag used as a tool doesn’t have to be deterministic. It can contain LLM or agent tasks itself, so an agent can call your existing localization pipeline instead of translating text with its own prompt.

Avoid accidental cycles. If Dag A contains an agent that can start Dag B, Dag B must not contain an agent that can start Dag A.

Root-cause analysis and self-healing pipelines

When a Dag run fails, an agent with access to your Airflow environment can investigate the failure, form a hypothesis, and emit a recommendation that downstream tasks act on. A typical implementation looks like this:

A Dag run fails.
An alert configured on Dag run failures triggers an auto_fix Dag.
A task in auto_fix calls Otto, Astronomer’s data engineering agent, to investigate the failure using logs, past runs, deployment configuration, and lineage.
When the failure is caused by a mistake in the Dag code, a @task.agent task generates a code suggestion.
A final task opens a pull request with the fix.
A human reviews and merges the pull request.

Because the agent’s output is a pull request, no fix is merged without human approval. For step-by-step instructions, see How to use Otto to automatically investigate Dag failures and PR a fix.

Multi-agent orchestration

Some work is better split across several agents than handled by one, either because it divides into distinct roles or because a decision benefits from more than one perspective. In Airflow, each agent is a separate task with its own retries, logs, usage limits, and connections, and results move between them through XCom.

Orchestrator-worker

One orchestrator agent on an expensive model breaks the work into independent pieces, and cheaper worker agents complete them in parallel. The orchestrator does three things: it writes the instructions for each piece of work, it decides which model and which tools each piece needs, and it defines the success criteria that a downstream review step checks against. Choosing a model per piece of work is also called intelligent model routing. Because prompt, model_id, and system_prompt are all template fields on AgentOperator, you can map over sets of keyword arguments to create one worker task instance per piece of work, each with its own model.

Restrict which models the orchestrator can pick in its output_type as a Literal of the model IDs available from your model provider.

Agentic council

Several agents assess the same input from different perspectives, then a downstream task combines their conclusions. For example, one agent evaluates a startup’s funding application as a CTO, one as a CFO, and one as a head of product. Give the agents a shared location to exchange information, such as a prefix in object storage that each one can write to and read from, and private prefixes when you want agents to address each other directly. A downstream LLM task consolidates the assessments into one report.

Agentic software development

The orchestrator-worker pattern applied to a codebase: one agent plans the work from a bug report or feature request, developer agents implement the pieces, reviewer agents check each other’s work, and a consolidation agent resolves conflicts between them. Output quality improves most when the agents have clear success criteria, which is what test suites, linters, and documented conventions in the repository supply. Run the full test suite in a deterministic task after the agent tasks finish, and leave the pull request approval to a human.

Next steps

Evaluate the output of the AI with AI model evals.
Measure whether the AI data product is useful and optimize for business value with AI product evals.
Add human decisions to any pipeline with Human-in-the-loop workflows with Apache Airflow®.
Explore more AI decorators in Orchestrate AI tasks with Apache Airflow® and the Common AI provider.

Overview

Get started

Airflow and AI

Airflow concepts

Tutorials

Integrations & connections

Reference Architectures

Observability concepts

Glossary

Agent orchestration with Apache Airflow®

Assumed knowledge

Why Airflow for agent orchestration

Anatomy of an agent orchestration pipeline

Run an agent with any harness

Run an agent with the Common AI provider

Restrict what an agent can do

Toolsets in the Common AI provider

Move tools between harnesses

Downstream review of agent outputs

Make agent tasks production-ready

Control

Observability

Recoverability

Agent orchestration patterns

Data exploration agent

Support ticket agent

Self-improving agents

Dag-as-a-tool

Root-cause analysis and self-healing pipelines

Multi-agent orchestration

Orchestrator-worker

Agentic council

Agentic software development

Next steps

​Assumed knowledge

​Why Airflow for agent orchestration

​Anatomy of an agent orchestration pipeline

​Run an agent with any harness

​Run an agent with the Common AI provider

​Restrict what an agent can do

​Toolsets in the Common AI provider

​Move tools between harnesses

​Downstream review of agent outputs

​Make agent tasks production-ready

​Control

​Observability

​Recoverability

​Agent orchestration patterns

​Data exploration agent

​Support ticket agent

​Self-improving agents

​Dag-as-a-tool

​Root-cause analysis and self-healing pipelines

​Multi-agent orchestration

​Orchestrator-worker

​Agentic council

​Agentic software development

​Next steps

Assumed knowledge

Why Airflow for agent orchestration

Anatomy of an agent orchestration pipeline

Run an agent with any harness

Run an agent with the Common AI provider

Restrict what an agent can do

Toolsets in the Common AI provider

Move tools between harnesses

Downstream review of agent outputs

Make agent tasks production-ready

Control

Observability

Recoverability

Agent orchestration patterns

Data exploration agent

Support ticket agent

Self-improving agents

Dag-as-a-tool

Root-cause analysis and self-healing pipelines

Multi-agent orchestration

Orchestrator-worker

Agentic council

Agentic software development

Next steps