Spaces:
Paused
Paused
| import pandas as pd | |
| import json | |
| import logging | |
| import time | |
| from pathlib import Path | |
| from src.ai.classification.xgboost_classifier import XGBoostClassifier | |
| from src.ai.explainability.shap_engine import SHAPEngine | |
| from src.ai.explainability.explainer import RuleExplanationEngine, BehaviorExplanationEngine, RiskExplanationEngine, NaturalLanguageGenerator | |
| from src.ai.explainability.evidence import EvidencePackageGenerator | |
| logger = logging.getLogger(__name__) | |
| class Phase5Pipeline: | |
| """ | |
| Phase 5 Orchestrator: Explainability (XAI), Evidence Generation & Audit Trail. | |
| """ | |
| def __init__(self, config: dict, model_dir: str, output_dir: str): | |
| self.config = config | |
| self.model_dir = Path(model_dir) | |
| self.output_dir = Path(output_dir) | |
| # Initializing sub-engines | |
| self.rule_explainer = RuleExplanationEngine() | |
| self.behavior_explainer = BehaviorExplanationEngine() | |
| self.risk_explainer = RiskExplanationEngine() | |
| self.nl_generator = NaturalLanguageGenerator() | |
| self.evidence_gen = EvidencePackageGenerator(config) | |
| # Load classification model for SHAP | |
| self.classifier = XGBoostClassifier(str(model_dir), config) | |
| self.classifier.load() | |
| self.shap_engine = SHAPEngine(self.classifier.model) if self.classifier.model else None | |
| def execute(self, features_path: str, risk_scores_path: str) -> pd.DataFrame: | |
| logger.info("Starting Phase 5 Explainability Pipeline...") | |
| features_df = pd.read_parquet(features_path) | |
| risk_df = pd.read_parquet(risk_scores_path) | |
| # Sync indices using event_id | |
| df = pd.merge(risk_df, features_df, on="event_id", how="inner") | |
| # Extract predictions for SHAP | |
| # Note: mapping category back to label index if possible, else default 1 | |
| predicted_classes = [] | |
| for cat in df["attack_category"]: | |
| try: | |
| idx = self.classifier.labels.index(cat) | |
| predicted_classes.append(idx) | |
| except ValueError: | |
| predicted_classes.append(1) | |
| # 1. SHAP Feature Importance | |
| # Limit background data for performance in SHAP (just an assumption) | |
| X_shap = df[self.classifier.features_used] if self.classifier else pd.DataFrame() | |
| if self.shap_engine and not X_shap.empty: | |
| shap_explanations = self.shap_engine.explain(X_shap, predicted_classes) | |
| else: | |
| shap_explanations = {i: [] for i in range(len(df))} | |
| output_records = [] | |
| for i, row in df.iterrows(): | |
| event_id = row["event_id"] | |
| # 2. Rule Explanations | |
| triggered_rules = row.get("triggered_rules", "") | |
| rule_exp = self.rule_explainer.explain(triggered_rules, row.to_dict()) | |
| # 3. Behavior Explanations | |
| behavior_score = row.get("behavioral_contribution", 0) | |
| behavior_exp = self.behavior_explainer.explain(behavior_score, row.to_dict()) | |
| # 4. Risk Breakdown | |
| risk_exp = self.risk_explainer.explain(row.to_dict()) | |
| # 5. Build Evidence Package First | |
| feature_values = row[self.classifier.features_used].to_dict() if self.classifier else {} | |
| package = { | |
| "Event Metadata": {"Event ID": event_id}, | |
| "Triggered Rules": rule_exp, | |
| "Behavior Deviations": behavior_exp, | |
| "Feature Contributions": shap_explanations.get(i, []), | |
| "Feature Values": feature_values, | |
| "Risk Breakdown": risk_exp, | |
| "Attack Classification": row.get("attack_category", "Unknown"), | |
| "Risk Level": row.get("risk_level", "Unknown"), | |
| "Confidence": row.get("prediction_confidence", 0.0), | |
| "Recommended Action": self.evidence_gen.recommender.recommend(row.get("risk_level", "Unknown")), | |
| "Audit Metadata": self.evidence_gen.audit.generate_audit_metadata( | |
| event_id, | |
| row.get("detection_engine_version", "1.0"), | |
| row.get("risk_engine_version", "1.0"), | |
| row.get("classification_model_version", "1.0") | |
| ) | |
| } | |
| # 6. Natural Language Generation | |
| nl_exp = self.nl_generator.generate(package) | |
| package["Explanation"] = nl_exp | |
| # Map back to flat schema per Output Contract | |
| output_records.append({ | |
| "event_id": event_id, | |
| "Explanation": json.dumps(nl_exp), | |
| "SHAP_Values": json.dumps(package["Feature Contributions"]), | |
| "Triggered_Rules": json.dumps(package["Triggered Rules"]), | |
| "Risk_Breakdown": json.dumps(package["Risk Breakdown"]), | |
| "Recommended_Action": package["Recommended Action"], | |
| "Audit_Metadata": json.dumps(package["Audit Metadata"]) | |
| }) | |
| out_df = pd.DataFrame(output_records) | |
| # Persist Output | |
| self.output_dir.mkdir(parents=True, exist_ok=True) | |
| out_path = self.output_dir / "explanations.parquet" | |
| out_df.to_parquet(out_path, index=False) | |
| logger.info(f"Phase 5 Complete. Explanations persisted to {out_path}") | |
| return out_df | |