Update app.py

app.py

@@ -3,7 +3,7 @@ import spaces
 import json
 import datetime
 import random
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import torch
 import time
 
@@ -26,40 +26,91 @@ custom_css = """
     padding: 10px;
     border-radius: 5px;
 }
-.status-warning {
-    background: #fff3cd;
-    border: 1px solid #ffeaa7;
-    color: #856404;
-    padding: 10px;
-    border-radius: 5px;
+.gpt-oss-badge {
+    background: linear-gradient(45deg, #00c6ff, #0072ff);
+    color: white;
+    padding: 5px 10px;
+    border-radius: 15px;
+    font-weight: bold;
 }
 """
 
-# Initialize the LLM pipeline with zeroGPU support
+# Global variables for model management
+model = None
+tokenizer = None
+model_status = "🔄 Initializing GPT-OSS-20B..."
+
+# Initialize GPT-OSS-20B with proper harmony format
 @spaces.GPU
-def initialize_llm():
+def initialize_gpt_oss():
+    """Initialize OpenAI GPT-OSS-20B with harmony response format"""
+    global model, tokenizer, model_status
+    
     try:
+        model_id = "openai/gpt-oss-20b"
+        print(f"🚀 Loading {model_id}...")
+        
         # Check GPU availability
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        print(f"Using device: {device}")
+        print(f"Device: {device}")
+        
+        if torch.cuda.is_available():
+            print(f"GPU: {torch.cuda.get_device_name()}")
+            print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f}GB")
+        
+        # Load tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+        print("✅ Tokenizer loaded")
+        
+        # Load model with optimized settings for zeroGPU
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype="auto",  # Let it choose best dtype (MXFP4)
+            device_map="auto",   # Automatic GPU placement
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+            # MXFP4 quantization is built-in
+        )
         
-        # Try to use a larger model with GPU acceleration
-        model_id = "microsoft/DialoGPT-medium"
-        pipe = pipeline(
-            "text-generation",
-            model=model_id,
-            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-            device_map="auto" if device == "cuda" else "cpu",
-            max_length=512,
-            pad_token_id=50256
+        print("✅ Model loaded with MXFP4 quantization")
+        print(f"Model device: {next(model.parameters()).device}")
+        
+        # Test generation to ensure everything works
+        test_messages = [
+            {"role": "user", "content": "Hello, test message."}
+        ]
+        
+        test_inputs = tokenizer.apply_chat_template(
+            test_messages,
+            add_generation_prompt=True,
+            return_tensors="pt",
+            return_dict=True,
         )
-        return pipe, f"✅ LLM Model loaded on {device}: {model_id}"
+        
+        if device == "cuda":
+            test_inputs = {k: v.to(model.device) for k, v in test_inputs.items()}
+        
+        with torch.no_grad():
+            test_output = model.generate(
+                **test_inputs,
+                max_new_tokens=10,
+                do_sample=False,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        
+        print("✅ Test generation successful")
+        
+        model_status = f"✅ OpenAI GPT-OSS-20B loaded successfully on {device} | MXFP4 Quantized | ~16GB Memory"
+        return model_status
+        
     except Exception as e:
-        return None, f"⚠️ LLM not available: {str(e)[:100]}... Using fallback analysis."
-
-pipe, model_status = initialize_llm()
+        error_msg = f"❌ Failed to load GPT-OSS-20B: {str(e)}"
+        print(error_msg)
+        model_status = error_msg
+        model, tokenizer = None, None
+        return model_status
 
-# Enhanced attack scenarios with more realistic data
+# Enhanced attack scenarios
 ATTACK_SCENARIOS = {
     "🔄 Lateral Movement": {
         "description": "Advanced Persistent Threat (APT) - Attacker moving laterally through network after initial compromise",
@@ -67,7 +118,7 @@ ATTACK_SCENARIOS = {
         "alerts": [
             {
                 "id": "ALR-001",
-                "timestamp": "2025-01-15 14:30:45",
+                "timestamp": "2025-08-10 14:30:45",
                 "source_ip": "192.168.1.100",
                 "destination_ip": "192.168.1.25", 
                 "user": "corp\\john.doe",
@@ -81,7 +132,7 @@ ATTACK_SCENARIOS = {
             },
             {
                 "id": "ALR-002", 
-                "timestamp": "2025-01-15 14:35:12",
+                "timestamp": "2025-08-10 14:35:12",
                 "source_ip": "192.168.1.100",
                 "destination_ip": "192.168.1.50",
                 "user": "corp\\john.doe",
@@ -95,7 +146,7 @@ ATTACK_SCENARIOS = {
             },
             {
                 "id": "ALR-003",
-                "timestamp": "2025-01-15 14:42:18", 
+                "timestamp": "2025-08-10 14:42:18", 
                 "source_ip": "192.168.1.100",
                 "destination_ip": "10.0.0.15",
                 "user": "SYSTEM",
@@ -115,7 +166,7 @@ ATTACK_SCENARIOS = {
         "alerts": [
             {
                 "id": "ALR-004",
-                "timestamp": "2025-01-15 09:15:30",
+                "timestamp": "2025-08-10 09:15:30",
                 "source_ip": "203.0.113.50",
                 "destination_ip": "192.168.1.75",
                 "user": "corp\\sarah.wilson",
@@ -129,7 +180,7 @@ ATTACK_SCENARIOS = {
             },
             {
                 "id": "ALR-005",
-                "timestamp": "2025-01-15 09:45:22",
+                "timestamp": "2025-08-10 09:45:22",
                 "source_ip": "192.168.1.75", 
                 "destination_ip": "203.0.113.50",
                 "user": "corp\\sarah.wilson",
@@ -149,7 +200,7 @@ ATTACK_SCENARIOS = {
         "alerts": [
             {
                 "id": "ALR-006",
-                "timestamp": "2025-01-15 16:20:10",
+                "timestamp": "2025-08-10 16:20:10",
                 "source_ip": "192.168.1.85",
                 "destination_ip": "192.168.1.85",
                 "user": "corp\\admin.backup",
@@ -163,7 +214,7 @@ ATTACK_SCENARIOS = {
             },
             {
                 "id": "ALR-007",
-                "timestamp": "2025-01-15 16:25:33",
+                "timestamp": "2025-08-10 16:25:33",
                 "source_ip": "192.168.1.85",
                 "destination_ip": "45.33.22.11", 
                 "user": "SYSTEM",
@@ -180,179 +231,213 @@ ATTACK_SCENARIOS = {
 }
 
 @spaces.GPU
-def generate_advanced_llm_analysis(alert_data, analyst_level):
-    """Generate comprehensive LLM-based analysis with enhanced prompting and GPU acceleration"""
+def generate_gpt_oss_analysis(alert_data, analyst_level):
+    """Generate analysis using OpenAI GPT-OSS-20B with harmony format"""
     
-    # Enhanced context with more structured prompting
-    system_context = f"""You are an expert cybersecurity analyst assistant specializing in SOC operations. 
-    Analyze the following security alert for a Level {analyst_level} analyst.
-
-    ALERT CONTEXT:
-    ID: {alert_data['id']}
-    Type: {alert_data['alert_type']} 
-    Severity: {alert_data['severity']}
-    Timestamp: {alert_data['timestamp']}
-    Network: {alert_data['source_ip']} → {alert_data['destination_ip']}
-    User: {alert_data['user']}
-    Description: {alert_data['description']}
-    Technical Details: {alert_data['raw_log']}
-    Threat Intelligence: {alert_data['threat_intel']}
-    MITRE ATT&CK: {alert_data['mitre_tactic']}
-    Confidence: {alert_data['confidence']}%
-
-    Provide analysis appropriate for {analyst_level} level:"""
-
-    if pipe:
-        try:
-            # Use GPU acceleration for faster inference
-            device = next(pipe.model.parameters()).device
-            print(f"LLM running on device: {device}")
-            
-            prompt = f"{system_context}\n\nAnalysis:"
-            response = pipe(
-                prompt, 
-                max_new_tokens=300, 
-                do_sample=True, 
-                temperature=0.7, 
+    if not model or not tokenizer:
+        return get_fallback_analysis(alert_data, analyst_level)
+    
+    # Enhanced prompts designed for GPT-OSS reasoning capabilities
+    security_prompts = {
+        "L1": f"""You are a Level 1 SOC analyst conducting initial triage. Analyze this security alert and provide immediate actionable recommendations.
+
+**SECURITY ALERT:**
+- ID: {alert_data['id']}
+- Type: {alert_data['alert_type']} 
+- Severity: {alert_data['severity']}
+- Source: {alert_data['source_ip']} → {alert_data['destination_ip']}
+- User: {alert_data['user']}
+- Evidence: {alert_data['raw_log']}
+- Intel: {alert_data['threat_intel']}
+- MITRE ATT&CK: {alert_data['mitre_tactic']}
+- Confidence: {alert_data['confidence']}%
+
+**PROVIDE L1 TRIAGE:**
+1. Immediate containment actions
+2. Risk assessment
+3. Escalation decision with reasoning
+4. Priority timeline
+
+Think step-by-step about the threat level and required response.""",
+
+        "L2": f"""You are a Level 2 SOC analyst conducting detailed investigation. Perform comprehensive analysis of this cybersecurity incident.
+
+**INCIDENT DETAILS:**
+- Alert: {alert_data['alert_type']} | Severity: {alert_data['severity']}
+- Network Flow: {alert_data['source_ip']} → {alert_data['destination_ip']}
+- User Context: {alert_data['user']}
+- Technical Evidence: {alert_data['raw_log']}
+- Threat Intelligence: {alert_data['threat_intel']}
+- MITRE ATT&CK Technique: {alert_data['mitre_tactic']}
+- Detection Confidence: {alert_data['confidence']}%
+
+**CONDUCT L2 INVESTIGATION:**
+1. Technical root cause analysis
+2. Evidence correlation and timeline
+3. Threat actor behavior analysis
+4. Impact assessment and containment strategy
+5. Investigation roadmap
+
+Use chain-of-thought reasoning to analyze the attack progression and recommend next steps.""",
+
+        "L3": f"""You are a senior cybersecurity expert analyzing a sophisticated threat. Provide strategic assessment and executive-level recommendations.
+
+**THREAT INTELLIGENCE:**
+- Attack Vector: {alert_data['description']}
+- Technical Indicators: {alert_data['raw_log']}
+- Attribution Context: {alert_data['threat_intel']}
+- MITRE Technique: {alert_data['mitre_tactic']}
+- Confidence Level: {alert_data['confidence']}%
+
+**DELIVER L3 EXPERT ANALYSIS:**
+1. Adversary attribution and campaign analysis
+2. Strategic threat landscape assessment
+3. Business impact and risk quantification
+4. Comprehensive response strategy
+5. Executive briefing points
+
+Apply deep reasoning to assess the broader implications and provide strategic recommendations."""
+    }
+    
+    try:
+        prompt = security_prompts.get(analyst_level, security_prompts["L2"])
+        
+        # Use proper harmony format for chat
+        messages = [
+            {"role": "user", "content": prompt}
+        ]
+        
+        # Apply chat template (automatically uses harmony format)
+        inputs = tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            return_tensors="pt",
+            return_dict=True,
+        )
+        
+        # Move to device if using GPU
+        if torch.cuda.is_available():
+            inputs = {k: v.to(model.device) for k, v in inputs.items()}
+        
+        # Generate with optimized parameters for reasoning
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=500,
+                do_sample=True,
+                temperature=0.2,  # Lower for focused analysis
                 top_p=0.9,
-                pad_token_id=pipe.tokenizer.eos_token_id
+                top_k=50,
+                repetition_penalty=1.1,
+                pad_token_id=tokenizer.eos_token_id,
+                eos_token_id=tokenizer.eos_token_id
             )
-            generated_text = response[0]['generated_text']
-            analysis = generated_text[len(prompt):].strip()
-            return analysis if analysis else get_fallback_analysis(alert_data, analyst_level)
-        except Exception as e:
-            print(f"LLM Error: {e}")
-            return f"LLM Processing Error: {str(e)}\n\n{get_fallback_analysis(alert_data, analyst_level)}"
-    
-    return get_fallback_analysis(alert_data, analyst_level)
+        
+        # Decode the response
+        input_length = inputs["input_ids"].shape[-1]
+        generated_tokens = outputs[0][input_length:]
+        analysis = tokenizer.decode(generated_tokens, skip_special_tokens=True)
+        
+        # Ensure quality
+        if len(analysis.strip()) < 100:
+            return get_fallback_analysis(alert_data, analyst_level)
+        
+        return f"""🤖 **OpenAI GPT-OSS-20B Analysis**
+<div class="gpt-oss-badge">Powered by GPT-OSS-20B • MoE Architecture • MXFP4 Quantized</div>
+
+━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+
+{analysis.strip()}
+
+━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+*Analysis generated using OpenAI's latest open-weight reasoning model*
+*21B parameters • 3.6B active per token • Apache 2.0 licensed*"""
+        
+    except Exception as e:
+        print(f"GPT-OSS Error: {e}")
+        return f"⚠️ GPT-OSS Error: {str(e)[:100]}\n\n{get_fallback_analysis(alert_data, analyst_level)}"
 
 def get_fallback_analysis(alert_data, analyst_level):
-    """Enhanced fallback analysis with detailed recommendations"""
+    """High-quality fallback when model fails"""
     
-    base_analysis = {
-        "L1": {
-            "icon": "🚨",
-            "title": "L1 TRIAGE ANALYSIS",
-            "focus": "Initial Assessment & Escalation",
-            "template": """
-{icon} {title}
+    templates = {
+        "L1": f"""🚨 **L1 SOC TRIAGE ANALYSIS**
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-🎯 THREAT SUMMARY: {alert_type} - {severity} severity
-⏰ OCCURRED: {timestamp}
-🌐 AFFECTED SYSTEM: {source_ip} (User: {user})
-🔍 CONFIDENCE LEVEL: {confidence}%
-
-🚀 IMMEDIATE ACTIONS:
-• Isolate affected system: {source_ip}
-• Verify user account status: {user}
-• Check for similar alerts in timeframe
-• Document incident ID: {id}
-
-⬆️ ESCALATION CRITERIA: 
-• Severity: {severity} - Meets L2 escalation threshold
-• MITRE Tactic: {mitre_tactic}
-• Recommend immediate L2 review
-
-📋 INITIAL NOTES:
-{threat_intel}
-            """
-        },
-        "L2": {
-            "icon": "🔍", 
-            "title": "L2 INVESTIGATION ANALYSIS",
-            "focus": "Detailed Investigation & Correlation",
-            "template": """
-{icon} {title}
+
+**🎯 THREAT ASSESSMENT:**
+Alert: {alert_data['alert_type']} | Severity: {alert_data['severity']}
+Confidence: {alert_data['confidence']}% | Source: {alert_data['source_ip']}
+
+**⚡ IMMEDIATE ACTIONS:**
+1. Isolate affected system: {alert_data['source_ip']}
+2. Disable user account: {alert_data['user']}
+3. Block connections to: {alert_data['destination_ip']}
+4. Preserve evidence for investigation
+
+**⬆️ ESCALATION DECISION:**
+Severity: {alert_data['severity']} → ESCALATE TO L2
+Technique: {alert_data['mitre_tactic']} requires deeper analysis
+
+**📝 INITIAL ASSESSMENT:**
+{alert_data['threat_intel']}""",
+
+        "L2": f"""🔍 **L2 INVESTIGATION ANALYSIS**
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-🎯 ATTACK VECTOR: {description}
-⚙️ TECHNICAL DETAILS: {raw_log}
-🧠 THREAT CONTEXT: {threat_intel}
-🎪 MITRE ATT&CK: {mitre_tactic}
-
-🔬 INVESTIGATION STEPS:
-1. Examine parent process tree for {source_ip}
-2. Correlate network connections in ±30min window  
-3. Review authentication logs for user: {user}
-4. Check for indicators across environment
-5. Analyze file system changes (if applicable)
-
-🎯 CORRELATION POINTS:
-• Source IP timeline analysis
-• User behavior baseline comparison
-• Similar TTPs in recent incidents
-• Network segmentation verification
-
-📊 RISK ASSESSMENT:
-• Technical Impact: {severity}
-• Business Risk: Review asset criticality
-• Containment Priority: High (based on {confidence}% confidence)
-
-⬆️ L3 ESCALATION IF:
-• Attack campaign indicators found
-• Critical asset involvement confirmed
-• Advanced persistent threat suspected
-            """
-        },
-        "L3": {
-            "icon": "🎯",
-            "title": "L3 EXPERT ANALYSIS", 
-            "focus": "Attribution & Strategic Response",
-            "template": """
-{icon} {title}
+
+**🎯 ATTACK VECTOR ANALYSIS:**
+Technique: {alert_data['mitre_tactic']}
+Evidence: {alert_data['raw_log']}
+Context: {alert_data['description']}
+
+**🔬 INVESTIGATION ROADMAP:**
+1. Timeline correlation: ±30min from {alert_data['timestamp']}
+2. User behavior analysis: {alert_data['user']} baseline comparison
+3. Network flow analysis: {alert_data['source_ip']} → {alert_data['destination_ip']}
+4. Process tree examination: Parent/child relationships
+5. Artifact collection: Memory dumps, logs, files
+
+**📊 THREAT ASSESSMENT:**
+Confidence Level: {alert_data['confidence']}%
+Business Impact: {alert_data['severity']}
+Attribution: {alert_data['threat_intel']}
+
+**🎯 RECOMMENDATIONS:**
+Deploy hunting queries for similar TTPs
+Review authentication logs for compromise
+Consider L3 escalation if campaign indicators found""",
+
+        "L3": f"""🎯 **L3 EXPERT STRATEGIC ANALYSIS**
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-🎭 ADVERSARY PROFILE: Advanced threat actor
-🎪 CAMPAIGN ANALYSIS: {threat_intel}
-💼 BUSINESS IMPACT: {severity} - Requires C-level awareness
-🛡️ DEFENSIVE POSTURE: Enhanced monitoring required
-
-🕵️ THREAT HUNTING PRIORITIES:
-1. Memory forensics on {source_ip}
-2. Network traffic deep packet inspection
-3. Endpoint artifact preservation
-4. Active Directory security log analysis
-5. Cloud infrastructure review (if applicable)
-
-🎯 ATTRIBUTION INDICATORS:
-• TTPs match: {mitre_tactic}
-• Technical sophistication: High
-• Targeting pattern: [Analyze organizational profile]
-• Infrastructure overlap: Review IOC databases
-
-🛠️ MITIGATION STRATEGY:
-• Immediate: Block C2 communications
-• Short-term: Deploy hunting queries 
-• Medium-term: Security architecture review
-• Long-term: Staff training and awareness
-
-📈 EXECUTIVE BRIEFING POINTS:
+
+**🎭 ADVERSARY ASSESSMENT:**
+Sophistication: Advanced (based on {alert_data['mitre_tactic']})
+Campaign Context: {alert_data['threat_intel']}
+Success Probability: {alert_data['confidence']}%
+
+**💼 BUSINESS IMPACT:**
+Severity Level: {alert_data['severity']}
+Executive Notification: Required
+Regulatory Implications: Under review
+
+**🛡️ STRATEGIC RESPONSE:**
+Immediate: Threat hunting deployment across environment
+Short-term: Enhanced monitoring and detection rules
+Medium-term: Security architecture review
+Long-term: Threat intelligence integration enhancement
+
+**📈 EXECUTIVE BRIEFING:**
 • Sophisticated attack requiring coordinated response
-• Potential for lateral movement and data exfiltration
+• High potential for lateral movement and data exfiltration
 • Recommend incident response team activation
-• Consider external forensics support
-
-🔮 PREDICTIVE ANALYSIS:
-• High probability of follow-up attacks
-• Recommend 48-72 hour enhanced monitoring
-• Consider threat landscape implications
-            """
-        }
+• Consider external forensics support engagement"""
     }
     
-    if analyst_level in base_analysis:
-        template = base_analysis[analyst_level]["template"]
-        return template.format(
-            icon=base_analysis[analyst_level]["icon"],
-            title=base_analysis[analyst_level]["title"],
-            **alert_data
-        )
-    
-    return "Analysis not available for specified level."
+    return templates.get(analyst_level, templates["L2"])
 
-def analyze_alert_comprehensive(scenario_name, alert_index, analyst_level):
-    """Enhanced main analysis function with timing and status updates"""
+def analyze_alert_with_gpt_oss(scenario_name, alert_index, analyst_level):
+    """Main analysis function using GPT-OSS-20B"""
     start_time = time.time()
     
-    # Validate inputs
     if scenario_name not in ATTACK_SCENARIOS:
         return "❌ Invalid scenario selected.", "", "Error: Invalid scenario"
     
@@ -364,94 +449,90 @@ def analyze_alert_comprehensive(scenario_name, alert_index, analyst_level):
     
     selected_alert = alerts[alert_index]
     
-    # Generate comprehensive analysis
-    analysis = generate_advanced_llm_analysis(selected_alert, analyst_level)
+    # Generate analysis using GPT-OSS-20B
+    analysis = generate_gpt_oss_analysis(selected_alert, analyst_level)
     
-    # Enhanced alert details formatting
-    alert_details = f"""
-🎫 ALERT ID: {selected_alert['id']} | 🕐 {selected_alert['timestamp']}
-━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+    # Format alert details
+    alert_details = f"""🎫 **ALERT {selected_alert['id']}** | 🕐 {selected_alert['timestamp']}
+━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 
-🌐 NETWORK FLOW:
+🌐 **NETWORK FLOW:**
    Source: {selected_alert['source_ip']} → Destination: {selected_alert['destination_ip']}
    
-👤 USER CONTEXT:
+👤 **USER CONTEXT:**
    Account: {selected_alert['user']}
    
-⚠️ ALERT CLASSIFICATION:
+⚠️ **CLASSIFICATION:**
    Type: {selected_alert['alert_type']}
    Severity: {selected_alert['severity']}
    Confidence: {selected_alert['confidence']}%
    
-📝 DESCRIPTION:
+📝 **DESCRIPTION:**
    {selected_alert['description']}
    
-🔍 TECHNICAL EVIDENCE:
+🔍 **TECHNICAL EVIDENCE:**
    {selected_alert['raw_log']}
    
-🧠 THREAT INTELLIGENCE:
+🧠 **THREAT INTELLIGENCE:**
    {selected_alert['threat_intel']}
    
-🎪 MITRE ATT&CK MAPPING:
+🎪 **MITRE ATT&CK:**
    {selected_alert['mitre_tactic']}
 
-━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-    """
+━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"""
     
     processing_time = round(time.time() - start_time, 2)
-    status_message = f"✅ {analyst_level} analysis completed in {processing_time}s | Model: {model_status}"
+    device_info = "GPU" if torch.cuda.is_available() else "CPU"
+    status = f"✅ {analyst_level} analysis completed in {processing_time}s | Device: {device_info} | {model_status}"
     
-    return alert_details, analysis, status_message
+    return alert_details, analysis, status
 
-def get_enhanced_scenario_info(scenario_name):
-    """Enhanced scenario information with threat overview"""
+def get_scenario_info(scenario_name):
+    """Get scenario information"""
     if scenario_name in ATTACK_SCENARIOS:
         scenario = ATTACK_SCENARIOS[scenario_name]
         
-        info = f"""
-## 🎭 **Attack Scenario: {scenario_name}**
+        info = f"""## 🎭 **Attack Scenario: {scenario_name}**
 
 **📋 Description:** {scenario['description']}
-**⚠️ Severity Level:** {scenario['severity']}
-**📊 Total Alerts:** {len(scenario['alerts'])} security events detected
+**⚠️ Severity:** {scenario['severity']}
+**📊 Total Alerts:** {len(scenario['alerts'])} security events
 
 ### 🔍 **Alert Timeline:**
 """
         
         for i, alert in enumerate(scenario['alerts']):
-            info += f"""
-**[{i+1}] {alert['timestamp']}** - {alert['alert_type']}
+            info += f"""**[{i+1}] {alert['timestamp']}** - {alert['alert_type']}
    └─ Severity: {alert['severity']} | Confidence: {alert['confidence']}%
 """
         
-        info += f"""
-### 🎯 **Analysis Capabilities:**
-- **L1 Triage:** Initial assessment and escalation decisions
-- **L2 Investigation:** Detailed technical analysis and correlation  
-- **L3 Expert:** Attribution, impact assessment, and strategic response
-        """
+        info += """
+### 🤖 **AI-Powered Analysis:**
+- **OpenAI GPT-OSS-20B:** Latest open-weight reasoning model
+- **MXFP4 Quantization:** Optimized for efficient inference
+- **Harmony Format:** Advanced response structure
+- **21B Parameters:** With 3.6B active per token (MoE)"""
         
         return info
-    return "⚠️ No scenario selected. Please choose an attack scenario to begin analysis."
+    return "⚠️ No scenario selected."
 
-# Create enhanced Gradio interface
-with gr.Blocks(title="SOC LLM Assistant - Advanced PoC", theme=gr.themes.Soft(), css=custom_css) as demo:
+# Create Gradio interface
+with gr.Blocks(title="SOC Assistant - GPT-OSS-20B", theme=gr.themes.Soft(), css=custom_css) as demo:
     
-    # Header
     gr.Markdown("""
-    # 🛡️ SOC LLM Assistant - Advanced Proof of Concept
-    **Intelligent Security Alert Analysis for Multi-Level SOC Operations**
+    # 🛡️ SOC LLM Assistant - OpenAI GPT-OSS-20B Edition
+    **Powered by OpenAI's Latest Open-Weight Reasoning Model**
     
-    *Demonstrating LLM-powered assistance for L1, L2, and L3 security analysts*
+    *First open-weight model from OpenAI since GPT-2 • Released August 8, 2025*
     """)
     
     # Model status display
-    gr.Markdown(f"🤖 **System Status:** {model_status}")
+    status_display = gr.Markdown("🔄 Loading OpenAI GPT-OSS-20B...")
     
     with gr.Row():
-        # Left Panel - Controls
+        # Left Panel
         with gr.Column(scale=1, min_width=300):
-            gr.Markdown("## 🎮 Attack Simulation Control")
+            gr.Markdown("## 🎮 Attack Simulation")
             
             scenario_dropdown = gr.Dropdown(
                 choices=list(ATTACK_SCENARIOS.keys()),
@@ -470,91 +551,86 @@ with gr.Blocks(title="SOC LLM Assistant - Advanced PoC", theme=gr.themes.Soft(),
                 maximum=2,
                 step=1,
                 value=0,
-                label="📋 Alert Selection",
-                info="Choose which alert from the scenario to analyze"
+                label="📋 Alert Selection"
             )
             
             analyst_level = gr.Radio(
                 choices=["L1", "L2", "L3"],
                 label="👤 Analyst Level",
                 value="L2",
-                info="L1: Triage | L2: Investigation | L3: Expert Analysis"
+                info="L1: Triage | L2: Investigation | L3: Expert"
             )
             
             analyze_btn = gr.Button(
-                "🔍 Analyze Alert", 
+                "🚀 Analyze with GPT-OSS-20B", 
                 variant="primary",
                 size="lg"
             )
             
+            init_btn = gr.Button(
+                "🔄 Reinitialize Model",
+                variant="secondary"
+            )
+            
             gr.Markdown("---")
-            gr.Markdown("## 📊 Quick Stats")
+            gr.Markdown("## 🤖 Model Information")
             gr.Markdown("""
-            **🎯 Demo Features:**
-            - 3 realistic attack scenarios
-            - Multi-level analysis (L1/L2/L3)
-            - MITRE ATT&CK mapping
-            - Threat intelligence integration
-            - Real-time LLM processing
+            **🎯 GPT-OSS-20B Features:**
+            - 21B parameters (3.6B active)
+            - MXFP4 quantization
+            - 128K context length
+            - Apache 2.0 licensed
+            - Harmony response format
+            - Reasoning capabilities
             """)
         
-        # Right Panel - Results
+        # Right Panel
         with gr.Column(scale=2):
             gr.Markdown("## 📋 Security Alert Details")
             alert_output = gr.Textbox(
-                label="🎫 Raw Alert Information",
+                label="🎫 Alert Information",
                 lines=15,
-                interactive=False,
-                placeholder="Alert details will appear here after analysis..."
+                interactive=False
             )
             
-            gr.Markdown("## 🤖 AI-Powered Analysis")
+            gr.Markdown("## 🤖 GPT-OSS-20B Analysis")
             analysis_output = gr.Textbox(
-                label="🧠 Intelligent Analysis & Recommendations",
-                lines=20,
-                interactive=False,
-                placeholder="LLM analysis will appear here after processing..."
+                label="🧠 AI-Powered Security Analysis",
+                lines=25,
+                interactive=False
             )
             
             status_output = gr.Textbox(
-                label="📊 Processing Status", 
-                interactive=False,
-                lines=1
+                label="📊 Processing Status",
+                lines=1,
+                interactive=False
             )
     
-    # Footer information
     gr.Markdown("""
     ---
-    ## 📖 **Usage Instructions:**
+    ## 🎉 **About OpenAI GPT-OSS-20B**
     
-    1. **📊 Select Scenario:** Choose from realistic cybersecurity attack scenarios
-    2. **🎯 Pick Alert:** Use the slider to select which alert in the sequence to analyze  
-    3. **👤 Choose Level:** Select analyst expertise level (L1/L2/L3) for tailored analysis
-    4. **🔍 Analyze:** Click the analyze button to get AI-powered insights and recommendations
+    Released August 8, 2025 - OpenAI's first open-weight model since GPT-2! This groundbreaking release features:
     
-    ## 🎯 **Key Capabilities Demonstrated:**
+    - **🧠 Advanced Reasoning:** Comparable to o3-mini performance
+    - **⚡ Efficient Architecture:** MoE with only 3.6B active parameters per token
+    - **🔧 Harmony Format:** New structured response system for better tool use
+    - **📱 Consumer Hardware:** Runs on just 16GB memory
+    - **🔓 Open License:** Apache 2.0 - fully permissive for commercial use
     
-    - **🎭 Realistic Scenarios:** Based on actual cybersecurity incidents and attack patterns
-    - **🧠 Contextual Analysis:** LLM considers all available metadata, threat intelligence, and historical patterns
-    - **👥 Role-Based Insights:** Tailored recommendations for different SOC analyst skill levels
-    - **⚡ Real-Time Processing:** Immediate analysis with actionable next steps
-    - **🎪 Industry Standards:** MITRE ATT&CK framework integration for standardized threat classification
-    
-    ## 🔬 **Research Value:**
-    This PoC demonstrates the feasibility of LLM integration in operational security environments, supporting research in automated threat analysis, human-AI collaboration, and intelligent SOC operations.
+    Perfect for cybersecurity analysis requiring sophisticated reasoning and chain-of-thought capabilities!
     
     ---
-    **👨‍🎓 Developed by:** Abdullah Alanazi | **🏛️ Institution:** KAUST | **👨‍🏫 Supervisor:** Prof. Ali Shoker
+    **👨‍🎓 Research:** Abdullah Alanazi | **🏛️ KAUST** | **👨‍🏫 Prof. Ali Shoker**
     """)
     
-    # Event handlers with enhanced functionality
+    # Event handlers
     scenario_dropdown.change(
-        fn=get_enhanced_scenario_info,
+        fn=get_scenario_info,
         inputs=[scenario_dropdown],
         outputs=[scenario_info]
     )
     
-    # Update slider maximum based on scenario
     def update_slider_max(scenario_name):
         if scenario_name in ATTACK_SCENARIOS:
             max_alerts = len(ATTACK_SCENARIOS[scenario_name]["alerts"]) - 1
@@ -568,19 +644,28 @@ with gr.Blocks(title="SOC LLM Assistant - Advanced PoC", theme=gr.themes.Soft(),
     )
     
     analyze_btn.click(
-        fn=analyze_alert_comprehensive,
+        fn=analyze_alert_with_gpt_oss,
         inputs=[scenario_dropdown, alert_slider, analyst_level],
         outputs=[alert_output, analysis_output, status_output]
     )
     
-    # Initialize with default scenario
+    init_btn.click(
+        fn=initialize_gpt_oss,
+        outputs=[status_display]
+    )
+    
+    # Initialize on startup
     demo.load(
-        fn=get_enhanced_scenario_info,
+        fn=get_scenario_info,
         inputs=[scenario_dropdown],
         outputs=[scenario_info]
     )
+    
+    demo.load(
+        fn=initialize_gpt_oss,
+        outputs=[status_display]
+    )
 
-# Launch configuration
 if __name__ == "__main__":
     demo.launch(
         share=True,