> ## 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.
# LLM orchestration with Apache Airflow®
*LLM orchestration* means running a single model inference as an [Apache Airflow®](https://airflow.apache.org/) task, with optional upstream tasks preparing the context and downstream tasks consuming the result. One prompt goes in, one (structured) response comes out, and the rest of the pipeline treats that response like any other piece of data. It's the smallest unit of AI orchestration, and where most use cases start.
This guide builds on the concepts in [AI orchestration](/docs/learn/ai-orchestration-overview): the harness, the micro and macro orchestration split, and when a single model call is the right unit of work rather than an agent.
This guide covers:
* The structure of an LLM orchestration pipeline
* How to run the AI framework you already use inside a Dag
* How the Common AI provider exposes model configuration as task parameters
* Common LLM orchestration patterns
* How to make an LLM 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®](/docs/learn/intro-to-airflow).
* Airflow decorators. See [Introduction to the TaskFlow API and Airflow decorators](/docs/learn/airflow-decorators).
* Airflow connections. See [Manage connections in Apache Airflow®](/docs/learn/connections).
* Basic familiarity with LLM prompting.
## Why Airflow for LLM orchestration
Running a single model inference as an Airflow task offers a lot of advantages: harness and model provider independence, credentials and tools governed through Airflow connections, dependencies between the AI task and the deterministic tasks around it, a choice of scheduling options and execution environments, dynamic task mapping, retries, and human-in-the-loop steps. See [Why Airflow for AI orchestration](/docs/learn/ai-orchestration-overview#why-airflow-for-ai-orchestration) for the full list.
## Anatomy of an LLM orchestration pipeline
An LLM orchestration pipeline has the same general structure as other pipelines, like ELT/ETL, with the model call sitting in the transform position.
1. **Upstream tasks assemble the context.** They query a database, read files from object storage, call an API, or pull results from an upstream Dag, then format that data into the user prompt. Deterministic checks, such as data quality tests, belong here too.
2. **One task calls the model API.** The system prompt contains the task instructions and only changes when you change the code. The user prompt contains the data that differs per run. The response is pushed to [XCom](/docs/learn/airflow-passing-data-between-tasks).
3. **Downstream tasks consume the result.** They load it into a warehouse, branch on it, hand it to an AI or human reviewer, or feed it into another model call.
Airflow handles the macro orchestration around that model call, and the harness inside the task handles the micro orchestration. See [Micro and macro orchestration](/docs/learn/ai-orchestration-overview#micro-and-macro-orchestration) for the split that the next section builds on.
LLM tasks aren't [idempotent](/docs/learn/dag-best-practices#review-idempotency), so a rerun or a backfill can produce a different result from the same input. See [Non-determinism and idempotency](/docs/learn/ai-orchestration-overview#non-determinism-and-idempotency) for how to design around that.
## Start with the harness you already have
If your team already utilizes LangChain, LangGraph, CrewAI, or Temporal, you don't have to rewrite your code to orchestrate it with Airflow. Any harness runs inside a regular `@task`, because a `@task` is just Python. Each of the following examples implements the same single model call pattern, a support ticket summarized into a typed object, using different frameworks.
While the following examples utilize four of the most popular harnesses, the same principle applies to many more including custom implementations, as long as they have a Python SDK.
The most direct case: one chat model, structured output, one call.
```python summarize_langchain.py icon="python" theme={null}
from airflow.sdk import task
from pydantic import BaseModel
class TicketSummary(BaseModel):
summary: str
priority: str
@task
def summarize_ticket(ticket: dict) -> dict:
from langchain_openai import ChatOpenAI
model = ChatOpenAI(model="gpt-5-mini").with_structured_output(TicketSummary)
result = model.invoke(
[
("system", SYSTEM_PROMPT),
("user", f"Summarize this ticket:\n\n{ticket['body']}"),
]
)
return result.model_dump()
```
LangGraph adds a graph around LangChain's integrations. A single model call step is a one-node graph. The graph extends in agent orchestration patterns that include several model and tool call steps.
```python summarize_langgraph.py icon="python" theme={null}
from airflow.sdk import task
from pydantic import BaseModel
from typing_extensions import TypedDict
class TicketSummary(BaseModel):
summary: str
priority: str
class State(TypedDict):
ticket: str
summary: dict
@task
def summarize_ticket(ticket: dict) -> dict:
from langchain_openai import ChatOpenAI
from langgraph.graph import END, START, StateGraph
def summarize(state: State) -> dict:
model = ChatOpenAI(model="gpt-5-mini").with_structured_output(TicketSummary)
result = model.invoke(
[
("system", SYSTEM_PROMPT),
("user", f"Summarize this ticket:\n\n{state['ticket']}"),
]
)
return {"summary": result.model_dump()}
builder = StateGraph(State)
builder.add_node("summarize", summarize)
builder.add_edge(START, "summarize")
builder.add_edge("summarize", END)
graph = builder.compile()
return graph.invoke({"ticket": ticket["body"]})["summary"]
```
CrewAI models everything as an agent with a role, a goal, and a backstory. Give that agent no tools and a single task, and it's one inference step.
```python summarize_crewai.py icon="python" theme={null}
from airflow.sdk import task
from pydantic import BaseModel
class TicketSummary(BaseModel):
summary: str
priority: str
@task
def summarize_ticket(ticket: dict) -> dict:
from crewai import Agent, Crew
from crewai import Task as CrewTask
triage_agent = Agent(
role="Support triage analyst",
goal="Summarize incoming tickets and assign a priority.",
backstory="Experienced support engineer who triages a large inbox daily.",
llm="openai/gpt-5-mini",
)
summarize = CrewTask(
description=f"Summarize this support ticket:\n\n{ticket['body']}",
expected_output="A TicketSummary with a summary and a priority.",
output_pydantic=TicketSummary,
agent=triage_agent,
)
result = Crew(agents=[triage_agent], tasks=[summarize]).kickoff()
return result.pydantic.model_dump()
```
Temporal is a durable execution engine rather than an AI framework. The model call goes in an activity. Temporal saves the state of the task in between activities independently of Airflow. Airflow also has an abstraction that allows you to save state at any point in the task, see [Task state store in Apache Airflow®](/docs/learn/airflow-task-state-store).
```python summarize_temporal.py icon="python" theme={null}
import asyncio
from datetime import timedelta
from airflow.sdk import task
from pydantic import BaseModel
from temporalio import activity, workflow
TASK_QUEUE = "support-ticket-queue"
class TicketSummary(BaseModel):
summary: str
priority: str
@activity.defn
async def summarize_activity(ticket_body: str) -> dict:
from langchain_openai import ChatOpenAI
model = ChatOpenAI(model="gpt-5-mini").with_structured_output(TicketSummary)
result = await model.ainvoke(
[
("system", SYSTEM_PROMPT),
("user", f"Summarize this ticket:\n\n{ticket_body}"),
]
)
return result.model_dump()
@workflow.defn
class SummarizeTicketWorkflow:
@workflow.run
async def run(self, ticket: dict) -> dict:
return await workflow.execute_activity(
summarize_activity,
ticket["body"],
start_to_close_timeout=timedelta(minutes=2),
)
@task
def summarize_ticket(ticket: dict) -> dict:
from temporalio.client import Client
async def _run() -> dict:
client = await Client.connect("temporal:7233")
return await client.execute_workflow(
SummarizeTicketWorkflow.run,
ticket,
id=f"summarize-{ticket['ticket_id']}",
task_queue=TASK_QUEUE,
)
return TicketSummary.model_validate(asyncio.run(_run())).model_dump()
```
In production the Temporal worker runs as its own long-lived service that the Airflow task connects to. Because a running workflow stays addressable by its ID, several Airflow tasks can interact with the same long-running process through Temporal signals and queries.
If you instead start a worker inside the task to keep the example self-contained, pass `workflow_runner=UnsandboxedWorkflowRunner()` to the `Worker`.
These are all single inference calls. The same frameworks are also how you'd run a tool-calling loop in a task, which is [agent orchestration](/docs/learn/ai-orchestration-agent) rather than LLM orchestration.
Running harness code directly inside an Airflow task already gives you many of the benefits of using a macro orchestration layer: scheduling, setting upstream and downstream dependencies, [dynamic task mapping](/docs/learn/dynamic-tasks), retries, logging, alerting, and more.
But there are also elements of macro orchestration you are missing when you are translating your existing harness code into a regular `@task` task. The model credentials are stored in the harness framework's own configuration instead of an Airflow connection, human review is done outside of Airflow and not easily accessible in downstream tasks, and changing model providers means needing significant edits to your task code.
The Common AI provider, covered in the following section, moves each of those into Airflow: the credentials into a connection, the review step into the Airflow UI, and the model choice into configuration.
If a harness pulls in heavy or conflicting dependencies, run it in an isolated environment rather than in your main Airflow image. See [Run tasks in isolated environments in Apache Airflow®](/docs/learn/airflow-isolated-environments).
## Move to the Common AI provider
The [Common AI provider](/docs/learn/airflow-common-ai-provider) is the Airflow-native harness. It's built on PydanticAI, works with any [compatible model provider](https://pydantic.dev/docs/ai/models/overview/), and exposes credentials, output validation, spend limits, human review, and tracing as task parameters rather than as framework code.
The same single model call as a `@task.llm` task:
```python summarize_native.py icon="python" theme={null}
from typing import Literal
from airflow.sdk import task
from pydantic import BaseModel, Field
from pydantic_ai.usage import UsageLimits
class TicketSummary(BaseModel):
summary: str = Field(description="Two sentences at most, in plain language.")
priority: Literal["P0", "P1", "P2", "P3", "P4"] = Field(
description="Incident-style severity, P0 most urgent, P4 least."
)
@task.llm(
llm_conn_id="pydanticai_default",
system_prompt="You triage incoming support tickets. Do not answer the ticket.",
output_type=TicketSummary,
usage_limits=UsageLimits(request_limit=2, total_tokens_limit=4_000),
)
def summarize_ticket(ticket: dict) -> str:
return f"Summarize this ticket:\n\n{ticket['body']}"
```
The string the decorated function returns is the user prompt. Everything else is configured using task parameters:
| Micro orchestration aspect | Common AI equivalent |
| ------------------------------------------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| Model API keys and base URLs in framework config | An [Airflow connection](/docs/learn/airflow-common-ai-provider#set-the-pydanticai-connection). Switching from OpenAI to Anthropic is a connection change, which does not need any code changes. |
| Output parsing and schema retries | `output_type` set to a Pydantic model. If the LLM output does not fit the `output_type`, the model call is retried according to the `output_retries` (one by default) set in `agent_params`, which are separate from Airflow retries. |
| Usage and spend limitations | `usage_limits` with a `UsageLimits` object containing caps on requests, tokens, tool calls, and [more](https://pydantic.dev/docs/ai/api/pydantic-ai/usage/#pydantic_ai.usage.UsageLimits) |
| Human-in-the-loop | `require_approval` and `allow_modifications`, reviewed in the Airflow UI |
| Framework-specific tracing | OpenTelemetry GenAI spans are available through Airflow's existing exporter |
See [Orchestrate AI tasks with Apache Airflow® and the Common AI provider](/docs/learn/airflow-common-ai-provider) for the full parameter reference.
To try a single AI 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()`](/docs/learn/testing-airflow#debug-interactively-with-dag-test).
### Keep your model configuration in Airflow
Moving to the Common AI provider isn't all or nothing. If you want to keep writing LangChain code but stop hardcoding model configuration in it, `LangChainHook` builds a LangChain chat model from an Airflow connection:
```python summarize_hook.py icon="python" theme={null}
from airflow.sdk import task
from pydantic import BaseModel
class TicketSummary(BaseModel):
summary: str
priority: str
@task
def summarize_ticket(ticket: dict) -> dict:
from airflow.providers.common.ai.hooks.langchain import LangChainHook
model = LangChainHook(llm_conn_id="langchain_default").get_chat_model()
result = model.with_structured_output(TicketSummary).invoke(
[
("system", SYSTEM_PROMPT),
("user", f"Summarize this ticket:\n\n{ticket['body']}"),
]
)
return result.model_dump()
```
Compared to the earlier LangChain example, the model and the credentials now come from a connection, so they're governed alongside every other credential in your environment and changing the model is a connection edit. The rest of the chain is unchanged. This needs the `langchain` extra:
```text theme={null}
apache-airflow-providers-common-ai[langchain]==
```
*Tools* bridge in both directions too: PydanticAI wraps existing LangChain tools as a toolset a Common AI agent can use, and the provider converts an Airflow toolset into LangChain tools. Tools belong to agents rather than single model calls, so see [Move tools between harnesses](/docs/learn/ai-orchestration-agent#move-tools-between-harnesses) for both bridges, and [Recoverability](/docs/learn/ai-orchestration-agent#recoverability) for how to enable durable execution with the Common AI provider, to save the output of model and tool calls in between Airflow retries of a task.
## Make LLM tasks production-ready
[From prototype to production](/docs/learn/ai-orchestration-overview#from-prototype-to-production) covers the three properties that separate a prototype from a production pipeline: control, observability, and recoverability. This section is what each one means for a single model call.
### Control
* **Cap the spend.** `usage_limits` restricts number of requests, tokens, and tool calls per run. Exceeding a limit raises `UsageLimitExceeded` and fails the task.
* **Constrain the output.** `output_type` is enforced by the harness. A task either returns an object matching your schema or it fails. Be specific about the field type: a `priority` field typed as `str` invites the model to answer `"Critical (P0)"` in one run and `"P0"` in the next, which potentially causes issues in downstream tasks. The same field typed as `Literal["P0", "P1", "P2", "P3", "P4"]` can only ever return one of those five values.
* **Add a human gate for anything outward facing.** Setting `require_approval`, `allow_modifications`, and `approval_timeout` allows output to be reviewed, edited, and approved before a downstream task uses it. See [Generation with human approval](#generation-with-human-approval).
### Observability
The Common AI provider emits OpenTelemetry GenAI spans for each model call and tool call, and routes them through the existing Airflow OpenTelemetry exporter. The spans nest under the task span, so they are associated with the Dag ID, run ID, and try number, and you can see token usage per task run in whatever backend you use for OpenTelemetry.
Enable tracing by setting the following [environment variables](/docs/astro/environment-manager-overview) for your exporter and collector:
```text theme={null}
AIRFLOW__TRACES__OTEL_ON=True
AIRFLOW__COMMON_AI__OTEL_EXPORT_ENABLED=True
OTEL_EXPORTER_OTLP_ENDPOINT=
OTEL_SERVICE_NAME=
```
The provider doesn't configure an exporter of its own; it reuses the one core tracing installs. If `AIRFLOW__TRACES__OTEL_ON` isn't enabled in the worker process, the provider emits no spans at all.
`[common.ai] capture_content = True` adds prompt and completion text to the spans. Airflow's secret masking applies to logs and rendered template fields, not to span attributes, so that content is exported unredacted. Enable it for debugging in a trusted environment only.
Tracing tells you what a run cost and how long it took. It is usually gathered in [AI model eval](/docs/learn/ai-orchestration-model-evals) pipelines to connect model actions and token use to AI output quality.
### Recoverability
At scale, some fraction of model calls fail for transient reasons: rate limits, provider outages, or timeouts. Airflow retries handle those, but not every failure needs the same response. Waiting fixes a rate limit but not an expired API key.
Airflow's retry policies let you decide that per exception type. `ExceptionRetryPolicy` maps an exception to an action, with an optional delay and a reason that gets logged. Use it first: it's deterministic, costs nothing, and behaves the same on every run. See [Retry policies](/docs/learn/rerunning-dags#retry-policies).
For AI tasks, the failures worth retrying aren't always ones you can enumerate in advance. Providers change their error strings, and the same status code can mean "wait two seconds" or "you are out of quota for the month". The Common AI provider's `LLMRetryPolicy` covers that case by asking a model to classify the failure:
```python llm_retry_policy.py icon="python" theme={null}
from datetime import timedelta
from airflow.providers.common.ai.policies.retry import LLMRetryPolicy
from airflow.sdk import task
from airflow.sdk.definitions.retry_policy import RetryAction, RetryRule
llm_policy = LLMRetryPolicy(
llm_conn_id="pydanticai_default",
timeout=30.0,
fallback_rules=[
RetryRule(exception=ConnectionError, action=RetryAction.RETRY, retry_delay=timedelta(seconds=10)),
RetryRule(exception=PermissionError, action=RetryAction.FAIL),
],
)
@task(retries=5, retry_delay=timedelta(seconds=5), retry_policy=llm_policy)
def call_external_api(): ...
```
The model receives the exception message and returns a category such as `rate_limit`, `auth`, or `data`, whether to retry, a suggested delay, and its reasoning, which is written to the task log. A rate limit is retried with the delay the model suggests, which overrides the task's own `retry_delay`. An error classified as `auth` fails immediately instead of burning five paid retries.
`fallback_rules` is not a fast path for errors you already anticipated. The model is consulted on every failure, and the rules are evaluated as an `ExceptionRetryPolicy` only when the classification call itself fails, for example when your model provider is the system that's down. Without them, the task falls back to its normal retry behavior.
Two things follow from classification running on every failed attempt. It sends that task's exception message to a model provider, so the policy pushes the message through Airflow's secret masker by default, with `redact_exception`, `redactor`, and `max_exception_length` to control what leaves your environment. It also spends a model call per failure, which is why `model_id` is worth pointing at a small, cheap model for this job.
Retry policies require Airflow 3.3 or later, and work with all Airflow decorators and operators.
## LLM orchestration patterns
The following patterns cover a lot of common single-call use cases, and they can be combined within the same Dag.
### Summarization
Condense long text into a short form for human review or downstream storage.
The following example illustrates how to summarize legal documents. Define the summary as a schema to be able to use individual fields in downstream tasks:
```python summarization.py icon="python" theme={null}
from pydantic import BaseModel, Field
class LegalSummary(BaseModel):
summary: str
parties: list[str] = Field(description="Named individuals or entities in the document.")
key_dates: list[str] = Field(description="Deadlines, effective dates, or hearing dates.")
requires_attorney_review: bool = Field(
description="True if the document has ambiguous language or high monetary exposure."
)
```
Because a batch usually contains many documents, [map the task dynamically](/docs/learn/dynamic-tasks) over the list of inputs to get one summary per document. If a source document exceeds the model's context window, add a chunking step upstream.
### Classification and routing
Classify unstructured input against criteria written in natural language, then use the classification to decide what the Dag does next.
When the classification is only used for routing, let the model pick the branch directly with `@task.llm_branch`:
```python routing.py icon="python" theme={null}
from airflow.sdk import chain, dag, task
@task.llm_branch(
llm_conn_id="pydanticai_default",
system_prompt=f"Pick a downstream branch based on email priority: {EMAIL_CRITERIA}",
)
def route_email(email: str) -> str:
return f"Email: {email}"
# placeholder downstream tasks
@task
def handle_p0_p1(): ...
@task
def handle_p2_p3(): ...
@task
def handle_p4(): ...
chain(
route_email(fetch_email()),
[handle_p0_p1(), handle_p2_p3(), handle_p4()],
)
```
When you need the classification itself as data, for example to tag a record *and* route on its priority, use `@task.llm` with a structured output and a regular [`@task.branch`](/docs/learn/airflow-branch-operator) downstream. That keeps the classification available to every downstream task instead of only to the branch.
### Extraction of structured data
Pull specific fields out of free-form text such as contracts, transcripts, or scanned documents, then load them into a database or use them as features for a traditional machine learning model. The AI step exists in service of an otherwise conventional ETL or MLOps pipeline.
When using an LLM to generate features, make sure to keep the extraction task in a shared module so a training Dag and an inference Dag use exactly the same schema and prompt:
```python extraction.py icon="python" theme={null}
# include/extraction.py, imported by both Dags
from airflow.sdk import task
from pydantic import BaseModel, Field
class DealFeatures(BaseModel):
objections_raised: list[str] = Field(description="Objections the prospect raised.")
competitors_mentioned: list[str] = Field(description="Competitor names mentioned.")
budget_confirmed: bool = Field(description="True if the prospect confirmed budget.")
@task.llm(
llm_conn_id="pydanticai_default",
system_prompt="Extract structured deal signals from sales calls, emails, and research.",
output_type=DealFeatures,
)
def extract_deal_features(deal_context: str) -> str:
return f"Extract deal signals from the following:\n\n{deal_context}"
```
Mark AI-derived features as such, for example with an `ai_` column prefix, because a backfill can produce different values from the same source text.
### Transformation
Rewrite content according to rules: convert between formats, translate between human languages, or port code from one programming language to another. Inputs and outputs usually match in length and general structure.
For localization use cases, you'll often want to translate several text chunks into several languages at the same time. In Airflow you can use dynamic task mapping to [map over multiple parameters](/docs/learn/dynamic-tasks#mapping-over-multiple-parameters):
```python transformation.py icon="python" theme={null}
from airflow.sdk import task
@task.llm(
llm_conn_id="pydanticai_default",
system_prompt=(
"Translate the given text into the target language. Preserve formatting, "
"tone, and technical terminology exactly. Return only the translated text."
),
)
def translate_chunk(text_chunk: str, target_language: str) -> str:
return f"Translate the following text into {target_language}:\n\n{text_chunk}"
translate_chunk.expand(text_chunk=chunks, target_language=languages)
```
### Generation with human approval
Ask a model to write something new: a drafted reply, an executive summary, a personalized recommendation, or a block of configuration. Unlike the preceding patterns, the output is usually the deliverable rather than an input to another task, and it's often read by someone outside your organization. That raises the cost of a bad output and the importance of requiring human approval with `require_approval=True`, and adding the option to change the output with `allow_modifications=True`.
```python approval.py icon="python" theme={null}
from datetime import timedelta
from airflow.sdk import task
@task.llm(
llm_conn_id="pydanticai_default",
system_prompt=(
"You draft replies to customer support tickets. Be brief and specific. "
"Never promise a delivery date, a refund, or a feature that isn't confirmed."
),
require_approval=True,
allow_modifications=True,
approval_timeout=timedelta(hours=4),
)
def draft_reply(ticket: dict, account: dict) -> str:
return (
f"Draft a reply to this ticket from {account['name']}, "
f"who is on the {account['tier']} plan.\n\n{ticket['body']}"
)
```
The task generates the draft, then waits in the `awaiting_input` state until it receives a response from a reviewer.
The reviewer can open the task instance in the Airflow UI, read the draft under **Required Action**, and approve or reject it. With `allow_modifications=True` they can edit the text first, and the edited version becomes the task result pushed to XCom, which is accessible in downstream tasks. Alternatively, required actions can be responded to using the [Airflow REST API](/docs/learn/airflow-human-in-the-loop#human-in-the-loop-api-endpoints).
Approval always round-trips the output through a string, because the review body has to render as text. With the plain `str` output above, the reviewer edits prose and that string is the result. Combine approval with a Pydantic `output_type` and the reviewer edits JSON instead, which the provider then validates back into your model before pushing it to XCom.
If a reviewer's edited text fits the model, it gets cast back into the model class, but if not, the task doesn't fail. In that case, it returns the raw edited string, so a downstream task that expects an object receives text instead.
If `approval_timeout` passes with no response, the task fails with a `TimeoutError`.
This is a single approve, reject, or edit decision on one finished output. It's distinct from `@task.agent`'s `enable_hitl_review`, which is an iterative loop where a reviewer sends feedback and the agent regenerates, and which can be accessed under its own **HITL Review** tab. See [Agent orchestration](/docs/learn/ai-orchestration-agent).
If the decision is more complex than approve or reject, for example picking one of several downstream paths or supplying free-text input, use the human-in-the-loop operators in the standard provider as a downstream task instead. See [Human-in-the-loop workflows with Apache Airflow®](/docs/learn/airflow-human-in-the-loop).
### Text-to-SQL
Text-to-SQL is a special case of generation where the generated asset is code that runs inside the same pipeline.
`@task.llm_sql` can read information from a database, generate SQL, and push it to XCom. This decorator does not execute the generated SQL. A separate operator runs the query, which is what keeps SQL generation and execution separate and allows you to use different connection IDs with different permissions, as well as to add a human-in-the-loop task in between, if necessary.
The `@task.llm_sql` decorator reads table schemas through a `DbApiHook` and parses the generated SQL, so it needs the `sql` extra:
```text theme={null}
apache-airflow-providers-common-ai[sql]==
```
```python sql_generation.py icon="python" theme={null}
from datetime import timedelta
from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
from airflow.sdk import chain, task
@task.llm_sql(
llm_conn_id="pydanticai_default",
db_conn_id="snowflake_analytics",
table_names=["bookings", "payments", "cancellations"],
validate_sql=True,
require_approval=True,
allow_modifications=True,
approval_timeout=timedelta(hours=1),
)
def generate_sql(question: str) -> str:
return question
run_sql = SQLExecuteQueryOperator(
task_id="run_sql",
conn_id="snowflake_analytics",
sql="{{ ti.xcom_pull(task_ids='generate_sql') }}",
)
chain(generate_sql(question="{{ params.question }}"), run_sql)
```
With `require_approval=True` the task enters the `awaiting_input` state after the model call has returned the generated SQL. A reviewer opens the task instance in the Airflow UI, reads the generated SQL under **Required Action**, edits it when `allow_modifications` is set, and approves or rejects it. If `approval_timeout` passes with no response, the task fails.
Prompt instructions are not a guardrail. The only reliable guardrail for AI-generated SQL is the permission scope of the connection that executes it. Use a role with read-only grants on only the tables, columns, and rows the stakeholder is allowed to see.
When using `@task.llm`, the model only gets one attempt at generating SQL with whatever schema context you gave it. When a question needs several queries and reasoning over their results, use an [agent](/docs/learn/ai-orchestration-agent) instead.
## Next steps
* Add tools and multi-step reasoning with [Agent orchestration](/docs/learn/ai-orchestration-agent).
* Evaluate the output of the AI with [AI model evals](/docs/learn/ai-orchestration-model-evals).
* Measure whether the AI data product is useful and optimize for business value with [AI product evals](/docs/learn/ai-orchestration-product-evals).
* Look up decorator and operator parameters in [Orchestrate AI tasks with Apache Airflow® and the Common AI provider](/docs/learn/airflow-common-ai-provider).