LLeMbas/tests/test_tools.py

"""Tool calling: reassembling calls, deciding what is offered, running it."""

from __future__ import annotations

import json

import pytest

from lembas.db.models import Chat, Connection, Model
from lembas.services import settings_store
from lembas.services import tools as tools_service
from lembas.services.llm.openai_client import delta_tool_calls, finish_reason
from lembas.services.search.base import SearchError, SearchResult


# --- Reading the stream ------------------------------------------------------
def test_delta_tool_calls_reads_the_normal_shape():
    chunk = {"choices": [{"delta": {"tool_calls": [{"index": 0, "id": "a"}]}}]}
    assert delta_tool_calls(chunk) == [{"index": 0, "id": "a"}]


@pytest.mark.parametrize(
    "chunk", [{}, {"choices": []}, {"choices": [{}]}, {"choices": [{"delta": {}}]}]
)
def test_delta_tool_calls_tolerates_junk(chunk):
    assert delta_tool_calls(chunk) == []


def test_finish_reason_is_read_when_present():
    assert finish_reason({"choices": [{"finish_reason": "tool_calls"}]}) == "tool_calls"
    assert finish_reason({"choices": [{}]}) == ""


# --- The accumulator ---------------------------------------------------------
def test_arguments_split_across_chunks_are_rejoined():
    """Arguments arrive one token at a time; the id and name arrive once, on
    the first fragment only. This is the real shape llama.cpp produces."""
    accumulator = tools_service.ToolCallAccumulator()
    accumulator.feed(
        [{"index": 0, "id": "c1", "function": {"name": "web_search", "arguments": "{"}}]
    )
    for piece in ['"query"', ":", '"mallorn', ' tree"', "}"]:
        accumulator.feed([{"index": 0, "function": {"arguments": piece}}])

    assert accumulator.calls == [
        {"id": "c1", "name": "web_search", "arguments": '{"query":"mallorn tree"}'}
    ]


def test_two_calls_are_kept_apart_by_index():
    """Not by name: a model calling the same tool twice in one turn is exactly
    the case that breaks if name is the key."""
    accumulator = tools_service.ToolCallAccumulator()
    accumulator.feed(
        [
            {"index": 0, "id": "a", "function": {"name": "web_search", "arguments": '{"q":1}'}},
            {"index": 1, "id": "b", "function": {"name": "web_search", "arguments": '{"q":2}'}},
        ]
    )
    assert [c["id"] for c in accumulator.calls] == ["a", "b"]
    assert accumulator.calls[1]["arguments"] == '{"q":2}'


def test_a_missing_index_is_treated_as_the_only_call():
    """Some servers omit index entirely when there is just one call."""
    accumulator = tools_service.ToolCallAccumulator()
    accumulator.feed([{"function": {"name": "web_search", "arguments": "{}"}}])
    assert len(accumulator.calls) == 1


def test_a_call_with_no_name_is_not_a_call():
    accumulator = tools_service.ToolCallAccumulator()
    accumulator.feed([{"index": 0, "function": {"arguments": "{}"}}])
    assert accumulator.calls == []
    assert not accumulator


def test_an_id_is_invented_when_the_server_supplies_none():
    """The id is required when the results are sent back, and not every server
    provides one."""
    accumulator = tools_service.ToolCallAccumulator()
    accumulator.feed([{"index": 0, "function": {"name": "web_search", "arguments": "{}"}}])
    assert accumulator.calls[0]["id"]


# --- What gets offered -------------------------------------------------------
def _chat_with(db, user_id, *, capabilities):
    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=capabilities)
    )
    db.commit()
    chat = Chat(user_id=user_id, model_id="m", connection_id=connection.id)
    db.add(chat)
    db.commit()
    return chat


def _user(db, user_id):
    from lembas.db.models import User

    return db.get(User, user_id)


def _names(offered):
    return {tool["function"]["name"] for tool in offered}


def test_web_search_is_absent_when_search_is_off(db, user_id):
    """The library tools do not depend on a search provider, so they stay."""
    chat = _chat_with(db, user_id, capabilities={"tools": True})
    offered = _names(tools_service.enabled_tools(db, chat, _user(db, user_id)))
    assert "web_search" not in offered
    assert "notes_search" in offered


def test_nothing_at_all_without_the_tools_capability(db, user_id):
    """Sending a tools array to an endpoint that does not implement tool calling
    fails the entire request, exactly as image parts do without vision."""
    settings_store.update(db, {"enabled": True}, key=settings_store.SEARCH)
    chat = _chat_with(db, user_id, capabilities={})
    assert tools_service.enabled_tools(db, chat, _user(db, user_id)) == []


def test_web_search_is_offered_when_everything_lines_up(db, user_id):
    settings_store.update(db, {"enabled": True}, key=settings_store.SEARCH)
    chat = _chat_with(db, user_id, capabilities={"tools": True})
    assert "web_search" in _names(tools_service.enabled_tools(db, chat, _user(db, user_id)))


def test_a_model_predating_the_split_keeps_its_tools(db, user_id):
    """Rows configured before the per-tool flags existed have no tool_* keys.
    Reading absent as off would silently take web search away from every model
    already set up for it."""
    settings_store.update(db, {"enabled": True}, key=settings_store.SEARCH)
    chat = _chat_with(db, user_id, capabilities={"tools": True})
    assert "web_search" in _names(tools_service.enabled_tools(db, chat, _user(db, user_id)))


def test_a_family_turned_off_for_the_model_is_withheld(db, user_id):
    settings_store.update(db, {"enabled": True}, key=settings_store.SEARCH)
    chat = _chat_with(
        db, user_id, capabilities={"tools": True, "tool_notes": False}
    )
    offered = _names(tools_service.enabled_tools(db, chat, _user(db, user_id)))
    assert "notes_search" not in offered
    assert "web_search" in offered, "turning one family off must not affect another"


def test_web_search_is_withheld_when_the_provider_cannot_run(db, user_id, monkeypatch):
    """Offering a tool that will fail on every call is worse than not offering
    it at all."""
    monkeypatch.setattr(
        "lembas.services.search.availability", lambda _key: "ddgs is not installed."
    )
    settings_store.update(db, {"enabled": True}, key=settings_store.SEARCH)
    chat = _chat_with(db, user_id, capabilities={"tools": True})
    assert "web_search" not in _names(
        tools_service.enabled_tools(db, chat, _user(db, user_id))
    )


def _context(**kwargs):
    return tools_service.ToolContext(owner_id="someone", **kwargs)


# --- Knowledge is scoped to the chat's bases ---------------------------------
async def test_knowledge_search_is_limited_to_the_attached_bases(db, user_id):
    """"Answer from the contracts folder" is a different question from "answer
    from everything I have ever uploaded"."""
    from lembas.db.models import User
    from lembas.services.library import documents as documents_service

    owner = db.get(User, user_id)
    trees = documents_service.create_base(db, owner=owner, name="Trees")
    contracts = documents_service.create_base(db, owner=owner, name="Contracts")
    documents_service.store_upload(
        db, owner=owner, payload=b"The mallorn is golden.", filename="a.txt",
        title="Mallorn", base=trees,
    )
    documents_service.store_upload(
        db, owner=owner, payload=b"The mallorn clause is void.", filename="b.txt",
        title="Clause", base=contracts,
    )

    everywhere = await tools_service.run_tool(
        tools_service.ToolContext(owner_id=user_id), "knowledge_search",
        '{"query": "mallorn"}',
    )
    assert {r["title"] for r in everywhere.event["results"]} == {"Mallorn", "Clause"}

    scoped = await tools_service.run_tool(
        tools_service.ToolContext(owner_id=user_id, base_ids=[contracts.id]),
        "knowledge_search",
        '{"query": "mallorn"}',
    )
    assert [r["title"] for r in scoped.event["results"]] == ["Clause"]


def test_a_chat_with_no_bases_searches_everything(db, user_id):
    """Empty means "everything the owner can see", not "nothing"."""
    from lembas.db.models import User

    chat = _chat_with(db, user_id, capabilities={"tools": True})
    context = tools_service.context_for(db, db.get(User, user_id), chat)
    assert context.base_ids == []


# --- Running one -------------------------------------------------------------
async def test_running_web_search_formats_results_for_the_model(monkeypatch):
    async def fake_run(_config, query, *, limit=None):
        return [SearchResult("A title", "https://a.test", "a snippet")]

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

    outcome = await tools_service.run_tool(_context(), "web_search", '{"query": "mallorn"}')
    assert "A title" in outcome.content
    assert "https://a.test" in outcome.content
    assert outcome.event["status"] == "ok"
    assert outcome.event["results"][0]["host"] == "a.test"


async def test_malformed_argument_json_is_treated_as_the_query(monkeypatch):
    """Small models emit broken argument JSON often enough that this is a normal
    path, not an exceptional one."""
    seen = {}

    async def fake_run(_config, query, *, limit=None):
        seen["query"] = query
        return []

    monkeypatch.setattr("lembas.services.search.run", fake_run)
    await tools_service.run_tool(_context(), "web_search", "mallorn tree")
    assert seen["query"] == "mallorn tree"


async def test_a_failed_search_hands_the_model_an_explanation(monkeypatch):
    """A failed search should produce "I could not look that up", not kill the
    whole reply."""

    async def fake_run(_config, _query, *, limit=None):
        raise SearchError("DuckDuckGo is rate limiting this instance.")

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

    outcome = await tools_service.run_tool(_context(), "web_search", '{"query": "x"}')
    assert "rate limiting" in outcome.content
    assert outcome.event["status"] == "error"


async def test_an_unknown_tool_is_reported_rather_than_raised():
    outcome = await tools_service.run_tool(_context(), "launch_missiles", "{}")
    assert "no tool called" in outcome.content
    assert outcome.event["status"] == "error"


async def test_a_call_with_no_query_is_reported():
    outcome = await tools_service.run_tool(_context(), "web_search", '{"query": "  "}')
    assert outcome.event["status"] == "error"


# --- The turns sent back -----------------------------------------------------
def test_the_assistant_turn_echoes_the_calls():
    """The endpoint needs its own tool_calls back before the tool replies, or it
    has nothing to match the tool_call_ids against."""
    turn = tools_service.assistant_turn(
        [{"id": "c1", "name": "web_search", "arguments": '{"query":"x"}'}], "Looking it up."
    )
    assert turn["role"] == "assistant"
    assert turn["content"] == "Looking it up."
    assert turn["tool_calls"][0]["id"] == "c1"
    assert turn["tool_calls"][0]["function"]["name"] == "web_search"


def test_an_empty_assistant_message_becomes_null():
    """Most endpoints reject an assistant turn whose content is an empty string
    alongside tool_calls."""
    turn = tools_service.assistant_turn([{"id": "c", "name": "n", "arguments": "{}"}], "")
    assert turn["content"] is None


def test_the_tool_turn_carries_the_call_id():
    turn = tools_service.tool_turn({"id": "c1", "name": "web_search"}, "results here")
    assert turn == {
        "role": "tool",
        "tool_call_id": "c1",
        "name": "web_search",
        "content": "results here",
    }


def test_every_schema_is_valid_json():
    """They are sent verbatim to the endpoint; a schema that will not serialise
    fails every request rather than one."""
    for name, tool in tools_service.REGISTRY.items():
        json.dumps(tool.schema)
        assert tool.schema["function"]["name"] == name
        assert tool.family in tools_service.FAMILIES


def test_every_tool_describes_when_to_use_it():
    """The description is all the model has to decide with."""
    for tool in tools_service.REGISTRY.values():
        assert len(tool.description) > 40, tool.name