STT-Assistant/tests/test_llm.py

"""Tests for the LLM client and chat session (no network)."""

import pytest

from stt_daemon.assistant.llm import ChatSession, LlmClient, LlmError
from stt_daemon.config import LlmConfig


def _make_client(responder, cfg=None):
    cfg = cfg or LlmConfig()
    calls: list[dict] = []

    def opener(payload, timeout):
        calls.append(payload)
        return responder(payload)

    return LlmClient(cfg, opener=opener), calls


def _reply(content="", reasoning=None, tool_calls=None, finish="stop"):
    message = {"role": "assistant", "content": content}
    if reasoning is not None:
        message["reasoning_content"] = reasoning
    if tool_calls:
        message["tool_calls"] = tool_calls
    return {"choices": [{"message": message, "finish_reason": finish}]}


def test_reasoning_content_is_ignored():
    client, _ = _make_client(
        lambda _p: _reply(content="It is noon.", reasoning="thinking hard about it")
    )
    response = client.chat([{"role": "user", "content": "hi"}])
    assert response.content == "It is noon."
    assert "thinking hard" not in response.content


def test_cache_prompt_and_endpoint_are_sent():
    client, calls = _make_client(lambda _p: _reply(content="ok"))
    client.chat([{"role": "user", "content": "hi"}])

    assert calls[0]["body"]["cache_prompt"] is True
    assert calls[0]["body"]["stream"] is False
    assert calls[0]["__url__"] == LlmConfig().endpoint


def test_tools_sent_only_when_provided():
    client, calls = _make_client(lambda _p: _reply(content="ok"))
    client.chat([{"role": "user", "content": "hi"}])
    assert "tools" not in calls[0]["body"]

    client.chat(
        [{"role": "user", "content": "hi"}],
        tools=[{"type": "function", "function": {"name": "get_datetime"}}],
    )
    assert calls[1]["body"]["tools"][0]["function"]["name"] == "get_datetime"
    assert calls[1]["body"]["tool_choice"] == "auto"


def test_tool_calls_are_parsed():
    client, _ = _make_client(
        lambda _p: _reply(
            tool_calls=[
                {
                    "id": "1",
                    "type": "function",
                    "function": {"name": "get_datetime", "arguments": "{}"},
                }
            ],
            finish="tool_calls",
        )
    )
    response = client.chat([{"role": "user", "content": "time?"}], tools=[{"type": "function"}])
    assert len(response.tool_calls) == 1
    assert response.tool_calls[0].name == "get_datetime"
    assert response.finish_reason == "tool_calls"


def test_missing_choices_raises():
    client, _ = _make_client(lambda _p: {"choices": []})
    with pytest.raises(LlmError):
        client.chat([{"role": "user", "content": "hi"}])


def test_chat_session_appends_and_keeps_prefix():
    client, calls = _make_client(lambda _p: _reply(content="hello"))
    session = ChatSession(LlmConfig(), client)

    assert session.send("hi") == "hello"
    session.send("again")

    assert [m["role"] for m in session.history] == [
        "user",
        "assistant",
        "user",
        "assistant",
    ]
    first = calls[0]["body"]["messages"]
    second = calls[1]["body"]["messages"]
    assert second[: len(first)] == first  # append-only, prefix preserved
    assert second[0]["content"] == LlmConfig().system_prompt_chat


def test_chat_session_trims_in_blocks_and_keeps_system_prompt():
    cfg = LlmConfig(max_prompt_tokens=50, trim_block_turns=1)
    client, calls = _make_client(lambda _p: _reply(content="ok"))
    session = ChatSession(cfg, client)

    for _ in range(10):
        session.send("x" * 40)

    assert len(session.history) <= 6
    assert calls[-1]["body"]["messages"][0]["content"] == cfg.system_prompt_chat


def test_chat_session_reset_forgets_history():
    client, _ = _make_client(lambda _p: _reply(content="ok"))
    session = ChatSession(LlmConfig(), client)
    session.send("hi")
    session.reset()
    assert session.history == []


def test_chat_session_drops_reply_when_reset_mid_flight():
    holder: list[ChatSession] = []

    def opener(_payload, _timeout):
        holder[0].reset()
        return _reply(content="late reply")

    holder.append(ChatSession(LlmConfig(), LlmClient(LlmConfig(), opener=opener)))

    assert holder[0].send("hi") == ""
    assert holder[0].history == []