For the past two years, the standard pattern for running LLM calls in Airflow was a PythonOperator that imported the OpenAI client, called the API, and returned the result as XCom. It worked. But when the call failed at 3 AM, Airflow retried the entire task — including the database query that produced the input. When a model call consumed 50,000 tokens more than expected, there was nothing in the task logs to show it. When someone asked “can we have a human approve this before the next step runs?” the answer was “not without custom callback plumbing.”
The apache-airflow-providers-common-ai package, released in April 2026 for Airflow 3.0+, changes all of that. Every LLM call becomes its own named, logged, retryable Airflow task. Token budgets are enforced at the operator level. Human-in-the-loop approval is a single parameter. And from version 0.3.0 onward, LLMRetryPolicy lets an LLM classify the error before deciding whether to retry — so a rate-limit gets a smart backoff and an expired API key fails immediately rather than wasting five retry attempts.
This guide covers the full operator surface, the patterns that are working in production, and the gotchas that aren’t in the quickstart.
TL;DR
- The
common.aiprovider (GA April 13 2026) ships five operators:LLMOperator,LLMBranchOperator,LLMSQLQueryOperator,AgentOperator, andDocumentLoaderOperator— each has a matching@taskdecorator. - All operators are backed by PydanticAI and work with any provider it supports: OpenAI, Anthropic, Google, Bedrock, Vertex, Ollama, or any OpenAI-compatible endpoint. You configure the model via a
pydanticaiAirflow connection. UsageLimitsenforces per-task token and request budgets at runtime — when the limit is hit, the task fails and Airflow’s standard retry policy applies on top.LLMRetryPolicy(v0.3.0+, requires Airflow ≥ 3.3.0) asks an LLM to classify each task failure and returns RETRY, FAIL, or SKIP — with a reason string that appears in the task logs.- Human-in-the-loop approval is built in: set
require_approval=Trueon any@task.llmand the DAG pauses inawaiting_inputstate until a reviewer approves or rejects the LLM output from the Airflow UI. - The provider emits OpenTelemetry GenAI spans for every model call and tool call, routed through Airflow’s existing OTel exporter — no separate tracing setup needed.
- For privacy-sensitive environments, point
LLMRetryPolicyat a local Ollama model so exception traces never leave your infrastructure.
Why Burying LLM Calls in PythonOperator Was Always Wrong
The old pattern wasn’t wrong because it failed — it was wrong because failures were invisible. A PythonOperator that calls an LLM is, from Airflow’s perspective, a black box. The scheduler knows the task started and whether it succeeded or failed. It knows nothing about how many tokens were consumed, what model was used, how long the LLM call took relative to the surrounding code, or what the model actually returned before XCom captured the final value.
📷 every capability the old pattern lacked is a first-class parameter in the new operators — nothing to build from scratch
The new operators expose each LLM interaction as a distinct unit of work in the DAG graph. Retry logic applies specifically to the LLM call. Token consumption appears in task metadata. The model response is visible in XCom before any downstream task consumes it. And when something goes wrong, the task log says why the model call failed — not just that the PythonOperator raised an exception.
The Operator Surface: Five Tools, Clear Boundaries
📷 five operators, five distinct jobs — picking the wrong one is the most common setup mistake
| Operator / Decorator | Best For | Avoid When | Since |
|---|---|---|---|
LLMOperator / @task.llm | Single-turn: classify, summarize, extract, structured Pydantic output | Agent needs tool calls or multi-turn reasoning | v0.1.0 |
LLMBranchOperator / @task.llm_branch | LLM picks the next task from a declared list of choices | Branch logic can be expressed as a simple conditional | v0.1.0 |
LLMSQLQueryOperator / @task.llm_sql | NL → SQL → execute against a DB connection, returns rows | Generated SQL must be reviewed by a human before execution | v0.1.0 |
AgentOperator / @task.agent | Multi-turn loop with HookToolset, SQLToolset, or custom tools | Task is single-turn — use LLMOperator, it’s simpler | v0.1.0 |
DocumentLoaderOperator | Parse text, CSV, JSON, PDF, DOCX into list[dict] for downstream embedding | You need to embed on the fly — pair with a vector store operator | v0.3.0 |
LLMSchemaCompareOperator | Compare two schema versions, return structured diff and compatibility flag | Schema drift monitoring at scale — runs per-table, not per-schema | v0.2.0 |
The @task.llm Pattern in Practice
The decorator pattern is the cleanest entry point. The function body returns the user prompt as a string. Everything else — model routing, structured output, token limits, retry policy, human approval — is configured as decorator arguments.
from airflow.sdk import dag, task
from pydantic import BaseModel, Field
from pydantic_ai.usage import UsageLimits
from typing import Literal
class TicketTriage(BaseModel):
summary: str = Field(description="Two sentences max, plain language.")
priority: Literal["P0","P1","P2","P3","P4"]
needs_human_review: bool
@dag(schedule=None, tags=["ai","support"])
def support_triage():
@task.llm(
llm_conn_id="pydanticai_default", # Airflow connection: model + API key
system_prompt=(
"You triage incoming support tickets. "
"Do NOT answer the ticket — only classify it."
),
output_type=TicketTriage, # structured Pydantic output
usage_limits=UsageLimits(
request_limit=2,
total_tokens_limit=4_000, # fail task if exceeded
),
require_approval=False, # set True to pause for human review
retries=3,
)
def triage_ticket(ticket: dict) -> str:
# function body returns the user prompt only
return f"Triage this ticket:\n\n{ticket['body']}"
@task
def route_ticket(triage: TicketTriage):
if triage.priority in ("P0", "P1"):
print(f"URGENT: {triage.summary}")
# downstream logic here
tickets = [{"body": "Production DB down, all writes failing"},
{"body": "Can you add dark mode to the dashboard?"}]
results = triage_ticket.expand(ticket=tickets) # Dynamic Task Mapping
route_ticket.expand(triage=results)
support_triage()
Two things to notice. First, output_type=TicketTriage tells the operator to parse the model response into a validated Pydantic object — the downstream task receives a TicketTriage instance, not a raw string. Second, .expand(ticket=tickets) uses Airflow 3’s Dynamic Task Mapping to fan out one LLM task per ticket. Each becomes its own named, independently retryable task in the DAG — exactly the observability pattern the old PythonOperator loop couldn’t provide.
LLMRetryPolicy: The Error Classification Layer
Static retry policies treat all failures identically. A rate-limit error gets the same 60-second wait as a malformed API key. A task that failed because the model returned invalid JSON for a structured output retries with the exact same prompt that just failed. None of this is useful. LLMRetryPolicy, introduced in v0.3.0 and requiring Airflow 3.3+, replaces this with a classification step.
📷 llmretrypolicy fires between the failure and the next attempt — the scheduler stays fully deterministic, the llm only advises
The critical design point: the LLM does not modify the DAG, interact with workers, or make any scheduling decisions. It reads the exception class, message, and traceback, then returns one of three actions: RETRY (with an optional delay), FAIL (immediately, no more retries), or SKIP (mark the task as skipped and continue the DAG). The scheduler enforces the action. The LLM is advisory only — this is not autonomous AI modifying production infrastructure.
from datetime import timedelta
from airflow.sdk import task
from airflow.providers.common.ai.retry_policies import (
LLMRetryPolicy, RetryRule, RetryAction
)
SNOWFLAKE_INSTRUCTIONS = """
You classify Snowflake query task failures. Rules:
- OperationalError with "Connection reset": RETRY after 30s
- ProgrammingError with "SQL compilation error": FAIL immediately — bad SQL won't fix itself
- ConnectTimeout: RETRY after 60s, up to 3 times
- Any auth/credential error: FAIL immediately — retrying costs tokens for no benefit
- Unknown errors: RETRY once after 30s, then FAIL
Return 0 for errors that should not retry.
"""
snowflake_policy = LLMRetryPolicy(
llm_conn_id="pydanticai_default",
instructions=SNOWFLAKE_INSTRUCTIONS,
# fallback_rules fire without an LLM call — cheap and deterministic
fallback_rules=[
RetryRule(
exception=ConnectionError,
action=RetryAction.RETRY,
retry_delay=timedelta(seconds=30),
),
],
)
@task(retries=5, retry_policy=snowflake_policy)
def run_snowflake_enrichment(batch_id: str):
# your Snowflake query logic here
...
Privacy note on LLMRetryPolicy: by default, the exception message and traceback are sent to whichever model your llm_conn_id points to. If your exceptions may contain sensitive data — customer IDs in query parameters, table names from restricted schemas — either sanitise the traceback in instructions or point the policy at a local Ollama instance so the data never leaves your infrastructure: LLMRetryPolicy(llm_conn_id="ollama_local", model_id="ollama:llama3.2").
Human-in-the-Loop Approval
One of the most-requested Airflow features for AI pipelines is the ability to pause a DAG and wait for a human to review an LLM output before proceeding. The common.ai provider builds this directly into the operator with a single parameter. When require_approval=True, the task completes its LLM call, writes the output to XCom, and then transitions to awaiting_input state. A reviewer opens the Airflow UI, reads the generated output, and either approves (DAG continues) or rejects (task marked failed).
@task.llm(
llm_conn_id="pydanticai_default",
system_prompt="Draft a customer-facing email response to this support ticket.",
require_approval=True, # pause here for human review
allow_modifications=True, # reviewer can edit the text before approving
approval_timeout=timedelta(hours=4), # auto-fail if no review after 4h
)
def draft_response(ticket: dict) -> str:
return f"Write a response to: {ticket['body']}"
On Airflow 3.3+, the awaiting_input task state is natively supported — the provider uses it directly. On earlier 3.x versions it falls back to a DEFERRED state with a sensor polling for the approval signal. The allow_modifications=True parameter lets the reviewer edit the LLM’s draft in the UI before approving, so the approved text (not the raw model output) flows into XCom for downstream tasks.
AgentOperator and SQLToolset
For multi-turn workflows where the model needs to decide which tools to call and when, AgentOperator is the right choice. The most practical out-of-the-box toolset is SQLToolset, which gives the agent the ability to discover schema, run queries, and check results — all against a configured Airflow DB connection.
from airflow.sdk import dag, task
from airflow.providers.common.ai.toolsets.sql import SQLToolset
from pydantic_ai.usage import UsageLimits
@dag(schedule="@daily", tags=["ai","analytics"])
def daily_anomaly_report():
@task.agent(
llm_conn_id="pydanticai_default",
system_prompt=(
"You are a data analyst. Investigate the anomalies table "
"and produce a concise summary with root causes and affected row counts."
),
toolsets=[
SQLToolset(
conn_id="snowflake_prod",
allowed_tables=["anomalies", "orders", "customers"], # enforce allowlist
)
],
usage_limits=UsageLimits(total_tokens_limit=20_000),
retries=2,
)
def investigate_anomalies(run_date: str) -> str:
return f"Investigate anomalies logged on {run_date}. Focus on order_id failures."
investigate_anomalies(run_date="{{ ds }}")
daily_anomaly_report()
The allowed_tables parameter on SQLToolset is not just advisory — as of v0.5.0 (June 2026), it parses the agent’s SQL with sqlglot and rejects any query that touches a table outside the allowlist before execution. This includes subqueries, CTEs, JOINs, and set operations. It is an application-level guardrail, not a substitute for least-privilege DB permissions — use both.
If you’re running these agents against Snowflake and want to understand what they’re consuming token-wise, the same CORTEX_AGENT_USAGE_HISTORY monitoring patterns from our Cortex AI monitoring guide apply — except here the calls originate from Airflow rather than from inside Snowflake. You’ll want to add QUERY_TAG or a custom logging step to tie Airflow task IDs to Snowflake session metadata.
Connecting the Model: PydanticAI Connections
Every operator takes a llm_conn_id that points to a pydanticai connection type in Airflow. The connection stores the model ID and API key. You set it up once; all operators share it.
# Via Airflow UI: Admin → Connections → Add
# Conn Type: pydanticai
# Conn ID: pydanticai_default
# Host: (leave blank for cloud providers)
# Schema: anthropic # or openai, google-gla, bedrock, ollama
# Password: your-api-key
# Extra: {"model_id": "claude-sonnet-5"}
# Or via environment variable (for CI/CD):
export AIRFLOW_CONN_PYDANTICAI_DEFAULT='pydanticai://\
:your-api-key@/?schema=anthropic&extra={"model_id":"claude-sonnet-5"}'
# Local / air-gapped: Ollama
# Schema: ollama Host: http://localhost:11434
# model_id: llama3.2
One thing the old PythonOperator pattern couldn’t do: swap models without touching DAG code. Because the model is in the connection, you can maintain separate pydanticai_prod and pydanticai_dev connections pointing at different models and switch via Airflow Variables or environment overrides. This is particularly useful for local Ollama setups where you want to prototype with a small model before promoting to a frontier one.
The Gotchas
The provider is Airflow 3.0+ only — no backport to 2.x.This is stated explicitly in every release note. If your team is on Airflow 2.10 or earlier, none of these operators are available. The migration path to Airflow 3.0 involves removing execution_date references and updating provider pins — substantial work if your DAG codebase is large. Don’t plan a common.ai migration without first auditing your Airflow version.
LLMRetryPolicy requires Airflow ≥ 3.3.0 specifically.The policy object installs fine on 3.0 and 3.1 — there’s no import error. But the retry_policy parameter on @task is only honoured from 3.3.0 onward. On earlier versions the parameter is silently ignored and the task falls back to standard retry behaviour. Check airflow version before building retry logic around it.
UsageLimits failures are task failures — they trigger standard retries.When a task exceeds its UsageLimits, PydanticAI raises UsageLimitExceeded and the task marks FAILED. Airflow’s standard retry policy then applies on top. If you set retries=3 and the prompt is structurally over the token budget, you’ll burn three more calls before the task finally fails. Set retries=0 for token-budget failures, or add a LLMRetryPolicy rule that returns RetryAction.FAIL for UsageLimitExceeded.
SQLToolset’s allowed_tables allowlist rejects quoted identifiers.This is a documented limitation from v0.5.0: while an allowlist is active, SQL with quoted identifiers ("my_table" rather than my_table), inline comments, cross-database references, SHOW statements, table-valued functions, and dynamic SQL are all rejected before execution. Agents must send unquoted, comment-free SQL. If your schema uses quoted identifiers consistently, you’ll hit this immediately and need to decide whether to drop the allowlist or change the schema naming convention.
OpenTelemetry spans only export if your Airflow instance has an OTel exporter configured.The provider emits GenAI spans unconditionally, but they’re routed through Airflow’s existing OTel exporter. If you haven’t configured AIRFLOW__METRICS__OTEL_ON=True and an exporter endpoint, the spans are emitted and immediately dropped. Check your metrics config before assuming traces are flowing to your observability backend.
The One Principle
“Every LLM call is a task. If it isn’t a task, it isn’t observable, retryable, or governable — and those three things are what separate a prototype from production.”
FAQ
What is the Apache Airflow common.ai provider?
It’s a first-party Airflow provider (apache-airflow-providers-common-ai) that ships LLM-native operators for Airflow 3.0+. Released in April 2026, it lets you run LLM calls, agentic tool loops, NL-to-SQL tasks, and document parsing as standard Airflow tasks — with structured output, token budgets, human-in-the-loop approval, and intelligent retry policies built in. It uses PydanticAI under the hood and works with OpenAI, Anthropic, Google, Bedrock, Vertex, Ollama, and any OpenAI-compatible endpoint.
How does LLMRetryPolicy differ from standard Airflow retries?
Standard retries treat every failure identically — wait the configured delay, try again. LLMRetryPolicy sends the exception class, message, and traceback to an LLM, which classifies the error and returns one of three actions: RETRY (with a custom delay), FAIL (stop retrying immediately), or SKIP (mark the task as skipped and continue the DAG). The scheduler remains fully deterministic — the LLM only advises. LLMRetryPolicy requires Airflow 3.3.0 or later.
How do I add human approval to an Airflow LLM task?
Set require_approval=True on any @task.llm or LLMOperator. The task completes its model call, writes the output to XCom, and transitions to awaiting_input state. A reviewer approves or rejects via the Airflow UI. You can also set allow_modifications=True so the reviewer can edit the LLM output before approving, and approval_timeout to auto-fail if no review arrives within a set window.
Can I use the common.ai provider with local models?
Yes. Create a pydanticai connection with Schema set to “ollama”, Host set to your Ollama endpoint (e.g. http://localhost:11434), and model_id set to the local model name. This is especially useful for LLMRetryPolicy in privacy-sensitive environments where exception traces should not leave your infrastructure.
When should I use AgentOperator instead of LLMOperator?
Use AgentOperator when the task requires multiple tool calls or multi-turn reasoning — for example, querying a database, checking a result, then querying again based on what it found. Use LLMOperator for single-turn tasks like classification, summarization, or extraction that produce one output and finish. AgentOperator without toolsets is valid but if you don’t need tools, LLMOperator is simpler and more explicit.
Does the common.ai provider work with Dynamic Task Mapping?
Yes. All task-decorator variants ((@task.llm, @task.agent, etc.) support .expand() for Dynamic Task Mapping, so you can fan out one LLM task per item in a list. Each mapped instance becomes its own named, independently retryable task in the DAG graph — this is the core observability advantage over running a loop inside a single PythonOperator.
Related reading: Snowflake Cortex AI token monitoring · Using MCP Servers with Snowflake · Snowflake Dynamic Data Masking & Row Access Policies · Ollama inside Airflow DAGs for PII tagging · AI coding agents and pipeline security · common.ai provider docs (official) · Agentic workloads on Airflow 3 (official blog)
