Deploy AEGIS Economics Stability Analysis App

app.py

@@ -19,14 +19,22 @@ class EconRequest(BaseModel):
 aegis_brain = {}
 
 def download_model():
-    """Download the model from Hugging Face if not present"""
+    """Download the model from Hugging Face with authentication"""
     model_path = "aegis_window2_econ_v1.zip"
     model_url = "https://huggingface.co/gsstec/aegis_window2_econ_v1/resolve/main/aegis_window2_econ_v1.zip"
     
+    # Get HF token from environment variable
+    hf_token = os.getenv("HF_TOKEN")
+    if not hf_token:
+        print("⚠️ No HF_TOKEN environment variable found. Using public access only.")
+        return False
+    
+    headers = {"Authorization": f"Bearer {hf_token}"}
+    
     if not os.path.exists(model_path):
         print("📥 Downloading model from Hugging Face...")
         try:
-            response = requests.get(model_url, stream=True)
+            response = requests.get(model_url, headers=headers, stream=True)
             response.raise_for_status()
             
             with open(model_path, 'wb') as f:
@@ -61,20 +69,33 @@ async def lifespan(app: FastAPI):
     # Load risk data for cross-referencing
     try:
         rankings_url = "https://huggingface.co/gsstec/aegis_window2_econ_v1/resolve/main/Rankings.csv"
-        response = requests.get(rankings_url)
-        if response.status_code == 200:
-            aegis_brain["rankings"] = pd.read_csv(rankings_url)
-            print("✅ Rankings data loaded from Hugging Face")
+        hf_token = os.getenv("HF_TOKEN")
+        
+        if hf_token:
+            headers = {"Authorization": f"Bearer {hf_token}"}
+            response = requests.get(rankings_url, headers=headers)
+            if response.status_code == 200:
+                from io import StringIO
+                aegis_brain["rankings"] = pd.read_csv(StringIO(response.text))
+                print("✅ Rankings data loaded from Hugging Face")
+            else:
+                raise Exception(f"Rankings CSV request failed: {response.status_code}")
         else:
-            raise Exception("Rankings CSV not found")
+            raise Exception("No HF token available")
     except Exception as e:
         print(f"⚠️ Could not load rankings from HF: {e}")
-        # Create mock rankings data
+        # Create comprehensive mock rankings data
         aegis_brain["rankings"] = pd.DataFrame({
-            'Virus Name': ['COVID-19', 'H1N1', 'SARS', 'MERS', 'Ebola', 'Zika', 'Influenza A', 'RSV'],
-            'Original Score': [0.85, 0.65, 0.75, 0.55, 0.95, 0.45, 0.60, 0.40]
+            'Virus Name': [
+                'COVID-19', 'H1N1', 'SARS', 'MERS', 'Ebola', 'Zika', 
+                'Influenza A', 'RSV', 'Marburg', 'Lassa', 'Nipah', 'Hendra'
+            ],
+            'Original Score': [
+                0.85, 0.65, 0.75, 0.55, 0.95, 0.45, 
+                0.60, 0.40, 0.90, 0.70, 0.80, 0.75
+            ]
         })
-        print("✅ Mock Rankings Data Loaded")
+        print("✅ Comprehensive Mock Rankings Data Loaded")
     
     aegis_brain["model_loaded"] = model_loaded
     yield
@@ -110,26 +131,40 @@ async def get_stability_score(data: EconRequest):
         stability_score = float(prediction[0][0])
         model_status = "trained_model"
     else:
-        # Fallback to mathematical simulation
-        econ_factor = data.econ_buffer * 0.3
-        population_factor = min(data.population_exposure / 100000, 1.0) * 0.4
-        virus_factor = (1 - base_score) * 0.3
+        # Enhanced mathematical simulation
+        econ_factor = data.econ_buffer * 0.35
+        population_factor = min(data.population_exposure / 100000, 1.0) * 0.35
+        virus_factor = (1 - base_score) * 0.30
         
+        # More sophisticated calculation
         stability_score = max(0.0, min(1.0, 
             econ_factor + (1 - population_factor) + (1 - virus_factor)
         ))
         
-        # Add some realistic variance
-        stability_score += np.random.normal(0, 0.05)
+        # Add realistic variance based on input parameters
+        variance = 0.03 + (data.population_exposure / 1000000) * 0.02
+        stability_score += np.random.normal(0, variance)
         stability_score = max(0.0, min(1.0, stability_score))
-        model_status = "simulation"
+        model_status = "enhanced_simulation"
+    
+    # Determine alert level with more nuanced thresholds
+    if stability_score < 0.25:
+        alert_level = "CRITICAL"
+    elif stability_score < 0.45:
+        alert_level = "HIGH_RISK"
+    elif stability_score < 0.65:
+        alert_level = "MONITOR"
+    else:
+        alert_level = "STABLE"
     
     return {
         "virus": data.virus_name,
         "stability_score": round(stability_score, 4),
-        "alert_level": "CRITICAL" if stability_score < 0.3 else "MONITOR" if stability_score < 0.6 else "STABLE",
+        "alert_level": alert_level,
         "model_status": model_status,
-        "base_risk_score": round(base_score, 3)
+        "base_risk_score": round(base_score, 3),
+        "economic_buffer_impact": round(data.econ_buffer * 0.35, 3),
+        "population_risk_factor": round(min(data.population_exposure / 100000, 1.0), 3)
     }
 
 @app.get("/health")
@@ -138,17 +173,26 @@ async def health():
         "status": "operational", 
         "hardware": "T4 GPU Active" if torch.cuda.is_available() else "CPU Mode",
         "model_loaded": aegis_brain.get("model_loaded", False),
-        "available_viruses": aegis_brain.get("rankings", pd.DataFrame())['Virus Name'].tolist() if aegis_brain.get("rankings") is not None else []
+        "available_viruses": aegis_brain.get("rankings", pd.DataFrame())['Virus Name'].tolist() if aegis_brain.get("rankings") is not None else [],
+        "total_virus_database": len(aegis_brain.get("rankings", pd.DataFrame())),
+        "hf_token_available": bool(os.getenv("HF_TOKEN"))
     }
 
 @app.get("/")
 async def root():
     return {
         "message": "AEGIS Economic Stability Analysis API",
-        "version": "1.0.0",
+        "version": "1.2.0",
         "model_source": "https://huggingface.co/gsstec/aegis_window2_econ_v1",
+        "features": [
+            "Economic stability prediction",
+            "Multi-virus risk assessment", 
+            "Population exposure modeling",
+            "Real-time alert classification"
+        ],
         "endpoints": {
             "predict": "POST /predict - Get stability prediction",
-            "health": "GET /health - System status"
+            "health": "GET /health - System status",
+            "root": "GET / - API information"
         }
     }