Ragnarok04/Email-Classification

1from fastapi import FastAPI, HTTPException, Body
2from pydantic import BaseModel
3from typing import List, Dict, Any, Optional
4import uvicorn
5import os
6
7# Import our utility functions and models
8from utils import PIIMasker, preprocess_text
9from models import EmailClassifier, prepare_dummy_data
10
11
12# Define request/response models
13class EmailRequest(BaseModel):
14    email_body: str
15
16
17class MaskedEntity(BaseModel):
18    position: List[int]
19    classification: str
20    entity: str
21
22
23class EmailResponse(BaseModel):
24    input_email_body: str
25    list_of_masked_entities: List[MaskedEntity]
26    masked_email: str
27    category_of_the_email: str
28
29
30# Initialize FastAPI app
31app = FastAPI(title="Email Classification API",
32              description="API for classifying emails and masking PII",
33              version="1.0.0")
34
35# Initialize PII masker
36pii_masker = PIIMasker()
37
38# Load or train email classifier
39model_path = "email_classifier.pkl"
40if os.path.exists(model_path):
41    email_classifier = EmailClassifier.load_model(model_path)
42else:
43    # Train with dummy data
44    texts, labels = prepare_dummy_data()
45    email_classifier = EmailClassifier(model_type="traditional")
46    email_classifier.train(texts, labels)
47    email_classifier.save_model(model_path)
48
49
50@app.post("/classify-email", response_model=EmailResponse)
51async def classify_email(email_request: EmailRequest = Body(...)):
52    """
53    Classify an email and mask PII
54
55    Args:
56        email_request: Email body to process
57
58    Returns:
59        dict: Classification results and masked email
60    """
61    try:
62        # Get email text
63        email_text = email_request.email_body
64
65        # Mask PII
66        masked_email, entities = pii_masker.mask_pii(email_text)
67
68        # Preprocess text for classification
69        processed_text = preprocess_text(masked_email)
70
71        # Classify email
72        category = email_classifier.predict(processed_text)
73
74        # Format response
75        response = {
76            "input_email_body": email_text,
77            "list_of_masked_entities": entities,
78            "masked_email": masked_email,
79            "category_of_the_email": category
80        }
81
82        return response
83    except Exception as e:
84        raise HTTPException(status_code=500, detail=f"Error processing email: {str(e)}")
85
86
87@app.get("/")
88async def root():
89    """Root endpoint with API information"""
90    return {
91        "message": "Email Classification API is running",
92        "endpoints": {
93            "/classify-email": "POST endpoint to classify emails and mask PII"
94        }
95    }