STT-Assistant/tests/test_executor.py

"""Tests for the LLM executor's tool-calling loop."""

from stt_daemon.assistant.executor import LlmExecutor
from stt_daemon.assistant.llm import LlmResponse, ToolCall
from stt_daemon.assistant.tools import default_registry
from stt_daemon.config import LlmConfig


class FakeClient:
    def __init__(self, responses):
        self.responses = list(responses)
        self.calls: list[dict] = []

    def chat(self, messages, tools=None):
        self.calls.append({"messages": messages, "tools": tools})
        return self.responses.pop(0)


def _executor(responses, cfg=None):
    cfg = cfg or LlmConfig(max_tool_iterations=3)
    client = FakeClient(responses)
    return LlmExecutor(cfg, client, default_registry()), client


def _collector():
    emitted: list[dict] = []
    return emitted, (lambda type_, **fields: emitted.append({"type": type_, **fields}))


def test_tool_call_then_final_answer():
    call = ToolCall(id="1", name="get_datetime", arguments="{}")
    executor, client = _executor(
        [
            LlmResponse(content="", tool_calls=[call], finish_reason="tool_calls"),
            LlmResponse(content="It is 12:00."),
        ]
    )
    emitted, emit = _collector()
    result = executor.run("what time is it", emit)

    assert result.ok is True
    assert result.tool == "get_datetime"
    assert result.result == "It is 12:00."
    assert any(e["type"] == "tool_call" and e["tool"] == "get_datetime" for e in emitted)
    assert client.calls[1]["messages"][-1]["role"] == "tool"


def test_no_tool_call_returns_answer():
    executor, _ = _executor([LlmResponse(content="Just an answer.")])
    result = executor.run("hello", lambda *a, **k: None)
    assert result.ok is True
    assert result.result == "Just an answer."


def test_unknown_tool_is_rejected():
    call = ToolCall(id="1", name="nope", arguments="{}")
    executor, _ = _executor(
        [LlmResponse(content="", tool_calls=[call], finish_reason="tool_calls")]
    )
    emitted, emit = _collector()
    result = executor.run("do it", emit)

    assert result.ok is False
    assert any(e["type"] == "error" and e["code"] == "invalid_tool_call" for e in emitted)
    assert not any(e["type"] == "tool_call" for e in emitted)


def test_malformed_arguments_are_rejected():
    call = ToolCall(id="1", name="get_datetime", arguments="{not json")
    executor, _ = _executor(
        [LlmResponse(content="", tool_calls=[call], finish_reason="tool_calls")]
    )
    emitted, emit = _collector()
    result = executor.run("do it", emit)

    assert result.ok is False
    assert any(e["type"] == "error" and e["code"] == "invalid_tool_call" for e in emitted)
    assert not any(e["type"] == "tool_call" for e in emitted)


def test_tool_loop_is_capped():
    call = ToolCall(id="1", name="get_datetime", arguments="{}")
    responses = [
        LlmResponse(content="", tool_calls=[call], finish_reason="tool_calls")
        for _ in range(10)
    ]
    executor, client = _executor(responses, cfg=LlmConfig(max_tool_iterations=3))
    emitted, emit = _collector()
    result = executor.run("keep going", emit)

    assert result.ok is False
    assert any(e["type"] == "error" and e["code"] == "tool_loop" for e in emitted)
    assert len(client.calls) == 3


def test_truncated_empty_answer_is_an_error():
    executor, _ = _executor([LlmResponse(content="", finish_reason="length")])
    emitted, emit = _collector()
    result = executor.run("do it", emit)

    assert result.ok is False
    assert result.result == "response truncated"
    assert any(e["type"] == "error" and e["code"] == "truncated" for e in emitted)


def test_truncated_partial_answer_reports_partial():
    executor, _ = _executor([LlmResponse(content="par", finish_reason="length")])
    emitted, emit = _collector()
    result = executor.run("do it", emit)

    assert result.ok is False
    assert result.result == "par"
    assert any(e["type"] == "error" and e["code"] == "truncated" for e in emitted)


class RecordingClient:
    """Fake that accepts (and records) every extra keyword passed to `chat`."""

    def __init__(self, responses):
        self.responses = list(responses)
        self.calls: list[dict] = []

    def chat(self, messages, tools=None, **kwargs):
        self.calls.append({"messages": messages, "tools": tools, **kwargs})
        return self.responses.pop(0)


def test_executor_forwards_configured_slot():
    cfg = LlmConfig(max_tool_iterations=3)
    client = RecordingClient([LlmResponse(content="answer")])
    executor = LlmExecutor(cfg, client, default_registry(), slot_id=4)
    result = executor.run("hello", lambda *a, **k: None)

    assert result.ok is True
    assert result.result == "answer"
    assert client.calls[0]["slot_id"] == 4


def test_executor_omits_slot_when_none_for_old_fake():
    """Misleading-success probe: an old fake with no `slot_id` parameter is called
    fine, proving the kwarg is omitted (passing it would raise TypeError)."""
    executor, client = _executor([LlmResponse(content="answer")])
    result = executor.run("hello", lambda *a, **k: None)

    assert result.ok is True
    assert result.result == "answer"
    assert len(client.calls) == 1