ds-agent-cli 0.1.0

package/src/tools/advanced_training.py

@@ -0,0 +1,906 @@
+"""
+Advanced Model Training Tools
+Tools for hyperparameter tuning, ensemble methods, and cross-validation.
+"""
+
+import polars as pl
+import pandas as pd
+import numpy as np
+from typing import Dict, Any, List, Optional, Tuple
+from pathlib import Path
+import sys
+import os
+import joblib
+import json
+import optuna
+from optuna.pruners import MedianPruner
+from optuna.samplers import TPESampler
+import warnings
+import tempfile
+
+warnings.filterwarnings('ignore')
+
+# Add parent directory to path for imports
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+# Import artifact store
+try:
+    from storage.helpers import save_model_with_store
+    ARTIFACT_STORE_AVAILABLE = True
+except ImportError:
+    ARTIFACT_STORE_AVAILABLE = False
+    print("⚠️  Artifact store not available, using local paths")
+
+from sklearn.model_selection import train_test_split, KFold, StratifiedKFold, TimeSeriesSplit, cross_val_score
+from sklearn.linear_model import LogisticRegression, Ridge, Lasso, ElasticNet
+from sklearn.ensemble import (
+    RandomForestClassifier, RandomForestRegressor,
+    GradientBoostingClassifier, GradientBoostingRegressor,
+    VotingClassifier, VotingRegressor,
+    StackingClassifier, StackingRegressor
+)
+from xgboost import XGBClassifier, XGBRegressor
+from sklearn.metrics import (
+    accuracy_score, precision_score, recall_score, f1_score, roc_auc_score,
+    mean_squared_error, mean_absolute_error, r2_score
+)
+
+from ds_agent.utils.polars_helpers import load_dataframe, get_numeric_columns, split_features_target
+from ds_agent.utils.validation import (
+    validate_file_exists, validate_file_format, validate_dataframe,
+    validate_column_exists, validate_target_column
+)
+
+
+def hyperparameter_tuning(
+    file_path: str,
+    target_col: str,
+    model_type: str = "random_forest",
+    task_type: str = "auto",
+    n_trials: int = 50,
+    cv_folds: int = 5,
+    optimization_metric: str = "auto",
+    test_size: float = 0.2,
+    random_state: int = 42,
+    output_path: Optional[str] = None
+) -> Dict[str, Any]:
+    """
+    Perform Bayesian hyperparameter optimization using Optuna.
+    
+    ⚠️ WARNING: This tool is VERY computationally expensive and can take 5-10 minutes!
+    For large datasets (>100K rows), n_trials is automatically reduced to prevent timeout.
+    
+    Args:
+        file_path: Path to prepared dataset
+        target_col: Target column name
+        model_type: Model to tune ('random_forest', 'xgboost', 'lightgbm', 'catboost', 'logistic', 'ridge')
+        task_type: 'classification', 'regression', or 'auto' (detect from target)
+        n_trials: Number of optimization trials (default 50, auto-reduced for large datasets)
+        cv_folds: Number of cross-validation folds
+        optimization_metric: Metric to optimize ('auto', 'accuracy', 'f1', 'roc_auc', 'rmse', 'r2')
+        test_size: Test set size for final evaluation
+        random_state: Random seed
+        output_path: Path to save best model
+        
+    Returns:
+        Dictionary with tuning results, best parameters, and performance
+    """
+    # ⚠️ CRITICAL FIX: Convert integer params (Gemini/LLMs pass floats)
+    n_trials = int(n_trials)
+    cv_folds = int(cv_folds)
+    random_state = int(random_state)
+    
+    # Validation
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    # Load data
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    validate_column_exists(df, target_col)
+    
+    # ⚠️ CRITICAL: Auto-reduce trials for large datasets to prevent memory crashes
+    n_rows = len(df)
+    if n_rows > 100000 and n_trials > 20:
+        original_trials = n_trials
+        n_trials = 20
+        print(f"   ⚠️ Large dataset ({n_rows:,} rows) - reducing trials from {original_trials} to {n_trials} to prevent timeout")
+    elif n_rows > 50000 and n_trials > 30:
+        original_trials = n_trials
+        n_trials = 30
+        print(f"   ⚠️ Medium dataset ({n_rows:,} rows) - reducing trials from {original_trials} to {n_trials}")
+    
+    # ⚠️ PERFORMANCE FIX: Sample large datasets for hyperparameter tuning
+    # Hyperparameters found on sample will be used to train final model on full dataset
+    MAX_TUNING_ROWS = 50000
+    sampled = False
+    if n_rows > MAX_TUNING_ROWS:
+        original_rows = n_rows
+        sample_frac = MAX_TUNING_ROWS / n_rows
+        df = df.sample(n=MAX_TUNING_ROWS, random_state=random_state)
+        sampled = True
+        print(f"   📊 Sampled {MAX_TUNING_ROWS:,} rows ({sample_frac:.1%}) from {original_rows:,} for faster tuning")
+        print(f"   💡 Hyperparameters found on sample will generalize well to full dataset")
+        print(f"   ⏱️ Expected speedup: 3-5x faster tuning")
+    
+    # ⚠️ Auto-reduce CV folds for very large datasets
+    original_cv_folds = cv_folds
+    if n_rows > 100000 and cv_folds > 3:
+        cv_folds = 3
+        print(f"   ⏱️ Using {cv_folds}-fold CV (instead of {original_cv_folds}) for faster tuning on large dataset")
+    
+    # ⚠️ SKIP DATETIME CONVERSION: Already handled by create_time_features() in workflow step 7
+    # The encoded.csv file should already have time features extracted
+    # If datetime columns still exist, they will be handled as regular features
+    
+    # ⚠️ CRITICAL FIX: Convert Polars to Pandas if needed (for XGBoost compatibility)
+    if hasattr(df, 'to_pandas'):
+        print(f"   🔄 Converting Polars DataFrame to Pandas for XGBoost compatibility...")
+        df = df.to_pandas()
+    
+    # ⚠️ CRITICAL: Drop any remaining datetime columns that weren't converted to features
+    # XGBoost cannot handle Timestamp objects in NumPy arrays
+    if isinstance(df, pd.DataFrame):
+        datetime_cols = df.select_dtypes(include=['datetime64', 'datetime64[ns]', 'datetime64[ns, UTC]']).columns.tolist()
+        if datetime_cols:
+            print(f"   ⚠️ Dropping {len(datetime_cols)} datetime columns that cannot be used directly: {datetime_cols}")
+            print(f"   💡 Time features should have been extracted in workflow step 7 (create_time_features)")
+            df = df.drop(columns=datetime_cols)
+        
+        # ⚠️ CRITICAL: Drop any remaining string/object columns (not encoded properly)
+        # XGBoost cannot handle string values like 'mb', 'ml', etc.
+        object_cols = df.select_dtypes(include=['object', 'string']).columns.tolist()
+        # Don't drop the target column if it's object type
+        object_cols = [col for col in object_cols if col != target_col]
+        if object_cols:
+            print(f"   ⚠️ Dropping {len(object_cols)} string columns that weren't encoded: {object_cols}")
+            print(f"   💡 Categorical encoding should have been done in workflow step 8 (encode_categorical)")
+            print(f"   💡 These columns likely weren't in the encoded file or encoding failed")
+            df = df.drop(columns=object_cols)
+    
+    # Prepare data - handle both Polars and Pandas
+    if target_col not in df.columns:
+        raise ValueError(f"Target column '{target_col}' not found in dataframe. Available columns: {list(df.columns)}")
+    
+    # Split features and target (works for both Polars and Pandas)
+    if hasattr(df, 'drop'):  # Both have drop method
+        X = df.drop(columns=[target_col]) if isinstance(df, pd.DataFrame) else df.drop(target_col)
+        y = df[target_col]
+    else:
+        X, y = split_features_target(df, target_col)
+    
+    # Convert to numpy for sklearn compatibility
+    if hasattr(X, 'to_numpy'):
+        X = X.to_numpy()
+        y = y.to_numpy()
+    elif hasattr(X, 'values'):
+        X = X.values
+        y = y.values
+    
+    X_train, X_test, y_train, y_test = train_test_split(
+        X, y, test_size=test_size, random_state=random_state, stratify=y if task_type == "classification" else None
+    )
+    
+    # Detect task type
+    if task_type == "auto":
+        unique_values = len(np.unique(y))
+        task_type = "classification" if unique_values < 20 else "regression"
+    
+    # Set default metric
+    if optimization_metric == "auto":
+        optimization_metric = "accuracy" if task_type == "classification" else "rmse"
+    
+    # Define objective function for Optuna
+    def objective(trial):
+        # Suggest hyperparameters based on model type
+        if model_type == "random_forest":
+            params = {
+                'n_estimators': trial.suggest_int('n_estimators', 50, 500),
+                'max_depth': trial.suggest_int('max_depth', 3, 20),
+                'min_samples_split': trial.suggest_int('min_samples_split', 2, 20),
+                'min_samples_leaf': trial.suggest_int('min_samples_leaf', 1, 10),
+                'max_features': trial.suggest_categorical('max_features', ['sqrt', 'log2', None]),
+                'random_state': random_state
+            }
+            if task_type == "classification":
+                model = RandomForestClassifier(**params)
+            else:
+                model = RandomForestRegressor(**params)
+                
+        elif model_type == "xgboost":
+            params = {
+                'n_estimators': trial.suggest_int('n_estimators', 50, 500),
+                'max_depth': trial.suggest_int('max_depth', 3, 10),
+                'learning_rate': trial.suggest_float('learning_rate', 0.01, 0.3, log=True),
+                'subsample': trial.suggest_float('subsample', 0.5, 1.0),
+                'colsample_bytree': trial.suggest_float('colsample_bytree', 0.5, 1.0),
+                'gamma': trial.suggest_float('gamma', 0, 5),
+                'reg_alpha': trial.suggest_float('reg_alpha', 0, 2),
+                'reg_lambda': trial.suggest_float('reg_lambda', 0, 2),
+                'random_state': random_state
+            }
+            if task_type == "classification":
+                model = XGBClassifier(**params, use_label_encoder=False, eval_metric='logloss')
+            else:
+                model = XGBRegressor(**params)
+                
+        elif model_type == "logistic":
+            params = {
+                'C': trial.suggest_float('C', 0.001, 100, log=True),
+                'penalty': trial.suggest_categorical('penalty', ['l1', 'l2', 'elasticnet']),
+                'solver': 'saga',
+                'max_iter': 1000,
+                'random_state': random_state
+            }
+            if params['penalty'] == 'elasticnet':
+                params['l1_ratio'] = trial.suggest_float('l1_ratio', 0, 1)
+            model = LogisticRegression(**params)
+            
+        elif model_type == "ridge":
+            params = {
+                'alpha': trial.suggest_float('alpha', 0.001, 100, log=True),
+                'solver': trial.suggest_categorical('solver', ['auto', 'svd', 'cholesky', 'lsqr']),
+                'random_state': random_state
+            }
+            model = Ridge(**params)
+            
+        elif model_type == "lightgbm":
+            from lightgbm import LGBMClassifier, LGBMRegressor
+            params = {
+                'n_estimators': trial.suggest_int('n_estimators', 50, 500),
+                'max_depth': trial.suggest_int('max_depth', 3, 12),
+                'learning_rate': trial.suggest_float('learning_rate', 0.01, 0.3, log=True),
+                'subsample': trial.suggest_float('subsample', 0.5, 1.0),
+                'colsample_bytree': trial.suggest_float('colsample_bytree', 0.5, 1.0),
+                'reg_alpha': trial.suggest_float('reg_alpha', 1e-8, 10.0, log=True),
+                'reg_lambda': trial.suggest_float('reg_lambda', 1e-8, 10.0, log=True),
+                'num_leaves': trial.suggest_int('num_leaves', 15, 127),
+                'min_child_samples': trial.suggest_int('min_child_samples', 5, 100),
+                'random_state': random_state,
+                'verbosity': -1
+            }
+            if task_type == "classification":
+                model = LGBMClassifier(**params)
+            else:
+                model = LGBMRegressor(**params)
+                
+        elif model_type == "catboost":
+            from catboost import CatBoostClassifier, CatBoostRegressor
+            params = {
+                'iterations': trial.suggest_int('iterations', 50, 500),
+                'depth': trial.suggest_int('depth', 3, 10),
+                'learning_rate': trial.suggest_float('learning_rate', 0.01, 0.3, log=True),
+                'l2_leaf_reg': trial.suggest_float('l2_leaf_reg', 1e-8, 10.0, log=True),
+                'bagging_temperature': trial.suggest_float('bagging_temperature', 0, 10),
+                'random_strength': trial.suggest_float('random_strength', 0, 10),
+                'random_seed': random_state,
+                'verbose': 0
+            }
+            if task_type == "classification":
+                model = CatBoostClassifier(**params)
+            else:
+                model = CatBoostRegressor(**params)
+        else:
+            raise ValueError(f"Unsupported model_type: {model_type}. Use 'random_forest', 'xgboost', 'lightgbm', 'catboost', 'logistic', or 'ridge'.")
+        
+        # Cross-validation
+        if task_type == "classification":
+            cv = StratifiedKFold(n_splits=cv_folds, shuffle=True, random_state=random_state)
+        else:
+            cv = KFold(n_splits=cv_folds, shuffle=True, random_state=random_state)
+        
+        # Select scoring metric
+        if optimization_metric == "accuracy":
+            scoring = 'accuracy'
+        elif optimization_metric == "f1":
+            scoring = 'f1_weighted'
+        elif optimization_metric == "roc_auc":
+            scoring = 'roc_auc_ovr_weighted'
+        elif optimization_metric == "rmse":
+            scoring = 'neg_root_mean_squared_error'
+        elif optimization_metric == "r2":
+            scoring = 'r2'
+        else:
+            scoring = 'accuracy' if task_type == "classification" else 'neg_root_mean_squared_error'
+        
+        # Cross-validation score
+        scores = cross_val_score(model, X_train, y_train, cv=cv, scoring=scoring, n_jobs=-1)
+        
+        # Return mean score (Optuna maximizes by default)
+        return scores.mean()
+    
+    # Run optimization
+    print(f"🔧 Starting hyperparameter tuning with {n_trials} trials...")
+    study = optuna.create_study(
+        direction='maximize',
+        sampler=TPESampler(seed=random_state),
+        pruner=MedianPruner(n_startup_trials=5, n_warmup_steps=10)
+    )
+    study.optimize(objective, n_trials=n_trials, show_progress_bar=True)
+    
+    # Get best parameters
+    best_params = study.best_params
+    best_score = study.best_value
+    
+    print(f"✅ Best {optimization_metric}: {best_score:.4f}")
+    print(f"📊 Best parameters: {best_params}")
+    
+    # Train final model with best parameters
+    if model_type == "random_forest":
+        if task_type == "classification":
+            final_model = RandomForestClassifier(**best_params)
+        else:
+            final_model = RandomForestRegressor(**best_params)
+    elif model_type == "xgboost":
+        if task_type == "classification":
+            final_model = XGBClassifier(**best_params, use_label_encoder=False, eval_metric='logloss')
+        else:
+            final_model = XGBRegressor(**best_params)
+    elif model_type == "logistic":
+        final_model = LogisticRegression(**best_params)
+    elif model_type == "ridge":
+        final_model = Ridge(**best_params)
+    
+    final_model.fit(X_train, y_train)
+    
+    # Evaluate on test set
+    y_pred = final_model.predict(X_test)
+    
+    if task_type == "classification":
+        test_metrics = {
+            'accuracy': float(accuracy_score(y_test, y_pred)),
+            'precision': float(precision_score(y_test, y_pred, average='weighted', zero_division=0)),
+            'recall': float(recall_score(y_test, y_pred, average='weighted', zero_division=0)),
+            'f1': float(f1_score(y_test, y_pred, average='weighted', zero_division=0))
+        }
+        if len(np.unique(y)) == 2:
+            y_pred_proba = final_model.predict_proba(X_test)[:, 1]
+            test_metrics['roc_auc'] = float(roc_auc_score(y_test, y_pred_proba))
+    else:
+        test_metrics = {
+            'rmse': float(np.sqrt(mean_squared_error(y_test, y_pred))),
+            'mae': float(mean_absolute_error(y_test, y_pred)),
+            'r2': float(r2_score(y_test, y_pred))
+        }
+    
+    # Save model if output path provided
+    actual_model_path = None
+    if output_path:
+        if ARTIFACT_STORE_AVAILABLE:
+            # Save using artifact store (returns internal storage path)
+            actual_model_path = save_model_with_store(
+                model_data=final_model,
+                filename=os.path.basename(output_path),
+                metadata={
+                    "model_type": model_type,
+                    "task_type": task_type,
+                    "best_params": best_params,
+                    "cv_score": float(best_score),
+                    "test_metrics": test_metrics
+                }
+            )
+            # Also save to user-requested path for LLM to find it
+            os.makedirs(os.path.dirname(output_path), exist_ok=True)
+            joblib.dump(final_model, output_path)
+            print(f"💾 Model saved to: {output_path} (artifact store: {actual_model_path})")
+        else:
+            os.makedirs(os.path.dirname(output_path), exist_ok=True)
+            joblib.dump(final_model, output_path)
+            actual_model_path = output_path
+            print(f"💾 Model saved to: {output_path}")
+    
+    return {
+        'status': 'success',
+        'model_type': model_type,
+        'task_type': task_type,
+        'n_trials': n_trials,
+        'best_params': best_params,
+        'best_cv_score': float(best_score),
+        'optimization_metric': optimization_metric,
+        'test_metrics': test_metrics,
+        'trials_summary': {
+            'total_trials': len(study.trials),
+            'best_trial': study.best_trial.number,
+            'completed_trials': len([t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE])
+        },
+        'model_path': output_path if output_path else None
+    }
+
+
+def train_ensemble_models(
+    file_path: str,
+    target_col: str,
+    ensemble_type: str = "voting",
+    task_type: str = "auto",
+    test_size: float = 0.2,
+    random_state: int = 42,
+    output_path: Optional[str] = None
+) -> Dict[str, Any]:
+    """
+    Train ensemble models using stacking, blending, or voting.
+    
+    Args:
+        file_path: Path to prepared dataset
+        target_col: Target column name
+        ensemble_type: 'voting', 'stacking', or 'blending'
+        task_type: 'classification', 'regression', or 'auto'
+        test_size: Test set size
+        random_state: Random seed
+        output_path: Path to save ensemble model
+        
+    Returns:
+        Dictionary with ensemble performance and comparison
+    """
+    # Validation
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    # Load data
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    validate_column_exists(df, target_col)
+    
+    # ⚠️ SKIP DATETIME CONVERSION: Already handled by create_time_features() in workflow step 7
+    # The encoded.csv file should already have time features extracted
+    
+    # ⚠️ CRITICAL FIX: Convert Polars to Pandas if needed (for XGBoost compatibility)
+    if hasattr(df, 'to_pandas'):
+        print(f"   🔄 Converting Polars DataFrame to Pandas for XGBoost compatibility...")
+        df = df.to_pandas()
+    
+    # ⚠️ CRITICAL: Drop remaining datetime columns BEFORE NumPy conversion
+    # XGBoost cannot handle Timestamp objects (causes TypeError: float() argument must be a string or a real number, not 'Timestamp')
+    if isinstance(df, pd.DataFrame):
+        datetime_cols = df.select_dtypes(include=['datetime64', 'datetime64[ns]', 'datetime64[ns, UTC]']).columns.tolist()
+        if datetime_cols:
+            print(f"   ⚠️ Dropping {len(datetime_cols)} datetime columns: {datetime_cols}")
+            print(f"   💡 Time features should have been extracted in workflow step 7 (create_time_features)")
+            df = df.drop(columns=datetime_cols)
+        
+        # ⚠️ CRITICAL: Drop any remaining string/object columns (not encoded properly)
+        object_cols = df.select_dtypes(include=['object', 'string']).columns.tolist()
+        object_cols = [col for col in object_cols if col != target_col]
+        if object_cols:
+            print(f"   ⚠️ Dropping {len(object_cols)} string columns that weren't encoded: {object_cols}")
+            print(f"   💡 Categorical encoding should have been done in workflow step 8")
+            df = df.drop(columns=object_cols)
+    
+    # Prepare data - handle both Polars and Pandas
+    if target_col not in df.columns:
+        raise ValueError(f"Target column '{target_col}' not found in dataframe. Available columns: {list(df.columns)}")
+    
+    # Split features and target (works for both Polars and Pandas)
+    if hasattr(df, 'drop'):
+        X = df.drop(columns=[target_col]) if isinstance(df, pd.DataFrame) else df.drop(target_col)
+        y = df[target_col]
+    else:
+        X, y = split_features_target(df, target_col)
+    
+    # Convert to numpy for sklearn compatibility
+    if hasattr(X, 'to_numpy'):
+        X = X.to_numpy()
+        y = y.to_numpy()
+    elif hasattr(X, 'values'):
+        X = X.values
+        y = y.values
+    
+    # Detect task type
+    if task_type == "auto":
+        unique_values = len(np.unique(y))
+        task_type = "classification" if unique_values < 20 else "regression"
+    
+    X_train, X_test, y_train, y_test = train_test_split(
+        X, y, test_size=test_size, random_state=random_state,
+        stratify=y if task_type == "classification" else None
+    )
+    
+    # Define base models
+    if task_type == "classification":
+        base_models = [
+            ('lr', LogisticRegression(max_iter=1000, random_state=random_state)),
+            ('rf', RandomForestClassifier(n_estimators=100, random_state=random_state)),
+            ('xgb', XGBClassifier(n_estimators=100, random_state=random_state, use_label_encoder=False, eval_metric='logloss'))
+        ]
+        meta_model = LogisticRegression(max_iter=1000, random_state=random_state)
+    else:
+        base_models = [
+            ('ridge', Ridge(random_state=random_state)),
+            ('rf', RandomForestRegressor(n_estimators=100, random_state=random_state)),
+            ('xgb', XGBRegressor(n_estimators=100, random_state=random_state))
+        ]
+        meta_model = Ridge(random_state=random_state)
+    
+    # Train individual models for comparison
+    individual_results = {}
+    for name, model in base_models:
+        model.fit(X_train, y_train)
+        y_pred = model.predict(X_test)
+        
+        if task_type == "classification":
+            individual_results[name] = {
+                'accuracy': float(accuracy_score(y_test, y_pred)),
+                'f1': float(f1_score(y_test, y_pred, average='weighted', zero_division=0))
+            }
+        else:
+            individual_results[name] = {
+                'rmse': float(np.sqrt(mean_squared_error(y_test, y_pred))),
+                'r2': float(r2_score(y_test, y_pred))
+            }
+    
+    # Create ensemble
+    print(f"🎯 Building {ensemble_type} ensemble...")
+    
+    if ensemble_type == "voting":
+        if task_type == "classification":
+            ensemble = VotingClassifier(estimators=base_models, voting='soft')
+        else:
+            ensemble = VotingRegressor(estimators=base_models)
+            
+    elif ensemble_type == "stacking":
+        if task_type == "classification":
+            ensemble = StackingClassifier(
+                estimators=base_models,
+                final_estimator=meta_model,
+                cv=5
+            )
+        else:
+            ensemble = StackingRegressor(
+                estimators=base_models,
+                final_estimator=meta_model,
+                cv=5
+            )
+    
+    elif ensemble_type == "blending":
+        # Split training data for blending
+        X_base_train, X_blend_train, y_base_train, y_blend_train = train_test_split(
+            X_train, y_train, test_size=0.3, random_state=random_state,
+            stratify=y_train if task_type == "classification" else None
+        )
+        
+        # Train base models on base training set
+        base_predictions_train = []
+        base_predictions_test = []
+        
+        for name, model in base_models:
+            model.fit(X_base_train, y_base_train)
+            base_predictions_train.append(model.predict(X_blend_train))
+            base_predictions_test.append(model.predict(X_test))
+        
+        # Stack predictions
+        X_blend = np.column_stack(base_predictions_train)
+        X_test_blend = np.column_stack(base_predictions_test)
+        
+        # Train meta-model
+        meta_model.fit(X_blend, y_blend_train)
+        y_pred = meta_model.predict(X_test_blend)
+        
+        # Calculate metrics
+        if task_type == "classification":
+            ensemble_metrics = {
+                'accuracy': float(accuracy_score(y_test, y_pred)),
+                'precision': float(precision_score(y_test, y_pred, average='weighted', zero_division=0)),
+                'recall': float(recall_score(y_test, y_pred, average='weighted', zero_division=0)),
+                'f1': float(f1_score(y_test, y_pred, average='weighted', zero_division=0))
+            }
+        else:
+            ensemble_metrics = {
+                'rmse': float(np.sqrt(mean_squared_error(y_test, y_pred))),
+                'mae': float(mean_absolute_error(y_test, y_pred)),
+                'r2': float(r2_score(y_test, y_pred))
+            }
+        
+        # Save for blending
+        actual_model_path = None
+        if output_path:
+            if ARTIFACT_STORE_AVAILABLE:
+                # Save using artifact store (returns internal storage path)
+                actual_model_path = save_model_with_store(
+                    model_data={
+                        'base_models': dict(base_models),
+                        'meta_model': meta_model,
+                        'ensemble_type': 'blending'
+                    },
+                    filename=os.path.basename(output_path),
+                    metadata={
+                        "ensemble_type": "blending",
+                        "task_type": task_type,
+                        "ensemble_metrics": ensemble_metrics,
+                        "num_base_models": len(base_models)
+                    }
+                )
+                # Also save to user-requested path for LLM to find it
+                os.makedirs(os.path.dirname(output_path), exist_ok=True)
+                joblib.dump({
+                    'base_models': dict(base_models),
+                    'meta_model': meta_model,
+                    'ensemble_type': 'blending'
+                }, output_path)
+            else:
+                os.makedirs(os.path.dirname(output_path), exist_ok=True)
+                joblib.dump({
+                    'base_models': dict(base_models),
+                    'meta_model': meta_model,
+                    'ensemble_type': 'blending'
+                }, output_path)
+                actual_model_path = output_path
+        
+        return {
+            'status': 'success',
+            'ensemble_type': ensemble_type,
+            'task_type': task_type,
+            'ensemble_metrics': ensemble_metrics,
+            'individual_models': individual_results,
+            'improvement': f"+{(ensemble_metrics.get('accuracy', ensemble_metrics.get('r2', 0)) - max([m.get('accuracy', m.get('r2', 0)) for m in individual_results.values()])) * 100:.2f}%",
+            'model_path': output_path if output_path else None
+        }
+    
+    else:
+        raise ValueError(f"Unsupported ensemble_type: {ensemble_type}")
+    
+    # Train ensemble (voting or stacking)
+    ensemble.fit(X_train, y_train)
+    y_pred = ensemble.predict(X_test)
+    
+    # Calculate ensemble metrics
+    if task_type == "classification":
+        ensemble_metrics = {
+            'accuracy': float(accuracy_score(y_test, y_pred)),
+            'precision': float(precision_score(y_test, y_pred, average='weighted', zero_division=0)),
+            'recall': float(recall_score(y_test, y_pred, average='weighted', zero_division=0)),
+            'f1': float(f1_score(y_test, y_pred, average='weighted', zero_division=0))
+        }
+        best_individual_metric = max([m['accuracy'] for m in individual_results.values()])
+        improvement = ensemble_metrics['accuracy'] - best_individual_metric
+    else:
+        ensemble_metrics = {
+            'rmse': float(np.sqrt(mean_squared_error(y_test, y_pred))),
+            'mae': float(mean_absolute_error(y_test, y_pred)),
+            'r2': float(r2_score(y_test, y_pred))
+        }
+        best_individual_metric = max([m['r2'] for m in individual_results.values()])
+        improvement = ensemble_metrics['r2'] - best_individual_metric
+    
+    # Save model
+    actual_model_path = None
+    if output_path:
+        if ARTIFACT_STORE_AVAILABLE:
+            # Save using artifact store (returns internal storage path)
+            actual_model_path = save_model_with_store(
+                model_data=ensemble,
+                filename=os.path.basename(output_path),
+                metadata={
+                    "ensemble_type": ensemble_type,
+                    "task_type": task_type,
+                    "ensemble_metrics": ensemble_metrics,
+                    "improvement_pct": float(improvement * 100)
+                }
+            )
+            # Also save to user-requested path for LLM to find it
+            os.makedirs(os.path.dirname(output_path), exist_ok=True)
+            joblib.dump(ensemble, output_path)
+            print(f"💾 Ensemble model saved to: {output_path} (artifact store: {actual_model_path})")
+        else:
+            os.makedirs(os.path.dirname(output_path), exist_ok=True)
+            joblib.dump(ensemble, output_path)
+            actual_model_path = output_path
+            print(f"💾 Ensemble model saved to: {output_path}")
+    
+    return {
+        'status': 'success',
+        'ensemble_type': ensemble_type,
+        'task_type': task_type,
+        'ensemble_metrics': ensemble_metrics,
+        'individual_models': individual_results,
+        'improvement': f"+{improvement * 100:.2f}%",
+        'model_path': output_path if output_path else None
+    }
+
+
+def perform_cross_validation(
+    file_path: str,
+    target_col: str,
+    model_type: str = "random_forest",
+    task_type: str = "auto",
+    cv_strategy: str = "kfold",
+    n_splits: int = 5,
+    random_state: int = 42,
+    save_oof: bool = False,
+    output_path: Optional[str] = None
+) -> Dict[str, Any]:
+    """
+    Perform comprehensive cross-validation with out-of-fold predictions.
+    
+    Args:
+        file_path: Path to prepared dataset
+        target_col: Target column name
+        model_type: 'random_forest', 'xgboost', 'logistic', 'ridge'
+        task_type: 'classification', 'regression', or 'auto'
+        cv_strategy: 'kfold', 'stratified', or 'timeseries'
+        n_splits: Number of CV folds
+        random_state: Random seed
+        save_oof: Whether to save out-of-fold predictions
+        output_path: Path to save OOF predictions
+        
+    Returns:
+        Dictionary with CV scores, statistics, and OOF predictions
+    """
+    # ⚠️ CRITICAL FIX: Convert n_splits and random_state to int (Gemini/LLMs pass floats)
+    n_splits = int(n_splits)
+    random_state = int(random_state)
+    
+    # Validation
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    # Load data
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    validate_column_exists(df, target_col)
+    
+    # ⚠️ SKIP DATETIME CONVERSION: Already handled by create_time_features() in workflow step 7
+    # The encoded.csv file should already have time features extracted
+    
+    # ⚠️ CRITICAL FIX: Convert Polars to Pandas if needed (for XGBoost compatibility)
+    if hasattr(df, 'to_pandas'):
+        print(f"   🔄 Converting Polars DataFrame to Pandas for XGBoost compatibility...")
+        df = df.to_pandas()
+    
+    # ⚠️ CRITICAL: Drop remaining datetime columns BEFORE NumPy conversion
+    # XGBoost cannot handle Timestamp objects (causes TypeError: float() argument must be a string or a real number, not 'Timestamp')
+    if isinstance(df, pd.DataFrame):
+        datetime_cols = df.select_dtypes(include=['datetime64', 'datetime64[ns]', 'datetime64[ns, UTC]']).columns.tolist()
+        if datetime_cols:
+            print(f"   ⚠️ Dropping {len(datetime_cols)} datetime columns: {datetime_cols}")
+            print(f"   💡 Time features should have been extracted in workflow step 7 (create_time_features)")
+            df = df.drop(columns=datetime_cols)
+        
+        # ⚠️ CRITICAL: Drop any remaining string/object columns (not encoded properly)
+        object_cols = df.select_dtypes(include=['object', 'string']).columns.tolist()
+        object_cols = [col for col in object_cols if col != target_col]
+        if object_cols:
+            print(f"   ⚠️ Dropping {len(object_cols)} string columns that weren't encoded: {object_cols}")
+            print(f"   💡 Categorical encoding should have been done in workflow step 8")
+            df = df.drop(columns=object_cols)
+    
+    # Prepare data - handle both Polars and Pandas
+    if target_col not in df.columns:
+        raise ValueError(f"Target column '{target_col}' not found in dataframe. Available columns: {list(df.columns)}")
+    
+    # Split features and target (works for both Polars and Pandas)
+    if hasattr(df, 'drop'):
+        X = df.drop(columns=[target_col]) if isinstance(df, pd.DataFrame) else df.drop(target_col)
+        y = df[target_col]
+    else:
+        X, y = split_features_target(df, target_col)
+    
+    # Convert to numpy for sklearn compatibility
+    if hasattr(X, 'to_numpy'):
+        X = X.to_numpy()
+        y = y.to_numpy()
+    elif hasattr(X, 'values'):
+        X = X.values
+        y = y.values
+    
+    # Detect task type    # Detect task type
+    if task_type == "auto":
+        unique_values = len(np.unique(y))
+        task_type = "classification" if unique_values < 20 else "regression"
+    
+    # Create model
+    if model_type == "random_forest":
+        if task_type == "classification":
+            model = RandomForestClassifier(n_estimators=100, random_state=random_state)
+        else:
+            model = RandomForestRegressor(n_estimators=100, random_state=random_state)
+    elif model_type == "xgboost":
+        if task_type == "classification":
+            model = XGBClassifier(n_estimators=100, random_state=random_state, use_label_encoder=False, eval_metric='logloss')
+        else:
+            model = XGBRegressor(n_estimators=100, random_state=random_state)
+    elif model_type == "logistic":
+        model = LogisticRegression(max_iter=1000, random_state=random_state)
+    elif model_type == "ridge":
+        model = Ridge(random_state=random_state)
+    else:
+        raise ValueError(f"Unsupported model_type: {model_type}")
+    
+    # Create CV splitter
+    # ⚠️ CRITICAL FIX: Auto-use StratifiedKFold for classification to avoid single-class folds
+    if cv_strategy == "timeseries":
+        cv = TimeSeriesSplit(n_splits=n_splits)
+    elif task_type == "classification":
+        # Always use stratified for classification (unless timeseries)
+        cv = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=random_state)
+        if cv_strategy != "stratified":
+            print(f"   💡 Auto-switching to StratifiedKFold for classification (prevents single-class folds)")
+    else:
+        # Regression: use regular KFold
+        cv = KFold(n_splits=n_splits, shuffle=True, random_state=random_state)
+    
+    print(f"🔄 Performing {n_splits}-fold cross-validation ({cv_strategy})...")
+    
+    # Perform cross-validation with detailed tracking
+    fold_scores = []
+    oof_predictions = np.zeros(len(y))
+    oof_indices = []
+    
+    for fold_idx, (train_idx, val_idx) in enumerate(cv.split(X, y if cv_strategy == "stratified" else None)):
+        X_train_fold, X_val_fold = X[train_idx], X[val_idx]
+        y_train_fold, y_val_fold = y[train_idx], y[val_idx]
+        
+        # Train model
+        model.fit(X_train_fold, y_train_fold)
+        
+        # Predict on validation fold
+        y_pred_fold = model.predict(X_val_fold)
+        
+        # Store OOF predictions
+        oof_predictions[val_idx] = y_pred_fold
+        oof_indices.extend(val_idx.tolist())
+        
+        # Calculate fold metrics
+        if task_type == "classification":
+            fold_score = {
+                'fold': fold_idx + 1,
+                'accuracy': float(accuracy_score(y_val_fold, y_pred_fold)),
+                'f1': float(f1_score(y_val_fold, y_pred_fold, average='weighted', zero_division=0)),
+                'samples': len(val_idx)
+            }
+        else:
+            fold_score = {
+                'fold': fold_idx + 1,
+                'rmse': float(np.sqrt(mean_squared_error(y_val_fold, y_pred_fold))),
+                'r2': float(r2_score(y_val_fold, y_pred_fold)),
+                'samples': len(val_idx)
+            }
+        
+        fold_scores.append(fold_score)
+        print(f"  Fold {fold_idx + 1}: {fold_score}")
+    
+    # Calculate overall OOF metrics
+    if task_type == "classification":
+        oof_metrics = {
+            'accuracy': float(accuracy_score(y, oof_predictions)),
+            'precision': float(precision_score(y, oof_predictions, average='weighted', zero_division=0)),
+            'recall': float(recall_score(y, oof_predictions, average='weighted', zero_division=0)),
+            'f1': float(f1_score(y, oof_predictions, average='weighted', zero_division=0))
+        }
+        mean_fold_metric = np.mean([f['accuracy'] for f in fold_scores])
+        std_fold_metric = np.std([f['accuracy'] for f in fold_scores])
+        metric_name = "accuracy"
+    else:
+        oof_metrics = {
+            'rmse': float(np.sqrt(mean_squared_error(y, oof_predictions))),
+            'mae': float(mean_absolute_error(y, oof_predictions)),
+            'r2': float(r2_score(y, oof_predictions))
+        }
+        mean_fold_metric = np.mean([f['rmse'] for f in fold_scores])
+        std_fold_metric = np.std([f['rmse'] for f in fold_scores])
+        metric_name = "rmse"
+    
+    print(f"\n✅ Overall OOF {metric_name}: {oof_metrics.get(metric_name):.4f} (±{std_fold_metric:.4f})")
+    
+    # Save OOF predictions if requested
+    if save_oof and output_path:
+        oof_df = pl.DataFrame({
+            'index': list(range(len(y))),
+            'true_values': y,
+            'oof_predictions': oof_predictions
+        })
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        oof_df.write_csv(output_path)
+        print(f"💾 OOF predictions saved to: {output_path}")
+    
+    return {
+        'status': 'success',
+        'model_type': model_type,
+        'task_type': task_type,
+        'cv_strategy': cv_strategy,
+        'n_splits': n_splits,
+        'fold_scores': fold_scores,
+        'oof_metrics': oof_metrics,
+        'mean_cv_score': float(mean_fold_metric),
+        'std_cv_score': float(std_fold_metric),
+        'confidence_interval_95': f"[{mean_fold_metric - 1.96 * std_fold_metric:.4f}, {mean_fold_metric + 1.96 * std_fold_metric:.4f}]",
+        'oof_path': output_path if save_oof and output_path else None
+    }