feat(core): introduce model abstraction layer with BaseModelWrapper

- Added BaseModelWrapper and LinearModelWrapper to support model-agnostic diagnostics
- Centralized model fitting using get_model_wrapper(), fit once for all checks
- Updated linearity check to use model wrapper and fallback gracefully
- Introduced --model-type CLI flag for model selection (currently supports 'linear')
- Enabled filtering of assumption checks by applicable model types via @register_assumption
- Printed model metadata in console report output (e.g. 'Model Type: Linear Regression')
- Reordered function parameters for clarity and consistency across checks
- Removed redundant classmethod import and updated docstrings for clarity

Author: W-Thurston

app/core/dispatcher.py

@@ -1,5 +1,5 @@
 # app/core/dispatcher.py
-from typing import Dict
+from typing import Dict, Tuple
 
 import pandas as pd
 
@@ -9,6 +9,8 @@
 from app.core import normality  # noqa: F401
 from app.core.registry import ASSUMPTION_CHECKS
 from app.core.types import AssumptionResult
+from app.models.base_model_wrapper import BaseModelWrapper
+from app.models.utils import get_model_wrapper
 
 __all__ = ["check_assumption", "run_all_checks"]
 
@@ -21,7 +23,7 @@ def check_assumption(
 
     Args:
         name (str): assumption name
-        X (pd.Series): Predictor (1D)
+        X (pd.Series or pd.DataFrame): Predictor values (1D or multivariate)
         y (pd.Series): Response (1D)
         return_plot (bool, optional): Whether to return base64-encoded
             PNG of the plot. Defaults to False.
@@ -42,13 +44,13 @@ def check_assumption(
 
 
 def run_all_checks(
-    X: pd.Series, y: pd.Series, return_plot: bool = False
-) -> Dict[str, AssumptionResult]:
+    X: pd.Series, y: pd.Series, model_type=None, return_plot: bool = False
+) -> Tuple[Dict[str, AssumptionResult], BaseModelWrapper]:
     """
     Run all registered assumption checks and return a dictionary of results.
 
     Args:
-        X (pd.Series): Predictor (1D)
+        X (pd.Series or pd.DataFrame): Predictor values (1D or multivariate)
         y (pd.Series): Response (1D)
         return_plot (bool, optional): Whether to return base64-encoded
             PNG of the plot. Defaults to False.
@@ -62,6 +64,10 @@ def run_all_checks(
     if isinstance(X, pd.Series):
         X = X.to_frame()
 
+    model_wrapper = get_model_wrapper(model_type, X, y)
+
     for name, func in ASSUMPTION_CHECKS.items():
-        results[name] = func(X, y, return_plot)
-    return results
+        if model_type not in getattr(func, "_model_types", ["linear"]):
+            continue
+        results[name] = func(X, y, model_wrapper=model_wrapper, return_plot=return_plot)
+    return results, model_wrapper

app/core/linearity.py

@@ -9,7 +9,6 @@
 
 import matplotlib.pyplot as plt
 import pandas as pd
-from sklearn.linear_model import LinearRegression
 from sklearn.metrics import r2_score
 
 from app.config import LINEARITY_R2_THRESHOLD, R2_SEVERITY_THRESHOLDS
@@ -20,9 +19,9 @@
 __all__ = ["check_linearity"]
 
 
-@register_assumption("linearity")
+@register_assumption("linearity", model_types=["linear"])
 def check_linearity(
-    X: pd.Series, y: pd.Series, return_plot: bool = False
+    X: pd.Series, y: pd.Series, return_plot: bool = False, model_wrapper=None
 ) -> AssumptionResult:
     """
     Check linearity assumption using:
@@ -59,11 +58,15 @@ def check_linearity(
             )
         X = X.iloc[:, 0]  # Convert to Series
 
+    # Guard for if model_wrapper is None
+    if model_wrapper is None:
+        from app.models.utils import get_model_wrapper
+
+        model_wrapper = get_model_wrapper("linear", X, y)
+
     # Fit simple linear model to input data
-    X_reshaped = X.values.reshape(-1, 1)
-    model = LinearRegression().fit(X_reshaped, y)
-    y_pred = model.predict(X_reshaped)
-    residuals = y - y_pred
+    residuals = model_wrapper.residuals()
+    y_pred = model_wrapper.fitted()
 
     # Coefficient of determination (R²) measures goodness of fit
     r2 = r2_score(y, y_pred)

app/core/registry.py

@@ -14,7 +14,9 @@
 AssumptionCheck = Callable[[pd.Series, pd.Series, bool], AssumptionResult]
 
 
-def register_assumption(name: str) -> Callable[[AssumptionCheck], AssumptionCheck]:
+def register_assumption(
+    name: str, model_types: list = ["linear"]
+) -> Callable[[AssumptionCheck], AssumptionCheck]:
     """
     Decorator to register an assumption check function under a given name.
 
@@ -26,6 +28,8 @@ def register_assumption(name: str) -> Callable[[AssumptionCheck], AssumptionChec
     """
 
     def decorator(func: AssumptionCheck) -> AssumptionCheck:
+        func._assumption_name = name
+        func._model_types = model_types
         ASSUMPTION_CHECKS[name] = func
         return func
 

app/models/base_model_wrapper.py

@@ -0,0 +1,22 @@
+from abc import ABC, abstractmethod
+
+
+class BaseModelWrapper(ABC):
+    def __init__(self, X, y):
+        self.X = X
+        self.y = y
+
+    @abstractmethod
+    def fit(self): ...
+
+    @abstractmethod
+    def predict(self): ...
+
+    @abstractmethod
+    def residuals(self): ...
+
+    @abstractmethod
+    def fitted(self): ...
+
+    def summary(self):
+        return {}

app/models/linear_model_wrapper.py

@@ -0,0 +1,21 @@
+import statsmodels.api as sm
+
+from app.models.base_model_wrapper import BaseModelWrapper
+
+
+class LinearModelWrapper(BaseModelWrapper):
+    def fit(self):
+        self.model = sm.OLS(self.y, sm.add_constant(self.X)).fit()
+        return self
+
+    def predict(self):
+        return self.model.predict(sm.add_constant(self.X))
+
+    def residuals(self):
+        return self.model.resid
+
+    def fitted(self):
+        return self.model.fittedvalues
+
+    def summary(self):
+        return {"model_type": "Linear Regression", "r_squared": self.model.rsquared}

app/models/utils.py

@@ -0,0 +1,11 @@
+from app.models.base_model_wrapper import BaseModelWrapper
+from app.models.linear_model_wrapper import LinearModelWrapper
+
+
+def get_model_wrapper(model_type: str, X, y) -> BaseModelWrapper:
+    if model_type == "linear":
+        return LinearModelWrapper(X, y).fit()
+    elif model_type == "PLACEHOLDER":
+        ...
+    else:
+        raise ValueError(f"Unsupported model type: {model_type}")

app/report.py

@@ -15,6 +15,7 @@
 def generate_report(
     X,
     y,
+    model_type=None,
     return_plot: bool = False,
     output_format: str = "console",
     verbose: bool = False,
@@ -23,7 +24,7 @@ def generate_report(
     Generate an assumption diagnostic report using the registered checks.
 
     Args:
-        X (pd.Series): Predictor values.
+        X (pd.Series or pd.DataFrame): Predictor values (1D or multivariate)
         y (pd.Series): Response values.
         return_plot (bool, optional): Include base64-encoded plots in results.
         output_format (str): 'console', 'json', or 'markdown'.
@@ -32,10 +33,12 @@ def generate_report(
     Raises:
         ValueError: If the output_format is not recognized.
     """
-    results = run_all_checks(X, y, return_plot=return_plot)
+    results, model_wrapper = run_all_checks(
+        X, y, model_type=model_type, return_plot=return_plot
+    )
 
     if output_format == "console":
-        print_console_report(results, verbose=verbose)
+        print_console_report(results, model_wrapper=model_wrapper, verbose=verbose)
     elif output_format == "json":
         export_to_json(results)
     elif output_format == "markdown":
@@ -44,7 +47,7 @@ def generate_report(
         raise ValueError("Unsupported output format")
 
 
-def print_console_report(results, verbose: bool = False):
+def print_console_report(results, model_wrapper, verbose: bool = False):
     """
     Print a structured Rich panel for each assumption result.
 
@@ -54,6 +57,11 @@ def print_console_report(results, verbose: bool = False):
     """
     console = Console()
     console.rule("[bold yellow]Assumption Check Report")
+
+    # Print mdoel metadata
+    model_info = model_wrapper.summary().get("model_type", "Unknown")
+    console.print(f"[bold cyan]Model Type:[/bold cyan] {model_info}")
+
     for name, result in results.items():
 
         # Determine pass/fail icon and panel title
@@ -198,13 +206,22 @@ def export_to_markdown(results, filename: str = None):
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Run regression assumption checks.")
+
+    parser = argparse.ArgumentParser(
+        description="Run statistical assumption checks for supervised models."
+    )
     parser.add_argument(
         "--data",
         choices=list_simulations().keys(),
         default="linear",
         help="Which simulated dataset to run assumption checks on.",
     )
+    parser.add_argument(
+        "--model-type",
+        choices=["linear"],
+        default="linear",
+        help="Which model to fit for diagnostics.",
+    )
     parser.add_argument(
         "--format",
         choices=["console", "json", "markdown"],
@@ -222,6 +239,10 @@ def export_to_markdown(results, filename: str = None):
 
     args = parser.parse_args()
 
+    diagnostic_context = {
+        "model_type": args.model_type,
+    }
+
     data_func = list_simulations()[args.data]
     df = data_func(seed=42)
 
@@ -230,5 +251,10 @@ def export_to_markdown(results, filename: str = None):
     y = df["y"]
 
     generate_report(
-        X, y, return_plot=args.plot, output_format=args.format, verbose=args.verbose
+        X,
+        y,
+        model_type=args.model_type,
+        return_plot=args.plot,
+        output_format=args.format,
+        verbose=args.verbose,
     )

tests/test_dispatcher.py

@@ -20,7 +20,7 @@ def test_dispatch_all_assumptions():
     Test dispatcher's run_all_checks().
     """
     df = simulated_data.generate_linear_data(n_samples=300, seed=42)
-    results = dispatcher.run_all_checks(df["x"], df["y"])
+    results, _ = dispatcher.run_all_checks(df["x"], df["y"], model_type="linear")
     assert "linearity" in results
     assert "homoscedasticity" in results
     assert results["linearity"].passed