Update hf_demo.py

hf_demo.py

@@ -36,6 +36,9 @@ from infrastructure import (
     RecommendedAction,
 )
 
+# ============== HMC LEARNER IMPORT ==============
+from hmc_learner import train_hmc_model   # new import
+
 # ============== CONFIGURATION (Pydantic V2) ==============
 class Settings(BaseSettings):
     """Application settings loaded from environment variables."""
@@ -257,6 +260,44 @@ class BayesianRiskEngine:
         except sqlite3.Error as e:
             logger.error(f"Failed to record outcome: {e}")
 
+    # ---------- NEW: Enhanced risk using HMC coefficients ----------
+    def enhanced_risk(self, action_text: str, context: Dict, hmc_coeffs: Optional[Dict] = None) -> float:
+        """
+        Compute a risk score using HMC coefficients if available.
+        Falls back to simple posterior score if no coefficients.
+        """
+        if hmc_coeffs is None:
+            return self.calculate_posterior(action_text, context)["score"]
+
+        # Build feature vector (same as in hmc_learner preprocessing)
+        action_cat = self.classify_action(action_text)
+        # Map category to code using saved mapping (if present)
+        cat_mapping = hmc_coeffs.get("action_cat_mapping", {})
+        # Invert mapping (category -> code)
+        cat_to_code = {v: k for k, v in cat_mapping.items()}
+        cat_code = cat_to_code.get(action_cat, 0)  # default to 0 if not found
+
+        env_prod = 1 if context.get('environment') == 'production' else 0
+        role_junior = 1 if context.get('user_role') == 'junior' else 0
+        hour = datetime.now().hour
+        # Use the simple posterior risk as a feature (normalized)
+        simple_risk = self.calculate_posterior(action_text, context)["score"]
+        confidence = context.get('confidence', 0.85)
+
+        # Linear predictor from HMC coefficients
+        logit = (
+            hmc_coeffs.get('α_cat', {}).get('mean', [0])[cat_code] +
+            hmc_coeffs.get('β_env', {}).get('mean', 0) * env_prod +
+            hmc_coeffs.get('β_role', {}).get('mean', 0) * role_junior +
+            hmc_coeffs.get('β_risk', {}).get('mean', 0) * (simple_risk - 0.5) +
+            hmc_coeffs.get('β_hour', {}).get('mean', 0) * ((hour - 12) / 12) +
+            hmc_coeffs.get('β_conf', {}).get('mean', 0) * (confidence - 0.5)
+        )
+        # Convert to probability
+        prob = 1 / (1 + np.exp(-logit))
+        return prob
+
+
 class PolicyEngine:
     """Deterministic OSS policies – advisory only."""
     def __init__(self):
@@ -399,6 +440,18 @@ class RAGMemory:
         self.embedding_cache[text] = embedding
         return embedding
 
+    def _ensure_columns(self, conn, columns):
+        """Add columns to incidents table if they do not exist."""
+        cursor = conn.execute("PRAGMA table_info(incidents)")
+        existing = [row[1] for row in cursor.fetchall()]
+        for col_name, col_type in columns:
+            if col_name not in existing:
+                try:
+                    conn.execute(f"ALTER TABLE incidents ADD COLUMN {col_name} {col_type}")
+                    logger.info(f"Added column {col_name} to incidents table")
+                except sqlite3.Error as e:
+                    logger.error(f"Failed to add column {col_name}: {e}")
+
     def _init_db(self):
         try:
             with self._get_db() as conn:
@@ -416,6 +469,15 @@ class RAGMemory:
                         embedding TEXT
                     )
                 ''')
+                # Add new columns if they don't exist
+                self._ensure_columns(conn, [
+                    ('environment', 'TEXT'),
+                    ('user_role', 'TEXT'),
+                    ('requires_human', 'BOOLEAN'),
+                    ('rollback_feasible', 'BOOLEAN'),
+                    ('hour_of_day', 'INTEGER'),
+                    ('action_category', 'TEXT')
+                ])
                 conn.execute('''
                     CREATE TABLE IF NOT EXISTS signals (
                         id TEXT PRIMARY KEY,
@@ -449,17 +511,19 @@ class RAGMemory:
                 conn.close()
 
     def store_incident(self, action: str, risk_score: float, risk_level: RiskLevel,
-                       confidence: float, allowed: bool, gates: List[Dict]):
+                       confidence: float, allowed: bool, gates: List[Dict],
+                       environment: str, user_role: str, requires_human: bool,
+                       rollback_feasible: bool, hour_of_day: int, action_category: str):
         action_hash = hashlib.sha256(action.encode()).hexdigest()[:50]
-        # Build a descriptive text and generate embedding
         incident_text = self._build_incident_text(action)
         embedding = json.dumps(self._simple_embedding(incident_text))
         try:
             with self._get_db() as conn:
                 conn.execute('''
                     INSERT INTO incidents 
-                    (id, action, action_hash, risk_score, risk_level, confidence, allowed, gates, timestamp, embedding)
-                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                    (id, action, action_hash, risk_score, risk_level, confidence, allowed, gates, timestamp, embedding,
+                     environment, user_role, requires_human, rollback_feasible, hour_of_day, action_category)
+                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                 ''', (
                     str(uuid.uuid4()),
                     action[:500],
@@ -470,7 +534,13 @@ class RAGMemory:
                     1 if allowed else 0,
                     json.dumps(gates),
                     datetime.utcnow().isoformat(),
-                    embedding
+                    embedding,
+                    environment,
+                    user_role,
+                    1 if requires_human else 0,
+                    1 if rollback_feasible else 0,
+                    hour_of_day,
+                    action_category
                 ))
                 conn.commit()
         except sqlite3.Error as e:
@@ -663,6 +733,23 @@ class InfrastructureEvaluationResponse(BaseModel):
     confidence_score: float
     evaluation_details: Dict[str, Any]
 
+# ============== GLOBAL HMC MODEL DATA ==============
+hmc_model_data = None
+
+def load_hmc_model():
+    global hmc_model_data
+    model_path = f"{settings.data_dir}/hmc_model.json"
+    if os.path.exists(model_path):
+        try:
+            with open(model_path, 'r') as f:
+                hmc_model_data = json.load(f)
+            logger.info("HMC model loaded successfully")
+        except Exception as e:
+            logger.error(f"Failed to load HMC model: {e}")
+            hmc_model_data = None
+    else:
+        logger.info("No HMC model found; using default risk engine")
+
 # ============== FASTAPI APP ==============
 app = FastAPI(
     title="ARF OSS Real Engine (API Only)",
@@ -686,6 +773,7 @@ app.add_middleware(
 risk_engine = BayesianRiskEngine()
 policy_engine = PolicyEngine()
 memory = RAGMemory()
+load_hmc_model()   # Load HMC model after memory init
 
 # ============== INFRASTRUCTURE SIMULATOR INSTANCE ==============
 # Corrected: RegionAllowedPolicy expects 'allowed_regions', not 'regions'
@@ -741,12 +829,28 @@ async def evaluate_action(request: ActionRequest):
             "environment": "production",
             "user_role": request.user_role,
             "backup_available": request.rollbackFeasible,
-            "requires_human": request.requiresHuman
+            "requires_human": request.requiresHuman,
+            "confidence": request.confidenceScore   # added for enhanced_risk
         }
-        risk = risk_engine.calculate_posterior(
-            action_text=request.proposedAction,
-            context=context
-        )
+        # Use HMC-enhanced risk if available
+        if hmc_model_data:
+            risk_score_val = risk_engine.enhanced_risk(request.proposedAction, context, hmc_model_data)
+            # Convert to a risk dict compatible with policy engine (needs level and interval)
+            # For simplicity, reuse the simple engine's level mapping based on enhanced score
+            risk = risk_engine.calculate_posterior(request.proposedAction, context)
+            risk["score"] = risk_score_val
+            if risk_score_val > 0.8:
+                risk["level"] = RiskLevel.CRITICAL
+            elif risk_score_val > 0.6:
+                risk["level"] = RiskLevel.HIGH
+            elif risk_score_val > 0.4:
+                risk["level"] = RiskLevel.MEDIUM
+            else:
+                risk["level"] = RiskLevel.LOW
+            # Recalculate credible interval? We'll keep the simple one for now.
+        else:
+            risk = risk_engine.calculate_posterior(request.proposedAction, context)
+
         policy = policy_engine.evaluate(
             action=request.proposedAction,
             risk=risk,
@@ -754,6 +858,14 @@ async def evaluate_action(request: ActionRequest):
         )
         similar = memory.find_similar(request.proposedAction, limit=3)
 
+        # Capture additional fields for logging
+        environment = context["environment"]
+        user_role = request.user_role
+        requires_human = request.requiresHuman
+        rollback_feasible = request.rollbackFeasible
+        hour_of_day = datetime.now().hour
+        action_category = risk_engine.classify_action(request.proposedAction)
+
         if not policy["allowed"] and risk["score"] > 0.7:
             memory.track_enterprise_signal(
                 signal_type=LeadSignal.HIGH_RISK_BLOCKED,
@@ -778,7 +890,13 @@ async def evaluate_action(request: ActionRequest):
             risk_level=risk["level"],
             confidence=request.confidenceScore,
             allowed=policy["allowed"],
-            gates=policy["gates"]
+            gates=policy["gates"],
+            environment=environment,
+            user_role=user_role,
+            requires_human=requires_human,
+            rollback_feasible=rollback_feasible,
+            hour_of_day=hour_of_day,
+            action_category=action_category
         )
         gates = []
         for g in policy["gates"]:
@@ -904,6 +1022,20 @@ async def evaluate_infrastructure_intent(request: InfrastructureIntentRequest):
         logger.error(f"Infrastructure evaluation failed: {e}", exc_info=True)
         raise HTTPException(500, detail=str(e))
 
+# ============== NEW HMC TRAINING ENDPOINT ==============
+@app.post("/api/v1/admin/train_hmc", dependencies=[Depends(verify_api_key)])
+async def train_hmc():
+    """Trigger HMC training on historical incident data."""
+    global hmc_model_data
+    try:
+        db_path = f"{settings.data_dir}/memory.db"
+        model_data = train_hmc_model(db_path, output_dir=settings.data_dir)
+        hmc_model_data = model_data
+        return {"status": "success", "message": "HMC model trained and loaded", "coefficients": model_data.get("coefficients")}
+    except Exception as e:
+        logger.error(f"HMC training failed: {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail=str(e))
+
 # ============== MAIN ENTRY POINT ==============
 if __name__ == "__main__":
     import uvicorn