MaverickMCP

"""A/B testing framework for comparing ML model performance.""" import logging import random from datetime import datetime from typing import Any import numpy as np import pandas as pd from scipy import stats from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score from .model_manager import ModelManager logger = logging.getLogger(__name__) class ABTestGroup: """Represents a group in an A/B test.""" def __init__( self, group_id: str, model_id: str, model_version: str, traffic_allocation: float, description: str = "", ): """Initialize A/B test group. Args: group_id: Unique identifier for the group model_id: Model identifier model_version: Model version traffic_allocation: Fraction of traffic allocated to this group (0-1) description: Description of the group """ self.group_id = group_id self.model_id = model_id self.model_version = model_version self.traffic_allocation = traffic_allocation self.description = description self.created_at = datetime.now() # Performance tracking self.predictions: list[Any] = [] self.actual_values: list[Any] = [] self.prediction_timestamps: list[datetime] = [] self.prediction_confidence: list[float] = [] def add_prediction( self, prediction: Any, actual: Any, confidence: float = 1.0, timestamp: datetime | None = None, ): """Add a prediction result to the group. Args: prediction: Model prediction actual: Actual value confidence: Prediction confidence score timestamp: Prediction timestamp """ self.predictions.append(prediction) self.actual_values.append(actual) self.prediction_confidence.append(confidence) self.prediction_timestamps.append(timestamp or datetime.now()) def get_metrics(self) -> dict[str, float]: """Calculate performance metrics for the group. Returns: Dictionary of performance metrics """ if not self.predictions or not self.actual_values: return {} try: predictions = np.array(self.predictions) actuals = np.array(self.actual_values) metrics = { "sample_count": len(predictions), "accuracy": accuracy_score(actuals, predictions), "precision": precision_score( actuals, predictions, average="weighted", zero_division=0 ), "recall": recall_score( actuals, predictions, average="weighted", zero_division=0 ), "f1_score": f1_score( actuals, predictions, average="weighted", zero_division=0 ), "avg_confidence": np.mean(self.prediction_confidence), } # Add confusion matrix for binary/multiclass unique_labels = np.unique(np.concatenate([predictions, actuals])) if len(unique_labels) <= 10: # Reasonable number of classes from sklearn.metrics import confusion_matrix cm = confusion_matrix(actuals, predictions, labels=unique_labels) metrics["confusion_matrix"] = cm.tolist() metrics["unique_labels"] = unique_labels.tolist() return metrics except Exception as e: logger.error(f"Error calculating metrics for group {self.group_id}: {e}") return {"error": str(e)} def to_dict(self) -> dict[str, Any]: """Convert group to dictionary representation.""" return { "group_id": self.group_id, "model_id": self.model_id, "model_version": self.model_version, "traffic_allocation": self.traffic_allocation, "description": self.description, "created_at": self.created_at.isoformat(), "metrics": self.get_metrics(), } class ABTest: """Manages an A/B test between different model versions.""" def __init__( self, test_id: str, name: str, description: str = "", random_seed: int | None = None, ): """Initialize A/B test. Args: test_id: Unique identifier for the test name: Human-readable name for the test description: Description of the test random_seed: Random seed for reproducible traffic splitting """ self.test_id = test_id self.name = name self.description = description self.created_at = datetime.now() self.started_at: datetime | None = None self.ended_at: datetime | None = None self.status = "created" # created, running, completed, cancelled # Groups in the test self.groups: dict[str, ABTestGroup] = {} # Traffic allocation self.traffic_splitter = TrafficSplitter(random_seed) # Test configuration self.min_samples_per_group = 100 self.confidence_level = 0.95 self.minimum_detectable_effect = 0.05 def add_group( self, group_id: str, model_id: str, model_version: str, traffic_allocation: float, description: str = "", ) -> bool: """Add a group to the A/B test. Args: group_id: Unique identifier for the group model_id: Model identifier model_version: Model version traffic_allocation: Fraction of traffic (0-1) description: Description of the group Returns: True if successful """ if self.status != "created": logger.error( f"Cannot add group to test {self.test_id} - test already started" ) return False if group_id in self.groups: logger.error(f"Group {group_id} already exists in test {self.test_id}") return False # Validate traffic allocation current_total = sum(g.traffic_allocation for g in self.groups.values()) if ( current_total + traffic_allocation > 1.0001 ): # Small tolerance for floating point logger.error( f"Traffic allocation would exceed 100%: {current_total + traffic_allocation}" ) return False group = ABTestGroup( group_id=group_id, model_id=model_id, model_version=model_version, traffic_allocation=traffic_allocation, description=description, ) self.groups[group_id] = group self.traffic_splitter.update_allocation( {gid: g.traffic_allocation for gid, g in self.groups.items()} ) logger.info(f"Added group {group_id} to test {self.test_id}") return True def start_test(self) -> bool: """Start the A/B test. Returns: True if successful """ if self.status != "created": logger.error( f"Cannot start test {self.test_id} - invalid status: {self.status}" ) return False if len(self.groups) < 2: logger.error(f"Cannot start test {self.test_id} - need at least 2 groups") return False # Validate traffic allocation sums to approximately 1.0 total_allocation = sum(g.traffic_allocation for g in self.groups.values()) if abs(total_allocation - 1.0) > 0.01: logger.error(f"Traffic allocation does not sum to 1.0: {total_allocation}") return False self.status = "running" self.started_at = datetime.now() logger.info(f"Started A/B test {self.test_id} with {len(self.groups)} groups") return True def assign_traffic(self, user_id: str | None = None) -> str | None: """Assign traffic to a group. Args: user_id: User identifier for consistent assignment Returns: Group ID or None if test not running """ if self.status != "running": return None return self.traffic_splitter.assign_group(user_id) def record_prediction( self, group_id: str, prediction: Any, actual: Any, confidence: float = 1.0, timestamp: datetime | None = None, ) -> bool: """Record a prediction result for a group. Args: group_id: Group identifier prediction: Model prediction actual: Actual value confidence: Prediction confidence timestamp: Prediction timestamp Returns: True if successful """ if group_id not in self.groups: logger.error(f"Group {group_id} not found in test {self.test_id}") return False self.groups[group_id].add_prediction(prediction, actual, confidence, timestamp) return True def get_results(self) -> dict[str, Any]: """Get current A/B test results. Returns: Dictionary with test results """ results = { "test_id": self.test_id, "name": self.name, "description": self.description, "status": self.status, "created_at": self.created_at.isoformat(), "started_at": self.started_at.isoformat() if self.started_at else None, "ended_at": self.ended_at.isoformat() if self.ended_at else None, "groups": {}, "statistical_analysis": {}, } # Group results for group_id, group in self.groups.items(): results["groups"][group_id] = group.to_dict() # Statistical analysis if len(self.groups) >= 2: results["statistical_analysis"] = self._perform_statistical_analysis() return results def _perform_statistical_analysis(self) -> dict[str, Any]: """Perform statistical analysis of A/B test results. Returns: Statistical analysis results """ analysis = { "ready_for_analysis": True, "sample_size_adequate": True, "statistical_significance": {}, "effect_sizes": {}, "recommendations": [], } # Check sample sizes sample_sizes = { group_id: len(group.predictions) for group_id, group in self.groups.items() } min_samples = min(sample_sizes.values()) if sample_sizes else 0 if min_samples < self.min_samples_per_group: analysis["ready_for_analysis"] = False analysis["sample_size_adequate"] = False analysis["recommendations"].append( f"Need at least {self.min_samples_per_group} samples per group (current min: {min_samples})" ) if not analysis["ready_for_analysis"]: return analysis # Pairwise comparisons group_ids = list(self.groups.keys()) for i, group_a_id in enumerate(group_ids): for group_b_id in group_ids[i + 1 :]: comparison_key = f"{group_a_id}_vs_{group_b_id}" try: group_a = self.groups[group_a_id] group_b = self.groups[group_b_id] # Compare accuracy scores accuracy_a = accuracy_score( group_a.actual_values, group_a.predictions ) accuracy_b = accuracy_score( group_b.actual_values, group_b.predictions ) # Perform statistical test # For classification accuracy, we can use a proportion test n_correct_a = sum( np.array(group_a.predictions) == np.array(group_a.actual_values) ) n_correct_b = sum( np.array(group_b.predictions) == np.array(group_b.actual_values) ) n_total_a = len(group_a.predictions) n_total_b = len(group_b.predictions) # Two-proportion z-test p_combined = (n_correct_a + n_correct_b) / (n_total_a + n_total_b) se = np.sqrt( p_combined * (1 - p_combined) * (1 / n_total_a + 1 / n_total_b) ) if se > 0: z_score = (accuracy_a - accuracy_b) / se p_value = 2 * (1 - stats.norm.cdf(abs(z_score))) # Effect size (Cohen's h for proportions) h = 2 * ( np.arcsin(np.sqrt(accuracy_a)) - np.arcsin(np.sqrt(accuracy_b)) ) analysis["statistical_significance"][comparison_key] = { "accuracy_a": accuracy_a, "accuracy_b": accuracy_b, "difference": accuracy_a - accuracy_b, "z_score": z_score, "p_value": p_value, "significant": p_value < (1 - self.confidence_level), "effect_size_h": h, } # Recommendations based on results if p_value < (1 - self.confidence_level): if accuracy_a > accuracy_b: analysis["recommendations"].append( f"Group {group_a_id} significantly outperforms {group_b_id} " f"(p={p_value:.4f}, effect_size={h:.4f})" ) else: analysis["recommendations"].append( f"Group {group_b_id} significantly outperforms {group_a_id} " f"(p={p_value:.4f}, effect_size={h:.4f})" ) else: analysis["recommendations"].append( f"No significant difference between {group_a_id} and {group_b_id} " f"(p={p_value:.4f})" ) except Exception as e: logger.error( f"Error in statistical analysis for {comparison_key}: {e}" ) analysis["statistical_significance"][comparison_key] = { "error": str(e) } return analysis def stop_test(self, reason: str = "completed") -> bool: """Stop the A/B test. Args: reason: Reason for stopping Returns: True if successful """ if self.status != "running": logger.error(f"Cannot stop test {self.test_id} - not running") return False self.status = "completed" if reason == "completed" else "cancelled" self.ended_at = datetime.now() logger.info(f"Stopped A/B test {self.test_id}: {reason}") return True def to_dict(self) -> dict[str, Any]: """Convert test to dictionary representation.""" return { "test_id": self.test_id, "name": self.name, "description": self.description, "status": self.status, "created_at": self.created_at.isoformat(), "started_at": self.started_at.isoformat() if self.started_at else None, "ended_at": self.ended_at.isoformat() if self.ended_at else None, "groups": {gid: g.to_dict() for gid, g in self.groups.items()}, "configuration": { "min_samples_per_group": self.min_samples_per_group, "confidence_level": self.confidence_level, "minimum_detectable_effect": self.minimum_detectable_effect, }, } class TrafficSplitter: """Handles traffic splitting for A/B tests.""" def __init__(self, random_seed: int | None = None): """Initialize traffic splitter. Args: random_seed: Random seed for reproducible splitting """ self.random_seed = random_seed self.group_allocation: dict[str, float] = {} self.cumulative_allocation: list[tuple[str, float]] = [] def update_allocation(self, allocation: dict[str, float]): """Update group traffic allocation. Args: allocation: Dictionary mapping group_id to allocation fraction """ self.group_allocation = allocation.copy() # Create cumulative distribution for sampling self.cumulative_allocation = [] cumulative = 0.0 for group_id, fraction in allocation.items(): cumulative += fraction self.cumulative_allocation.append((group_id, cumulative)) def assign_group(self, user_id: str | None = None) -> str | None: """Assign a user to a group. Args: user_id: User identifier for consistent assignment Returns: Group ID or None if no groups configured """ if not self.cumulative_allocation: return None # Generate random value if user_id is not None: # Hash user_id for consistent assignment import hashlib hash_object = hashlib.md5(user_id.encode()) hash_int = int(hash_object.hexdigest(), 16) rand_value = (hash_int % 10000) / 10000.0 # Normalize to [0, 1) else: if self.random_seed is not None: random.seed(self.random_seed) rand_value = random.random() # Find group based on cumulative allocation for group_id, cumulative_threshold in self.cumulative_allocation: if rand_value <= cumulative_threshold: return group_id # Fallback to last group return self.cumulative_allocation[-1][0] if self.cumulative_allocation else None class ABTestManager: """Manages multiple A/B tests.""" def __init__(self, model_manager: ModelManager): """Initialize A/B test manager. Args: model_manager: Model manager instance """ self.model_manager = model_manager self.tests: dict[str, ABTest] = {} def create_test( self, test_id: str, name: str, description: str = "", random_seed: int | None = None, ) -> ABTest: """Create a new A/B test. Args: test_id: Unique identifier for the test name: Human-readable name description: Description random_seed: Random seed for reproducible splitting Returns: ABTest instance """ if test_id in self.tests: raise ValueError(f"Test {test_id} already exists") test = ABTest(test_id, name, description, random_seed) self.tests[test_id] = test logger.info(f"Created A/B test {test_id}: {name}") return test def get_test(self, test_id: str) -> ABTest | None: """Get an A/B test by ID. Args: test_id: Test identifier Returns: ABTest instance or None """ return self.tests.get(test_id) def list_tests(self, status_filter: str | None = None) -> list[dict[str, Any]]: """List all A/B tests. Args: status_filter: Filter by status (created, running, completed, cancelled) Returns: List of test summaries """ tests = [] for test in self.tests.values(): if status_filter is None or test.status == status_filter: tests.append( { "test_id": test.test_id, "name": test.name, "status": test.status, "groups_count": len(test.groups), "created_at": test.created_at.isoformat(), "started_at": test.started_at.isoformat() if test.started_at else None, } ) # Sort by creation time (newest first) tests.sort(key=lambda x: x["created_at"], reverse=True) return tests def run_model_comparison( self, test_name: str, model_versions: list[tuple[str, str]], # (model_id, version) test_data: pd.DataFrame, feature_extractor: Any, target_extractor: Any, traffic_allocation: list[float] | None = None, test_duration_hours: int = 24, ) -> str: """Run a model comparison A/B test. Args: test_name: Name for the test model_versions: List of (model_id, version) tuples to compare test_data: Test data for evaluation feature_extractor: Function to extract features target_extractor: Function to extract targets traffic_allocation: Custom traffic allocation (defaults to equal split) test_duration_hours: Duration to run the test Returns: Test ID """ # Generate unique test ID test_id = f"comparison_{datetime.now().strftime('%Y%m%d_%H%M%S')}" # Create test test = self.create_test( test_id=test_id, name=test_name, description=f"Comparing {len(model_versions)} model versions", ) # Default equal traffic allocation if traffic_allocation is None: allocation_per_group = 1.0 / len(model_versions) traffic_allocation = [allocation_per_group] * len(model_versions) # Add groups for i, (model_id, version) in enumerate(model_versions): group_id = f"group_{i}_{model_id}_{version}" test.add_group( group_id=group_id, model_id=model_id, model_version=version, traffic_allocation=traffic_allocation[i], description=f"Model {model_id} version {version}", ) # Start test test.start_test() # Extract features and targets features = feature_extractor(test_data) targets = target_extractor(test_data) # Simulate predictions for each group for _, row in features.iterrows(): # Assign traffic group_id = test.assign_traffic(str(row.name)) # Use row index as user_id if group_id is None: continue # Get corresponding group's model group = test.groups[group_id] model_version = self.model_manager.load_model( group.model_id, group.model_version ) if model_version is None or model_version.model is None: logger.warning(f"Could not load model for group {group_id}") continue try: # Make prediction X = row.values.reshape(1, -1) if model_version.scaler is not None: X = model_version.scaler.transform(X) prediction = model_version.model.predict(X)[0] # Get confidence if available confidence = 1.0 if hasattr(model_version.model, "predict_proba"): proba = model_version.model.predict_proba(X)[0] confidence = max(proba) # Get actual value actual = targets.loc[row.name] # Record prediction test.record_prediction(group_id, prediction, actual, confidence) except Exception as e: logger.warning(f"Error making prediction for group {group_id}: {e}") logger.info(f"Completed model comparison test {test_id}") return test_id def get_test_summary(self) -> dict[str, Any]: """Get summary of all A/B tests. Returns: Summary dictionary """ total_tests = len(self.tests) status_counts = {} for test in self.tests.values(): status_counts[test.status] = status_counts.get(test.status, 0) + 1 recent_tests = sorted( [ { "test_id": test.test_id, "name": test.name, "status": test.status, "created_at": test.created_at.isoformat(), } for test in self.tests.values() ], key=lambda x: x["created_at"], reverse=True, )[:10] return { "total_tests": total_tests, "status_distribution": status_counts, "recent_tests": recent_tests, }