A/B Testing AI Agents

You've built a new prompt that seems better in evals. But will it actually improve the customer experience in production? A/B testing lets you find out safely.

However, testing AI agents is trickier than testing button colors. Let's learn how to do it right.

Why Agent A/B Testing is Different

A/B Testing Framework

from dataclasses import dataclass
from typing import Dict, Any, Optional, List
from enum import Enum
import hashlib
import random

class ExperimentStatus(Enum):
    DRAFT = "draft"
    RUNNING = "running"
    PAUSED = "paused"
    COMPLETED = "completed"

@dataclass
class AgentVariant:
    """A single variant in an A/B test"""
    id: str
    name: str
    description: str
    config: Dict[str, Any]  # Agent configuration
    traffic_percentage: float

@dataclass
class Experiment:
    """An A/B test experiment"""
    id: str
    name: str
    hypothesis: str
    primary_metric: str
    secondary_metrics: List[str]
    variants: List[AgentVariant]
    status: ExperimentStatus
    start_date: Optional[str] = None
    end_date: Optional[str] = None
    min_sample_size: int = 1000
    min_duration_days: int = 7


class ExperimentManager:
    """Manage A/B test experiments for agents"""

    def __init__(self):
        self.experiments: Dict[str, Experiment] = {}
        self.assignments: Dict[str, str] = {}  # user_id -> variant_id

    def create_experiment(
        self,
        name: str,
        hypothesis: str,
        control_config: Dict,
        treatment_config: Dict,
        traffic_split: float = 0.5,
        primary_metric: str = "resolution_rate"
    ) -> Experiment:
        """Create a new A/B test"""

        experiment = Experiment(
            id=self._generate_id(name),
            name=name,
            hypothesis=hypothesis,
            primary_metric=primary_metric,
            secondary_metrics=[
                "user_satisfaction",
                "response_time",
                "escalation_rate",
                "cost_per_conversation"
            ],
            variants=[
                AgentVariant(
                    id="control",
                    name="Control",
                    description="Current production agent",
                    config=control_config,
                    traffic_percentage=1 - traffic_split
                ),
                AgentVariant(
                    id="treatment",
                    name="Treatment",
                    description="New agent variant",
                    config=treatment_config,
                    traffic_percentage=traffic_split
                )
            ],
            status=ExperimentStatus.DRAFT
        )

        self.experiments[experiment.id] = experiment
        return experiment

    def get_variant(
        self,
        experiment_id: str,
        user_id: str,
        session_id: str
    ) -> AgentVariant:
        """Deterministically assign user to variant"""

        experiment = self.experiments[experiment_id]

        if experiment.status != ExperimentStatus.RUNNING:
            # Return control if experiment not running
            return experiment.variants[0]

        # Check for existing assignment (sticky sessions)
        assignment_key = f"{experiment_id}:{user_id}"
        if assignment_key in self.assignments:
            variant_id = self.assignments[assignment_key]
            return next(v for v in experiment.variants if v.id == variant_id)

        # Deterministic assignment based on user_id hash
        hash_input = f"{experiment_id}:{user_id}"
        hash_value = int(hashlib.md5(hash_input.encode()).hexdigest(), 16)
        bucket = (hash_value % 100) / 100.0

        # Assign to variant based on traffic split
        cumulative = 0
        for variant in experiment.variants:
            cumulative += variant.traffic_percentage
            if bucket < cumulative:
                self.assignments[assignment_key] = variant.id
                return variant

        return experiment.variants[-1]

Experiment Configuration

# Example: Testing a new prompt strategy
experiment = manager.create_experiment(
    name="concise_responses_v2",
    hypothesis="Shorter, more direct responses will improve resolution rate and satisfaction",
    control_config={
        "model": "gpt-4o",
        "system_prompt": CURRENT_PROMPT,
        "max_tokens": 500,
        "temperature": 0.7
    },
    treatment_config={
        "model": "gpt-4o",
        "system_prompt": CONCISE_PROMPT,  # New shorter prompt
        "max_tokens": 300,  # Encourage brevity
        "temperature": 0.5  # More focused
    },
    traffic_split=0.2,  # 20% to treatment initially
    primary_metric="resolution_rate"
)

# Example: Testing model upgrade
model_experiment = manager.create_experiment(
    name="gpt4o_mini_migration",
    hypothesis="GPT-4o-mini can handle simple queries with equivalent quality at 90% lower cost",
    control_config={
        "model": "gpt-4o",
        "routing": "all_to_same_model"
    },
    treatment_config={
        "model": "gpt-4o-mini",
        "routing": "simple_queries_only",
        "complexity_threshold": 0.3
    },
    traffic_split=0.1,
    primary_metric="user_satisfaction"
)

Metrics Collection

@dataclass
class ExperimentMetrics:
    """Metrics for an experiment variant"""
    variant_id: str
    sample_size: int

    # Primary metrics
    resolution_rate: float  # Conversations resolved without escalation
    user_satisfaction: float  # Average rating

    # Secondary metrics
    avg_response_time_ms: float
    escalation_rate: float
    avg_turns_per_conversation: float
    cost_per_conversation: float
    followup_rate: float  # Questions after initial response

    # Quality metrics
    accuracy_score: float  # From eval pipeline
    helpfulness_score: float
    safety_incidents: int

    # Statistical
    confidence_interval: tuple
    p_value: Optional[float] = None


class MetricsCollector:
    """Collect and aggregate experiment metrics"""

    async def collect_for_experiment(
        self,
        experiment_id: str,
        date_range: tuple
    ) -> Dict[str, ExperimentMetrics]:
        """Collect metrics for all variants in an experiment"""

        experiment = self.experiments[experiment_id]
        results = {}

        for variant in experiment.variants:
            conversations = await self._get_conversations(
                experiment_id=experiment_id,
                variant_id=variant.id,
                date_range=date_range
            )

            metrics = self._calculate_metrics(conversations)
            results[variant.id] = metrics

        return results

    def _calculate_metrics(
        self,
        conversations: List[Conversation]
    ) -> ExperimentMetrics:
        """Calculate metrics from conversation data"""

        n = len(conversations)
        if n == 0:
            return None

        resolved = [c for c in conversations if c.resolved_without_escalation]
        ratings = [c.user_rating for c in conversations if c.user_rating]

        return ExperimentMetrics(
            variant_id=conversations[0].variant_id,
            sample_size=n,
            resolution_rate=len(resolved) / n,
            user_satisfaction=np.mean(ratings) if ratings else None,
            avg_response_time_ms=np.mean([c.response_time_ms for c in conversations]),
            escalation_rate=len([c for c in conversations if c.escalated]) / n,
            avg_turns_per_conversation=np.mean([c.turn_count for c in conversations]),
            cost_per_conversation=np.mean([c.total_cost for c in conversations]),
            followup_rate=len([c for c in conversations if c.had_followup]) / n,
            accuracy_score=np.mean([c.eval_accuracy for c in conversations if c.eval_accuracy]),
            helpfulness_score=np.mean([c.eval_helpfulness for c in conversations if c.eval_helpfulness]),
            safety_incidents=len([c for c in conversations if c.safety_flag]),
            confidence_interval=self._calculate_ci(resolved, n)
        )

Statistical Analysis

from scipy import stats
import numpy as np

class ExperimentAnalyzer:
    """Statistical analysis for experiments"""

    def analyze(
        self,
        control_metrics: ExperimentMetrics,
        treatment_metrics: ExperimentMetrics,
        primary_metric: str = "resolution_rate"
    ) -> Dict[str, Any]:
        """Compare control vs treatment"""

        control_value = getattr(control_metrics, primary_metric)
        treatment_value = getattr(treatment_metrics, primary_metric)

        # Calculate relative lift
        lift = (treatment_value - control_value) / control_value

        # Statistical significance (two-proportion z-test for rates)
        if primary_metric in ["resolution_rate", "escalation_rate", "followup_rate"]:
            p_value = self._proportion_test(
                control_metrics.sample_size,
                int(control_value * control_metrics.sample_size),
                treatment_metrics.sample_size,
                int(treatment_value * treatment_metrics.sample_size)
            )
        else:
            # t-test for continuous metrics
            p_value = self._t_test(control_metrics, treatment_metrics, primary_metric)

        # Determine if significant
        significant = p_value < 0.05

        # Calculate required sample size if not significant
        if not significant:
            required_n = self._calculate_required_sample(
                control_value, treatment_value,
                control_metrics.sample_size + treatment_metrics.sample_size
            )
        else:
            required_n = None

        return {
            "primary_metric": primary_metric,
            "control_value": control_value,
            "treatment_value": treatment_value,
            "lift": lift,
            "lift_percentage": f"{lift * 100:.2f}%",
            "p_value": p_value,
            "significant": significant,
            "required_sample_size": required_n,
            "recommendation": self._get_recommendation(
                lift, p_value, treatment_metrics.safety_incidents
            )
        }

    def _proportion_test(self, n1, x1, n2, x2) -> float:
        """Two-proportion z-test"""
        p1 = x1 / n1
        p2 = x2 / n2
        p_pool = (x1 + x2) / (n1 + n2)

        se = np.sqrt(p_pool * (1 - p_pool) * (1/n1 + 1/n2))
        z = (p1 - p2) / se

        return 2 * (1 - stats.norm.cdf(abs(z)))

    def _get_recommendation(
        self,
        lift: float,
        p_value: float,
        safety_incidents: int
    ) -> str:
        """Generate recommendation based on results"""

        if safety_incidents > 0:
            return "DO_NOT_SHIP - Safety incidents detected"

        if p_value >= 0.05:
            return "CONTINUE_TEST - Not yet significant"

        if lift > 0.05:  # >5% improvement
            return "SHIP - Significant positive impact"
        elif lift > 0:
            return "CONSIDER_SHIPPING - Small but significant improvement"
        elif lift > -0.02:
            return "NEUTRAL - No meaningful difference"
        else:
            return "DO_NOT_SHIP - Significant negative impact"

Experiment Results Dashboard

Progressive Rollout

class ProgressiveRollout:
    """Safely roll out winning variants"""

    ROLLOUT_STAGES = [
        {"percentage": 5, "duration_hours": 24, "auto_advance": False},
        {"percentage": 25, "duration_hours": 48, "auto_advance": True},
        {"percentage": 50, "duration_hours": 48, "auto_advance": True},
        {"percentage": 100, "duration_hours": None, "auto_advance": False},
    ]

    async def start_rollout(
        self,
        experiment_id: str,
        winning_variant_id: str
    ):
        """Begin progressive rollout of winning variant"""

        for stage in self.ROLLOUT_STAGES:
            # Update traffic allocation
            await self._set_traffic(
                experiment_id,
                winning_variant_id,
                stage["percentage"]
            )

            self.logger.info(
                f"Rollout stage: {stage['percentage']}% traffic to {winning_variant_id}"
            )

            # Monitor for issues
            if stage["duration_hours"]:
                issue_detected = await self._monitor_rollout(
                    experiment_id,
                    duration_hours=stage["duration_hours"]
                )

                if issue_detected:
                    await self._rollback(experiment_id)
                    return {"status": "rolled_back", "stage": stage["percentage"]}

            # Check for auto-advance
            if not stage["auto_advance"]:
                # Wait for manual approval
                await self._wait_for_approval(experiment_id, stage["percentage"])

        return {"status": "completed", "percentage": 100}

    async def _monitor_rollout(
        self,
        experiment_id: str,
        duration_hours: int
    ) -> bool:
        """Monitor for issues during rollout"""

        end_time = datetime.now() + timedelta(hours=duration_hours)

        while datetime.now() < end_time:
            metrics = await self.collector.get_recent_metrics(
                experiment_id,
                window_minutes=30
            )

            # Check for degradation
            if metrics.error_rate > 0.05:
                return True  # Issue detected
            if metrics.escalation_rate > metrics.baseline_escalation * 1.5:
                return True
            if metrics.safety_incidents > 0:
                return True

            await asyncio.sleep(300)  # Check every 5 minutes

        return False  # No issues

Key Takeaways

Next up: Security & Guardrails - Protecting your agent from malicious inputs and ensuring safe outputs.