How to Design a Modern Data + AI System: Control, Data, and Decision Planes

from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.sensors.external_task import ExternalTaskSensor
from datetime import datetime, timedelta

default_args = {
    "owner": "data-platform",
    "retries": 3,
    "retry_delay": timedelta(minutes=5),
    "email_on_failure": True,
}

with DAG(
    "ai_recommendation_pipeline",
    default_args=default_args,
    schedule_interval="0 2 * * *",  # 2 AM daily
    start_date=datetime(2026, 1, 1),
    catchup=False,
    tags=["ai", "recommendations"],
) as dag:

    # Wait for upstream data to be ready
    wait_for_dbt = ExternalTaskSensor(
        task_id="wait_for_dbt_completion",
        external_dag_id="dbt_production",
        external_task_id="dbt_run",
        timeout=3600,
        poke_interval=60,
    )

    prepare_features = PythonOperator(
        task_id="prepare_features",
        python_callable=run_feature_prep,
        op_kwargs={"execution_date": "{{ ds }}"},
    )

    run_inference = PythonOperator(
        task_id="run_llm_inference",
        python_callable=batch_llm_inference,
        op_kwargs={"model": "claude-haiku-4-5-20251001", "batch_size": 100},
    )

    write_results = PythonOperator(
        task_id="write_results",
        python_callable=write_to_warehouse,
    )

    wait_for_dbt >> prepare_features >> run_inference >> write_results

-- 1. Staging: rename + cast raw source fields
-- models/staging/stg_user_events.sql
SELECT
    id                                    AS event_id,
    user_id,
    event_type,
    CAST(created_at AS TIMESTAMP)         AS event_at,
    CAST(properties AS JSON)              AS properties
FROM {{ source('raw', 'events') }}
WHERE event_type IS NOT NULL

-- 2. Intermediate: business logic lives here
-- models/intermediate/int_user_engagement_rolling.sql
SELECT
    user_id,
    COUNT(*) FILTER (WHERE event_at >= CURRENT_DATE - 7)   AS event_count_7d,
    COUNT(DISTINCT event_type)
      FILTER (WHERE event_at >= CURRENT_DATE - 7)          AS distinct_event_types_7d,
    DATE_DIFF('day', MAX(event_at), CURRENT_DATE)          AS days_since_last_event
FROM {{ ref('stg_user_events') }}
GROUP BY 1

-- 3. Mart: the clean, AI-ready feature table
-- models/marts/fct_user_features_for_ai.sql
{{ config(
    materialized='table',
    contract={"enforced": true}
) }}

SELECT
    u.user_id,
    u.plan_tier,
    u.days_since_signup,
    e.event_count_7d,
    e.days_since_last_event,
    -- Pre-computed context string — inspect this in any SQL client
    CONCAT(
        'User has been on the ', u.plan_tier, ' plan for ',
        u.days_since_signup, ' days. ',
        'Last active ', e.days_since_last_event, ' days ago. ',
        'Performed ', e.event_count_7d, ' events in the last week.'
    ) AS user_summary_for_llm
FROM {{ ref('dim_users') }} u
LEFT JOIN {{ ref('int_user_engagement_rolling') }} e USING (user_id)

# models/marts/fct_user_features_for_ai.yml
version: 2

models:
  - name: fct_user_features_for_ai
    contract:
      enforced: true
    columns:
      - name: user_id
        data_type: varchar
        constraints:
          - type: not_null
          - type: unique
      - name: user_summary_for_llm
        data_type: varchar
        constraints:
          - type: not_null
    tests:
      - not_null:
          column_name: user_summary_for_llm
      - dbt_expectations.expect_column_value_lengths_to_be_between:
          column_name: user_summary_for_llm
          min_value: 50
          max_value: 500  # guard against runaway context strings

import anthropic
import json
import concurrent.futures
from dataclasses import dataclass

client = anthropic.Anthropic()

@dataclass
class UserRecommendation:
    user_id: str
    recommendation_type: str
    recommendation_text: str
    confidence_score: float
    reasoning: str

SYSTEM_PROMPT = """You are a data-driven product advisor. Given a user engagement summary,
recommend the single most impactful action for this user.

Respond in JSON with this exact structure:
{
  "recommendation_type": "upgrade_prompt|re_engagement|feature_discovery|churn_risk",
  "recommendation_text": "specific, actionable text",
  "confidence_score": 0.0-1.0,
  "reasoning": "one sentence explaining why"
}"""

def get_recommendation(user_id: str, user_summary: str) -> UserRecommendation:
    message = client.messages.create(
        model="claude-haiku-4-5-20251001",
        max_tokens=256,
        system=SYSTEM_PROMPT,
        messages=[{"role": "user", "content": user_summary}],
    )
    result = json.loads(message.content[0].text)
    return UserRecommendation(
        user_id=user_id,
        recommendation_type=result["recommendation_type"],
        recommendation_text=result["recommendation_text"],
        confidence_score=result["confidence_score"],
        reasoning=result["reasoning"],
    )

def batch_llm_inference(execution_date: str, batch_size: int = 100) -> None:
    """
    Airflow callable: reads features from DWH, runs inference, writes back.
    Idempotent — safe to retry on partial failure.
    """
    rows = read_pending_rows(execution_date, batch_size)

    if not rows:
        print(f"No rows to process for {execution_date}")
        return

    results = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
        futures = {
            executor.submit(get_recommendation, row["user_id"], row["user_summary_for_llm"]): row["user_id"]
            for row in rows
        }
        for future in concurrent.futures.as_completed(futures):
            try:
                results.append(future.result())
            except Exception as e:
                print(f"Inference failed for user {futures[future]}: {e}")

    write_recommendations_to_staging(results, execution_date)

def select_model(task_complexity: str) -> str:
    """Choose the cheapest model that meets the quality bar."""
    high_stakes = {"churn_risk", "upgrade_prompt"}
    if task_complexity in high_stakes:
        return "claude-sonnet-4-6"      # more capable, ~4x cost
    return "claude-haiku-4-5-20251001"  # fast, cheap, sufficient for simple tasks

Control Plane (Airflow)
│
├── ExternalTaskSensor ──→ blocks until dbt tests pass
│
├── prepare_features ────→ reads from mart, adds execution_date column
│
├── run_llm_inference ───→ Decision Plane: reads features, calls LLM, stages output
│
└── write_results ───────→ writes staged results to final mart table
                              └── triggers: alerts, emails, dashboard refreshes