LLeMbas/tests/test_generation_tools.py

"""The tool loop: one reply, several requests."""

from __future__ import annotations

from lembas.db.models import ROLE_ASSISTANT, Chat, Connection, Message, Model
from lembas.services import generation as generation_service
from lembas.services import settings_store
from lembas.services import tools as tools_service
from lembas.services.search.base import SearchResult


def _chat_with_tools(db, user_id):
    connection = Connection(name="c", base_url="http://127.0.0.1:1", api_key_encrypted="")
    db.add(connection)
    db.commit()
    db.add(Model(connection_id=connection.id, model_id="m", capabilities_json={"tools": True}))
    db.commit()
    chat = Chat(user_id=user_id, model_id="m", connection_id=connection.id)
    db.add(chat)
    db.commit()
    db.add(Message(chat_id=chat.id, role="user", content="What is a mallorn?", complete=True))
    db.commit()
    assistant = Message(chat_id=chat.id, role=ROLE_ASSISTANT, content="", complete=False)
    db.add(assistant)
    db.commit()
    return chat.id, assistant.id


def _tool_call_chunk(name: str, arguments: str) -> dict:
    return {
        "choices": [
            {"delta": {"tool_calls": [{"index": 0, "id": "c1", "function": {
                "name": name, "arguments": arguments}}]}}
        ]
    }


def _text_chunk(text: str) -> dict:
    return {"choices": [{"delta": {"content": text}}]}


def _stub_stream(rounds, seen_payloads):
    """A stream_chat that returns a different scripted round each time."""

    async def stream_chat(_endpoint, payload):
        seen_payloads.append(payload)
        for chunk in rounds[min(len(seen_payloads) - 1, len(rounds) - 1)]:
            yield chunk

    return stream_chat


async def test_a_tool_call_produces_a_second_request(db, user_id, monkeypatch):
    """The whole point: one reply, two round trips, with the search result in
    the second one's messages."""
    settings_store.update(db, {"enabled": True}, key=settings_store.SEARCH)
    chat_id, message_id = _chat_with_tools(db, user_id)

    async def fake_search(_config, _query, *, limit=None):
        return [SearchResult("Mallorn", "https://tolkien.test/mallorn", "A golden tree.")]

    monkeypatch.setattr("lembas.services.search.run", fake_search)

    payloads = []
    monkeypatch.setattr(
        generation_service,
        "stream_chat",
        _stub_stream(
            [
                [_tool_call_chunk("web_search", '{"query": "mallorn"}')],
                [_text_chunk("A mallorn is a golden tree.")],
            ],
            payloads,
        ),
    )
    monkeypatch.setattr(
        "lembas.services.chat.generate_title", _never_called_title
    )

    generation = generation_service.Generation(chat_id=chat_id, message_id=message_id)
    await generation_service._run(generation)

    assert len(payloads) == 2, "the model asked for a tool, so it must be asked again"
    assert generation.text == "A mallorn is a golden tree."

    # The second request carries the assistant's own call back, then the result.
    followups = payloads[1]["messages"][-2:]
    assert followups[0]["tool_calls"][0]["function"]["name"] == "web_search"
    assert followups[1]["role"] == "tool"
    assert "https://tolkien.test/mallorn" in followups[1]["content"]

    # And the reader gets to see what it looked up.
    assert generation.tool_events[0]["query"] == "mallorn"
    assert generation.tool_events[0]["results"][0]["url"] == "https://tolkien.test/mallorn"


async def test_the_tools_array_is_absent_without_the_capability(db, user_id, monkeypatch):
    chat_id, message_id = _chat_with_tools(db, user_id)
    # Search enabled, but the model is not marked as supporting tools.
    settings_store.update(db, {"enabled": True}, key=settings_store.SEARCH)
    model = db.query(Model).first()
    model.capabilities_json = {}
    db.commit()

    payloads = []
    monkeypatch.setattr(
        generation_service, "stream_chat", _stub_stream([[_text_chunk("hi")]], payloads)
    )
    monkeypatch.setattr("lembas.services.chat.generate_title", _never_called_title)

    await generation_service._run(
        generation_service.Generation(chat_id=chat_id, message_id=message_id)
    )
    assert "tools" not in payloads[0]


async def test_text_before_a_tool_call_is_kept(db, user_id, monkeypatch):
    """A model that narrates what it is about to look up must not lose that
    when the results come back."""
    settings_store.update(db, {"enabled": True}, key=settings_store.SEARCH)
    chat_id, message_id = _chat_with_tools(db, user_id)

    monkeypatch.setattr(
        "lembas.services.search.run", _empty_search
    )
    monkeypatch.setattr(
        generation_service,
        "stream_chat",
        _stub_stream(
            [
                [
                    _text_chunk("Let me look that up. "),
                    _tool_call_chunk("web_search", '{"query": "mallorn"}'),
                ],
                [_text_chunk("Nothing found.")],
            ],
            [],
        ),
    )
    monkeypatch.setattr("lembas.services.chat.generate_title", _never_called_title)

    generation = generation_service.Generation(chat_id=chat_id, message_id=message_id)
    await generation_service._run(generation)
    assert generation.text == "Let me look that up. Nothing found."


async def test_a_model_that_only_ever_calls_tools_is_stopped(db, user_id, monkeypatch):
    """Otherwise a small model that has decided searching is the answer keeps
    searching until the context runs out, at a full request each time."""
    settings_store.update(db, {"enabled": True}, key=settings_store.SEARCH)
    chat_id, message_id = _chat_with_tools(db, user_id)

    monkeypatch.setattr("lembas.services.search.run", _empty_search)
    # Set explicitly rather than read from the constant: a test that reads the
    # number under test passes whatever the number becomes, which is the
    # assertion nobody wanted.
    settings_store.update(db, {"max_chat_rounds": 3})

    payloads = []
    monkeypatch.setattr(
        generation_service,
        "stream_chat",
        _stub_stream([[_tool_call_chunk("web_search", '{"query": "x"}')]], payloads),
    )
    monkeypatch.setattr("lembas.services.chat.generate_title", _never_called_title)

    generation = generation_service.Generation(chat_id=chat_id, message_id=message_id)
    await generation_service._run(generation)

    # Three rounds that may call tools, the one that notices the ceiling, and
    # the one asked for an answer with the tools withdrawn.
    assert len(payloads) == 5
    # And the last request carried no tools at all, which is what makes it a
    # round the model can only answer.
    assert "tools" not in payloads[-1]
    # Recorded rather than silently dropped: an answer that stops here has to
    # be explicable.
    assert generation.tool_events[-1]["status"] == "error"
    assert "3 rounds" in generation.tool_events[-1]["error"]


async def test_running_out_of_rounds_still_produces_an_answer(db, user_id, monkeypatch):
    """The bug this exists for.

    A model that goes straight to tool calls has written no prose at all by the
    time a budget runs out, so ending the reply there produced an empty bubble
    with an error line under it -- a good piece of research, six searches deep,
    thrown away. The tools are withdrawn and it is asked once more instead: what
    it gathered is in the transcript either way, and one request turns it into
    an answer.
    """
    settings_store.update(db, {"enabled": True}, key=settings_store.SEARCH)
    chat_id, message_id = _chat_with_tools(db, user_id)
    monkeypatch.setattr("lembas.services.search.run", _empty_search)
    settings_store.update(db, {"max_chat_rounds": 2})

    payloads: list[dict] = []

    async def stream_chat(_endpoint, payload):
        payloads.append(payload)
        # Exactly what a real model does: call tools while it has them, and
        # answer when it has none.
        if payload.get("tools"):
            yield _tool_call_chunk("web_search", '{"query": "x"}')
        else:
            yield {"choices": [{"delta": {"content": "Here is what I found."}}]}

    monkeypatch.setattr(generation_service, "stream_chat", stream_chat)
    monkeypatch.setattr("lembas.services.chat.generate_title", _never_called_title)

    generation = generation_service.Generation(chat_id=chat_id, message_id=message_id)
    await generation_service._run(generation)

    assert generation.text == "Here is what I found."
    assert not generation.error
    # And the reader can still tell this apart from an answer the model chose
    # to give, which is what the event is for.
    assert generation.tool_events[-1]["status"] == "error"
    assert "2 rounds" in generation.tool_events[-1]["error"]


async def test_no_ceiling_by_default(db, user_id):
    """A number low enough to be reached by ordinary work is a schedule, not a
    ceiling. What bounds an ordinary chat is the context window."""
    assert settings_store.chat_rounds(db) == 0


async def test_without_a_ceiling_the_model_is_told_to_keep_working(db, user_id):
    """Exactly one of the two fragments ever appears. With no budget the model
    must not be left with nothing said about when to stop -- and must certainly
    not be told it has a budget of two hundred, which it would ration."""
    from lembas.db.models import User
    from lembas.services import harness
    from lembas.services import tools as tools_service

    chat_id, _message_id = _chat_with_tools(db, user_id)
    chat = db.get(Chat, chat_id)
    user = db.get(User, user_id)
    offered = tools_service.resolve_tools(db, chat, user).schemas

    values = harness.context_variables(db, user, offered, chat)
    assert values["round_budget"] == ""
    assert values["unbounded"] == "yes"

    text = harness.compose(db, user, offered, chat)
    assert "Keep working until the task is actually done" in text
    assert "rounds of tool calls before you have to" not in text


async def test_with_a_ceiling_the_model_is_told_the_budget(db, user_id):
    from lembas.db.models import User
    from lembas.services import harness
    from lembas.services import tools as tools_service

    settings_store.update(db, {"max_chat_rounds": 3})
    chat_id, _message_id = _chat_with_tools(db, user_id)
    chat = db.get(Chat, chat_id)
    user = db.get(User, user_id)
    offered = tools_service.resolve_tools(db, chat, user).schemas

    values = harness.context_variables(db, user, offered, chat)
    assert values["round_budget"] == "3"
    assert values["unbounded"] == ""

    text = harness.compose(db, user, offered, chat)
    assert "at most 3 rounds" in text
    assert "Keep working until the task is actually done" not in text


async def test_a_prompt_queued_during_a_one_round_reply_waits_for_its_own(
    db, user_id, monkeypatch
):
    """`_inject` only takes a prompt in while there is a round left to answer
    in. With a ceiling of one there never is, so a queued message is not
    swallowed into a reply that then has no chance to address it -- it waits for
    `_drain`, which always gives it a reply of its own.

    One is no longer the default, but it is still a setting somebody can choose,
    and "it happens to work" and "it is meant to work" look the same until
    somebody changes the guard.
    """
    settings_store.update(db, {"max_chat_rounds": 1})
    from lembas.db.models import Message
    from lembas.services import chat as chat_service

    settings_store.update(db, {"enabled": True}, key=settings_store.SEARCH)
    chat_id, message_id = _chat_with_tools(db, user_id)
    chat = db.get(Chat, chat_id)
    queued = chat_service.create_message(db, chat, "user", "actually, do it the other way",
                                         queued=True)
    queued_id = queued.id

    monkeypatch.setattr("lembas.services.search.run", _empty_search)
    monkeypatch.setattr(
        generation_service,
        "stream_chat",
        _stub_stream(
            [[_tool_call_chunk("web_search", '{"query": "x"}')], [_text_chunk("Done.")]],
            [],
        ),
    )
    monkeypatch.setattr("lembas.services.chat.generate_title", _never_called_title)

    generation = generation_service.Generation(chat_id=chat_id, message_id=message_id)
    await generation_service._run(generation)

    # The row, not the payload: it was handed to a fresh reply by `_drain`,
    # which is what clears `queued`.
    db.expire_all()
    assert db.get(Message, queued_id).queued is False
    assert generation.drained is True


async def test_tool_activity_is_stored_with_the_message(db, user_id, monkeypatch):
    settings_store.update(db, {"enabled": True}, key=settings_store.SEARCH)
    chat_id, message_id = _chat_with_tools(db, user_id)

    async def fake_search(_config, _query, *, limit=None):
        return [SearchResult("Mallorn", "https://tolkien.test/m", "A tree.")]

    monkeypatch.setattr("lembas.services.search.run", fake_search)
    monkeypatch.setattr(
        generation_service,
        "stream_chat",
        _stub_stream(
            [
                [_tool_call_chunk("web_search", '{"query": "mallorn"}')],
                [_text_chunk("Done.")],
            ],
            [],
        ),
    )
    monkeypatch.setattr("lembas.services.chat.generate_title", _never_called_title)

    await generation_service._run(
        generation_service.Generation(chat_id=chat_id, message_id=message_id)
    )

    stored = db.get(Message, message_id)
    db.refresh(stored)
    assert stored.tool_calls_json[0]["query"] == "mallorn"
    assert stored.complete is True


async def _empty_search(_config, _query, *, limit=None):
    return []


async def _never_called_title(*_args, **_kwargs):
    """Auto-titling makes its own request; these tests are about the tool loop."""
    return "A title"


# --- Progress and concurrency ------------------------------------------------
async def test_the_status_names_the_running_tool_and_is_cleared(db, user_id, monkeypatch):
    """A remote tool can take seconds with nothing streaming, and a silent
    pause is exactly what a hang looks like."""
    settings_store.update(db, {"enabled": True}, key=settings_store.SEARCH)
    chat_id, message_id = _chat_with_tools(db, user_id)

    seen: list[str] = []

    async def fake_search(_config, _query, *, limit=None):
        seen.append(generation.status)
        return []

    monkeypatch.setattr("lembas.services.search.run", fake_search)
    monkeypatch.setattr(
        generation_service,
        "stream_chat",
        _stub_stream(
            [[_tool_call_chunk("web_search", '{"query": "mallorn"}')], [_text_chunk("Done.")]],
            [],
        ),
    )

    generation = generation_service.Generation(chat_id=chat_id, message_id=message_id)
    await generation_service._run(generation)

    # In words, from services/tool_labels.py -- the same table the transcript
    # and the approval card read. It used to say "Running web_search…".
    assert seen == ["Running Web search…"]
    assert generation.status == "", "and it is cleared once they are done"


async def test_results_stay_paired_with_their_calls_when_run_together(db, user_id, monkeypatch):
    """Indexed rather than appended as they finish: an endpoint matching on
    tool_call_id would otherwise pair the right id with the wrong content."""
    import asyncio

    settings_store.update(db, {"enabled": True}, key=settings_store.SEARCH)
    chat_id, message_id = _chat_with_tools(db, user_id)

    async def slow_first(_config, query, *, limit=None):
        # The first call finishes last, which is the whole point of the test.
        await asyncio.sleep(0.02 if query == "first" else 0)
        return [SearchResult(f"result for {query}", f"https://t.test/{query}", "")]

    monkeypatch.setattr("lembas.services.search.run", slow_first)

    two_calls = {
        "choices": [
            {"delta": {"tool_calls": [
                {"index": 0, "id": "a", "function": {
                    "name": "web_search", "arguments": '{"query": "first"}'}},
                {"index": 1, "id": "b", "function": {
                    "name": "web_search", "arguments": '{"query": "second"}'}},
            ]}}
        ]
    }
    payloads: list[dict] = []
    monkeypatch.setattr(
        generation_service,
        "stream_chat",
        _stub_stream([[two_calls], [_text_chunk("Done.")]], payloads),
    )

    generation = generation_service.Generation(chat_id=chat_id, message_id=message_id)
    await generation_service._run(generation)

    turns = [m for m in payloads[1]["messages"] if m.get("role") == "tool"]
    assert [turn["tool_call_id"] for turn in turns] == ["a", "b"]
    assert "first" in turns[0]["content"] and "second" in turns[1]["content"]
    # And the transcript keeps the same order.
    assert [event["query"] for event in generation.tool_events] == ["first", "second"]


async def test_a_custom_tool_runs_inside_the_loop(db, user_id, monkeypatch, mock_http):
    """End to end: a row becomes an offered schema, the model calls it, and the
    result comes back in the next request's messages."""
    import httpx

    from lembas.db.models import CustomTool

    monkeypatch.setattr(
        "socket.getaddrinfo", lambda *a, **k: [(2, 1, 6, "", ("93.184.216.34", 80))]
    )
    mock_http(lambda _r: httpx.Response(200, json={"summary": "Sunny in Minas Tirith."}))

    db.add(
        CustomTool(
            slug="weather",
            name="Weather",
            description="Look up the weather.",
            url_template="https://api.test/{{city}}",
            parameters_json={"type": "object", "properties": {"city": {"type": "string"}}},
            response_mode="json",
            response_path="summary",
        )
    )
    db.commit()

    chat_id, message_id = _chat_with_tools(db, user_id)
    payloads: list[dict] = []
    monkeypatch.setattr(
        generation_service,
        "stream_chat",
        _stub_stream(
            [
                [_tool_call_chunk("weather", '{"city": "Minas Tirith"}')],
                [_text_chunk("It is sunny.")],
            ],
            payloads,
        ),
    )

    generation = generation_service.Generation(chat_id=chat_id, message_id=message_id)
    await generation_service._run(generation)

    offered = {tool["function"]["name"] for tool in payloads[0]["tools"]}
    assert "weather" in offered

    tool_turns = [m for m in payloads[1]["messages"] if m.get("role") == "tool"]
    assert tool_turns[0]["content"] == "Sunny in Minas Tirith."
    assert generation.tool_events[0]["kind"] == "custom"
    assert generation.tool_events[0]["label"] == "Weather"


def test_the_default_and_the_fallback_cannot_drift(db):
    """`tools.MAX_ROUNDS` exists for callers with no session; the setting is
    what the loop and the harness read. Two numbers meaning one thing is how
    a model gets told a budget it does not have."""
    assert tools_service.MAX_ROUNDS == settings_store.DEFAULT_CHAT_ROUNDS
    assert settings_store.chat_rounds(db) == tools_service.MAX_ROUNDS


async def test_a_ceiling_of_zero_does_not_mean_zero_rounds(db, user_id, monkeypatch):
    """It means no ceiling. Read carelessly it would mean the model never gets
    to call anything, which is the opposite."""
    settings_store.update(db, {"enabled": True}, key=settings_store.SEARCH)
    settings_store.update(db, {"max_chat_rounds": 0})
    chat_id, message_id = _chat_with_tools(db, user_id)

    monkeypatch.setattr("lembas.services.search.run", _empty_search)
    payloads = []
    monkeypatch.setattr(
        generation_service,
        "stream_chat",
        _stub_stream(
            [[_tool_call_chunk("web_search", '{"query": "x"}')], [_text_chunk("Done.")]],
            payloads,
        ),
    )
    monkeypatch.setattr("lembas.services.chat.generate_title", _never_called_title)

    generation = generation_service.Generation(chat_id=chat_id, message_id=message_id)
    await generation_service._run(generation)

    assert generation.text == "Done."
    assert len(payloads) == 2, "it called a tool and then answered, uninterrupted"


def test_the_ceiling_is_clamped(db):
    settings_store.update(db, {"max_chat_rounds": 9999})
    assert settings_store.chat_rounds(db) == 100
    settings_store.update(db, {"max_chat_rounds": -5})
    assert settings_store.chat_rounds(db) == 0