Page History

    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}}