Irannas/Masked_Email_Classification

1# pylint: disable=line-too-long
2"""This module defines an API endpoint for classifying email content."""
3
4from fastapi import APIRouter
5from pydantic import BaseModel
6
7from utils import mask_pii
8from models import predict_category
9
10router = APIRouter()
11
12
13class EmailInput(BaseModel):
14    # pylint: disable=too-few-public-methods
15    """Represents the input model for the classify_email endpoint.
16
17    input_email_body (str): The raw email body text provided by the user.
18    """
19
20    input_email_body: str
21
22
23@router.post("/classify")
24def classify_email(input_data: EmailInput):
25    """
26    Classifies the content of an email after masking any PII.
27
28    Args:
29        input_data (EmailInput): The input data containing the raw email body.
30
31    Returns:
32        dict: A dictionary containing:
33            - input_email_body (str): The original email body.
34            - list_of_masked_entities (list): The list of detected and masked PII entities.
35            - masked_email (str): The email body after masking PII.
36            - category_of_the_email (str): The predicted category of the email.
37    """
38    original_email = input_data.input_email_body
39    masked_email, entities = mask_pii(original_email)
40    category = predict_category(masked_email)
41
42    return {
43        "input_email_body": original_email,
44        "list_of_masked_entities": entities,
45        "masked_email": masked_email,
46        "category_of_the_email": category,
47    }
48