src/app/main.py

import os
import logging
import time
from datetime import datetime
from typing import Dict, List, Optional
import uvicorn
from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import numpy as np
from contextlib import asynccontextmanager

# Create logs directory if it doesn't exist
os.makedirs('logs', exist_ok=True)

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('logs/app.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# Global variables for model and tokenizer
model = None
tokenizer = None
model_loaded = False
model_info = {
    "model_name": "songhieng/roberta-phishing-content-detector-5.0",
    "loaded_at": None,
    "version": "5.0",
    "framework": "transformers"
}

class PredictionRequest(BaseModel):
    text: str = Field(..., description="Text content to analyze for phishing", min_length=1, max_length=10000)
    
class PredictionResponse(BaseModel):
    text: str
    score: float
    description: str
    processing_time_ms: float
    timestamp: str

class HealthResponse(BaseModel):
    status: str
    model_loaded: bool
    timestamp: str
    uptime_seconds: float

class BatchPredictionRequest(BaseModel):
    texts: List[str] = Field(..., description="List of texts to analyze", max_items=100)

# Application startup and shutdown events
@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup
    logger.info("Starting up the application...")
    await load_model()
    yield
    # Shutdown
    logger.info("Shutting down the application...")

app = FastAPI(
    title="RoBERTa Phishing Content Detector API",
    description="MLOps deployment of RoBERTa model for phishing content detection",
    version="5.0.0",
    lifespan=lifespan
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Startup time for uptime calculation
startup_time = time.time()

async def load_model():
    """Load the model and tokenizer"""
    global model, tokenizer, model_loaded, model_info
    
    try:
        logger.info("Loading model and tokenizer...")
        model_path = "models/roberta-phishing-detector"
        
        if not os.path.exists(model_path):
            logger.error(f"Model path {model_path} does not exist!")
            raise FileNotFoundError(f"Model not found at {model_path}")
        
        # Load tokenizer and model
        tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)
        model = AutoModelForSequenceClassification.from_pretrained(model_path, local_files_only=True)
        
        # Set model to evaluation mode
        model.eval()
        
        model_loaded = True
        model_info["loaded_at"] = datetime.now().isoformat()
        
        logger.info("Model and tokenizer loaded successfully!")
        
    except Exception as e:
        logger.error(f"Failed to load model: {str(e)}")
        model_loaded = False
        raise e

def predict_phishing(text: str) -> float:
    """Predict if text is phishing content and return a phishing score
    A higher score (closer to 1) indicates more likely to be phishing
    A lower score (closer to 0) indicates more likely to be legitimate
    """
    if not model_loaded:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    try:
        # Tokenize the input text
        inputs = tokenizer(
            text,
            truncation=True,
            padding=True,
            max_length=4096,
            return_tensors="pt"
        )
        
        # Make prediction
        with torch.no_grad():
            outputs = model(**inputs)
            predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
            
            # Get the phishing score (class 1 probability)
            phishing_score = float(predictions[0][1])
            
            return phishing_score
        
    except Exception as e:
        logger.error(f"Prediction error: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")

@app.middleware("http")
async def log_requests(request: Request, call_next):
    """Log all requests"""
    start_time = time.time()
    
    response = await call_next(request)
    
    process_time = time.time() - start_time
    logger.info(
        f"Request: {request.method} {request.url} - "
        f"Status: {response.status_code} - "
        f"Time: {process_time:.4f}s"
    )
    
    return response

@app.get("/", response_model=Dict)
async def root():
    """Root endpoint"""
    return {
        "message": "RoBERTa Phishing Content Detector API",
        "version": "1.0.0",
        "model": model_info["model_name"],
        "status": "healthy" if model_loaded else "unhealthy"
    }

@app.get("/health", response_model=HealthResponse)
async def health_check():
    """Health check endpoint for monitoring"""
    uptime = time.time() - startup_time
    
    return HealthResponse(
        status="healthy" if model_loaded else "unhealthy",
        model_loaded=model_loaded,
        timestamp=datetime.now().isoformat(),
        uptime_seconds=uptime
    )

@app.get("/model/info")
async def model_info_endpoint():
    """Get model information"""
    return {
        "model_info": model_info,
        "model_loaded": model_loaded,
        "torch_version": torch.__version__
    }

@app.post("/predict", response_model=PredictionResponse)
async def predict(request: PredictionRequest):
    """Predict if text content is phishing"""
    start_time = time.time()
    
    try:
        phishing_score = predict_phishing(request.text)
        
        # Generate description based on score
        if phishing_score < 0.2:
            classification = "Legitimate (Very Low Risk)"
        elif phishing_score < 0.4:
            classification = "Likely Legitimate (Low Risk)"
        elif phishing_score < 0.6:
            classification = "Uncertain (Medium Risk)"
        elif phishing_score < 0.8:
            classification = "Likely Phishing (High Risk)"
        else:
            classification = "Phishing (Very High Risk)"
            
        description = f"{classification}: Score {phishing_score:.4f} - Lower scores (closer to 0) indicate legitimate content, higher scores (closer to 1) indicate phishing/malicious content"
        
        processing_time = (time.time() - start_time) * 1000  # Convert to milliseconds
        
        response = PredictionResponse(
            text=request.text[:100] + "..." if len(request.text) > 100 else request.text,
            score=phishing_score,
            description=description,
            processing_time_ms=round(processing_time, 2),
            timestamp=datetime.now().isoformat()
        )
        
        logger.info(f"Prediction made: (phishing score: {phishing_score:.4f}, classification: {classification})")
        return response
        
    except Exception as e:
        logger.error(f"Prediction endpoint error: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/predict/batch")
async def predict_batch(request: BatchPredictionRequest):
    """Batch prediction endpoint"""
    start_time = time.time()
    
    try:
        results = []
        
        for text in request.texts:
            phishing_score = predict_phishing(text)
            
            # Generate classification based on score
            if phishing_score < 0.2:
                classification = "Legitimate (Very Low Risk)"
            elif phishing_score < 0.4:
                classification = "Likely Legitimate (Low Risk)"
            elif phishing_score < 0.6:
                classification = "Uncertain (Medium Risk)"
            elif phishing_score < 0.8:
                classification = "Likely Phishing (High Risk)"
            else:
                classification = "Phishing (Very High Risk)"
                
            results.append({
                "text": text[:50] + "..." if len(text) > 50 else text,
                "score": phishing_score,
                "classification": classification
            })
        
        processing_time = (time.time() - start_time) * 1000
        
        return {
            "results": results,
            "total_processed": len(results),
            "processing_time_ms": round(processing_time, 2),
            "timestamp": datetime.now().isoformat(),
            "note": "Lower scores (closer to 0) indicate legitimate content, higher scores (closer to 1) indicate phishing/malicious content"
        }
        
    except Exception as e:
        logger.error(f"Batch prediction error: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/metrics")
async def metrics():
    """Basic metrics endpoint for monitoring"""
    uptime = time.time() - startup_time
    
    return {
        "uptime_seconds": uptime,
        "model_loaded": model_loaded,
        "model_info": model_info,
        "memory_usage": "Not implemented",  # Could add psutil for real memory usage
        "timestamp": datetime.now().isoformat()
    }

if __name__ == "__main__":
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=8000,
        reload=False,
        log_level="info"
    )