Implementation Architecture for Email Threading Algorithms in eDiscovery Pipelines

import asyncio
import logging
import json
import sys
from typing import Dict, List, Optional, Set
from dataclasses import dataclass, field
from collections import defaultdict
from datetime import datetime, timezone

# Structured JSON logging configuration for audit trails
class JSONFormatter(logging.Formatter):
    def format(self, record):
        log_obj = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "level": record.levelname,
            "logger": record.name,
            "message": record.getMessage(),
        }
        if hasattr(record, "extra_data"):
            log_obj.update(record.extra_data)
        return json.dumps(log_obj)

logger = logging.getLogger("email_threading_engine")
handler = logging.StreamHandler(sys.stdout)
handler.setFormatter(JSONFormatter())
logger.addHandler(handler)
logger.setLevel(logging.INFO)

@dataclass
class EmailNode:
    message_id: str
    in_reply_to: Optional[str] = None
    references: List[str] = field(default_factory=list)
    timestamp: Optional[datetime] = None
    subject: Optional[str] = None
    thread_root: Optional[str] = None
    is_orphan: bool = False

class AsyncThreadBuilder:
    def __init__(self, batch_size: int = 10000):
        self.nodes: Dict[str, EmailNode] = {}
        self.adjacency: Dict[str, List[str]] = defaultdict(list)
        self.batch_size = batch_size
        self.audit_log: List[Dict] = []

    def normalize_header(self, raw_id: str) -> Optional[str]:
        """Strip angle brackets, whitespace, and validate format."""
        if not raw_id:
            return None
        cleaned = raw_id.strip().strip("<>")
        return cleaned if "@" in cleaned else None

    async def ingest_batch(self, raw_emails: List[Dict]) -> None:
        """Memory-constrained ingestion with header normalization."""
        for raw in raw_emails:
            msg_id = self.normalize_header(raw.get("message_id"))
            if not msg_id:
                continue

            node = EmailNode(
                message_id=msg_id,
                in_reply_to=self.normalize_header(raw.get("in_reply_to")),
                references=[self.normalize_header(r) for r in raw.get("references", []) if self.normalize_header(r)],
                timestamp=raw.get("timestamp"),
                subject=raw.get("subject")
            )
            self.nodes[msg_id] = node
            parent = node.in_reply_to or (node.references[-1] if node.references else None)
            if parent and parent in self.nodes:
                self.adjacency[parent].append(msg_id)

    def _resolve_cycles(self) -> Set[str]:
        """Detect and break circular references deterministically."""
        visited, stack, broken_edges = set(), set(), set()

        def dfs(node_id: str, path: List[str]):
            if node_id in stack:
                cycle_start = path.index(node_id)
                cycle = path[cycle_start:]
                # Break at the node with the latest timestamp; fall back to the
                # lexicographically largest Message-ID so the choice is
                # deterministic even when timestamps are missing.
                epoch = datetime.min.replace(tzinfo=timezone.utc)
                break_node = max(
                    cycle,
                    key=lambda n: (self.nodes[n].timestamp or epoch, n),
                )
                broken_edges.add(break_node)
                return
            if node_id in visited:
                return
            visited.add(node_id)
            stack.add(node_id)
            path.append(node_id)
            for child in self.adjacency.get(node_id, []):
                dfs(child, path.copy())
            stack.discard(node_id)

        for root in self.nodes:
            dfs(root, [])
        return broken_edges

    async def build_thread_families(self) -> Dict[str, List[str]]:
        """Assemble DAG, resolve cycles, and assign deterministic roots."""
        broken = self._resolve_cycles()
        families: Dict[str, List[str]] = defaultdict(list)

        # Detach broken nodes from their (cyclic) parents before walking up.
        for node_id in broken:
            node = self.nodes[node_id]
            node.in_reply_to = None
            node.references = []

        # Walk each node up to its thread root.
        for msg_id, node in self.nodes.items():
            current = msg_id
            visited_path = {current}
            while True:
                parent = self.nodes[current].in_reply_to
                if not parent or parent not in self.nodes or parent in visited_path:
                    break  # Reached a root (or guarded a residual cycle).
                visited_path.add(parent)
                current = parent

            root_id = current
            node.thread_root = root_id
            families[root_id].append(msg_id)
            logger.info("Thread assigned", extra={"extra_data": {"root": root_id, "child": msg_id}})

        return dict(families)

    async def run(self, raw_batches: List[List[Dict]]) -> Dict[str, List[str]]:
        """Pipeline orchestrator with explicit telemetry."""
        logger.info("Pipeline started", extra={"extra_data": {"total_batches": len(raw_batches)}})
        for i, batch in enumerate(raw_batches):
            await self.ingest_batch(batch)
            logger.info(f"Batch {i+1} ingested", extra={"extra_data": {"count": len(batch)}})

        families = await self.build_thread_families()
        logger.info("Pipeline completed", extra={"extra_data": {"families_generated": len(families)}})
        return families