FabledCurator/backend/app/services/discord_grouping.py

"""Discord drop grouping — FC authors the post that Discord never wrote.

Milestone 388, step E2.

Discord is a delivery CHANNEL, not a publisher. A creator drops a set of
near-variants — the same piece with different hair colour, accessories, an
outfit swap — across a handful of messages, and today each of those messages
lands as its own `post` row, so chat lines compete with authored work for the
same surface. The fix is not to demote them into a second-class feed; it is to
let FC write the post: one row per DROP, its images the drop's images, its body
the messages' text in arrival order.

The result is post-shaped by construction, which is the entire reason to
synthesise a `Post` rather than invent a parallel entity — feed, provenance,
translation, attachments and series all keep working on it unchanged.

## The predicate: three axes, ANDed, and the time one does the real work

**Similarity alone over-groups, and that is the failure that would make this
useless.** Any two pieces of the same character by the same artist sit close in
SigLIP space; a cosine-only rule collapses a month of one character into a
single "post". What makes a variant set a set is that it was dropped TOGETHER.

    same source  AND  cosine distance <= threshold  AND  no gap > window

Two details in there are load-bearing:

* **Distance is measured to the group's SEED, never to the previous member.**
  Chaining to the previous member lets a group DRIFT: twenty small steps walk
  from one piece to a completely different one, each hop individually within
  threshold. Anchoring on the seed bounds the whole group to one neighbourhood.
* **The window is measured between CONSECUTIVE messages, not from the first.**
  An artist trickling variants out over an evening is one drop; a window
  anchored on the first message would cut it in half at an arbitrary point.

## Why this is a post-import sweep and not part of ingest

The obvious alternative was to migrate Discord to the native post-first
ingester (#1266) and group at capture time. **That cannot work**, and the
reason is worth recording: the grouping signal is `siglip_embedding`, which is
produced ASYNCHRONOUSLY after import (`tasks/ml.py`, the GPU queue backfill).
At capture time the embedding does not exist yet, so an ingester has nothing to
group on. Grouping is necessarily something that happens once the vectors have
caught up — which also means this sweep must be re-runnable and must simply
skip what it cannot yet place. It does: a post whose image has no embedding is
left alone and picked up on a later run.

## The honesty rule

A synthetic post must never pretend an artist authored it. It carries
`synthesized_by`, records what it was built from in `synthesis_details`
(members, count, and the thresholds in force at the time), and leaves its
member posts intact and reachable. Deleting the synthetic post releases the
members back into the feed — one DELETE, no repair step. FC invented this
grouping; the operator has to be able to see that, inspect it, and undo it.
"""

from __future__ import annotations

import logging
import math
from dataclasses import dataclass, field
from datetime import UTC, datetime, timedelta

from sqlalchemy import Select, func, select, update
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.ext.asyncio import AsyncSession

from ..models import ImageProvenance, ImageRecord, MLSettings, Post, Source

log = logging.getLogger(__name__)

# The value that lands in `post.synthesized_by`. One grouper today; a second
# would be another value here, which is exactly why the column has no CHECK.
DROP_GROUPER = "discord_drop"

PLATFORM = "discord"

# Ceiling on member posts examined per source per run. A first sweep over an
# established library would otherwise pull every Discord message's 1152-float
# vector into memory at once. The sweep is re-runnable and works oldest-first,
# so a backlog simply drains over successive runs rather than needing one
# heroic pass.
MAX_CANDIDATES_PER_SOURCE = 500


@dataclass
class DropGroup:
    """One drop: the member posts, in arrival order, that will become a post."""

    member_ids: list[int] = field(default_factory=list)
    seed: list[float] | None = None
    last_at: datetime | None = None


def cosine_distance(a, b) -> float:
    """Cosine distance between two embeddings, in the same units pgvector's
    `cosine_distance` operator returns (0 = identical, 1 = orthogonal).

    Computed in Python rather than SQL because the comparison is against a
    group seed held in a loop, not against a column — and pure arithmetic keeps
    numpy off this path entirely. pgvector may hand back a numpy array or a
    list depending on driver version, so both are coerced.
    """
    va = [float(x) for x in a]
    vb = [float(x) for x in b]
    # strict=True: two embeddings of different length is a corrupted row or a
    # model swap that skipped the re-embed, and silently truncating to the
    # shorter one would score it as a near match.
    dot = sum(x * y for x, y in zip(va, vb, strict=True))
    na = math.sqrt(sum(x * x for x in va))
    nb = math.sqrt(sum(y * y for y in vb))
    if na == 0.0 or nb == 0.0:
        # A zero vector has no direction, so no meaningful distance. Return the
        # maximum so it can never pull anything into a group.
        return 1.0
    return 1.0 - (dot / (na * nb))


def _candidate_stmt(source_id: int, *, not_after: datetime) -> Select:
    """Ungrouped Discord message-posts, one representative image each, OLDEST
    FIRST — which is the order `build_groups` requires.

    DISTINCT ON the post picks the lowest-id embedded image as that post's
    representative: a Discord message carrying several attachments is still one
    point in the drop, and comparing every attachment would let one incidental
    image drag an unrelated message into the group.

    The DISTINCT ON is wrapped in a subquery rather than ordered directly,
    because Postgres requires a DISTINCT ON query's ORDER BY to LEAD with the
    distinct expression — so the inner query must sort by `post.id`, which is
    insertion order and not arrival order at all once a backfill has imported
    anything out of sequence. Sorting outside is what makes the caller's LIMIT
    take the OLDEST candidates instead of the lowest-numbered ones.
    """
    sort_key = func.coalesce(Post.post_date, Post.downloaded_at)
    inner = (
        select(
            Post.id.label("post_id"),
            sort_key.label("occurred_at"),
            ImageRecord.siglip_embedding.label("embedding"),
        )
        .join(ImageRecord, ImageRecord.primary_post_id == Post.id)
        .where(
            Post.source_id == source_id,
            # Never absorb a post FC wrote, and never re-absorb one already
            # taken — both would build groups out of groups.
            Post.synthesized_by.is_(None),
            Post.absorbed_by_post_id.is_(None),
            ImageRecord.siglip_embedding.is_not(None),
            sort_key <= not_after,
        )
        .distinct(Post.id)
        .order_by(Post.id, ImageRecord.id)
        .subquery()
    )
    return (
        select(inner.c.post_id, inner.c.occurred_at, inner.c.embedding)
        .order_by(inner.c.occurred_at, inner.c.post_id)
    )


def build_groups(
    rows: list[tuple[int, datetime, list[float]]],
    *,
    max_distance: float,
    window: timedelta,
) -> list[DropGroup]:
    """Walk candidates in arrival order and cut them into drops.

    `rows` must be sorted oldest-first — the whole predicate is about
    adjacency in time, so an unsorted input would silently produce nonsense
    rather than fail.
    """
    groups: list[DropGroup] = []
    current: DropGroup | None = None

    for post_id, occurred_at, embedding in rows:
        if current is not None:
            gap_ok = occurred_at - current.last_at <= window
            # Distance to the SEED, not to the previous member — see the module
            # docstring on drift.
            near = cosine_distance(current.seed, embedding) <= max_distance
            if gap_ok and near:
                current.member_ids.append(post_id)
                current.last_at = occurred_at
                continue
            groups.append(current)
        current = DropGroup(
            member_ids=[post_id], seed=embedding, last_at=occurred_at,
        )

    if current is not None:
        groups.append(current)
    return groups


async def _synthesize(
    session: AsyncSession,
    *,
    source: Source,
    group: DropGroup,
    max_distance: float,
    window_minutes: float,
) -> Post | None:
    """Write one synthetic post for `group` and absorb its members."""
    members = (await session.execute(
        select(Post)
        .where(Post.id.in_(group.member_ids))
        .order_by(func.coalesce(Post.post_date, Post.downloaded_at), Post.id)
    )).scalars().all()
    if not members:
        return None

    first = members[0]
    # Deterministic key, so a re-run cannot mint a second post for the same
    # drop: the unique (source_id, external_post_id) constraint would reject it
    # even if the member filter somehow let the drop through twice.
    external_id = f"fc-drop:{first.external_post_id}"[:128]

    # The messages' own text, in arrival order, IS the post's body — that is
    # what the operator asked for and it is the only text a drop has. Blank
    # messages (an attachment with no caption) contribute nothing rather than a
    # run of empty lines.
    body = "\n\n".join(m.description.strip() for m in members if m.description and m.description.strip())

    post = Post(
        source_id=source.id,
        artist_id=source.artist_id,
        external_post_id=external_id,
        # post_title stays NULL DELIBERATELY. A synthesised title is the one
        # place this feature could accidentally put words in a creator's mouth;
        # the UI labels the row from `synthesized_by` instead, which cannot be
        # mistaken for something the artist wrote.
        post_title=None,
        post_url=first.post_url,
        post_date=first.post_date or first.downloaded_at,
        description=body or None,
        synthesized_by=DROP_GROUPER,
        synthesis_details={
            "member_post_ids": [m.id for m in members],
            "message_count": len(members),
            # The thresholds AS THEY WERE. They are operator-tunable, so
            # without this "why did it group these" is unanswerable later.
            "max_distance": max_distance,
            "window_minutes": window_minutes,
            "grouped_at": datetime.now(UTC).isoformat(),
            # Growth accumulated since the post last moved in the feed (E3).
            # Seeded here so the joiner never meets its absence — creation IS
            # a surfacing, so the count starts at zero.
            "images_since_surface": 0,
        },
    )
    session.add(post)
    await session.flush()

    await session.execute(
        update(Post)
        .where(Post.id.in_([m.id for m in members]))
        .values(absorbed_by_post_id=post.id)
    )

    await _link_member_images(
        session, post_id=post.id, member_ids=[m.id for m in members],
        source_id=source.id,
    )
    return post


async def _link_member_images(
    session: AsyncSession, *, post_id: int, member_ids: list[int], source_id: int | None,
) -> int:
    """Attach every member's images to the synthetic post. Returns how many.

    The feed and detail views already union provenance with `primary_post_id`
    (post_feed_service._thumbnails_for), so this alone makes the drop's images
    show up under the post FC wrote — no second render path.

    `primary_post_id` is deliberately NOT rewritten: the message post remains
    the image's true origin, and the synthetic post is an ADDITIONAL claim on
    it, which is what keeps the grouping reversible.

    Shared by creation and by the E3 joiner rather than written twice, because
    the two would otherwise be free to drift on exactly the detail (which post
    owns the image) that makes a grouping reversible.
    """
    if not member_ids:
        return 0
    image_rows = (await session.execute(
        select(ImageRecord.id).where(ImageRecord.primary_post_id.in_(member_ids))
    )).scalars().all()
    if not image_rows:
        return 0
    await session.execute(
        pg_insert(ImageProvenance)
        .values([
            {"image_record_id": iid, "post_id": post_id, "source_id": source_id}
            for iid in image_rows
        ])
        # (image, post) is unique. This is a BACKSTOP against a re-run that
        # raced itself, not the correctness argument: callers only ever pass
        # members that were unabsorbed a moment ago, so a conflict here means
        # concurrency, not a logic error.
        .on_conflict_do_nothing(constraint="uq_image_provenance_image_post")
    )
    return len(image_rows)


async def group_source(
    session: AsyncSession,
    source: Source,
    *,
    max_distance: float,
    window_minutes: float,
    now: datetime | None = None,
) -> int:
    """Group one Discord source's ungrouped messages. Returns posts created."""
    window = timedelta(minutes=window_minutes)
    now = now or datetime.now(UTC)
    # Leave the most recent window alone: a drop that is still arriving would
    # otherwise be cut in half by whichever sweep happened to land mid-drop,
    # and the second half would become a separate post claiming to be its own
    # drop. Waiting one window costs nothing (the sweep re-runs) and is the E2
    # side of "keep the grouping open"; E3 handles the harder case where a
    # matching drop resumes after the gap has already passed.
    rows = (await session.execute(
        _candidate_stmt(source.id, not_after=now - window)
        .limit(MAX_CANDIDATES_PER_SOURCE)
    )).all()
    if not rows:
        return 0

    groups = build_groups(
        [(pid, occurred, emb) for pid, occurred, emb in rows],
        max_distance=max_distance, window=window,
    )
    if len(rows) == MAX_CANDIDATES_PER_SOURCE and len(groups) > 1:
        # The cap may have fallen INSIDE the last drop, and synthesising a
        # truncated group would publish a post that claims to be the whole drop
        # while the rest of it sits one row past the limit. Leave it for the
        # next run, which starts from the same place and sees the remainder.
        # Guarded on len > 1 so a single oversized group is not dropped
        # forever — it would make no progress at all.
        groups = groups[:-1]

    created = 0
    for group in groups:
        post = await _synthesize(
            session, source=source, group=group,
            max_distance=max_distance, window_minutes=window_minutes,
        )
        if post is not None:
            created += 1
    return created


# ---------------------------------------------------------------------------
# E3: an open grouping — a later drop joins its post and updates it.
# ---------------------------------------------------------------------------
#
# A synthetic post is not sealed at creation. A creator who adds two more
# variants the next day extends the existing post rather than starting a new
# one, and its body grows with the new messages. That is what makes chat
# capture read as content TRICKLING IN rather than as a stream of separate
# arrivals.
#
# Three hard problems, each answered deliberately below: bridging (a candidate
# near two groups), re-surfacing without thrashing the feed, and groups that
# stay open forever.

# How much closer the nearest group must be than the runner-up before a
# candidate is assigned to it at all.
#
# NOT a setting, deliberately. It is not a quality dial the operator would tune
# toward a better feed — it expresses "these two are too close to call", and
# exposing it would invite turning it to zero, which is precisely the silent
# arbitrary choice it exists to prevent. When a candidate is genuinely between
# two groups the recoverable answer is to leave it out and let it start its
# own; the unrecoverable one is to merge, because a merge rewrites history —
# two posts the operator may already have seen become one, and anything
# pointing at the absorbed post dangles.
AMBIGUITY_MARGIN = 0.02


def should_resurface(
    *,
    images_since_surface: int,
    last_surface_at: datetime,
    now: datetime,
    min_images: int,
    cooldown: timedelta,
) -> bool:
    """Has this group grown enough, and waited long enough, to move in the feed?

    An updated post SHOULD be visible — that is the point of keeping it open —
    but a group gaining one image a day must not sit permanently at the top.
    Both conditions have to hold: enough new images that the update is worth an
    interruption, and enough time since the last one that a steady drip cannot
    chain bumps together.
    """
    if images_since_surface < min_images:
        return False
    return now - last_surface_at >= cooldown


async def _group_seed(session: AsyncSession, post_id: int) -> list[float] | None:
    """The embedding a group is measured against — its FIRST member's image.

    Derived rather than stored, and derived by the same definition `build_groups`
    used (earliest member, lowest-id embedded image). Storing it at creation
    would have meant a backfill for groups already written and two definitions
    free to disagree; this way there is one.
    """
    sort_key = func.coalesce(Post.post_date, Post.downloaded_at)
    return (await session.execute(
        select(ImageRecord.siglip_embedding)
        .join(Post, ImageRecord.primary_post_id == Post.id)
        .where(
            Post.absorbed_by_post_id == post_id,
            ImageRecord.siglip_embedding.is_not(None),
        )
        .order_by(sort_key, Post.id, ImageRecord.id)
        .limit(1)
    )).scalar_one_or_none()


async def open_groups(
    session: AsyncSession, source_id: int, *, now: datetime, close_after: timedelta,
) -> list[tuple[Post, list[float]]]:
    """This source's synthetic posts that are still accepting members.

    Openness is DERIVED, not stored: a group is open if it grew — or started —
    within `close_after`. A group left open forever would eventually absorb
    something it shouldn't, because artists reuse characters for years; a
    stored `closed_at` would need a sweep to set it and a repair path to ever
    change the policy. This way the policy IS the query.
    """
    cutoff = now - close_after
    posts = (await session.execute(
        select(Post).where(
            Post.source_id == source_id,
            Post.synthesized_by == DROP_GROUPER,
            # A synthetic post that was itself absorbed is not a thing today
            # (nothing absorbs one), but joining into one would nest groups.
            Post.absorbed_by_post_id.is_(None),
            func.coalesce(Post.last_grew_at, Post.post_date, Post.downloaded_at) >= cutoff,
        )
    )).scalars().all()

    out: list[tuple[Post, list[float]]] = []
    for post in posts:
        seed = await _group_seed(session, post.id)
        if seed is not None:
            out.append((post, seed))
    return out


def assign_to_group(
    embedding: list[float],
    groups: list[tuple[Post, list[float]]],
    *,
    max_distance: float,
) -> Post | None:
    """Pick the one group this image belongs to, or None to leave it alone.

    Returns None in two cases that mean different things and are deliberately
    treated the same: nothing is close enough (so E2 will start a new group
    from it), or two groups are BOTH close and too near each other to choose
    between (so E2 will start a new group from it). The second is the bridging
    case, and letting it start its own post is the recoverable failure —
    merging two existing posts is not.
    """
    # key= on the distance ALONE. Sorting bare tuples falls through to the
    # second element when two distances tie, and `Post` has no ordering — so a
    # perfectly symmetric bridge (the exact case this function exists for)
    # would raise TypeError instead of declining to choose.
    scored = sorted(
        ((cosine_distance(seed, embedding), post) for post, seed in groups),
        key=lambda pair: pair[0],
    )
    within = [(d, p) for d, p in scored if d <= max_distance]
    if not within:
        return None
    if len(within) >= 2 and (within[1][0] - within[0][0]) < AMBIGUITY_MARGIN:
        return None
    return within[0][1]


async def _absorb_into(
    session: AsyncSession,
    *,
    group: Post,
    member_ids: list[int],
    source_id: int | None,
    now: datetime,
    min_images: int,
    cooldown: timedelta,
) -> int:
    """Extend an existing synthetic post with new members. Returns images added."""
    members = (await session.execute(
        select(Post)
        .where(Post.id.in_(member_ids))
        .order_by(func.coalesce(Post.post_date, Post.downloaded_at), Post.id)
    )).scalars().all()
    if not members:
        return 0

    added_images = await _link_member_images(
        session, post_id=group.id, member_ids=[m.id for m in members],
        source_id=source_id,
    )
    await session.execute(
        update(Post)
        .where(Post.id.in_([m.id for m in members]))
        .values(absorbed_by_post_id=group.id)
    )

    # The new messages' text joins the body, in arrival order, exactly as at
    # creation — the post's body is the drop's text and the drop just grew.
    new_text = "\n\n".join(
        m.description.strip() for m in members if m.description and m.description.strip()
    )
    if new_text:
        group.description = f"{group.description}\n\n{new_text}" if group.description else new_text

    details = dict(group.synthesis_details or {})
    existing_ids = list(details.get("member_post_ids") or [])
    details["member_post_ids"] = existing_ids + [
        m.id for m in members if m.id not in existing_ids
    ]
    details["message_count"] = len(details["member_post_ids"])
    since = int(details.get("images_since_surface") or 0) + added_images
    details["last_grew_at"] = now.isoformat()

    # The feed position moves only when the anti-thrash rule fires. Measured
    # from the last time the post actually MOVED (resurfaced_at), falling back
    # to when the drop started — creation is itself a surfacing.
    last_surface = group.resurfaced_at or group.post_date or group.downloaded_at
    if should_resurface(
        images_since_surface=since, last_surface_at=last_surface, now=now,
        min_images=min_images, cooldown=cooldown,
    ):
        group.resurfaced_at = now
        since = 0
    details["images_since_surface"] = since

    group.synthesis_details = details
    group.last_grew_at = now
    return added_images


async def join_open_groups(
    session: AsyncSession,
    source: Source,
    *,
    max_distance: float,
    window_minutes: float,
    close_after_hours: float,
    resurface_min_images: int,
    resurface_cooldown_hours: float,
    now: datetime | None = None,
) -> int:
    """Offer this source's ungrouped messages to its open groups.

    Runs BEFORE `group_source` in the sweep: a message that belongs to an
    existing drop must join it rather than found a rival post, and whichever
    runs first wins that message.
    """
    now = now or datetime.now(UTC)
    window = timedelta(minutes=window_minutes)
    groups = await open_groups(
        session, source.id, now=now, close_after=timedelta(hours=close_after_hours),
    )
    if not groups:
        return 0

    # Same quarantine as E2: a drop still arriving is left for the next run.
    rows = (await session.execute(
        _candidate_stmt(source.id, not_after=now - window)
        .limit(MAX_CANDIDATES_PER_SOURCE)
    )).all()
    if not rows:
        return 0

    claimed: dict[int, list[int]] = {}
    for post_id, _occurred_at, embedding in rows:
        target = assign_to_group(embedding, groups, max_distance=max_distance)
        if target is not None:
            claimed.setdefault(target.id, []).append(post_id)

    by_id = {post.id: post for post, _seed in groups}
    joined = 0
    for group_id, member_ids in claimed.items():
        joined += await _absorb_into(
            session, group=by_id[group_id], member_ids=member_ids,
            source_id=source.id, now=now,
            min_images=resurface_min_images,
            cooldown=timedelta(hours=resurface_cooldown_hours),
        )
    return joined


async def sweep(session: AsyncSession, *, now: datetime | None = None) -> dict:
    """Group every enabled Discord source. No-op when the switch is off.

    Two passes per source, and the ORDER is load-bearing: offer new messages to
    the groups that are still open (E3) BEFORE founding new ones (E2). Whichever
    runs first claims a message, and a variant that belongs to yesterday's drop
    must extend that post rather than found a rival to it.
    """
    settings = await MLSettings.load(session)
    if not settings.discord_grouping_enabled:
        return {
            "enabled": False, "sources": 0, "posts_created": 0, "images_joined": 0,
        }

    sources = (await session.execute(
        select(Source).where(
            Source.platform == PLATFORM,
            Source.enabled.is_(True),
        )
    )).scalars().all()

    max_distance = float(settings.discord_group_max_distance)
    window_minutes = float(settings.discord_group_window_minutes)

    created = 0
    joined = 0
    for source in sources:
        joined += await join_open_groups(
            session, source,
            max_distance=max_distance,
            window_minutes=window_minutes,
            close_after_hours=float(settings.discord_group_close_after_hours),
            resurface_min_images=int(settings.discord_group_resurface_min_images),
            resurface_cooldown_hours=float(
                settings.discord_group_resurface_cooldown_hours
            ),
            now=now,
        )
        created += await group_source(
            session, source,
            max_distance=max_distance,
            window_minutes=window_minutes,
            now=now,
        )
    log.info(
        "discord drop grouping: %d source(s), %d synthetic post(s) created, "
        "%d image(s) joined to open groups",
        len(sources), created, joined,
    )
    return {
        "enabled": True, "sources": len(sources),
        "posts_created": created, "images_joined": joined,
    }