LLeMbas/src/lembas/services/crowd.py

"""Several models answering one turn, in order, then again in reverse.

The shape the owner asked for: the chat's own model answers, then each other
member in order; then the order runs **backwards**, each member asked whether it
disagrees with anything; and it ends at the main model, which decides whether to
go round again or stop.

## Why N chained replies and not one clever one

One `Generation` per speaker, one `Message` per speaker, chained where `_drain`
already chains a queued turn. That is not the cheapest shape, it is the only one
in which every existing invariant keeps holding for the reason it already holds:

* `Generation` is **one reply's** state and `_follow` streams **per message**,
  keyed on `generation.message_id`. One generation cannot stream into nine
  bubbles without a second streaming protocol, and `ensure(chat_id, message_id)`
  would have no answer to "which of the nine am I" after a restart.
* Exactly one incomplete assistant row exists at any moment, so
  `_reply_in_flight` needs no teaching and the composer queues for the whole
  round.
* Each speaker gets its own `steps_json`, `usage_json` and `model_id`, so the
  avatar, the metrics chip and the regenerate button are per speaker with no new
  rendering.

A subagent per speaker was rejected outright: a helper is handed a *serialisation*
of the conversation, its answer comes back as a tool result, and tool results are
never replayed -- so speaker 3 could not see speaker 2, which is the entire point
of a crowd. That feature already exists and is called `ask_friend`.

## Where the round lives

On the **message row**, in `Message.crowd_json`, and not on the chat. "The row is
the authority, not the registry" is the rule the reload story was won with, and
round state on the chat reintroduces exactly the split it was won against: a
restart between speakers, or a rewind that deletes the rows, would leave
chat-level state describing turns that no longer exist -- which is the problem
`Chat.compacted_through_id` already documents.

`Message.parent_id` is **not** used for grouping. It is reserved for conversation
branching and says so in its own comment.

## The scheduler is a pure function

`next_turn` takes numbers and returns numbers. Every refusal -- out of rounds, out
of time, nobody to ask, not the newest message -- is therefore testable without an
endpoint, which matters because the refusals are the interesting half.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, replace
from datetime import UTC, datetime
from typing import Any

from sqlalchemy import select
from sqlalchemy.orm import Session as DBSession

from lembas.db.models import Chat, Message

log = logging.getLogger(__name__)

# The forward pass: everybody answers in order.
PHASE_OUT = "out"
# The way back: each member is asked whether it disagrees, in reverse order,
# stopping one short of the main model.
PHASE_BACK = "back"
# The main model's last word, where it decides whether to go round again.
PHASE_CLOSE = "close"

PHASES = (PHASE_OUT, PHASE_BACK, PHASE_CLOSE)

# Why a round ended, when it ended for a reason rather than by finishing.
STOPPED_ROUNDS = "rounds"
STOPPED_TIME = "time"
STOPPED_ERRORS = "errors"

# How many speaker errors in a row end the round. One is skipped: the commonest
# failure in a crowd is not a dead endpoint but a small member's context window
# overflowing on a transcript several models have been writing into, and killing
# the round at whichever member is smallest is the wrong answer. Two in a row is
# an endpoint that has actually gone, which is what `_drain`'s refusal protects
# against and is worth keeping.
MAX_CONSECUTIVE_ERRORS = 2


@dataclass(frozen=True)
class Turn:
    """Where one crowd round has got to, as it is stored on a message."""

    turn: str
    round: int
    phase: str
    index: int
    of: int
    started_at: str
    errors: int = 0
    stopped: str = ""

    def as_json(self) -> dict[str, Any]:
        return {
            "turn": self.turn,
            "round": self.round,
            "phase": self.phase,
            "index": self.index,
            "of": self.of,
            "started_at": self.started_at,
            "errors": self.errors,
            "stopped": self.stopped,
        }

    @property
    def is_main(self) -> bool:
        return self.index == 0


def state_of(message: Message | None) -> Turn | None:
    """The round state on a message, or None if it is not part of one."""
    raw = getattr(message, "crowd_json", None) or None
    if not raw or not isinstance(raw, dict):
        return None
    try:
        return Turn(
            turn=str(raw.get("turn") or ""),
            round=int(raw.get("round") or 1),
            phase=str(raw.get("phase") or PHASE_OUT),
            index=int(raw.get("index") or 0),
            of=int(raw.get("of") or 1),
            started_at=str(raw.get("started_at") or ""),
            errors=int(raw.get("errors") or 0),
            stopped=str(raw.get("stopped") or ""),
        )
    except (TypeError, ValueError):  # pragma: no cover - a hand-edited row
        return None


def is_opening(state: Turn | None) -> bool:
    """Whether this state is the main model's opening reply.

    `phase=out, index=0` is **display state and never scheduling state**. The
    opening reply is not started by the crowd -- the composer starts it, exactly
    as it starts every other reply, and a round only begins when it *finishes*.
    Stamping it afterwards is what lets the transcript say `1 of 3` on the bubble
    that opened the round; before that it was the one contribution with no chip,
    so a two-model round read as an ordinary reply followed by a crowd.

    Everything that asks "is a round already in progress?" has to skip it, or the
    stamp changes behaviour it was never meant to touch -- see `scheduling_state`.
    """
    return state is not None and state.phase == PHASE_OUT and state.index == 0


def scheduling_state(message: Message | None) -> Turn | None:
    """The round state the scheduler should act on: `state_of`, minus the opening.

    Two things would break if the opening stamp were fed to `next_turn` as real
    state, and both are silent:

    * **`started_at` would be inherited on a regenerate.** Regenerating the
      opening reply an hour later would hand `next_turn` an hour-old clock and the
      round would stop with "out of time" before anybody spoke.
    * **The once-per-turn gates key off "no state at all"** -- compaction, the
      title, the unread push. A stamped opening reads as a later speaker, and each
      of them would be skipped for the turn that is supposed to have them.

    So the stamp is written where the transcript reads it and nowhere else.
    """
    state = state_of(message)
    return None if is_opening(state) else state


def now_stamp() -> str:
    return datetime.now(UTC).isoformat()


def elapsed(started_at: str) -> float:
    """Seconds since a round began, or 0.0 if the stamp is unreadable.

    Unreadable reads as "no time has passed" rather than as "out of time": a
    round abandoned because of a bad timestamp would be a feature failing for a
    reason nobody could see.
    """
    try:
        began = datetime.fromisoformat(started_at)
    except (TypeError, ValueError):
        return 0.0
    if began.tzinfo is None:
        began = began.replace(tzinfo=UTC)
    return max(0.0, (datetime.now(UTC) - began).total_seconds())


def next_turn(
    *,
    speakers: int,
    state: Turn | None,
    turn_id: str,
    again: bool = False,
    errored: bool = False,
    max_rounds: int = 2,
    wall_seconds: int = 900,
) -> Turn | None:
    """Who speaks next, or None when the round is over.

    Pure: numbers in, numbers out, no session and no clock beyond the stamp it is
    handed. `speakers` counts the main model as one of them.

    `state=None` means the reply that has just finished was the ordinary first
    one, started by the composer as it always is -- so this is where a round
    begins rather than continues.
    """
    if speakers < 2:
        return None

    if state is None:
        return Turn(
            turn=turn_id,
            round=1,
            phase=PHASE_OUT,
            index=1,
            of=speakers,
            started_at=now_stamp(),
        )

    # Errors are counted consecutively, so one member timing out is skipped and
    # an endpoint that has gone ends the round.
    errors = state.errors + 1 if errored else 0
    if errors >= MAX_CONSECUTIVE_ERRORS:
        return replace(state, stopped=STOPPED_ERRORS)

    if wall_seconds and elapsed(state.started_at) >= wall_seconds:
        return replace(state, errors=errors, stopped=STOPPED_TIME)

    carry = {
        "turn": state.turn,
        "of": speakers,
        "started_at": state.started_at,
        "errors": errors,
    }

    if state.phase == PHASE_OUT:
        if state.index + 1 <= speakers - 1:
            return Turn(round=state.round, phase=PHASE_OUT, index=state.index + 1, **carry)
        # The forward pass is done. The way back starts one short of the speaker
        # that has just finished -- asking it whether it disagrees with itself is
        # a round spent on nothing.
        if speakers - 2 >= 1:
            return Turn(round=state.round, phase=PHASE_BACK, index=speakers - 2, **carry)
        return Turn(round=state.round, phase=PHASE_CLOSE, index=0, **carry)

    if state.phase == PHASE_BACK:
        if state.index - 1 >= 1:
            return Turn(round=state.round, phase=PHASE_BACK, index=state.index - 1, **carry)
        return Turn(round=state.round, phase=PHASE_CLOSE, index=0, **carry)

    # The main model has had its last word. Another round only if it asked for
    # one *and* there is one left.
    if not again:
        return None
    if state.round + 1 > max_rounds:
        return replace(state, errors=errors, stopped=STOPPED_ROUNDS)
    return Turn(round=state.round + 1, phase=PHASE_OUT, index=1, **carry)


# --- Resolving the membership --------------------------------------------------
def member_speakers(db: DBSession, chat: Chat, user=None) -> list:
    """Every member that can actually be reached, in order, main model first.

    Filtered through `permissions.models_visible_to` by way of
    `chat_service.roster_models`, so a member whose access has been revoked, whose
    model has been disabled, or whose row has gone is skipped rather than
    attempted -- and the skip is visible in the transcript rather than silent.

    Deduplicated against the main model: adding the chat's own model to the crowd
    would have it answer twice in a row, which is not what anybody meant by it.
    """
    from lembas.services import chat as chat_service

    reachable = {
        model.model_id: model for model in chat_service.roster_models(db, user, exclude="")
    }
    speakers = [chat_service.Speaker(chat.model_id, chat.connection_id)]
    seen = {chat.model_id}
    for member in sorted(chat.crowd, key=lambda row: (row.position, row.model_id)):
        if member.model_id in seen or member.model_id not in reachable:
            continue
        seen.add(member.model_id)
        speakers.append(chat_service.Speaker(member.model_id, member.connection_id))
    return speakers


def unreachable_members(db: DBSession, chat: Chat, user=None) -> list[str]:
    """Members that will be skipped, so a screen can say so rather than lie."""
    from lembas.services import chat as chat_service

    reachable = {
        model.model_id for model in chat_service.roster_models(db, user, exclude="")
    }
    return [
        member.model_id
        for member in chat.crowd
        if member.model_id not in reachable or member.model_id == chat.model_id
    ]


def is_newest(db: DBSession, message: Message) -> bool:
    """Whether this is the last message in its chat.

    The guard that stops a regenerate from forking the round. `restart` re-runs
    `_run`, whose `finally` advances the crowd again -- and speakers further down
    already exist, so without this, regenerating member 2 creates a second member
    3 and two chains race down one turn. `_drain` never needed it, because a
    queued row only ever exists *forward* of the reply.
    """
    latest = db.scalars(
        select(Message)
        .where(Message.chat_id == message.chat_id)
        .order_by(Message.created_at.desc(), Message.id.desc())
        .limit(1)
    ).first()
    return latest is not None and latest.id == message.id


# --- Asking for another round ---------------------------------------------------
async def _run_crowd_again(context, args: dict[str, Any]):
    """Record that the main model wants the crowd to go round again.

    Written onto the running `Generation` rather than onto the row, because it is
    a fact about *this* reply and dies with it -- and onto a field rather than
    parsed back out of the prose, for the reason `plan_json` exists: a sentinel
    phrase in an answer is a decision nobody can see and a wording nobody can
    change.

    Offered only on the main model's closing turn and only while a round is left,
    so a call arriving anywhere else is a call that was never on the table.
    """
    from lembas.services import generation as generation_service
    from lembas.services.tools import ToolOutcome

    reason = str(args.get("focus") or "").strip()
    running = generation_service.running_for(context.chat_id) if context.chat_id else None
    if running is None:
        return ToolOutcome(
            "There is no round to continue.",
            {"name": "crowd_again", "status": "error", "error": "no round"},
        )

    running.crowd_again = True
    return ToolOutcome(
        "The others will answer again."
        + (f" You have asked them to focus on: {reason}" if reason else "")
        + " Finish your answer now: what you write is what the person reads for "
        "this round.",
        {
            "name": "crowd_again",
            "status": "ok",
            "query": reason[:160],
            "detail": "another round",
        },
    )


def tool_defs() -> list:
    """The one tool, offered only to the closing speaker of a crowd round."""
    from lembas.services.tools import FAMILY_CROWD, RISK_READ, ToolDef

    return [
        ToolDef(
            name="crowd_again",
            family=FAMILY_CROWD,
            description=(
                "Send the other models round again, because the disagreement is "
                "real and another pass would settle it. Say what they should focus "
                "on. Use it sparingly: every round costs the person another wait, "
                "and a crowd asked to go round because the discussion was "
                "interesting will keep finding things to discuss. If the answers "
                "have converged, or the disagreement is a matter of taste, or "
                "nobody has said anything new on the way back, do not call this -- "
                "write the answer instead."
            ),
            parameters={
                "type": "object",
                "properties": {
                    "focus": {
                        "type": "string",
                        "description": (
                            "What the next round should settle, in one sentence."
                        ),
                    }
                },
                "required": [],
            },
            run=_run_crowd_again,
            # It changes nothing in the world; what it costs is more replies, and
            # that is bounded by `crowd.max_rounds` rather than by an approval.
            risk=RISK_READ,
        ),
    ]


__all__ = [
    "MAX_CONSECUTIVE_ERRORS",
    "PHASES",
    "PHASE_BACK",
    "PHASE_CLOSE",
    "PHASE_OUT",
    "STOPPED_ERRORS",
    "STOPPED_ROUNDS",
    "STOPPED_TIME",
    "Turn",
    "elapsed",
    "is_newest",
    "is_opening",
    "member_speakers",
    "next_turn",
    "now_stamp",
    "scheduling_state",
    "state_of",
    "tool_defs",
    "unreachable_members",
]