Building AI Data Pipelines: From Raw Data to Production-Ready AI Systems

from sqlalchemy import create_engine, text
from datetime import datetime, timedelta
import pandas as pd

class DatabaseExtractor:
    def __init__(self, connection_string):
        self.engine = create_engine(connection_string)

    def extract_incremental(self, table, timestamp_column, last_run_time):
        """Extract only new/updated records since last run."""
        query = text(f"""
            SELECT *
            FROM {table}
            WHERE {timestamp_column} > :last_run
            ORDER BY {timestamp_column}
        """)

        with self.engine.connect() as conn:
            df = pd.read_sql(query, conn, params={"last_run": last_run_time})

        self.logger.info(f"Extracted {len(df)} rows from {table}")
        return df.to_dict('records')

    def extract_full(self, query):
        """Extract full dataset."""
        with self.engine.connect() as conn:
            df = pd.read_sql(text(query), conn)
        return df.to_dict('records')

# Usage
extractor = DatabaseExtractor("postgresql://user:pass@localhost/db")

# Incremental: only get new data
last_run = datetime.now() - timedelta(days=1)
new_docs = extractor.extract_incremental(
    table="documents",
    timestamp_column="updated_at",
    last_run_time=last_run
)

import re
from html import unescape

class TextCleaner:
    @staticmethod
    def clean(text: str) -> str:
        """Clean and normalize text."""
        # Decode HTML entities
        text = unescape(text)

        # Remove excessive whitespace
        text = re.sub(r'\s+', ' ', text)

        # Remove special characters (keep punctuation)
        text = re.sub(r'[^\w\s.,!?-]', '', text)

        # Normalize quotes
        text = text.replace('"', '"').replace('"', '"')
        text = text.replace(''', "'").replace(''', "'")

        return text.strip()

    @staticmethod
    def remove_boilerplate(text: str) -> str:
        """Remove common boilerplate text."""
        # Remove email signatures
        text = re.sub(r'Best regards,.*?$', '', text, flags=re.DOTALL | re.IGNORECASE)

        # Remove confidentiality notices
        text = re.sub(r'This email is confidential.*?$', '', text, flags=re.DOTALL | re.IGNORECASE)

        return text.strip()

from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime, timedelta

# Define default arguments
default_args = {
    'owner': 'data_team',
    'depends_on_past': False,
    'email': ['alerts@example.com'],
    'email_on_failure': True,
    'email_on_retry': False,
    'retries': 2,
    'retry_delay': timedelta(minutes=5),
}

# Create DAG
dag = DAG(
    'ai_data_pipeline',
    default_args=default_args,
    description='Daily AI data pipeline for RAG system',
    schedule='0 2 * * *',  # Run at 2 AM daily
    start_date=datetime(2025, 1, 1),
    catchup=False,
    tags=['ai', 'embeddings', 'rag']
)

# Define tasks
def extract_data(**context):
    """Extract data from sources."""
    from pipelines.extractors import extract_all_sources

    data = extract_all_sources()

    # Store in XCom for next task
    context['task_instance'].xcom_push(key='raw_data', value=data)

    return f"Extracted {len(data)} documents"

def transform_data(**context):
    """Transform and clean data."""
    from pipelines.transformers import clean_and_chunk

    # Get data from previous task
    raw_data = context['task_instance'].xcom_pull(key='raw_data')

    chunks = clean_and_chunk(raw_data)

    context['task_instance'].xcom_push(key='chunks', value=chunks)

    return f"Created {len(chunks)} chunks"

def generate_embeddings(**context):
    """Generate embeddings."""
    from pipelines.embeddings import generate_batch_embeddings

    chunks = context['task_instance'].xcom_pull(key='chunks')

    embedded = generate_batch_embeddings(chunks)

    context['task_instance'].xcom_push(key='embedded', value=embedded)

    return f"Generated {len(embedded)} embeddings"

def load_to_vector_db(**context):
    """Load to vector database."""
    from pipelines.loaders import load_to_qdrant

    embedded = context['task_instance'].xcom_pull(key='embedded')

    load_to_qdrant(embedded)

    return f"Loaded {len(embedded)} vectors"

def validate_pipeline(**context):
    """Validate pipeline results."""
    from pipelines.validators import validate_data_quality

    results = validate_data_quality()

    if not results['passed']:
        raise ValueError(f"Validation failed: {results['errors']}")

    return "Validation passed"

# Create task instances
extract_task = PythonOperator(
    task_id='extract_data',
    python_callable=extract_data,
    dag=dag
)

transform_task = PythonOperator(
    task_id='transform_data',
    python_callable=transform_data,
    dag=dag
)

embed_task = PythonOperator(
    task_id='generate_embeddings',
    python_callable=generate_embeddings,
    dag=dag
)

load_task = PythonOperator(
    task_id='load_to_vector_db',
    python_callable=load_to_vector_db,
    dag=dag
)

validate_task = PythonOperator(
    task_id='validate_pipeline',
    python_callable=validate_pipeline,
    dag=dag
)

# Define task dependencies
extract_task >> transform_task >> embed_task >> load_task >> validate_task

from dataclasses import dataclass
from datetime import datetime

@dataclass
class PipelineMetrics:
    pipeline_name: str
    run_id: str
    start_time: datetime
    end_time: datetime
    records_processed: int
    records_failed: int
    bytes_processed: int
    embeddings_generated: int
    cost_usd: float
    errors: list

class MetricsCollector:
    def __init__(self):
        self.metrics = []

    def record_pipeline_run(self, metrics: PipelineMetrics):
        """Record pipeline metrics."""
        # Send to monitoring system (DataDog, CloudWatch, etc.)
        self._send_to_datadog(metrics)

        # Store in database for analysis
        self._store_in_db(metrics)

        # Check for anomalies
        self._check_anomalies(metrics)

    def _check_anomalies(self, metrics: PipelineMetrics):
        """Detect anomalies in pipeline performance."""
        # Get historical average
        avg_duration = self._get_avg_duration(metrics.pipeline_name)
        current_duration = (metrics.end_time - metrics.start_time).total_seconds()

        # Alert if 2x slower than average
        if current_duration > avg_duration * 2:
            self._send_alert(f"Pipeline running slow: {current_duration:.0f}s vs {avg_duration:.0f}s avg")

        # Alert if high failure rate
        failure_rate = metrics.records_failed / metrics.records_processed if metrics.records_processed > 0 else 0
        if failure_rate > 0.05:  # > 5% failure
            self._send_alert(f"High failure rate: {failure_rate:.1%}")