app.py

from flask import Flask, render_template, request, jsonify
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import re
import gc
import json
import threading
import time
from datetime import datetime
from typing import Dict, List, Optional
from pydantic import BaseModel
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = Flask(__name__)

# Global variables for model and tokenizer
model = None
tokenizer = None
model_loaded = False
loading_status = "Initializing AEGIS BIO LAB 10 CONDUCTOR Multi-Domain Expert System..."

# AEGIS BIO LAB 10 CONDUCTOR Configuration
MODEL_REPO = "Gaston895/aegisconduct"
AEGIS_VERSION = "10.0"
GLOBAL_REGIONS = [
    "North America", "Europe", "Asia", "Africa", 
    "South America", "Middle East", "Oceania", "Arctic Region"
]

def clear_memory():
    """Clear memory aggressively"""
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

def load_model():
    """Load the AEGIS Conduct model with 8-bit quantization"""
    global model, tokenizer, model_loaded, loading_status
    
    try:
        # Initial memory cleanup
        clear_memory()
        
        loading_status = "Loading AEGIS BIO LAB 10 CONDUCTOR tokenizer..."
        print("🔄 Loading AEGIS BIO LAB 10 CONDUCTOR tokenizer...")
        
        # Load tokenizer first
        tokenizer = AutoTokenizer.from_pretrained(
            MODEL_REPO,
            trust_remote_code=True,
            use_fast=True
        )
        
        # Set padding token if not set
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        
        loading_status = "Configuring 8-bit quantization (16GB → 8GB)..."
        print("⚙️ Configuring 8-bit quantization (reduces 16GB model to ~8GB)...")
        
        # Configure 8-bit quantization for CPU
        quantization_config = BitsAndBytesConfig(
            load_in_8bit=True,
            llm_int8_threshold=6.0,
            llm_int8_has_fp16_weight=False,
            bnb_8bit_compute_dtype=torch.float32,  # CPU compatible
            bnb_8bit_use_double_quant=False,
        )
        
        loading_status = "Loading 16GB model with 8-bit quantization..."
        print("📥 Loading 16GB model with 8-bit quantization...")
        
        # Load model with 8-bit quantization
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_REPO,
            quantization_config=quantization_config,
            device_map="auto",
            trust_remote_code=True,
            low_cpu_mem_usage=True,
            torch_dtype=torch.float32,
            max_memory={"cpu": "12GB"},
        )
        
        # Set model to evaluation mode
        model.eval()
        
        # Final memory cleanup
        clear_memory()
        
        print(f"✅ AEGIS BIO LAB 10 CONDUCTOR loaded successfully!")
        print(f"⚡ 8-bit quantization: 16GB → ~8GB memory usage")
        print(f"🚀 Ready for fast CPU inference")
        
        loading_status = "AEGIS BIO LAB 10 CONDUCTOR Multi-Domain Expert loaded successfully!"
        model_loaded = True
        return True
        
    except Exception as e:
        loading_status = f"Error loading AEGIS BIO LAB 10 CONDUCTOR model: {str(e)}"
        print(f"❌ Error loading model: {e}")
        print(f"💡 Tip: Ensure bitsandbytes>=0.41.0 is installed")
        model_loaded = False
        return False

def format_response(text):
    """Clean and format the model response"""
    # Remove thinking tags if present
    text = re.sub(r'<thinking>.*?</thinking>', '', text, flags=re.DOTALL)
    
    # Clean up extra whitespace
    text = re.sub(r'\n\s*\n', '\n\n', text)
    text = text.strip()
    
    return text

def analyze_with_aegis_conductor(prompt: str, analysis_type: str = "general") -> str:
    """Analyze using AEGIS BIO LAB 10 CONDUCTOR - Multi-Domain Expert System"""
    global model, tokenizer, model_loaded
    
    if not model_loaded or model is None or tokenizer is None:
        return "AEGIS BIO LAB 10 CONDUCTOR is still loading... Please wait a moment and try again."
    
    # Enhanced prompts for AEGIS BIO LAB 10 CONDUCTOR multi-domain analysis
    system_prompts = {
        "general": f"You are the AEGIS BIO LAB {AEGIS_VERSION} CONDUCTOR - an advanced multi-domain analysis system. You can provide expert analysis on ANY topic including economics, technology, science, politics, health, environment, security, and more. Provide comprehensive, well-reasoned responses with global perspective across all 8 regions: {', '.join(GLOBAL_REGIONS)}.",
        "economic": f"You are the AEGIS BIO LAB {AEGIS_VERSION} CONDUCTOR Economics Expert. Provide comprehensive economic analysis covering market dynamics, financial implications, GDP impacts, inflation effects, trade relationships, and policy recommendations across all 8 global regions: {', '.join(GLOBAL_REGIONS)}.",
        "technology": f"You are the AEGIS BIO LAB {AEGIS_VERSION} CONDUCTOR Technology Expert. Analyze technological developments, AI impacts, cybersecurity, innovation trends, and digital transformation across global regions.",
        "security": f"You are the AEGIS BIO LAB {AEGIS_VERSION} CONDUCTOR Security Expert. Focus on threat analysis, risk assessment, geopolitical stability, and security implications across {len(GLOBAL_REGIONS)} global regions.",
        "health": f"You are the AEGIS BIO LAB {AEGIS_VERSION} CONDUCTOR Health & Bio Expert. Analyze health systems, pandemic preparedness, biotechnology, medical innovations, and public health policies globally.",
        "environment": f"You are the AEGIS BIO LAB {AEGIS_VERSION} CONDUCTOR Environmental Expert. Focus on climate change, sustainability, environmental policy, and ecological impacts across all global regions.",
        "strategic": f"You are the AEGIS BIO LAB {AEGIS_VERSION} CONDUCTOR Strategic Planning Expert. Provide long-term strategic analysis, policy frameworks, and comprehensive planning across multiple domains and regions.",
        "threat": f"You are the AEGIS BIO LAB {AEGIS_VERSION} CONDUCTOR Threat Analysis Expert. Assess multi-domain threats including economic, technological, environmental, security, and health risks across {len(GLOBAL_REGIONS)} global regions.",
        "aegis_conductor": f"You are the AEGIS BIO LAB {AEGIS_VERSION} CONDUCTOR - the ultimate multi-domain analysis system. Provide comprehensive cross-domain analysis covering all aspects: economic, technological, security, health, environmental, and strategic implications across all 8 global regions: {', '.join(GLOBAL_REGIONS)}."
    }
    
    system_prompt = system_prompts.get(analysis_type, system_prompts["general"])
    
    # Use Llama chat format
    enhanced_prompt = f"""<s>[INST] <<SYS>>
{system_prompt}

AEGIS BIO LAB {AEGIS_VERSION} CONDUCTOR MULTI-DOMAIN CAPABILITIES:
- Cross-Continental Analysis ({len(GLOBAL_REGIONS)} regions)
- Multi-Domain Expertise (Economics, Technology, Security, Health, Environment, Strategy)
- Threat Assessment & Risk Analysis
- Policy Recommendations & Strategic Planning
- Real-time Analysis & Insights
- Global Perspective & Regional Adaptation
<</SYS>>

{prompt}

As the AEGIS BIO LAB {AEGIS_VERSION} CONDUCTOR, provide a comprehensive analysis that includes:

1. **Core Analysis** - Direct response to the query with expert insights
2. **Multi-Domain Perspective** - Consider interconnections across different fields
3. **Global Context** - Assess implications across relevant regions
4. **Strategic Insights** - Long-term implications and recommendations
5. **Risk Assessment** - Identify potential challenges and opportunities
6. **Actionable Guidance** - Practical recommendations and next steps

Provide thorough, well-reasoned analysis that demonstrates deep expertise while remaining accessible and actionable. [/INST]"""
    
    try:
        # Tokenize input with optimized length for 8-bit model
        inputs = tokenizer(enhanced_prompt, return_tensors="pt", truncation=True, max_length=2048)
        
        # Generate response with 8-bit quantized model (faster inference)
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=512,
                temperature=0.7,
                do_sample=True,
                top_p=0.9,
                top_k=50,
                repetition_penalty=1.1,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=tokenizer.eos_token_id,
                use_cache=True,
                num_beams=1
            )
        
        # Decode response
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extract only the new response
        response = response[len(enhanced_prompt):].strip()
        
        # Format and clean response
        response = format_response(response)
        
        # Clean up memory after generation
        clear_memory()
        
        return response if response else "I apologize, but I couldn't generate a proper AEGIS BIO LAB 10 CONDUCTOR analysis. Please try rephrasing your question."
        
    except Exception as e:
        return f"AEGIS BIO LAB 10 CONDUCTOR analysis error: {str(e)}. Please try a shorter question."

def conduct_aegis_threat_analysis(tech_scores: Dict[str, float], year: str = None) -> Dict:
    """Conduct comprehensive AEGIS BIO LAB 10 CONDUCTOR threat analysis"""
    if year is None:
        year = str(datetime.now().year)
    
    # Filter critical threats (scores > 6.0)
    critical_threats = {k: v for k, v in tech_scores.items() if v > 6.0}
    
    # Enhanced threat analysis prompt
    analysis_prompt = f"""AEGIS BIO LAB {AEGIS_VERSION} CONDUCTOR - COMPREHENSIVE THREAT ANALYSIS - Year {year}

TECHNOLOGY THREAT ASSESSMENT:
Critical Threats: {len(critical_threats)} detected from {len(tech_scores)} total threat categories
High-Impact Threat Categories: {list(critical_threats.keys())}
Technology Threat Scores: {dict(list(tech_scores.items()))}

REQUIRED CALCULATIONS:
1. Market Shock Index (0-1 scale): Calculate based on threat interaction effects
2. Impact Classification: Determine impact level (Limited/Moderate/Major/Crisis)
3. Threat Level: Assess overall threat (Low/Medium/High/Extreme Risk)

REGIONAL VULNERABILITIES (0-10 scale for each region):
4. North America: Technology and financial sector resilience
5. Europe: Manufacturing and energy security
6. Asia: Trade diversification and supply chain adaptation
7. Africa: Agricultural and resource sector protection
8. South America: Climate adaptation and economic diversification
9. Middle East: Energy transition and modernization
10. Oceania: Resource security and climate resilience
11. Arctic Region: Sustainable development

Provide comprehensive analysis with specific numerical values for all calculated metrics."""
    
    # Get analysis from the AEGIS model
    full_analysis = analyze_with_aegis_conductor(analysis_prompt, "aegis_conductor")
    
    # Parse metrics from the response
    result = {
        "reasoning_analysis": full_analysis,
        "market_shock_index": 0.0,
        "impact_classification": "Analysis in Progress",
        "threat_level": "Assessment Pending",
        "regional_vulnerabilities": {},
        "contagion_metrics": {},
        "tech_scores": tech_scores,
        "year": year,
        "analysis_timestamp": datetime.now().isoformat()
    }
    
    # Extract metrics from model response
    lines = full_analysis.split('\n')
    for line in lines:
        line = line.strip()
        if 'Market Shock Index:' in line or 'market shock index' in line.lower():
            try:
                import re
                numbers = re.findall(r'(\d+\.?\d*)', line)
                if numbers:
                    value = float(numbers[0])
                    if value <= 1.0:
                        result["market_shock_index"] = value
            except:
                pass
        elif 'Impact Classification:' in line or 'impact classification' in line.lower():
            parts = line.split(':')
            if len(parts) > 1:
                result["impact_classification"] = parts[1].strip()
        elif 'Threat Level:' in line or 'threat level' in line.lower():
            parts = line.split(':')
            if len(parts) > 1:
                result["threat_level"] = parts[1].strip()
    
    return result

@app.route('/')
def index():
    """Main AEGIS BIO LAB 10 CONDUCTOR interface"""
    return render_template('index.html')

@app.route('/status')
def status():
    """Get AEGIS model loading status"""
    return jsonify({
        'loaded': model_loaded,
        'status': loading_status,
        'model': MODEL_REPO,
        'version': AEGIS_VERSION,
        'regions': len(GLOBAL_REGIONS),
        'quantization': '8-bit' if model_loaded else 'Loading'
    })

@app.route('/chat', methods=['POST'])
def chat():
    """Handle AEGIS multi-domain chat messages"""
    data = request.json
    message = data.get('message', '').strip()
    history = data.get('history', [])
    temperature = float(data.get('temperature', 0.7))
    max_tokens = int(data.get('max_tokens', 256))
    analysis_type = data.get('analysis_type', 'general')
    
    if not message:
        return jsonify({'error': 'No message provided'}), 400
    
    # Generate response using AEGIS Multi-Domain System
    response = analyze_with_aegis_conductor(message, analysis_type)
    
    return jsonify({
        'response': response,
        'timestamp': time.time(),
        'model': f"AEGIS BIO LAB {AEGIS_VERSION} CONDUCTOR (8-bit)",
        'analysis_type': analysis_type
    })

@app.route('/aegis_analysis', methods=['POST'])
def aegis_analysis():
    """Handle comprehensive AEGIS BIO LAB 10 CONDUCTOR threat analysis"""
    data = request.json
    
    # Get technology threat scores
    tech_scores = {
        'AI': float(data.get('ai_score', 7.0)),
        'Cyber': float(data.get('cyber_score', 6.5)),
        'Bio': float(data.get('bio_score', 8.0)),
        'Nuclear': float(data.get('nuclear_score', 4.0)),
        'Climate': float(data.get('climate_score', 7.5)),
        'Space': float(data.get('space_score', 5.0))
    }
    
    year = data.get('year', str(datetime.now().year))
    
    # Conduct comprehensive AEGIS analysis
    analysis_result = conduct_aegis_threat_analysis(tech_scores, year)
    
    return jsonify(analysis_result)

@app.route('/clear', methods=['POST'])
def clear_chat():
    """Clear chat history and free memory"""
    clear_memory()
    return jsonify({'status': 'AEGIS BIO LAB 10 CONDUCTOR memory cleared'})

# Start model loading in background thread
def start_model_loading():
    load_model()

if __name__ == '__main__':
    # Start model loading in background
    loading_thread = threading.Thread(target=start_model_loading)
    loading_thread.daemon = True
    loading_thread.start()
    
    app.run(host='0.0.0.0', port=7860, debug=False)