mosaic/utils/text_utils.py

"""Text chunking, cleaning, and token estimation for Mosaic."""
import re
import hashlib
import logging
from typing import Optional

logger = logging.getLogger(__name__)

# Rough token estimation: ~4 chars per token for English text
CHARS_PER_TOKEN = 4

# Default chunk sizes
DEFAULT_CHUNK_TOKENS = 8000
DEFAULT_OVERLAP_TOKENS = 500
EMAIL_MAX_TOKENS = 4000
PDF_MAX_TOKENS = 8000


def estimate_tokens(text: str) -> int:
    """Estimate token count from text length."""
    return len(text) // CHARS_PER_TOKEN


def clean_text(text: str) -> str:
    """Clean extracted text: normalize whitespace, remove control chars."""
    if not text:
        return ""
    # Remove null bytes and other control characters (keep newlines, tabs)
    text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]', '', text)
    # Normalize line endings
    text = text.replace('\r\n', '\n').replace('\r', '\n')
    # Collapse excessive blank lines (3+ → 2)
    text = re.sub(r'\n{3,}', '\n\n', text)
    # Strip trailing whitespace per line
    text = '\n'.join(line.rstrip() for line in text.split('\n'))
    return text.strip()


def chunk_text(text: str, doc_type: str = "text",
               max_tokens: int = DEFAULT_CHUNK_TOKENS,
               overlap_tokens: int = DEFAULT_OVERLAP_TOKENS,
               metadata: Optional[dict] = None) -> list[dict]:
    """Chunk text using per-document-type strategy.

    Args:
        text: The text to chunk
        doc_type: Document type (html, pdf, email, cable, text)
        max_tokens: Maximum tokens per chunk
        overlap_tokens: Token overlap between chunks
        metadata: Optional metadata (e.g., cable headers)

    Returns:
        List of dicts with 'text', 'chunk_index', 'total_chunks', 'metadata'
    """
    if not text:
        return []

    if doc_type == 'cable':
        return _chunk_cable(text, max_tokens, overlap_tokens, metadata)
    elif doc_type == 'pdf':
        return _chunk_pdf(text, max_tokens, overlap_tokens)
    elif doc_type == 'email':
        return _chunk_email(text, max_tokens, overlap_tokens, metadata)
    elif doc_type == 'html':
        return _chunk_html(text, max_tokens, overlap_tokens)
    else:
        return _chunk_default(text, max_tokens, overlap_tokens)


def _chunk_cable(text: str, max_tokens: int, overlap_tokens: int,
                  metadata: Optional[dict] = None) -> list[dict]:
    """Chunk cable: headers as structured context, body chunked independently."""
    headers = metadata or {}
    header_text = ""
    if headers:
        header_lines = []
        for key in ('SUBJECT', 'ORIGIN', 'CLASSIFICATION', 'TAGS', 'DATE'):
            if key in headers:
                header_lines.append(f"{key}: {headers[key]}")
        header_text = '\n'.join(header_lines)

    # If small enough, return as one chunk
    total_tokens = estimate_tokens(header_text + '\n\n' + text)
    if total_tokens <= max_tokens:
        return [{'text': (header_text + '\n\n' + text).strip(),
                 'chunk_index': 0, 'total_chunks': 1, 'metadata': headers}]

    # Chunk body, prepend headers to each chunk
    body_chunks = _chunk_default(text, max_tokens - estimate_tokens(header_text) - 50,
                                  overlap_tokens)
    for i, chunk in enumerate(body_chunks):
        chunk['text'] = header_text + '\n\n' + chunk['text']
        chunk['metadata'] = headers
        chunk['chunk_index'] = i
        chunk['total_chunks'] = len(body_chunks)
    return body_chunks


def _chunk_pdf(text: str, max_tokens: int, overlap_tokens: int) -> list[dict]:
    """Chunk PDF by page boundaries (indicated by form feeds or page markers)."""
    # Split by form feed or common page break markers
    pages = re.split(r'\f|\n-{3,}\n|\n={3,}\n', text)
    pages = [p.strip() for p in pages if p.strip()]

    if not pages:
        return _chunk_default(text, max_tokens, overlap_tokens)

    chunks = []
    current_text = ""
    for page in pages:
        combined_tokens = estimate_tokens(current_text + '\n\n' + page)
        if combined_tokens > max_tokens and current_text:
            chunks.append(current_text.strip())
            # Keep overlap from end of previous chunk
            overlap_chars = overlap_tokens * CHARS_PER_TOKEN
            current_text = current_text[-overlap_chars:] + '\n\n' + page
        else:
            current_text = (current_text + '\n\n' + page).strip()

    if current_text.strip():
        chunks.append(current_text.strip())

    return [{'text': c, 'chunk_index': i, 'total_chunks': len(chunks), 'metadata': {}}
            for i, c in enumerate(chunks)]


def _chunk_email(text: str, max_tokens: int, overlap_tokens: int,
                  metadata: Optional[dict] = None) -> list[dict]:
    """Chunk email: headers + body as atomic unit unless body is too large."""
    headers = metadata or {}
    total_tokens = estimate_tokens(text)

    if total_tokens <= EMAIL_MAX_TOKENS:
        return [{'text': text, 'chunk_index': 0, 'total_chunks': 1, 'metadata': headers}]

    # Split headers from body
    header_section = ""
    body = text
    if '\n\n' in text:
        parts = text.split('\n\n', 1)
        header_section = parts[0]
        body = parts[1] if len(parts) > 1 else ""

    # Chunk body, prepend headers
    body_chunks = _chunk_default(body, max_tokens - estimate_tokens(header_section) - 50,
                                  overlap_tokens)
    for i, chunk in enumerate(body_chunks):
        chunk['text'] = header_section + '\n\n' + chunk['text']
        chunk['metadata'] = headers
        chunk['chunk_index'] = i
        chunk['total_chunks'] = len(body_chunks)
    return body_chunks


def _chunk_html(text: str, max_tokens: int, overlap_tokens: int) -> list[dict]:
    """Chunk HTML: section-aware chunking preserving heading hierarchy."""
    # Split on headings
    sections = re.split(r'(?=\n#{1,6}\s|\n[A-Z][A-Za-z\s]{3,}\n[=-]+\n)', text)
    sections = [s for s in sections if s.strip()]

    if not sections or len(sections) == 1:
        return _chunk_default(text, max_tokens, overlap_tokens)

    chunks = []
    current_text = ""
    for section in sections:
        combined_tokens = estimate_tokens(current_text + '\n\n' + section)
        if combined_tokens > max_tokens and current_text:
            chunks.append(current_text.strip())
            overlap_chars = overlap_tokens * CHARS_PER_TOKEN
            current_text = current_text[-overlap_chars:] + '\n\n' + section
        else:
            current_text = (current_text + '\n\n' + section).strip()

    if current_text.strip():
        chunks.append(current_text.strip())

    return [{'text': c, 'chunk_index': i, 'total_chunks': len(chunks), 'metadata': {}}
            for i, c in enumerate(chunks)]


def _chunk_default(text: str, max_tokens: int, overlap_tokens: int) -> list[dict]:
    """Default chunking: paragraph-aware with token overlap."""
    total_tokens = estimate_tokens(text)
    if total_tokens <= max_tokens:
        return [{'text': text, 'chunk_index': 0, 'total_chunks': 1, 'metadata': {}}]

    max_chars = max_tokens * CHARS_PER_TOKEN
    overlap_chars = overlap_tokens * CHARS_PER_TOKEN

    # Split into paragraphs
    paragraphs = text.split('\n\n')
    chunks = []
    current = ""

    for para in paragraphs:
        if len(current) + len(para) + 2 > max_chars and current:
            chunks.append(current.strip())
            # Keep overlap
            current = current[-overlap_chars:] + '\n\n' + para
        else:
            current = (current + '\n\n' + para) if current else para

    if current.strip():
        chunks.append(current.strip())

    # If still too large (single giant paragraph), force-split
    final_chunks = []
    for chunk in chunks:
        if len(chunk) > max_chars * 1.5:
            # Force split at sentence boundaries or character limit
            for sub in _force_split(chunk, max_chars, overlap_chars):
                final_chunks.append(sub)
        else:
            final_chunks.append(chunk)

    return [{'text': c, 'chunk_index': i, 'total_chunks': len(final_chunks), 'metadata': {}}
            for i, c in enumerate(final_chunks)]


def _force_split(text: str, max_chars: int, overlap_chars: int) -> list[str]:
    """Force-split text at sentence boundaries or character limits."""
    sentences = re.split(r'(?<=[.!?])\s+', text)
    chunks = []
    current = ""

    for sent in sentences:
        if len(current) + len(sent) + 1 > max_chars and current:
            chunks.append(current.strip())
            current = current[-overlap_chars:] + ' ' + sent
        else:
            current = (current + ' ' + sent) if current else sent

    if current.strip():
        chunks.append(current.strip())

    return chunks


def dedup_text(texts: list[str], threshold: float = 0.95) -> list[str]:
    """Remove near-duplicate texts using hash comparison.

    For exact duplicates. Near-duplicate detection would require
    more sophisticated approaches (minhash, simhash, etc.)
    """
    seen = set()
    unique = []
    for text in texts:
        h = hashlib.md5(text.encode()).hexdigest()
        if h not in seen:
            seen.add(h)
            unique.append(text)
    return unique