Update api.py

api.py

@@ -1,108 +1,80 @@
-import gradio as gr
-import os
+from fastapi import FastAPI, HTTPException, Body
+from pydantic import BaseModel
+import uvicorn
 from models.pii_masker import PIIMasker
 from models.classifier import EmailClassifier
-from utils.utils import (preprocess_email, create_sample_dataset,
-                         parse_emails_dataset)
+from utils.utils import preprocess_email
 
-# Check if model exists, if not train it
-model_path = "models/email_classifier.joblib"
-if not os.path.exists(model_path):
-    print("Training model as it doesn't exist yet...")
-    data_path = "data/emails.csv"
+# Initialize FastAPI app
+app = FastAPI(title="Email Classification API",
+              description="API for classifying support emails and masking PII",
+              version="1.0.0")
 
-    # Create sample dataset if needed
-    if not os.path.exists(data_path):
-        print("Creating sample dataset...")
-        create_sample_dataset(data_path)
-
-    # Train model
-    df = parse_emails_dataset(data_path)
-    X = df['email'].tolist()
-    y = df['type'].tolist()
+# Initialize PII masker and classifier
+pii_masker = PIIMasker()
+classifier = EmailClassifier(model_path="models/email_classifier.joblib")
 
-    classifier = EmailClassifier()
-    classifier.train(X, y)
 
-    # Save trained model
-    os.makedirs(os.path.dirname(model_path), exist_ok=True)
-    classifier.save_model(model_path)
-    print("Model trained and saved successfully!")
+class EmailRequest(BaseModel):
+    email_body: str
 
 
-# Initialize components
-pii_masker = PIIMasker()
-classifier = EmailClassifier(model_path=model_path)
+class EmailResponse(BaseModel):
+    input_email_body: str
+    list_of_masked_entities: list
+    masked_email: str
+    category_of_the_email: str
 
 
-def process_email(email_body):
+@app.post("/classify-email", response_model=EmailResponse)
+async def classify_email(request: EmailRequest = Body(...)):
     """
-    Process email by masking PII and classifying it.
-        Args:
-        email_body: Raw email text
+    Classify an email and mask PII information.
+
+    Args:
+        request: Email request object containing the email body
+
     Returns:
-        tuple: (masked_email, entities_text, category)
+        dict: Response with masked email and classification
     """
-    # Mask PII
-    masked_email, entities = pii_masker.mask_pii(email_body)
-
-    # Preprocess for classification
-    processed_email = preprocess_email(masked_email)
-
-    # Classify email
-    category = classifier.classify(processed_email)
-
-    # Format entities for display
-    entities_text = "\n".join([
-        f"- {entity['classification']}: {entity['entity']}"
-        for entity in entities
-    ])
-
-    return masked_email, entities_text, category
-
-
-def test_masking():
-    """Example function to demonstrate PII masking"""
-    test_email = (
-        "Hello, my name is John Doe, and my email is johndoe@example.com.\n"
-        "My phone number is 555-123-4567 and I was born on 15/04/1985.\n"
-        "My Aadhar number is 1234 5678 9012 and my credit card number is "
-        "4111 1111 1111 1111 with CVV 123 expiring on 12/25."
-    )
-    return test_email
-
-
-# Create Gradio interface
-demo = gr.Interface(
-    fn=process_email,
-    inputs=gr.Textbox(
-        lines=10,
-        label="Email Content",
-        placeholder="Enter email text to classify and mask PII..."
-    ),
-    outputs=[
-        gr.Textbox(label="Masked Email"),
-        gr.Textbox(label="Detected PII Entities"),
-        gr.Textbox(label="Email Category")
-    ],
-    title="Email Classification System",
-    description=(
-        "This application classifies support emails and masks personally "
-        "identifiable information (PII)."
-    ),
-    examples=[
-        ["Hello, my name is John Doe, and my email is johndoe@example.com. "
-         "I need help with my account."],
-        ["I'm having trouble logging in to my account. My username is user"
-         "123."]
-    ],
-    article="""
-    ## How It Works
-   1. **PII Masking**: The system identifies and masks personal information
-    2. **Email Classification**: The masked email is classified into categories
-    3. **Results**: View the masked version, detected PII, and email category
+    try:
+        email_body = request.email_body
+
+        # Mask PII
+        masked_email, entities = pii_masker.mask_pii(email_body)
+
+        # Preprocess for classification
+        processed_email = preprocess_email(masked_email)
+
+        # Classify email
+        category = classifier.classify(processed_email)
+
+        # Prepare response
+        response = {
+            "input_email_body": email_body,
+            "list_of_masked_entities": entities,
+            "masked_email": masked_email,
+            "category_of_the_email": category
+        }
+
+        return response
+
+    except Exception as e:
+        error_msg = f"Error processing request: {str(e)}"
+        raise HTTPException(status_code=500, detail=error_msg)
+
+
+@app.get("/health")
+async def health_check():
     """
-)
+    Health check endpoint.
+
+    Returns:
+        dict: Status message
+    """
+    return {"status": "healthy"}
+
 
-# Launch the app
-demo.launch()
+# For local development
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)