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)

    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)

    # One round that may call tools, then one that has to answer with words.
    # Spelled out rather than derived from the constant: a test that reads
    # MAX_ROUNDS passes whatever MAX_ROUNDS becomes, which is exactly the
    # assertion nobody wanted.
    assert tools_service.MAX_ROUNDS == 1
    assert len(payloads) == 2
    # Recorded rather than silently dropped: an answer that stops here has to
    # be explicable.
    assert generation.tool_events[-1]["status"] == "error"
    assert "one round" in generation.tool_events[-1]["error"]


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,
    and with one round 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.

    No code change went with this; it falls out of the guard. The test is here
    because "it happens to work" and "it is meant to work" look the same until
    somebody changes the guard.
    """
    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"