churnkit 0.75.0a1__py3-none-any.whl

customer_retention/core/config/pipeline_config.py

@@ -0,0 +1,117 @@
+from enum import Enum
+from typing import Any, Optional
+
+from pydantic import BaseModel
+
+from .source_config import DataSourceConfig
+
+
+class TimestampStrategy(str, Enum):
+    PRODUCTION = "production"
+    SYNTHETIC_RANDOM = "synthetic_random"
+    SYNTHETIC_INDEX = "synthetic_index"
+    SYNTHETIC_FIXED = "synthetic_fixed"
+    DERIVED = "derived"
+
+
+class DedupStrategy(str, Enum):
+    KEEP_FIRST = "keep_first"
+    KEEP_LAST = "keep_last"
+    KEEP_MOST_COMPLETE = "keep_most_complete"
+
+
+class BronzeConfig(BaseModel):
+    deduplicate: bool = True
+    dedup_strategy: DedupStrategy = DedupStrategy.KEEP_LAST
+    dedup_keys: list[str] = ["custid"]
+    max_missing_pct: float = 0.95
+    min_distinct_values: int = 2
+
+
+class SilverConfig(BaseModel):
+    entity_key: str = "custid"
+    reference_date_column: Optional[str] = None
+    auto_detect_encoding: bool = True
+    auto_detect_scaling: bool = True
+
+
+class GoldConfig(BaseModel):
+    feature_store_catalog: str = "main"
+    feature_store_schema: str = "feature_store"
+    feature_table_name: str = "customer_features"
+    version: str = "v1"
+
+    def get_full_feature_table_name(self) -> str:
+        return f"{self.feature_store_catalog}.{self.feature_store_schema}.{self.feature_table_name}"
+
+
+class ModelingConfig(BaseModel):
+    target_column: str = "retained"
+    positive_class: int = 1
+    test_size: float = 0.2
+    stratify: bool = True
+    primary_metric: str = "average_precision"
+    cost_false_negative: float = 100.0
+    cost_false_positive: float = 10.0
+
+    def get_cost_ratio(self) -> float:
+        return self.cost_false_negative / self.cost_false_positive
+
+
+class ValidationConfig(BaseModel):
+    fail_on_critical: bool = True
+    fail_on_high: bool = False
+    leakage_correlation_threshold: float = 0.90
+    max_overfit_gap: float = 0.15
+
+
+class TemporalConfig(BaseModel):
+    timestamp_strategy: TimestampStrategy = TimestampStrategy.PRODUCTION
+    feature_timestamp_column: Optional[str] = None
+    label_timestamp_column: Optional[str] = None
+    observation_window_days: int = 90
+    synthetic_base_date: str = "2024-01-01"
+    synthetic_range_days: int = 365
+    snapshot_prefix: str = "ml_training_snapshot"
+    enforce_point_in_time: bool = True
+    max_feature_target_correlation: float = 0.90
+    block_future_features: bool = True
+    derive_label_from_feature: bool = False
+    derivation_config: Optional[dict[str, Any]] = None
+
+
+class PathConfig(BaseModel):
+    landing_zone: Optional[str] = None
+    bronze: Optional[str] = None
+    silver: Optional[str] = None
+    gold: Optional[str] = None
+
+
+class PipelineConfig(BaseModel):
+    project_name: str
+    project_description: Optional[str] = None
+    version: str = "1.0.0"
+
+    data_sources: list[DataSourceConfig] = []
+    bronze: BronzeConfig = BronzeConfig()
+    silver: SilverConfig = SilverConfig()
+    gold: GoldConfig = GoldConfig()
+    modeling: ModelingConfig = ModelingConfig()
+    validation: ValidationConfig = ValidationConfig()
+    temporal: TemporalConfig = TemporalConfig()
+    paths: PathConfig = PathConfig()
+
+    def get_source_by_name(self, name: str) -> Optional[DataSourceConfig]:
+        return next((s for s in self.data_sources if s.name == name), None)
+
+    def get_target_source(self) -> Optional[DataSourceConfig]:
+        for source in self.data_sources:
+            if any(c.column_type.value == "target" for c in source.columns):
+                return source
+        return None
+
+    def get_all_feature_columns(self) -> list[str]:
+        feature_cols = []
+        for source in self.data_sources:
+            feature_cols.extend([c.name for c in source.get_feature_columns()])
+        return feature_cols

customer_retention/core/config/source_config.py

@@ -0,0 +1,83 @@
+from enum import Enum
+from typing import Optional
+
+from pydantic import BaseModel, model_validator
+
+from .column_config import ColumnConfig
+
+
+class SourceType(str, Enum):
+    BATCH_FILE = "batch_file"
+    BATCH_TABLE = "batch_table"
+    STREAM = "stream"
+
+
+class FileFormat(str, Enum):
+    CSV = "csv"
+    PARQUET = "parquet"
+    DELTA = "delta"
+    JSON = "json"
+    ORC = "orc"
+    AVRO = "avro"
+
+
+class Grain(str, Enum):
+    CUSTOMER = "customer"
+    TRANSACTION = "transaction"
+    EVENT = "event"
+
+
+class DataSourceConfig(BaseModel):
+    name: str
+    source_type: SourceType
+    primary_key: str
+
+    path: Optional[str] = None
+    file_format: Optional[FileFormat] = None
+
+    catalog: Optional[str] = None
+    schema_name: Optional[str] = None
+    table: Optional[str] = None
+
+    delimiter: str = ","
+    header: bool = True
+    quote_char: str = '"'
+    encoding: str = "utf-8"
+
+    columns: list[ColumnConfig] = []
+    timestamp_column: Optional[str] = None
+    customer_key: Optional[str] = None
+    grain: Optional[Grain] = None
+
+    expected_row_count_min: Optional[int] = None
+    expected_row_count_max: Optional[int] = None
+    expected_columns: Optional[list[str]] = None
+    freshness_sla_hours: Optional[int] = None
+
+    @model_validator(mode='after')
+    def validate_source_requirements(self):
+        if self.source_type == SourceType.BATCH_FILE:
+            if not self.path:
+                raise ValueError("path required for batch_file source_type")
+            if not self.file_format:
+                raise ValueError("file_format required for batch_file source_type")
+        if self.source_type == SourceType.BATCH_TABLE and not self.table:
+            raise ValueError("table required for batch_table source_type")
+        return self
+
+    def get_full_table_name(self) -> str:
+        if self.source_type != SourceType.BATCH_TABLE:
+            raise ValueError("full_table_name only applicable for batch_table")
+        parts = [p for p in [self.catalog, self.schema_name, self.table] if p]
+        return ".".join(parts)
+
+    def get_column_by_name(self, name: str) -> Optional[ColumnConfig]:
+        return next((c for c in self.columns if c.name == name), None)
+
+    def get_feature_columns(self) -> list[ColumnConfig]:
+        return [c for c in self.columns if c.should_be_used_as_feature()]
+
+    def is_cloud_path(self) -> bool:
+        if not self.path:
+            return False
+        return any(self.path.startswith(prefix) for prefix in ["s3://", "abfss://", "gs://", "wasb://", "adl://"])

customer_retention/core/utils/__init__.py

@@ -0,0 +1,28 @@
+from ..components.enums import Severity
+from .leakage import (
+    DEFAULT_THRESHOLDS,
+    TEMPORAL_METADATA_COLUMNS,
+    LeakageThresholds,
+    calculate_class_overlap,
+    classify_correlation,
+    classify_separation,
+    get_valid_feature_columns,
+)
+from .severity import ThresholdConfig, classify_by_thresholds, severity_recommendation
+from .statistics import (
+    compute_chi_square,
+    compute_effect_size,
+    compute_ks_statistic,
+    compute_psi_categorical,
+    compute_psi_from_series,
+    compute_psi_numeric,
+)
+
+__all__ = [
+    "Severity",
+    "compute_psi_numeric", "compute_psi_categorical", "compute_psi_from_series", "compute_ks_statistic", "compute_chi_square",
+    "compute_effect_size",
+    "LeakageThresholds", "classify_correlation", "classify_separation", "calculate_class_overlap", "DEFAULT_THRESHOLDS",
+    "ThresholdConfig", "classify_by_thresholds", "severity_recommendation",
+    "TEMPORAL_METADATA_COLUMNS", "get_valid_feature_columns",
+]

customer_retention/core/utils/leakage.py

@@ -0,0 +1,85 @@
+from dataclasses import dataclass
+from typing import FrozenSet, List, Optional, Set, Tuple
+
+from customer_retention.core.compat import DataFrame, Series
+
+from ..components.enums import Severity
+
+TEMPORAL_METADATA_COLUMNS: FrozenSet[str] = frozenset({
+    "feature_timestamp",
+    "label_timestamp",
+    "label_available_flag",
+    "event_timestamp",
+})
+
+
+def _build_exclusion_set(entity_column: Optional[str], target_column: Optional[str], additional_exclude: Optional[Set[str]]) -> Set[str]:
+    exclude = set(TEMPORAL_METADATA_COLUMNS)
+    if entity_column:
+        exclude.add(entity_column)
+    if target_column:
+        exclude.add(target_column)
+    if additional_exclude:
+        exclude.update(additional_exclude)
+    return exclude
+
+
+def get_valid_feature_columns(
+    df: DataFrame,
+    entity_column: Optional[str] = None,
+    target_column: Optional[str] = None,
+    additional_exclude: Optional[Set[str]] = None,
+) -> List[str]:
+    """Filter DataFrame columns to those valid as model features."""
+    exclude = _build_exclusion_set(entity_column, target_column, additional_exclude)
+    exclude.update(c for c in df.columns if c.startswith("original_"))
+    return [c for c in df.columns if c not in exclude]
+
+
+@dataclass
+class LeakageThresholds:
+    correlation_critical: float = 0.90
+    correlation_high: float = 0.70
+    correlation_medium: float = 0.50
+    separation_critical: float = 0.0
+    separation_high: float = 1.0
+    separation_medium: float = 5.0
+    auc_critical: float = 0.90
+    auc_high: float = 0.80
+
+
+DEFAULT_THRESHOLDS = LeakageThresholds()
+
+
+def classify_correlation(corr: float, thresholds: LeakageThresholds = DEFAULT_THRESHOLDS) -> Tuple[Severity, str]:
+    abs_corr = abs(corr)
+    if abs_corr >= thresholds.correlation_critical:
+        return Severity.CRITICAL, "high_correlation"
+    if abs_corr >= thresholds.correlation_high:
+        return Severity.HIGH, "suspicious_correlation"
+    if abs_corr >= thresholds.correlation_medium:
+        return Severity.MEDIUM, "elevated_correlation"
+    return Severity.INFO, "normal"
+
+
+def classify_separation(overlap_pct: float, thresholds: LeakageThresholds = DEFAULT_THRESHOLDS) -> Tuple[Severity, str]:
+    if overlap_pct <= thresholds.separation_critical:
+        return Severity.CRITICAL, "perfect_separation"
+    if overlap_pct < thresholds.separation_high:
+        return Severity.HIGH, "near_perfect_separation"
+    if overlap_pct < thresholds.separation_medium:
+        return Severity.MEDIUM, "high_separation"
+    return Severity.INFO, "normal"
+
+
+def calculate_class_overlap(feature: Series, target: Series) -> float:
+    class_0, class_1 = feature[target == 0].dropna(), feature[target == 1].dropna()
+    if len(class_0) == 0 or len(class_1) == 0:
+        return 100.0
+    min_0, max_0 = class_0.min(), class_0.max()
+    min_1, max_1 = class_1.min(), class_1.max()
+    total_range = max(max_0, max_1) - min(min_0, min_1)
+    if total_range == 0:
+        return 100.0
+    overlap = max(0, min(max_0, max_1) - max(min_0, min_1))
+    return (overlap / total_range) * 100

customer_retention/core/utils/severity.py

@@ -0,0 +1,53 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from ..components.enums import Severity
+
+
+@dataclass
+class ThresholdConfig:
+    critical: Optional[float] = None
+    high: Optional[float] = None
+    warning: Optional[float] = None
+    medium: Optional[float] = None
+    low: Optional[float] = None
+    ascending: bool = True
+
+
+def classify_by_thresholds(value: float, config: ThresholdConfig) -> Severity:
+    if config.ascending:
+        if config.critical is not None and value >= config.critical:
+            return Severity.CRITICAL
+        if config.high is not None and value >= config.high:
+            return Severity.HIGH
+        if config.warning is not None and value >= config.warning:
+            return Severity.WARNING
+        if config.medium is not None and value >= config.medium:
+            return Severity.MEDIUM
+        if config.low is not None and value >= config.low:
+            return Severity.LOW
+    else:
+        if config.critical is not None and value <= config.critical:
+            return Severity.CRITICAL
+        if config.high is not None and value <= config.high:
+            return Severity.HIGH
+        if config.warning is not None and value <= config.warning:
+            return Severity.WARNING
+        if config.medium is not None and value <= config.medium:
+            return Severity.MEDIUM
+        if config.low is not None and value <= config.low:
+            return Severity.LOW
+    return Severity.INFO
+
+
+def severity_recommendation(severity: Severity, context: str, action_critical: str = "investigate immediately",
+                            action_warning: str = "monitor closely", action_info: str = "no action needed") -> str:
+    recommendations = {
+        Severity.CRITICAL: f"CRITICAL: {context}. {action_critical}.",
+        Severity.HIGH: f"HIGH: {context}. {action_critical}.",
+        Severity.WARNING: f"WARNING: {context}. {action_warning}.",
+        Severity.MEDIUM: f"MEDIUM: {context}. {action_warning}.",
+        Severity.LOW: f"LOW: {context}. {action_info}.",
+        Severity.INFO: f"INFO: {context}. {action_info}.",
+    }
+    return recommendations.get(severity, f"INFO: {context}.")

customer_retention/core/utils/statistics.py

@@ -0,0 +1,90 @@
+from typing import Dict, List, Tuple, Union
+
+import numpy as np
+from scipy import stats
+
+from customer_retention.core.compat import Series
+
+
+def _ensure_array(obj: Union[np.ndarray, List[float]]) -> np.ndarray:
+    return obj if isinstance(obj, np.ndarray) else np.array(obj)
+
+
+def compute_effect_size(group1: Union[np.ndarray, List[float]], group2: Union[np.ndarray, List[float]]) -> Tuple[float, str]:
+    arr1 = _ensure_array(group1)
+    arr2 = _ensure_array(group2)
+    if len(arr1) < 2 or len(arr2) < 2:
+        return 0.0, "Negligible"
+    pooled_std = np.sqrt((np.var(arr1) + np.var(arr2)) / 2)
+    if pooled_std == 0:
+        return 0.0, "Negligible"
+    d = float((np.mean(arr1) - np.mean(arr2)) / pooled_std)
+    abs_d = abs(d)
+    if abs_d >= 0.8:
+        return d, "Large effect"
+    if abs_d >= 0.5:
+        return d, "Medium effect"
+    if abs_d >= 0.2:
+        return d, "Small effect"
+    return d, "Negligible"
+
+
+def compute_psi_numeric(current: Series, reference_hist_edges: List[float], reference_hist_counts: List[int], epsilon: float = 1e-10) -> float:
+    edges = np.array(reference_hist_edges)
+    baseline_counts = np.array(reference_hist_counts)
+    current_counts, _ = np.histogram(current.dropna(), bins=edges)
+    baseline_prop = baseline_counts / baseline_counts.sum()
+    current_prop = current_counts / current_counts.sum() if current_counts.sum() > 0 else np.zeros_like(current_counts, dtype=float)
+    baseline_prop = np.maximum(baseline_prop, epsilon)
+    current_prop = np.maximum(current_prop, epsilon)
+    return float(np.sum((current_prop - baseline_prop) * np.log(current_prop / baseline_prop)))
+
+
+def _is_categorical_dtype(dtype) -> bool:
+    return dtype in ['object', 'category', 'bool']
+
+
+def compute_psi_from_series(reference: Series, current: Series, n_bins: int = 10, epsilon: float = 1e-10) -> float:
+    ref_clean, curr_clean = reference.dropna(), current.dropna()
+    if _is_categorical_dtype(ref_clean.dtype) or _is_categorical_dtype(curr_clean.dtype):
+        return compute_psi_categorical(ref_clean, curr_clean, epsilon)
+    min_val = min(ref_clean.min(), curr_clean.min())
+    max_val = max(ref_clean.max(), curr_clean.max())
+    bins = np.linspace(min_val, max_val, n_bins + 1)
+    ref_hist, _ = np.histogram(ref_clean, bins=bins)
+    curr_hist, _ = np.histogram(curr_clean, bins=bins)
+    ref_pct = ref_hist / len(ref_clean) + epsilon
+    curr_pct = curr_hist / len(curr_clean) + epsilon if len(curr_clean) > 0 else np.full_like(ref_hist, epsilon, dtype=float)
+    return float(np.sum((curr_pct - ref_pct) * np.log(curr_pct / ref_pct)))
+
+
+def compute_psi_categorical(reference: Series, current: Series, epsilon: float = 1e-10) -> float:
+    ref_counts = reference.value_counts(normalize=True)
+    curr_counts = current.value_counts(normalize=True)
+    all_categories = set(ref_counts.index) | set(curr_counts.index)
+    psi = 0.0
+    for cat in all_categories:
+        ref_pct = ref_counts.get(cat, epsilon)
+        curr_pct = curr_counts.get(cat, epsilon)
+        psi += (curr_pct - ref_pct) * np.log((curr_pct + epsilon) / (ref_pct + epsilon))
+    return float(psi)
+
+
+def compute_ks_statistic(reference: Series, current: Series) -> Tuple[float, float]:
+    ref_clean, curr_clean = reference.dropna(), current.dropna()
+    statistic, pvalue = stats.ks_2samp(ref_clean, curr_clean)
+    return float(statistic), float(pvalue)
+
+
+def compute_chi_square(current: Series, baseline_proportions: Dict[str, float]) -> Tuple[float, float]:
+    current_counts = current.value_counts()
+    all_categories = sorted(set(list(current_counts.index) + list(baseline_proportions.keys())))
+    observed, expected = [], []
+    total_current = len(current)
+    for cat in all_categories:
+        observed.append(current_counts.get(cat, 0))
+        expected.append(max(baseline_proportions.get(cat, 0) * total_current, 1e-10))
+    expected_arr = np.array(expected)
+    expected_arr = expected_arr * (sum(observed) / sum(expected_arr)) if sum(expected_arr) > 0 else expected_arr
+    chi_square, pvalue = stats.chisquare(observed, expected_arr)
+    return float(chi_square), float(pvalue)

customer_retention/generators/notebook_generator/__init__.py

@@ -0,0 +1,167 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, List, Optional
+
+from .base import NotebookGenerator, NotebookStage
+
+if TYPE_CHECKING:
+    from customer_retention.analysis.auto_explorer import ExplorationFindings
+from .cell_builder import CellBuilder
+from .config import FeatureStoreConfig, MLflowConfig, NotebookConfig, OutputFormat, Platform
+from .databricks_generator import DatabricksNotebookGenerator
+from .local_generator import LocalNotebookGenerator
+from .project_init import ProjectInitializer, initialize_project
+from .runner import (
+    NotebookRunner,
+    NotebookValidationResult,
+    ScriptRunner,
+    ValidationReport,
+    validate_generated_notebooks,
+)
+from .script_generator import DatabricksScriptGenerator, LocalScriptGenerator, ScriptGenerator
+
+
+@dataclass
+class GenerationResult:
+    platform: Platform
+    notebook_paths: List[str]
+    validation_report: Optional[ValidationReport] = None
+
+    @property
+    def all_valid(self) -> bool:
+        return self.validation_report.all_passed if self.validation_report else True
+
+
+def generate_orchestration_notebooks(
+    findings_path: Optional[str] = None,
+    output_dir: str = "./generated_pipelines",
+    platforms: Optional[List[Platform]] = None,
+    config: Optional[NotebookConfig] = None,
+    validate: bool = False,
+) -> Dict[Platform, List[str]]:
+    if platforms is None:
+        platforms = [Platform.LOCAL, Platform.DATABRICKS]
+    if config is None:
+        config = NotebookConfig()
+
+    findings = None
+    if findings_path:
+        from customer_retention.analysis.auto_explorer import ExplorationFindings
+        findings = ExplorationFindings.load(findings_path)
+
+    results = {}
+    for platform in platforms:
+        generator = create_notebook_generator(platform, findings, config)
+        platform_dir = str(Path(output_dir) / platform.value)
+        saved_paths = generator.save_all(platform_dir)
+        results[platform] = saved_paths
+
+    return results
+
+
+def generate_and_validate_notebooks(
+    findings_path: Optional[str] = None,
+    output_dir: str = "./generated_pipelines",
+    platforms: Optional[List[Platform]] = None,
+    config: Optional[NotebookConfig] = None,
+) -> Dict[Platform, GenerationResult]:
+    if platforms is None:
+        platforms = [Platform.LOCAL, Platform.DATABRICKS]
+    if config is None:
+        config = NotebookConfig()
+
+    findings = None
+    if findings_path:
+        from customer_retention.analysis.auto_explorer import ExplorationFindings
+        findings = ExplorationFindings.load(findings_path)
+
+    results = {}
+    runner = NotebookRunner(dry_run=True)
+
+    for platform in platforms:
+        generator = create_notebook_generator(platform, findings, config)
+        platform_dir = str(Path(output_dir) / platform.value)
+        saved_paths = generator.save_all(platform_dir)
+        validation_report = runner.validate_sequence(platform_dir, platform.value)
+        results[platform] = GenerationResult(platform, saved_paths, validation_report)
+        save_validation_report(platform_dir, validation_report)
+
+    return results
+
+
+def save_validation_report(output_dir: str, report: ValidationReport) -> str:
+    report_path = Path(output_dir) / "VALIDATION_REPORT.md"
+    report_path.write_text(report.to_markdown())
+    return str(report_path)
+
+
+def create_notebook_generator(
+    platform: Platform,
+    findings: Optional["ExplorationFindings"] = None,
+    config: Optional[NotebookConfig] = None,
+) -> NotebookGenerator:
+    if config is None:
+        config = NotebookConfig()
+
+    if platform == Platform.LOCAL:
+        return LocalNotebookGenerator(config, findings)
+    elif platform == Platform.DATABRICKS:
+        return DatabricksNotebookGenerator(config, findings)
+    else:
+        raise ValueError(f"Unsupported platform: {platform}")
+
+
+def create_script_generator(
+    platform: Platform,
+    findings: Optional["ExplorationFindings"] = None,
+    config: Optional[NotebookConfig] = None,
+) -> ScriptGenerator:
+    if config is None:
+        config = NotebookConfig()
+
+    if platform == Platform.LOCAL:
+        return LocalScriptGenerator(config, findings)
+    elif platform == Platform.DATABRICKS:
+        return DatabricksScriptGenerator(config, findings)
+    else:
+        raise ValueError(f"Unsupported platform: {platform}")
+
+
+def generate_orchestration_scripts(
+    findings_path: Optional[str] = None,
+    output_dir: str = "./generated_pipelines/scripts",
+    platforms: Optional[List[Platform]] = None,
+    config: Optional[NotebookConfig] = None,
+) -> Dict[Platform, List[str]]:
+    if platforms is None:
+        platforms = [Platform.LOCAL, Platform.DATABRICKS]
+    if config is None:
+        config = NotebookConfig()
+
+    findings = None
+    if findings_path:
+        from customer_retention.analysis.auto_explorer import ExplorationFindings
+        findings = ExplorationFindings.load(findings_path)
+
+    results = {}
+    for platform in platforms:
+        generator = create_script_generator(platform, findings, config)
+        platform_dir = str(Path(output_dir) / platform.value)
+        saved_paths = generator.save_all(platform_dir)
+        results[platform] = saved_paths
+
+    return results
+
+
+__all__ = [
+    "NotebookGenerator", "NotebookStage", "NotebookConfig", "Platform",
+    "MLflowConfig", "FeatureStoreConfig", "CellBuilder", "OutputFormat",
+    "LocalNotebookGenerator", "DatabricksNotebookGenerator",
+    "NotebookRunner", "NotebookValidationResult", "ValidationReport", "ScriptRunner",
+    "GenerationResult", "generate_orchestration_notebooks",
+    "generate_and_validate_notebooks", "create_notebook_generator",
+    "validate_generated_notebooks", "save_validation_report",
+    "ScriptGenerator", "LocalScriptGenerator", "DatabricksScriptGenerator",
+    "create_script_generator", "generate_orchestration_scripts",
+    "ProjectInitializer", "initialize_project",
+]

customer_retention/generators/notebook_generator/base.py

@@ -0,0 +1,55 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, List, Optional
+
+import nbformat
+
+from .config import NotebookConfig
+
+if TYPE_CHECKING:
+    from customer_retention.analysis.auto_explorer import ExplorationFindings
+
+
+class NotebookStage(str, Enum):
+    INGESTION = "01_ingestion"
+    PROFILING = "02_profiling"
+    CLEANING = "03_cleaning"
+    TRANSFORMATION = "04_transformation"
+    FEATURE_ENGINEERING = "05_feature_engineering"
+    FEATURE_SELECTION = "06_feature_selection"
+    MODEL_TRAINING = "07_model_training"
+    DEPLOYMENT = "08_deployment"
+    MONITORING = "09_monitoring"
+    BATCH_INFERENCE = "10_batch_inference"
+    FEATURE_STORE = "11_feature_store"
+
+
+class NotebookGenerator(ABC):
+    def __init__(self, config: NotebookConfig, findings: Optional["ExplorationFindings"]):
+        self.config = config
+        self.findings = findings
+
+    @abstractmethod
+    def generate_stage(self, stage: NotebookStage) -> nbformat.NotebookNode:
+        pass
+
+    @property
+    def available_stages(self) -> List[NotebookStage]:
+        if hasattr(self, "stage_generators"):
+            return list(self.stage_generators.keys())
+        return list(NotebookStage)
+
+    def generate_all(self) -> Dict[NotebookStage, nbformat.NotebookNode]:
+        return {stage: self.generate_stage(stage) for stage in self.available_stages}
+
+    def save_all(self, output_dir: str) -> List[str]:
+        output_path = Path(output_dir)
+        output_path.mkdir(parents=True, exist_ok=True)
+        saved_paths = []
+        for stage, notebook in self.generate_all().items():
+            file_path = output_path / f"{stage.value}.ipynb"
+            with open(file_path, "w", encoding="utf-8") as f:
+                nbformat.write(notebook, f)
+            saved_paths.append(str(file_path))
+        return saved_paths