Page History

Status

Background & Motivation

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

1. Schema Drift Detection: Breaking changes between upstream and downstream systems consume significant engineering time
2. Multi-Cloud Complexity: Data scattered across AWS, GCP, Azure with different formats (Iceberg, Delta Lake, Parquet, PostgreSQL, etc.)
3. Data Quality at Scale: Context-aware validation that understands business rules, not just syntax
4. AI Context Requirements: Providing accurate data context to AI/ML models and agents for reliable insights

Real-World Pain Points:

• Schema mismatches between producers and consumers causing pipeline failures
• Manual data quality checks that don't scale with data volume or complexity
• Fragmented tooling for accessing data across cloud providers and storage formats
• Lack of intelligent validation that understands business context

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:

1. Airflow Production Integration: No native connection management, XCom, DAG context, or retry logic
2. Context-Aware Safety: No built-in protection against dangerous SQL operations or file modifications
3. Automatic Context Injection: No integration with Airflow's 500+ hooks for schema/metadata discovery
4. Workflow-Native Features: No approval workflows, asset integration, or Airflow monitoring

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:

from airflow.providers.ai.decorators import task

# Schema comparison with custom logic

@task.llm_schema_compare(data_sources=[s3_asset, postgres_asset])

def intelligent_schema_validation():

    # Pre-processing: check business calendar

    is_migration_window = check_migration_window()

    

    if is_migration_window:

        prompt = "Compare schemas and generate migration plan for scheduled maintenance window"

    else:

        prompt = "Compare schemas and flag breaking changes - no migrations allowed"

    return {

        "prompt": prompt,

        "migration_allowed": is_migration_window,

        "additional_context": {"maintenance_window": is_migration_window}

    }

# Data quality with dynamic rules

@task.llm_data_quality(data_sources=[customer_asset])

def adaptive_quality_checks():

    # Pre-processing: get current business rules

    current_rules = fetch_business_rules()

    seasonal_adjustments = get_seasonal_data_patterns()

    prompt = f"""

    Validate customer data against current business rules:

    {current_rules}

    Apply seasonal adjustments for data volume expectations:

    {seasonal_adjustments}

    Generate appropriate validation queries.

    """

    return {

        "prompt": prompt,

        "business_rules": current_rules,

        "seasonal_context": seasonal_adjustments

    }


# File analysis with preprocessing

@task.llm_file_analysis(data_sources=[log_files_asset])

def analyze_logs_with_context():

    # Pre-processing: get system context

    recent_deployments = get_recent_deployments()

    system_alerts = get_active_alerts()

    prompt = f"""

    Analyze log files for anomalies, considering:

    - Recent deployments: {recent_deployments}

    - Active system alerts: {system_alerts}

    Focus on correlation between deployment events and error patterns.

    """

    return {

        "prompt": prompt,

        "deployment_context": recent_deployments,

        "alert_context": system_alerts

    }

Benefits of Decorator Approach:

• 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:

class LLMSQLQueryOperator(BaseOperator):

    # Built-in dangerous operation blocking


    BLOCKED_KEYWORDS = ["DROP", "TRUNCATE", "DELETE FROM", "ALTER TABLE", "GRANT", "REVOKE"]


    DEFAULT_SYSTEM_PROMPT = """You are a SQL expert integrated with {database_type}.

    SAFETY RULES:

    - NEVER generate DROP, TRUNCATE, DELETE without WHERE, or ALTER statements

    - Always use proper JOIN syntax for this database type

    - Respect column types and constraints provided in schema

    - Use database-specific functions when available

    CONTEXT: You have access to schema info, sample data, and dialect features.

    Generate optimized, safe queries that work with {database_type} version {version}."""

File Operator Built-in Protection:

class LLMFileAnalysisOperator(BaseOperator):

    ALLOWED_OPERATIONS = ["read", "analyze", "summarize", "validate"]


    DEFAULT_SYSTEM_PROMPT = """You are a file analysis expert for {storage_type} {file_format} files.

    SAFETY RULES:

    - ONLY read and analyze files, NEVER modify or delete

    - Respect file size limits and memory constraints

    - Generate efficient queries for large datasets

    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

schema_drift = LLMSchemaCompareOperator(

    task_id="detect_schema_drift",

    data_sources=[customer_s3, customer_postgres, customer_snowflake],

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

    # Automatically gets context from each system:

    # - S3: Parquet schema, partitioning, file stats

    # - PostgreSQL: Table schema, constraints, indexes  

    # - 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:

from airflow.providers.ai.operators import AnalyticsOperator

# Execute queries across different storage systems uniformly

analytics_task = AnalyticsOperator(

    task_id="cross_cloud_analysis",

    query="{{ ti.xcom_pull(task_ids='quality_check') }}",  # LLM-generated query

    data_sources=[customer_s3, orders_gcs, inventory_azure],  # Multi-cloud

    engine="datafusion",  # High-performance query engine

    output_format="parquet",

    output_location="s3://results/analysis/"

)

Why Apache DataFusion for AI Workloads:

Based on Wren AI's experience, DataFusion provides significant advantages for AI-driven data analysis:

• 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 HITL 

(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)

)

Option B: Separate HITL Steps

# 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:

# Future: General AI task abstraction

process_customer_churn = AITask(

    task_id="analyze_churn_patterns",

    objective="Identify customers at risk of churning and recommend actions",

    resources={

        "customer_data": snowflake_conn,

        "email_system": sendgrid_conn,

        "ml_platform": sagemaker_conn

    },

    constraints={"budget": "$50", "privacy": "pii_protected"},

    output_model=ChurnAnalysisReport

)

Technical Implementation Details

Core Components

1. Specialized LLM Operators & Decorators

• LLMSQLQueryOperator / @task.llm_sql_query: Natural language to SQL generation

• LLMSchemaCompareOperator / @task.llm_schema_compare: Schema drift detection between systems
• LLMDataQualityOperator / @task.llm_data_quality: Context-aware data validation query generation
• LLMFileAnalysisOperator / @task.llm_file_analysis: File content analysis and processing

Decorator Benefits:

• Dynamic prompt generation based on runtime conditions
• Custom pre-processing logic before LLM calls
• Context enrichment from external systems (business rules, calendars, alerts)
• Conditional logic for different operational scenarios

Enhanced Asset System

# Proposed Asset structure evolution

class Asset:

    name: str

    uri: str

    conn_id: Optional[str] = None  # Direct connection reference

    schema: Optional[Dict[str, str]] = None  # Structured schema

    sensitivity: Optional[str] = None  # pii, sensitive, public

    format: Optional[str] = None  # parquet, json, csv, etc.

    statistics: Optional[Dict] = None  # Row counts, update times

    extra: Dict[str, Any] = field(default_factory=dict)  # Backward compatibility

AnalyticsOperator with DataFusion

• Unified interface for multi-cloud data access across S3, GCS, Azure Blob Storage
• Native handling of Parquet, Iceberg, Delta Lake, JSON, CSV, Avro formats
• High-performance processing: 50M+ records in 10-15 seconds on single node(This is from my experiments) 
• Cost-effective alternative to Spark for AI-generated analytical queries
• SQL dialect unification - write once, run anywhere
• DataFusion table provider supports integrating existing databases like sqlite, postgres. A good option here is using datafusion-table-providers. https://github.com/datafusion-contrib/datafusion-table-providers 

Flexible HITL Integration

• 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 Users:

• Democratizes data access - business users describe intent in natural language
• Data engineers maintain governance, reliability, and safety through built-in protections
• Rich data context: Assets become intelligent, carrying schema, business rules, and sensitivity information

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)