mcp-maildir/src/indexer.py

"""
Indexer script to parse emails from Maildir and push them to Qdrant.
"""

import os
import email
import mailbox
from datetime import datetime
from email.utils import parsedate_to_datetime, parseaddr
from email.header import decode_header
from typing import List, Dict, Any, Tuple
import uuid

from dotenv import load_dotenv
from qdrant_client import QdrantClient
from qdrant_client.http import models
from fastembed import TextEmbedding
from bs4 import BeautifulSoup

# Load .env config
load_dotenv()

# Configuration
MAILDIR_PATH = os.environ.get("MAILDIR_PATH", "")
QDRANT_URL = os.environ.get("QDRANT_URL", "")
COLLECTION_NAME = os.environ.get("COLLECTION_NAME", "")

if not MAILDIR_PATH:
    raise ValueError("MAILDIR_PATH environment variable is required.")
if not QDRANT_URL:
    raise ValueError("QDRANT_URL environment variable is required.")
if not COLLECTION_NAME:
    raise ValueError("COLLECTION_NAME environment variable is required.")

EMBEDDING_MODEL_NAME = os.environ.get("EMBEDDING_MODEL_NAME", "BAAI/bge-small-en-v1.5")
BATCH_SIZE = 50
METADATA_COLLECTION = "mcp_indexer_metadata"
INCREMENTAL_DAYS = int(os.environ.get("INCREMENTAL_DAYS", "7"))

def decode_mime_words(s: str) -> str:
    """Decodes MIME encoded strings (e.g. subjects, filenames)."""
    if not s:
        return ""
    decoded_words = decode_header(s)
    result = []
    for word, encoding in decoded_words:
        if isinstance(word, bytes):
            try:
                result.append(word.decode(encoding or 'utf-8', errors='replace'))
            except LookupError:
                result.append(word.decode('utf-8', errors='replace'))
        else:
            result.append(word)
    return "".join(result)

def extract_text_from_html(html_content: str) -> str:
    """Extracts plain text from HTML content."""
    try:
        soup = BeautifulSoup(html_content, "html.parser")
        return soup.get_text(separator=" ", strip=True)
    except Exception:
        return html_content

def normalize_email_address(value: str) -> str:
    """Extracts and normalizes the bare email address from a header value."""
    if not value:
        return ""
    _, addr = parseaddr(value)
    return (addr or value).strip().lower()


def parse_email_message(msg: mailbox.Message) -> Tuple[str, List[str]]:
    """Extracts plain text body and a list of attachment filenames."""
    body_parts = []
    attachments = []

    for part in msg.walk():
        # Skip multiparts, we only care about leaf nodes
        if part.is_multipart():
            continue
        
        content_type = part.get_content_type()
        content_disposition = str(part.get("Content-Disposition", ""))

        # Check for attachments
        if "attachment" in content_disposition or part.get_filename():
            filename = part.get_filename()
            if filename:
                attachments.append(decode_mime_words(filename))
            continue
            
        # Extract text body
        if content_type in ["text/plain", "text/html"]:
            try:
                payload = part.get_payload(decode=True)
                if payload:
                    charset = part.get_content_charset('utf-8') or 'utf-8'
                    if isinstance(payload, bytes):
                        text = payload.decode(charset, errors='replace')
                    else:
                        text = str(payload)
                        
                    if content_type == "text/html":
                        text = extract_text_from_html(text)
                    body_parts.append(text)
            except Exception as e:
                print(f"Error extracting payload: {e}")
                pass

    return "\n".join(body_parts).strip(), attachments

def init_qdrant_collection(client: QdrantClient, vector_size: int):
    """Ensures Qdrant collection exists and payload indexes are created."""
    
    # Check if collection exists
    collections = client.get_collections().collections
    if not any(c.name == COLLECTION_NAME for c in collections):
        print(f"Creating collection '{COLLECTION_NAME}' with vector size {vector_size}...")
        client.create_collection(
            collection_name=COLLECTION_NAME,
            vectors_config=models.VectorParams(size=vector_size, distance=models.Distance.COSINE),
        )
    else:
        print(f"Collection '{COLLECTION_NAME}' already exists.")

    # Create payload indexes for filtering metadata deterministically
    print("Ensuring payload indexes exist...")
    
    # Date index (DATETIME)
    client.create_payload_index(
        collection_name=COLLECTION_NAME,
        field_name="date",
        field_schema=models.PayloadSchemaType.DATETIME,
    )
    
    # Sender index (KEYWORD)
    client.create_payload_index(
        collection_name=COLLECTION_NAME,
        field_name="sender",
        field_schema=models.PayloadSchemaType.KEYWORD,
    )
    
    # Receiver index (KEYWORD)
    client.create_payload_index(
        collection_name=COLLECTION_NAME,
        field_name="receiver",
        field_schema=models.PayloadSchemaType.KEYWORD,
    )

def init_metadata_collection(client: QdrantClient):
    """Ensures the indexer metadata collection exists in Qdrant."""
    collections = client.get_collections().collections
    if not any(c.name == METADATA_COLLECTION for c in collections):
        print(f"Creating metadata collection '{METADATA_COLLECTION}'...")
        client.create_collection(
            collection_name=METADATA_COLLECTION,
            # Minimal vector (size=1) — we only use this collection for payload storage
            vectors_config=models.VectorParams(size=1, distance=models.Distance.COSINE),
        )


def is_bootstrap_done(client: QdrantClient) -> bool:
    """Returns True if a successful full bootstrap has already been recorded."""
    try:
        results, _ = client.scroll(
            collection_name=METADATA_COLLECTION,
            scroll_filter=models.Filter(
                must=[
                    models.FieldCondition(
                        key="event",
                        match=models.MatchValue(value="bootstrap_complete"),
                    )
                ]
            ),
            limit=1,
        )
        return len(results) > 0
    except Exception as e:
        print(f"Warning: could not check bootstrap state: {e}")
        return False


def mark_bootstrap_done(client: QdrantClient):
    """Records a bootstrap_complete event in the metadata collection."""
    point_id = str(uuid.uuid5(uuid.NAMESPACE_OID, "bootstrap_complete"))
    client.upsert(
        collection_name=METADATA_COLLECTION,
        points=[
            models.PointStruct(
                id=point_id,
                vector=[0.0],  # placeholder — collection is payload-only
                payload={
                    "event": "bootstrap_complete",
                    "timestamp": datetime.now().isoformat(),
                },
            )
        ],
    )
    print("Bootstrap state recorded in Qdrant metadata collection.")


def get_recent_keys(mbox: mailbox.Maildir, days: int) -> set:
    """
    Returns the set of Maildir keys whose backing file has been modified
    within the last `days` days (based on filesystem mtime).
    """
    from datetime import timezone, timedelta

    cutoff = datetime.now(tz=timezone.utc) - timedelta(days=days)
    recent = set()
    maildir_root = mbox._path  # type: ignore[attr-defined]

    for key in mbox.keys():
        # Maildir keys are stored under cur/ or new/ — try both
        for subdir in ("cur", "new"):
            file_path = os.path.join(maildir_root, subdir, key)
            if os.path.isfile(file_path):
                try:
                    mtime = os.path.getmtime(file_path)
                    mtime_dt = datetime.fromtimestamp(mtime, tz=timezone.utc)
                    if mtime_dt >= cutoff:
                        recent.add(key)
                except OSError:
                    pass
                break  # found the file, no need to check the other subdir

    return recent


def main():
    """
    Main ingestion function.
    Reads Maildir, extracts text, generates local embeddings, and pushes to Qdrant.
    """
    print(f"Indexing emails from {MAILDIR_PATH} into {QDRANT_URL}...")

    if not os.path.exists(MAILDIR_PATH):
        print(f"Error: Maildir path not found: {MAILDIR_PATH}")
        return

    # Initialize model
    print(f"Loading embedding model: {EMBEDDING_MODEL_NAME}...")
    model = TextEmbedding(model_name=EMBEDDING_MODEL_NAME)
    vector_size = len(next(iter(model.embed(["dimension_probe"]))))

    # Initialize Qdrant
    print("Connecting to Qdrant...")
    qdrant_client = QdrantClient(url=QDRANT_URL)
    init_qdrant_collection(qdrant_client, vector_size)
    init_metadata_collection(qdrant_client)

    # Determine indexing mode: full bootstrap or incremental update
    if is_bootstrap_done(qdrant_client):
        mode = "incremental"
        print(f"[MODE] INCREMENTAL — scanning only files modified in the last {INCREMENTAL_DAYS} days.")
    else:
        mode = "full"
        print("[MODE] BOOTSTRAP — full scan of all emails (first-time indexing).")

    points = []
    has_error = False

    # Iterate and parse over all maildir directories found in MAILDIR_PATH
    for root, dirs, files in os.walk(MAILDIR_PATH):
        # A valid Maildir has 'cur', 'new', and 'tmp' subdirectories
        if not all(subdir in dirs for subdir in ['cur', 'new', 'tmp']):
            continue

        print(f"Processing Maildir found at: {root}")
        mbox = mailbox.Maildir(root)
        total_emails_in_dir = len(mbox)

        if mode == "incremental":
            keys_to_process = get_recent_keys(mbox, INCREMENTAL_DAYS)
            print(
                f"Found {total_emails_in_dir} emails total, "
                f"{len(keys_to_process)} modified in the last {INCREMENTAL_DAYS} days."
            )
        else:
            keys_to_process = set(mbox.keys())
            print(f"Found {total_emails_in_dir} emails — indexing all.")

        for idx, key in enumerate(keys_to_process):
            try:
                msg = mbox[key]

                # Parse headers
                subject = decode_mime_words(msg.get("Subject", "No Subject"))
                sender_raw = decode_mime_words(msg.get("From", "Unknown"))
                receiver_raw = decode_mime_words(msg.get("To", "Unknown"))
                sender = normalize_email_address(sender_raw)
                receiver = normalize_email_address(receiver_raw)
                message_id = msg.get("Message-ID", str(uuid.uuid4()))

                # Parse date
                date_str = msg.get("Date")
                dt_obj = None
                if date_str:
                    try:
                        dt_obj = parsedate_to_datetime(date_str)
                    except Exception:
                        pass

                if dt_obj is None:
                    dt_obj = datetime.now()

                # Format to ISO 8601 for Qdrant DATETIME index
                iso_date = dt_obj.isoformat()

                # Parse Body and Attachments
                body_text, attachments = parse_email_message(msg)

                # Prepare Vector text
                attachments_str = ", ".join(attachments) if attachments else "None"
                vector_text = (
                    f"Date: {iso_date}\n"
                    f"From: {sender}\n"
                    f"To: {receiver}\n"
                    f"Subject: {subject}\n\n"
                    f"{body_text}\n\n"
                    f"Attachments: {attachments_str}"
                )

                # Embed the text — fastembed returns an iterable of numpy arrays
                embeddings = list(model.embed([vector_text]))
                vector = embeddings[0].tolist()

                # Prepare payload (metadata)
                payload = {
                    "message_id": message_id,
                    "date": iso_date,
                    "sender": sender,
                    "sender_raw": sender_raw,
                    "receiver": receiver,
                    "receiver_raw": receiver_raw,
                    "subject": subject,
                    "body_text": body_text,
                    "attachments": attachments,
                }

                # Assign deterministic UUID point ID based on message_id
                point_id = str(uuid.uuid5(uuid.NAMESPACE_OID, message_id))

                points.append(models.PointStruct(
                    id=point_id,
                    vector=vector,
                    payload=payload,
                ))

                # Push in batches
                if len(points) >= BATCH_SIZE:
                    qdrant_client.upsert(
                        collection_name=COLLECTION_NAME,
                        points=points,
                    )
                    print(f"  Upserted batch — {idx + 1}/{len(keys_to_process)} emails processed in current directory.")
                    points = []

            except Exception as e:
                print(f"Error processing email key={key}: {e}")
                has_error = True

    # Push remaining points
    if points:
        qdrant_client.upsert(
            collection_name=COLLECTION_NAME,
            points=points,
        )

    # Record bootstrap completion so subsequent runs use incremental mode
    if mode == "full" and not has_error:
        mark_bootstrap_done(qdrant_client)

    print("Indexing completed successfully!")

if __name__ == "__main__":
    main()