DUE TO SPAM, SIGN-UP IS DISABLED. Goto Selfserve wiki signup and request an account.
Status
| State | |
| Accepted | |
| Discussion Thread | https://lists.apache.org/thread/dgpgvoszh52vxxszmg65wmcgxnj9zwby |
| Vote Thread | https://lists.apache.org/thread/vhy4bynwmqvrxw0cqgwkvmx6vhlzbqnr |
| Vote Result Thread | https://lists.apache.org/thread/7g2gw53lf9yf8mmt9g6mnf29rj2yzgpr |
| Progress Tracking (PR/GitHub Project/Issue Label) | https://github.com/orgs/apache/projects/586/views/1 |
| Date Created |
|
| Version Released | |
| Authors |
Background & Motivation
Background & Motivation
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
| Code Block | ||||
|---|---|---|---|---|
| ||||
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:
| Code Block | ||||
|---|---|---|---|---|
| ||||
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
} # 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):
| Code Block | ||||
|---|---|---|---|---|
| ||||
# 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):
...
| Code Block | ||||
|---|---|---|---|---|
| ||||
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
...
File Operator Built-in Protection:
| Code Block | ||||
|---|---|---|---|---|
| ||||
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
...
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)
| Code Block | ||||
|---|---|---|---|---|
| ||||
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
| Code Block |
|---|
# Generate queries
generate_queries = LLMDataQualityOperator(
| ||||
# 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 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:
| Code Block | ||||
|---|---|---|---|---|
| ||||
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", |
...
"Validate data against business rules"
], |
...
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
- Specialized LLM Operators & Decorators
...
- 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
...
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 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:
- Embedded HITL as mentioned in the Option A example.
- Progress reporting the task running in cycliness, eg: see comment from Unknown User (potiuk)