added db to store email for later demasking

.env

@@ -0,0 +1,8 @@
+# Database configuration - using SQLite instead of PostgreSQL
+DATABASE_PATH=emails.db
+
+# API configuration
+PORT=8000
+
+# Model configuration
+MODEL_PATH=Sparkonix/email-classifier-model

Dockerfile

@@ -14,13 +14,17 @@ COPY . .
 # Set environment variables
 ENV PORT=7860
 ENV MODEL_PATH="Sparkonix/email-classifier-model" 
-# Replace YOUR_ACTUAL_USERNAME with your Hugging Face username after uploading the model
+
+# Change this to point to SQLite database location
+ENV DATABASE_PATH="/data/emails.db"
 
 # Add this line to set cache location to a writable directory
 ENV HF_HOME="/app/.cache/huggingface"
 
 # Create the Hugging Face cache directory and set permissions
 RUN mkdir -p /app/.cache/huggingface && chmod -R 777 /app/.cache/huggingface
+# Create data directory for SQLite
+RUN mkdir -p /data && chmod -R 777 /data
 
 # Expose the port
 EXPOSE 7860

Email_Classification_API_Tests.postman_collection.json

@@ -0,0 +1,596 @@
+{
+    "info": {
+      "_postman_id": "7e5f2c6b-3e5d-4e40-a8f6-abc9f0c92a72",
+      "name": "Email Classification API Tests",
+      "description": "Tests for the Email Classification API hosted on Hugging Face Spaces",
+      "schema": "https://schema.getpostman.com/json/collection/v2.1.0/collection.json"
+    },
+    "variable": [
+      {
+        "key": "baseUrl",
+        "value": "https://sparkonix-email-classification-model.hf.space",
+        "type": "string"
+      },
+      {
+        "key": "emailId",
+        "value": ""
+      },
+      {
+        "key": "accessKey",
+        "value": ""
+      },
+      {
+        "key": "piiEmailId",
+        "value": ""
+      },
+      {
+        "key": "piiAccessKey",
+        "value": ""
+      }
+    ],
+    "item": [
+      {
+        "name": "Health Check",
+        "request": {
+          "method": "GET",
+          "header": [],
+          "url": {
+            "raw": "{{baseUrl}}/health",
+            "host": [
+              "{{baseUrl}}"
+            ],
+            "path": [
+              "health"
+            ]
+          },
+          "description": "Check if the API is running"
+        },
+        "event": [
+          {
+            "listen": "test",
+            "script": {
+              "exec": [
+                "// Check if the response status is 200 OK",
+                "pm.test(\"Status code is 200\", function() {",
+                "    pm.response.to.have.status(200);",
+                "});",
+                "",
+                "// Check if the health check returns the correct structure",
+                "pm.test(\"Health check returns correct structure\", function() {",
+                "    var jsonData = pm.response.json();",
+                "    pm.expect(jsonData).to.have.property(\"status\");",
+                "    pm.expect(jsonData).to.have.property(\"message\");",
+                "    pm.expect(jsonData.status).to.eql(\"healthy\");",
+                "});"
+              ],
+              "type": "text/javascript"
+            }
+          }
+        ],
+        "response": []
+      },
+      {
+        "name": "Basic Email Classification",
+        "request": {
+          "method": "POST",
+          "header": [],
+          "body": {
+            "mode": "raw",
+            "raw": "{\n    \"input_email_body\": \"Hi, I am experiencing a problem with my login. The system is showing an error when I try to log in. Please help.\"\n}",
+            "options": {
+              "raw": {
+                "language": "json"
+              }
+            }
+          },
+          "url": {
+            "raw": "{{baseUrl}}/classify",
+            "host": [
+              "{{baseUrl}}"
+            ],
+            "path": [
+              "classify"
+            ]
+          },
+          "description": "Basic email classification test"
+        },
+        "event": [
+          {
+            "listen": "test",
+            "script": {
+              "exec": [
+                "// Check if the response status is 200 OK",
+                "pm.test(\"Status code is 200\", function() {",
+                "    pm.response.to.have.status(200);",
+                "});",
+                "",
+                "// Check if the response structure is correct",
+                "pm.test(\"Response has correct structure\", function() {",
+                "    var jsonData = pm.response.json();",
+                "    pm.expect(jsonData).to.have.property(\"input_email_body\");",
+                "    pm.expect(jsonData).to.have.property(\"masked_email\");",
+                "    pm.expect(jsonData).to.have.property(\"list_of_masked_entities\");",
+                "    pm.expect(jsonData).to.have.property(\"category_of_the_email\");",
+                "});",
+                "",
+                "// Check if the input email matches what we sent",
+                "pm.test(\"Input email matches request\", function() {",
+                "    var jsonData = pm.response.json();",
+                "    pm.expect(jsonData.input_email_body).to.eql(",
+                "        \"Hi, I am experiencing a problem with my login. The system is showing an error when I try to log in. Please help.\"",
+                "    );",
+                "});",
+                "",
+                "// Check if the category is one of the expected values",
+                "pm.test(\"Category is valid\", function() {",
+                "    var jsonData = pm.response.json();",
+                "    pm.expect([\"Incident\", \"Request\", \"Change\", \"Problem\"]).to.include(jsonData.category_of_the_email);",
+                "});",
+                "",
+                "// Store the email_id and access_key for database retrieval tests",
+                "var responseJson = pm.response.json();",
+                "if (responseJson.email_id && responseJson.access_key) {",
+                "    pm.collectionVariables.set(\"emailId\", responseJson.email_id);",
+                "    pm.collectionVariables.set(\"accessKey\", responseJson.access_key);",
+                "    console.log(\"Stored email_id and access_key for retrieval tests\");",
+                "}"
+              ],
+              "type": "text/javascript"
+            }
+          }
+        ],
+        "response": []
+      },
+      {
+        "name": "Email with PII Classification",
+        "request": {
+          "method": "POST",
+          "header": [],
+          "body": {
+            "mode": "raw",
+            "raw": "{\n    \"input_email_body\": \"Hello, my name is John Smith and my email is john.smith@example.com. I am having issues with my account. My phone number is 555-123-4567 and my credit card ending in 1234 seems to have been charged twice. Please help.\"\n}",
+            "options": {
+              "raw": {
+                "language": "json"
+              }
+            }
+          },
+          "url": {
+            "raw": "{{baseUrl}}/classify",
+            "host": [
+              "{{baseUrl}}"
+            ],
+            "path": [
+              "classify"
+            ]
+          },
+          "description": "Email classification with PII masking test"
+        },
+        "event": [
+          {
+            "listen": "test",
+            "script": {
+              "exec": [
+                "// Check if the response status is 200 OK",
+                "pm.test(\"Status code is 200\", function() {",
+                "    pm.response.to.have.status(200);",
+                "});",
+                "",
+                "// Check if PII is being masked",
+                "pm.test(\"PII is masked properly\", function() {",
+                "    var jsonData = pm.response.json();",
+                "    pm.expect(jsonData.masked_email).to.not.include(\"John Smith\");",
+                "    pm.expect(jsonData.masked_email).to.not.include(\"john.smith@example.com\");",
+                "    pm.expect(jsonData.masked_email).to.not.include(\"555-123-4567\");",
+                "    pm.expect(jsonData.masked_email).to.not.include(\"1234\");",
+                "});",
+                "",
+                "// Check if masked entities list is populated",
+                "pm.test(\"Masked entities list is populated\", function() {",
+                "    var jsonData = pm.response.json();",
+                "    pm.expect(jsonData.list_of_masked_entities).to.be.an(\"array\").that.is.not.empty;",
+                "});",
+                "",
+                "// Store the email_id and access_key for database retrieval tests specifically for PII emails",
+                "var responseJson = pm.response.json();",
+                "if (responseJson.email_id && responseJson.access_key) {",
+                "    pm.collectionVariables.set(\"piiEmailId\", responseJson.email_id);",
+                "    pm.collectionVariables.set(\"piiAccessKey\", responseJson.access_key);",
+                "    console.log(\"Stored PII email_id and access_key for retrieval tests\");",
+                "}"
+              ],
+              "type": "text/javascript"
+            }
+          }
+        ],
+        "response": []
+      },
+      {
+        "name": "Incident Email Classification",
+        "request": {
+          "method": "POST",
+          "header": [],
+          "body": {
+            "mode": "raw",
+            "raw": "{\n    \"input_email_body\": \"URGENT: The application is down. None of our users can access the system. This is causing severe business impact and we need immediate attention. Our customer operations are halted.\"\n}",
+            "options": {
+              "raw": {
+                "language": "json"
+              }
+            }
+          },
+          "url": {
+            "raw": "{{baseUrl}}/classify",
+            "host": [
+              "{{baseUrl}}"
+            ],
+            "path": [
+              "classify"
+            ]
+          },
+          "description": "Test for incident classification"
+        },
+        "event": [
+          {
+            "listen": "test",
+            "script": {
+              "exec": [
+                "// Check if the response status is 200 OK",
+                "pm.test(\"Status code is 200\", function() {",
+                "    pm.response.to.have.status(200);",
+                "});",
+                "",
+                "// This should likely be classified as an Incident",
+                "pm.test(\"Should be classified as an Incident\", function() {",
+                "    var jsonData = pm.response.json();",
+                "    pm.expect(jsonData.category_of_the_email).to.eql(\"Incident\");",
+                "});"
+              ],
+              "type": "text/javascript"
+            }
+          }
+        ],
+        "response": []
+      },
+      {
+        "name": "Request Email Classification",
+        "request": {
+          "method": "POST",
+          "header": [],
+          "body": {
+            "mode": "raw",
+            "raw": "{\n    \"input_email_body\": \"I would like to request access to the financial reporting system for my new team member. Their details are as follows: Name: Jane Doe, Department: Finance, Employee ID: 12345. Please provide access by end of week.\"\n}",
+            "options": {
+              "raw": {
+                "language": "json"
+              }
+            }
+          },
+          "url": {
+            "raw": "{{baseUrl}}/classify",
+            "host": [
+              "{{baseUrl}}"
+            ],
+            "path": [
+              "classify"
+            ]
+          },
+          "description": "Test for request classification"
+        },
+        "event": [
+          {
+            "listen": "test",
+            "script": {
+              "exec": [
+                "// Check if the response status is 200 OK",
+                "pm.test(\"Status code is 200\", function() {",
+                "    pm.response.to.have.status(200);",
+                "});",
+                "",
+                "// This should likely be classified as a Request",
+                "pm.test(\"Should be classified as a Request\", function() {",
+                "    var jsonData = pm.response.json();",
+                "    pm.expect(jsonData.category_of_the_email).to.eql(\"Request\");",
+                "});",
+                "",
+                "// Check if PII is masked",
+                "pm.test(\"Name should be masked\", function() {",
+                "    var jsonData = pm.response.json();",
+                "    pm.expect(jsonData.masked_email).to.not.include(\"Jane Doe\");",
+                "});"
+              ],
+              "type": "text/javascript"
+            }
+          }
+        ],
+        "response": []
+      },
+      {
+        "name": "Change Email Classification",
+        "request": {
+          "method": "POST",
+          "header": [],
+          "body": {
+            "mode": "raw",
+            "raw": "{\n    \"input_email_body\": \"We need to change the configuration of the production server to increase memory allocation. Please schedule this change for the next maintenance window this Sunday at 2 AM. Approval has been granted by the IT Director.\"\n}",
+            "options": {
+              "raw": {
+                "language": "json"
+              }
+            }
+          },
+          "url": {
+            "raw": "{{baseUrl}}/classify",
+            "host": [
+              "{{baseUrl}}"
+            ],
+            "path": [
+              "classify"
+            ]
+          },
+          "description": "Test for change classification"
+        },
+        "event": [
+          {
+            "listen": "test",
+            "script": {
+              "exec": [
+                "// Check if the response status is 200 OK",
+                "pm.test(\"Status code is 200\", function() {",
+                "    pm.response.to.have.status(200);",
+                "});",
+                "",
+                "// This should likely be classified as a Change",
+                "pm.test(\"Should be classified as a Change\", function() {",
+                "    var jsonData = pm.response.json();",
+                "    pm.expect(jsonData.category_of_the_email).to.eql(\"Change\");",
+                "});"
+              ],
+              "type": "text/javascript"
+            }
+          }
+        ],
+        "response": []
+      },
+      {
+        "name": "Problem Email Classification",
+        "request": {
+          "method": "POST",
+          "header": [],
+          "body": {
+            "mode": "raw",
+            "raw": "{\n    \"input_email_body\": \"We have noticed that the application has been running slow for the past week. This happens consistently during peak hours (10 AM - 2 PM). Can you investigate the root cause of this ongoing performance issue?\"\n}",
+            "options": {
+              "raw": {
+                "language": "json"
+              }
+            }
+          },
+          "url": {
+            "raw": "{{baseUrl}}/classify",
+            "host": [
+              "{{baseUrl}}"
+            ],
+            "path": [
+              "classify"
+            ]
+          },
+          "description": "Test for problem classification"
+        },
+        "event": [
+          {
+            "listen": "test",
+            "script": {
+              "exec": [
+                "// Check if the response status is 200 OK",
+                "pm.test(\"Status code is 200\", function() {",
+                "    pm.response.to.have.status(200);",
+                "});",
+                "",
+                "// This should likely be classified as a Problem",
+                "pm.test(\"Should be classified as a Problem\", function() {",
+                "    var jsonData = pm.response.json();",
+                "    pm.expect(jsonData.category_of_the_email).to.eql(\"Problem\");",
+                "});"
+              ],
+              "type": "text/javascript"
+            }
+          }
+        ],
+        "response": []
+      },
+      {
+        "name": "Empty Email Test",
+        "request": {
+          "method": "POST",
+          "header": [],
+          "body": {
+            "mode": "raw",
+            "raw": "{\n    \"input_email_body\": \"\"\n}",
+            "options": {
+              "raw": {
+                "language": "json"
+              }
+            }
+          },
+          "url": {
+            "raw": "{{baseUrl}}/classify",
+            "host": [
+              "{{baseUrl}}"
+            ],
+            "path": [
+              "classify"
+            ]
+          },
+          "description": "Test with empty email body"
+        },
+        "event": [
+          {
+            "listen": "test",
+            "script": {
+              "exec": [
+                "// The API should either return an error or a classification",
+                "pm.test(\"Response status is valid\", function() {",
+                "    pm.expect(pm.response.code).to.be.oneOf([200, 400, 422, 500]);",
+                "    ",
+                "    if (pm.response.code === 200) {",
+                "        var jsonData = pm.response.json();",
+                "        pm.expect(jsonData).to.have.property(\"category_of_the_email\");",
+                "    } else {",
+                "        // If it is an error response, make sure it has a proper structure",
+                "        var jsonData = pm.response.json();",
+                "        pm.expect(jsonData).to.have.property(\"detail\");",
+                "    }",
+                "});"
+              ],
+              "type": "text/javascript"
+            }
+          }
+        ],
+        "response": []
+      },
+      {
+        "name": "Retrieve Original Email",
+        "request": {
+          "method": "POST",
+          "header": [],
+          "body": {
+            "mode": "raw",
+            "raw": "{\n    \"email_id\": \"{{piiEmailId}}\",\n    \"access_key\": \"{{piiAccessKey}}\"\n}",
+            "options": {
+              "raw": {
+                "language": "json"
+              }
+            }
+          },
+          "url": {
+            "raw": "{{baseUrl}}/api/v1/original-email/retrieve",
+            "host": [
+              "{{baseUrl}}"
+            ],
+            "path": [
+              "api",
+              "v1",
+              "original-email",
+              "retrieve"
+            ]
+          },
+          "description": "Retrieve the original unmasked email from the database"
+        },
+        "event": [
+          {
+            "listen": "test",
+            "script": {
+              "exec": [
+                "// Check if the response status is 200 OK",
+                "pm.test(\"Status code is 200\", function() {",
+                "    pm.response.to.have.status(200);",
+                "});",
+                "",
+                "// Check if the response has the correct structure",
+                "pm.test(\"Response has correct structure\", function() {",
+                "    var jsonData = pm.response.json();",
+                "    pm.expect(jsonData).to.have.property(\"status\");",
+                "    pm.expect(jsonData).to.have.property(\"data\");",
+                "    pm.expect(jsonData).to.have.property(\"message\");",
+                "    pm.expect(jsonData.status).to.eql(\"success\");",
+                "});",
+                "",
+                "// Check if the data contains the original email with PII",
+                "pm.test(\"Data contains original email with PII\", function() {",
+                "    var jsonData = pm.response.json();",
+                "    pm.expect(jsonData.data).to.have.property(\"original_email\");",
+                "    pm.expect(jsonData.data).to.have.property(\"masked_email\");",
+                "    pm.expect(jsonData.data).to.have.property(\"masked_entities\");",
+                "    ",
+                "    // Check that the original email contains the PII (for the PII test email)",
+                "    if (pm.collectionVariables.get(\"piiEmailId\")) {",
+                "        pm.expect(jsonData.data.original_email).to.include(\"John Smith\");",
+                "        pm.expect(jsonData.data.original_email).to.include(\"john.smith@example.com\");",
+                "        pm.expect(jsonData.data.original_email).to.include(\"555-123-4567\");",
+                "    }",
+                "});"
+              ],
+              "type": "text/javascript"
+            }
+          }
+        ],
+        "response": []
+      },
+      {
+        "name": "Retrieve With Invalid Access Key",
+        "request": {
+          "method": "POST",
+          "header": [],
+          "body": {
+            "mode": "raw",
+            "raw": "{\n    \"email_id\": \"{{piiEmailId}}\",\n    \"access_key\": \"invalid_access_key_123456\"\n}",
+            "options": {
+              "raw": {
+                "language": "json"
+              }
+            }
+          },
+          "url": {
+            "raw": "{{baseUrl}}/api/v1/original-email/retrieve",
+            "host": [
+              "{{baseUrl}}"
+            ],
+            "path": [
+              "api",
+              "v1",
+              "original-email",
+              "retrieve"
+            ]
+          },
+          "description": "Test security by attempting to retrieve email with invalid access key"
+        },
+        "event": [
+          {
+            "listen": "test",
+            "script": {
+              "exec": [
+                "// Check that we get an error (404) for invalid access key",
+                "pm.test(\"Should return error for invalid access key\", function() {",
+                "    pm.expect(pm.response.code).to.equal(404);",
+                "});",
+                "",
+                "// Check error message",
+                "pm.test(\"Response contains appropriate error message\", function() {",
+                "    var jsonData = pm.response.json();",
+                "    pm.expect(jsonData).to.have.property(\"detail\");",
+                "    pm.expect(jsonData.detail).to.include(\"not found\");",
+                "});"
+              ],
+              "type": "text/javascript"
+            }
+          }
+        ],
+        "response": []
+      }
+    ],
+    "event": [
+      {
+        "listen": "prerequest",
+        "script": {
+          "type": "text/javascript",
+          "exec": [
+            "// Reset the tests passed flag before each request",
+            "pm.variables.set(\"testsPassed\", true);"
+          ]
+        }
+      },
+      {
+        "listen": "test",
+        "script": {
+          "type": "text/javascript",
+          "exec": [
+            "// Update the tests passed flag if any tests failed",
+            "if (pm.test.allTests.filter(test => !test.passed).length > 0) {",
+            "    pm.variables.set(\"testsPassed\", false);",
+            "}"
+          ]
+        }
+      }
+    ]
+  }

database.py

@@ -0,0 +1,207 @@
+"""
+Database module for handling email storage operations with SQLite.
+"""
+import os
+import json
+import sqlite3
+from typing import Dict, Any, Optional, List, Tuple
+from datetime import datetime
+import uuid
+import hashlib
+
+
+class EmailDatabase:
+    """
+    Database class for storing and retrieving email data with PII masking information.
+    Uses SQLite for storage in Hugging Face's persistent directory.
+    """
+    
+    def __init__(self, connection_string: str = None):
+        """
+        Initialize the database connection.
+        
+        Args:
+            connection_string: Database connection string or path. 
+                              For SQLite, this will be treated as a file path.
+        """
+        # Hugging Face Spaces has a /data directory that persists between restarts
+        self.db_path = connection_string or os.environ.get(
+            "DATABASE_PATH", 
+            "/data/emails.db"  # This path persists in Hugging Face Spaces
+        )
+        
+        # Ensure the data directory exists
+        self._ensure_data_directory()
+        
+        self._create_tables()
+    
+    def _ensure_data_directory(self):
+        """Ensure the data directory exists, and use a fallback if needed."""
+        try:
+            data_dir = os.path.dirname(self.db_path)
+            if data_dir and not os.path.exists(data_dir):
+                os.makedirs(data_dir, exist_ok=True)
+        except (OSError, PermissionError):
+            # If we can't write to /data, fall back to the current directory
+            self.db_path = "emails.db"
+            print(f"Warning: Using fallback database path: {self.db_path}")
+    
+    def _get_connection(self):
+        """Get a database connection."""
+        return sqlite3.connect(self.db_path)
+    
+    def _create_tables(self):
+        """Create the necessary tables if they don't exist."""
+        conn = self._get_connection()
+        try:
+            cursor = conn.cursor()
+            
+            # Create the emails table to store original emails and their masked versions
+            cursor.execute('''
+            CREATE TABLE IF NOT EXISTS emails (
+                id TEXT PRIMARY KEY,
+                original_email TEXT NOT NULL,
+                masked_email TEXT NOT NULL,
+                masked_entities TEXT NOT NULL, 
+                category TEXT,
+                created_at TEXT NOT NULL,
+                access_key TEXT NOT NULL
+            )
+            ''')
+            
+            # Create an index on the access_key field
+            cursor.execute('CREATE INDEX IF NOT EXISTS idx_access_key ON emails (access_key)')
+            
+            conn.commit()
+        except Exception as e:
+            conn.rollback()
+            raise e
+        finally:
+            conn.close()
+    
+    def _generate_id(self) -> str:
+        """Generate a unique ID for the email record."""
+        return str(uuid.uuid4())
+    
+    def _generate_access_key(self, email_id: str) -> str:
+        """
+        Generate an access key for retrieving the original email.
+        This acts as a security measure to prevent unauthorized access.
+        """
+        # Use a combination of the email ID and a timestamp, hashed
+        data = f"{email_id}:{datetime.now().isoformat()}:{os.urandom(8).hex()}"
+        return hashlib.sha256(data.encode()).hexdigest()
+    
+    def store_email(self, original_email: str, masked_email: str, 
+                   masked_entities: List[Dict[str, Any]], category: Optional[str] = None) -> Tuple[str, str]:
+        """
+        Store the original email along with its masked version and related information.
+        
+        Args:
+            original_email: The original email with PII
+            masked_email: The masked version of the email
+            masked_entities: List of entities that were masked
+            category: Optional category of the email
+            
+        Returns:
+            Tuple of (email_id, access_key) for future reference
+        """
+        conn = self._get_connection()
+        try:
+            cursor = conn.cursor()
+            
+            email_id = self._generate_id()
+            access_key = self._generate_access_key(email_id)
+            
+            # Store the email data
+            cursor.execute(
+                'INSERT INTO emails (id, original_email, masked_email, masked_entities, category, created_at, access_key) '
+                'VALUES (?, ?, ?, ?, ?, ?, ?)',
+                (
+                    email_id,
+                    original_email,
+                    masked_email,
+                    json.dumps(masked_entities),  # Convert to JSON string for SQLite
+                    category,
+                    datetime.now().isoformat(),
+                    access_key
+                )
+            )
+            
+            conn.commit()
+            return email_id, access_key
+        except Exception as e:
+            conn.rollback()
+            raise e
+        finally:
+            conn.close()
+    
+    def get_original_email(self, email_id: str, access_key: str) -> Optional[Dict[str, Any]]:
+        """
+        Retrieve the original email with PII using the access key.
+        
+        Args:
+            email_id: The ID of the email record
+            access_key: The security key required to access the original email
+            
+        Returns:
+            Dictionary with email data or None if not found or access_key is invalid
+        """
+        conn = self._get_connection()
+        try:
+            cursor = conn.cursor()
+            
+            cursor.execute(
+                'SELECT id, original_email, masked_email, masked_entities, category, created_at '
+                'FROM emails WHERE id = ? AND access_key = ?',
+                (email_id, access_key)
+            )
+            
+            row = cursor.fetchone()
+            if not row:
+                return None
+                
+            return {
+                "id": row[0],
+                "original_email": row[1],
+                "masked_email": row[2],
+                "masked_entities": json.loads(row[3]),  # Convert from JSON string back to Python dict
+                "category": row[4],
+                "created_at": row[5]
+            }
+        finally:
+            conn.close()
+    
+    def get_email_by_id(self, email_id: str) -> Optional[Dict[str, Any]]:
+        """
+        Retrieve the masked email data without the original PII-containing email.
+        
+        Args:
+            email_id: The ID of the email
+            
+        Returns:
+            Dictionary with masked email data or None if not found
+        """
+        conn = self._get_connection()
+        try:
+            cursor = conn.cursor()
+            
+            cursor.execute(
+                'SELECT id, masked_email, masked_entities, category, created_at '
+                'FROM emails WHERE id = ?',
+                (email_id,)
+            )
+            
+            row = cursor.fetchone()
+            if not row:
+                return None
+                
+            return {
+                "id": row[0],
+                "masked_email": row[1],
+                "masked_entities": json.loads(row[2]),  # Convert from JSON string back to Python dict
+                "category": row[3],
+                "created_at": row[4]
+            }
+        finally:
+            conn.close()

docker-compose.yml

@@ -7,7 +7,12 @@ services:
       - "8000:7860"
     volumes:
       - .:/app
+      - sqlite_data:/data  # Mount a volume for persistent SQLite data
     environment:
       - PORT=7860
       - HF_HOME=/app/.cache/huggingface
-    restart: unless-stopped
+      - DATABASE_PATH=/data/emails.db
+    restart: unless-stopped
+
+volumes:
+  sqlite_data:  # Define the volume for SQLite database

main.py

@@ -3,17 +3,30 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from typing import Dict, Any, List, Tuple, Optional
 import uvicorn
+from dotenv import load_dotenv
 
 from utils import PIIMasker
 from models import EmailClassifier
 
+# Load environment variables from .env file if available
+try:
+    load_dotenv()
+except ImportError:
+    pass  # dotenv might not be installed in production
+
+# Set database path for Hugging Face, using persistent storage
+if os.path.exists('/data'):
+    db_path = "/data/emails.db"
+else:
+    db_path = "emails.db"  # Fallback to local directory
+
 # Initialize the FastAPI application
 app = FastAPI(title="Email Classification API", 
               description="API for classifying support emails and masking PII",
               version="1.0.0")
 
 # Initialize the PII masker and email classifier
-pii_masker = PIIMasker()
+pii_masker = PIIMasker(db_path=db_path)
 email_classifier = EmailClassifier()
 
 class EmailInput(BaseModel):
@@ -33,6 +46,11 @@ class EmailOutput(BaseModel):
     masked_email: str
     category_of_the_email: str
 
+class EmailRetrievalInput(BaseModel):
+    """Input model for retrieving original email"""
+    email_id: str
+    access_key: str
+
 @app.post("/classify", response_model=EmailOutput)
 async def classify_email(email_input: EmailInput) -> Dict[str, Any]:
     """
@@ -51,10 +69,46 @@ async def classify_email(email_input: EmailInput) -> Dict[str, Any]:
         # Classify the masked email
         classified_data = email_classifier.process_email(processed_data)
         
-        return classified_data
+        # Make sure we return only the fields expected in the response model
+        return {
+            "input_email_body": email_input.input_email_body,
+            "list_of_masked_entities": classified_data["list_of_masked_entities"],
+            "masked_email": classified_data["masked_email"],
+            "category_of_the_email": classified_data["category_of_the_email"]
+        }
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error processing email: {str(e)}")
 
+@app.post("/api/v1/original-email/retrieve", response_model=Dict[str, Any])
+async def retrieve_original_email_v1(retrieval_input: EmailRetrievalInput) -> Dict[str, Any]:
+    """
+    New API endpoint to retrieve the original unmasked email from SQLite database.
+    
+    Args:
+        retrieval_input: The email ID and access key
+        
+    Returns:
+        The original email data with PII information
+    """
+    try:
+        email_data = pii_masker.get_original_email(
+            retrieval_input.email_id, 
+            retrieval_input.access_key
+        )
+        
+        if not email_data:
+            raise HTTPException(status_code=404, detail="Email not found or invalid access key")
+        
+        return {
+            "status": "success",
+            "data": email_data,
+            "message": "Original email retrieved successfully"
+        }
+    except Exception as e:
+        if isinstance(e, HTTPException):
+            raise e
+        raise HTTPException(status_code=500, detail=f"Error retrieving email: {str(e)}")
+
 @app.get("/health")
 async def health_check():
     """

requirements.in

@@ -1,52 +1,74 @@
 # This file was autogenerated by uv via the following command:
 #    uv pip compile requirements.in -o requirements.txt
 annotated-types==0.7.0
-    # via pydantic
+    # via
+    #   -r requirements.in
+    #   pydantic
 anyio==4.9.0
-    # via starlette
+    # via
+    #   -r requirements.in
+    #   starlette
 blis==1.3.0
-    # via thinc
+    # via
+    #   -r requirements.in
+    #   thinc
 catalogue==2.0.10
     # via
+    #   -r requirements.in
     #   spacy
     #   srsly
     #   thinc
 certifi==2025.4.26
-    # via requests
+    # via
+    #   -r requirements.in
+    #   requests
 charset-normalizer==3.4.2
-    # via requests
+    # via
+    #   -r requirements.in
+    #   requests
 click==8.2.0
     # via
+    #   -r requirements.in
     #   typer
     #   uvicorn
 cloudpathlib==0.21.1
-    # via weasel
+    # via
+    #   -r requirements.in
+    #   weasel
 confection==0.1.5
     # via
+    #   -r requirements.in
     #   thinc
     #   weasel
 cymem==2.0.11
     # via
+    #   -r requirements.in
     #   preshed
     #   spacy
     #   thinc
 en-core-web-sm @ https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl
     # via -r requirements.in
 exceptiongroup==1.3.0
-    # via anyio
+    # via
+    #   -r requirements.in
+    #   anyio
 fastapi==0.115.12
     # via -r requirements.in
 filelock==3.18.0
     # via
+    #   -r requirements.in
     #   huggingface-hub
     #   torch
     #   transformers
 fsspec==2025.3.2
     # via
+    #   -r requirements.in
     #   huggingface-hub
     #   torch
 h11==0.16.0
-    # via uvicorn
+    # via
+    #   -r requirements.in
+    #   uvicorn
 huggingface-hub==0.31.2
     # via
     #   -r requirements.in
@@ -54,79 +76,130 @@ huggingface-hub==0.31.2
     #   transformers
 idna==3.10
     # via
+    #   -r requirements.in
     #   anyio
     #   requests
 jinja2==3.1.6
     # via
+    #   -r requirements.in
     #   spacy
     #   torch
 langcodes==3.5.0
-    # via spacy
+    # via
+    #   -r requirements.in
+    #   spacy
 language-data==1.3.0
-    # via langcodes
+    # via
+    #   -r requirements.in
+    #   langcodes
 marisa-trie==1.2.1
-    # via language-data
+    # via
+    #   -r requirements.in
+    #   language-data
 markdown-it-py==3.0.0
-    # via rich
+    # via
+    #   -r requirements.in
+    #   rich
 markupsafe==3.0.2
-    # via jinja2
+    # via
+    #   -r requirements.in
+    #   jinja2
 mdurl==0.1.2
-    # via markdown-it-py
+    # via
+    #   -r requirements.in
+    #   markdown-it-py
 mpmath==1.3.0
-    # via sympy
+    # via
+    #   -r requirements.in
+    #   sympy
 murmurhash==1.0.12
     # via
+    #   -r requirements.in
     #   preshed
     #   spacy
     #   thinc
 networkx==3.4.2
-    # via torch
+    # via
+    #   -r requirements.in
+    #   torch
 numpy==2.2.5
     # via
+    #   -r requirements.in
     #   blis
     #   spacy
     #   spacy-transformers
     #   thinc
     #   transformers
+
+# SQLite is included in Python standard library
+python-dotenv
+    # for environment variable management
+
 nvidia-cublas-cu12==12.6.4.1
     # via
+    #   -r requirements.in
     #   nvidia-cudnn-cu12
     #   nvidia-cusolver-cu12
     #   torch
 nvidia-cuda-cupti-cu12==12.6.80
-    # via torch
+    # via
+    #   -r requirements.in
+    #   torch
 nvidia-cuda-nvrtc-cu12==12.6.77
-    # via torch
+    # via
+    #   -r requirements.in
+    #   torch
 nvidia-cuda-runtime-cu12==12.6.77
-    # via torch
+    # via
+    #   -r requirements.in
+    #   torch
 nvidia-cudnn-cu12==9.5.1.17
-    # via torch
+    # via
+    #   -r requirements.in
+    #   torch
 nvidia-cufft-cu12==11.3.0.4
-    # via torch
+    # via
+    #   -r requirements.in
+    #   torch
 nvidia-cufile-cu12==1.11.1.6
-    # via torch
+    # via
+    #   -r requirements.in
+    #   torch
 nvidia-curand-cu12==10.3.7.77
-    # via torch
+    # via
+    #   -r requirements.in
+    #   torch
 nvidia-cusolver-cu12==11.7.1.2
-    # via torch
+    # via
+    #   -r requirements.in
+    #   torch
 nvidia-cusparse-cu12==12.5.4.2
     # via
+    #   -r requirements.in
     #   nvidia-cusolver-cu12
     #   torch
 nvidia-cusparselt-cu12==0.6.3
-    # via torch
+    # via
+    #   -r requirements.in
+    #   torch
 nvidia-nccl-cu12==2.26.2
-    # via torch
+    # via
+    #   -r requirements.in
+    #   torch
 nvidia-nvjitlink-cu12==12.6.85
     # via
+    #   -r requirements.in
     #   nvidia-cufft-cu12
     #   nvidia-cusolver-cu12
     #   nvidia-cusparse-cu12
     #   torch
 nvidia-nvtx-cu12==12.6.77
-    # via torch
+    # via
+    #   -r requirements.in
+    #   torch
 packaging==25.0
     # via
+    #   -r requirements.in
     #   huggingface-hub
     #   spacy
     #   thinc
@@ -134,6 +207,7 @@ packaging==25.0
     #   weasel
 preshed==3.0.9
     # via
+    #   -r requirements.in
     #   spacy
     #   thinc
 pydantic==2.11.4
@@ -145,74 +219,109 @@ pydantic==2.11.4
     #   thinc
     #   weasel
 pydantic-core==2.33.2
-    # via pydantic
+    # via
+    #   -r requirements.in
+    #   pydantic
 pygments==2.19.1
-    # via rich
+    # via
+    #   -r requirements.in
+    #   rich
 python-multipart==0.0.20
     # via -r requirements.in
 pyyaml==6.0.2
     # via
+    #   -r requirements.in
     #   huggingface-hub
     #   transformers
 regex==2024.11.6
-    # via transformers
+    # via
+    #   -r requirements.in
+    #   transformers
 requests==2.32.3
     # via
+    #   -r requirements.in
     #   huggingface-hub
     #   spacy
     #   transformers
     #   weasel
 rich==14.0.0
-    # via typer
+    # via
+    #   -r requirements.in
+    #   typer
 safetensors==0.5.3
     # via
     #   -r requirements.in
     #   transformers
+sentencepiece==0.2.0
+    # via -r requirements.in
 setuptools==80.7.1
     # via
+    #   -r requirements.in
     #   marisa-trie
     #   spacy
     #   thinc
     #   triton
 shellingham==1.5.4
-    # via typer
+    # via
+    #   -r requirements.in
+    #   typer
 smart-open==7.1.0
-    # via weasel
+    # via
+    #   -r requirements.in
+    #   weasel
 sniffio==1.3.1
-    # via anyio
+    # via
+    #   -r requirements.in
+    #   anyio
 spacy==3.8.5
     # via
     #   -r requirements.in
     #   spacy-transformers
 spacy-alignments==0.9.1
-    # via spacy-transformers
+    # via
+    #   -r requirements.in
+    #   spacy-transformers
 spacy-legacy==3.0.12
-    # via spacy
+    # via
+    #   -r requirements.in
+    #   spacy
 spacy-loggers==1.0.5
-    # via spacy
+    # via
+    #   -r requirements.in
+    #   spacy
 spacy-transformers==1.3.8
     # via -r requirements.in
 srsly==2.5.1
     # via
+    #   -r requirements.in
     #   confection
     #   spacy
     #   spacy-transformers
     #   thinc
     #   weasel
 starlette==0.46.2
-    # via fastapi
+    # via
+    #   -r requirements.in
+    #   fastapi
 sympy==1.14.0
-    # via torch
+    # via
+    #   -r requirements.in
+    #   torch
 thinc==8.3.6
-    # via spacy
+    # via
+    #   -r requirements.in
+    #   spacy
 tokenizers==0.21.1
-    # via transformers
+    # via
+    #   -r requirements.in
+    #   transformers
 torch==2.7.0
     # via
     #   -r requirements.in
     #   spacy-transformers
 tqdm==4.67.1
     # via
+    #   -r requirements.in
     #   huggingface-hub
     #   spacy
     #   transformers
@@ -221,13 +330,17 @@ transformers==4.49.0
     #   -r requirements.in
     #   spacy-transformers
 triton==3.3.0
-    # via torch
+    # via
+    #   -r requirements.in
+    #   torch
 typer==0.15.3
     # via
+    #   -r requirements.in
     #   spacy
     #   weasel
 typing-extensions==4.13.2
     # via
+    #   -r requirements.in
     #   anyio
     #   cloudpathlib
     #   exceptiongroup
@@ -241,20 +354,28 @@ typing-extensions==4.13.2
     #   typing-inspection
     #   uvicorn
 typing-inspection==0.4.0
-    # via pydantic
+    # via
+    #   -r requirements.in
+    #   pydantic
 urllib3==2.4.0
-    # via requests
+    # via
+    #   -r requirements.in
+    #   requests
 uvicorn==0.34.2
     # via -r requirements.in
 wasabi==1.1.3
     # via
+    #   -r requirements.in
     #   spacy
     #   thinc
     #   weasel
 weasel==0.4.1
-    # via spacy
+    # via
+    #   -r requirements.in
+    #   spacy
 wrapt==1.17.2
-    # via smart-open
+    # via
+    #   -r requirements.in
+    #   smart-open
 xx-ent-wiki-sm @ https://github.com/explosion/spacy-models/releases/download/xx_ent_wiki_sm-3.8.0/xx_ent_wiki_sm-3.8.0-py3-none-any.whl
     # via -r requirements.in
-sentencepiece 

requirements.txt

@@ -221,6 +221,8 @@ pygments==2.19.1
     # via
     #   -r requirements.in
     #   rich
+python-dotenv==1.1.0
+    # via -r requirements.in
 python-multipart==0.0.20
     # via -r requirements.in
 pyyaml==6.0.2

utils.py

@@ -1,6 +1,7 @@
 import re
 import spacy
-from typing import List, Dict, Tuple, Any
+from typing import List, Dict, Tuple, Any, Optional
+from database import EmailDatabase
 
 class Entity:
     def __init__(self, start: int, end: int, entity_type: str, value: str):
@@ -20,7 +21,7 @@ class Entity:
         return f"Entity(type='{self.entity_type}', value='{self.value}', start={self.start}, end={self.end})"
 
 class PIIMasker:
-    def __init__(self, spacy_model_name: str = "xx_ent_wiki_sm"): # Allow model choice
+    def __init__(self, spacy_model_name: str = "xx_ent_wiki_sm", db_path: str = None): # Allow model choice
         # Load SpaCy model
         try:
             self.nlp = spacy.load(spacy_model_name)
@@ -39,7 +40,9 @@ class PIIMasker:
                     spacy.cli.download("en_core_web_sm")
                     self.nlp = spacy.load("en_core_web_sm")
 
-
+        # Initialize database connection with SQLite path
+        self.db = EmailDatabase(connection_string=db_path)
+        
         # Initialize regex patterns
         self._initialize_patterns()
 
@@ -320,12 +323,50 @@ class PIIMasker:
 
     def process_email(self, email_text: str) -> Dict[str, Any]:
         """
-        Process an email by detecting and masking PII entities
+        Process an email by detecting and masking PII entities.
+        The original email is stored in the database for later retrieval if needed.
         """
+        # Mask the email
         masked_email, entity_info = self.mask_text(email_text)
+        
+        # Store the email in the SQLite database
+        email_id, access_key = self.db.store_email(
+            original_email=email_text,
+            masked_email=masked_email,
+            masked_entities=entity_info
+        )
+        
+        # Return the processed data with database references
         return {
-            "input_email_body": email_text,
+            "input_email_body": email_text,  # Return original input for compatibility
             "list_of_masked_entities": entity_info,
             "masked_email": masked_email,
-            "category_of_the_email": ""
-        }
+            "category_of_the_email": "",
+            "email_id": email_id,
+            "access_key": access_key  # Include access key for immediate retrieval if needed
+        }
+    
+    def get_original_email(self, email_id: str, access_key: str) -> Optional[Dict[str, Any]]:
+        """
+        Retrieve the original email with PII using the email ID and access key.
+        
+        Args:
+            email_id: The ID of the stored email
+            access_key: The security key for accessing the original email
+            
+        Returns:
+            The original email data or None if not found or access_key is invalid
+        """
+        return self.db.get_original_email(email_id, access_key)
+    
+    def get_masked_email_by_id(self, email_id: str) -> Optional[Dict[str, Any]]:
+        """
+        Retrieve a masked email by its ID (without the original PII-containing email).
+        
+        Args:
+            email_id: The ID of the stored email
+            
+        Returns:
+            The masked email data or None if not found
+        """
+        return self.db.get_email_by_id(email_id)