STT-Assistant/src/stt_daemon/assistant/llm.py

"""llama.cpp (OpenAI-compatible) client and the chat session.

Reasoning models return a separate `reasoning_content` field, which is ignored
by field name — never by scraping it out of `content`.
"""

import json
import logging
import threading
import urllib.error
import urllib.request
from dataclasses import dataclass, field
from typing import Any, Callable

logger = logging.getLogger(__name__)


class LlmError(RuntimeError):
    """Raised when the LLM endpoint cannot be reached or returns garbage."""


@dataclass
class ToolCall:
    """One tool call requested by the model."""

    id: str
    name: str
    arguments: str

    def message(self) -> dict:
        return {
            "id": self.id,
            "type": "function",
            "function": {"name": self.name, "arguments": self.arguments},
        }


@dataclass
class LlmResponse:
    """Parsed completion."""

    content: str
    tool_calls: list[ToolCall] = field(default_factory=list)
    finish_reason: str = "stop"


def _default_opener(payload: dict, timeout: float) -> dict:
    request = urllib.request.Request(
        payload["__url__"],
        data=json.dumps(payload["body"]).encode("utf-8"),
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(request, timeout=timeout) as response:
        return json.loads(response.read().decode("utf-8"))


def _safe_io(cb, event) -> None:
    try:
      cb(event)
    except Exception as e:  # noqa: BLE001 - a dump hook must not break chat
      logger.warning("on_io callback failed: %s", e)


class LlmClient:
    """Minimal OpenAI-compatible chat client (stdlib only)."""

    def __init__(self, cfg, opener: Callable[[dict, float], dict] | None = None):
        """Create the client.

        Args:
            cfg: An `LlmConfig`.
            opener: Transport override for tests; receives
                `{"__url__":..., "body":...}` and a timeout, returns the parsed JSON.
        """
        self.cfg = cfg
        self._opener = opener or _default_opener

    def chat(
        self,
        messages: list[dict],
        tools: list[dict] | None = None,
        max_tokens: int | None = None,
        timeout_s: float | None = None,
        on_io: Callable[[dict], None] | None = None,
        *,
        slot_id: int | None = None,
    ) -> LlmResponse:
        """Send one completion request and parse the reply.

        Args:
            messages: The conversation messages to send.
            tools: Optional tool schemas.
            max_tokens: Per-request output cap; uses cfg.max_tokens when None.
            timeout_s: Per-request timeout; uses cfg.timeout_s when None.
            on_io: Optional hook receiving `{"phase": "request", ...}` before
                the call and `{"phase": "response", ...}` after it.
            slot_id: Optional llama.cpp slot id (`id_slot`); omitted when None.
        """
        effective_max_tokens = self.cfg.max_tokens if max_tokens is None else max_tokens
        effective_timeout   = self.cfg.timeout_s  if timeout_s is None else timeout_s

        body: dict[str, Any] = {
            "messages": messages,
            "temperature": self.cfg.temperature,
            "max_tokens": effective_max_tokens,
            "cache_prompt": self.cfg.cache_prompt,
            "stream": False,
        }
        if slot_id is not None:
            body["id_slot"] = slot_id
        if tools:
            body["tools"] = tools
            body["tool_choice"] = "auto"

        if on_io is not None:
            _safe_io(
                on_io, {"phase": "request", "url": self.cfg.endpoint, "body": body}
            )

        try:
            payload = self._opener(
                {"__url__": self.cfg.endpoint, "body": body}, effective_timeout
            )
        except (urllib.error.URLError, OSError, json.JSONDecodeError) as e:
            if on_io is not None:
                _safe_io(on_io, {"phase": "response", "payload": None, "error": str(e)})
            raise LlmError(str(e)) from e

        if on_io is not None:
            _safe_io(on_io, {"phase": "response", "payload": payload, "error": None})

        choices = payload.get("choices") or []
        if not choices:
            raise LlmError("response contained no choices")
        choice = choices[0]
        message = choice.get("message") or {}

        # `reasoning_content` is deliberately never read.
        content = (message.get("content") or "").strip()
        tool_calls = []
        for call in message.get("tool_calls") or []:
            function = call.get("function") or {}
            tool_calls.append(
                ToolCall(
                    id=call.get("id") or "",
                    name=function.get("name") or "",
                    arguments=function.get("arguments") or "",
                )
            )

        return LlmResponse(
            content=content,
            tool_calls=tool_calls,
            finish_reason=choice.get("finish_reason") or "stop",
        )


class ConversationEngine:
    """Conversation history with rolling-summary compaction and block-trim fallback.

    State:
        summary  — compressed text of evicted turns (empty string when none).
        history  — verbatim recent window (append-only between compactions).

    Prompt assembly:
        [system: system_prompt]
        [system: "Earlier conversation summary:\\n{summary}"]  # only when summary is set
        *history

    Lock discipline: self._lock (RLock) guards history/summary mutations and reads.
    It is NEVER held across an HTTP call or a finalize callback.
    """

    def __init__(
        self,
        cfg,
        client,
        system_prompt: str,
        compaction_system_prompt: str | None = None,
        max_tokens: int | None = None,
        timeout_s: float | None = None,
        notify: Callable[..., None] | None = None,
        slot_id: int | None = None,
    ):
        self.cfg                      = cfg
        self.client                   = client
        self.system                   = system_prompt
        self._compaction_system_prompt = compaction_system_prompt
        self._max_tokens              = max_tokens
        self._timeout_s               = timeout_s
        self.notify                   = notify
        self._slot_id                 = slot_id

        self.history: list[dict] = []
        self.summary: str        = ""
        self._generation         = 0
        self._lock               = threading.RLock()

    def reset(self) -> None:
        """Forget the conversation (volatile by design)."""
        with self._lock:
            self.history     = []
            self.summary     = ""
            self._generation += 1

    def last_reply(self) -> str:
        """Return the most recent assistant turn, or "" when there is none."""
        with self._lock:
            for msg in reversed(self.history):
                if msg["role"] == "assistant":
                    return msg["content"]
        return ""

    def send(
        self,
        user_text: str,
        on_io: Callable[[dict], None] | None = None,
        finalize: Callable[[str, str], str | None] | None = None,
    ) -> str:
        """Append a user turn, ask the model, append and return the reply.

        Args:
            user_text: The user turn to append.
            on_io: Optional per-call IO hook, forwarded to the client only when set.
            finalize: Optional gate called as `finalize(reply, finish_reason)` after
                the staleness check. Returning None rejects the turn: history and
                summary are restored from the pre-call snapshot and "" is returned.
                Returning a string appends THAT string as the assistant turn.
        """
        with self._lock:
            generation       = self._generation
            snapshot_history = list(self.history)
            snapshot_summary = self.summary
            self.history.append({"role": "user", "content": user_text})

        self._maybe_compact()

        with self._lock:
            messages = [{"role": "system", "content": self.system}]
            if self.summary:
                messages.append(
                    {"role": "system", "content": f"Earlier conversation summary:\n{self.summary}"}
                )
            messages += self.history

        # HTTP call — no lock held. Optional kwargs are included only when set so
        # that old fakes (which accept neither) keep working.
        chat_kwargs: dict[str, Any] = {}
        if on_io is not None:
            chat_kwargs["on_io"] = on_io
        if self._slot_id is not None:
            chat_kwargs["slot_id"] = self._slot_id
        response = self.client.chat(
            messages,
            max_tokens=self._max_tokens,
            timeout_s=self._timeout_s,
            **chat_kwargs,
        )
        reply    = response.content
        finish_reason = response.finish_reason
        truncated = finish_reason == "length"

        with self._lock:
            if generation != self._generation:
                # The session was reset while the model was thinking: the reply belongs
                # to a conversation that no longer exists, so drop it.
                return ""

        # Finalize callback — no lock held, so a reset() inside it is not undone.
        if finalize is not None:
            final = finalize(reply, finish_reason)
            with self._lock:
                if generation != self._generation:
                    # A reset() intervened (possibly inside finalize itself); the session
                    # is already clean — restoring would undo the reset.
                    return ""
                if final is None:
                    self.history = snapshot_history
                    self.summary = snapshot_summary
                    return ""
                if final is not None:
                    self.history.append({"role": "assistant", "content": final})
                return final

        with self._lock:
            if generation != self._generation:
                return ""
            if reply:
                self.history.append({"role": "assistant", "content": reply})

        if truncated and self.notify is not None:
            self.notify("truncated", message="response truncated")

        return reply

    def _estimate_tokens(self) -> int:
        # Approximate: 4 chars ≈ 1 token (good enough for budget gating).
        characters = (
            len(self.system)
            + len(self.summary)
            + sum(len(m["content"]) for m in self.history)
        )
        return characters // 4

    def _trim(self) -> None:
        block = max(1, self.cfg.trim_block_turns) * 2
        while self._estimate_tokens() > self.cfg.prompt_budget and self.history:
            del self.history[: min(block, len(self.history))]

    def _maybe_compact(self) -> None:
        """Compact history when the prompt approaches the budget.

        Compaction path (compaction_enabled=True):
          1. Return early if still under the trigger threshold.
          2. If nothing to evict (window alone is over budget), fall back to _trim().
          3. Snapshot evicted/recent under the lock, release, call the summarizer.
          4. Re-acquire, re-check generation; apply or fall back to _trim().

        Legacy path (compaction_enabled=False):
          Trim under the lock (history stays bounded).
        """
        cfg = self.cfg

        if not cfg.compaction_enabled:
            with self._lock:
                self._trim()
            return

        with self._lock:
            if self._estimate_tokens() <= cfg.compaction_trigger_ratio * cfg.prompt_budget:
                return

            keep = max(0, cfg.compaction_keep_turns) * 2
            if len(self.history) <= keep:
                # The recent window alone is over budget; no turns to evict.
                # Fall back to block trimming (Oracle R5).
                self._trim()
                return

            evicted    = list(self.history[:-keep])
            recent     = list(self.history[-keep:])
            generation = self._generation
            previous   = self.summary

        # Summarizer HTTP call — no lock held.
        new_summary = self._summarize(previous, evicted)

        with self._lock:
            if generation != self._generation:
                # A reset() intervened; the session is already clean.
                return

            if new_summary is None:
                # Summarizer failed; fall back to block trimming (Oracle R7/R8).
                self._trim()
                return

            self.summary = new_summary
            self.history = recent

        # Notify outside the lock to avoid re-entry issues.
        if self.notify is not None:
            self.notify("compacted", turns=len(evicted))

    def _summarize(self, previous: str, evicted: list[dict]) -> str | None:
        """Summarize evicted turns, optionally folding over a previous summary.

        Bounded input: if the combined text exceeds compaction_input_tokens * 4
        characters, split evicted into chronological chunks and fold sequentially.

        Returns the summary string, or None on any failure.
        """
        cfg = self.cfg
        char_limit = cfg.compaction_input_tokens * 4

        # Build chunks that each fit within the character limit.
        chunks: list[list[dict]] = []
        current_chunk: list[dict] = []
        current_chars = 0

        for msg in evicted:
            msg_chars = len(msg["content"])
            if current_chunk and current_chars + msg_chars > char_limit:
                chunks.append(current_chunk)
                current_chunk = [msg]
                current_chars = msg_chars
            else:
                current_chunk.append(msg)
                current_chars += msg_chars

        if current_chunk:
            chunks.append(current_chunk)

        # Fold sequentially: s = summarize_once(s, chunk) for each chunk.
        running = previous
        for chunk in chunks:
            result = self._summarize_once(running, chunk)
            if result is None:
                return None
            running = result

        return running if running != previous else None

    def _summarize_once(self, previous: str, chunk: list[dict]) -> str | None:
        """Call the LLM to summarize one chunk of evicted turns.

        Returns the summary string, or None on failure (LlmError, empty content,
        or finish_reason == "length" — Oracle R8).
        """
        cfg = self.cfg

        lines = []
        if previous:
            lines.append(f"Previous summary:\n{previous}\n")
        lines.append("Conversation to summarize:")
        for msg in chunk:
            lines.append(f"{msg['role']}: {msg['content']}")
        user_text = "\n".join(lines)

        compaction_prompt = self._compaction_system_prompt or cfg.compaction_system_prompt
        messages = [
            {"role": "system", "content": compaction_prompt},
            {"role": "user",   "content": user_text},
        ]

        try:
            # Never pin the summarizer to a slot: compaction must not evict the
            # engine's own conversation slot.
            response = self.client.chat(
                messages,
                max_tokens=cfg.compaction_max_tokens,
                timeout_s=cfg.compaction_timeout_s,
            )
        except LlmError:
            logger.warning("Summarizer LLM call failed; falling back to block trim")
            return None

        if not response.content or response.finish_reason == "length":
            logger.warning(
                "Summarizer returned unusable reply (empty=%s, finish_reason=%r); "
                "falling back to block trim",
                not response.content,
                response.finish_reason,
            )
            return None

        return response.content


class ChatSession(ConversationEngine):
    """Back-compat chat session: a ConversationEngine on the chat system prompt."""

    def __init__(self, cfg, client, notify: Callable[..., None] | None = None):
        super().__init__(cfg, client, system_prompt=cfg.system_prompt_chat, notify=notify)


class BriefEngine:
    """Stateless-per-call brief engine that pins the last accepted brief.

    Unlike ``ConversationEngine`` this never accumulates history: each call
    sends the same static system prompt plus ONE user turn carrying the current
    brief inline (built by the analyzer). The previous brief is therefore the
    single source of truth, the prompt stays flat (~2-3k tokens) instead of
    growing linearly, the model never competes with stale copies of its own
    output, and there is nothing for compaction to evict or mangle.

    Lock discipline: ``self._lock`` guards the pinned brief and generation and
    is never held across the HTTP call or the finalize callback.
    """

    def __init__(
        self,
        cfg,
        client,
        system_prompt: str,
        max_tokens: int | None = None,
        timeout_s: float | None = None,
        slot_id: int | None = None,
        notify: Callable[..., None] | None = None,
    ):
        self.cfg = cfg
        self.client = client
        self.system = system_prompt
        self._max_tokens = max_tokens
        self._timeout_s = timeout_s
        self._slot_id = slot_id
        self.notify = notify
        self._brief = ""
        self._generation = 0
        self._lock = threading.RLock()

    def reset(self) -> None:
        """Invalidate any in-flight reply (the service owns the actual brief state)."""
        with self._lock:
            self._generation += 1

    def send(
        self,
        user_text: str,
        on_io: Callable[[dict], None] | None = None,
        finalize: Callable[[str, str], str | None] | None = None,
    ) -> str:
        """Ask the model to update the brief from one user turn.

        The engine is stateless with respect to brief content: the caller owns
        both the volatile narrative and the code-owned ledger, and passes them in
        via ``user_text``. ``finalize`` behaves as on ``ConversationEngine.send``:
        returning None rejects the turn; returning a string returns it as the
        accepted reply. A truncated (``finish_reason == "length"``) reply is still
        returned so the caller's parser can salvage it, but the caller must not
        treat a missing section as an empty fact list for a truncated reply.
        """
        with self._lock:
            generation = self._generation
            messages = [
                {"role": "system", "content": self.system},
                {"role": "user", "content": user_text},
            ]

        chat_kwargs: dict[str, Any] = {}
        if on_io is not None:
            chat_kwargs["on_io"] = on_io
        if self._slot_id is not None:
            chat_kwargs["slot_id"] = self._slot_id
        response = self.client.chat(
            messages,
            max_tokens=self._max_tokens,
            timeout_s=self._timeout_s,
            **chat_kwargs,
        )
        reply = response.content
        finish_reason = response.finish_reason

        with self._lock:
            if generation != self._generation:
                return ""

        if finalize is not None:
            final = finalize(reply, finish_reason)
            with self._lock:
                if generation != self._generation:
                    return ""
                return final or ""

        if finish_reason == "length" and self.notify is not None:
            self.notify("truncated", message="response truncated")

        return reply