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_context_tokens=1330,
        max_tokens=1024,
        compaction_safety_margin=256,
        trim_block_turns=1,
        compaction_enabled=False,
    )  # prompt_budget == 50
    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 == []


# ---------------------------------------------------------------------------
# Phase C tests
# ---------------------------------------------------------------------------

def _compact_cfg(
    *,
    budget_chars=200,
    trigger_ratio=0.8,
    keep_turns=1,
    compaction_max_tokens=64,
    compaction_timeout_s=5.0,
    compaction_input_tokens=8192,
    compaction_enabled=True,
    trim_block_turns=1,
):
    """Return a small LlmConfig whose prompt_budget == budget_chars // 4 tokens.

    prompt_budget = max_context_tokens - max_tokens - compaction_safety_margin
    We want budget_chars // 4 tokens, so:
        max_context_tokens = (budget_chars // 4) + max_tokens + compaction_safety_margin
    """
    max_tokens              = 16
    compaction_safety_margin = 0
    max_context_tokens      = (budget_chars // 4) + max_tokens + compaction_safety_margin
    return LlmConfig(
        max_context_tokens=max_context_tokens,
        max_tokens=max_tokens,
        compaction_safety_margin=compaction_safety_margin,
        compaction_enabled=compaction_enabled,
        compaction_trigger_ratio=trigger_ratio,
        compaction_keep_turns=keep_turns,
        compaction_max_tokens=compaction_max_tokens,
        compaction_timeout_s=compaction_timeout_s,
        compaction_input_tokens=compaction_input_tokens,
        trim_block_turns=trim_block_turns,
        system_prompt_chat="SYS",
        system_prompt_command="CMD",
        compaction_system_prompt="COMPACT",
    )


def test_compaction_fires_and_window_shrinks():
    """Crossing the trigger compacts: window shrinks, summary present in next prompt."""
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
    # prompt_budget = 200//4 = 50 tokens; trigger at 40 tokens (0.8 * 50)

    summary_text = "summary of old turns"
    call_index   = 0

    def responder(payload):
        nonlocal call_index
        msgs = payload["body"]["messages"]
        call_index += 1
        # Summarizer calls have the compaction system prompt.
        if msgs[0]["content"] == "COMPACT":
            return _reply(content=summary_text)
        return _reply(content="ok")

    client, calls = _make_client(responder, cfg)
    session = ChatSession(cfg, client)

    # Send enough turns to exceed the trigger (each "x"*40 = 40 chars = 10 tokens).
    # After 4 turns (user+assistant pairs = 8 messages * 10 tokens = 80 tokens > 40 trigger).
    for _ in range(4):
        session.send("x" * 40)

    # After compaction the window must be smaller than the full history.
    assert len(session.history) <= 4  # keep_turns=1 → keep=2 messages

    # The summary must be set.
    assert session.summary != ""

    # The next chat call must include the summary as a second system message,
    # with the chat system prompt FIRST.
    session.send("after compaction")
    chat_calls = [c for c in calls if c["body"]["messages"][0]["content"] == "SYS"]
    last_chat  = chat_calls[-1]["body"]["messages"]
    assert last_chat[0]["role"] == "system"
    assert last_chat[0]["content"] == "SYS"
    assert last_chat[1]["role"] == "system"
    assert "Earlier conversation summary:" in last_chat[1]["content"]
    assert summary_text in last_chat[1]["content"]


def test_summary_never_appended_to_history():
    """The summary text must not appear as a history entry."""
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
    summary_text = "SUMMARY_SENTINEL"

    def responder(payload):
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            return _reply(content=summary_text)
        return _reply(content="ok")

    client, _ = _make_client(responder, cfg)
    session   = ChatSession(cfg, client)

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

    for msg in session.history:
        assert summary_text not in msg["content"]


def test_summarizer_failure_falls_back_to_trim():
    """On summarizer failure the session falls back to _trim() and stays within budget."""
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)

    def responder(payload):
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            return _reply(content="", finish="stop")  # empty → failure
        return _reply(content="ok")

    client, _ = _make_client(responder, cfg)
    session   = ChatSession(cfg, client)

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

    # Summary must remain empty (fallback used trim, not compaction).
    assert session.summary == ""
    # History must be bounded within the budget.
    assert session._estimate_tokens() <= cfg.prompt_budget


def test_summarizer_length_finish_reason_falls_back_to_trim():
    """finish_reason == 'length' on the summarizer is treated as failure."""
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)

    def responder(payload):
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            return _reply(content="partial summary", finish="length")
        return _reply(content="ok")

    client, _ = _make_client(responder, cfg)
    session   = ChatSession(cfg, client)

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

    assert session.summary == ""
    assert session._estimate_tokens() <= cfg.prompt_budget


def test_summarizer_llmerror_falls_back_to_trim():
    """LlmError from the summarizer falls back to _trim()."""
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)

    def responder(payload):
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            raise LlmError("network error")
        return _reply(content="ok")

    client, _ = _make_client(responder, cfg)
    session   = ChatSession(cfg, client)

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

    assert session.summary == ""
    assert session._estimate_tokens() <= cfg.prompt_budget


def test_reset_clears_summary():
    """reset() clears both history and summary."""
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
    summary_text = "old summary"

    def responder(payload):
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            return _reply(content=summary_text)
        return _reply(content="ok")

    client, _ = _make_client(responder, cfg)
    session   = ChatSession(cfg, client)

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

    assert session.summary != ""

    session.reset()
    assert session.history == []
    assert session.summary == ""


def test_compaction_disabled_never_calls_summarizer():
    """compaction_enabled=False never calls the summarizer and still trims."""
    cfg = _compact_cfg(
        budget_chars=200,
        trigger_ratio=0.8,
        keep_turns=1,
        compaction_enabled=False,
    )

    summarizer_called = False

    def responder(payload):
        nonlocal summarizer_called
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            summarizer_called = True
            return _reply(content="should not happen")
        return _reply(content="ok")

    client, _ = _make_client(responder, cfg)
    session   = ChatSession(cfg, client)

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

    assert not summarizer_called
    assert session.summary == ""
    # History must still be bounded (trim ran).
    assert session._estimate_tokens() <= cfg.prompt_budget


def test_nothing_to_evict_skips_summarizer():
    """When len(history) <= keep, compaction is skipped and trim runs instead."""
    # keep_turns=10 → keep=20 messages; we only send 2 turns (4 messages).
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=10)

    summarizer_called = False

    def responder(payload):
        nonlocal summarizer_called
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            summarizer_called = True
            return _reply(content="should not happen")
        return _reply(content="ok")

    client, _ = _make_client(responder, cfg)
    session   = ChatSession(cfg, client)

    # Push over the trigger but keep history shorter than keep=20.
    for _ in range(4):
        session.send("x" * 40)

    assert not summarizer_called


def test_chunk_and_fold_large_evicted():
    """Large evicted list triggers chunk-and-fold (more than one summarizer call).

    budget_chars=160 → prompt_budget=40 tokens, trigger at 32 tokens (0.8 * 40).
    Each turn: user='x'*40 (40 chars) + assistant='ok' (2 chars) = ~10 tokens.
    Compaction fires at turn 4 with 6 evicted messages.
    compaction_input_tokens=1 → char_limit=4 chars per chunk, so each 40-char
    or 2-char message is its own chunk → 6 summarizer calls for 6 chunks.
    """
    cfg = _compact_cfg(
        budget_chars=160,
        trigger_ratio=0.8,
        keep_turns=1,
        compaction_input_tokens=1,  # 1 token * 4 = 4 chars per chunk → each msg is a chunk
    )

    summarizer_calls = []

    def responder(payload):
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            summarizer_calls.append(payload)
            return _reply(content="folded")
        return _reply(content="ok")

    client, _ = _make_client(responder, cfg)
    session   = ChatSession(cfg, client)

    # 4 turns → 8 messages → fires compaction with 6 evicted (keep=2).
    for _ in range(4):
        session.send("x" * 40)

    # With compaction_input_tokens=1 (4 chars/chunk) and 40-char messages,
    # each message is its own chunk → multiple summarizer calls.
    assert len(summarizer_calls) > 1


def test_chat_per_request_max_tokens_override():
    """LlmClient.chat sends per-request max_tokens when given, cfg default otherwise."""
    cfg = LlmConfig(max_tokens=1024, timeout_s=60.0)
    client, calls = _make_client(lambda _p: _reply(content="ok"), cfg)

    # Default: uses cfg.max_tokens.
    client.chat([{"role": "user", "content": "hi"}])
    assert calls[0]["body"]["max_tokens"] == 1024

    # Override: uses the supplied value.
    client.chat([{"role": "user", "content": "hi"}], max_tokens=128)
    assert calls[1]["body"]["max_tokens"] == 128


def test_chat_per_request_timeout_override():
    """LlmClient.chat uses per-request timeout_s when given, cfg default otherwise."""
    cfg = LlmConfig(max_tokens=1024, timeout_s=60.0)
    recorded_timeouts: list[float] = []

    def opener(payload, timeout):
        recorded_timeouts.append(timeout)
        return _reply(content="ok")

    client = LlmClient(cfg, opener=opener)

    # Default timeout.
    client.chat([{"role": "user", "content": "hi"}])
    assert recorded_timeouts[0] == 60.0

    # Per-request override.
    client.chat([{"role": "user", "content": "hi"}], timeout_s=5.0)
    assert recorded_timeouts[1] == 5.0


def test_summarizer_uses_compaction_max_tokens_and_timeout():
    """The summarizer call uses compaction_max_tokens and compaction_timeout_s."""
    cfg = _compact_cfg(
        budget_chars=200,
        trigger_ratio=0.8,
        keep_turns=1,
        compaction_max_tokens=42,
        compaction_timeout_s=7.0,
    )

    recorded: list[dict] = []

    def opener(payload, timeout):
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            recorded.append({"max_tokens": payload["body"]["max_tokens"], "timeout": timeout})
            return _reply(content="summary")
        return _reply(content="ok")

    client  = LlmClient(cfg, opener=opener)
    session = ChatSession(cfg, client)

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

    assert len(recorded) >= 1
    assert recorded[0]["max_tokens"] == 42
    assert recorded[0]["timeout"] == 7.0


def test_notify_called_on_compaction():
    """notify('compacted', turns=N) is called after a successful compaction."""
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)

    notifications: list[tuple] = []

    def notify(kind, **fields):
        notifications.append((kind, fields))

    def responder(payload):
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            return _reply(content="summary")
        return _reply(content="ok")

    client  = LlmClient(cfg, opener=lambda p, t: responder(p))
    session = ChatSession(cfg, client, notify=notify)

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

    assert any(kind == "compacted" for kind, _ in notifications)
    compacted = [(k, f) for k, f in notifications if k == "compacted"]
    assert compacted[0][1]["turns"] > 0


def test_notify_not_called_when_none():
    """ChatSession works without a notify callback (default None)."""
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)

    def responder(payload):
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            return _reply(content="summary")
        return _reply(content="ok")

    client  = LlmClient(cfg, opener=lambda p, t: responder(p))
    session = ChatSession(cfg, client)  # no notify

    # Must not raise.
    for _ in range(4):
        session.send("x" * 40)


def test_system_prompt_is_first_message_after_compaction():
    """After compaction the chat system prompt is always the first message."""
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)

    def responder(payload):
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            return _reply(content="summary text")
        return _reply(content="ok")

    client, calls = _make_client(responder, cfg)
    session       = ChatSession(cfg, client)

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

    # Find the last chat call (system prompt == "SYS", not "COMPACT").
    chat_calls = [c for c in calls if c["body"]["messages"][0]["content"] == "SYS"]
    assert chat_calls, "no chat calls found after compaction"
    last_msgs = chat_calls[-1]["body"]["messages"]
    assert last_msgs[0]["role"] == "system"
    assert last_msgs[0]["content"] == "SYS"


def test_chat_session_truncated_empty_reply_notifies_and_is_not_appended():
    """A length-limited empty reply is reported and never stored as an empty turn."""
    notifications: list[tuple] = []

    def notify(kind, **fields):
        notifications.append((kind, fields))

    client, _ = _make_client(lambda _p: _reply(content="", finish="length"))
    session = ChatSession(LlmConfig(), client, notify=notify)

    assert session.send("hi") == ""
    assert session.history == [{"role": "user", "content": "hi"}]
    assert notifications == [("truncated", {"message": "response truncated"})]


def test_chat_session_truncated_partial_reply_is_appended_and_notified():
    """A truncated but non-empty reply is kept and reported as truncated."""
    notifications: list[tuple] = []

    def notify(kind, **fields):
        notifications.append((kind, fields))

    client, _ = _make_client(lambda _p: _reply(content="partial answer", finish="length"))
    session = ChatSession(LlmConfig(), client, notify=notify)

    assert session.send("hi") == "partial answer"
    assert session.history[-1] == {"role": "assistant", "content": "partial answer"}
    assert any(kind == "truncated" for kind, _ in notifications)


def test_on_io_receives_request_then_response():
    """chat(..., on_io=cb) reports the request (url+body) then the response payload."""
    client, _ = _make_client(lambda _p: _reply(content="ok"))
    seen: list[dict] = []

    response = client.chat([{"role": "user", "content": "hi"}], on_io=seen.append)

    assert response.content == "ok"
    assert [event["phase"] for event in seen] == ["request", "response"]
    assert seen[0]["url"] == LlmConfig().endpoint
    assert seen[0]["body"]["messages"] == [{"role": "user", "content": "hi"}]
    assert seen[0]["body"]["max_tokens"] == LlmConfig().max_tokens
    assert seen[1]["payload"]["choices"][0]["message"]["content"] == "ok"
    assert seen[1]["error"] is None


def test_on_io_response_phase_reports_opener_error():
    """An opener raising still yields a response phase with the error set."""
    import urllib.error

    def opener(_payload, _timeout):
        raise urllib.error.URLError("down")

    client = LlmClient(LlmConfig(), opener=opener)
    seen: list[dict] = []

    with pytest.raises(LlmError):
        client.chat([{"role": "user", "content": "hi"}], on_io=seen.append)

    assert [event["phase"] for event in seen] == ["request", "response"]
    assert seen[1]["payload"] is None
    assert "down" in seen[1]["error"]


def test_raising_on_io_does_not_break_chat():
    """A throwing on_io callback is swallowed and the reply still parses."""

    def bad(_event):
        raise RuntimeError("boom")

    client, _ = _make_client(lambda _p: _reply(content="ok"))

    assert client.chat([{"role": "user", "content": "hi"}], on_io=bad).content == "ok"