feat: add probability-aware evaluation metrics and reporting

plexe/CODE_INDEX.md

@@ -1,6 +1,6 @@
 # Code Index: plexe
 
-> Generated on 2026-02-26 19:02:04
+> Generated on 2026-02-27 14:44:45
 
 Code structure and public interface documentation for the **plexe** package.
 
@@ -154,6 +154,7 @@ Configuration for plexe.
 - `parse_otel_headers_from_env(self) -> 'Config'` - Parse OTEL_EXPORTER_OTLP_HEADERS (comma-separated key=value pairs).
 
 **Functions:**
+- `detect_installed_frameworks() -> list[str]` - Detect which ML frameworks are installed and importable.
 - `get_routing_for_model(config: RoutingConfig | None, model_id: str) -> tuple[str | None, dict[str, str]]` - Get routing configuration for a specific model ID.
 - `setup_logging(config: Config) -> logging.Logger` - Configure logging for the plexe package.
 - `setup_litellm(config: Config) -> None` - Configure LiteLLM global settings.
@@ -222,6 +223,7 @@ Helper functions for workflow.
 - `select_viable_model_types(data_layout: DataLayout, selected_frameworks: list[str] | None) -> list[str]` - Select viable model types using three-tier filtering.
 - `evaluate_on_sample(spark: SparkSession, sample_uri: str, model_artifacts_path: Path, model_type: str, metric: str, target_columns: list[str], group_column: str | None) -> float` - Evaluate model on sample (fast).
 - `compute_metric_hardcoded(y_true, y_pred, metric_name: str) -> float` - Compute metric using hardcoded sklearn implementations.
+- `compute_metric_proba(y_true, y_proba, metric_name: str) -> float` - Compute metrics that require probability estimates.
 - `compute_metric(y_true, y_pred, metric_name: str, group_ids) -> float` - Compute metric value.
 
 ---
@@ -442,6 +444,7 @@ Standard CatBoost predictor - NO Plexe dependencies.
 **`CatBoostPredictor`** - Standalone CatBoost predictor.
 - `__init__(self, model_dir: str)`
 - `predict(self, x: pd.DataFrame) -> pd.DataFrame` - Make predictions on input DataFrame.
+- `predict_proba(self, x: pd.DataFrame) -> pd.DataFrame` - Predict per-class probabilities on input DataFrame.
 
 ---
 ## `templates/inference/keras_predictor.py`
@@ -450,6 +453,7 @@ Standard Keras predictor - NO Plexe dependencies.
 **`KerasPredictor`** - Standalone Keras predictor.
 - `__init__(self, model_dir: str)`
 - `predict(self, x: pd.DataFrame) -> pd.DataFrame` - Make predictions on input DataFrame.
+- `predict_proba(self, x: pd.DataFrame) -> pd.DataFrame` - Predict per-class probabilities on input DataFrame.
 
 ---
 ## `templates/inference/lightgbm_predictor.py`
@@ -458,6 +462,7 @@ Standard LightGBM predictor - NO Plexe dependencies.
 **`LightGBMPredictor`** - Standalone LightGBM predictor.
 - `__init__(self, model_dir: str)`
 - `predict(self, x: pd.DataFrame) -> pd.DataFrame` - Make predictions on input DataFrame.
+- `predict_proba(self, x: pd.DataFrame) -> pd.DataFrame` - Predict per-class probabilities on input DataFrame.
 
 ---
 ## `templates/inference/pytorch_predictor.py`
@@ -466,6 +471,7 @@ Standard PyTorch predictor - NO Plexe dependencies.
 **`PyTorchPredictor`** - Standalone PyTorch predictor.
 - `__init__(self, model_dir: str)`
 - `predict(self, x: pd.DataFrame) -> pd.DataFrame` - Make predictions on input DataFrame.
+- `predict_proba(self, x: pd.DataFrame) -> pd.DataFrame` - Predict per-class probabilities on input DataFrame.
 
 ---
 ## `templates/inference/xgboost_predictor.py`
@@ -474,6 +480,7 @@ Standard XGBoost predictor - NO Plexe dependencies.
 **`XGBoostPredictor`** - Standalone XGBoost predictor.
 - `__init__(self, model_dir: str)`
 - `predict(self, x: pd.DataFrame) -> pd.DataFrame` - Make predictions on input DataFrame.
+- `predict_proba(self, x: pd.DataFrame) -> pd.DataFrame` - Predict per-class probabilities on input DataFrame.
 
 ---
 ## `templates/training/train_catboost.py`

plexe/agents/model_evaluator.py

@@ -9,6 +9,7 @@
 import logging
 from typing import Any, TYPE_CHECKING
 
+import numpy as np
 import pandas as pd
 from plexe.utils.litellm_wrapper import PlexeLiteLLMModel
 from smolagents import CodeAgent
@@ -67,6 +68,63 @@ def __init__(
         self.config = config
         self.llm_model = config.evaluation_llm
 
+    @staticmethod
+    def _is_classification_task(task_type: str | None) -> bool:
+        """Check whether a task type string represents classification."""
+        # TODO(task-type-enum): Switch this helper to the canonical TaskType enum once merged.
+        if not task_type:
+            return False
+        normalized = task_type.strip().lower()
+        return normalized in {"classification", "binary_classification", "multiclass_classification"}
+
+    def _can_run_probability_analysis(self, predictor: Any, test_sample_df: pd.DataFrame) -> tuple[bool, str]:
+        """Fail-closed gate for probability analysis phase."""
+        task_type = (self.context.task_analysis or {}).get("task_type")
+        if not self._is_classification_task(task_type):
+            return False, f"task_type '{task_type}' is not classification"
+
+        proba_fn = getattr(predictor, "predict_proba", None)
+        if not callable(proba_fn):
+            return False, f"predictor '{type(predictor).__name__}' has no callable predict_proba()"
+
+        if not self.context.output_targets:
+            return False, "output_targets is empty"
+        target_col = self.context.output_targets[0]
+        if target_col not in test_sample_df.columns:
+            return False, f"target column '{target_col}' missing from test sample"
+
+        feature_cols = [col for col in test_sample_df.columns if col not in self.context.output_targets]
+        if not feature_cols:
+            return False, "no feature columns available for probability analysis"
+
+        smoke_df = test_sample_df.head(min(16, len(test_sample_df)))
+        if smoke_df.empty:
+            return False, "test sample is empty"
+
+        try:
+            proba_df = proba_fn(smoke_df[feature_cols])
+        except Exception as exc:
+            return False, f"predict_proba smoke test failed: {exc}"
+
+        if not isinstance(proba_df, pd.DataFrame):
+            return False, "predict_proba output is not a pandas DataFrame"
+        if len(proba_df) != len(smoke_df):
+            return False, "predict_proba output row count does not match input rows"
+        if proba_df.shape[1] < 2:
+            return False, "predict_proba output must have at least two probability columns"
+        if not all(str(col).startswith("proba_") for col in proba_df.columns):
+            return False, "predict_proba columns must start with 'proba_'"
+
+        proba_vals = proba_df.to_numpy(dtype=float)
+        if not np.isfinite(proba_vals).all():
+            return False, "predict_proba output contains NaN/Inf"
+        if (proba_vals < 0).any() or (proba_vals > 1).any():
+            return False, "predict_proba output contains values outside [0, 1]"
+        if not np.allclose(proba_vals.sum(axis=1), 1.0, atol=1e-3):
+            return False, "predict_proba rows do not sum to 1"
+
+        return True, "gate checks passed"
+
     def _build_agent(self, phase_name: str, phase_prompt: str, tools: list) -> CodeAgent:
         """
         Build CodeAgent for a specific evaluation phase.
@@ -111,7 +169,9 @@ def _build_agent(self, phase_name: str, phase_prompt: str, tools: list) -> CodeA
                 f"{phase_prompt}\\n\\n"
                 "CRITICAL: Always register your results using the specified tool.\\n"
                 "CRITICAL: Provide interpretation, not just numbers.\\n"
-                "IMPORTANT: Do not create plots or visualizations (headless environment).\\n"
+                "IMPORTANT: In the probability analysis phase, if the predictor has predict_proba(), compute ROC "
+                "curve data, calibration curve, and Brier score. Use matplotlib with Agg backend to generate plots, "
+                "encode as base64 PNG, and include in the visualizations field.\\n"
             ),
             model=PlexeLiteLLMModel(
                 model_id=self.llm_model,
@@ -122,7 +182,18 @@ def _build_agent(self, phase_name: str, phase_prompt: str, tools: list) -> CodeA
             tools=tools,
             add_base_tools=False,
             additional_authorized_imports=self.config.allowed_base_imports
-            + ["collections", "pandas", "pandas.*", "numpy", "numpy.*", "sklearn", "sklearn.*"],
+            + [
+                "collections",
+                "pandas",
+                "pandas.*",
+                "numpy",
+                "numpy.*",
+                "sklearn",
+                "sklearn.*",
+                "matplotlib",
+                "matplotlib.*",
+                "base64",
+            ],
             max_steps=20,
             planning_interval=5,
         )
@@ -213,6 +284,23 @@ def run(
             logger.error("Core Metrics phase failed - cannot continue evaluation")
             return None
 
+        # Phase 1.5: Probability Analysis
+        should_run_proba_phase, skip_reason = self._can_run_probability_analysis(predictor, test_sample_df)
+        if should_run_proba_phase:
+            core_metrics = self.context.scratch.get("_core_metrics_report")
+            proba_args = {**additional_args, "core_metrics_report": core_metrics}
+            success = self._run_phase(
+                phase_name="ProbabilityAnalysis",
+                phase_prompt=self._get_phase_probability_prompt(self.context.intent),
+                tools=[get_register_core_metrics_tool(self.context)],
+                additional_args=proba_args,
+                registry_key="_core_metrics_report",
+            )
+            if not success:
+                logger.warning("Probability analysis phase failed - continuing with partial evaluation")
+        else:
+            logger.info(f"Skipping ProbabilityAnalysis phase: {skip_reason}")
+
         # Phase 2: Error Analysis
         success = self._run_phase(
             phase_name="ErrorAnalysis",
@@ -335,12 +423,55 @@ def _get_phase_1_prompt(task: str, primary_metric_name: str) -> str:
             f"    all_metrics={{...}},\\n"
             f"    statistical_notes='Your interpretation',\\n"
             f"    metric_confidence_intervals=None,  # Optional\\n"
-            f"    visualizations=None  # Optional (headless)\\n"
+            f"    visualizations=None  # Optional\\n"
             f")\\n\\n"
             f"After successful registration, call final_answer('Phase 1 complete').\\n\\n"
             f"IMPORTANT: Focus on rigorous computation and thoughtful interpretation."
         )
 
+    @staticmethod
+    def _get_phase_probability_prompt(task: str) -> str:
+        return (
+            f"PHASE 1.5: PROBABILITY ANALYSIS\\n\\n"
+            f"Task Context: {task}\\n\\n"
+            f"Your mission: Use probability outputs to assess model calibration and discrimination.\\n\\n"
+            f"You have access to core_metrics_report from Phase 1 (reuse its fields).\\n"
+            f"If the task is classification AND predictor has predict_proba(), compute:\\n"
+            f"1. ROC curve data (fpr, tpr, thresholds)\\n"
+            f"2. Calibration curve data (prob_true, prob_pred, n_bins)\\n"
+            f"3. Brier score\\n"
+            f"4. ROC AUC computed from probabilities\\n"
+            f"5. ROC + calibration plots (matplotlib Agg backend), base64-encode PNGs\\n"
+            f"   (set matplotlib.use('Agg') before importing pyplot)\\n\\n"
+            f"Data prep reminder:\\n"
+            f"- feature_cols = [col for col in test_sample_df.columns if col not in output_targets]\\n"
+            f"- X_test = test_sample_df[feature_cols]\\n"
+            f"- y_true = test_sample_df[output_targets[0]].values\\n"
+            f"- y_proba_df = predictor.predict_proba(X_test)\\n\\n"
+            f"Recommended data format:\\n"
+            f"- roc_curve_data={{'fpr': [...], 'tpr': [...], 'thresholds': [...]}}\\n"
+            f"- calibration_data={{'prob_true': [...], 'prob_pred': [...], 'n_bins': 10}}\\n\\n"
+            f"If task is regression or predict_proba is unavailable, keep probability fields as None.\\n"
+            f"Always merge new visualizations into existing core_metrics_report.visualizations if present.\\n\\n"
+            f"Register using:\\n"
+            f"register_core_metrics_report(\\n"
+            f"    task_type=core_metrics_report.task_type,\\n"
+            f"    primary_metric_name=core_metrics_report.primary_metric_name,\\n"
+            f"    primary_metric_value=core_metrics_report.primary_metric_value,\\n"
+            f"    primary_metric_ci_lower=core_metrics_report.primary_metric_ci_lower,\\n"
+            f"    primary_metric_ci_upper=core_metrics_report.primary_metric_ci_upper,\\n"
+            f"    all_metrics=core_metrics_report.all_metrics,\\n"
+            f"    statistical_notes=core_metrics_report.statistical_notes,\\n"
+            f"    metric_confidence_intervals=core_metrics_report.metric_confidence_intervals,\\n"
+            f"    visualizations=merged_visualizations,\\n"
+            f"    roc_auc_from_proba=...,\\n"
+            f"    brier_score=...,\\n"
+            f"    calibration_data={{...}},\\n"
+            f"    roc_curve_data={{...}}\\n"
+            f")\\n\\n"
+            f"After successful registration, call final_answer('Probability analysis complete').\\n"
+        )
+
     @staticmethod
     def _get_phase_2_prompt(task: str) -> str:
         return (

plexe/config.py

@@ -122,6 +122,7 @@ class StandardMetric(str, Enum):
     ROC_AUC_OVR = "roc_auc_ovr"
     ROC_AUC_OVO = "roc_auc_ovo"
     LOG_LOSS = "log_loss"
+    BRIER_SCORE = "brier_score"
 
     # Classification - Other
     MATTHEWS_CORRCOEF = "matthews_corrcoef"