airflow.providers.common.ai.decorators.llm_sql¶

TaskFlow decorator for LLM SQL generation.

The user writes a function that returns the prompt. The decorator handles the LLM call, schema introspection, and safety validation. The decorated task’s XCom output is the generated SQL string.

Functions¶

llm_sql_task([python_callable])

Wrap a function that returns a natural language prompt into an LLM SQL task.

Module Contents¶

airflow.providers.common.ai.decorators.llm_sql.llm_sql_task(python_callable=None, **kwargs)[source]¶

Wrap a function that returns a natural language prompt into an LLM SQL task.

The function body constructs the prompt (can use Airflow context, XCom, etc.). The decorator handles: LLM connection, schema introspection, SQL generation, and safety validation.

Usage:

@task.llm_sql(
    llm_conn_id="openai_default",
    db_conn_id="postgres_default",
    table_names=["customers", "orders"],
)
def build_query(ds=None):
    return f"Find top 10 customers by revenue in {ds}"

Parameters:: python_callable (collections.abc.Callable | None) – Function to decorate.