Spaces:
Paused
Paused
File size: 5,474 Bytes
5a3b9db | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | 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
|