Prepare HF Docker deployment: add root app.py, Dockerfile, startup handling

Dockerfile

@@ -0,0 +1,17 @@
+FROM python:3.10-slim
+
+# set workdir
+WORKDIR /code
+
+# copy repo
+COPY . /code
+
+# upgrade pip then install
+RUN pip install --upgrade pip
+RUN pip install --no-cache-dir -r requirements.txt
+
+# expose the HF required port
+EXPOSE 7860
+
+# run uvicorn pointing to the app object exposed at repo-root/app.py -> app.app:app would also work
+CMD ["uvicorn", "app.app:app", "--host", "0.0.0.0", "--port", "7860", "--workers", "1"]

Readme.md

@@ -0,0 +1,8 @@
+---
+title: Email Security AI
+emoji: "🛡️"
+colorFrom: "blue"
+colorTo: "purple"
+sdk: docker
+app_port: 7860
+---

app.py

@@ -0,0 +1,2 @@
+# repo-root/app.py
+from app.app import app

app/app.py

@@ -0,0 +1,49 @@
+# app/main.py
+import logging
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+
+# import the router (assumes app/routes/analyze.py exists)
+from app.routes.analyze import router as analyze_router
+
+# import pipeline to warm up singletons on startup
+import app.pipelines.security_pipeline as security_pipeline
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("email-security-ai")
+
+app = FastAPI(
+    title="Email Security AI",
+    version="1.0",
+    description="Model-1 -> Model-2 -> RAG -> Model-3 email security pipeline"
+)
+
+# Allow frontend access during development; tighten in prod.
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],          # change to specific origins in production
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+app.include_router(analyze_router, prefix="")
+
+@app.on_event("startup")
+async def startup_event():
+    """
+    Ensure model singletons are created and log basic info.
+    Note: model initialization (importing security_pipeline) may trigger HF model downloads
+    on first run — that is expected.
+    """
+    logger.info("Starting Email Security AI — warming models (may take time on first run).")
+    try:
+        # Load the newly decoupled pipeline models to memory before serving
+        security_pipeline.load_models()
+        logger.info("Model singletons initialized and ready for inference.")
+    except Exception as e:
+        logger.exception("Error while warming models: %s", e)
+
+@app.get("/health", tags=["health"])
+def health():
+    return {"status": "ok", "ready": True}

app/embeddings/embedding_model.py

@@ -0,0 +1,15 @@
+from langchain_huggingface import HuggingFaceEmbeddings
+
+
+class EmbeddingModel:
+
+    def __init__(self):
+
+        self.model_name = "sentence-transformers/all-MiniLM-L6-v2"
+
+        self.embedding = HuggingFaceEmbeddings(
+            model_name=self.model_name
+        )
+
+    def get(self):
+        return self.embedding

app/models/model1_classifier.py

@@ -0,0 +1,112 @@
+import time
+import torch
+from langchain_core.runnables import RunnableLambda
+from transformers import pipeline
+
+CANDIDATE_LABELS = [
+    "phishing",
+    "scam",
+    "fraud",
+    "malware",
+    "spam",
+    "benign"
+]
+
+class Model1Classifier:
+    def __init__(self):
+        self.model_name = "facebook/bart-large-mnli"
+        self.model_version = "model1-v1"
+        
+        # Determine device
+        if torch.backends.mps.is_available():
+            device = "mps"
+        elif torch.cuda.is_available():
+            device = "cuda"
+        else:
+            device = "cpu"
+
+        # Load HuggingFace pipeline with zero-shot classification
+        self.hf_pipeline = pipeline(
+            "zero-shot-classification",
+            model=self.model_name,
+            device=device
+        )
+
+    def _run_inference(self, input_text: str) -> dict:
+        """Internal method called by the LangChain Runnable"""
+        # Run inference using softmax over all classes and a targeted hypothesis template
+        # `multi_label=False` forces the probabilities to sum to 1.0 (Softmax)
+        result = self.hf_pipeline(
+            input_text,
+            candidate_labels=CANDIDATE_LABELS,
+            multi_label=False,
+            hypothesis_template="This email is {}."
+        )
+
+        labels = result["labels"]
+        scores = result["scores"]
+
+        # --- Heuristic Overrides for extremely short scam messages ---
+        text_lower = input_text.lower()
+        if "won an iphone" in text_lower or (
+            "congratulations" in text_lower and ("won" in text_lower or "prize" in text_lower)
+        ):
+            target_phishing_prob = 0.65
+            
+            try:
+                p_index = labels.index("phishing")
+                current_p_prob = scores[p_index]
+            except ValueError:
+                p_index = -1
+                current_p_prob = 0.0
+
+            if current_p_prob < target_phishing_prob:
+                remaining_target = 1.0 - target_phishing_prob
+                remaining_current = 1.0 - current_p_prob
+                scale_factor = (remaining_target / remaining_current) if remaining_current > 0 else 0
+                
+                for i in range(len(scores)):
+                    if labels[i] == "phishing":
+                        scores[i] = target_phishing_prob
+                    else:
+                        scores[i] = scores[i] * scale_factor
+
+        # Re-sort lists post-override
+        sorted_pairs = sorted(zip(labels, scores), key=lambda x: x[1], reverse=True)
+        labels = [p[0] for p in sorted_pairs]
+        scores = [p[1] for p in sorted_pairs]
+
+        top_labels = []
+        for label, score in list(zip(labels, scores))[:5]:
+            top_labels.append({
+                "label": label,
+                "probability": round(float(score), 4)
+            })
+
+        return {"top_labels": top_labels}
+
+    def predict(self, payload):
+        start = time.perf_counter()
+
+        # Safely extract text fields
+        subject = getattr(payload, "subject", "") or ""
+        
+        text = ""
+        if getattr(payload, "content", None):
+            text = getattr(payload.content, "text", "") or ""
+
+        # Construct input text string
+        input_text = f"{subject} {text}".strip()
+        if not input_text:
+            input_text = "empty message"
+
+        # Invoke inference explicitly without failing LCEL wrapper
+        result_dict = self._run_inference(input_text)
+        
+        elapsed = int((time.perf_counter() - start) * 1000)
+        
+        return {
+            "top_labels": result_dict["top_labels"],
+            "model_version": self.model_version,
+            "time_ms": elapsed
+        }

app/models/model2_risk.py

@@ -0,0 +1,143 @@
+import time
+import torch
+from typing import Dict, Any
+from langchain_core.runnables import RunnableLambda
+from transformers import pipeline
+
+class Model2RiskEngine:
+    def __init__(self):
+        self.model_name = "microsoft/deberta-v3-base"
+        self.model_version = "model2-v3"
+
+        if torch.backends.mps.is_available():
+            device = "mps"
+        elif torch.cuda.is_available():
+            device = "cuda"
+        else:
+            device = "cpu"
+
+        # Initialize the DeBERTa model for sequence classification to identify
+        # specific risk factors inside the email body 
+        try:
+            self.analyzer = pipeline(
+                "zero-shot-classification", 
+                model=self.model_name, 
+                device=device
+            )
+        except Exception as e:
+            print(f"Could not load zero-shot pipeline for {self.model_name}: {e}. Falling back to default.")
+            self.analyzer = None
+
+        self.candidate_reasons = [
+            "Contains urgent financial request",
+            "Suspicious domain detected",
+            "Requests sensitive information"
+        ]
+
+    def _run_nlp_analysis(self, text: str) -> list:
+        reasons = []
+        if self.analyzer is not None and text.strip():
+            try:
+                result = self.analyzer(text, candidate_labels=self.candidate_reasons)
+                # Parse through probabilities
+                for label, score in zip(result['labels'], result['scores']):
+                    if score > 0.4:
+                        reasons.append(label)
+            except Exception as e:
+                print(f"Model 2 inference error: {e}")
+        return reasons
+
+    def analyze(self, payload: Any, model1_result: Dict) -> Dict[str, Any]:
+        start = time.perf_counter()
+
+        reasons = []
+        risk_score = 0
+
+        # Construct input text string
+        text = ""
+        subject = ""
+        if payload.content and payload.content.text:
+            text = payload.content.text.lower()
+        if payload.subject:
+            subject = payload.subject.lower()
+            
+        full_text = f"{subject} {text}".strip()
+
+        # -------------------------------------------------------------------------
+        # 1. Model 1 Integration Score
+        # -------------------------------------------------------------------------
+        if model1_result and "top_labels" in model1_result and len(model1_result["top_labels"]) > 0:
+            top_label = model1_result["top_labels"][0]["label"]
+            top_prob = model1_result["top_labels"][0]["probability"]
+
+            if top_label == "benign":
+                # Inverse penalty if benign
+                risk_score = max(0, 20 * (1 - top_prob))
+            else:
+                risk_score = top_prob * 60
+                reasons.append(f"Message classified as potential {top_label}")
+
+        # -------------------------------------------------------------------------
+        # 2. DeBERTa NLP Analysis
+        # -------------------------------------------------------------------------
+        nlp_reasons = self._run_nlp_analysis(full_text)
+        for r in nlp_reasons:
+            if r not in reasons:
+                reasons.append(r)
+                risk_score += 10 # Bump risk for each NLP factor identified
+
+        # -------------------------------------------------------------------------
+        # 3. Structural Heuristics
+        # -------------------------------------------------------------------------
+        if payload.metadata and getattr(payload.metadata, "has_attachment", False):
+            risk_score += 15
+            reasons.append("Message contains attachment")
+
+        # Link Analysis
+        if payload.links:
+            for link in payload.links:
+                domain = (getattr(link, "domain", "") or "").lower()
+                for keyword in ["login", "verify", "secure", "update", "account"]:
+                    if keyword in domain:
+                        if "Suspicious verification-style domain detected" not in reasons:
+                            risk_score += 15
+                            reasons.append("Suspicious verification-style domain detected")
+                        break
+
+        # Sender Mismatch
+        if payload.sender and getattr(payload.sender, "email", None) and getattr(payload.sender, "domain", None):
+            email_domain = payload.sender.email.split("@")[-1].lower()
+            declared_domain = payload.sender.domain.lower()
+            if email_domain != declared_domain:
+                risk_score += 15
+                reasons.append("Sender domain mismatch detected")
+
+        # -------------------------------------------------------------------------
+        # Clamp Risk Score and Classify Strategy
+        # -------------------------------------------------------------------------
+        risk_score = min(100, int(risk_score))
+
+        if risk_score >= 75:
+            risk_level = "high"
+        elif risk_score >= 50:
+            risk_level = "risky"
+        elif risk_score >= 25:
+            risk_level = "suspicious"
+        else:
+            risk_level = "safe"
+
+        elapsed = int((time.perf_counter() - start) * 1000)
+
+        # Ensure reasons list handles duplicates
+        unique_reasons = []
+        for r in reasons:
+            if r not in unique_reasons:
+                unique_reasons.append(r)
+
+        return {
+            "risk_score": risk_score,
+            "risk_level": risk_level,
+            "reasons": unique_reasons,
+            "model_version": self.model_version,
+            "time_ms": elapsed,
+        }

app/models/model3_action.py

@@ -0,0 +1,105 @@
+import time
+import logging
+from typing import List, Dict, Any, Optional
+
+from langchain_core.prompts import PromptTemplate
+from langchain_core.runnables import RunnableLambda
+from transformers import pipeline
+
+logger = logging.getLogger("email-security-ai.model3")
+logger.addHandler(logging.NullHandler())
+
+class Model3ActionGenerator:
+    """
+    Generates user safety actions using google/flan-t5-base
+    via LangChain LCEL implementation
+    """
+    def __init__(self, model_name: str = "google/flan-t5-base", device: str = "cpu"):
+        self.model_name = model_name
+        self.model_version = "model3-v2"
+        
+        try:
+            self.generator = pipeline(
+                task="text2text-generation",
+                model=self.model_name,
+                device=device
+            )
+            logger.info("Model3 loaded successfully: %s", self.model_name)
+        except Exception as e:
+            logger.exception("Failed to load Model3: %s", e)
+            self.generator = None
+
+        # Build prompt template
+        self.prompt = PromptTemplate.from_template(
+            "Based on a {risk_level} email risk with these reasons: {reasons_text}. "
+            "Relevant cybersecurity knowledge: {evidence_text}. "
+            "What should the user do? List 3 specific short actionable steps separated by commas."
+        )
+
+    def _run_generation(self, prompt_text: str) -> str:
+        if self.generator is None:
+            return ""
+        try:
+            result = self.generator(
+                prompt_text.text if hasattr(prompt_text, "text") else str(prompt_text),
+                max_new_tokens=60,
+                do_sample=False,
+                num_return_sequences=1,
+                repetition_penalty=1.2,
+                no_repeat_ngram_size=3
+            )
+            return result[0]['generated_text']
+        except Exception as e:
+            logger.exception("Model 3 generation failed: %s", e)
+            return ""
+
+    def _parse_output(self, generated_text: str) -> List[str]:
+        if not generated_text:
+            return []
+        
+        # Split by both commas and periods to break up run-on and repeating sentences
+        raw_actions = [a.strip() for a in generated_text.replace(".", ",").split(",") if a.strip()]
+        
+        # Deduplicate while preserving order
+        actions = []
+        for a in raw_actions:
+            if a.lower() not in [existing.lower() for existing in actions]:
+                actions.append(a)
+                
+        return actions[:5]
+
+    def _fallback_actions(self) -> List[str]:
+        return [
+            "Do not click any links in the message.",
+            "Verify the sender through the official website.",
+            "Report the message as phishing.",
+            "Delete the suspicious message."
+        ]
+
+    def generate(self, risk_analysis: Dict[str, Any], rag_evidence: List[str]) -> Dict[str, Any]:
+        start = time.perf_counter()
+
+        risk_level = risk_analysis.get("risk_level", "unknown")
+        reasons = risk_analysis.get("reasons", [])
+        
+        reasons_text = ", ".join(reasons) if reasons else "None"
+        evidence_text = " ".join(rag_evidence) if rag_evidence else "None"
+
+        # Format prompt using LangChain and execute
+        prompt_value = self.prompt.format(
+            risk_level=risk_level, 
+            reasons_text=reasons_text, 
+            evidence_text=evidence_text
+        )
+        actions = self._parse_output(self._run_generation(prompt_value))
+
+        if not actions:
+            actions = self._fallback_actions()
+
+        elapsed = int((time.perf_counter() - start) * 1000)
+
+        return {
+            "actions": actions,
+            "model_version": self.model_version,
+            "time_ms": elapsed
+        }

app/pipelines/security_pipeline.py

@@ -0,0 +1,104 @@
+# app/pipeline/security_pipeline.py
+import time
+
+from app.models.model1_classifier import Model1Classifier
+from app.models.model2_risk import Model2RiskEngine
+from app.models.model3_action import Model3ActionGenerator
+from app.rag.rag_engine import RAGEngine
+
+from app.schemas.response_schema import (
+    AnalyzeResponse,
+    Classification,
+    RiskAnalysis,
+    Metadata,
+)
+
+# -------- LAZY MODEL LOADING --------
+
+model1 = None
+model2 = None
+model3 = None
+rag_engine = None
+
+def load_models():
+    global model1, model2, model3, rag_engine
+    if model1 is None:
+        import os
+        os.makedirs("data/vector_store", exist_ok=True)
+        model1 = Model1Classifier()
+        model2 = Model2RiskEngine()
+        model3 = Model3ActionGenerator()
+        rag_engine = RAGEngine()
+        # Optionally ensure index is loaded
+        if hasattr(rag_engine, "rebuild_index"):
+            rag_engine.rebuild_index()
+
+
+def analyze(payload):
+    load_models()
+    start_total = time.perf_counter()
+
+    # -------------------------------------------------------------
+    # 1. RUN MODEL 1 (SPAM / INTENT CLASSIFICATION)
+    # -------------------------------------------------------------
+    model1_result = model1.predict(payload)
+
+    # -------------------------------------------------------------
+    # 2. RUN MODEL 2 (RISK ENGINE)
+    # -------------------------------------------------------------
+    model2_result = model2.analyze(
+        payload,
+        model1_result
+    )
+
+    # -------------------------------------------------------------
+    # 3. RAG RETRIEVAL (EVIDENCE GATHERING)
+    # -------------------------------------------------------------
+    text_content = payload.content.text if payload.content and payload.content.text else ""
+    subject = payload.subject if payload.subject else ""
+    query = f"{subject} {text_content}".strip()
+    rag_evidence = rag_engine.get_evidence(query) if getattr(rag_engine, "get_evidence", None) else []
+
+    # -------------------------------------------------------------
+    # 4. RUN MODEL 3 (ACTION GENERATION) 
+    # -------------------------------------------------------------
+    model3_result = model3.generate(
+        {"risk_level": model2_result.get("risk_level", "unknown"), "reasons": model2_result.get("reasons", [])},
+        rag_evidence
+    )
+    recommended_actions = model3_result.get("actions", [])
+
+    # -------------------------------------------------------------
+    # TOTAL LATENCY METRICS
+    # -------------------------------------------------------------
+    total_time = int((time.perf_counter() - start_total) * 1000)
+
+    metadata = Metadata(
+        processing_time_ms=total_time,
+        models={
+            "model1": "facebook/bart-large-mnli",
+            "model2": "microsoft/deberta-v3-base",
+            "model3": "google/flan-t5-base"
+        }
+    )
+
+    # -------------------------------------------------------------
+    # BUILD RESPONSE
+    # -------------------------------------------------------------
+    response = AnalyzeResponse(
+        classification=Classification(
+            top_labels=model1_result["top_labels"],
+            model_version=model1_result["model_version"]
+        ),
+        risk_analysis=RiskAnalysis(
+            risk_score=model2_result["risk_score"],
+            risk_level=model2_result["risk_level"],
+            reasons=model2_result["reasons"],
+            model_version=model2_result["model_version"]
+        ),
+        recommended_actions=recommended_actions,
+        rag_evidence=rag_evidence,
+        metadata=metadata
+    )
+
+    return response