webmail-client/backend/app/completion.py

"""Sentence completion in the compose window: phrases the user writes, finished.

No AI: a per-user index of the word sequences the user has written at least
twice, built from their own sent mail (quoted replies and signatures cut off),
plus a few common email phrases per language to start from. Typing "I look
forward" finds what usually follows those words — "to hearing from you." — and
the compose window shows it in grey until Tab takes it.

It works in any language the user writes in: matching is on the words just
typed, so a Polish sentence only ever meets Polish phrases.

The index lives in ``completion_phrases`` (context → continuation, with how
often it was seen) and is rebuilt in the background at most once a day, when
the compose window asks for suggestions (:func:`build_index_job`).
"""

from __future__ import annotations

import datetime as dt
import html as html_lib
import logging
import re
import uuid
from collections import Counter
from collections.abc import Iterable
from dataclasses import dataclass
from typing import Any

from sqlalchemy import delete, select
from sqlalchemy.dialects.postgresql import insert

from app.db import session_scope
from app.models import (
    Account,
    BodyPart,
    CompletionIndexState,
    CompletionPhrase,
    Mailbox,
    Message,
)

logger = logging.getLogger(__name__)

START = "^"  # a sentence's start, as a word: "^ pozdr" completes the first word
MAX_CONTEXT = 3  # words before the cursor matched on
MAX_CONTINUATION = 6  # words suggested at most
MIN_COUNT = 2  # seen this often to be offered; one-offs are noise (and private)
EXTEND_SHARE = 0.6  # a longer suggestion only while it is what usually follows
CORPUS_MESSAGES = 500  # most recent sent messages read
CORPUS_CHARS = 3000  # of each
REBUILD_AFTER = dt.timedelta(days=1)

# Where the user's own words end and quoted mail or a signature begins: a reply
# header ("On … wrote:", Polish Gmail's "… napisał(a):", …), a forwarded or
# original-message marker, or an Outlook-style header block.
_QUOTE_START = re.compile(
    r"^\s*(?:"
    r"-{2,}\s*(?:original message|forwarded message|wiadomość oryginalna"
    r"|wiadomość przekazana dalej|ursprüngliche nachricht|weitergeleitete nachricht)"
    r"|(?:from|od|von|de):\s.+"
    r"|sent from my\s"
    r")",
    re.IGNORECASE,
)
_QUOTE_HEADER_END = re.compile(
    # German names the writer after the verb: "Am … schrieb Anna:".
    r"(?:wrote|napisał\(a\)|napisała?|pisze|schrieb(?:\s+.{0,100})?|a écrit|escribió)\s*:\s*$",
    re.IGNORECASE,
)
_SIGNATURE = re.compile(r"^--\s*$")
_SENTENCE_END = re.compile(r"(?<=[.!?])\s+|\n+")
_WORD = re.compile(r"\S+")
_EDGE_PUNCT = re.compile(r"^[\"'“”„«»(\[–—-]+|[\"'“”„«»)\],.;:!?…–—-]+$")
# Words a phrase is never learned across: numbers (phones, dates, prices),
# addresses and links — useless to complete, and often private.
_NOT_A_WORD = re.compile(r"\d|@|www\.|https?:", re.IGNORECASE)


def html_to_text(html: str) -> str:
    """An HTML signature as text lines (block tags and <br> break lines)."""
    text = re.sub(r"(?i)<br\s*/?>|</(?:p|div|li|tr|h[1-6])\s*>", "\n", html)
    text = re.sub(r"<[^>]+>", "", text)
    return html_lib.unescape(text)


def own_text(text: str, signatures: Iterable[str] = ()) -> str:
    """The part of a sent message the user wrote: no quoted mail, no signature
    (the ``--`` kind, or one of ``signatures`` — the accounts' own, as text)."""
    sig_lines = {norm_line(x) for sig in signatures for x in sig.splitlines() if norm_line(x)}
    kept: list[str] = []
    for line in text.splitlines():
        if _SIGNATURE.match(line) or _QUOTE_START.match(line) or _QUOTE_HEADER_END.search(line):
            break
        if norm_line(line) in sig_lines:
            continue
        if line.lstrip().startswith(">"):
            continue  # an inline-quoted line
        kept.append(line)
    return "\n".join(kept)


def norm_line(line: str) -> str:
    return " ".join(line.lower().split())


def norm(word: str) -> str:
    """A word as matched: lower case, outer punctuation off ("Hello," -> "hello")."""
    return _EDGE_PUNCT.sub("", word.lower())


def sentences(text: str) -> Iterable[list[str]]:
    for chunk in _SENTENCE_END.split(text):
        words = _WORD.findall(chunk)
        if words:
            yield words


@dataclass(frozen=True)
class Phrase:
    context: str  # normalized words before the cursor, space-joined ("^ i look")
    continuation: str  # the words that follow, as written ("forward to hearing")
    count: int


def build_phrases(texts: Iterable[str], weight: int = 1) -> dict[tuple[str, str], int]:
    """(context, continuation) -> count, for every sequence seen MIN_COUNT times.

    Counts each sentence's word sequences (with the start marker), keeps the
    frequent ones, and splits each into "what was typed" and "what follows".
    """
    counts: Counter[tuple[str, ...]] = Counter()
    surface: dict[tuple[str, ...], tuple[str, ...]] = {}
    longest = MAX_CONTEXT + MAX_CONTINUATION
    for text in texts:
        for words in sentences(text):
            seq = [START, *words]
            keys = [START, *("" if _NOT_A_WORD.search(w) else norm(w) for w in words)]
            for i in range(len(seq)):
                for n in range(2, min(longest, len(seq) - i) + 1):
                    gram = tuple(keys[i : i + n])
                    if "" in gram:
                        break  # a lone punctuation mark: no phrase across it
                    counts[gram] += weight
                    surface.setdefault(gram, tuple(seq[i : i + n]))
    out: dict[tuple[str, str], int] = {}
    for gram, count in counts.items():
        if count < MIN_COUNT:
            continue
        written = surface[gram]
        for k in range(1, min(MAX_CONTEXT, len(gram) - 1) + 1):
            rest = len(gram) - k
            if rest > MAX_CONTINUATION:
                continue
            key = (" ".join(gram[:k]), " ".join(written[k:]))
            out[key] = max(out.get(key, 0), count)
    return out


# A few common phrases per language, for users with little sent mail yet.
STARTER_PHRASES: dict[str, list[str]] = {
    "en": [
        "I look forward to hearing from you.",
        "I'm looking forward to hearing from you.",
        "Looking forward to hearing from you.",
        "Looking forward to working with you.",
        "Thank you for your email.",
        "Thank you for your help.",
        "Thank you for your quick reply.",
        "Thanks in advance.",
        "Please let me know if you have any questions.",
        "Let me know if you have any questions.",
        "Let me know what you think.",
        "Please find attached the",
        "I hope this email finds you well.",
        "I hope you're doing well.",
        "Sorry for the late reply.",
        "I would be happy to help.",
        "Feel free to contact me if you need anything else.",
        "Could you please send me the",
        "Please confirm at your earliest convenience.",
        "Have a great day!",
        "Have a nice weekend!",
        "Best regards,",
        "Kind regards,",
    ],
    "pl": [
        "Z poważaniem,",
        "Pozdrawiam serdecznie,",
        "Dziękuję za wiadomość.",
        "Dziękuję za szybką odpowiedź.",
        "Dziękuję za pomoc.",
        "Z góry dziękuję.",
        "W załączniku przesyłam",
        "Proszę o informację zwrotną.",
        "W razie pytań proszę o kontakt.",
        "Czekam na odpowiedź.",
        "Przepraszam za opóźnienie w odpowiedzi.",
        "Będę wdzięczny za odpowiedź.",
        "Będę wdzięczna za odpowiedź.",
        "Zgodnie z naszą rozmową",
        "Proszę o potwierdzenie.",
        "Daj znać, co myślisz.",
        "Miłego dnia!",
        "Miłego weekendu!",
    ],
    "de": [
        "Mit freundlichen Grüßen",
        "Viele Grüße",
        "Beste Grüße",
        "Vielen Dank für Ihre E-Mail.",
        "Vielen Dank für Ihre schnelle Antwort.",
        "Vielen Dank im Voraus.",
        "Ich freue mich auf Ihre Antwort.",
        "Bei Fragen stehe ich Ihnen gerne zur Verfügung.",
        "Im Anhang finden Sie",
        "Entschuldigen Sie die späte Antwort.",
        "Ich hoffe, es geht Ihnen gut.",
        "Schönen Tag noch!",
        "Schönes Wochenende!",
    ],
}

# Built once: every starter phrase counts as seen MIN_COUNT times.
_STARTERS: dict[str, Counter[str]] = {}
for (_ctx, _cont), _n in build_phrases(
    (p for phrases in STARTER_PHRASES.values() for p in phrases), weight=MIN_COUNT
).items():
    _STARTERS.setdefault(_ctx, Counter())[_cont] = _n

_LEADING_PUNCT = re.compile(r"^[\"'“”„«»(\[]+")


def _candidates(phrases: dict[str, Counter[str]], context: str, partial: str) -> Counter[str]:
    """Continuations of ``context`` whose first word starts with ``partial``."""
    out: Counter[str] = Counter()
    for continuation, count in phrases.get(context, Counter()).items():
        first = continuation.split(" ", 1)[0]
        if partial and not first.lower().startswith(partial.lower()):
            continue
        if partial and len(first) <= len(partial):
            continue  # that word is typed in full already
        out[continuation] = count
    return out


def choose(candidates: Counter[str]) -> str | None:
    """The suggestion: the most frequent next word, extended word by word only
    while the longer phrase is what usually follows (EXTEND_SHARE)."""
    if not candidates:
        return None
    by_first: Counter[str] = Counter()
    for continuation, count in candidates.items():
        first = continuation.split(" ", 1)[0]
        by_first[first] = max(by_first[first], count)
    chosen, chosen_count = by_first.most_common(1)[0]
    while True:
        depth = chosen.count(" ") + 1
        longer = [
            (c, n)
            for c, n in candidates.items()
            if c.startswith(chosen + " ") and c.count(" ") == depth
        ]
        if not longer:
            return chosen
        best, best_count = max(longer, key=lambda x: x[1])
        if best_count < max(MIN_COUNT, EXTEND_SHARE * chosen_count):
            return chosen
        chosen, chosen_count = best, best_count


def contexts(before: str) -> tuple[list[str], str]:
    """For the text before the cursor: the contexts to try, longest first, and
    the partly typed word ('' after a space). No contexts: nothing to suggest
    (an empty sentence, or the cursor right after a punctuation mark)."""
    if not before or not before.strip():
        return [], ""
    last = before[-1]
    if not last.isspace() and not last.isalnum() and last not in "'’-":
        return [], ""  # "you," or "end." — not mid-word, and no space typed yet
    tail = _SENTENCE_END.split(before)[-1]
    words = _WORD.findall(tail)
    partial = ""
    if words and not last.isspace():
        partial = _LEADING_PUNCT.sub("", words.pop())
    if not words and not partial:
        return [], ""  # a new sentence with nothing in it yet
    keys = [START, *(norm(w) for w in words)]
    if "" in keys:  # a lone punctuation mark: the phrase starts after it
        keys = keys[len(keys) - keys[::-1].index("") :]
    if not keys:
        return [], ""
    return [" ".join(keys[-k:]) for k in range(min(MAX_CONTEXT, len(keys)), 0, -1)], partial


def suggest(personal: dict[str, Counter[str]], before: str) -> str | None:
    """What to show after ``before`` (the text up to the cursor), or None.

    The returned text starts where the cursor is: the rest of a partly typed
    word, or (after a space) the next words. The user's own phrases and the
    starter phrases count together.
    """
    tried, partial = contexts(before)
    for context in tried:
        merged = _candidates(personal, context, partial)
        for cont, count in _candidates(_STARTERS, context, partial).items():
            merged[cont] = merged.get(cont, 0) + count
        if " " not in context and not partial:
            # A single word before a space is a weak hint: only a strong habit.
            merged = Counter({c: n for c, n in merged.items() if n >= 2 * MIN_COUNT})
        best = choose(merged)
        if best:
            return best[len(partial) :] if partial else best
    return None


def complete(personal: dict[str, Counter[str]], before: str) -> str | None:
    """:func:`suggest`, and for a word typed in full but no space yet ("I look
    forward"), what follows it — with the space ("␣to hearing from you.")."""
    found = suggest(personal, before)
    if found is None and before[-1:].isalnum():
        after_space = suggest(personal, before + " ")
        if after_space:
            return " " + after_space
    return found


def keys_for(before: str) -> list[str]:
    """The contexts :func:`complete` may look up, to fetch them from the index."""
    keys = contexts(before)[0]
    if before[-1:].isalnum():
        keys += contexts(before + " ")[0]
    return keys


async def personal_phrases(
    session: Any, user_id: uuid.UUID, keys: list[str]
) -> dict[str, Counter[str]]:
    rows = await session.execute(
        select(
            CompletionPhrase.context, CompletionPhrase.continuation, CompletionPhrase.count
        ).where(CompletionPhrase.user_id == user_id, CompletionPhrase.context.in_(keys))
    )
    out: dict[str, Counter[str]] = {}
    for context, continuation, count in rows:
        out.setdefault(context, Counter())[continuation] = count
    return out


async def build_index_job(ctx: dict[str, Any], user_id: str) -> int:
    """Rebuild a user's phrase index from their recent sent mail. Returns its size."""
    uid = uuid.UUID(user_id)
    async with session_scope() as session:
        texts = list(
            await session.scalars(
                select(BodyPart.content)
                .join(Message, Message.id == BodyPart.message_id)
                .join(Mailbox, Mailbox.id == Message.mailbox_id)
                .where(
                    Message.user_id == uid,
                    Mailbox.special_use == "sent",
                    BodyPart.kind == "text",
                )
                .order_by(Message.internal_date.desc())
                .limit(CORPUS_MESSAGES)
            )
        )
        signatures = [
            html_to_text(sig)
            for sig in await session.scalars(
                select(Account.signature).where(Account.user_id == uid, Account.signature != "")
            )
        ]
        phrases = build_phrases(own_text(t, signatures)[:CORPUS_CHARS] for t in texts)
        await session.execute(delete(CompletionPhrase).where(CompletionPhrase.user_id == uid))
        rows = [
            {"user_id": uid, "context": c, "continuation": cont, "count": n}
            for (c, cont), n in phrases.items()
        ]
        for start in range(0, len(rows), 1000):
            await session.execute(insert(CompletionPhrase), rows[start : start + 1000])
        stmt = insert(CompletionIndexState).values(user_id=uid, built_at=dt.datetime.now(dt.UTC))
        await session.execute(
            stmt.on_conflict_do_update(
                index_elements=["user_id"], set_={"built_at": stmt.excluded.built_at}
            )
        )
        await session.commit()
    logger.info(
        "completion index for %s: %d phrases from %d messages", user_id, len(rows), len(texts)
    )
    return len(rows)