Page History

Status

Background & Motivation

Background & Motivation

In In today's evolving data landscape, organizations face significant challenges:

...

This proposal leverages Airflow's strengths (production reliability, 1000+ integrations, governance) while adding AI-native capabilities for intelligent data operations.

Core Proposal: Specialized LLM Operators with Rich Context Integration

Based on Pydantic AI's capabilities, while it provides excellent agent framework, multi-model support, and structured outputs, it lacks:

...

An example view of how Operators look like below:

1. Specialized LLM Operators with Built-in Protection & Context

from airflow.providers.ai.operators import (

    LLMSQLQueryOperator,

    LLMSchemaCompareOperator, 

    LLMDataQualityOperator,

    LLMFileAnalysisOperator

)


from airflow.sdk import Asset

# Enhanced Asset with structured metadata

customer_postgres = Asset(

    name="customer_data_postgres", 

    uri="postgres://warehouse/public/customers",

    conn_id="postgres_warehouse",

    schema={

        "customer_id": "integer PRIMARY KEY",

        "full_name": "varchar(255) NOT NULL",

        "email_address": "varchar(255) UNIQUE", 

        "created_at": "timestamp DEFAULT now()",

        "total_revenue": "decimal(10,2)"

    },

    sensitivity="pii"
)


# SQL operator with automatic context injection and safety

# Option 1: Traditional operator approach


sql_analysis = LLMSQLQueryOperator(

    task_id="analyze_customer_segments",

    prompt="Find top 10 customers by revenue, include their email and signup date",

    data_sources=[customer_postgres],

    # Automatic context injection includes:

    # - Database type: PostgreSQL 

    # - Available tables and schemas from DbApiHook

    # - Column types and constraints

    # - Sample data (first 5 rows) -> explicit approvals with htil to read data if its PII or PCI

    # - Built-in SQL safety (blocks DROP, DELETE without WHERE, etc.)

)




# Option 2: Decorator approach with pre-processing

@task.llm_sql_query(data_sources=[customer_postgres])

def analyze_customer_segments_with_preprocessing():

    # Custom pre-processing logic

    current_date = datetime.now().strftime('%Y-%m-%d')

    business_hours = get_business_hours()

   
    # Dynamic prompt generation based on context

    prompt = f"""

    Find top 10 customers by revenue as of {current_date}.

    Include their email and signup date.

    Filter for customers active during business hours: {business_hours}

    """

    return {"prompt": prompt, "additional_context": {"analysis_date": current_date}}

2. Task Decorators for Dynamic AI Workflows

Each LLM operator has a corresponding decorator for more flexible, Pythonic workflows:

...

• Dynamic prompts based on runtime conditions
• Custom pre-processing logic before LLM calls
• Context enrichment from external systems
• Conditional logic for different scenarios
• Pythonic workflow familiar to Airflow users
• Full XCom integration for passing data between tasks

3. Rich Context Injection Examples

SQL Operator Context (via DbApiHook integration):

...

# Automatically injected context for PostgreSQL:


{

    "database_type": "postgresql",

    "version": "15.2",

    "available_tables": ["customers", "orders", "products"],

    "schema_info": {

        "customers": {

            "customer_id": {"type": "integer", "nullable": False, "primary_key": True},

            "email_address": {"type": "varchar(255)", "nullable": False, "unique": True},

            "total_revenue": {"type": "decimal(10,2)", "nullable": True}

        }

    },

    "sample_data": {

        "customers": [




            {"customer_id": 1, "email_address": "john@example.com", "total_revenue": 1250.00},

            {"customer_id": 2, "email_address": "jane@example.com", "total_revenue": 890.50}

        ]

    },

    "dialect_features": {

        "supports_window_functions": True,

        "supports_cte": True,

        "date_functions": ["DATE_TRUNC", "EXTRACT", "AGE"]

    }

}

File Operator Context (via S3Hook/GCSHook integration):

# Automatically injected context for S3 Parquet files:

{




    "storage_type": "s3",

    "file_format": "parquet", 

    "file_size_mb": 245,

    "estimated_rows": 1000000,

    "schema_info": {

        "id": "int64",

        "name": "string", 

        "email": "string",

        "signup_date": "timestamp[ns]"

    },

    "sample_data": [

        {"id": 1, "name": "John Doe", "email": "john@example.com"},

        {"id": 2, "name": "Jane Smith", "email": "jane@example.com"}

    ],

    "partitioning": ["year", "month"],

    "compression": "snappy"

}

3. Operator-Specific Safety & System Prompts

SQL Operator Built-in Protection:

...

    CONTEXT: File has {estimated_rows} rows, {file_size_mb}MB, partitioned by {partitioning}."""

4. Schema Comparison with Multi-Database Context

# Cross-system schema drift detection

...

    # - Snowflake: Column types, clustering keys, data sharing info

)

5. Unified Data Access with Apache DataFusion

A new AnalyticsOperator that provides unified access to multi-cloud data:

...

• Exceptional Performance: Query and perform aggregation operations on approximately 50 million records in under 10-15 seconds on single-node operations
• Multi-Cloud Native: Built-in support for S3, GCS, Azure Blob Storage without additional configuration
• Multi-Format Support: Native handling of Parquet, JSON, CSV, Avro, Iceberg, Delta Lake formats
• Cost Effective: Eliminates need for expensive distributed compute frameworks like Spark for many AI use cases
• SQL Dialect Unification: Provides unified SQL interface across different storage systems, crucial for AI agents
• Rust Performance: High-performance query engine optimized for the analytical workloads AI agents typically generate

6. Human-in-the-Loop Integration Options

We propose both embedded and separate HITL patterns:

Option A: Embedded HITLHITL 

(Out of scope this scenario moving this to New AIP)

quality_check = LLMDataQualityOperator(

    task_id="customer_quality_analysis",

    data_sources=[customer_s3],

    prompt="Generate data quality validation queries",

    require_approval=True,  # Built-in HITL

    approval_timeout=timedelta(hours=2)

)

...

# Generate queries

generate_queries = LLMDataQualityOperator(

    task_id="generate_quality_queries",

    data_sources=[customer_s3],

    prompt="Generate data quality validation queries",

    dry_run=True  # Don't execute, just generate

)

# Human approval step

approve_queries = ApprovalOperator(

    task_id="approve_queries",

    body="{{ ti.xcom_pull(task_ids='generate_quality_queries') }}",

    allow_modifications=True  # Users can edit generated queries

)


# Execute approved queries

execute_analysis = AnalyticsOperator(

    task_id="execute_quality_checks",

    query="{{ ti.xcom_pull(task_ids='approve_queries') }}",

    data_sources=[customer_s3]

)


generate_queries >> approve_queries >> execute_analysis

Complete Workflow Example

Here's a real-world scenario combining all components:

...

from datetime import datetime, timedelta




from airflow.sdk import DAG, Asset

from airflow.providers.ai.operators import (

    LLMSchemaCompareOperator,

    LLMDataQualityOperator, 

    AnalyticsOperator

)

from airflow.operators.approval import ApprovalOperator




# Define multi-cloud assets with structured metadata




customer_s3 = Asset(

    name="customer_feed_s3",

    uri="s3://data-lake/customer/",

    conn_id="aws_default",

    schema={"id": "int32", "name": "string", "email": "string"},

    sensitivity="pii",

    format="parquet"

)


customer_postgres = Asset(

    name="customer_master_postgres", 

    uri="postgres://warehouse/public/customers",

    conn_id="postgres_default",

    schema={"customer_id": "integer", "full_name": "varchar", "email_address": "varchar"},

    sensitivity="pii"

)


with DAG(

    "intelligent_data_validation",

    start_date=datetime(2024, 1, 1),

    schedule=timedelta(hours=6),

) as dag:




    # 1. Detect schema drift between S3 feed and PostgreSQL master




    schema_drift = LLMSchemaCompareOperator(

        task_id="detect_schema_drift",

        data_sources=[customer_s3, customer_postgres],

        prompt="Identify schema mismatches that would break data loading",

        output_format="structured_report"
    )   


    # 2. Generate data quality queries for new S3 data


    generate_quality_checks = LLMDataQualityOperator(

        task_id="generate_quality_queries",

        data_sources=[customer_s3],

        prompts=[

            "Generate summary statistics queries",

            "Check for duplicate email addresses",

        ],

        dry_run=True

    )


    # 3. Human approval for generated queries (with edit capability)


    approve_queries = ApprovalOperator(

        task_id="approve_quality_queries",

        body="{{ ti.xcom_pull(task_ids='generate_quality_queries') }}",

        allow_modifications=True,

        timeout=timedelta(hours=2)

    )


    # 4. Execute approved quality checks using DataFusion

    execute_quality_checks = AnalyticsOperator(

        task_id="run_quality_analysis",

        query="{{ ti.xcom_pull(task_ids='approve_quality_queries') }}",

        data_sources=[customer_s3],

        engine="datafusion",

        output_location="s3://results/quality-reports/"

    )

    schema_drift >> generate_quality_checks >> approve_queries >> execute_quality_checks

Evolution Path: From LLMOperator to AITask

While we start with SQL generation, the architecture supports broader AI workflows:

...

    output_model=ChurnAnalysisReport

)

Technical Implementation Details

Core Components

1. Specialized LLM Operators & Decorators

...

• Embedded approval within operators (require_approval=True)
• Separate ApprovalOperator with query modification capabilities
• Configurable timeout and escalation policies

Implementation Approach

Phase 1: Standalone Provider - No Core Changes Required

• Provider: apache-airflow-providers-ai (0.x releases for iteration)
• Core LLM operators with Pydantic AI integration
• Schema detection using existing connection types and hooks (no core modifications needed)
• DataFusion-based AnalyticsOperator for unified data access ( Users can extend this interface and provide their own implementation for querying if they don't prefer datafusion, eg use; duckdb)
• Asset integration using current Asset.extra for metadata (backward compatible)
• Validations: Query validation and safety analysis

...

Key Benefit: Can iterate rapidly on operator design and context injection without any core Airflow changes. All functionality works through existing connection and hook infrastructure.

Phase 2: Production Features

• Advanced HITL workflows with modification capabilities
• Performance optimization and caching
• Enhanced Asset metadata (if community feedback supports core changes)

Phase 3: Core Integration

• Only if Phase 1 proves valuable: Consider core Asset enhancements
• Provider-specific AI capabilities in existing providers
• Community tool ecosystem and standards

Why This Matters

For Airflow:

• Positions Airflow at the center of AI-powered data infrastructure
• Makes Airflow more data-aware: Enhanced Asset metadata creates richer data context and lineage
• Drives Asset adoption: Users get immediate value from defining Assets with schema and sensitivity metadata
• Leverages our unique strengths in production reliability and 1000+ provider ecosystem

...

For the Industry: Creates the missing bridge between AI capabilities and production data infrastructure while making data assets first-class citizens in AI workflows.

Conclusion:

Overall the proposal is to build LLM's based operators by leveraging the built in airflow connections and assets functionality 

Are there any downsides to this change?

None

Which users are affected by the change?

None

What defines this AIP as "done"?

By completing all the phases

Future Scope:

1. Embedded HITL as mentioned in the Option A example.
2. Progress reporting the task running in cycliness, eg: see comment from Unknown User (potiuk)