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

"""Single-lane inference queue for serialized LLM work.

All LLM calls (chat, command execution, compaction) run through one worker thread,
so there are no concurrent requests to the server.
"""

import collections
import logging
import threading
from dataclasses import dataclass
from typing import Callable

logger = logging.getLogger(__name__)


@dataclass
class InferenceJob:
    """A unit of work for the inference queue."""

    name: str  # "chat" | "command" | "summary"
    run: Callable[[], None]  # the work to do
    priority: int = 0  # lower values run first; commands use -1


class InferenceQueue:
    """Single worker thread owning every LLM inference call."""

    def __init__(self, size: int = 8):
        """Create the queue.

        Args:
            size: Max queued jobs (bounded backpressure).
        """
        self.size = size
        self._queue: collections.deque[InferenceJob] = collections.deque()
        self._in_flight: InferenceJob | None = None
        self._lock = threading.Lock()
        self._cond = threading.Condition(self._lock)
        self._stop = threading.Event()
        self._worker: threading.Thread | None = None

    def start(self) -> None:
        """Spawn the worker thread."""
        if self._worker is not None:
            return
        self._worker = threading.Thread(
            target=self._worker_loop, name="assistant-inference", daemon=True
        )
        self._worker.start()

    def stop(self) -> None:
        """Discard queued jobs, let in-flight finish, join the worker."""
        with self._cond:
            self._queue.clear()
            self._stop.set()
            self._cond.notify_all()
        if self._worker is not None:
            self._worker.join(timeout=1.0)

    def submit(self, job: InferenceJob) -> bool:
        """Queue a job for execution.

        Returns:
            False only when a priority job (priority < 0) is rejected due to overflow.
            Normal jobs drop the oldest normal job on overflow and return True.
        """
        with self._cond:
            if self._stop.is_set():
                return False

            # Capacity includes the in-flight job (if any) and the queue
            capacity_used = (1 if self._in_flight else 0) + len(self._queue)

            # Check if at capacity
            if capacity_used >= self.size:
                if job.priority < 0 or not self._drop_oldest_normal():
                    return False

            # Insert the job: priority jobs at front, normal jobs at back
            if job.priority < 0:
                self._queue.appendleft(job)
            else:
                self._queue.append(job)

            self._cond.notify_all()
        return True

    def drop(self, predicate: Callable[[InferenceJob], bool]) -> int:
        """Remove queued jobs matching the predicate.

        Does not touch the in-flight job. Returns the count of jobs removed.
        """
        with self._cond:
            to_drop = [j for j in self._queue if predicate(j)]
            for job in to_drop:
                self._queue.remove(job)
            return len(to_drop)

    def depth(self) -> int:
        """Return the number of queued jobs (not counting in-flight)."""
        with self._lock:
            return len(self._queue)

    def _drop_oldest_normal(self) -> bool:
        """Remove the oldest normal (priority >= 0) job; False if none exists.

        Assumes the lock is held.
        """
        for i, job in enumerate(self._queue):
            if job.priority >= 0:
                del self._queue[i]
                return True
        return False

    def _worker_loop(self) -> None:
        """Worker thread: dequeue and execute jobs until stop."""
        while True:
            with self._cond:
                while not self._queue and not self._stop.is_set():
                    self._cond.wait()

                if self._stop.is_set() and not self._queue:
                    break

                if not self._queue:
                    continue

                job = self._queue.popleft()
                self._in_flight = job

            # Release the lock during execution so a concurrent drop() can proceed.
            try:
                job.run()
            except Exception as e:  # noqa: BLE001 - one bad job must not kill the lane
                logger.error("Inference job %r failed: %s", job.name, e)
            finally:
                with self._lock:
                    self._in_flight = None