Update hf_demo.py

hf_demo.py

@@ -1,10 +1,12 @@
-# hf_demo.py – ARF v4 dashboard for Hugging Face Spaces
+# hf_demo.py – ARF v4 API with Gradio frontend (optional)
 import logging
 import uuid
 from datetime import datetime, timezone
+from typing import Dict, Any, Optional
 
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
 import gradio as gr
 
 # ARF v4 imports
@@ -15,9 +17,10 @@ from agentic_reliability_framework.runtime.memory.constants import MemoryConstan
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-app = FastAPI(title="ARF v4 API with Memory")
+# ========================= FASTAPI APP =========================
+app = FastAPI(title="ARF v4 API")
 
-# Enable CORS
+# Enable CORS for your frontend
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["https://arf-frontend-sandy.vercel.app"],
@@ -25,62 +28,55 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# ---------------------------------------------------------------------------
-# Initialize ARF components
-# ---------------------------------------------------------------------------
-
+# ========================= ARF COMPONENTS =========================
 risk_engine = RiskEngine()
-
 faiss_index = create_faiss_index(dim=MemoryConstants.VECTOR_DIM)
 memory = RAGGraphMemory(faiss_index)
 
-# ---------------------------------------------------------------------------
-# Decision Memory (NEW)
-# ---------------------------------------------------------------------------
-
+# In‑memory storage for demo purposes (replace with a real DB later)
 decision_history = []
-risk_history = []
 
-# ---------------------------------------------------------------------------
-# API Endpoints
-# ---------------------------------------------------------------------------
+# ========================= PYDANTIC MODELS =========================
+class EvaluateRequest(BaseModel):
+    service_name: str
+    event_type: str
+    severity: str
+    metrics: Dict[str, float] = {}
+
+class EvaluateResponse(BaseModel):
+    risk_score: float
+    base_risk: float
+    memory_risk: Optional[float] = None
+    weight: float
+    similar_events: list = []
+    confidence: float
+
+# ========================= HELPER: Demo Intent =========================
+class _DemoIntent:
+    environment = "dev"
+    deployment_target = "dev"
+    service_name = "demo"
+
+# ========================= API ENDPOINTS =========================
 
 @app.get("/")
 async def root():
-    return {
-        "service": "ARF OSS API",
-        "version": "4.0.0",
-        "status": "operational",
-        "memory_stats": memory.get_graph_stats() if memory.has_historical_data() else "empty",
-    }
+    """Root endpoint – returns a welcome message."""
+    return {"message": "ARF v4 API. See /docs for documentation."}
 
 @app.get("/health")
 async def health():
     return {"status": "ok", "version": "4.0.0"}
 
-# ---------------------------------------------------------------------------
-# Demo Intent Object
-# ---------------------------------------------------------------------------
-
-class _DemoIntent:
-    environment = "dev"
-    deployment_target = "dev"
-    service_name = "demo"
-
-# ---------------------------------------------------------------------------
-# Risk Calculation
-# ---------------------------------------------------------------------------
-
-def _calculate_demo_risk():
-
+@app.get("/api/v1/get_risk")
+async def get_risk():
+    """Return the current demo risk."""
     intent = _DemoIntent()
-
     risk_value, explanation, contributions = risk_engine.calculate_risk(
         intent=intent,
         cost_estimate=None,
         policy_violations=[],
     )
-
     decision = "approve"
     if risk_value > 0.8:
         decision = "deny"
@@ -88,167 +84,138 @@ def _calculate_demo_risk():
         decision = "escalate"
 
     decision_id = str(uuid.uuid4())
-
-    decision_record = {
-        "id": decision_id,
+    decision_history.append({
+        "decision_id": decision_id,
         "timestamp": datetime.now(timezone.utc).isoformat(),
-        "risk": float(risk_value),
-        "decision": decision
-    }
-
-    decision_history.append(decision_record)
-
-    risk_history.append(
-        {
-            "timestamp": decision_record["timestamp"],
-            "risk": float(risk_value)
-        }
-    )
+        "risk_score": float(risk_value),
+        "outcome": None,  # will be filled when feedback is given
+    })
 
     return {
-        "risk": float(risk_value),
+        "system_risk": float(risk_value),
         "status": "critical" if risk_value > 0.8 else "normal",
         "explanation": explanation,
         "contributions": contributions,
         "decision_id": decision_id,
-        "decision": decision
+        "decision": decision,
+        "timestamp": datetime.now(timezone.utc).isoformat()
     }
 
-# ---------------------------------------------------------------------------
-# Outcome Feedback
-# ---------------------------------------------------------------------------
-
-def record_outcome(success: bool):
-
-    if not decision_history:
-        return {"error": "no decisions yet"}
-
-    last_decision = decision_history[-1]
-
-    intent = _DemoIntent()
-
-    try:
-        risk_engine.update_outcome(intent, success)
-    except Exception as e:
-        logger.exception("Outcome update failed")
-
-    last_decision["outcome"] = "success" if success else "failure"
+@app.get("/api/v1/history")
+async def get_history():
+    """Return the last 10 decisions."""
+    return decision_history[-10:]
 
-    return {
-        "decision_id": last_decision["id"],
-        "outcome": last_decision["outcome"],
-        "timestamp": datetime.now(timezone.utc).isoformat()
-    }
+@app.post("/api/v1/incidents/evaluate", response_model=EvaluateResponse)
+async def evaluate_incident(request: EvaluateRequest):
+    """
+    Evaluate an incident and return a risk score with explainability.
+    This is a placeholder – replace with actual call to your risk engine.
+    """
+    # For now, return a dummy response
+    return EvaluateResponse(
+        risk_score=0.23,
+        base_risk=0.15,
+        memory_risk=0.3,
+        weight=0.5,
+        similar_events=[],
+        confidence=0.9
+    )
 
-# ---------------------------------------------------------------------------
-# Snapshot functions
-# ---------------------------------------------------------------------------
+@app.post("/api/v1/feedback")
+async def record_outcome(decision_id: str, success: bool):
+    """Record the outcome of a decision (success/failure)."""
+    for dec in decision_history:
+        if dec["decision_id"] == decision_id:
+            dec["outcome"] = "success" if success else "failure"
+            # Update the risk engine (optional)
+            intent = _DemoIntent()
+            try:
+                risk_engine.update_outcome(intent, success)
+            except Exception as e:
+                logger.exception("Outcome update failed")
+            return {"status": "ok", "decision_id": decision_id, "outcome": dec["outcome"]}
+    return {"error": "decision not found"}
+
+# ========================= GRADIO UI (Optional) =========================
+# This mounts the Gradio interface at the root path ("/").
 
 def get_risk_snapshot():
-
     try:
-        snapshot = _calculate_demo_risk()
-
-        snapshot["timestamp"] = datetime.now(timezone.utc).isoformat()
-
-        return snapshot
-
-    except Exception as e:
-
-        logger.exception("Failed to compute risk snapshot")
-
-        return {
-            "status": "error",
-            "error": str(e),
+        intent = _DemoIntent()
+        risk_value, explanation, contributions = risk_engine.calculate_risk(
+            intent=intent,
+            cost_estimate=None,
+            policy_violations=[],
+        )
+        decision = "approve"
+        if risk_value > 0.8:
+            decision = "deny"
+        elif risk_value > 0.2:
+            decision = "escalate"
+        decision_id = str(uuid.uuid4())
+        decision_history.append({
+            "decision_id": decision_id,
             "timestamp": datetime.now(timezone.utc).isoformat(),
+            "risk_score": float(risk_value),
+            "outcome": None,
+        })
+        return {
+            "risk": float(risk_value),
+            "status": "critical" if risk_value > 0.8 else "normal",
+            "explanation": explanation,
+            "contributions": contributions,
+            "decision_id": decision_id,
+            "decision": decision,
+            "timestamp": datetime.now(timezone.utc).isoformat()
         }
+    except Exception as e:
+        logger.exception("Failed to compute risk snapshot")
+        return {"error": str(e)}
 
 def get_health_snapshot():
-
-    return {
-        "status": "ok",
-        "version": "4.0.0",
-        "service": "ARF OSS API",
-        "timestamp": datetime.now(timezone.utc).isoformat(),
-    }
+    return {"status": "ok", "version": "4.0.0", "service": "ARF OSS API", "timestamp": datetime.now(timezone.utc).isoformat()}
 
 def get_memory_snapshot():
+    if memory.has_historical_data():
+        return {"status": "ok", "memory_stats": memory.get_graph_stats(), "timestamp": datetime.now(timezone.utc).isoformat()}
+    return {"status": "empty", "memory_stats": "No historical memory yet.", "timestamp": datetime.now(timezone.utc).isoformat()}
 
+def record_outcome_ui(success: bool):
+    if not decision_history:
+        return {"error": "no decisions yet"}
+    last = decision_history[-1]
+    last["outcome"] = "success" if success else "failure"
+    intent = _DemoIntent()
     try:
-
-        if memory.has_historical_data():
-
-            stats = memory.get_graph_stats()
-
-            return {
-                "status": "ok",
-                "memory_stats": stats,
-                "timestamp": datetime.now(timezone.utc).isoformat(),
-            }
-
-        return {
-            "status": "empty",
-            "memory_stats": "No historical memory yet.",
-            "timestamp": datetime.now(timezone.utc).isoformat(),
-        }
-
+        risk_engine.update_outcome(intent, success)
     except Exception as e:
-
-        logger.exception("Failed to compute memory snapshot")
-
-        return {
-            "status": "error",
-            "error": str(e),
-            "timestamp": datetime.now(timezone.utc).isoformat(),
-        }
-
-def get_decision_history():
-    return decision_history[-10:]
-
-# ---------------------------------------------------------------------------
-# Gradio Dashboard
-# ---------------------------------------------------------------------------
+        logger.exception("Outcome update failed")
+    return {"decision_id": last["decision_id"], "outcome": last["outcome"], "timestamp": datetime.now(timezone.utc).isoformat()}
 
 with gr.Blocks(title="ARF v4 Demo") as demo:
-
-    gr.Markdown("# Agentic Reliability Framework v4")
-    gr.Markdown("### Probabilistic Infrastructure Governance")
-
+    gr.Markdown("# Agentic Reliability Framework v4\n### Probabilistic Infrastructure Governance")
     with gr.Row():
         health_output = gr.JSON(label="Health")
         risk_output = gr.JSON(label="Current Risk")
-
     with gr.Row():
         memory_output = gr.JSON(label="Memory Stats")
-
     with gr.Row():
         decision_output = gr.JSON(label="Recent Decisions")
-
     with gr.Row():
         refresh_btn = gr.Button("Evaluate Intent")
         success_btn = gr.Button("Action Succeeded")
         fail_btn = gr.Button("Action Failed")
-
     refresh_btn.click(fn=get_risk_snapshot, outputs=risk_output)
-    success_btn.click(fn=lambda: record_outcome(True), outputs=decision_output)
-    fail_btn.click(fn=lambda: record_outcome(False), outputs=decision_output)
-
+    success_btn.click(fn=lambda: record_outcome_ui(True), outputs=decision_output)
+    fail_btn.click(fn=lambda: record_outcome_ui(False), outputs=decision_output)
     with gr.Row():
         health_btn = gr.Button("Refresh Health")
         memory_btn = gr.Button("Refresh Memory")
         history_btn = gr.Button("Show Decision History")
-
     health_btn.click(fn=get_health_snapshot, outputs=health_output)
     memory_btn.click(fn=get_memory_snapshot, outputs=memory_output)
-    history_btn.click(fn=get_decision_history, outputs=decision_output)
-
-    demo.load(fn=get_health_snapshot, outputs=health_output)
-    demo.load(fn=get_risk_snapshot, outputs=risk_output)
-    demo.load(fn=get_memory_snapshot, outputs=memory_output)
-
-# ---------------------------------------------------------------------------
-# Entry point
-# ---------------------------------------------------------------------------
+    history_btn.click(fn=lambda: decision_history[-10:], outputs=decision_output)
 
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0")
+# Mount Gradio app at the root path
+app = gr.mount_gradio_app(app, demo, path="/")