452 lines
18 KiB
Python
452 lines
18 KiB
Python
"""
|
|
Indexer script to parse emails from Maildir and push them to Qdrant.
|
|
"""
|
|
|
|
import os
|
|
import email
|
|
import mailbox
|
|
import warnings
|
|
from concurrent.futures import ThreadPoolExecutor, Future
|
|
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", "")
|
|
MAILDIR_FOLDERS = os.environ.get("MAILDIR_FOLDERS", "")
|
|
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 = int(os.environ.get("BATCH_SIZE", "100"))
|
|
EMBEDDING_BATCH_SIZE = int(os.environ.get("EMBEDDING_BATCH_SIZE", "64"))
|
|
METADATA_COLLECTION = "mcp_indexer_metadata"
|
|
INCREMENTAL_DAYS = int(os.environ.get("INCREMENTAL_DAYS", "7"))
|
|
FORCE_REINDEX = os.environ.get("FORCE_REINDEX", "").lower() in ("1", "true", "yes")
|
|
|
|
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, email_id: str = "") -> 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):
|
|
try:
|
|
text = payload.decode(charset, errors='replace')
|
|
except (LookupError, UnicodeDecodeError):
|
|
# Unknown or broken charset — fall back to utf-8
|
|
print(f" Warning: unknown charset '{charset}', falling back to utf-8 [{email_id}]")
|
|
text = payload.decode('utf-8', errors='replace')
|
|
else:
|
|
text = str(payload)
|
|
|
|
if content_type == "text/html":
|
|
with warnings.catch_warnings(record=True) as caught:
|
|
warnings.simplefilter("always")
|
|
text = extract_text_from_html(text)
|
|
for w in caught:
|
|
print(f" Warning: {w.category.__name__}: {w.message} [{email_id}]")
|
|
body_parts.append(text)
|
|
except Exception as e:
|
|
print(f" Warning: error extracting payload: {e} [{email_id}]")
|
|
|
|
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)
|
|
cutoff_ts = cutoff.timestamp()
|
|
recent = set()
|
|
maildir_root = mbox._path # type: ignore[attr-defined]
|
|
|
|
for subdir in ("cur", "new"):
|
|
subdir_path = os.path.join(maildir_root, subdir)
|
|
try:
|
|
filenames = os.listdir(subdir_path)
|
|
except OSError:
|
|
continue
|
|
|
|
for filename in filenames:
|
|
file_path = os.path.join(subdir_path, filename)
|
|
try:
|
|
if os.stat(file_path).st_mtime >= cutoff_ts:
|
|
# Maildir keys are the filename, but the actual file on disk
|
|
# often has a suffix (like ":2,S"). mailbox.Maildir treats
|
|
# the part BEFORE the first colon as the key.
|
|
key = filename.split(":")[0] if ":" in filename else filename
|
|
recent.add(key)
|
|
except OSError:
|
|
continue
|
|
|
|
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)
|
|
|
|
# Force reindex: wipe existing collections to start from scratch
|
|
if FORCE_REINDEX:
|
|
print("[FORCE_REINDEX] Deleting existing collections for a clean re-bootstrap...")
|
|
for col_name in (COLLECTION_NAME, METADATA_COLLECTION):
|
|
try:
|
|
qdrant_client.delete_collection(collection_name=col_name)
|
|
print(f" Deleted collection '{col_name}'.")
|
|
except Exception:
|
|
print(f" Collection '{col_name}' did not exist, skipping.")
|
|
|
|
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).")
|
|
|
|
# Pipeline:
|
|
# parse phase → accumulates (vector_text, payload, point_id) into a batch
|
|
# embed phase → model.embed() called once per batch (vectorized ONNX inference)
|
|
# upsert phase → submitted to a background thread so the next batch can be
|
|
# parsed+embedded while the previous one is in-flight to Qdrant
|
|
#
|
|
# pending_batch: accumulates parsed email metadata + vector_text until BATCH_SIZE
|
|
# pending_future: the in-flight ThreadPoolExecutor Future for the previous upsert
|
|
pending_batch: List[Dict[str, Any]] = []
|
|
pending_future: Future | None = None
|
|
has_error = False
|
|
total_processed = 0
|
|
|
|
def _flush_batch(executor: ThreadPoolExecutor, batch: List[Dict[str, Any]]) -> Future:
|
|
"""Embed a batch of pre-parsed emails and submit an async upsert to Qdrant."""
|
|
vector_texts = [item["vector_text"] for item in batch]
|
|
vectors = [v.tolist() for v in model.embed(vector_texts, batch_size=EMBEDDING_BATCH_SIZE)]
|
|
points = [
|
|
models.PointStruct(
|
|
id=item["point_id"],
|
|
vector=vectors[i],
|
|
payload=item["payload"],
|
|
)
|
|
for i, item in enumerate(batch)
|
|
]
|
|
return executor.submit(
|
|
qdrant_client.upsert,
|
|
collection_name=COLLECTION_NAME,
|
|
points=points,
|
|
)
|
|
|
|
with ThreadPoolExecutor(max_workers=1) as executor:
|
|
# Determine which directories to process
|
|
maildir_roots = []
|
|
if MAILDIR_FOLDERS:
|
|
folders = [f.strip() for f in MAILDIR_FOLDERS.split(",") if f.strip()]
|
|
for folder in folders:
|
|
folder_path = os.path.join(MAILDIR_PATH, folder)
|
|
if os.path.isdir(folder_path):
|
|
maildir_roots.append(folder_path)
|
|
else:
|
|
print(f"Warning: Specified folder not found: {folder_path}")
|
|
else:
|
|
for root_dir, dirs, _ in os.walk(MAILDIR_PATH):
|
|
if all(subdir in dirs for subdir in ['cur', 'new', 'tmp']):
|
|
maildir_roots.append(root_dir)
|
|
|
|
# Iterate and parse over selected maildir directories
|
|
for root in maildir_roots:
|
|
# A valid Maildir has 'cur', 'new', and 'tmp' subdirectories
|
|
if not all(os.path.isdir(os.path.join(root, subdir)) 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 key in 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_raw = msg.get("Message-ID")
|
|
message_id = str(message_id_raw) if message_id_raw is not None else str(uuid.uuid4())
|
|
|
|
# Parse date — msg.get() may return an email.header.Header
|
|
# object instead of str when the header contains non-ASCII
|
|
# bytes (e.g. timezone comments like "heure d'été").
|
|
# We must coerce to str before parsing.
|
|
date_raw = msg.get("Date")
|
|
date_str = str(date_raw) if date_raw is not None else None
|
|
dt_obj = None
|
|
if date_str:
|
|
try:
|
|
dt_obj = parsedate_to_datetime(date_str)
|
|
except Exception:
|
|
pass
|
|
if dt_obj is None:
|
|
# Fallback: warn and use current time
|
|
print(f" Warning: could not parse Date header: {repr(date_raw)} [key={key}, subject={subject}]")
|
|
dt_obj = datetime.now()
|
|
iso_date = dt_obj.isoformat()
|
|
|
|
# Parse body and attachments
|
|
body_text, attachments = parse_email_message(msg, email_id=f"key={key}, subject={subject}")
|
|
|
|
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}"
|
|
)
|
|
|
|
pending_batch.append({
|
|
"vector_text": vector_text,
|
|
"point_id": str(uuid.uuid5(uuid.NAMESPACE_OID, message_id)),
|
|
"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,
|
|
},
|
|
})
|
|
|
|
if len(pending_batch) >= BATCH_SIZE:
|
|
# Wait for the previous upsert to complete before submitting the next
|
|
if pending_future is not None:
|
|
pending_future.result()
|
|
pending_future = _flush_batch(executor, pending_batch)
|
|
total_processed += len(pending_batch)
|
|
print(f" Embedded+upserted batch — {total_processed} emails total so far.")
|
|
pending_batch = []
|
|
|
|
except Exception as e:
|
|
print(f"Error processing email key={key}: {e}")
|
|
has_error = True
|
|
|
|
# Flush remaining emails
|
|
if pending_batch:
|
|
if pending_future is not None:
|
|
pending_future.result()
|
|
pending_future = _flush_batch(executor, pending_batch)
|
|
total_processed += len(pending_batch)
|
|
pending_batch = []
|
|
|
|
# Wait for the last upsert to finish before exiting the executor context
|
|
if pending_future is not None:
|
|
pending_future.result()
|
|
|
|
print(f" Total emails indexed: {total_processed}")
|
|
|
|
# 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()
|