server.py

from fastapi import FastAPI, Query
from pydantic import BaseModel
from typing import List
from utils import classify_text, fetch_latest_messages_and_classify, to_dataframe, load_model

load_model()

app = FastAPI(title="Email Threat Detector API")

class PredictRequest(BaseModel):
    text: str

@app.post("/predict_email")
def predict_email(req: PredictRequest):
    res = classify_text(req.text)
    return {
        "label": res["label"],
        "confidence": round(res["confidence"] * 100, 2),
        "probs": res["probs"],
        "labels": res["labels"]
    }

@app.get("/scan_gmail")
def scan_gmail(limit: int = Query(10, ge=1, le=50)):
    items = fetch_latest_messages_and_classify(limit=limit)
    return {"count": len(items), "emails": items}