LLeMbas/tests/test_crowd_chain.py

"""One user turn, several speakers, chained.

`_advance_crowd` is the shell around the pure scheduler, so what is asserted here
is the part the scheduler cannot see: which rows exist, when, and how many. The
producer is replaced, so nothing here talks to an endpoint — what matters is the
chat each speaker is handed and the invariant that holds between them.

**Exactly one incomplete assistant row at every observation.** That is the whole
reason this shape was chosen over one generation writing many bubbles: it is what
`_reply_in_flight`, `_too_many_replies`, `wake.lock_for` and the superseded guards
in `_persist`/`_drain` all already rely on, and its symptom when broken is a Stop
button pointing at whichever bubble comes first in the document.
"""

from __future__ import annotations

import pytest
from sqlalchemy import select

from lembas.db.models import (
    ROLE_ASSISTANT,
    ROLE_USER,
    Chat,
    Connection,
    CrowdMember,
    Message,
    Model,
    User,
)
from lembas.services import chat as chat_service
from lembas.services import crowd as crowd_service
from lembas.services import generation as generation_service
from lembas.services import settings_store
from lembas.services.crypto import encrypt

MEMBERS = ("second-model", "third-model")


@pytest.fixture(autouse=True)
def empty_registry():
    yield
    generation_service._RUNNING.clear()
    generation_service._TASKS.clear()


@pytest.fixture(autouse=True)
def crowd_on(db, registered):
    settings_store.update(db, {"enabled": True}, key=settings_store.CROWD)
    connection = Connection(
        name="Test", base_url="http://127.0.0.1:1", api_key_encrypted=encrypt("")
    )
    db.add(connection)
    db.commit()
    for index, name in enumerate(("main-model", *MEMBERS)):
        db.add(
            Model(
                connection_id=connection.id,
                model_id=name,
                display_name=name,
                position=index,
                capabilities_json={"tools": True},
            )
        )
    db.commit()


@pytest.fixture
def started(monkeypatch):
    """Every chat id `ensure` was asked to start a reply in, in order."""
    calls: list[tuple[str, str]] = []

    def _fake_ensure(chat_id, message_id):
        calls.append((chat_id, message_id))

    monkeypatch.setattr(generation_service, "ensure", _fake_ensure)
    return calls


def _user(db) -> User:
    return db.scalars(select(User).order_by(User.created_at)).first()


def _crowd_chat(db, members=MEMBERS) -> Chat:
    connection = db.scalars(select(Connection)).first()
    chat = Chat(
        user_id=_user(db).id, title="t", model_id="main-model", connection_id=connection.id
    )
    db.add(chat)
    db.commit()
    for index, name in enumerate(members):
        db.add(CrowdMember(chat_id=chat.id, model_id=name, position=index))
    db.commit()
    return chat


def _opening_reply(db, chat) -> Message:
    """The main model's first answer: a user turn and a finished assistant one."""
    chat_service.create_message(db, chat, ROLE_USER, "What should we do?")
    reply = chat_service.create_message(
        db, chat, ROLE_ASSISTANT, "Rewrite it.", model_id=chat.model_id
    )
    return reply


def _advance(db, chat, message, **kwargs) -> bool:
    generation = generation_service.Generation(chat_id=chat.id, message_id=message.id)
    for key, value in kwargs.items():
        setattr(generation, key, value)
    generation_service._RUNNING[message.id] = generation
    try:
        return generation_service._advance_crowd(generation)
    finally:
        generation_service._RUNNING.pop(message.id, None)


def _incomplete(db, chat) -> list[Message]:
    return list(
        db.scalars(
            select(Message).where(
                Message.chat_id == chat.id, Message.complete.is_(False)
            )
        )
    )


def _run_round(db, chat, started, *, answers: int = 12) -> list[Message]:
    """Walk a whole round by finishing each speaker as it is created."""
    order: list[Message] = []
    message = _opening_reply(db, chat)
    for _ in range(answers):
        assert len(_incomplete(db, chat)) == 0, "a row was left incomplete"
        if not _advance(db, chat, message):
            break
        db.expire_all()
        fresh = _incomplete(db, chat)
        assert len(fresh) == 1, f"{len(fresh)} replies in flight at once"
        message = fresh[0]
        order.append(message)
        message.content = "Something."
        message.complete = True
        db.commit()
    return order


# --- The chain ----------------------------------------------------------------
def test_a_whole_round_speaks_in_order(db, started):
    chat = _crowd_chat(db)
    order = _run_round(db, chat, started)

    assert [m.model_id for m in order] == [
        "second-model",   # out
        "third-model",    # out
        "second-model",   # back
        "main-model",     # close
    ]


def test_each_speaker_is_started_through_ensure(db, started):
    chat = _crowd_chat(db)
    order = _run_round(db, chat, started)
    assert [message_id for _chat_id, message_id in started] == [m.id for m in order]


def test_the_round_is_recorded_on_every_row(db, started):
    chat = _crowd_chat(db)
    order = _run_round(db, chat, started)

    phases = [crowd_service.state_of(m).phase for m in order]
    assert phases == [
        crowd_service.PHASE_OUT,
        crowd_service.PHASE_OUT,
        crowd_service.PHASE_BACK,
        crowd_service.PHASE_CLOSE,
    ]
    # And they all belong to the same question.
    anchors = {crowd_service.state_of(m).turn for m in order}
    assert len(anchors) == 1


def test_each_speaker_carries_its_own_connection(db, started):
    """So `speaker_for` resolves the pair rather than guessing at the id."""
    chat = _crowd_chat(db)
    order = _run_round(db, chat, started)
    for message in order:
        assert chat_service.speaker_for(db, chat, message).model_id == message.model_id


def test_a_chat_with_no_crowd_is_not_chained(db, started):
    chat = _crowd_chat(db, members=())
    message = _opening_reply(db, chat)
    assert _advance(db, chat, message) is False
    assert started == []


def test_the_feature_switch_holds_the_whole_thing(db, started):
    settings_store.update(db, {"enabled": False}, key=settings_store.CROWD)
    chat = _crowd_chat(db)
    message = _opening_reply(db, chat)
    assert _advance(db, chat, message) is False


def test_the_member_cap_trims_the_crowd(db, started):
    settings_store.update(db, {"max_models": 1}, key=settings_store.CROWD)
    chat = _crowd_chat(db)
    order = _run_round(db, chat, started)
    # One member: out to it, then straight back to the main model.
    assert [m.model_id for m in order] == ["second-model", "main-model"]


def test_a_member_whose_model_has_gone_is_skipped(db, started):
    """Membership is text with no foreign key, so a model that disappears upstream
    leaves a row behind. Skipping it is the point -- a cascade would have deleted
    the crowd out of every chat on the next Test & refresh."""
    chat = _crowd_chat(db)
    gone = db.scalar(select(Model).where(Model.model_id == "third-model"))
    db.delete(gone)
    db.commit()

    order = _run_round(db, chat, started)

    assert "third-model" not in [m.model_id for m in order]
    assert [m.model_id for m in order] == ["second-model", "main-model"]
    # The row is still there, so a screen can say it was skipped.
    assert any(row.model_id == "third-model" for row in db.get(Chat, chat.id).crowd)


def test_a_disabled_model_is_skipped_too(db, started):
    chat = _crowd_chat(db)
    off = db.scalar(select(Model).where(Model.model_id == "third-model"))
    off.enabled = False
    db.commit()
    assert "third-model" not in [m.model_id for m in _run_round(db, chat, started)]


# --- The refusals -------------------------------------------------------------
def test_stop_ends_the_round(db, started):
    """Not just the speaker writing at the time. `_drain` refuses after a stop for
    the same reason: somebody asked for it to end."""
    chat = _crowd_chat(db)
    message = _opening_reply(db, chat)
    assert _advance(db, chat, message, stopped=True) is False
    assert started == []


def test_a_superseded_generation_advances_nothing(db, started):
    """The guard `_persist` and `_drain` both carry."""
    chat = _crowd_chat(db)
    message = _opening_reply(db, chat)
    other = generation_service.Generation(chat_id=chat.id, message_id=message.id)
    generation_service._RUNNING[message.id] = other
    try:
        mine = generation_service.Generation(chat_id=chat.id, message_id=message.id)
        assert generation_service._advance_crowd(mine) is False
    finally:
        generation_service._RUNNING.pop(message.id, None)
    assert started == []


def test_regenerating_a_speaker_does_not_fork_the_round(db, started):
    """`restart` re-runs `_run`, whose `finally` advances the crowd again -- and the
    speakers after it already exist. Without the newest-message guard, regenerating
    member 2 creates a second member 3 and two chains race down one turn."""
    chat = _crowd_chat(db)
    order = _run_round(db, chat, started)
    before = len(list(db.scalars(select(Message).where(Message.chat_id == chat.id))))

    # The second speaker is regenerated: it is no longer the newest row.
    assert _advance(db, chat, order[0]) is False

    db.expire_all()
    after = len(list(db.scalars(select(Message).where(Message.chat_id == chat.id))))
    assert after == before


def test_one_speaker_failing_is_skipped(db, started):
    chat = _crowd_chat(db)
    message = _opening_reply(db, chat)
    assert _advance(db, chat, message) is True
    db.expire_all()
    second = _incomplete(db, chat)[0]
    second.complete = True
    second.error = "the endpoint fell over"
    db.commit()

    assert _advance(db, chat, second, error="the endpoint fell over") is True
    db.expire_all()
    assert _incomplete(db, chat)[0].model_id == "third-model"


def test_two_failures_in_a_row_abandon_the_round(db, started):
    chat = _crowd_chat(db)
    message = _opening_reply(db, chat)
    _advance(db, chat, message)
    db.expire_all()
    second = _incomplete(db, chat)[0]
    second.complete = True
    db.commit()
    _advance(db, chat, second, error="down")
    db.expire_all()
    third = _incomplete(db, chat)[0]
    third.complete = True
    db.commit()

    assert _advance(db, chat, third, error="down") is False

    db.expire_all()
    assert crowd_service.state_of(db.get(Message, third.id)).stopped == (
        crowd_service.STOPPED_ERRORS
    )


def test_why_a_round_stopped_is_written_where_it_stopped(db, started):
    """So the transcript can say a round ended rather than simply ending."""
    settings_store.update(db, {"max_rounds": 1}, key=settings_store.CROWD)
    chat = _crowd_chat(db)
    order = _run_round(db, chat, started)
    closing = order[-1]

    assert _advance(db, chat, closing, crowd_again=True) is False

    db.expire_all()
    assert crowd_service.state_of(db.get(Message, closing.id)).stopped == (
        crowd_service.STOPPED_ROUNDS
    )


# --- Going round again --------------------------------------------------------
def test_the_main_model_can_send_them_round_again(db, started):
    settings_store.update(db, {"max_rounds": 2}, key=settings_store.CROWD)
    chat = _crowd_chat(db)
    order = _run_round(db, chat, started)
    closing = order[-1]

    assert _advance(db, chat, closing, crowd_again=True) is True

    db.expire_all()
    fresh = _incomplete(db, chat)[0]
    state = crowd_service.state_of(fresh)
    assert state.round == 2
    assert state.phase == crowd_service.PHASE_OUT
    assert fresh.model_id == "second-model"


def test_without_asking_the_round_is_over(db, started):
    chat = _crowd_chat(db)
    order = _run_round(db, chat, started)
    assert _advance(db, chat, order[-1], crowd_again=False) is False