> ## 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.
# Orchestrate AI tasks with Apache Airflow® and the Common AI provider
The [Airflow Common AI provider](https://airflow.apache.org/docs/apache-airflow-providers-common-ai/stable/index.html) is an [Airflow provider package](https://airflow.apache.org/docs/apache-airflow-providers/index.html) that contains several operators and other modules to add AI-based tasks to your Dags, from simple LLM calls to AI agents with access to Airflow-based tools. It is built on top of [PydanticAI](https://pydantic.dev/docs/ai/overview/) and can be used with any [compatible model provider](https://pydantic.dev/docs/ai/models/overview/), including OpenAI, Anthropic, Gemini, AWS Bedrock, HuggingFace, and more.
In this guide you'll learn:
* Basic AI concepts to understand how to use the Common AI provider.
* How to install the Common AI provider and connect Airflow to your model provider.
* How to use the Common AI decorators and operators.
* How to add toolsets, durable execution, and human-in-the-loop review to your agentic tasks.
## Assumed knowledge
To get the most out of this guide, you should have an existing knowledge of:
* Basic Airflow concepts. See [Introduction to Apache Airflow](/docs/learn/intro-to-airflow).
* Airflow operators and decorators. See [Airflow operators](/docs/learn/what-is-an-operator) and [Airflow decorators](/docs/learn/airflow-decorators).
* Airflow hooks. See [Airflow hooks](/docs/learn/what-is-a-hook).
## Concepts
The Common AI provider abstracts calls to LLMs (large language models) and LMMs (large multimodal models), often as an AI agent with tool access. The calls are made through PydanticAI, which provides a consistent interface for [all compatible model providers](https://pydantic.dev/docs/ai/models/overview/).
| Concept | Description |
| -------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| LLM | Large language model. A model that takes text input and generates text output. |
| LMM | Large multimodal model. A model that takes input from multiple modalities (text, images, audio, video, and other inputs) and generates output in one or more modalities. |
| [AI agent](#@task-agent) | An LLM or LMM that has access to a set of tools. Typically, agents perform multiple steps of output generation and tool calls to achieve a goal. The goal of an agent can range from providing output to performing complex tasks involving multiple systems. |
| [Tools](#toolsets) | Functions that an AI agent can call. Agents typically use tools to interact with an MCP or API to perform an action in another system. For example, retrieving data from a database or writing a file to object storage. |
| [MCP](#pre-built-toolsets) | [Model Context Protocol](https://modelcontextprotocol.io/docs/getting-started/intro), a standard way for agents to connect to tools. You can think of an MCP server as an agent-callable interface for a tool or data source, often wrapping an API. |
As soon as [tools](#toolsets) are added to an [AI agent](#@task-agent), the agent can independently call that tool and perform any action that is possible through the tool. If you are adding tools, especially custom tooling with access to production systems, make sure to scope the access AI agents have to prevent destructive actions like dropping a table or deleting a file. **Never rely on instructions in a prompt** to restrict your agent from performing actions, always enforce access control programmatically.
Keep in mind that LLMs and LMMs aren't deterministic and therefore tasks using these models aren't **idempotent**. This means you might get different results when rerunning or backfilling Dags that contain AI-based tasks, even if all inputs are the same.
## Install the Common AI provider
To use the Common AI provider, you need to be at least on Airflow 3.0 and install it by adding it to your `requirements.txt` file. Make sure to pin the latest [version](https://airflow.apache.org/docs/apache-airflow-providers-common-ai/stable/changelog.html).
```text wrap theme={null}
apache-airflow-providers-common-ai==
```
Most operators in the Common AI provider depend on modules from the [Airflow standard provider](https://airflow.apache.org/docs/apache-airflow-providers-standard/stable/index.html), which is pre-installed when using the Astro CLI. Additionally, you'll need to install the [Airflow Common SQL provider](https://airflow.apache.org/docs/apache-airflow-providers-common-sql/stable/changelog.html) when using [`SQLToolset`](#pre-built-toolsets), [`@task.llm_sql`](#@task-llm_sql), [`@task.llm_schema_compare`](#@task-llm_schema_compare), or other features that read database metadata through a `DbApiHook`.
```text wrap theme={null}
apache-airflow-providers-standard==
apache-airflow-providers-common-sql==
```
## Set the PydanticAI connection
The Common AI provider uses the same operators and decorators across many model providers by interacting with them through PydanticAI. You can you switch providers by updating the [Airflow connection](/docs/learn/connections).
The connection has the following format:
```text wrap theme={null}
AIRFLOW_CONN_PYDANTICAI_DEFAULT='{
"conn_type": "pydanticai",
"host": "",
"password": "",
"extra": {
"model": ":"
}
}'
```
You can set a default model for the connection by specifying `model` in the connection `extra` field in the format `:` (for example `anthropic:claude-opus-4-7`). Models can be overridden at the task level, as long as the model is accessible through the provided Airflow connection.
## Decide which decorator to use
The following table lists each decorator and operator in the Common AI provider and describes typical use cases.
| Decorator or operator | Description |
| ------------------------------------------------------------------------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| [`@task.agent`](#@task-agent) (`AgentOperator`) | Multi-step agent with toolsets, optional durable caching, or iterative human-in-the-loop review. Use it when a model needs access to external systems to perform a task. For example, an agent could assemble a feature proposal based on retrieving recent chat messages, the current product roadmap, and a support ticket with a feature request. |
| [`@task.llm`](#@task-llm) (`LLMOperator`) | Single-turn model call with optional approval and optional direct edits to the output in the Airflow UI. Use this decorator when you only need a single model call. |
| [`@task.llm_sql`](#@task-llm_sql) (`LLMSQLQueryOperator`) | Natural language to SQL with optional access to the table schema, and SQL validation. This decorator works well to create Dags that can run ad-hoc queries [on demand](/docs/learn/airflow-event-driven-scheduling). |
| [`@task.llm_branch`](#@task-llm_branch) (`LLMBranchOperator`) | Let the model pick one or more downstream tasks from [branches](/docs/learn/airflow-branch-operator). Use it when you want to pick different workflow paths based on unstructured input, for example routing support tickets to different teams. |
| [`@task.llm_file_analysis`](#@task-llm_file_analysis) (`LLMFileAnalysisOperator`) | Analyze files from object or local storage. This decorator is useful for document summarization across a set of files. |
| [`@task.llm_schema_compare`](#@task-llm_schema_compare) (`LLMSchemaCompareOperator`) | Compare table schemas across databases or mixed sources and summarize differences between them. Use this decorator when you have frequent schema changes that you need to detect to adjust downstream tasks. |
## `@task.agent`
The `@task.agent` decorator and `AgentOperator` let you run a PydanticAI agent as a task in your Airflow Dag.
```python expandable wrap theme={null}
from datetime import timedelta
from airflow.sdk import task
from pydantic import BaseModel
from pydantic_ai import FunctionToolset
class MyToolset(FunctionToolset): ... # your custom toolset
class MyOutputClass(BaseModel): ... # your custom output class
@task.agent(
llm_conn_id="pydanticai_default", # required: your Airflow connection ID
model_id=":", # default: None (uses the model set in the connection)
# Optional parameters:
system_prompt="", # default: ""
output_type=MyOutputClass, # can be a primitive type or a subclass of pydantic.BaseModel, default: str
toolsets=[MyToolset()], # can be set to a list of pydantic_ai.toolset instances and/or Airflow hooks, default: None
enable_tool_logging=True, # default: True
agent_params={"tool_timeout": 60.0}, # default: None
durable=False, # default: False
enable_hitl_review=True, # default: False
max_hitl_iterations=5, # default: 5
hitl_timeout=timedelta(minutes=5), # default: None
hitl_poll_interval=10.0, # default: 10.0
)
def my_agentic_task(my_input: str) -> str:
return f"This is the user prompt! {my_input}"
my_agentic_task(my_input="Say hello!")
```
The string returned by the `@task.agent` decorated function is the *user prompt input* to the agent. When using the `AgentOperator` directly, providing a user prompt with the `prompt` parameter is required.
The only other mandatory parameter is `llm_conn_id`, the [Airflow connection ID](#set-the-pydanticai-connection) to use for the LLM.
The following optional parameters are available for the `@task.agent` decorator:
* `model_id`: The model to use for the agent. If not given, the operator uses the model set in the [Airflow connection](#set-the-pydanticai-connection). You need to set the model in the format `:`, for example `anthropic:claude-opus-4-7`.
* `system_prompt`: The system prompt to use for the agent, which is loaded into context before the user prompt. It is common to give general instructions to the agent in the system prompt about their role and task, as well as information about available tools. Note that instructions given in the prompt aren't guaranteed to be followed.
* [`output_type`](#output-type): The format for the output from the agent. You can use primitive types like `str`, `int`, `float`, `bool`, or a subclass of [Pydantic BaseModel](https://pydantic.dev/docs/validation/latest/api/pydantic/base_model/) for more complex structured outputs. Defaults to `str`. The `output_type` format is enforced, in contrast to instructions in the prompt.
* [`toolsets`](#toolsets): The toolsets the agent has access to as a list of [`pydantic_ai`.toolset](https://pydantic.dev/docs/ai/api/pydantic-ai/toolsets/) instances. Defaults to `None`.
* `enable_tool_logging`: If `True`, every toolset is wrapped in a `LoggingToolset` that logs tool calls with timing at INFO level and arguments at DEBUG level. Defaults to `True`.
* `agent_params`: Additional keyword arguments passed to the agent constructor. See the [PydanticAI agent documentation](https://pydantic.dev/docs/ai/api/pydantic-ai/agent/#parameters) for a list of available parameters.
* `durable`: Whether to enable step-level caching of model responses and tool results. Note that in order to use durable execution you need to set `AIRFLOW__COMMON_AI__DURABLE_CACHE_PATH`. See [durable execution](#durable-execution). Defaults to `False`. Can't be used with human-in-the-loop review.
* [`enable_hitl_review`](#human-in-the-loop-review): Whether to enable human-in-the-loop review through an Airflow plugin. Defaults to `False`. Needs Airflow 3.1+ and can't be used with durable execution.
* `max_hitl_iterations`: The maximum number of iterations of human-in-the-loop review. Defaults to `5`.
* `hitl_timeout`: The timeout for the human-in-the-loop review as a timedelta object. Defaults to `None`.
* `hitl_poll_interval`: The interval between polling for human-in-the-loop review. Defaults to `10.0`.
### Toolsets
AI agents can use tools to perform actions in another system. Tools are typically implemented as functions that can be called by the agent. A toolset is a collection of such functions in a class.
When using the Common AI provider, your agents can use three types of toolsets:
* Pre-built toolsets included in the Common AI provider: `SQLToolset` and `MCPToolset`.
* Use any [Airflow hook](/docs/learn/what-is-a-hook) as a toolset by wrapping it in the `HookToolset` class, the hook methods are the tools the agent can call.
* Custom toolsets you can implement yourself by subclassing a [`pydantic_ai`.toolset](https://pydantic.dev/docs/ai/api/pydantic-ai/toolsets/#functiontoolset) class.
#### Pre-built toolsets
The [SQLToolset](https://airflow.apache.org/docs/apache-airflow-providers-common-ai/stable/_api/airflow/providers/common/ai/toolsets/sql/index.html) is a curated toolset that gives your agent access to a SQL database, only allowing the following methods: `list_tables`, `get_schema`, `query` (`SELECT` only by default), and `check_query`. It can be used with any SQL database that is supported by the `DbApiHook`.
```python wrap theme={null}
from airflow.providers.common.ai.toolsets.sql import SQLToolset
toolsets=[
SQLToolset(
db_conn_id="my_sql_conn",
allowed_tables=["my_table"], # restrict which tables the agent can access through list_tables and get_schema
schema="my_schema", # Database schema/namespace for table listing and introspection.
allow_writes=False, # If True, modify operations are allowed (insert, update, delete). Default: False.
max_rows=1000 # Maximum number of rows to return from a query. Default: 50.
)
],
```
The `allowed_tables` parameter does *not* parse or validate table references in SQL queries. An LLM can still query tables outside this list if it guesses the name. For query-level restrictions, use database-level permissions (for example a read-only role with grants limited to specific tables).
Setting `allow_writes=True` allows the agent to perform modify operations (insert, update, delete) on the database. Use at your own risk.
The [MCPToolset](https://airflow.apache.org/docs/apache-airflow-providers-common-ai/stable/_api/airflow/providers/common/ai/toolsets/mcp/index.html) is a toolset that gives your agent access to an MCP server. It uses the `MCPHook` to retrieve credentials from an [Airflow connection](/docs/learn/connections) and creates a [PydanticAI MCP server instance](https://pydantic.dev/docs/ai/api/pydantic-ai/mcp/).
```python wrap theme={null}
from airflow.providers.common.ai.toolsets.mcp import MCPToolset
toolsets=[MCPToolset(mcp_conn_id="mcp_default")],
```
Set the [MCP server connection](https://airflow.apache.org/docs/apache-airflow-providers-common-ai/stable/connections/mcp.html) as an Airflow connection. There are two main types of transports for MCP servers:
* **Streamable HTTP**: Uses HTTP to stream responses, recommended for remote servers. If your `host` doesn't require authentication, you can omit the `password` field.
```text wrap theme={null}
AIRFLOW_CONN_MCP_DEFAULT='{
"conn_type": "mcp",
"host": "http://localhost:3001/mcp",
"password": "",
"extra": {
"transport": "http"
}
}'
```
* **stdio**: Runs the MCP server as a subprocess communicating over stdin/stdout.
```text wrap theme={null}
AIRFLOW_CONN_MCP_DEFAULT='{
"conn_type": "mcp",
"extra": {
"transport": "stdio",
"command": "uvx",
"args": ["-m", "mcp_server"]
}
}'
```
#### Use Airflow hooks as toolsets
You can use any [Airflow hook](/docs/learn/what-is-a-hook) as a toolset by wrapping it in the `HookToolset` class. The hook methods are the tools the agent can call. The `allowed_methods` parameter is required to make methods available to the agent, auto-discovery is intentionally disabled for safety purposes.
```python wrap theme={null}
from airflow.providers.common.ai.toolsets.hook import HookToolset
toolsets=[
HookToolset(
hook=MyHook(conn_id="my_conn"),
allowed_methods=["method1", "method2"],
tool_name_prefix="my_prefix_", # optional, default: ""
)
],
```
You can explore available hooks in the [Airflow registry](https://airflow.apache.org/registry/).
### Human-in-the-loop review
Human-in-the-loop review lets you iterate on the output of an agentic task by reviewing the agent output and providing feedback. An Airflow plugin adds a **HITL Review** tab to the task instance page in the Airflow UI. This is different from [human-in-the-loop operators](/docs/learn/airflow-human-in-the-loop), which surface decisions under **Required Actions**.
```python wrap theme={null}
from airflow.sdk import task
from datetime import timedelta
@task.agent(
enable_hitl_review=True,
max_hitl_iterations=5,
hitl_timeout=timedelta(minutes=5),
hitl_poll_interval=10.0,
)
def my_agentic_task(my_input: str) -> str:
return f"This is the user prompt! {my_input}"
my_agentic_task(my_input="Say hello!")
```
On each iteration you have three options:
* **Approve**: The task succeeds with the current assistant output and the HITL review ends.
* Provide feedback to the agent and click **Send**: The agent will take the feedback as user prompt and regenerate the output. The maximum number of iterations is controlled by the `max_hitl_iterations` parameter.
* **Reject**: The task fails and the next iteration isn't started.
The `hitl_timeout` parameter provides an optional time limit for the entire HITL review phase, that is, all review rounds combined. When that limit is exceeded, the task fails. The `hitl_poll_interval` parameter is the number of seconds to wait between polls for the reviewer action on the human-action XCom key.
You can have human-in-the-loop steps in Airflow Dags outside of agentic tasks by using standalone [HITL operators](/docs/learn/airflow-human-in-the-loop).
### Durable execution
The `durable` parameter allows you to enable step-level caching of model responses and tool results. When a task is retried, steps that have already finished don't run again; the task reads their results from the cache instead.
To enable durable mode you need to set the `AIRFLOW__COMMON_AI__DURABLE_CACHE_PATH` [environment variable](/docs/astro/manage-env-vars) to the path where the cache will be stored. The destination can be a local temporary directory on the worker or remote object storage like S3 or GCS. After successful task execution the cache is deleted.
```text wrap theme={null}
AIRFLOW__COMMON_AI__DURABLE_CACHE_PATH=/path/to/cache
```
To use remote object storage, add your connection ID to the path and provide the credentials in the connection, similar to the configuration of an [Object Storage XCom Backend](/docs/learn/custom-xcom-backends-tutorial). To use S3 as the durable cache destination, set the following:
```text wrap theme={null}
AIRFLOW__COMMON_AI__DURABLE_CACHE_PATH=s3://@my-bucket/some/prefix/
AIRFLOW_CONN_='{
"conn_type": "aws",
"login": "",
"password": "",
"extra": {
"region_name": ""
}
}'
```
To run an agent in durable mode set `durable=True`.
```python wrap theme={null}
from airflow.sdk import task
@task.agent(
durable=True,
)
```
When a durable agent retries and uses previously cached results you'll see a log message like this:
```text wrap theme={null}
[2026-04-26 20:37:33] INFO - Durable: replayed 3 cached steps (2 model, 1 tool), executed 4 new steps (2 model, 2 tool)
```
You can only use `durable=True` if `enable_hitl_review=False`.
### Output type
Pass `output_type` to set the format of the agent output. It can be a primitive type like `str`, `int`, `float`, `bool`, or a subclass of [Pydantic BaseModel](https://pydantic.dev/docs/validation/latest/api/pydantic/base_model/) for more complex structured outputs. Use the `Field` class to add descriptions to the fields for the agent to use.
```python wrap theme={null}
from typing import Literal
from pydantic import BaseModel, Field
class MyOutputClass(BaseModel):
my_field_one: Literal["A", "B", "C", "D", "F"] = Field(
description=(
"Letter grade for xyz..."
)
)
my_field_two: str = Field(
description="String that describes xyz..."
)
my_field_three: list[str] = Field(
description=(
"List of strings that describe xyz..."
)
)
my_field_four: int = Field(
description=(
"Number that describes xyz..."
)
)
```
Using the above `MyOutputClass` as the `output_type` parameter the agentic task will always produce a JSON output with the fields and their values.
```python wrap theme={null}
{
"my_field_one": "A",
"my_field_two": "In summary...",
"my_field_three": ["In detail...", "In detail..."],
"my_field_four": 10
}
```
## `@task.llm`
The [`@task.llm` decorator and `LLMOperator`](https://airflow.apache.org/docs/apache-airflow-providers-common-ai/stable/operators/llm.html) make a single turn LLM call and return the model output.
```python wrap theme={null}
from datetime import timedelta
from airflow.sdk import task
@task.llm(
llm_conn_id="pydanticai_default", # required: your Airflow connection ID
model_id=":", # default: None (uses the model set in the connection)
# Optional parameters:
system_prompt="", # default: ""
output_type=str, # default: str; use a BaseModel subclass for structured output
agent_params=None, # default: None; passed to the PydanticAI Agent constructor
require_approval=False, # default: False; defer for human approve or reject in the UI
approval_timeout=timedelta(minutes=10), # default: None; max wait for approval
allow_modifications=False, # default: False; reviewer may edit text before approve
)
def my_llm_task(user_context: str) -> str:
return user_context
my_llm_task()
```
The string your `@task.llm` callable returns is the user prompt sent to the model. When you using `LLMOperator` directly, the `prompt` argument is required.
Optional parameters:
* `model_id`: Overrides the model in the connection `extra` (format: `:`).
* `system_prompt`: System instructions loaded before the user prompt.
* `output_type`: Return type for the run. Defaults to `str`. For structured JSON, set a subclass of [Pydantic `BaseModel`](https://pydantic.dev/docs/validation/latest/api/pydantic/base_model/), see [Output type](#output-type).
* `agent_params`: Extra keyword arguments for the PydanticAI `Agent` constructor (for example `model_settings`). See the [PydanticAI Agent parameters](https://pydantic.dev/docs/ai/api/pydantic-ai/agent/#parameters). Despite performing a single turn LLM call, you can still pass arguments to the agent constructor when using `@task.llm`.
* `require_approval`: When `True`, the task defers after generation until a human approves or rejects through the approval UI. Defaults to `False`.
* `approval_timeout`: Time limit to wait for a review when `require_approval` is `True`. When it is exceeded, the task fails. Defaults to `None`.
* `allow_modifications`: When `True` with approval enabled, the reviewer can edit the generated text before approval; that edited value becomes the task result. Defaults to `False`.
When `require_approval` is `True`, the task defers after the model returns its output. In the Airflow UI, open the task instance, select the **Required Action** tab, read the generated text, edit the output if needed, and then click **Approve** or **Reject**.
Human-in-the-loop review behaves differently for `@task.llm` and `@task.agent`. With `@task.llm`, the task [defers](/docs/learn/deferrable-operators) after generation until a human approves or rejects through the approval UI. There is only one approval step, and the reviewer can edit the model output before approval when `allow_modifications` is `True`. With `@task.agent` and [`enable_hitl_review=True`](#human-in-the-loop-review), the task doesn't defer; it keeps running after output generation and waits on the **HITL Review** tab for approval, rejection, or feedback that triggers another iteration. Use `max_hitl_iterations` to cap how many review rounds run.
You can have human-in-the-loop steps in Airflow Dags outside of agentic tasks by using standalone [HITL operators](/docs/learn/airflow-human-in-the-loop).
## `@task.llm_sql`
The [`@task.llm_sql` decorator and `LLMSQLQueryOperator`](https://airflow.apache.org/docs/apache-airflow-providers-common-ai/stable/operators/llm_sql.html) turn natural language into SQL. The operator can pull table metadata through a `DbApiHook` from `db_conn_id`, or you can supply a full schema string yourself. It generates SQL only; it doesn't run queries. Downstream tasks (for example [`SQLExecuteQueryOperator`](https://airflow.apache.org/docs/apache-airflow-providers-common-sql/stable/_api/airflow/providers/common/sql/operators/sql/index.html#airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator)) can execute the string returned in XCom.
```python wrap theme={null}
from datetime import timedelta
from airflow.sdk import task
@task.llm_sql(
llm_conn_id="pydanticai_default", # required: your Airflow connection ID
db_conn_id="postgres_default", # optional: connection that resolves to DbApiHook for accessing table schema
table_names=["orders", "customers"], # optional: tables to describe when using db_conn_id
schema_context=None, # optional: manual schema text; when set, skips db_conn_id access
validate_sql=True, # default True: validate generated SQL with sqlglot
dialect=None, # optional: for example "postgres"; inferred from the hook when None
model_id=":", # optional, default None
system_prompt="", # optional: appended to the built-in SQL safety instructions
agent_params=None, # optional: passed to the PydanticAI Agent constructor
require_approval=False,
approval_timeout=timedelta(minutes=10),
allow_modifications=False,
)
def my_nl_sql_task(question: str) -> str:
return question
my_nl_sql_task()
```
The string your callable returns is the natural language `prompt` that describes the query you want. When you use `LLMSQLQueryOperator` directly, pass that text with the `prompt` argument. You must set `llm_conn_id` to the [PydanticAI connection](#set-the-pydanticai-connection).
Additional parameters:
* `db_conn_id`: Airflow connection used to access metadata about your database. The hook must be a `DbApiHook`. Omit when you fully describe the schema with `schema_context`.
* `table_names`: List of table names to include when accessing the database with `db_conn_id`.
* `schema_context`: Free-form schema description. When you set this, the operator doesn't access the database for metadata.
* `validate_sql`: When `True` (default), generated SQL is checked with [sqlglot](https://sqlglot.com/sqlglot.html) before the task finishes.
* `allowed_sql_types`: Tuple of allowed statement roots (defaults to read-only shapes such as `Select`, `Union`, `Intersect`, and `Except` in sqlglot). Be careful when adding writing statements like `Insert`, `Update`, or `Delete`.
* `dialect`: sqlglot dialect name (for example `postgres`, `mysql`). When `None`, the operator tries to infer it from the database hook.
* `datasource_config`: Optional extra configuration for the data source side of generation (see the provider source for structure when you need it).
Parameters inherited from `LLMOperator` (`model_id`, `system_prompt`, `output_type`, `agent_params`, `require_approval`, `approval_timeout`, `allow_modifications`) behave like they do for `@task.llm`. When `require_approval` is `True` and `allow_modifications` is `True`, a reviewer can edit the generated SQL; the provider re-validates edited SQL against the `allowed_sql_types` rules before returning it.
## `@task.llm_branch`
The [`@task.llm_branch` decorator and `LLMBranchOperator`](https://airflow.apache.org/docs/apache-airflow-providers-common-ai/stable/operators/llm_branch.html) extend the `LLMOperator` with [branching](/docs/learn/airflow-branch-operator). At run time the operator reads downstream task IDs from the Dag, exposes them to the model as a constrained Enum through PydanticAI structured output, and skips tasks the model doesn't select. Note that you need to create a [dependency](/docs/learn/managing-dependencies) between the branch task and its downstream candidates.
```python expandable wrap theme={null}
from airflow.sdk import dag, task
@dag
def example_llm_branch():
@task.llm_branch(
llm_conn_id="pydanticai_default",
model_id=":", # optional, default None
allow_multiple_branches=False, # default False
system_prompt="Route support tickets to the right team.",
agent_params=None, # optional: passed to the PydanticAI Agent constructor
)
def route_ticket(message: str) -> str:
return f"Route this support ticket: {message}"
@task
def handle_billing():
return "Handling billing issue"
@task
def handle_auth():
return "Handling auth issue"
@task
def handle_general():
return "Handling general issue"
chain(
route_ticket("I was charged twice for my subscription."),
[
handle_billing(),
handle_auth(),
handle_general(),
]
)
example_llm_branch()
```
The string your callable returns is the user `prompt`. When you use `LLMBranchOperator` directly, pass that text with the `prompt` argument. Set `llm_conn_id` to the [PydanticAI connection](#set-the-pydanticai-connection). Additionally, you can set `allow_multiple_branches` to `True` to allow the model to return multiple downstream task IDs.
Parameters inherited from `LLMOperator` (`model_id`, `system_prompt`, `agent_params`, and the same optional human-in-the-loop approval fields as [`@task.llm`](#@task-llm)) behave the same way as for a plain LLM task.
## `@task.llm_file_analysis`
The [`@task.llm_file_analysis` decorator and `LLMFileAnalysisOperator`](https://airflow.apache.org/docs/apache-airflow-providers-common-ai/stable/operators/llm_file_analysis.html) analyze one file, a prefix, or a small set of files through a single LLM call. The operator resolves `file_path` with [`ObjectStoragePath`](https://airflow.apache.org/docs/apache-airflow/stable/core-concepts/objectstorage.html), normalizes supported formats into text context, and optionally attaches PNG, JPG, or PDF inputs as multimodal payloads when `multi_modal=True`. For object storage URIs you can embed the connection ID in the path (for example `s3://@my-bucket/prefix/`) or set `file_conn_id` separately.
```python wrap theme={null}
from datetime import timedelta
from airflow.sdk import , task
@task.llm_file_analysis(
llm_conn_id="pydanticai_default",
file_path="s3://aws_default@my-bucket/reports/quarterly.pdf",
file_conn_id=None, # optional: overrides connection embedded in file_path
multi_modal=True, # default False; set True for vision-capable models on images/PDF
max_files=20, # default 20; caps files resolved from a prefix
max_file_size_bytes=5 * 1024 * 1024, # default 5 MiB per file
max_total_size_bytes=20 * 1024 * 1024, # default 20 MiB total
max_text_chars=100_000, # default 100000; how much text to read from the files
sample_rows=10, # default 10; preview rows for CSV, Parquet, Avro
model_id=":",
system_prompt="",
output_type=str, # default str
agent_params=None,
require_approval=False,
approval_timeout=timedelta(minutes=10),
allow_modifications=False,
)
def review_quarterly_report() -> str:
return "Extract the key revenue, risk, and compliance findings from this report."
review_quarterly_report()
```
The string your callable returns is the analysis `prompt`. When you use `LLMFileAnalysisOperator` directly, pass that text with the `prompt` argument.
Additional parameters:
* `file_path`: File or prefix to analyze (local paths or object storage paths supported by Airflow object storage).
* `file_conn_id`: Optional Airflow connection for the storage backend when it isn't embedded in `file_path`.
* `multi_modal`: When `True`, PNG, JPG, JPEG, and PDF inputs can be sent as binary attachments; requires a multimodal-capable model.
* `max_files`, `max_file_size_bytes`, `max_total_size_bytes`, `max_text_chars`, `sample_rows`: Guardrails for listing, reading, and how much normalized text or row samples reach the model. Extra files under a large prefix are omitted and the operator notes that in the prompt context.
Optional dependencies: Parquet and Avro handling need the provider extras described in the [Airflow documentation](https://airflow.apache.org/docs/apache-airflow-providers-common-ai/stable/operators/llm_file_analysis.html).
## `@task.llm_schema_compare`
The [`@task.llm_schema_compare` decorator and `LLMSchemaCompareOperator`](https://airflow.apache.org/docs/apache-airflow-providers-common-ai/stable/operators/llm_schema_compare.html) read schema metadata from two or more systems and ask an LLM to flag differences. The task result is a dict `SchemaCompareResult` with fields such as `compatible`, `mismatches`, and `summary`. Each entry in `mismatches` is a `SchemaMismatch` with severity, column, types, and suggested actions.
You can supply databases in two ways:
* `db_conn_ids` together with `table_names`: shorthand to compare the same logical table across connections. Each connection must resolve to a `DbApiHook`.
* `data_sources`: a list of [`DataSourceConfig`](https://airflow.apache.org/docs/apache-airflow-providers-common-sql/stable/_api/airflow/providers/common/sql/config/index.html#airflow.providers.common.sql.config.DataSourceConfig) objects for more complex setups (for example object storage or catalog-backed sources combined with `db_conn_ids`).
```python wrap theme={null}
from datetime import timedelta
from airflow.sdk import task
@task.llm_schema_compare(
llm_conn_id="pydanticai_default",
db_conn_ids=["postgres_source", "snowflake_target"],
table_names=["customers"],
context_strategy="full", # "full" (default) or "basic"; full adds keys and indexes
model_id=":",
agent_params=None,
require_approval=False,
approval_timeout=timedelta(minutes=10),
allow_modifications=False,
)
def check_migration_readiness() -> str:
return (
"Compare schemas and flag breaking changes for nightly ETL. "
"Suggest migration actions where they help."
)
check_migration_readiness()
```
The string your callable returns is the comparison `prompt`. When you use `LLMSchemaCompareOperator` directly, pass that text with the `prompt` argument.
Additional parameters:
* `db_conn_ids` and `table_names`: Use together for the same table name across multiple database connections.
* `data_sources`: Optional list of `DataSourceConfig` for mixed database and object storage comparisons.
* `context_strategy`: `"basic"` sends column names and types only; `"full"` (default) adds primary keys, foreign keys, and indexes to the context sent to the model.
* `system_prompt`: The operator ships a default prompt that encodes cross-system type rules and severity levels. If you set `system_prompt` to any string, it replaces that default; import [`DEFAULT_SYSTEM_PROMPT`](https://airflow.apache.org/docs/apache-airflow-providers-common-ai/stable/_api/airflow/providers/common/ai/operators/llm_schema_compare/index.html#airflow.providers.common.ai.operators.llm_schema_compare.DEFAULT_SYSTEM_PROMPT) and concatenate when you want to extend rather than replace.
Parameters inherited from `LLMOperator` (`model_id`, `agent_params`, and optional approval fields) match [`@task.llm`](#@task-llm). Downstream tasks can branch on `comparison_result["compatible"]` or inspect `mismatches` as shown in the [Airflow documentation](https://airflow.apache.org/docs/apache-airflow-providers-common-ai/stable/operators/llm_schema_compare.html).