Update app.py

app.py

@@ -1,7 +1,12 @@
 import gradio as gr
-import pandas as pd
+import requests
+import os
 
-# Custom CSS to force the "Sidebar" look in Gradio 4
+# --- 1. DYNAMIC API ROUTING ---
+# Pulls the live AMD IP from Hugging Face Settings
+API_URL = os.environ.get("INTELLIGUARD_API", "http://127.0.0.1:8000/scan")
+
+# Custom CSS
 custom_css = """
 .gradio-container { background-color: #06090f !important; color: white !important; }
 .fake-sidebar { 
@@ -18,23 +23,82 @@ custom_css = """
     background: #111827; padding: 20px; border-radius: 12px; text-align: center;
     border-top: 3px solid #3b82f6; margin-bottom: 10px;
 }
-.establish-btn { background: linear-gradient(90deg, #3b82f6 0%, #00d4ff 100%) !important; color: white !important; }
+.establish-btn { background: linear-gradient(90deg, #3b82f6 0%, #00d4ff 100%) !important; color: white !important; font-weight: bold !important;}
+/* Fix for dark mode inputs */
+input, textarea { color: white !important; background-color: #161b22 !important; }
 """
 
+# --- 2. THE BRAIN LOGIC ---
+def process_query(query, tot_count, thr_count, saf_count, current_log):
+    if not query.strip():
+        return tot_count, thr_count, saf_count, "", "Please enter a query.", current_log, current_log, gr.update(), gr.update(), gr.update()
+
+    tot_count += 1
+    
+    try:
+        # Send payload to AMD Server
+        res = requests.post(API_URL, json={"text": query}, timeout=15)
+        data = res.json()
+        verdict = data.get("verdict", "ERROR")
+        category = data.get("attack_category", "Unknown")
+        score = data.get("score", 0.0)
+    except Exception as e:
+        verdict = "ERROR"
+        category = "Connection Failed"
+        score = 0.0
+
+    # Logic for metrics and UI updating
+    if verdict == "INJECTION":
+        thr_count += 1
+        color = "#ef4444" # Red
+        ui_status = f"🛑 THREAT BLOCKED: {category} (Confidence: {score:.2f})"
+        new_log = f"<div style='border-left: 3px solid #ef4444; padding-left: 10px; margin-bottom: 10px; font-size: 12px;'><b>Threat:</b> {category}<br><span style='color: #9ca3af;'>{query[:30]}...</span></div>"
+    elif verdict == "ERROR":
+        color = "#eab308" # Yellow
+        ui_status = f"⚠️ CONNECTION ERROR: Ensure AMD backend is running and Port 8000 is open."
+        new_log = f"<div style='border-left: 3px solid #eab308; padding-left: 10px; margin-bottom: 10px; font-size: 12px;'><b>Error:</b> API Timeout</div>"
+    else:
+        saf_count += 1
+        color = "#10b981" # Green
+        ui_status = f"✅ SAFE: Query permitted to internal LLM."
+        new_log = f"<div style='border-left: 3px solid #10b981; padding-left: 10px; margin-bottom: 10px; font-size: 12px;'><b>Safe Query</b><br><span style='color: #9ca3af;'>{query[:30]}...</span></div>"
+
+    # Compile the updated HTML components
+    updated_log = new_log + current_log
+    
+    tot_html = f"<div class='card' style='border-color:#3b82f6'><div style='font-size:24px'>{tot_count}</div><div style='font-size:10px; color:#9ca3af'>Total Queries</div></div>"
+    thr_html = f"<div class='card' style='border-color:#ef4444'><div style='font-size:24px'>{thr_count}</div><div style='font-size:10px; color:#9ca3af'>Threats Blocked</div></div>"
+    saf_html = f"<div class='card' style='border-color:#10b981'><div style='font-size:24px'>{saf_count}</div><div style='font-size:10px; color:#9ca3af'>Safe Queries</div></div>"
+    
+    out_box = f"<div style='padding: 15px; border-radius: 8px; background: #161b22; border: 1px solid {color}; color: {color}; text-align: center; font-weight: bold;'>{ui_status}</div>"
+
+    # Return order MUST match the outputs list below
+    return tot_count, thr_count, saf_count, "", out_box, updated_log, updated_log, tot_html, thr_html, saf_html
+
+
+# --- 3. THE UI WIRING ---
 with gr.Blocks(css=custom_css) as demo:
+    # State variables to track numbers
+    tot_state = gr.State(0)
+    thr_state = gr.State(0)
+    saf_state = gr.State(0)
+    log_state = gr.State("")
+
     with gr.Row():
-        # --- THE FAKE SIDEBAR ---
+        # --- SIDEBAR ---
         with gr.Column(scale=1, elem_classes="fake-sidebar"):
             gr.HTML("<div style='font-size: 40px; text-align:center;'>🛡️</div>")
             gr.Markdown("## **INTELLIGUARD**\nSECURITY PLATFORM")
             gr.HTML("<div class='status-badge'>● INTELLIGUARD PROTECTED</div>")
             
-            gr.Markdown("### 📋 AUDIT LOG\n<p style='color: #6b7280;'>No activity recorded.</p>")
+            gr.Markdown("### 📋 AUDIT LOG")
+            audit_log_ui = gr.HTML("<p style='color: #6b7280;'>No activity recorded.</p>")
+            
             gr.Markdown("---")
             gr.Markdown("### 📄 DOCUMENT UPLOAD")
             gr.File(label=None)
 
-        # --- THE MAIN CONTENT ---
+        # --- MAIN CONTENT ---
         with gr.Column(scale=4):
             gr.HTML("""
                 <div style="text-align: center; margin-top: 20px;">
@@ -43,20 +107,35 @@ with gr.Blocks(css=custom_css) as demo:
                 </div>
             """)
 
+            # Dynamic Metric Cards
             with gr.Row():
-                gr.HTML("<div class='card' style='border-color:#3b82f6'><div style='font-size:24px'>0</div><div style='font-size:10px; color:#9ca3af'>Total Queries</div></div>")
-                gr.HTML("<div class='card' style='border-color:#ef4444'><div style='font-size:24px'>0</div><div style='font-size:10px; color:#9ca3af'>Threats Blocked</div></div>")
-                gr.HTML("<div class='card' style='border-color:#10b981'><div style='font-size:24px'>0</div><div style='font-size:10px; color:#9ca3af'>Safe Queries</div></div>")
+                tot_ui = gr.HTML("<div class='card' style='border-color:#3b82f6'><div style='font-size:24px'>0</div><div style='font-size:10px; color:#9ca3af'>Total Queries</div></div>")
+                thr_ui = gr.HTML("<div class='card' style='border-color:#ef4444'><div style='font-size:24px'>0</div><div style='font-size:10px; color:#9ca3af'>Threats Blocked</div></div>")
+                saf_ui = gr.HTML("<div class='card' style='border-color:#10b981'><div style='font-size:24px'>0</div><div style='font-size:10px; color:#9ca3af'>Safe Queries</div></div>")
 
             with gr.Accordion("📧 Enterprise Email Gateway", open=True):
-                email = gr.Textbox(label="Target Email", value="security-ops@techcorp.com")
+                email = gr.Textbox(label="Target Email", value="security-ops@techcorp.com", interactive=True)
                 connect_btn = gr.Button("🔗 Establish Secure Connection", elem_classes="establish-btn")
 
-            gr.Markdown("<div style='height: 150px;'></div>")
+            gr.Markdown("<div style='height: 80px;'></div>")
+            
+            # Action Area
+            status_box = gr.HTML("<div style='padding: 15px; text-align: center; color: #6b7280;'>Awaiting Payload...</div>")
             
             with gr.Row():
-                chat = gr.Textbox(show_label=False, placeholder="Ask about company policies...", scale=10)
+                chat = gr.Textbox(show_label=False, placeholder="Ask about company policies...", scale=10, interactive=True)
                 submit = gr.Button("↑", scale=1, variant="primary")
 
-# No extra arguments - HF handles the server name and port automatically
-demo.launch()
+    # Connect the UI to the Python Function
+    submit.click(
+        fn=process_query,
+        inputs=[chat, tot_state, thr_state, saf_state, log_state],
+        outputs=[tot_state, thr_state, saf_state, chat, status_box, log_state, audit_log_ui, tot_ui, thr_ui, saf_ui]
+    )
+    chat.submit(
+        fn=process_query,
+        inputs=[chat, tot_state, thr_state, saf_state, log_state],
+        outputs=[tot_state, thr_state, saf_state, chat, status_box, log_state, audit_log_ui, tot_ui, thr_ui, saf_ui]
+    )
+
+demo.launch()