Data Flywheel for AI: Turn Production Outputs Into Better

from dataclasses import dataclass, asdict
from datetime import datetime, timezone
from typing import Optional
import json
import asyncpg

@dataclass
class PredictionLog:
    """Log entry for model prediction."""
    
    prediction_id: str
    model_version: str
    input_text: str
    output_text: str
    
    # Model metadata
    confidence_score: Optional[float] = None
    latency_ms: float = 0.0
    
    # User feedback
    user_feedback: Optional[str] = None  # "thumbs_up", "thumbs_down", "edited"
    user_correction: Optional[str] = None
    
    # Context
    user_id: str = None
    session_id: str = None
    timestamp: datetime = None
    
    def __post_init__(self):
        if self.timestamp is None:
            self.timestamp = datetime.now(timezone.utc)

class ProductionLogger:
    """Log production predictions for data flywheel."""
    
    def __init__(self, db_pool: asyncpg.Pool):
        self.db = db_pool
    
    async def log_prediction(self, log: PredictionLog) -> None:
        """Log prediction to database."""
        
        await self.db.execute(
            """
            INSERT INTO prediction_logs (
                prediction_id, model_version, input_text, output_text,
                confidence_score, latency_ms, user_feedback, user_correction,
                user_id, session_id, timestamp
            ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
            """,
            log.prediction_id,
            log.model_version,
            log.input_text,
            log.output_text,
            log.confidence_score,
            log.latency_ms,
            log.user_feedback,
            log.user_correction,
            log.user_id,
            log.session_id,
            log.timestamp,
        )
    
    async def update_user_feedback(
        self,
        prediction_id: str,
        feedback: str,
        correction: Optional[str] = None,
    ) -> None:
        """Update prediction with user feedback."""
        
        await self.db.execute(
            """
            UPDATE prediction_logs
            SET user_feedback = $1,
                user_correction = $2
            WHERE prediction_id = $3
            """,
            feedback,
            correction,
            prediction_id,
        )

# Usage in production API
from fastapi import FastAPI, Body
from uuid import uuid4

app = FastAPI()
logger = ProductionLogger(db_pool)

@app.post("/classify")
async def classify_document(text: str, user_id: str):
    """Classify document and log prediction."""
    
    import time
    prediction_id = str(uuid4())
    
    # Run model
    start = time.perf_counter()
    prediction = await model.predict(text)
    latency_ms = (time.perf_counter() - start) * 1000
    
    # Log prediction
    await logger.log_prediction(PredictionLog(
        prediction_id=prediction_id,
        model_version="v2.1.0",
        input_text=text,
        output_text=prediction.label,
        confidence_score=prediction.confidence,
        latency_ms=latency_ms,
        user_id=user_id,
    ))
    
    return {
        "prediction_id": prediction_id,
        "label": prediction.label,
        "confidence": prediction.confidence,
    }

@app.post("/feedback")
async def submit_feedback(
    prediction_id: str = Body(...),
    feedback: str = Body(...),  # "thumbs_up", "thumbs_down"
    correction: Optional[str] = Body(None),
):
    """Submit user feedback on prediction."""
    
    await logger.update_user_feedback(prediction_id, feedback, correction)
    
    return {"status": "success"}

-- Schema for production logs
CREATE TABLE prediction_logs (
    prediction_id UUID PRIMARY KEY,
    model_version VARCHAR(50) NOT NULL,
    
    input_text TEXT NOT NULL,
    output_text TEXT NOT NULL,
    
    confidence_score FLOAT,
    latency_ms FLOAT NOT NULL,
    
    user_feedback VARCHAR(20),  -- thumbs_up, thumbs_down, edited
    user_correction TEXT,
    
    user_id VARCHAR(255),
    session_id VARCHAR(255),
    timestamp TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    
    -- For training data pipeline
    labeled BOOLEAN DEFAULT FALSE,
    correct_label TEXT,
    labeled_at TIMESTAMPTZ,
    labeled_by VARCHAR(255),
    
    -- Sampling metadata
    sampling_priority FLOAT DEFAULT 0.0,
    sampled_for_labeling BOOLEAN DEFAULT FALSE
);

CREATE INDEX idx_prediction_logs_feedback 
ON prediction_logs (timestamp DESC, user_feedback);

CREATE INDEX idx_prediction_logs_unlabeled 
ON prediction_logs (labeled, sampling_priority DESC)
WHERE labeled = FALSE;

# sample_selection.py
from typing import List
import numpy as np

class IntelligentSampler:
    """Select high-value samples for labeling."""
    
    def __init__(self, db_pool):
        self.db = db_pool
    
    async def select_samples_for_labeling(
        self,
        budget: int = 500,
    ) -> List[dict]:
        """Select samples to label based on value."""
        
        # Strategy 1: Low confidence predictions (uncertainty sampling)
        uncertain_samples = await self._get_uncertain_samples(budget // 3)
        
        # Strategy 2: Negative user feedback
        negative_feedback = await self._get_negative_feedback_samples(budget // 3)
        
        # Strategy 3: Recent distribution shifts
        distribution_shift = await self._get_distribution_shift_samples(budget // 3)
        
        # Combine and deduplicate
        all_samples = uncertain_samples + negative_feedback + distribution_shift
        unique_samples = self._deduplicate(all_samples)
        
        # Update database
        await self._mark_as_sampled([s["prediction_id"] for s in unique_samples])
        
        return unique_samples[:budget]
    
    async def _get_uncertain_samples(self, limit: int) -> List[dict]:
        """Get predictions with low confidence."""
        
        rows = await self.db.fetch(
            """
            SELECT 
                prediction_id,
                model_version,
                input_text,
                output_text,
                confidence_score,
                'uncertain' as reason
            FROM prediction_logs
            WHERE labeled = FALSE
              AND confidence_score IS NOT NULL
              AND confidence_score < 0.7  -- Low confidence threshold
            ORDER BY confidence_score ASC
            LIMIT $1
            """,
            limit,
        )
        
        return [dict(row) for row in rows]
    
    async def _get_negative_feedback_samples(self, limit: int) -> List[dict]:
        """Get predictions with negative user feedback."""
        
        rows = await self.db.fetch(
            """
            SELECT 
                prediction_id,
                model_version,
                input_text,
                output_text,
                confidence_score,
                user_feedback,
                user_correction,
                'negative_feedback' as reason
            FROM prediction_logs
            WHERE labeled = FALSE
              AND user_feedback IN ('thumbs_down', 'edited')
            ORDER BY timestamp DESC
            LIMIT $1
            """,
            limit,
        )
        
        return [dict(row) for row in rows]
    
    async def _get_distribution_shift_samples(self, limit: int) -> List[dict]:
        """Get samples from recent distribution shifts."""
        
        # Detect input distribution shift (simplified)
        rows = await self.db.fetch(
            """
            WITH recent_inputs AS (
                SELECT 
                    prediction_id,
                    model_version,
                    input_text,
                    output_text,
                    confidence_score,
                    timestamp
                FROM prediction_logs
                WHERE labeled = FALSE
                  AND timestamp >= NOW() - INTERVAL '7 days'
            ),
            input_stats AS (
                SELECT 
                    prediction_id,
                    LENGTH(input_text) as input_length,
                    -- More sophisticated feature extraction in production
                    AVG(LENGTH(input_text)) OVER () as avg_length,
                    STDDEV(LENGTH(input_text)) OVER () as stddev_length
                FROM recent_inputs
            )
            SELECT 
                ri.prediction_id,
                ri.model_version,
                ri.input_text,
                ri.output_text,
                ri.confidence_score,
                'distribution_shift' as reason
            FROM recent_inputs ri
            JOIN input_stats ist ON ri.prediction_id = ist.prediction_id
            WHERE ABS(ist.input_length - ist.avg_length) > 2 * ist.stddev_length
            ORDER BY ri.timestamp DESC
            LIMIT $1
            """,
            limit,
        )
        
        return [dict(row) for row in rows]
    
    def _deduplicate(self, samples: List[dict]) -> List[dict]:
        """Remove duplicate samples."""
        
        seen_ids = set()
        unique = []
        
        for sample in samples:
            if sample["prediction_id"] not in seen_ids:
                seen_ids.add(sample["prediction_id"])
                unique.append(sample)
        
        return unique
    
    async def _mark_as_sampled(self, prediction_ids: List[str]) -> None:
        """Mark samples as selected for labeling."""
        
        await self.db.execute(
            """
            UPDATE prediction_logs
            SET sampled_for_labeling = TRUE
            WHERE prediction_id = ANY($1)
            """,
            prediction_ids,
        )

# Usage
sampler = IntelligentSampler(db_pool)

# Weekly: Select 500 samples for labeling
samples = await sampler.select_samples_for_labeling(budget=500)
print(f"Selected {len(samples)} samples for labeling")

# Distribution by reason
from collections import Counter
reasons = Counter(s["reason"] for s in samples)
print(f"Sampling distribution: {dict(reasons)}")
# Output: {'uncertain': 167, 'negative_feedback': 167, 'distribution_shift': 166}

# labeling_pipeline.py
from dataclasses import dataclass
from datetime import datetime
from typing import Optional

@dataclass
class LabelingTask:
    """Task for human labeler."""
    
    task_id: str
    prediction_id: str
    input_text: str
    model_output: str
    model_confidence: float
    
    # Context
    reason_selected: str
    user_feedback: Optional[str] = None
    user_correction: Optional[str] = None
    
    # Labeling metadata
    assigned_to: Optional[str] = None
    labeled_at: Optional[datetime] = None
    correct_label: Optional[str] = None
    labeler_confidence: Optional[str] = None  # "high", "medium", "low"

class LabelingPipeline:
    """Manage human labeling workflow."""
    
    def __init__(self, db_pool):
        self.db = db_pool
    
    async def create_labeling_batch(
        self,
        samples: List[dict],
        batch_name: str,
    ) -> str:
        """Create batch of labeling tasks."""
        
        import uuid
        batch_id = str(uuid.uuid4())
        
        tasks = []
        for sample in samples:
            task = LabelingTask(
                task_id=str(uuid.uuid4()),
                prediction_id=sample["prediction_id"],
                input_text=sample["input_text"],
                model_output=sample["output_text"],
                model_confidence=sample.get("confidence_score", 0.0),
                reason_selected=sample["reason"],
                user_feedback=sample.get("user_feedback"),
                user_correction=sample.get("user_correction"),
            )
            tasks.append(task)
        
        # Insert tasks
        await self.db.executemany(
            """
            INSERT INTO labeling_tasks (
                task_id, batch_id, batch_name, prediction_id,
                input_text, model_output, model_confidence,
                reason_selected, user_feedback, user_correction,
                status
            ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, 'pending')
            """,
            [
                (
                    t.task_id, batch_id, batch_name, t.prediction_id,
                    t.input_text, t.model_output, t.model_confidence,
                    t.reason_selected, t.user_feedback, t.user_correction,
                )
                for t in tasks
            ],
        )
        
        print(f"✓ Created batch {batch_name} with {len(tasks)} tasks")
        return batch_id
    
    async def assign_task(self, labeler_id: str) -> Optional[LabelingTask]:
        """Assign next task to labeler."""
        
        # Get next pending task
        row = await self.db.fetchrow(
            """
            UPDATE labeling_tasks
            SET assigned_to = $1,
                status = 'in_progress',
                assigned_at = NOW()
            WHERE task_id = (
                SELECT task_id
                FROM labeling_tasks
                WHERE status = 'pending'
                ORDER BY created_at ASC
                LIMIT 1
                FOR UPDATE SKIP LOCKED
            )
            RETURNING *
            """,
            labeler_id,
        )
        
        if not row:
            return None
        
        return LabelingTask(**dict(row))
    
    async def submit_label(
        self,
        task_id: str,
        correct_label: str,
        labeler_confidence: str,
        labeler_id: str,
    ) -> None:
        """Submit labeled result."""
        
        # Update task
        await self.db.execute(
            """
            UPDATE labeling_tasks
            SET correct_label = $1,
                labeler_confidence = $2,
                status = 'completed',
                labeled_at = NOW()
            WHERE task_id = $3
              AND assigned_to = $4
            """,
            correct_label,
            labeler_confidence,
            task_id,
            labeler_id,
        )
        
        # Update original prediction log
        await self.db.execute(
            """
            UPDATE prediction_logs
            SET labeled = TRUE,
                correct_label = $1,
                labeled_at = NOW(),
                labeled_by = $2
            WHERE prediction_id = (
                SELECT prediction_id FROM labeling_tasks WHERE task_id = $3
            )
            """,
            correct_label,
            labeler_id,
            task_id,
        )
    
    async def get_labeling_stats(self, batch_id: str) -> dict:
        """Get statistics for labeling batch."""
        
        row = await self.db.fetchrow(
            """
            SELECT 
                COUNT(*) as total_tasks,
                COUNT(*) FILTER (WHERE status = 'completed') as completed,
                COUNT(*) FILTER (WHERE status = 'in_progress') as in_progress,
                COUNT(*) FILTER (WHERE status = 'pending') as pending
            FROM labeling_tasks
            WHERE batch_id = $1
            """,
            batch_id,
        )
        
        stats = dict(row)
        stats["completion_rate"] = stats["completed"] / stats["total_tasks"] if stats["total_tasks"] > 0 else 0.0
        
        return stats

# Labeling UI endpoint (FastAPI)
@app.get("/labeling/next-task")
async def get_next_task(labeler_id: str):
    """Get next labeling task for labeler."""
    
    pipeline = LabelingPipeline(db_pool)
    task = await pipeline.assign_task(labeler_id)
    
    if not task:
        return {"status": "no_tasks"}
    
    return {
        "task_id": task.task_id,
        "input": task.input_text,
        "model_output": task.model_output,
        "confidence": task.model_confidence,
        "user_feedback": task.user_feedback,
        "user_correction": task.user_correction,
    }

@app.post("/labeling/submit")
async def submit_label(
    task_id: str = Body(...),
    correct_label: str = Body(...),
    confidence: str = Body(...),  # "high", "medium", "low"
    labeler_id: str = Body(...),
):
    """Submit labeled result."""
    
    pipeline = LabelingPipeline(db_pool)
    await pipeline.submit_label(task_id, correct_label, confidence, labeler_id)
    
    return {"status": "success"}

# retraining_pipeline.py
from datetime import datetime, timedelta
import asyncio

class RetrainingPipeline:
    """Automated model retraining pipeline."""
    
    def __init__(self, db_pool):
        self.db = db_pool
    
    async def should_trigger_retraining(self) -> tuple[bool, str]:
        """Check if should trigger retraining."""
        
        # Get stats on new labeled data
        stats = await self.db.fetchrow(
            """
            SELECT 
                COUNT(*) as new_labels,
                MIN(labeled_at) as oldest_label,
                MAX(labeled_at) as newest_label
            FROM prediction_logs
            WHERE labeled = TRUE
              AND labeled_at > (
                  SELECT COALESCE(MAX(training_data_cutoff), '1970-01-01'::timestamptz)
                  FROM model_versions
                  WHERE status = 'production'
              )
            """
        )
        
        new_labels = stats["new_labels"]
        
        # Trigger if:
        # 1. At least 1000 new labels
        if new_labels >= 1000:
            return True, f"{new_labels} new labels available"
        
        # 2. OR at least 500 labels and 7 days since last training
        if new_labels >= 500:
            oldest = stats["oldest_label"]
            if oldest and (datetime.now(timezone.utc) - oldest).days >= 7:
                return True, f"{new_labels} new labels, 7+ days old"
        
        return False, f"Only {new_labels} new labels"
    
    async def prepare_training_dataset(self) -> str:
        """Export training dataset with new labels."""
        
        import pandas as pd
        
        # Get all labeled data
        rows = await self.db.fetch(
            """
            SELECT 
                input_text,
                correct_label as label,
                model_version,
                labeled_at
            FROM prediction_logs
            WHERE labeled = TRUE
            ORDER BY labeled_at ASC
            """
        )
        
        df = pd.DataFrame([dict(r) for r in rows])
        
        # Save to training data location
        output_path = f"s3://training-data/dataset-{datetime.now().isoformat()}.parquet"
        df.to_parquet(output_path)
        
        print(f"✓ Exported {len(df)} training samples to {output_path}")
        return output_path
    
    async def trigger_training_job(self, dataset_path: str) -> str:
        """Trigger model training job."""
        
        import boto3
        
        # Trigger SageMaker training job
        sagemaker = boto3.client("sagemaker")
        
        training_job_name = f"retraining-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
        
        response = sagemaker.create_training_job(
            TrainingJobName=training_job_name,
            RoleArn="arn:aws:iam::ACCOUNT:role/SageMakerRole",
            AlgorithmSpecification={
                "TrainingImage": "your-training-image:latest",
                "TrainingInputMode": "File",
            },
            InputDataConfig=[
                {
                    "ChannelName": "training",
                    "DataSource": {
                        "S3DataSource": {
                            "S3Uri": dataset_path,
                            "S3DataType": "S3Prefix",
                        }
                    },
                }
            ],
            OutputDataConfig={
                "S3OutputPath": "s3://model-artifacts/",
            },
            ResourceConfig={
                "InstanceType": "ml.p3.2xlarge",
                "InstanceCount": 1,
                "VolumeSizeInGB": 50,
            },
            StoppingCondition={
                "MaxRuntimeInSeconds": 3600,
            },
        )
        
        print(f"✓ Started training job: {training_job_name}")
        return training_job_name
    
    async def run_retraining_cycle(self) -> None:
        """Run full retraining cycle."""
        
        # Check if should retrain
        should_retrain, reason = await self.should_trigger_retraining()
        
        if not should_retrain:
            print(f"⏸️  Skipping retraining: {reason}")
            return
        
        print(f"▶️  Triggering retraining: {reason}")
        
        # Prepare dataset
        dataset_path = await self.prepare_training_dataset()
        
        # Trigger training
        training_job_name = await self.trigger_training_job(dataset_path)
        
        # Wait for training to complete (in production, use async callback)
        # await self.wait_for_training_completion(training_job_name)
        
        # After training completes:
        # 1. Evaluate new model
        # 2. If improvement, deploy to staging
        # 3. Run A/B test
        # 4. Promote to production
        
        print(f"✓ Retraining cycle complete: {training_job_name}")

# Scheduled retraining (runs weekly)
async def scheduled_retraining():
    """Run scheduled retraining check."""
    
    pipeline = RetrainingPipeline(db_pool)
    await pipeline.run_retraining_cycle()

# Cron job: Run every Monday at 2 AM
# 0 2 * * 1 python -c "import asyncio; asyncio.run(scheduled_retraining())"

# quality_monitoring.py
import pandas as pd
from datetime import datetime, timedelta

class QualityMonitor:
    """Monitor model quality improvement from flywheel."""
    
    def __init__(self, db_pool):
        self.db = db_pool
    
    async def get_quality_over_time(self, days: int = 90) -> pd.DataFrame:
        """Get quality metrics over time."""
        
        rows = await self.db.fetch(
            """
            WITH weekly_stats AS (
                SELECT 
                    DATE_TRUNC('week', timestamp) as week,
                    model_version,
                    COUNT(*) as total_predictions,
                    AVG(CASE WHEN confidence_score > 0.8 THEN 1 ELSE 0 END) as high_confidence_rate,
                    COUNT(*) FILTER (WHERE user_feedback = 'thumbs_up') as positive_feedback,
                    COUNT(*) FILTER (WHERE user_feedback = 'thumbs_down') as negative_feedback,
                    AVG(latency_ms) as avg_latency_ms
                FROM prediction_logs
                WHERE timestamp >= NOW() - INTERVAL '$1 days'
                GROUP BY week, model_version
                ORDER BY week ASC
            )
            SELECT 
                week,
                model_version,
                total_predictions,
                high_confidence_rate,
                CASE 
                    WHEN (positive_feedback + negative_feedback) > 0
                    THEN positive_feedback::float / (positive_feedback + negative_feedback)
                    ELSE NULL
                END as satisfaction_rate,
                avg_latency_ms
            FROM weekly_stats
            """,
            days,
        )
        
        return pd.DataFrame([dict(r) for r in rows])
    
    async def calculate_flywheel_impact(self) -> dict:
        """Calculate impact of data flywheel on model quality."""
        
        # Get baseline (first model version)
        baseline = await self.db.fetchrow(
            """
            SELECT 
                AVG(CASE WHEN confidence_score > 0.8 THEN 1 ELSE 0 END) as high_confidence_rate
            FROM prediction_logs
            WHERE model_version = (SELECT MIN(model_version) FROM prediction_logs)
            LIMIT 10000
            """
        )
        
        # Get current (latest model version)
        current = await self.db.fetchrow(
            """
            SELECT 
                AVG(CASE WHEN confidence_score > 0.8 THEN 1 ELSE 0 END) as high_confidence_rate
            FROM prediction_logs
            WHERE model_version = (SELECT MAX(model_version) FROM prediction_logs)
            LIMIT 10000
            """
        )
        
        improvement = (current["high_confidence_rate"] - baseline["high_confidence_rate"]) / baseline["high_confidence_rate"]
        
        return {
            "baseline_quality": baseline["high_confidence_rate"],
            "current_quality": current["high_confidence_rate"],
            "improvement_percentage": improvement * 100,
        }

# Prometheus metrics
from prometheus_client import Gauge

flywheel_labels_collected = Gauge(
    "flywheel_labels_total",
    "Total labels collected from production",
)

flywheel_model_accuracy = Gauge(
    "flywheel_model_accuracy",
    "Model accuracy over time",
    ["model_version"],
)

flywheel_training_iterations = Gauge(
    "flywheel_training_iterations_total",
    "Number of retraining iterations",
)

┌─────────────────────────────────────────────────────────────┐
│                    Data Flywheel Cycle                       │
│                                                              │
│  1. Production Usage                                         │
│     ↓                                                        │
│  2. Log Predictions ──→ Database                            │
│     ↓                                                        │
│  3. Intelligent Sampling ──→ Select High-Value Samples      │
│     ↓                                                        │
│  4. Human Labeling ──→ Correct Labels                       │
│     ↓                                                        │
│  5. Automated Retraining ──→ New Model Version              │
│     ↓                                                        │
│  6. Deployment ──→ Replace Production Model                 │
│     │                                                        │
│     └────────────────────────────┐                          │
│                                  ↓                          │
│                          Better Model                        │
│                                  ↓                          │
│                    Attracts More Usage (repeat)             │
└─────────────────────────────────────────────────────────────┘