Update app.py

app.py

@@ -1,36 +1,27 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI
 from pydantic import BaseModel
 import joblib
 import re
 
-# Initialize the FastAPI app with default docs enabled
 app = FastAPI(
     title="Email Classification API",
     version="1.0.0",
-    description="Classifies emails and masks PII/PCI information.",
-    docs_url="/swagger",       # Try custom path
-    redoc_url="/redoc"         # Optional: enable ReDoc
+    description="Classifies support emails into categories and masks personal information.",
+    docs_url="/docs",       # Swagger UI enabled here
+    redoc_url="/redoc"      # Optional ReDoc UI
 )
 
-
-# Load the model once when the app starts
+# Load model
 model = joblib.load("model.joblib")
 
-# Define the root endpoint
-@app.get("/")
-def root():
-    return {"message": "Email Classification API is running."}
-
-# Define input data schema
+# Email input structure
 class EmailInput(BaseModel):
     subject: str = ""
     email: str
 
-# Function to mask and store PII
+# PII masker
 def mask_and_store_all_pii(text):
-    text = str(text)
     pii_map = {}
-
     patterns = {
         "email": r"\b[a-zA-Z0-9_.+-]+@[a-zA-Z0-9-]+\.[a-zA-Z0-9-.]+\b",
         "phone_number": r"\b\d{10}\b",
@@ -51,7 +42,7 @@ def mask_and_store_all_pii(text):
 
     return text, pii_map
 
-# Endpoint to classify email
+# Main endpoint
 @app.post("/classify")
 def classify_email(data: EmailInput):
     raw_text = f"{data.subject} {data.email}"
@@ -59,7 +50,11 @@ def classify_email(data: EmailInput):
     prediction = model.predict([masked_text])[0]
 
     return {
-        "predicted_category": prediction,
         "masked_text": masked_text,
+        "predicted_category": prediction,
         "pii_map": pii_map
     }
+
+@app.get("/")
+def root():
+    return {"message": "Email Classification API is running."}