Midhran/Email-classification

1from fastapi import FastAPI
2from pydantic import BaseModel
3from models import predict_category
4from pii_masker import PIIMasker
5
6app = FastAPI()
7masker = PIIMasker()
8
9class EmailRequest(BaseModel):
10    email_body: str
11
12@app.get("/healthcheck")
13def healthcheck():
14    return {"status": "ok"}
15
16@app.post("/classify-email")
17async def classify_email(request: EmailRequest):
18    original_email = request.email_body
19
20    masked_email, entities = masker.mask(original_email)
21    category, confidence = predict_category(original_email)
22    token_count = len(original_email.split())
23
24    return {
25        "input_email_body": original_email,
26        "list_of_masked_entities": entities,
27        "masked_email": masked_email,
28        "category_of_the_email": category,
29        "category_confidence": confidence,
30        "token_count": token_count,
31        "model_type": "logistic"
32    }
33