Orchestrate AI tasks with Apache Airflow® and the Common AI provider

apache-airflow-providers-common-ai==<version>

apache-airflow-providers-standard==<version>
apache-airflow-providers-common-sql==<version>

AIRFLOW_CONN_PYDANTICAI_DEFAULT='{
    "conn_type": "pydanticai",
    "host": "<your_host>",
    "password": "<your_api_key>",
    "extra": {
        "model": "<your_provider>:<your_model>"
    }
}'

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="<your_provider>:<your_model>",  # default: None (uses the model set in the connection)
    # Optional parameters:
    system_prompt="<your_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!")

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.
    )
],

from airflow.providers.common.ai.toolsets.mcp import MCPToolset

toolsets=[MCPToolset(mcp_conn_id="mcp_default")],

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: ""
    )
],

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!")

AIRFLOW__COMMON_AI__DURABLE_CACHE_PATH=/path/to/cache

AIRFLOW__COMMON_AI__DURABLE_CACHE_PATH=s3://<your_conn_id>@my-bucket/some/prefix/

AIRFLOW_CONN_<your_conn_id>='{
    "conn_type": "aws",
    "login": "<your-aws-access-key>",
    "password": "<your-aws-secret-key>",
    "extra": {
        "region_name": "<your-region>"
    }
}'

from airflow.sdk import task

@task.agent(
    durable=True,
)

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

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..."
        )
    )

{
    "my_field_one": "A",
    "my_field_two": "In summary...",
    "my_field_three": ["In detail...", "In detail..."],
    "my_field_four": 10
}

from datetime import timedelta
from airflow.sdk import task

@task.llm(
    llm_conn_id="pydanticai_default",  # required: your Airflow connection ID
    model_id="<your_provider>:<your_model>",  # 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()

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="<your_provider>:<your_model>",  # 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()

from airflow.sdk import dag, task


@dag
def example_llm_branch():
    @task.llm_branch(
        llm_conn_id="pydanticai_default",
        model_id="<your_provider>:<your_model>",  # 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()

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="<your_provider>:<your_model>",
    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()

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="<your_provider>:<your_model>",
    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()