Update app.py

app.py

@@ -7,44 +7,75 @@ import os
 import torch
 import numpy as np
 from datetime import datetime
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from sentence_transformers import SentenceTransformer, util
-from diffusers import StableDiffusionPipeline
-import librosa
-import soundfile as sf
-import tempfile
 
 # ARF components
 from agentic_reliability_framework.runtime.engine import EnhancedReliabilityEngine
-from hallucination_detective import HallucinationDetectiveAgent
-from memory_drift_diagnostician import MemoryDriftDiagnosticianAgent
-from image_detector import ImageQualityDetector
-from audio_detector import AudioQualityDetector
+from agentic_reliability_framework.core.models.event import ReliabilityEvent
+
+# Custom AI components
 from ai_event import AIEvent
 from ai_risk_engine import AIRiskEngine
+from hallucination_detective import HallucinationDetectiveAgent
+from memory_drift_diagnostician import MemoryDriftDiagnosticianAgent
 from nli_detector import NLIDetector
 from retrieval import SimpleRetriever
+from image_detector import ImageQualityDetector
+from audio_detector import AudioQualityDetector
+from iot_simulator import IoTSimulator
+from robotics_diagnostician import RoboticsDiagnostician
+from iot_event import IoTEvent
 
-logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+# ----------------------------------------------------------------------
+# Logging setup
+# ----------------------------------------------------------------------
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
 # ----------------------------------------------------------------------
-# Infrastructure engine (optional)
+# ARF infrastructure engine (optional)
 # ----------------------------------------------------------------------
 try:
     logger.info("Initializing EnhancedReliabilityEngine...")
-    engine = EnhancedReliabilityEngine()
+    infra_engine = EnhancedReliabilityEngine()
 except Exception as e:
-    logger.error(f"Engine init failed: {e}")
-    engine = None
+    logger.error(f"Infrastructure engine init failed: {e}")
+    infra_engine = None
 
 # ----------------------------------------------------------------------
-# Generative model for text (DialoGPT-small)
+# Text generation model (DialoGPT-small) with logprobs
 # ----------------------------------------------------------------------
+from transformers import AutoTokenizer, AutoModelForCausalLM
 gen_model_name = "microsoft/DialoGPT-small"
-tokenizer = AutoTokenizer.from_pretrained(gen_model_name)
-model = AutoModelForCausalLM.from_pretrained(gen_model_name)
-logger.info(f"Generator {gen_model_name} loaded.")
+try:
+    tokenizer = AutoTokenizer.from_pretrained(gen_model_name)
+    model = AutoModelForCausalLM.from_pretrained(gen_model_name)
+    model.eval()
+    logger.info(f"Generator {gen_model_name} loaded.")
+except Exception as e:
+    logger.error(f"Generator load failed: {e}")
+    tokenizer = model = None
+
+def generate_with_logprobs(prompt, max_new_tokens=100):
+    if tokenizer is None or model is None:
+        return "[Model not loaded]", -10.0
+    inputs = tokenizer(prompt, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            return_dict_in_generate=True,
+            output_scores=True
+        )
+    scores = outputs.scores
+    log_probs = [torch.log_softmax(score, dim=-1) for score in scores]
+    generated_ids = outputs.sequences[0][inputs['input_ids'].shape[1]:]
+    token_log_probs = []
+    for i, lp in enumerate(log_probs):
+        token_id = generated_ids[i]
+        token_log_probs.append(lp[0, token_id].item())
+    avg_log_prob = sum(token_log_probs) / len(token_log_probs) if token_log_probs else -10.0
+    generated_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
+    return generated_text, avg_log_prob
 
 # ----------------------------------------------------------------------
 # NLI detector
@@ -52,14 +83,15 @@ logger.info(f"Generator {gen_model_name} loaded.")
 nli_detector = NLIDetector()
 
 # ----------------------------------------------------------------------
-# Sentence‑Transformer retriever
+# Retrieval (sentence‑transformers + ChromaDB)
 # ----------------------------------------------------------------------
 retriever = SimpleRetriever()
-logger.info("Retriever loaded.")
 
 # ----------------------------------------------------------------------
-# Image generation (tiny model for demo)
+# Image generation (tiny diffusion model)
 # ----------------------------------------------------------------------
+from diffusers import StableDiffusionPipeline
+image_pipe = None
 try:
     image_pipe = StableDiffusionPipeline.from_pretrained(
         "hf-internal-testing/tiny-stable-diffusion-torch"
@@ -68,12 +100,13 @@ try:
         image_pipe.to("cpu")
     logger.info("Image pipeline loaded.")
 except Exception as e:
-    logger.error(f"Image pipeline failed: {e}")
-    image_pipe = None
+    logger.warning(f"Image pipeline load failed (will be disabled): {e}")
 
 # ----------------------------------------------------------------------
 # Audio transcription (Whisper tiny)
 # ----------------------------------------------------------------------
+from transformers import pipeline
+audio_pipe = None
 try:
     audio_pipe = pipeline(
         "automatic-speech-recognition",
@@ -82,8 +115,7 @@ try:
     )
     logger.info("Audio pipeline loaded.")
 except Exception as e:
-    logger.error(f"Audio pipeline failed: {e}")
-    audio_pipe = None
+    logger.warning(f"Audio pipeline load failed (will be disabled): {e}")
 
 # ----------------------------------------------------------------------
 # AI agents
@@ -92,6 +124,7 @@ hallucination_detective = HallucinationDetectiveAgent(nli_detector=nli_detector)
 memory_drift_diagnostician = MemoryDriftDiagnosticianAgent()
 image_quality_detector = ImageQualityDetector()
 audio_quality_detector = AudioQualityDetector()
+robotics_diagnostician = RoboticsDiagnostician()
 
 # ----------------------------------------------------------------------
 # Bayesian risk engine
@@ -99,37 +132,32 @@ audio_quality_detector = AudioQualityDetector()
 ai_risk_engine = AIRiskEngine()
 
 # ----------------------------------------------------------------------
-# Generation helper with log probabilities
+# IoT simulator
 # ----------------------------------------------------------------------
-def generate_with_logprobs(prompt, max_new_tokens=100):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=max_new_tokens,
-            return_dict_in_generate=True,
-            output_scores=True
-        )
-    scores = outputs.scores
-    log_probs = [torch.log_softmax(score, dim=-1) for score in scores]
-    generated_ids = outputs.sequences[0][inputs['input_ids'].shape[1]:]
-    token_log_probs = []
-    for i, lp in enumerate(log_probs):
-        token_id = generated_ids[i]
-        token_log_probs.append(lp[0, token_id].item())
-    avg_log_prob = sum(token_log_probs) / len(token_log_probs) if token_log_probs else 0.0
-    generated_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
-    return generated_text, avg_log_prob
+iot_sim = IoTSimulator()
+iot_history = []  # store recent readings for prediction
 
 # ----------------------------------------------------------------------
-# Task handlers
+# Helper: update risk with feedback
+# ----------------------------------------------------------------------
+last_task_category = None
+
+def feedback(thumbs_up: bool):
+    global last_task_category
+    if last_task_category is None:
+        return "No previous analysis to rate."
+    ai_risk_engine.update_outcome(last_task_category, success=thumbs_up)
+    return f"Feedback recorded: {'👍' if thumbs_up else '👎'} for {last_task_category}."
+
+# ----------------------------------------------------------------------
+# Async handlers for each tab
 # ----------------------------------------------------------------------
 async def handle_text(task_type, prompt):
+    global last_task_category
+    last_task_category = task_type
     try:
         response, avg_log_prob = generate_with_logprobs(prompt)
-        # Get retrieval score
         retrieval_score = retriever.get_similarity(prompt)
-        # Create event
         event = AIEvent(
             timestamp=datetime.utcnow(),
             component="ai",
@@ -145,13 +173,12 @@ async def handle_text(task_type, prompt):
             prompt=prompt,
             response=response,
             response_length=len(response),
-            confidence=float(np.exp(avg_log_prob)),  # convert to probability scale
+            confidence=float(np.exp(avg_log_prob)),  # convert to [0,1] scale (approx)
             perplexity=None,
             retrieval_scores=[retrieval_score],
             user_feedback=None,
             latency_ms=0
         )
-        # Analyze
         hallu_result = await hallucination_detective.analyze(event)
         drift_result = await memory_drift_diagnostician.analyze(event)
         risk_metrics = ai_risk_engine.risk_score(task_type)
@@ -169,14 +196,15 @@ async def handle_text(task_type, prompt):
         return {"error": str(e)}
 
 async def handle_image(prompt):
+    global last_task_category
+    last_task_category = "image"
     if image_pipe is None:
-        return {"error": "Image model not loaded"}
+        return {"error": "Image model not loaded"}, None
     try:
         import time
         start = time.time()
-        image = image_pipe(prompt, num_inference_steps=2).images[0]  # tiny steps for speed
+        image = image_pipe(prompt, num_inference_steps=2).images[0]  # minimal steps
         gen_time = time.time() - start
-        # Mock retrieval score (you could use CLIP similarity)
         retrieval_score = retriever.get_similarity(prompt)
         event = AIEvent(
             timestamp=datetime.utcnow(),
@@ -191,7 +219,7 @@ async def handle_image(prompt):
             model_name="tiny-sd",
             model_version="latest",
             prompt=prompt,
-            response="",  # image not text
+            response="",  # not text
             response_length=0,
             confidence=1.0 / (gen_time + 1),  # heuristic
             perplexity=None,
@@ -208,20 +236,24 @@ async def handle_image(prompt):
         }
     except Exception as e:
         logger.error(f"Image task error: {e}")
-        return {"error": str(e)}
+        return {"error": str(e)}, None
 
 async def handle_audio(audio_file):
+    global last_task_category
+    last_task_category = "audio"
     if audio_pipe is None:
         return {"error": "Audio model not loaded"}
     try:
-        # Load audio (Gradio provides file path)
+        import librosa
+        import soundfile as sf
+        import tempfile
         audio, sr = librosa.load(audio_file, sr=16000)
         with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
             sf.write(tmp.name, audio, sr)
             result = audio_pipe(tmp.name, return_timestamps=False)
         text = result["text"]
-        # Whisper does not output log probs easily; we'll use a placeholder
-        avg_log_prob = -2.0  # placeholder
+        # Whisper does not output log probs easily; use placeholder
+        avg_log_prob = -2.0
         event = AIEvent(
             timestamp=datetime.utcnow(),
             component="audio",
@@ -254,47 +286,96 @@ async def handle_audio(audio_file):
         logger.error(f"Audio task error: {e}")
         return {"error": str(e)}
 
-# ----------------------------------------------------------------------
-# Feedback handling
-# ----------------------------------------------------------------------
-last_event_category = None
-def feedback(thumbs_up: bool):
-    global last_event_category
-    if last_event_category is None:
-        return "No previous analysis to rate."
-    ai_risk_engine.update_outcome(last_event_category, success=thumbs_up)
-    return f"Feedback recorded: {'👍' if thumbs_up else '👎'} for {last_event_category}."
+async def read_iot_sensors(fault_type):
+    global last_task_category, iot_history
+    last_task_category = "iot"
+    iot_sim.set_fault(fault_type if fault_type != "none" else None)
+    data = iot_sim.read()
+    iot_history.append(data)
+    if len(iot_history) > 100:
+        iot_history.pop(0)
+
+    # Create IoTEvent
+    event = IoTEvent(
+        timestamp=datetime.utcnow(),
+        component="robotic_arm",
+        service_mesh="factory",
+        latency_p99=0,
+        error_rate=0.0,
+        throughput=1,
+        cpu_util=None,
+        memory_util=None,
+        temperature=data['temperature'],
+        vibration=data['vibration'],
+        motor_current=data['motor_current'],
+        position_error=data['position_error']
+    )
+    # Run diagnostician
+    diag_result = await robotics_diagnostician.analyze(event)
+
+    # Simple failure prediction
+    prediction = None
+    if len(iot_history) >= 5:
+        temps = [h['temperature'] for h in iot_history[-5:]]
+        x = np.arange(len(temps))
+        slope, intercept = np.polyfit(x, temps, 1)
+        next_temp = slope * len(temps) + intercept
+        if slope > 0.1:
+            time_to_threshold = (40.0 - next_temp) / slope if slope > 0 else None
+            prediction = {
+                "predicted_temperature": next_temp,
+                "time_to_overheat_min": time_to_threshold
+            }
+
+    return data, diag_result, prediction
 
 # ----------------------------------------------------------------------
 # Gradio UI
 # ----------------------------------------------------------------------
 with gr.Blocks(title="ARF v4 – AI Reliability Lab", theme="soft") as demo:
-    gr.Markdown("# 🧠 ARF v4 – AI Reliability Lab\n**Detect hallucinations, drift, and failures across text, image, and audio**")
-    
+    gr.Markdown("# 🧠 ARF v4 – AI Reliability Lab\n**Detect hallucinations, drift, and failures across text, image, audio, and robotics**")
+
     with gr.Tabs():
+        # Tab 1: Text Generation
         with gr.TabItem("Text Generation"):
             text_task = gr.Dropdown(["chat", "code", "summary"], value="chat", label="Task")
-            text_prompt = gr.Textbox(label="Prompt", value="What is the capital of France?")
+            text_prompt = gr.Textbox(label="Prompt", value="What is the capital of France?", lines=3)
             text_btn = gr.Button("Generate")
             text_output = gr.JSON(label="Analysis")
-        
+
+        # Tab 2: Image Generation
         with gr.TabItem("Image Generation"):
             img_prompt = gr.Textbox(label="Prompt", value="A cat wearing a hat")
             img_btn = gr.Button("Generate")
             img_output = gr.Image(label="Generated Image")
             img_json = gr.JSON(label="Analysis")
-        
+
+        # Tab 3: Audio Transcription
         with gr.TabItem("Audio Transcription"):
             audio_input = gr.Audio(type="filepath", label="Upload audio file")
             audio_btn = gr.Button("Transcribe")
             audio_output = gr.JSON(label="Analysis")
-    
+
+        # Tab 4: Robotics / IoT
+        with gr.TabItem("Robotics / IoT"):
+            gr.Markdown("### Simulated Robotic Arm Monitoring")
+            fault_type = gr.Dropdown(
+                ["none", "overheat", "vibration", "stall", "drift"],
+                value="none",
+                label="Inject Fault"
+            )
+            refresh_btn = gr.Button("Read Sensors")
+            sensor_display = gr.JSON(label="Sensor Readings")
+            diag_display = gr.JSON(label="Diagnosis")
+            pred_display = gr.JSON(label="Failure Prediction")
+
+    # Feedback row
     with gr.Row():
         feedback_up = gr.Button("👍 Correct")
         feedback_down = gr.Button("👎 Incorrect")
         feedback_msg = gr.Textbox(label="Feedback", interactive=False)
-    
-    # Wire up events
+
+    # Wire events
     text_btn.click(
         fn=lambda task, p: asyncio.run(handle_text(task, p)),
         inputs=[text_task, text_prompt],
@@ -310,6 +391,11 @@ with gr.Blocks(title="ARF v4 – AI Reliability Lab", theme="soft") as demo:
         inputs=audio_input,
         outputs=audio_output
     )
+    refresh_btn.click(
+        fn=lambda f: asyncio.run(read_iot_sensors(f)),
+        inputs=fault_type,
+        outputs=[sensor_display, diag_display, pred_display]
+    )
     feedback_up.click(fn=lambda: feedback(True), outputs=feedback_msg)
     feedback_down.click(fn=lambda: feedback(False), outputs=feedback_msg)