classifier.py

import json
import os
import google.generativeai as genai

_client = None


def _get_client():
    global _client
    if _client is None:
        genai.configure(api_key=os.environ["GEMINI_API_KEY"])
        _client = genai.GenerativeModel("gemini-3.1-flash-lite-preview")
    return _client


_INSTRUCTIONS = """You are an email triage assistant. Classify each email into exactly one of three categories:

CRITICAL — Requires immediate attention or action:
  - Government agencies (tax, immigration, licensing, courts)
  - Banks, financial institutions, investment accounts
  - Security alerts or account verification requests
  - Legal notices, contracts, or compliance matters
  - Medical appointments or health-related notifications
  - Insurance or utility notices requiring action
  - Any email explicitly requesting a response or action soon

MEDIUM — Good to know but not urgent:
  - Order confirmations and shipping updates
  - Newsletters or publications you're subscribed to
  - Social media and app notifications
  - Non-urgent work or professional correspondence
  - Messages from friends or family that don't need immediate reply

TRIVIAL — Promotional noise requiring no attention:
  - Marketing emails, discount offers, sales
  - Unsolicited promotional content
  - Bulk advertising from brands or services"""


def classify_emails_batch(emails: list) -> list:
    """
    Classify a batch of emails with Gemini.
    Returns a list of dicts: [{index, category, reason}, ...] ordered 1-based.
    """
    email_blocks = []
    for i, email in enumerate(emails, start=1):
        email_blocks.append(
            f"Email {i}:\n"
            f"From: {email['from']}\n"
            f"Subject: {email['subject']}\n"
            f"Preview: {email['snippet']}"
        )

    prompt = (
        _INSTRUCTIONS
        + "\n\nClassify the following emails. "
        "Respond with a JSON array only — no markdown, no extra text.\n"
        "Each element: {\"index\": <1-based int>, \"category\": \"CRITICAL\"|\"MEDIUM\"|\"TRIVIAL\", \"reason\": \"<one line>\"}\n\n"
        + "\n\n".join(email_blocks)
    )

    response = _get_client().generate_content(prompt)
    raw = response.text.strip()

    # Strip markdown code fences if Gemini wrapped the JSON
    if raw.startswith("```"):
        lines = raw.splitlines()
        raw = "\n".join(lines[1:-1] if lines[-1] == "```" else lines[1:])

    return json.loads(raw)