Funk-lab/stt/server/tests/test_tool_loop.py

"""Tests de la boucle de function-calling (`brain.ask_with_tools`).

Hors-ligne : on monkeypatche `brain._post` (l'appel LiteLLM) pour simuler les réponses
du modèle, et on fournit un `dispatch` factice. Lançable avec `pytest` (pas de plugin
async requis : chaque test pilote sa propre boucle via `asyncio.run`).
"""

from __future__ import annotations

import asyncio

from stt_server import brain


def test_loop_calls_tool_then_answers(monkeypatch):
    calls: list[dict] = []

    async def fake_post(payload):
        calls.append(payload)
        if len(calls) == 1:  # 1er tour : le modèle demande un outil
            return {
                "role": "assistant", "content": "",
                "tool_calls": [{
                    "id": "t1", "type": "function",
                    "function": {"name": "host_health", "arguments": '{"host": "gpu-01"}'},
                }],
            }
        return {"role": "assistant", "content": "gpu-01 est en ligne et se porte bien."}

    monkeypatch.setattr(brain, "_post", fake_post)
    executed: list[tuple[str, dict]] = []

    async def dispatch(name, args):
        executed.append((name, args))
        return "gpu-01 : en ligne, charge 0.16, RAM libre 48%."

    async def go():
        schemas = [{"type": "function", "function": {"name": "host_health", "parameters": {}}}]
        return await brain.ask_with_tools(
            "gpu-01 tourne bien ?", "Tu es Asa.",
            schemas=schemas, dispatch=dispatch,
        )

    reply, trace = asyncio.run(go())
    assert executed == [("host_health", {"host": "gpu-01"})]
    assert "en ligne" in reply
    assert len(trace) == 1 and trace[0]["name"] == "host_health"
    # Le 2e appel doit réinjecter l'assistant (avec tool_calls) puis le résultat d'outil.
    assert [m["role"] for m in calls[1]["messages"]] == ["system", "user", "assistant", "tool"]


def test_no_schemas_falls_back_to_plain_ask(monkeypatch):
    async def fake_post(payload):
        assert "tools" not in payload  # repli sur ask() : pas d'outils dans le payload
        return {"role": "assistant", "content": "réponse simple"}

    monkeypatch.setattr(brain, "_post", fake_post)

    async def dispatch(name, args):  # ne doit jamais être appelé
        raise AssertionError("dispatch ne doit pas être invoqué sans schémas")

    async def go():
        return await brain.ask_with_tools(
            "salut", "sys", schemas=[], dispatch=dispatch,
        )

    reply, trace = asyncio.run(go())
    assert reply == "réponse simple" and trace == []


def test_confirm_tool_stops_loop(monkeypatch):
    """Un outil de `confirm_tools` (ex. admin_action) arrête la boucle et surface sa réponse."""
    calls: list[dict] = []

    async def fake_post(payload):
        calls.append(payload)
        return {  # le modèle veut agir → appelle admin_action
            "role": "assistant", "content": "",
            "tool_calls": [{
                "id": "a1", "type": "function",
                "function": {"name": "admin_action",
                             "arguments": '{"description": "redémarrer open-webui"}'},
            }],
        }

    monkeypatch.setattr(brain, "_post", fake_post)

    async def dispatch(name, args):
        # simule le handshake : renvoie la question de confirmation (pas d'exécution)
        return f"Tu veux que j'exécute : « {args['description']} » ? Dis « confirme »."

    async def go():
        schemas = [{"type": "function", "function": {"name": "admin_action", "parameters": {}}}]
        return await brain.ask_with_tools(
            "redémarre open-webui", "sys",
            schemas=schemas, dispatch=dispatch, confirm_tools=("admin_action",),
        )

    reply, trace = asyncio.run(go())
    assert "confirme" in reply
    assert len(calls) == 1  # boucle stoppée après 1 appel LLM (pas de 2e tour)
    assert trace[0]["name"] == "admin_action"


def test_loop_forces_answer_after_max_iters(monkeypatch):
    """Si le modèle boucle indéfiniment sur des outils, un dernier appel force une réponse."""
    calls: list[dict] = []

    async def fake_post(payload):
        calls.append(payload)
        if payload.get("tool_choice") == "none":  # appel final forcé
            return {"role": "assistant", "content": "Réponse finale forcée."}
        return {  # tour normal : redemande toujours un outil
            "role": "assistant", "content": "",
            "tool_calls": [{
                "id": f"t{len(calls)}", "type": "function",
                "function": {"name": "cluster_status", "arguments": "{}"},
            }],
        }

    monkeypatch.setattr(brain, "_post", fake_post)

    async def dispatch(name, args):
        return "ok"

    async def go():
        schemas = [{"type": "function", "function": {"name": "cluster_status", "parameters": {}}}]
        return await brain.ask_with_tools(
            "état ?", "sys", schemas=schemas, dispatch=dispatch,
        )

    reply, trace = asyncio.run(go())
    assert reply == "Réponse finale forcée."
    # tool_max_iters tours d'outils + 1 appel final forcé.
    assert calls[-1]["tool_choice"] == "none"
    assert len(trace) == brain.settings.tool_max_iters