dataset.py

"""
Email dataset for all three tasks.
Each email has ground truth labels hidden from the agent.
"""
from models import Email

# ─── TASK 1: Easy — Binary Spam Detection (10 emails) ─────────────────────────

TASK_EASY_EMAILS = [
    Email(
        id="e001",
        subject="URGENT: You've won $1,000,000! Claim now",
        sender="lottery@prize-winners-global.xyz",
        sender_domain="prize-winners-global.xyz",
        body="Congratulations! You have been selected as our lucky winner. Send your bank details to claim your prize immediately. Limited time offer!!!",
        timestamp="2024-01-15T08:00:00Z",
        _gt_urgency="ignore",
        _gt_category="spam",
        _gt_action="delete",
        _gt_keywords=["lottery", "prize", "bank details"],
    ),
    Email(
        id="e002",
        subject="Q4 Budget Review — Action Required by Friday",
        sender="sarah.chen@acmecorp.com",
        sender_domain="acmecorp.com",
        body="Hi team, please review the attached Q4 budget spreadsheet and send your department's corrections by Friday EOD. We need this finalized before the board meeting Monday.",
        timestamp="2024-01-15T09:15:00Z",
        has_attachment=True,
        _gt_urgency="high",
        _gt_category="finance",
        _gt_action="reply",
        _gt_keywords=["budget", "review", "deadline"],
    ),
    Email(
        id="e003",
        subject="Make $5000/week from home — No experience needed",
        sender="info@work-from-home-secrets.net",
        sender_domain="work-from-home-secrets.net",
        body="Our system guarantees you $5000 per week working just 2 hours a day! Click the link below to get started. Over 50,000 people are already earning big. DON'T MISS OUT!!",
        timestamp="2024-01-15T09:30:00Z",
        _gt_urgency="ignore",
        _gt_category="spam",
        _gt_action="delete",
        _gt_keywords=["earn money", "work from home", "guarantee"],
    ),
    Email(
        id="e004",
        subject="Your invoice #INV-2847 is overdue",
        sender="billing@cloudservices.io",
        sender_domain="cloudservices.io",
        body="This is a reminder that invoice #INV-2847 for $2,340.00 was due on January 10th and remains unpaid. Please process payment within 5 business days to avoid service interruption.",
        timestamp="2024-01-15T10:00:00Z",
        _gt_urgency="high",
        _gt_category="finance",
        _gt_action="escalate",
        _gt_keywords=["invoice", "overdue", "payment"],
    ),
    Email(
        id="e005",
        subject="FREE Viagra and Cialis — No prescription needed",
        sender="pharma@discount-meds-online.ru",
        sender_domain="discount-meds-online.ru",
        body="Get genuine prescription medications at 90% off! No doctor needed. Discreet worldwide shipping. Order today and receive bonus pills!",
        timestamp="2024-01-15T10:15:00Z",
        _gt_urgency="ignore",
        _gt_category="spam",
        _gt_action="delete",
        _gt_keywords=["medication", "no prescription", "free"],
    ),
    Email(
        id="e006",
        subject="Welcome to the team — First day logistics",
        sender="hr@acmecorp.com",
        sender_domain="acmecorp.com",
        body="Welcome aboard! Please find attached your first-day agenda, building access card instructions, and IT setup guide. Your manager Alex will meet you at reception at 9am Monday.",
        timestamp="2024-01-15T11:00:00Z",
        has_attachment=True,
        _gt_urgency="medium",
        _gt_category="hr",
        _gt_action="reply",
        _gt_keywords=["welcome", "onboarding", "first day"],
    ),
    Email(
        id="e007",
        subject="RE: Your account has been compromised — Verify immediately",
        sender="security@bankofamerica-secure-verify.com",
        sender_domain="bankofamerica-secure-verify.com",
        body="We detected suspicious activity on your account. Click here immediately to verify your identity and prevent account suspension: http://bit.ly/verify-now-urgent",
        timestamp="2024-01-15T11:30:00Z",
        _gt_urgency="ignore",
        _gt_category="spam",
        _gt_action="delete",
        _gt_keywords=["phishing", "suspicious link", "fake bank"],
    ),
    Email(
        id="e008",
        subject="Client escalation — Apex Industries threatening to cancel contract",
        sender="james.rodriguez@acmecorp.com",
        sender_domain="acmecorp.com",
        body="Urgent: Apex Industries (our $2M annual contract) called to say they're experiencing repeated service outages. Their CTO is furious and mentioned cancellation. I need executive involvement ASAP.",
        timestamp="2024-01-15T12:00:00Z",
        _gt_urgency="critical",
        _gt_category="customer_complaint",
        _gt_action="escalate",
        _gt_keywords=["contract cancellation", "executive", "urgent"],
    ),
    Email(
        id="e009",
        subject="Newsletter: Top 10 productivity tips for January",
        sender="newsletter@productivityhacks.com",
        sender_domain="productivityhacks.com",
        body="Happy New Year! Here are our top 10 productivity tips to kickstart 2024. Tip 1: Time-block your calendar. Tip 2: Use the Pomodoro technique...",
        timestamp="2024-01-15T13:00:00Z",
        _gt_urgency="low",
        _gt_category="other",
        _gt_action="archive",
        _gt_keywords=["newsletter", "tips", "low priority"],
    ),
    Email(
        id="e010",
        subject="Legal notice — GDPR compliance audit findings",
        sender="compliance@legalpartners.co.uk",
        sender_domain="legalpartners.co.uk",
        body="Following our audit, we have identified 3 areas of non-compliance with GDPR Article 17 and Article 30. Please review the attached report and schedule a call with our team within 10 business days to discuss remediation.",
        timestamp="2024-01-15T14:00:00Z",
        has_attachment=True,
        _gt_urgency="critical",
        _gt_category="legal",
        _gt_action="escalate",
        _gt_keywords=["GDPR", "compliance", "legal"],
    ),
]

# ─── TASK 2: Medium — Priority Inbox Triage (15 emails) ──────────────────────

TASK_MEDIUM_EMAILS = TASK_EASY_EMAILS + [
    Email(
        id="m001",
        subject="Interview scheduling request — Senior Engineer candidate",
        sender="recruiting@talent-pipeline.io",
        sender_domain="talent-pipeline.io",
        body="Hi, I'm reaching out on behalf of candidate Maria Santos for the Senior Software Engineer role. She's available Tuesday 2-5pm or Thursday 10am-1pm. Could you confirm which slot works for the hiring panel?",
        timestamp="2024-01-16T08:30:00Z",
        _gt_urgency="medium",
        _gt_category="hr",
        _gt_action="forward",
        _gt_keywords=["interview", "scheduling", "candidate"],
    ),
    Email(
        id="m002",
        subject="CRITICAL: Production database down — all services affected",
        sender="alerts@monitoring.acmecorp.com",
        sender_domain="acmecorp.com",
        body="ALERT: Primary database cluster db-prod-01 is unreachable. All customer-facing services are returning 500 errors. Incident started at 14:32 UTC. On-call engineer has been paged. Status page updated.",
        timestamp="2024-01-16T14:35:00Z",
        _gt_urgency="critical",
        _gt_category="internal_ops",
        _gt_action="escalate",
        _gt_keywords=["outage", "production", "database down"],
    ),
    Email(
        id="m003",
        subject="Partnership proposal — Co-marketing opportunity",
        sender="partnerships@techstartup.ai",
        sender_domain="techstartup.ai",
        body="Hi, we're a Series A startup in the AI space and would love to explore a co-marketing partnership. We have 50k email subscribers and active social presence. Would your team be open to a 30-min exploratory call?",
        timestamp="2024-01-16T10:00:00Z",
        _gt_urgency="low",
        _gt_category="sales_inquiry",
        _gt_action="forward",
        _gt_keywords=["partnership", "marketing", "proposal"],
    ),
    Email(
        id="m004",
        subject="Performance improvement plan — follow-up",
        sender="jennifer.walsh@acmecorp.com",
        sender_domain="acmecorp.com",
        body="Following our conversation last week, I've drafted the formal PIP documentation for Tom B. in the engineering team. I need you to review and sign before we proceed with the 30-day timeline. Attached for your review.",
        timestamp="2024-01-16T15:00:00Z",
        has_attachment=True,
        _gt_urgency="high",
        _gt_category="hr",
        _gt_action="reply",
        _gt_keywords=["PIP", "HR", "performance"],
    ),
    Email(
        id="m005",
        subject="Press inquiry — TechCrunch wants comment on recent layoffs",
        sender="reporter@techcrunch.com",
        sender_domain="techcrunch.com",
        body="Hi, I'm a reporter at TechCrunch writing a story about tech layoffs in Q1 2024. I have information that Acme Corp reduced headcount by 12% last month. Can you provide a comment or connect me with your PR team? Deadline is 5pm today.",
        timestamp="2024-01-16T11:00:00Z",
        _gt_urgency="critical",
        _gt_category="pr",
        _gt_action="escalate",
        _gt_keywords=["press", "media", "deadline"],
    ),
]

# ─── TASK 3: Hard — Full Triage with Response Drafting (20 emails) ────────────

TASK_HARD_EMAILS = TASK_MEDIUM_EMAILS + [
    Email(
        id="h001",
        subject="Refund request — Order #78432 damaged in transit",
        sender="customer.angry@gmail.com",
        sender_domain="gmail.com",
        body="I am absolutely furious. I ordered a $450 laptop stand 3 weeks ago and it arrived completely smashed. I've attached photos. Your customer service has ignored my last 2 emails. I want a full refund immediately or I'm disputing with my credit card company and leaving reviews everywhere. This is unacceptable.",
        timestamp="2024-01-17T09:00:00Z",
        has_attachment=True,
        _gt_urgency="high",
        _gt_category="customer_complaint",
        _gt_action="reply",
        _gt_keywords=["refund", "damaged", "angry customer", "apologize", "resolution", "replacement"],
    ),
    Email(
        id="h002",
        subject="Enterprise license inquiry — 500 seat deployment",
        sender="procurement@fortune500corp.com",
        sender_domain="fortune500corp.com",
        body="Our organization is evaluating your enterprise software for a 500-seat deployment. We need volume pricing, SSO/SAML support confirmation, SOC 2 Type II certification, and a data processing agreement. Our evaluation timeline is 6 weeks. Who should we work with?",
        timestamp="2024-01-17T10:30:00Z",
        _gt_urgency="critical",
        _gt_category="sales_inquiry",
        _gt_action="reply",
        _gt_keywords=["enterprise", "500 seats", "pricing", "SSO", "SOC2", "sales team"],
    ),
    Email(
        id="h003",
        subject="Workplace harassment complaint — formal report",
        sender="employee.confidential@acmecorp.com",
        sender_domain="acmecorp.com",
        body="I'm submitting this formal complaint regarding repeated inappropriate comments made by my direct manager over the past 3 months. I have documented 7 incidents with dates and witnesses. I'm requesting a formal HR investigation and would like to keep this confidential from my manager during the process.",
        timestamp="2024-01-17T11:00:00Z",
        _gt_urgency="critical",
        _gt_category="hr",
        _gt_action="reply",
        _gt_keywords=["harassment", "confidential", "investigation", "HR", "acknowledge", "safe"],
    ),
    Email(
        id="h004",
        subject="Contract renewal — 30 days until expiration",
        sender="accounts@supplier-logistics.com",
        sender_domain="supplier-logistics.com",
        body="This is a reminder that your Master Service Agreement with Supplier Logistics expires on February 15th. To avoid service disruption, please initiate the renewal process. Updated terms with a 15% price increase are attached. Please confirm your intent to renew.",
        timestamp="2024-01-17T14:00:00Z",
        has_attachment=True,
        _gt_urgency="high",
        _gt_category="finance",
        _gt_action="forward",
        _gt_keywords=["contract", "renewal", "deadline", "legal review", "procurement"],
    ),
    Email(
        id="h005",
        subject="Security vulnerability disclosure — CVE-2024-0142",
        sender="security@bugcrowd.com",
        sender_domain="bugcrowd.com",
        body="A security researcher has identified a critical SQL injection vulnerability in your public API endpoint /api/v2/search. Proof-of-concept exploit is attached. The researcher is requesting $5,000 in bug bounty. This is being held for 90 days before public disclosure per responsible disclosure policy.",
        timestamp="2024-01-17T15:30:00Z",
        has_attachment=True,
        _gt_urgency="critical",
        _gt_category="support",
        _gt_action="escalate",
        _gt_keywords=["security", "vulnerability", "SQL injection", "engineering", "CISO"],
    ),
]

# ─── Task lookup ──────────────────────────────────────────────────────────────

TASK_EMAILS = {
    "task_easy":   TASK_EASY_EMAILS,
    "task_medium": TASK_MEDIUM_EMAILS,
    "task_hard":   TASK_HARD_EMAILS,
}