LLeMbas/tests/test_roster.py

"""The list of other models a model is given, and what decides it is there.

The roster is one `{{variable}}` and one fragment, so the interesting assertions
are about *absence*: it is missing on a single-model instance, missing for a
model that may not ask anyone anything, and missing a model this account cannot
reach. A list that is merely wrong would be bad; a list naming something the
reader has no access to is a leak and a dead end at once, because asking it
anything is refused by the same check.
"""

from __future__ import annotations

import pytest
from sqlalchemy import select

from lembas.db.models import ROLE_USER, Chat, Connection, Group, Model, User
from lembas.services import chat as chat_service
from lembas.services import harness as harness_service
from lembas.services import settings_store
from lembas.services.crypto import encrypt


@pytest.fixture(autouse=True)
def three_models(db, registered):
    settings_store.update(db, {"enabled": True}, key=settings_store.SUBAGENTS)
    settings_store.update(db, {"default_permissions": {"tools.friend": True}})
    connection = Connection(
        name="Test", base_url="http://127.0.0.1:1", api_key_encrypted=encrypt("")
    )
    db.add(connection)
    db.commit()
    rows = [
        ("test-model", "The asker", "", ""),
        ("big-model", "Big", "Long reasoning problems", "70B, Q4, MMLU 82"),
        ("small-model", "Small", "Quick summaries", ""),
    ]
    for index, (name, label, description, notes) in enumerate(rows):
        db.add(
            Model(
                connection_id=connection.id,
                model_id=name,
                display_name=label,
                description=description,
                notes=notes,
                position=index,
                capabilities_json={"tools": True},
            )
        )
    db.commit()


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


def _chat(db, model_id: str = "test-model") -> Chat:
    chat = Chat(user_id=_user(db).id, title="t", model_id=model_id)
    db.add(chat)
    db.commit()
    return chat


def _offered(db, families: list[str]) -> list[dict]:
    """Tool schemas for the families named, built from the real definitions so a
    family that stops existing takes these tests with it rather than passing on
    a hand-written string."""
    from lembas.services import tools as tools_service

    return [
        tool.schema
        for tool in tools_service.registry(db).values()
        if tools_service.gate_of(tool.family) in families
    ]


def _values(db, chat: Chat, *, families: list[str]) -> dict[str, str]:
    return harness_service.context_variables(db, _user(db), _offered(db, families), chat)


def _preamble(db, chat: Chat, *, families: list[str]) -> str:
    """The whole harness, through the path a request actually takes."""
    return harness_service.compose(db, _user(db), _offered(db, families), chat)


def test_every_other_model_is_listed_with_its_id(db):
    block = chat_service.roster_block(db, _user(db), exclude="test-model")
    assert "big-model" in block
    assert "small-model" in block
    assert "Big" in block


def test_the_asking_model_is_not_in_its_own_roster(db):
    block = chat_service.roster_block(db, _user(db), exclude="test-model")
    assert "test-model" not in block


def test_the_description_and_the_notes_both_reach_it(db):
    """Two fields on purpose: the description says what a model is for and is
    also shown to people, the notes say what it *is* and are for this alone. A
    model choosing whom to ask wants both."""
    block = chat_service.roster_block(db, _user(db), exclude="test-model")
    assert "Long reasoning problems" in block
    assert "70B, Q4, MMLU 82" in block


def test_a_model_this_account_cannot_reach_is_absent(db):
    group = Group(name="Wheel")
    db.add(group)
    restricted = db.scalar(select(Model).where(Model.model_id == "big-model"))
    restricted.public = False
    restricted.groups = [group]
    user = _user(db)
    user.role = ROLE_USER
    db.commit()

    block = chat_service.roster_block(db, user, exclude="test-model")
    assert "big-model" not in block
    assert "small-model" in block


def test_a_disabled_model_is_absent(db):
    off = db.scalar(select(Model).where(Model.model_id == "small-model"))
    off.enabled = False
    db.commit()
    assert "small-model" not in chat_service.roster_block(db, _user(db), exclude="test-model")


def test_the_block_is_bounded(db):
    """Every model an instance has multiplies this, and the harness has a budget
    the whole of it shares."""
    connection = db.scalars(select(Connection)).first()
    for index in range(60):
        db.add(
            Model(
                connection_id=connection.id,
                model_id=f"filler-{index}",
                display_name=f"Filler {index}",
                notes="x" * 400,
                position=10 + index,
            )
        )
    db.commit()

    block = chat_service.roster_block(db, _user(db), exclude="test-model")
    assert len(block) <= chat_service.MAX_ROSTER_CHARS + chat_service.MAX_ROSTER_ENTRY
    assert len(block.splitlines()) <= chat_service.MAX_ROSTER_MODELS


# --- Whether it is sent at all ------------------------------------------------
def test_the_variable_is_empty_for_a_model_that_cannot_ask_anyone(db):
    """Gated on the family, exactly as the memories block is gated on memory. A
    list of peers a model cannot reach is context spent on nothing, and it is why
    the roster and the tool are one switch rather than two."""
    chat = _chat(db)
    assert _values(db, chat, families=["memory"])["model_roster"] == ""
    assert _values(db, chat, families=["friend"])["model_roster"] != ""


def test_the_fragment_vanishes_on_a_single_model_instance(db):
    """`requires` rather than a conditional in the text: a heading above an empty
    list reads as "there is nobody", which is a different and wrong claim."""
    for extra in db.scalars(select(Model).where(Model.model_id != "test-model")):
        db.delete(extra)
    db.commit()

    chat = _chat(db)
    assert _values(db, chat, families=["friend"])["model_roster"] == ""
    assert "The other models here" not in _preamble(db, chat, families=["friend"])


def test_the_fragment_carries_the_list_when_there_is_one(db):
    chat = _chat(db)
    assembled = _preamble(db, chat, families=["friend"])
    assert "The other models here" in assembled
    assert "big-model" in assembled