Wave 2: Context-grounded security chatbot — /api/v2/chat upgrade

Dockerfile

@@ -0,0 +1,38 @@
+# Use Python 3.11 to avoid the gradio_client Python 3.13 bug
+FROM python:3.11-slim
+
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    git \
+    git-lfs \
+    ffmpeg \
+    libsm6 \
+    libxext6 \
+    && rm -rf /var/lib/apt/lists/* \
+    && git lfs install
+
+# Create user for HF Spaces
+RUN useradd -m -u 1000 user
+USER user
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH
+
+WORKDIR $HOME/app
+
+# Copy requirements and install dependencies
+COPY --chown=user requirements.txt .
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir -r requirements.txt && \
+    pip install --no-cache-dir ipykernel && \
+    python -m ipykernel install --user --name python3 --display-name "Python 3"
+
+# Copy application code
+COPY --chown=user . .
+
+# Expose port
+EXPOSE 7860
+
+# Run the application via uvicorn (FastAPI + Gradio mounted)
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -0,0 +1,45 @@
+---
+title: CyberForge AI
+emoji: 🔐
+colorFrom: blue
+colorTo: purple
+sdk: docker
+app_file: app.py
+pinned: false
+license: mit
+short_description: CyberForge AI - ML Training & Inference Platform
+---
+
+# 🔐 CyberForge AI - ML Training & Inference Platform
+
+Gradio UI for notebook execution, model training, and inference + FastAPI REST API for the Heroku backend.
+
+## REST API Endpoints (for backend mlService.js)
+
+| Endpoint | Method | Description |
+|----------|--------|-------------|
+| `/health` | GET | Health check |
+| `/analyze` | POST | AI chat (Gemini) |
+| `/analyze-url` | POST | URL threat analysis |
+| `/scan-threats` | POST | Threat scanning |
+| `/api/insights/generate` | POST | AI insights |
+| `/api/models/predict` | POST | ML model prediction |
+| `/models` | GET | List available models |
+| `/api/analysis/network` | POST | Network traffic analysis |
+| `/api/ai/execute-task` | POST | AI task execution |
+
+## Environment Secrets
+
+Set these in your Space settings → Secrets:
+
+- `GEMINI_API_KEY` - Google Gemini API key
+- `HF_TOKEN` - HuggingFace token (for model downloads)
+- `HF_MODEL_REPO` - Model repository (default: `Che237/cyberforge-models`)
+
+## Backend Integration
+
+Your Heroku backend connects to this Space:
+
+```bash
+heroku config:set AI_SERVICE_URL=https://che237-cyberforge.hf.space -a cyberforge
+```

app.py

@@ -0,0 +1,1945 @@
+"""
+CyberForge AI - ML Training & Inference Platform
+Hugging Face Spaces deployment with:
+  1) Gradio UI for notebook execution, training, and inference
+  2) FastAPI REST endpoints for the Heroku backend (mlService.js)
+"""
+
+import gradio as gr
+import pandas as pd
+import numpy as np
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+from datetime import datetime
+import logging
+from typing import Dict, List, Any, Optional, Tuple
+from urllib.parse import urlparse
+
+# ML Libraries
+from sklearn.model_selection import train_test_split, cross_val_score
+from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier, IsolationForest
+from sklearn.linear_model import LogisticRegression
+from sklearn.preprocessing import StandardScaler, LabelEncoder
+from sklearn.metrics import classification_report, confusion_matrix, accuracy_score, f1_score
+import joblib
+
+# Hugging Face Hub
+from huggingface_hub import HfApi, hf_hub_download, upload_file
+
+# FastAPI for REST endpoints
+from fastapi import FastAPI, Request
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+
+# Gemini AI  (new SDK: google-genai)
+try:
+    from google import genai
+    GEMINI_AVAILABLE = True
+except ImportError:
+    GEMINI_AVAILABLE = False
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# ============================================================================
+# CONFIGURATION
+# ============================================================================
+
+APP_DIR = Path(__file__).parent.absolute()
+MODELS_DIR = APP_DIR / "trained_models"
+MODELS_DIR.mkdir(exist_ok=True)
+DATASETS_DIR = APP_DIR / "datasets"
+DATASETS_DIR.mkdir(exist_ok=True)
+NOTEBOOKS_DIR = APP_DIR / "notebooks"
+KNOWLEDGE_BASE_DIR = APP_DIR / "knowledge_base"
+KNOWLEDGE_BASE_DIR.mkdir(exist_ok=True)
+TRAINING_DATA_DIR = APP_DIR / "training_data"
+TRAINING_DATA_DIR.mkdir(exist_ok=True)
+
+# Environment
+GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
+GEMINI_MODEL = os.environ.get("GEMINI_MODEL", "gemini-2.5-flash")
+HF_TOKEN = os.environ.get("HF_TOKEN", "")
+HF_MODEL_REPO = os.environ.get("HF_MODEL_REPO", "Che237/cyberforge-models")
+
+logger.info(f"APP_DIR: {APP_DIR}")
+logger.info(f"NOTEBOOKS_DIR: {NOTEBOOKS_DIR}")
+logger.info(f"NOTEBOOKS_DIR exists: {NOTEBOOKS_DIR.exists()}")
+
+# Model types available for training
+MODEL_TYPES = {
+    "Random Forest": RandomForestClassifier,
+    "Gradient Boosting": GradientBoostingClassifier,
+    "Logistic Regression": LogisticRegression,
+    "Isolation Forest (Anomaly)": IsolationForest,
+}
+
+SECURITY_TASKS = [
+    "Malware Detection", "Phishing Detection", "Network Intrusion Detection",
+    "Anomaly Detection", "Botnet Detection", "Web Attack Detection",
+    "Spam Detection", "Vulnerability Assessment", "DNS Tunneling Detection",
+    "Cryptomining Detection",
+]
+
+
+# ============================================================================
+# GEMINI SERVICE  (for REST API)
+# ============================================================================
+
+class GeminiService:
+    """Google Gemini AI for cybersecurity chat and analysis"""
+
+    SYSTEM_PROMPT = """You are CyberForge AI, an advanced cybersecurity expert. You specialize in:
+- Real-time threat detection and analysis
+- Malware and phishing identification
+- Network security assessment
+- Browser security monitoring
+- Risk assessment and mitigation
+
+When analyzing security queries, provide:
+1. Risk Level (Critical/High/Medium/Low)
+2. Threat Types identified
+3. Confidence Score (0.0-1.0)
+4. Detailed technical analysis
+5. Specific actionable recommendations
+
+Always be precise, professional, and actionable."""
+
+    def __init__(self):
+        self.client = None
+        self.ready = False
+        self.custom_knowledge = {}
+        self.training_examples = []
+
+    def initialize(self):
+        if not GEMINI_AVAILABLE:
+            logger.warning("google-genai not installed, Gemini unavailable")
+            return
+        if not GEMINI_API_KEY:
+            logger.warning("GEMINI_API_KEY not set, Gemini unavailable")
+            return
+        try:
+            self.client = genai.Client(api_key=GEMINI_API_KEY)
+            resp = self.client.models.generate_content(
+                model=GEMINI_MODEL,
+                contents="Test. Respond with OK."
+            )
+            if resp.text:
+                self.ready = True
+                logger.info(f"✅ Gemini initialized (model: {GEMINI_MODEL})")
+            self._load_knowledge()
+        except Exception as e:
+            logger.error(f"❌ Gemini init failed: {e}")
+            self.ready = False
+
+    def _load_knowledge(self):
+        try:
+            for f in KNOWLEDGE_BASE_DIR.glob("*.json"):
+                with open(f) as fh:
+                    self.custom_knowledge[f.stem] = json.load(fh)
+            logger.info(f"Loaded {len(self.custom_knowledge)} knowledge files")
+        except Exception as e:
+            logger.warning(f"Knowledge load error: {e}")
+        try:
+            for f in TRAINING_DATA_DIR.glob("*.json"):
+                with open(f) as fh:
+                    data = json.load(fh)
+                    if isinstance(data, list):
+                        self.training_examples.extend(data)
+                    else:
+                        self.training_examples.append(data)
+            logger.info(f"Loaded {len(self.training_examples)} training examples")
+        except Exception as e:
+            logger.warning(f"Training data load error: {e}")
+
+    def analyze(self, query: str, context: Dict = None, history: List = None) -> Dict:
+        if not self.ready:
+            return self._fallback(query)
+        try:
+            knowledge_str = json.dumps(self.custom_knowledge, indent=1)[:2000] if self.custom_knowledge else "None"
+            examples_str = "\n".join(
+                f"Q: {ex.get('input','')}\nA: {ex.get('output','')}"
+                for ex in self.training_examples[-3:]
+            ) if self.training_examples else "None"
+            context_str = f"\nCONTEXT:\n{json.dumps(context, indent=1)[:3000]}\n" if context else ""
+
+            prompt = f"""{self.SYSTEM_PROMPT}
+
+KNOWLEDGE BASE (summary):
+{knowledge_str}
+
+TRAINING EXAMPLES:
+{examples_str}
+{context_str}
+USER QUERY:
+{query}
+
+Provide a comprehensive cybersecurity analysis:"""
+
+            response = self.client.models.generate_content(
+                model=GEMINI_MODEL,
+                contents=prompt,
+                config={"temperature": 0.3, "max_output_tokens": 2048},
+            )
+            text = response.text if response.text else "No response generated"
+            text_lower = text.lower()
+            if "critical" in text_lower:
+                risk_level, risk_score = "Critical", 9.0
+            elif "high" in text_lower:
+                risk_level, risk_score = "High", 7.0
+            elif "medium" in text_lower:
+                risk_level, risk_score = "Medium", 5.0
+            else:
+                risk_level, risk_score = "Low", 3.0
+
+            return {
+                "response": text,
+                "confidence": 0.85,
+                "risk_level": risk_level,
+                "risk_score": risk_score,
+                "insights": [f"Analysis performed by Gemini ({GEMINI_MODEL})"],
+                "recommendations": [],
+                "model_used": GEMINI_MODEL,
+                "timestamp": datetime.utcnow().isoformat(),
+            }
+        except Exception as e:
+            logger.error(f"Gemini analysis error: {e}")
+            return self._fallback(query)
+
+    def _fallback(self, query: str) -> Dict:
+        """ML-powered analysis when Gemini is unavailable"""
+        import re
+        global ml_loader
+
+        url_pattern = re.compile(r'https?://[^\s"\'<>]+')
+        urls = url_pattern.findall(query)
+
+        risk_level = "Unknown"
+        risk_score = 0.0
+        insights = []
+        response_lines = ["## CyberForge ML Security Analysis\n"]
+
+        try:
+            loader = ml_loader
+            n_loaded = len(loader.models)
+        except Exception:
+            loader = None
+            n_loaded = 0
+
+        if urls and loader and n_loaded > 0:
+            threat_models_fired = []
+            all_scores = []
+
+            for url in urls[:3]:
+                features = extract_url_features(url)
+                url_threats = []
+                url_scores = []
+
+                for model_name in ["phishing_detection", "malware_detection", "web_attack_detection"]:
+                    pred = loader.predict(model_name, features)
+                    if pred.get("prediction", 0) == 1:
+                        label = model_name.replace("_detection", "").replace("_", " ").title()
+                        url_threats.append(label)
+                        url_scores.append(pred.get("confidence", 0.5))
+                        threat_models_fired.append(model_name)
+
+                avg = sum(url_scores) / len(url_scores) if url_scores else 0.15
+                all_scores.append(avg)
+                lvl = "HIGH" if avg > 0.6 else "MEDIUM" if avg > 0.35 else "LOW"
+                threats_str = ", ".join(url_threats) if url_threats else "None detected"
+                display_url = url if len(url) <= 70 else url[:67] + "..."
+                response_lines.append(f"**URL:** `{display_url}`")
+                response_lines.append(f"- Risk: **{lvl}** | Threats: {threats_str} | Score: {avg:.0%}\n")
+
+            overall = sum(all_scores) / len(all_scores) if all_scores else 0.2
+            if overall > 0.65:
+                risk_level, risk_score = "High", 7.5
+            elif overall > 0.4:
+                risk_level, risk_score = "Medium", 5.0
+            else:
+                risk_level, risk_score = "Low", 2.5
+
+            insights = [f"{n_loaded} ML models active"] + [
+                f"Threat model triggered: {m.replace('_detection', '').replace('_', ' ').title()}"
+                for m in set(threat_models_fired)
+            ]
+
+            if risk_level == "High":
+                response_lines.append("### Recommendation\n⚠️ **Block immediately.** Phishing or malware indicators detected — do not visit this URL.")
+            elif risk_level == "Medium":
+                response_lines.append("### Recommendation\n⚡ **Exercise caution.** Validate with additional threat intelligence before accessing.")
+            else:
+                response_lines.append("### Recommendation\n✅ **URL appears structurally safe** based on ML analysis.")
+
+        else:
+            query_lower = query.lower()
+            malware_kws = ["malware", "virus", "ransomware", "trojan", "spyware", "backdoor", "worm"]
+            phishing_kws = ["phishing", "credential", "fake login", "spoof", "scam", "social engineering"]
+            safe_kws = ["safe", "legitimate", "trusted", "secure", "verify"]
+
+            if any(k in query_lower for k in malware_kws):
+                risk_level, risk_score = "High", 7.0
+                response_lines.append(
+                    "**Malware indicators detected in your query.**\n\n"
+                    "Recommended actions:\n"
+                    "- Isolate the affected system immediately\n"
+                    "- Run a full endpoint detection scan\n"
+                    "- Review recently installed software and browser extensions\n"
+                    "- Check startup processes and scheduled tasks for persistence\n"
+                    "- Rotate credentials if any exposure is suspected"
+                )
+            elif any(k in query_lower for k in phishing_kws):
+                risk_level, risk_score = "High", 7.0
+                response_lines.append(
+                    "**Phishing threat indicators detected.**\n\n"
+                    "Recommended actions:\n"
+                    "- Do not submit credentials to the suspected site\n"
+                    "- Verify sender identity through a secondary channel\n"
+                    "- Report to your IT security team immediately\n"
+                    "- Enable MFA on all accounts that may be affected\n"
+                    "- Review email headers for spoofing indicators"
+                )
+            elif any(k in query_lower for k in safe_kws):
+                risk_level, risk_score = "Low", 2.0
+                response_lines.append("No immediate threat indicators detected. Continue standard monitoring procedures.")
+            else:
+                response_lines.append(
+                    f"CyberForge ML is operational with **{n_loaded}/4 models** loaded.\n\n"
+                    "For best results, include a URL or specific threat indicators in your query.\n\n"
+                    "**Available analysis capabilities:**\n"
+                    "- URL threat analysis (phishing, malware, web attacks)\n"
+                    "- Network anomaly detection\n"
+                    "- Real-time threat event monitoring\n\n"
+                    "*AI chat (Gemini) is currently unavailable. Provide a URL for full ML-based analysis.*"
+                )
+            insights = [f"{n_loaded} ML models active"]
+
+        response_lines.append("\n---\n*Powered by CyberForge ML models — Gemini AI offline.*")
+
+        return {
+            "response": "\n".join(response_lines),
+            "confidence": 0.65 if urls else 0.4,
+            "risk_level": risk_level,
+            "risk_score": risk_score,
+            "insights": insights,
+            "recommendations": [],
+            "model_used": "cyberforge-ml-fallback",
+            "timestamp": datetime.utcnow().isoformat(),
+        }
+
+
+# Singleton
+gemini_service = GeminiService()
+
+
+# ============================================================================
+# ML MODEL LOADER  (loads .pkl from HF Hub for REST predictions)
+# ============================================================================
+
+class MLModelLoader:
+    MODEL_NAMES = [
+        "phishing_detection", "malware_detection",
+        "anomaly_detection", "web_attack_detection",
+    ]
+
+    def __init__(self):
+        self.models: Dict[str, Any] = {}
+        self.scalers: Dict[str, Any] = {}
+        self.ready = False
+
+    def initialize(self):
+        loaded = 0
+        # All directories where models might exist (notebooks save to ../models)
+        search_dirs = [
+            MODELS_DIR,                          # trained_models/
+            APP_DIR / "models",                  # models/ (where notebooks output)
+            APP_DIR.parent / "models",           # one level up fallback
+        ]
+
+        for name in self.MODEL_NAMES:
+            if name in self.models:
+                continue
+            try:
+                # 1. Try HuggingFace Hub first
+                model_file = f"{name}/best_model.pkl"
+                scaler_file = f"{name}/scaler.pkl"
+                try:
+                    model_path = hf_hub_download(
+                        repo_id=HF_MODEL_REPO, filename=model_file,
+                        token=HF_TOKEN or None, cache_dir=str(MODELS_DIR),
+                    )
+                    scaler_path = hf_hub_download(
+                        repo_id=HF_MODEL_REPO, filename=scaler_file,
+                        token=HF_TOKEN or None, cache_dir=str(MODELS_DIR),
+                    )
+                    self.models[name] = joblib.load(model_path)
+                    self.scalers[name] = joblib.load(scaler_path)
+                    loaded += 1
+                    logger.info(f"✅ Loaded model from Hub: {name}")
+                    continue
+                except Exception:
+                    pass
+
+                # 2. Try all local search directories
+                for sdir in search_dirs:
+                    if name in self.models:
+                        break
+                    for model_fname in [f"{name}/best_model.pkl", f"{name}/model.pkl", f"{name}_model.pkl"]:
+                        candidate = sdir / model_fname
+                        if candidate.exists():
+                            self.models[name] = joblib.load(candidate)
+                            # Try to find matching scaler
+                            for scaler_fname in [f"{name}/scaler.pkl", f"{name}_scaler.pkl"]:
+                                sc = sdir / scaler_fname
+                                if sc.exists():
+                                    self.scalers[name] = joblib.load(sc)
+                                    break
+                            loaded += 1
+                            logger.info(f"✅ Loaded model from {sdir.name}/{model_fname}: {name}")
+                            break
+
+            except Exception as e:
+                logger.warning(f"Error loading model {name}: {e}")
+
+        # Sweep all search dirs for any .pkl files not yet loaded
+        for sdir in search_dirs:
+            if not sdir.exists():
+                continue
+            for pkl in sdir.glob("*.pkl"):
+                stem = pkl.stem.replace("_model", "").replace("_best", "")
+                if stem not in self.models:
+                    try:
+                        self.models[stem] = joblib.load(pkl)
+                        loaded += 1
+                        logger.info(f"✅ Loaded model sweep: {stem} from {sdir.name}")
+                    except Exception:
+                        pass
+
+        self.ready = loaded > 0
+        logger.info(f"ML Models: {loaded} loaded — {list(self.models.keys())}")
+
+    def predict(self, model_name: str, features: Dict) -> Dict:
+        if model_name not in self.models:
+            return self._heuristic_predict(model_name, features)
+        try:
+            model = self.models[model_name]
+            scaler = self.scalers.get(model_name)
+            X = np.array([list(features.values())])
+            if scaler:
+                X = scaler.transform(X)
+            prediction = int(model.predict(X)[0])
+            confidence = 0.5
+            probabilities = {}
+            if hasattr(model, "predict_proba"):
+                proba = model.predict_proba(X)[0]
+                confidence = float(max(proba))
+                probabilities = {str(i): round(float(p), 4) for i, p in enumerate(proba)}
+            # Derive a consistent threat_score: probability of class-1 (threat)
+            # so that the /analyze-url aggregation has a uniform field across all models
+            if "1" in probabilities:
+                threat_score = probabilities["1"]
+            elif prediction == 1:
+                threat_score = confidence
+            else:
+                threat_score = 1.0 - confidence
+            return {
+                "model": model_name,
+                "prediction": prediction,
+                "prediction_label": "threat" if prediction == 1 else "benign",
+                "confidence": round(confidence, 4),
+                "threat_score": round(threat_score, 4),
+                "probabilities": probabilities,
+                "inference_source": "ml_model",
+                "timestamp": datetime.utcnow().isoformat(),
+            }
+        except Exception as e:
+            logger.error(f"Prediction error {model_name}: {e}")
+            return self._heuristic_predict(model_name, features)
+
+    def feature_importance(self, model_name: str) -> Dict:
+        """Return feature importances for tree-based models (explainability)."""
+        feature_names = [
+            "url_length", "hostname_length", "path_length", "is_https",
+            "has_ip_address", "has_suspicious_tld", "subdomain_count",
+            "has_port", "query_params_count", "has_at_symbol",
+            "has_double_slash", "special_char_count",
+        ]
+        if model_name not in self.models:
+            return {"error": f"Model '{model_name}' not loaded", "available": list(self.models.keys())}
+        model = self.models[model_name]
+        try:
+            if hasattr(model, "feature_importances_"):
+                importances = model.feature_importances_
+                pairs = sorted(
+                    zip(feature_names, [round(float(v), 4) for v in importances]),
+                    key=lambda x: x[1], reverse=True,
+                )
+                return {
+                    "model": model_name,
+                    "method": "gini_importance",
+                    "feature_importances": dict(pairs),
+                    "top_features": [p[0] for p in pairs[:5]],
+                }
+            elif hasattr(model, "coef_"):
+                coefs = model.coef_[0] if model.coef_.ndim > 1 else model.coef_
+                pairs = sorted(
+                    zip(feature_names, [round(float(v), 4) for v in coefs]),
+                    key=lambda x: abs(x[1]), reverse=True,
+                )
+                return {
+                    "model": model_name,
+                    "method": "logistic_coefficients",
+                    "feature_importances": dict(pairs),
+                    "top_features": [p[0] for p in pairs[:5]],
+                }
+            else:
+                return {
+                    "model": model_name,
+                    "method": "not_available",
+                    "note": f"Model type {type(model).__name__} does not expose feature importances",
+                }
+        except Exception as e:
+            return {"model": model_name, "error": str(e)}
+
+    def _heuristic_predict(self, model_name: str, features: Dict) -> Dict:
+        score = 0.0
+        reasons = []
+        is_https = features.get("is_https", features.get("has_https", 1))
+        has_ip = features.get("has_ip_address", 0)
+        suspicious_tld = features.get("has_suspicious_tld", 0)
+        url_length = features.get("url_length", 0)
+        special_chars = features.get("special_char_count", 0)
+        if not is_https:
+            score += 0.3; reasons.append("No HTTPS")
+        if has_ip:
+            score += 0.25; reasons.append("IP address in URL")
+        if suspicious_tld:
+            score += 0.2; reasons.append("Suspicious TLD")
+        if url_length > 100:
+            score += 0.15; reasons.append("Very long URL")
+        if special_chars > 10:
+            score += 0.1; reasons.append("Many special characters")
+        is_threat = score >= 0.5
+        return {
+            "model": model_name,
+            "prediction": 1 if is_threat else 0,
+            "prediction_label": "threat" if is_threat else "benign",
+            "confidence": min(score + 0.3, 0.95) if is_threat else max(0.6, 1.0 - score),
+            "threat_score": score,
+            "reasons": reasons,
+            "inference_source": "heuristic",
+            "timestamp": datetime.utcnow().isoformat(),
+        }
+
+
+ml_loader = MLModelLoader()
+
+
+# ============================================================================
+# URL FEATURE EXTRACTION
+# ============================================================================
+
+def extract_url_features(url: str) -> Dict:
+    parsed = urlparse(url)
+    hostname = parsed.hostname or ""
+    return {
+        "url_length": len(url),
+        "hostname_length": len(hostname),
+        "path_length": len(parsed.path or ""),
+        "is_https": 1 if parsed.scheme == "https" else 0,
+        "has_ip_address": 1 if all(p.isdigit() for p in hostname.split(".")) and len(hostname.split(".")) == 4 else 0,
+        "has_suspicious_tld": 1 if any(hostname.endswith(t) for t in [".xyz", ".tk", ".ml", ".ga", ".cf", ".top", ".buzz"]) else 0,
+        "subdomain_count": max(0, len(hostname.split(".")) - 2),
+        "has_port": 1 if parsed.port else 0,
+        "query_params_count": len(parsed.query.split("&")) if parsed.query else 0,
+        "has_at_symbol": 1 if "@" in url else 0,
+        "has_double_slash": 1 if "//" in (parsed.path or "") else 0,
+        "special_char_count": sum(1 for c in url if c in "!@#$%^&*()+={}[]|\\:;<>?,"),
+    }
+
+
+# ============================================================================
+# TRANSFORMER MODEL LOADER  (Phase 3 — real HF transformer models)
+# ============================================================================
+#
+# Loads pretrained BERT-based classifiers from the Hugging Face Hub.
+# Models are loaded lazily on first request to keep cold-start fast.
+# A 7B Security LLM is NOT loaded locally (too big for free tier) —
+# it's accessed via the HF Inference API on demand.
+
+class TransformerModelLoader:
+    """Loads pretrained Transformer classifiers from the HF Hub on demand."""
+
+    # Model registry — name → HF repo + task description
+    # NOTE: DGA detector uses an inline entropy heuristic, not a transformer
+    #       (YangYang-Research/dga-detection has a non-standard model config that
+    #       isn't loadable via transformers.pipeline()).
+    REGISTRY = {
+        "url_phishing_bert": {
+            "repo":  "elftsdmr/malware-url-detect",
+            "task":  "text-classification",
+            "labels": ["benign", "malicious"],
+            "desc":  "BERT-based URL phishing/malware classifier",
+        },
+    }
+    # Model used via HF Inference API. ZySec-AI/SecurityLLM isn't served on the
+    # free Inference API tier (404). Mistral-7B-Instruct-v0.3 is widely available
+    # and works well for cyber Q&A. Override via SECURITY_LLM_MODEL env var.
+    SECURITY_LLM_REPO = os.environ.get("SECURITY_LLM_MODEL", "mistralai/Mistral-7B-Instruct-v0.3")
+
+    def __init__(self):
+        self.pipelines = {}     # name → transformers.Pipeline (lazy-loaded)
+        self.load_errors = {}   # name → last error message
+        self.transformers_available = False
+        try:
+            import transformers  # noqa: F401
+            import torch         # noqa: F401
+            self.transformers_available = True
+            logger.info("✅ transformers + torch available")
+        except ImportError as e:
+            logger.warning(f"⚠️ transformers/torch not installed: {e}")
+
+    def _ensure(self, name: str):
+        """Load a pipeline on first use. Returns the pipeline or None on failure."""
+        if name in self.pipelines:
+            return self.pipelines[name]
+        if not self.transformers_available:
+            self.load_errors[name] = "transformers/torch not installed"
+            return None
+        if name not in self.REGISTRY:
+            self.load_errors[name] = f"Unknown model: {name}"
+            return None
+        spec = self.REGISTRY[name]
+        try:
+            from transformers import pipeline
+            logger.info(f"⏳ Loading {name} from {spec['repo']}...")
+            pipe = pipeline(spec["task"], model=spec["repo"], device=-1, top_k=None)
+            self.pipelines[name] = pipe
+            logger.info(f"✅ Loaded {name}")
+            return pipe
+        except Exception as e:
+            err = f"{type(e).__name__}: {str(e)[:200]}"
+            self.load_errors[name] = err
+            logger.error(f"❌ Failed to load {name}: {err}")
+            return None
+
+    def predict_url_phishing(self, url: str) -> Dict:
+        """Classify a URL as benign or malicious using elftsdmr/malware-url-detect."""
+        pipe = self._ensure("url_phishing_bert")
+        if pipe is None:
+            return self._unavailable("url_phishing_bert")
+        try:
+            # Strip the protocol — model was trained on bare URLs
+            text = url.replace("https://", "").replace("http://", "")[:512]
+            result = pipe(text)
+            scores = result[0] if isinstance(result[0], list) else result
+            return self._format_classification(scores, "url_phishing_bert")
+        except Exception as e:
+            return self._error("url_phishing_bert", e)
+
+    def predict_dga(self, domain: str) -> Dict:
+        """Detect DGA-generated domains using a Shannon-entropy + character-pattern heuristic.
+        Real DGA domains have high entropy, low pronounceability, no real word substrings."""
+        import math
+        d = domain.lower().strip().split('.')[0][:60]  # SLD only, ignore TLD
+        if not d:
+            return {"model": "dga_detector", "error": "empty domain"}
+        # Shannon entropy of character distribution
+        from collections import Counter
+        freq = Counter(d)
+        n = len(d)
+        entropy = -sum((c / n) * math.log2(c / n) for c in freq.values())
+        # Vowel ratio — DGAs typically have very few vowels
+        vowels = sum(1 for c in d if c in 'aeiou')
+        vowel_ratio = vowels / n if n else 0
+        # Digit ratio — DGAs often mix digits in
+        digits = sum(1 for c in d if c.isdigit())
+        digit_ratio = digits / n if n else 0
+        # Length signal — DGAs are usually 10-25 chars
+        length_signal = 1.0 if 12 <= n <= 30 else 0.5 if 8 <= n <= 40 else 0.2
+        # Consonant runs — DGAs often have 4+ consonants in a row
+        max_consonant_run = 0
+        run = 0
+        for c in d:
+            if c.isalpha() and c not in 'aeiou':
+                run += 1
+                max_consonant_run = max(max_consonant_run, run)
+            else:
+                run = 0
+        # Score combination — empirically tuned thresholds
+        score = 0.0
+        if entropy > 3.5: score += 0.35
+        if vowel_ratio < 0.25: score += 0.20
+        if digit_ratio > 0.15: score += 0.15
+        if max_consonant_run >= 4: score += 0.20
+        score *= length_signal
+        score = min(1.0, score)
+        is_dga = score >= 0.45
+        return {
+            "model":           "dga_detector",
+            "prediction":      "dga" if is_dga else "legit",
+            "is_threat":       is_dga,
+            "confidence":      round(score * 100, 2) if is_dga else round((1 - score) * 100, 2),
+            "threat_score":    round(score, 4),
+            "features": {
+                "entropy":            round(entropy, 3),
+                "vowel_ratio":        round(vowel_ratio, 3),
+                "digit_ratio":        round(digit_ratio, 3),
+                "max_consonant_run":  max_consonant_run,
+                "length":             n,
+            },
+            "inference_source": "entropy-heuristic",
+        }
+
+    def security_chat(self, query: str, max_tokens: int = 512) -> Dict:
+        """Cybersecurity Q&A via ZySec-AI/SecurityLLM hosted on HF Inference API.
+        Falls back to Gemini when the LLM is rate-limited or HF token is missing."""
+        if not HF_TOKEN:
+            return {
+                "model": "security-llm",
+                "response": None,
+                "error": "HF_TOKEN env var required to call ZySec-AI/SecurityLLM via Inference API",
+                "fallback_available": gemini_service.ready,
+            }
+        try:
+            import requests
+            url = f"https://api-inference.huggingface.co/models/{self.SECURITY_LLM_REPO}"
+            headers = {"Authorization": f"Bearer {HF_TOKEN}", "Content-Type": "application/json"}
+            payload = {
+                "inputs": query[:2000],
+                "parameters": {"max_new_tokens": max_tokens, "temperature": 0.3, "return_full_text": False},
+                "options": {"wait_for_model": True},
+            }
+            r = requests.post(url, headers=headers, json=payload, timeout=45)
+            if r.status_code == 200:
+                data = r.json()
+                text = data[0].get("generated_text") if isinstance(data, list) and data else (data.get("generated_text") if isinstance(data, dict) else str(data))
+                return {
+                    "model":    "security-llm",
+                    "source":   "huggingface_inference_api",
+                    "response": text,
+                    "model_id": self.SECURITY_LLM_REPO,
+                }
+            return {
+                "model":  "security-llm",
+                "error":  f"HF Inference API HTTP {r.status_code}: {r.text[:200]}",
+                "fallback_available": gemini_service.ready,
+            }
+        except Exception as e:
+            return {
+                "model":  "security-llm",
+                "error":  f"{type(e).__name__}: {str(e)[:200]}",
+                "fallback_available": gemini_service.ready,
+            }
+
+    def status(self) -> Dict:
+        return {
+            "transformers_available": self.transformers_available,
+            "loaded": list(self.pipelines.keys()),
+            "available": list(self.REGISTRY.keys()) + ["security_llm (via HF Inference API)"],
+            "load_errors": self.load_errors,
+        }
+
+    @staticmethod
+    def _format_classification(scores, model_name) -> Dict:
+        """Normalize HF text-classification output to the cyberforge schema."""
+        if not scores:
+            return {"model": model_name, "error": "Empty scores"}
+        # scores is a list of {label, score} dicts. Find the threat label.
+        threat_labels = {"malicious", "phishing", "malware", "dga", "label_1", "1"}
+        # Top prediction
+        top = max(scores, key=lambda s: s["score"]) if isinstance(scores, list) else scores
+        is_threat = str(top["label"]).lower() in threat_labels
+        # Threat score: probability of the threat class (not just the top class)
+        threat_score = top["score"] if is_threat else 1.0 - top["score"]
+        return {
+            "model":           model_name,
+            "prediction":      top["label"],
+            "is_threat":       is_threat,
+            "confidence":      round(top["score"] * 100, 2),
+            "threat_score":    round(threat_score, 4),
+            "all_scores":      scores,
+            "inference_source": "huggingface_transformer",
+        }
+
+    @staticmethod
+    def _unavailable(model_name) -> Dict:
+        return {"model": model_name, "error": "Model unavailable — see /api/v2/status"}
+
+    @staticmethod
+    def _error(model_name, e: Exception) -> Dict:
+        return {"model": model_name, "error": f"{type(e).__name__}: {str(e)[:200]}"}
+
+
+transformer_loader = TransformerModelLoader()
+
+
+# ============================================================================
+# NOTEBOOK EXECUTION  (existing Gradio functionality)
+# ============================================================================
+
+def get_available_notebooks() -> List[str]:
+    if not NOTEBOOKS_DIR.exists():
+        return []
+    return sorted([f.name for f in NOTEBOOKS_DIR.glob("*.ipynb")])
+
+
+def read_notebook_content(notebook_name: str) -> str:
+    """Read and display notebook content as markdown"""
+    notebook_path = NOTEBOOKS_DIR / notebook_name
+    if not notebook_path.exists():
+        return f"Notebook not found: {notebook_name}"
+    try:
+        with open(notebook_path, "r") as f:
+            nb = json.load(f)
+        output = f"# {notebook_name}\n\n"
+        for i, cell in enumerate(nb.get("cells", []), 1):
+            cell_type = cell.get("cell_type", "code")
+            source = "".join(cell.get("source", []))
+            if cell_type == "markdown":
+                output += f"{source}\n\n"
+            else:
+                output += f"### Cell {i} (Python)\n```python\n{source}\n```\n\n"
+        return output
+    except Exception as e:
+        return f"Error reading notebook: {str(e)}"
+
+
+def execute_notebook(notebook_name: str, progress=gr.Progress()) -> Tuple[str, str]:
+    """Execute a notebook and return output"""
+    notebook_path = NOTEBOOKS_DIR / notebook_name
+    output_path = NOTEBOOKS_DIR / f"output_{notebook_name}"
+    if not notebook_path.exists():
+        available = list(NOTEBOOKS_DIR.glob("*.ipynb")) if NOTEBOOKS_DIR.exists() else []
+        return f"Error: Notebook not found: {notebook_path}\nAvailable: {available}", ""
+    progress(0.1, desc="Starting notebook execution...")
+    try:
+        cmd = [
+            sys.executable, "-m", "nbconvert",
+            "--to", "notebook", "--execute",
+            "--output", str(output_path.absolute()),
+            "--ExecutePreprocessor.timeout=600",
+            "--ExecutePreprocessor.kernel_name=python3",
+            str(notebook_path.absolute()),
+        ]
+        progress(0.3, desc="Executing cells...")
+        result = subprocess.run(cmd, capture_output=True, text=True, cwd=str(NOTEBOOKS_DIR), timeout=900)
+        progress(0.8, desc="Processing output...")
+        if result.returncode == 0:
+            if output_path.exists():
+                with open(output_path, "r") as f:
+                    executed_nb = json.load(f)
+                outputs = []
+                for i, cell in enumerate(executed_nb.get("cells", []), 1):
+                    if cell.get("cell_type") == "code":
+                        for out in cell.get("outputs", []):
+                            if "text" in out:
+                                outputs.append(f"Cell {i}:\n{''.join(out['text'])}")
+                            elif "data" in out and "text/plain" in out["data"]:
+                                outputs.append(f"Cell {i}:\n{''.join(out['data']['text/plain'])}")
+                progress(1.0, desc="Complete!")
+                return "Notebook executed successfully!", "\n\n".join(outputs)
+            else:
+                return "Notebook executed but output file not found", result.stdout
+        else:
+            return f"Execution failed:\n{result.stderr}", result.stdout
+    except subprocess.TimeoutExpired:
+        return "Error: Notebook execution timed out (15 min limit)", ""
+    except Exception as e:
+        return f"Error executing notebook: {str(e)}", ""
+
+
+def run_notebook_cell(notebook_name: str, cell_number: int) -> str:
+    """Execute a single cell from a notebook"""
+    notebook_path = NOTEBOOKS_DIR / notebook_name
+    if not notebook_path.exists():
+        return f"Error: Notebook not found at {notebook_path}"
+    try:
+        original_cwd = os.getcwd()
+        os.chdir(NOTEBOOKS_DIR)
+        with open(notebook_path, "r") as f:
+            nb = json.load(f)
+        cells = [c for c in nb.get("cells", []) if c.get("cell_type") == "code"]
+        if cell_number < 1 or cell_number > len(cells):
+            os.chdir(original_cwd)
+            return f"Error: Cell {cell_number} not found. Available: 1-{len(cells)}"
+        cell = cells[int(cell_number) - 1]
+        source = "".join(cell.get("source", []))
+        import io
+        from contextlib import redirect_stdout, redirect_stderr
+
+        namespace = {"__name__": "__main__", "__file__": str(notebook_path)}
+        stdout_capture = io.StringIO()
+        stderr_capture = io.StringIO()
+        with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
+            try:
+                exec(source, namespace)
+            except Exception as e:
+                os.chdir(original_cwd)
+                return f"Error: {str(e)}"
+        os.chdir(original_cwd)
+        output = stdout_capture.getvalue()
+        errors = stderr_capture.getvalue()
+        result_text = f"### Cell {int(cell_number)} Output:\n"
+        if output:
+            result_text += f"```\n{output}\n```\n"
+        if errors:
+            result_text += f"\n**Warnings/Errors:**\n```\n{errors}\n```"
+        if not output and not errors:
+            result_text += "*(No output)*"
+        return result_text
+    except Exception as e:
+        try:
+            os.chdir(original_cwd)
+        except Exception:
+            pass
+        return f"Error: {str(e)}"
+
+
+# ============================================================================
+# MODEL TRAINING  (existing Gradio functionality)
+# ============================================================================
+
+class SecurityModelTrainer:
+    def __init__(self):
+        self.scaler = StandardScaler()
+        self.label_encoder = LabelEncoder()
+
+    def prepare_data(self, df: pd.DataFrame, target_col: str = "label") -> Tuple:
+        if target_col not in df.columns:
+            raise ValueError(f"Target column '{target_col}' not found")
+        X = df.drop(columns=[target_col])
+        y = df[target_col]
+        X = X.select_dtypes(include=[np.number]).fillna(0)
+        if y.dtype == "object":
+            y = self.label_encoder.fit_transform(y)
+        X_scaled = self.scaler.fit_transform(X)
+        return train_test_split(X_scaled, y, test_size=0.2, random_state=42)
+
+    def train_model(self, model_type: str, X_train, y_train):
+        if model_type not in MODEL_TYPES:
+            raise ValueError(f"Unknown model type: {model_type}")
+        model_class = MODEL_TYPES[model_type]
+        if model_type == "Isolation Forest (Anomaly)":
+            model = model_class(contamination=0.1, random_state=42)
+        else:
+            model = model_class(random_state=42)
+        model.fit(X_train, y_train)
+        return model
+
+    def evaluate_model(self, model, X_test, y_test) -> Dict:
+        y_pred = model.predict(X_test)
+        return {
+            "accuracy": accuracy_score(y_test, y_pred),
+            "f1_score": f1_score(y_test, y_pred, average="weighted", zero_division=0),
+        }
+
+
+trainer = SecurityModelTrainer()
+
+
+def train_model_from_data(data_file, model_type: str, task: str, progress=gr.Progress()):
+    """Train model from uploaded data"""
+    if data_file is None:
+        return "Please upload a CSV file", None, None
+    progress(0.1, desc="Loading data...")
+    try:
+        df = pd.read_csv(data_file.name)
+        progress(0.3, desc="Preparing data...")
+        X_train, X_test, y_train, y_test = trainer.prepare_data(df)
+        progress(0.5, desc=f"Training {model_type}...")
+        model = trainer.train_model(model_type, X_train, y_train)
+        progress(0.8, desc="Evaluating model...")
+        metrics = trainer.evaluate_model(model, X_test, y_test)
+        model_name = f"{task.lower().replace(' ', '_')}_{model_type.lower().replace(' ', '_')}"
+        model_path = MODELS_DIR / f"{model_name}.pkl"
+        joblib.dump(model, model_path)
+        progress(1.0, desc="Complete!")
+        result = f"""
+## Training Complete!
+
+**Task:** {task}
+**Model:** {model_type}
+**Samples:** {len(df)}
+
+### Metrics
+- Accuracy: {metrics['accuracy']:.4f}
+- F1 Score: {metrics['f1_score']:.4f}
+
+**Model saved to:** {model_path}
+"""
+        return result, str(model_path), json.dumps(metrics, indent=2)
+    except Exception as e:
+        return f"Error: {str(e)}", None, None
+
+
+def run_inference(model_file, features_text: str):
+    """Run inference with a trained model"""
+    if model_file is None:
+        return "Please upload a model file"
+    try:
+        model = joblib.load(model_file.name)
+        features = json.loads(features_text)
+        X = np.array([list(features.values())])
+        prediction = model.predict(X)[0]
+        result = {"prediction": int(prediction), "features_used": len(features)}
+        if hasattr(model, "predict_proba"):
+            proba = model.predict_proba(X)[0]
+            result["confidence"] = float(max(proba))
+            result["probabilities"] = {str(i): float(p) for i, p in enumerate(proba)}
+        return json.dumps(result, indent=2)
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+
+def list_trained_models():
+    models = list(MODELS_DIR.glob("*.pkl"))
+    if not models:
+        return "No trained models found"
+    output = "## Trained Models\n\n"
+    for model_path in models:
+        size_kb = model_path.stat().st_size / 1024
+        output += f"- **{model_path.name}** ({size_kb:.1f} KB)\n"
+    return output
+
+
+# ============================================================================
+# FASTAPI APP  (REST endpoints for Heroku backend mlService.js)
+# ============================================================================
+
+api = FastAPI(title="CyberForge AI API", version="1.0.0")
+
+api.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+@api.get("/health")
+async def api_health():
+    return {
+        "status": "healthy",
+        "timestamp": datetime.utcnow().isoformat(),
+        "services": {
+            "gemini": gemini_service.ready,
+            "ml_models": ml_loader.ready,
+            "models_loaded": list(ml_loader.models.keys()),
+            "transformers": transformer_loader.transformers_available,
+            "transformer_models_loaded": list(transformer_loader.pipelines.keys()),
+            "gradio_ui": True,
+        },
+        "version": "2.0.0",
+        "new_endpoints": [
+            "POST /api/v2/batch-analyze",
+            "GET  /api/v2/explain/{model_name}",
+            "POST /api/v2/ioc-scan",
+            "POST /api/v2/url-enrich",
+            "POST /api/v2/chat (context-grounded security chatbot, Wave 2)",
+            "POST /api/v2/security-chat",
+        ],
+    }
+
+
+@api.post("/analyze")
+async def api_analyze(request: Request):
+    """Main analysis endpoint – called by backend mlService.chatWithAI()"""
+    body = await request.json()
+    query = body.get("query", "")
+    context = body.get("context", {})
+    history = body.get("conversation_history", [])
+    result = gemini_service.analyze(query, context, history)
+    return result
+
+
+@api.post("/analyze-url")
+async def api_analyze_url(request: Request):
+    """URL analysis – called by backend mlService.analyzeWebsite()
+    Returns per-model predictions + unified aggregate with risk_score (0-100).
+    All models now consistently include threat_score (probability of class-1).
+    """
+    body = await request.json()
+    url = body.get("url", "")
+    if not url:
+        return JSONResponse(status_code=400, content={"detail": "URL required"})
+    features = extract_url_features(url)
+    predictions = {}
+    for model_name in ml_loader.MODEL_NAMES:
+        predictions[model_name] = ml_loader.predict(model_name, features)
+
+    # threat_score is now always present on every prediction (ml_model + heuristic)
+    scores = [p.get("threat_score", 0.2) for p in predictions.values()]
+    avg_score = sum(scores) / len(scores) if scores else 0
+    max_score = max(scores) if scores else 0
+
+    # Heuristic risk factors derived from URL features
+    risk_factors = []
+    if not features.get("is_https"):
+        risk_factors.append("no_https")
+    if features.get("has_ip_address"):
+        risk_factors.append("ip_in_url")
+    if features.get("has_suspicious_tld"):
+        risk_factors.append("suspicious_tld")
+    if features.get("url_length", 0) > 100:
+        risk_factors.append("long_url")
+    if features.get("has_at_symbol"):
+        risk_factors.append("at_symbol_in_url")
+    if features.get("subdomain_count", 0) > 2:
+        risk_factors.append("excessive_subdomains")
+    if features.get("special_char_count", 0) > 10:
+        risk_factors.append("high_special_chars")
+    if features.get("has_double_slash"):
+        risk_factors.append("double_slash_in_path")
+
+    # Overall risk level using the max score (conservative)
+    if max_score > 0.8:
+        overall_risk = "critical"
+    elif max_score > 0.6:
+        overall_risk = "high"
+    elif max_score > 0.4:
+        overall_risk = "medium"
+    else:
+        overall_risk = "low"
+
+    # risk_score 0-100 for the Heroku backend (used in riskScore field)
+    risk_score_100 = round(max_score * 100, 1)
+
+    return {
+        "url": url,
+        "aggregate": {
+            "average_threat_score": round(avg_score, 3),
+            "max_threat_score": round(max_score, 3),
+            "risk_score": risk_score_100,
+            "overall_risk_level": overall_risk,
+            "risk_factors": risk_factors,
+            "models_flagged": sum(1 for p in predictions.values() if p.get("prediction", 0) == 1),
+            "models_total": len(predictions),
+        },
+        "model_predictions": predictions,
+        "features_analyzed": features,
+        "timestamp": datetime.utcnow().isoformat(),
+    }
+
+
+@api.post("/scan-threats")
+async def api_scan_threats(request: Request):
+    """Threat scanning – called by backend mlService.scanForThreats()"""
+    body = await request.json()
+    query = json.dumps(body.get("data", body), indent=1)[:3000]
+    result = gemini_service.analyze(f"Scan these indicators for threats:\n{query}")
+    return result
+
+
+@api.post("/api/insights/generate")
+async def api_generate_insights(request: Request):
+    """AI insights – called by backend mlService.getAIInsights()"""
+    body = await request.json()
+    query = body.get("query", "")
+    context = body.get("context", {})
+    result = gemini_service.analyze(query, context)
+    return {
+        "insights": result.get("response", ""),
+        "confidence": result.get("confidence", 0),
+        "timestamp": datetime.utcnow().isoformat(),
+    }
+
+
+@api.post("/api/models/predict")
+async def api_model_predict(request: Request):
+    """Model prediction – called by backend mlService.getModelPrediction()"""
+    body = await request.json()
+    model_type = body.get("model_type", "phishing_detection")
+    input_data = body.get("input_data", {})
+    result = ml_loader.predict(model_type, input_data)
+    return result
+
+
+@api.get("/api/models/list")
+@api.get("/models")
+async def api_list_models():
+    result = []
+    for name in ml_loader.MODEL_NAMES:
+        result.append({
+            "name": name,
+            "loaded": name in ml_loader.models,
+            "source": "ml_model" if name in ml_loader.models else "heuristic",
+        })
+    return {"models": result, "total": len(ml_loader.MODEL_NAMES), "loaded": len(ml_loader.models)}
+
+
+# ============================================================================
+# PHASE-3 ENDPOINTS — Real HF Transformer models
+# ============================================================================
+
+@api.post("/api/v2/url-classify")
+async def api_v2_url_classify(request: Request):
+    """URL phishing/malware classification using elftsdmr/malware-url-detect (BERT).
+    Body: { "url": "https://..." }"""
+    body = await request.json()
+    url = body.get("url", "").strip()
+    if not url:
+        return JSONResponse(status_code=400, content={"detail": "url required"})
+    return transformer_loader.predict_url_phishing(url)
+
+
+@api.post("/api/v2/dga-detect")
+async def api_v2_dga_detect(request: Request):
+    """DGA-generated domain detection using YangYang-Research/dga-detection.
+    Body: { "domain": "abc123xyz.com" }"""
+    body = await request.json()
+    domain = body.get("domain", "").strip()
+    if not domain:
+        return JSONResponse(status_code=400, content={"detail": "domain required"})
+    return transformer_loader.predict_dga(domain)
+
+
+@api.post("/api/v2/security-chat")
+async def api_v2_security_chat(request: Request):
+    """Cybersecurity Q&A via ZySec-AI/SecurityLLM (HF Inference API).
+    Body: { "query": "...", "max_tokens": 512 }"""
+    body = await request.json()
+    query = body.get("query", "").strip()
+    if not query:
+        return JSONResponse(status_code=400, content={"detail": "query required"})
+    max_tokens = int(body.get("max_tokens", 512))
+    result = transformer_loader.security_chat(query, max_tokens=max_tokens)
+    # Auto-fallback to Gemini when LLM unavailable
+    if result.get("error") and gemini_service.ready:
+        gemini_result = gemini_service.analyze(query)
+        gemini_result["source"] = "gemini-fallback"
+        gemini_result["llm_error"] = result["error"]
+        return gemini_result
+    return result
+
+
+@api.get("/api/v2/status")
+async def api_v2_status():
+    """Status of phase-3 transformer models — what's loaded, what failed, what's available."""
+    return transformer_loader.status()
+
+
+@api.post("/api/v2/batch-analyze")
+async def api_v2_batch_analyze(request: Request):
+    """Batch URL analysis for the distributed agent system.
+    Accepts up to 20 URLs in one call — avoids the per-URL HTTP overhead
+    when multiple agents run concurrent scans.
+
+    Body:
+      { "urls": ["https://...", "https://...", ...] }
+
+    Response:
+      { "results": [ { ...same shape as /analyze-url... }, ... ], "total": N, "elapsed_ms": N }
+    """
+    import time
+    body = await request.json()
+    urls = body.get("urls", [])
+    if not urls:
+        return JSONResponse(status_code=400, content={"detail": "urls array required"})
+    if len(urls) > 20:
+        return JSONResponse(status_code=400, content={"detail": "Maximum 20 URLs per batch"})
+
+    t0 = time.time()
+    results = []
+    for url in urls:
+        url = str(url).strip()
+        if not url:
+            results.append({"url": url, "error": "empty url"})
+            continue
+        features = extract_url_features(url)
+        predictions = {}
+        for model_name in ml_loader.MODEL_NAMES:
+            predictions[model_name] = ml_loader.predict(model_name, features)
+        scores = [p.get("threat_score", 0.2) for p in predictions.values()]
+        avg_score = sum(scores) / len(scores) if scores else 0
+        max_score = max(scores) if scores else 0
+        risk_factors = []
+        if not features.get("is_https"):
+            risk_factors.append("no_https")
+        if features.get("has_ip_address"):
+            risk_factors.append("ip_in_url")
+        if features.get("has_suspicious_tld"):
+            risk_factors.append("suspicious_tld")
+        if features.get("url_length", 0) > 100:
+            risk_factors.append("long_url")
+        if features.get("has_at_symbol"):
+            risk_factors.append("at_symbol_in_url")
+        if features.get("subdomain_count", 0) > 2:
+            risk_factors.append("excessive_subdomains")
+        if max_score > 0.8:
+            overall_risk = "critical"
+        elif max_score > 0.6:
+            overall_risk = "high"
+        elif max_score > 0.4:
+            overall_risk = "medium"
+        else:
+            overall_risk = "low"
+        results.append({
+            "url": url,
+            "aggregate": {
+                "average_threat_score": round(avg_score, 3),
+                "max_threat_score": round(max_score, 3),
+                "risk_score": round(max_score * 100, 1),
+                "overall_risk_level": overall_risk,
+                "risk_factors": risk_factors,
+                "models_flagged": sum(1 for p in predictions.values() if p.get("prediction", 0) == 1),
+            },
+            "model_predictions": predictions,
+            "features_analyzed": features,
+        })
+
+    elapsed_ms = round((time.time() - t0) * 1000, 1)
+    return {
+        "results": results,
+        "total": len(results),
+        "elapsed_ms": elapsed_ms,
+        "timestamp": datetime.utcnow().isoformat(),
+    }
+
+
+@api.get("/api/v2/explain/{model_name}")
+async def api_v2_explain(model_name: str):
+    """Feature importance / explainability for a loaded ML model.
+    Supports RandomForest, GradientBoosting (gini importance) and LogisticRegression (coefficients).
+
+    Path param: model_name — one of phishing_detection | malware_detection |
+                              anomaly_detection | web_attack_detection
+    """
+    return ml_loader.feature_importance(model_name)
+
+
+@api.post("/api/v2/ioc-scan")
+async def api_v2_ioc_scan(request: Request):
+    """Multi-field Indicator of Compromise scanner.
+    Runs all available ML models + DGA detector + BERT classifier against
+    multiple IOC types in a single call.
+
+    Body:
+      {
+        "url":    "https://...",           # optional
+        "domain": "example.com",           # optional
+        "ip":     "1.2.3.4",              # optional
+        "hash":   "abc123...",             # optional (future use)
+        "context": { ... }                 # optional — passed to Gemini if available
+      }
+
+    Response: per-indicator results with unified risk summary.
+    """
+    body = await request.json()
+    url     = body.get("url", "").strip()
+    domain  = body.get("domain", "").strip()
+    context = body.get("context", {})
+
+    ioc_results: Dict[str, Any] = {}
+    all_scores: List[float] = []
+
+    # --- URL analysis ---
+    if url:
+        features = extract_url_features(url)
+        url_preds = {name: ml_loader.predict(name, features) for name in ml_loader.MODEL_NAMES}
+        url_scores = [p.get("threat_score", 0.2) for p in url_preds.values()]
+        url_max = max(url_scores) if url_scores else 0
+        all_scores.append(url_max)
+
+        # BERT-based URL classifier (lazy-loaded)
+        bert_result = transformer_loader.predict_url_phishing(url)
+        ioc_results["url"] = {
+            "value": url,
+            "ml_predictions": url_preds,
+            "bert_classification": bert_result,
+            "max_threat_score": round(url_max, 3),
+        }
+
+    # --- Domain analysis ---
+    effective_domain = domain
+    if not effective_domain and url:
+        from urllib.parse import urlparse as _urlparse
+        effective_domain = _urlparse(url).hostname or ""
+
+    if effective_domain:
+        dga_result = transformer_loader.predict_dga(effective_domain)
+        dga_score = dga_result.get("threat_score", 0.0)
+        all_scores.append(dga_score)
+        ioc_results["domain"] = {
+            "value": effective_domain,
+            "dga_detection": dga_result,
+        }
+
+    # --- Overall risk summary ---
+    overall_score = max(all_scores) if all_scores else 0.0
+    if overall_score > 0.8:
+        risk_level = "critical"
+    elif overall_score > 0.6:
+        risk_level = "high"
+    elif overall_score > 0.4:
+        risk_level = "medium"
+    elif overall_score > 0.0:
+        risk_level = "low"
+    else:
+        risk_level = "unknown"
+
+    # Optional Gemini enrichment when available
+    ai_analysis = None
+    if gemini_service.ready and (url or domain):
+        target = url or domain
+        ai_analysis = gemini_service.analyze(
+            f"Scan these IOC indicators for threats: url={target}, domain={effective_domain}",
+            context=context,
+        )
+
+    return {
+        "ioc_results": ioc_results,
+        "summary": {
+            "overall_risk_level": risk_level,
+            "overall_threat_score": round(overall_score, 3),
+            "risk_score": round(overall_score * 100, 1),
+            "indicators_analyzed": len(ioc_results),
+        },
+        "ai_analysis": ai_analysis,
+        "timestamp": datetime.utcnow().isoformat(),
+    }
+
+
+@api.post("/api/v2/url-enrich")
+async def api_v2_url_enrich(request: Request):
+    """Enriched URL analysis: combines heuristic + ML models + BERT transformer
+    into a single scored response with explanations.
+    This is the richer alternative to /analyze-url for the AI Deep Scan mode.
+
+    Body: { "url": "https://..." }
+    """
+    body = await request.json()
+    url = body.get("url", "").strip()
+    if not url:
+        return JSONResponse(status_code=400, content={"detail": "url required"})
+
+    features = extract_url_features(url)
+
+    # Core ML predictions
+    ml_predictions = {name: ml_loader.predict(name, features) for name in ml_loader.MODEL_NAMES}
+    ml_scores = [p.get("threat_score", 0.2) for p in ml_predictions.values()]
+    ml_max = max(ml_scores) if ml_scores else 0
+
+    # BERT transformer
+    bert_result = transformer_loader.predict_url_phishing(url)
+    bert_score = bert_result.get("threat_score", 0.0) if not bert_result.get("error") else None
+
+    # DGA check on hostname
+    from urllib.parse import urlparse as _urlparse
+    hostname = _urlparse(url).hostname or ""
+    dga_result = transformer_loader.predict_dga(hostname) if hostname else None
+    dga_score = dga_result.get("threat_score", 0.0) if dga_result and not dga_result.get("error") else 0.0
+
+    # Feature importances for top model
+    top_model = ml_loader.MODEL_NAMES[0]
+    top_score = ml_scores[0]
+    for i, score in enumerate(ml_scores):
+        if score > top_score:
+            top_score = score
+            top_model = ml_loader.MODEL_NAMES[i]
+    explainability = ml_loader.feature_importance(top_model)
+
+    # Fuse scores: weighted average of ML (60%), BERT (30% if available), DGA (10%)
+    if bert_score is not None:
+        fused_score = 0.60 * ml_max + 0.30 * bert_score + 0.10 * dga_score
+    else:
+        fused_score = 0.80 * ml_max + 0.20 * dga_score
+
+    if fused_score > 0.8:
+        risk_level = "critical"
+    elif fused_score > 0.6:
+        risk_level = "high"
+    elif fused_score > 0.4:
+        risk_level = "medium"
+    else:
+        risk_level = "low"
+
+    risk_factors = []
+    if not features.get("is_https"):
+        risk_factors.append("no_https")
+    if features.get("has_ip_address"):
+        risk_factors.append("ip_in_url")
+    if features.get("has_suspicious_tld"):
+        risk_factors.append("suspicious_tld")
+    if features.get("url_length", 0) > 100:
+        risk_factors.append("long_url")
+    if features.get("has_at_symbol"):
+        risk_factors.append("at_symbol_in_url")
+    if features.get("subdomain_count", 0) > 2:
+        risk_factors.append("excessive_subdomains")
+    if dga_result and dga_result.get("is_threat"):
+        risk_factors.append("dga_domain")
+    if bert_result and bert_result.get("is_threat"):
+        risk_factors.append("bert_flagged_malicious")
+
+    return {
+        "url": url,
+        "risk_level": risk_level,
+        "risk_score": round(fused_score * 100, 1),
+        "fused_threat_score": round(fused_score, 4),
+        "components": {
+            "ml_max_threat_score": round(ml_max, 4),
+            "bert_threat_score": round(bert_score, 4) if bert_score is not None else None,
+            "dga_threat_score": round(dga_score, 4),
+        },
+        "ml_predictions": ml_predictions,
+        "bert_classification": bert_result,
+        "dga_detection": dga_result,
+        "explainability": explainability,
+        "risk_factors": risk_factors,
+        "features_analyzed": features,
+        "timestamp": datetime.utcnow().isoformat(),
+    }
+
+
+@api.post("/api/v2/chat")
+async def api_v2_chat(request: Request):
+    """Security Chatbot — Wave 2 implementation.
+
+    Context-grounded security assistant that:
+      1. Accepts a live context blob (telemetry, scan history, IOCs) and
+         grounds every answer in it (RAG-lite over the current request — no
+         persistent vector store, which requires paid GPU on free tier).
+      2. Translates raw ML signals / IOCs into plain human language when the
+         user includes raw data and asks for an explanation.
+      3. Falls back to the Security LLM (Mistral-7B via HF Inference API)
+         or the ML-powered heuristic analyzer when Gemini is offline.
+
+    NOTE ON REAL-TIME "LEARNING": HuggingFace Spaces cpu-basic does NOT
+    support fine-tuning or persistent model updates at runtime. What this
+    endpoint does instead is context injection — every turn receives the
+    latest system telemetry + scan history in the prompt, so the chatbot's
+    answers are grounded in live data without weight updates. This is the
+    correct RAG-lite pattern for free-tier inference.
+
+    Request body (stable contract):
+      {
+        "message": "What is the risk of this domain: evil.xyz?",
+        "session_id": "optional-uuid-for-conversation-continuity",
+        "conversation_history": [ {"role": "user", "content": "..."}, ... ],
+        "context": {
+          "telemetry": { "cpu": 42, "ram": 67, "net_in_kbps": 120, ... },
+          "recent_scans": [ { "url": "...", "risk_score": 78, "category": "phishing" }, ... ],
+          "active_threats": [ { "type": "phishing", "severity": "high", "source": "..." } ],
+          "behavioral_alerts": [ { "pattern": "...", "score": 0.9 } ],
+          "translate": true   # if true → force plain-language IOC translation mode
+        }
+      }
+
+    Response contract (preserved from Wave 1):
+      {
+        "response": "...",          # natural language answer grounded in context
+        "confidence": 0.0-1.0,
+        "risk_level": "low|medium|high|critical",
+        "risk_score": 0-100,
+        "model_used": "...",
+        "session_id": "...",
+        "sources": [ { "type": "...", "label": "...", "value": "..." } ],
+        "timestamp": "..."
+      }
+    """
+    body = await request.json()
+    message = body.get("message", "").strip()
+    session_id = body.get("session_id", "")
+    history = body.get("conversation_history", [])
+    context = body.get("context", {})
+
+    if not message:
+        return JSONResponse(status_code=400, content={"detail": "message required"})
+
+    # ── Build grounding context string and source citations ──────────────
+    grounding_lines: List[str] = []
+    sources: List[Dict] = []
+
+    # System telemetry
+    telemetry = context.get("telemetry", {})
+    if telemetry:
+        cpu = telemetry.get("cpu", telemetry.get("cpu_percent"))
+        ram = telemetry.get("ram", telemetry.get("ram_percent"))
+        net_in = telemetry.get("net_in_kbps", telemetry.get("net_in"))
+        telem_parts = []
+        if cpu is not None:
+            telem_parts.append(f"CPU {cpu}%")
+        if ram is not None:
+            telem_parts.append(f"RAM {ram}%")
+        if net_in is not None:
+            telem_parts.append(f"Net↓ {net_in} kbps")
+        if telem_parts:
+            grounding_lines.append(f"SYSTEM TELEMETRY (live, last 5s): {', '.join(telem_parts)}")
+            sources.append({"type": "telemetry", "label": "System Telemetry", "value": ", ".join(telem_parts)})
+
+    # Recent scan history
+    recent_scans = context.get("recent_scans", [])
+    if recent_scans:
+        scan_summary = "; ".join(
+            f"{s.get('url', 'unknown')} → {s.get('category', '?')} (score {s.get('risk_score', '?')}/100)"
+            for s in recent_scans[:5]
+        )
+        grounding_lines.append(f"RECENT URL SCANS (last 5): {scan_summary}")
+        sources.append({"type": "scan_history", "label": "Recent Scans", "value": f"{len(recent_scans)} scans"})
+
+    # Active threats
+    active_threats = context.get("active_threats", [])
+    if active_threats:
+        threat_summary = "; ".join(
+            f"{t.get('type', '?')} [{t.get('severity', '?')}] from {t.get('source', '?')}"
+            for t in active_threats[:5]
+        )
+        grounding_lines.append(f"ACTIVE THREATS: {threat_summary}")
+        sources.append({"type": "threats", "label": "Active Threats", "value": f"{len(active_threats)} active"})
+
+    # Behavioral alerts
+    behavioral_alerts = context.get("behavioral_alerts", [])
+    if behavioral_alerts:
+        alert_summary = "; ".join(
+            f"'{a.get('pattern', a.get('type', '?'))}' (score {a.get('score', '?')})"
+            for a in behavioral_alerts[:3]
+        )
+        grounding_lines.append(f"BEHAVIORAL ALERTS: {alert_summary}")
+        sources.append({"type": "behavioral", "label": "Behavioral Alerts", "value": f"{len(behavioral_alerts)} alerts"})
+
+    # Machine→Human translation mode
+    translate_mode = context.get("translate", False)
+    raw_data = context.get("raw_data", "")
+    if translate_mode and raw_data:
+        grounding_lines.append(
+            f"RAW DATA TO TRANSLATE INTO PLAIN LANGUAGE:\n{str(raw_data)[:1500]}"
+        )
+        sources.append({"type": "raw_translation", "label": "Raw Signal", "value": str(raw_data)[:100]})
+
+    # Construct enriched prompt
+    grounding_block = "\n".join(grounding_lines)
+    translate_instruction = (
+        "\n\nIMPORTANT: The user has requested plain-language translation of raw machine data. "
+        "Explain all technical signals, scores, and IOCs in simple, non-technical language first, "
+        "then give actionable security advice." if translate_mode else ""
+    )
+
+    conversation_instruction = ""
+    if history:
+        conversation_instruction = "\n\nCONVERSATION CONTEXT (recent turns):\n" + "\n".join(
+            f"{h.get('role','?').upper()}: {str(h.get('content',''))[:300]}"
+            for h in history[-6:]
+        )
+
+    enriched_message = message
+    if grounding_block:
+        enriched_message = (
+            f"[LIVE SYSTEM CONTEXT — use this to ground your answer]\n"
+            f"{grounding_block}\n"
+            f"{translate_instruction}"
+            f"{conversation_instruction}\n\n"
+            f"[USER MESSAGE]\n{message}"
+        )
+
+    # ── Run ML analysis on any URLs in the message (always available) ──
+    import re as _re
+    url_matches = _re.findall(r'https?://[^\s"\'<>]+', message)
+    ml_url_context: List[str] = []
+    detected_risk_score = 0.0
+    detected_risk_level = "unknown"
+
+    if url_matches and ml_loader.ready:
+        for url in url_matches[:3]:
+            features = extract_url_features(url)
+            preds = {name: ml_loader.predict(name, features) for name in ml_loader.MODEL_NAMES}
+            scores = [p.get("threat_score", 0.2) for p in preds.values()]
+            mx = max(scores) if scores else 0
+            if mx > detected_risk_score:
+                detected_risk_score = mx
+            flagged = [n for n, p in preds.items() if p.get("prediction", 0) == 1]
+            ml_url_context.append(
+                f"ML scan of {url[:60]}: max_threat={mx:.2%}, flagged_by={flagged or 'none'}"
+            )
+            sources.append({"type": "ml_scan", "label": f"ML scan: {url[:50]}", "value": f"risk {mx:.0%}"})
+
+        if ml_url_context:
+            enriched_message += "\n\n[ML PRE-SCAN RESULTS]\n" + "\n".join(ml_url_context)
+
+    if detected_risk_score > 0.8:
+        detected_risk_level = "critical"
+    elif detected_risk_score > 0.6:
+        detected_risk_level = "high"
+    elif detected_risk_score > 0.4:
+        detected_risk_level = "medium"
+    elif detected_risk_score > 0.0:
+        detected_risk_level = "low"
+
+    # ── Grounding: add ML model state to context ──────────────────────
+    if ml_loader.ready:
+        n_models = len(ml_loader.models)
+        sources.append({"type": "ml_models", "label": "ML Models", "value": f"{n_models}/4 loaded"})
+
+    # ── LLM cascade: Gemini → Security LLM → ML heuristic ────────────
+    # Path 1: Gemini (premium, requires GEMINI_API_KEY env var)
+    if gemini_service.ready:
+        result = gemini_service.analyze(enriched_message, context=context, history=history)
+        # Merge detected risk if ML gave a stronger signal
+        if detected_risk_score > 0 and detected_risk_level != "unknown":
+            ml_score_100 = round(detected_risk_score * 100, 1)
+            if ml_score_100 > result.get("risk_score", 0):
+                result["risk_score"] = ml_score_100
+                result["risk_level"] = detected_risk_level
+        result["session_id"] = session_id
+        result["sources"] = sources
+        return result
+
+    # Path 2: Security LLM via HF Inference API (Mistral-7B)
+    llm_result = transformer_loader.security_chat(enriched_message, max_tokens=600)
+    if llm_result.get("response") and not llm_result.get("error"):
+        text = llm_result["response"]
+        text_lower = text.lower()
+        if "critical" in text_lower:
+            risk_level, risk_score = "critical", 85.0
+        elif "high" in text_lower:
+            risk_level, risk_score = "high", 70.0
+        elif "medium" in text_lower:
+            risk_level, risk_score = "medium", 45.0
+        else:
+            risk_level, risk_score = "low", 20.0
+        # Override with ML scan if stronger
+        if detected_risk_score > 0:
+            ml_score_100 = round(detected_risk_score * 100, 1)
+            if ml_score_100 > risk_score:
+                risk_score = ml_score_100
+                risk_level = detected_risk_level
+        return {
+            "response": text,
+            "confidence": 0.72,
+            "risk_level": risk_level,
+            "risk_score": risk_score,
+            "model_used": llm_result.get("model_id", transformer_loader.SECURITY_LLM_REPO),
+            "session_id": session_id,
+            "sources": sources,
+            "timestamp": datetime.utcnow().isoformat(),
+        }
+
+    # Path 3: ML heuristic fallback (always available)
+    fallback = gemini_service._fallback(enriched_message)
+    # Override with ML scan risk if available
+    if detected_risk_score > 0:
+        ml_score_100 = round(detected_risk_score * 100, 1)
+        if ml_score_100 > fallback.get("risk_score", 0):
+            fallback["risk_score"] = ml_score_100
+            fallback["risk_level"] = detected_risk_level
+    fallback["session_id"] = session_id
+    fallback["sources"] = sources
+    return fallback
+
+
+@api.post("/api/analysis/network")
+async def api_analyze_network(request: Request):
+    """Network traffic analysis – called by backend mlService.analyzeNetworkTraffic()"""
+    body = await request.json()
+    traffic_data = body.get("traffic_data", {})
+    result = gemini_service.analyze(
+        f"Analyze this network traffic for security threats:\n{json.dumps(traffic_data, indent=1)[:3000]}"
+    )
+    return result
+
+
+@api.post("/api/ai/execute-task")
+async def api_execute_task(request: Request):
+    """AI task execution – called by backend mlService.executeAITask()"""
+    body = await request.json()
+    task_type = body.get("task_type", "")
+    task_data = body.get("task_data", {})
+    result = gemini_service.analyze(
+        f"Execute cybersecurity task '{task_type}':\n{json.dumps(task_data, indent=1)[:3000]}"
+    )
+    return result
+
+
+@api.post("/api/browser/analyze")
+async def api_analyze_browser(request: Request):
+    """Browser session analysis"""
+    body = await request.json()
+    session = body.get("session_data", body)
+    result = gemini_service.analyze(
+        f"Analyze this browser session for security threats:\n{json.dumps(session, indent=1)[:3000]}"
+    )
+    return result
+
+
+@api.post("/api/threat-feeds/analyze")
+async def api_analyze_threat_feed(request: Request):
+    """Threat feed analysis"""
+    body = await request.json()
+    result = gemini_service.analyze(
+        f"Analyze this threat feed:\n{json.dumps(body, indent=1)[:3000]}"
+    )
+    return result
+
+
+@api.post("/api/datasets/process")
+async def api_process_dataset(request: Request):
+    body = await request.json()
+    return {
+        "status": "processed",
+        "dataset_name": body.get("dataset_name", ""),
+        "timestamp": datetime.utcnow().isoformat(),
+    }
+
+
+# ============================================================================
+# GRADIO INTERFACE  (UI tabs – notebooks, training, inference, models, API)
+# ============================================================================
+
+def create_interface():
+    with gr.Blocks(title="CyberForge AI") as demo:
+        gr.Markdown("""
+        # 🔐 CyberForge AI - ML Training Platform
+
+        Train cybersecurity ML models and run Jupyter notebooks on Hugging Face.
+        """)
+
+        with gr.Tabs():
+            # ============ NOTEBOOKS TAB ============
+            with gr.TabItem("📓 Notebooks"):
+                gr.Markdown("### Run ML Pipeline Notebooks\nExecute the CyberForge ML notebooks directly in the cloud.")
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        notebook_dropdown = gr.Dropdown(
+                            choices=get_available_notebooks(),
+                            label="Select Notebook",
+                            value=get_available_notebooks()[0] if get_available_notebooks() else None,
+                        )
+                        refresh_btn = gr.Button("🔄 Refresh List")
+                        view_btn = gr.Button("👁 View Content", variant="secondary")
+                        execute_btn = gr.Button("▶ Execute Notebook", variant="primary")
+                        gr.Markdown("### Run Single Cell")
+                        cell_number = gr.Number(label="Cell Number", value=1, minimum=1)
+                        run_cell_btn = gr.Button("Run Cell")
+                    with gr.Column(scale=2):
+                        notebook_status = gr.Markdown("Select a notebook to view or execute.")
+                        notebook_output = gr.Markdown("", label="Output")
+
+                def refresh_notebooks():
+                    notebooks = get_available_notebooks()
+                    return gr.update(choices=notebooks, value=notebooks[0] if notebooks else None)
+
+                refresh_btn.click(refresh_notebooks, outputs=notebook_dropdown)
+                view_btn.click(read_notebook_content, inputs=notebook_dropdown, outputs=notebook_output)
+                execute_btn.click(execute_notebook, inputs=notebook_dropdown, outputs=[notebook_status, notebook_output])
+                run_cell_btn.click(run_notebook_cell, inputs=[notebook_dropdown, cell_number], outputs=notebook_output)
+
+            # ============ TRAIN MODEL TAB ============
+            with gr.TabItem("🎯 Train Model"):
+                gr.Markdown("### Train a Security ML Model\nUpload your dataset and train a model for threat detection.")
+                with gr.Row():
+                    with gr.Column():
+                        task_dropdown = gr.Dropdown(choices=SECURITY_TASKS, label="Security Task", value="Phishing Detection")
+                        model_dropdown = gr.Dropdown(choices=list(MODEL_TYPES.keys()), label="Model Type", value="Random Forest")
+                        data_upload = gr.File(label="Upload Training Data (CSV)", file_types=[".csv"])
+                        train_btn = gr.Button("🚀 Train Model", variant="primary")
+                    with gr.Column():
+                        train_output = gr.Markdown("Upload data and click Train to begin.")
+                        model_path_output = gr.Textbox(label="Model Path", visible=False)
+                        metrics_output = gr.Textbox(label="Metrics JSON", visible=False)
+                train_btn.click(train_model_from_data, inputs=[data_upload, model_dropdown, task_dropdown], outputs=[train_output, model_path_output, metrics_output])
+
+            # ============ INFERENCE TAB ============
+            with gr.TabItem("🔍 Inference"):
+                gr.Markdown("### Run Model Inference\nLoad a trained model and make predictions.")
+                with gr.Row():
+                    with gr.Column():
+                        model_upload = gr.File(label="Upload Model (.pkl)")
+                        features_input = gr.Textbox(label="Features (JSON)", value='{"url_length": 50, "has_https": 1, "digit_count": 5}', lines=5)
+                        predict_btn = gr.Button("🎯 Predict", variant="primary")
+                    with gr.Column():
+                        prediction_output = gr.Textbox(label="Prediction Result", lines=10)
+                predict_btn.click(run_inference, inputs=[model_upload, features_input], outputs=prediction_output)
+
+            # ============ MODELS TAB ============
+            with gr.TabItem("📦 Models"):
+                gr.Markdown("### Trained Models")
+                models_list = gr.Markdown(list_trained_models())
+                refresh_models_btn = gr.Button("🔄 Refresh")
+                refresh_models_btn.click(list_trained_models, outputs=models_list)
+
+            # ============ API TAB ============
+            with gr.TabItem("🔌 API"):
+                gr.Markdown("""
+## API Integration — v2.0
+
+### Core Endpoints (stable, used by Heroku backend)
+
+| Method | Endpoint | Description |
+|--------|----------|-------------|
+| GET | `/health` | Health check — version, models, services |
+| POST | `/analyze` | AI chat / general analysis (Gemini or ML fallback) |
+| POST | `/analyze-url` | URL threat analysis — 4 ML models + risk factors |
+| POST | `/scan-threats` | Deep threat scan (Gemini or ML fallback) |
+| POST | `/api/insights/generate` | AI-generated security insights |
+| POST | `/api/models/predict` | Direct ML model prediction |
+| GET | `/models` | List loaded models with source info |
+
+### v2 Endpoints (new — for distributed agents + AI Deep Scan)
+
+| Method | Endpoint | Description |
+|--------|----------|-------------|
+| POST | `/api/v2/batch-analyze` | Batch analyze up to 20 URLs (for multi-agent) |
+| GET | `/api/v2/explain/{model_name}` | Feature importances / explainability |
+| POST | `/api/v2/url-classify` | BERT-based URL phishing classifier |
+| POST | `/api/v2/dga-detect` | DGA domain detection (entropy heuristic) |
+| POST | `/api/v2/ioc-scan` | Multi-IOC scanner: URL + domain + DGA + BERT |
+| POST | `/api/v2/url-enrich` | Enriched analysis: ML + BERT + DGA + explainability |
+| POST | `/api/v2/chat` | Security chatbot — context-grounded, RAG-lite, machine→human translation |
+| POST | `/api/v2/security-chat` | Security LLM Q&A (Mistral-7B via HF Inference API) |
+| GET | `/api/v2/status` | Transformer model load status |
+
+### Quick Examples
+
+```bash
+# Batch analyze URLs (for distributed agents)
+curl -X POST https://che237-cyberforge.hf.space/api/v2/batch-analyze \\
+  -H "Content-Type: application/json" \\
+  -d '{"urls": ["https://example.com", "http://evil.xyz/login?id=1"]}'
+
+# Enriched URL analysis (Deep Scan mode)
+curl -X POST https://che237-cyberforge.hf.space/api/v2/url-enrich \\
+  -H "Content-Type: application/json" \\
+  -d '{"url": "https://example.com"}'
+
+# Feature explainability
+curl https://che237-cyberforge.hf.space/api/v2/explain/phishing_detection
+
+# IOC multi-field scan
+curl -X POST https://che237-cyberforge.hf.space/api/v2/ioc-scan \\
+  -H "Content-Type: application/json" \\
+  -d '{"url": "http://192.168.1.1/login", "domain": "192.168.1.1"}'
+
+# Chatbot (Wave 2 extension point)
+curl -X POST https://che237-cyberforge.hf.space/api/v2/chat \\
+  -H "Content-Type: application/json" \\
+  -d '{"message": "Is this URL suspicious: http://login.g00gle.com.xyz/verify"}'
+```
+                """)
+
+        gr.Markdown("---\n**CyberForge AI** | [GitHub](https://github.com/Che237/cyberforge) | [Datasets](https://huggingface.co/datasets/Che237/cyberforge-datasets)")
+
+    return demo
+
+
+# ============================================================================
+# MAIN – Mount Gradio on FastAPI
+# ============================================================================
+
+# Initialize services at startup
+logger.info("🚀 Initializing CyberForge AI services...")
+gemini_service.initialize()
+ml_loader.initialize()
+logger.info("✅ Services initialized")
+
+# Create Gradio app and mount it on FastAPI
+demo = create_interface()
+app = gr.mount_gradio_app(api, demo, path="/")
+
+if __name__ == "__main__":
+    import uvicorn
+    port = int(os.environ.get("PORT", 7860))
+    uvicorn.run(app, host="0.0.0.0", port=port)

requirements.txt

@@ -0,0 +1,45 @@
+# CyberForge AI - HuggingFace Space
+# Combined Gradio UI + FastAPI REST API
+
+# Core ML
+scikit-learn>=1.5.0
+pandas>=2.2.0
+numpy>=2.0.0
+joblib>=1.4.0
+
+# Transformers (for phase-3 BERT URL classifier + DGA detector)
+torch>=2.2.0
+transformers>=4.40.0
+tokenizers>=0.19.0
+
+# Gradio UI
+gradio>=5.20.0
+
+# FastAPI REST endpoints (for Heroku backend)
+fastapi>=0.100.0
+uvicorn[standard]>=0.22.0
+
+# HF Hub
+huggingface_hub>=0.25.0
+
+# Feature Engineering
+tldextract>=5.1.0
+
+# Gemini AI (new SDK)
+google-genai>=1.0.0
+
+# Notebook execution
+nbformat>=5.10.0
+nbconvert>=7.16.0
+nest_asyncio>=1.6.0
+
+# HTTP client for notebooks
+httpx>=0.27.0
+
+# Data handling
+datasets>=2.20.0
+pyarrow>=17.0.0
+
+# Utilities
+requests>=2.31.0
+python-dotenv>=1.0.0