# main.py import os import asyncio from dotenv import load_dotenv from fastapi import FastAPI, BackgroundTasks from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel from typing import Optional from datetime import datetime from contextlib import asynccontextmanager from supabase import create_client, Client # Import your agent functions from agents.security_agent import run_security_check, load_security_model, clean_email_text from agents.summary_agent import run_summarization, load_summary_model from agents.classification_agent import load_classification_model, run_classification from agents.extraction_agent import run_extraction from agents.translation_agent import run_translation from agents.reply_agent import generate_smart_reply load_dotenv() SUPABASE_URL = os.getenv("VITE_SUPABASE_URL") SUPABASE_KEY = os.getenv("VITE_SUPABASE_ANON_KEY") supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY) # --- Global Lock for the Sweeper --- is_sweeping = False # 1. The Startup Loader @asynccontextmanager async def lifespan(app: FastAPI): print("🚀 Booting up the PFE Intelligence Backend...") # --- DYNAMIC PATH RESOLUTION --- # This works automatically on both Windows (Local) and Linux (Hugging Face) BASE_DIR = os.path.dirname(os.path.abspath(__file__)) MODEL_PATH = os.path.join(BASE_DIR, "agents", "security_model_v2") # ------------------------------- load_security_model(MODEL_PATH) load_summary_model() load_classification_model() yield app = FastAPI(lifespan=lifespan) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # 2. Schemas class BoundingBox(BaseModel): x0: int y0: int x1: int y1: int text: str class EmailPayload(BaseModel): id: str user_id: str threadId: str sender: str date: datetime body: str historyId: Optional[str] = None subject: Optional[str] = None body_plain: Optional[str] = None snippet: Optional[str] = None class TranslationRequest(BaseModel): id: str body: str body_plain: str target_lang: str = "English" class SecurityFeedback(BaseModel): email_id: str is_secure: bool text: str class ReplyRequest(BaseModel): email_id: str text: str sender: str tone: str def update_agent_status(agent_id: str, status: str, current_task: str, increment: bool = False): try: update_data = { "status": status, "currentTask": current_task, "lastActivity": datetime.utcnow().isoformat() } if increment: res = supabase.table("agent_status").select("emailsProcessed").eq("id", agent_id).execute() if res.data: current_count = res.data[0].get("emailsProcessed", 0) update_data["emailsProcessed"] = current_count + 1 supabase.table("agent_status").update(update_data).eq("id", agent_id).execute() except Exception as e: print(f"⚠️ Failed to update agent {agent_id}: {e}") # 3. Core Processing Logic async def process_single_email(email_data: EmailPayload, raw_text: str, force_safe: bool = False): try: # --- 1. GATEKEEPER (Security / mapped to agent-1) --- if force_safe: update_agent_status("agent-1", "processing", f"Human override for {email_data.sender}...") clean_text = clean_email_text(raw_text) sec_result = { "is_secure": True, "phishing_score": 0.0, "clean_text": clean_text } update_agent_status("agent-1", "idle", "Bypassed (Human Override)", increment=True) else: update_agent_status("agent-1", "processing", f"Scanning {email_data.sender} for threats...") sec_result = await run_security_check(raw_text) update_agent_status("agent-1", "idle", "Standing by...", increment=True) clean_text = sec_result["clean_text"] ai_summary, ai_category, extracted_entities, ai_urgency = None, None, None, None performance_metrics = {"phishing_score": sec_result["phishing_score"]} if sec_result["is_secure"]: # --- 3. CLASSIFIER (mapped to agent-2) --- update_agent_status("agent-2", "processing", "Categorizing email content...") class_result = await run_classification(subject=email_data.subject or "", text=clean_text) ai_category = class_result["category"] update_agent_status("agent-2", "idle", "Standing by...", increment=True) # --- 4. EXTRACTOR / PRIORITY (mapped to agent-3) --- update_agent_status("agent-3", "processing", "Extracting entities & urgency level...") ext_result = await run_extraction(clean_text, ai_category) extracted_entities = ext_result["extracted_data"] ai_urgency = ext_result["urgency"] update_agent_status("agent-3", "idle", "Standing by...", increment=True) # --- 2. PROCESSOR / SUMMARIZER (mapped to agent-4) --- update_agent_status("agent-4", "processing", "Generating concise summary...") sum_result = await run_summarization(clean_text) ai_summary = sum_result["summary"] update_agent_status("agent-4", "idle", "Standing by...", increment=True) performance_metrics.update({ "summary": {"latency": sum_result["latency"]}, "classification": class_result["metrics"], "extraction": { "latency_s": ext_result["latency_s"], "model": "llama-3.1-8b-instant" } }) else: ai_summary = "⚠️ Processing aborted: Security threat detected." # 5. SAVE TO SUPABASE supabase.table("Email").update({ "is_secure": sec_result["is_secure"], "ai_summary": ai_summary, "ai_category": ai_category, "ai_extracted_entities": extracted_entities, "ai_urgency_score": ai_urgency, "performance_metrics": performance_metrics, "is_analyzed": True }).eq("id", email_data.id).execute() print(f"✅ Full pipeline complete for {email_data.id}") except Exception as e: print(f"⚠️ Pipeline crashed for {email_data.id}: {e}") # 4. The Continuous Sweeper Task async def continuous_sweeper_task(): global is_sweeping if is_sweeping: print("⏩ Sweeper is already running in the background. Skipping new trigger.") return is_sweeping = True print("🚀 Master Sweeper Engine started...") try: while True: response = supabase.table("Email").select("*").eq("is_analyzed", False).limit(10).execute() pending_emails = response.data if not pending_emails: print("✨ All emails have been analyzed! Sweeper going to sleep.") break print(f"📦 Processing new batch of {len(pending_emails)} emails...") for email_row in pending_emails: email_payload = EmailPayload( id=email_row["id"], user_id=email_row["user_id"], threadId=email_row["threadId"], sender=email_row["sender"], date=email_row["date"], body=email_row["body"] or "", historyId=email_row.get("historyId"), subject=email_row.get("subject"), body_plain=email_row.get("body_plain"), snippet=email_row.get("snippet") ) raw_text = email_row.get("body_plain") or email_row.get("body") or email_row.get("snippet") or "" await process_single_email(email_payload, raw_text) if len(pending_emails) < 10: print("✨ Final batch complete! Sweeper going to sleep.") break print("⏳ Batch complete. Resting for 15 seconds...") await asyncio.sleep(15) except Exception as e: print(f"⚠️ Master Sweeper crashed: {e}") finally: is_sweeping = False # 5. API Endpoints @app.post("/api/analyze-email") async def analyze_email(email: EmailPayload, background_tasks: BackgroundTasks): # --- NEW: Check the Feedback memory before scanning --- feedback_res = supabase.table("SecurityFeedback").select("is_phishing").eq("email_id", email.id).execute() force_safe = False if feedback_res.data and feedback_res.data[0].get("is_phishing") is False: force_safe = True print(f"🧠 Memory triggered: Bypassing security scan for {email.id} (Human Override)") db_insert_data = { "id": email.id, "user_id": email.user_id, "threadId": email.threadId, "historyId": email.historyId, "subject": email.subject, "sender": email.sender, "date": email.date.isoformat(), "body": email.body, "body_plain": email.body_plain, "snippet": email.snippet, "is_analyzed": False } supabase.table("Email").upsert(db_insert_data).execute() raw_text = email.body_plain or email.body or email.snippet or "" background_tasks.add_task(process_single_email, email, raw_text, force_safe) return {"status": "success", "message": "Email queued for full processing"} @app.post("/api/process-pending") async def process_pending_emails(background_tasks: BackgroundTasks): background_tasks.add_task(continuous_sweeper_task) return {"status": "success", "message": "Sweeper engine started in the background."} @app.post("/api/translate-email") async def translate_email_endpoint(payload: TranslationRequest): try: update_agent_status("agent-translation", "processing", f"Translating to {payload.target_lang}...") result = await run_translation(payload.body, payload.body_plain, payload.target_lang) if not result: update_agent_status("agent-translation", "idle", "Skipped - Content too short") return {"status": "skipped", "message": "Content too short or failed"} email_record = supabase.table("Email").select("performance_metrics").eq("id", payload.id).execute() current_metrics = email_record.data[0].get("performance_metrics", {}) if email_record.data else {} current_metrics["translation"] = result["metrics"] supabase.table("Email").update({ "body_translated": result["translated_html"], "body_plain": result["refined_plain"], "performance_metrics": current_metrics }).eq("id", payload.id).execute() update_agent_status("agent-translation", "idle", "Standing by...", increment=True) return {"status": "success", "data": result} except Exception as e: print(f"Translation Error: {e}") update_agent_status("agent-translation", "idle", "Error encountered") return {"status": "error", "message": str(e)} @app.post("/api/feedback/security") async def log_security_feedback(feedback: SecurityFeedback, background_tasks: BackgroundTasks): try: cleaned_text = clean_email_text(feedback.text) db_insert_data = { "email_id": feedback.email_id, "corrected_text": cleaned_text, "is_phishing": not feedback.is_secure, "processed": False } supabase.table("SecurityFeedback").upsert( db_insert_data, on_conflict="email_id" ).execute() print(f"🛡️ Security feedback logged for email {feedback.email_id}") email_res = supabase.table("Email").select("*").eq("id", feedback.email_id).execute() if email_res.data: email_row = email_res.data[0] email_payload = EmailPayload( id=email_row["id"], user_id=email_row["user_id"], threadId=email_row["threadId"], sender=email_row["sender"], date=email_row["date"], body=email_row["body"] or "", historyId=email_row.get("historyId"), subject=email_row.get("subject"), body_plain=email_row.get("body_plain"), snippet=email_row.get("snippet") ) raw_text = email_row.get("body_plain") or email_row.get("body") or email_row.get("snippet") or "" supabase.table("Email").update({ "is_analyzed": False, "ai_summary": None, "is_secure": True }).eq("id", feedback.email_id).execute() background_tasks.add_task(process_single_email, email_payload, raw_text, True) return {"status": "success", "message": "Feedback logged and email queued for full processing."} except Exception as e: print(f"⚠️ Failed to process security feedback: {e}") return {"status": "error", "message": str(e)} @app.post("/api/generate-reply") async def generate_reply_endpoint(payload: ReplyRequest): try: update_agent_status("agent-reply", "processing", f"Drafting a {payload.tone} reply...") draft = await generate_smart_reply( text=payload.text, sender=payload.sender, tone=payload.tone ) update_agent_status("agent-reply", "idle", "Standing by...", increment=True) return {"status": "success", "reply": draft} except Exception as e: update_agent_status("agent-reply", "idle", "Error encountered") return {"status": "error", "message": str(e)} @app.get("/") def read_root(): return {"status": "The PFE Intelligence Backend is awake and ready!"}