UmbraAudio/umbraaudio/voices/lmstudio_gate.py

"""Queueing gate for LM Studio LLM requests.

LM Studio on this Mac can only serve one LLM at a time and the model
coordinator swaps between the Qwen TTS models and the LLM. LLM calls that
arrive while LM Studio is busy -- loading a model, swapping models, or
warming up -- previously failed with connection / HTTP errors. This module
serializes requests and waits until the server is available again, then lets
the queued request through (the exact behavior requested: *if the local LM
Studio is busy the request is queued and waits for availability, then
processed*).

Only requests aimed at a local LM Studio (:1234) are gated; OpenAI provider
clients pass straight through untouched.
"""

from __future__ import annotations

import json
import logging
import threading
import time
from collections.abc import Callable
from contextlib import contextmanager
from typing import Any
from urllib.error import URLError
from urllib.request import Request, urlopen

logger = logging.getLogger(__name__)

DEFAULT_POLL_INTERVAL = 1.0
DEFAULT_WAIT_TIMEOUT = 900.0

_ML_LOCAL_HOSTS = {"localhost", "127.0.0.1", "0.0.0.0", "::1"}
_PLACEHOLDER_MODELS = {"", "local-model"}


class LMStudioBusyError(RuntimeError):
    """Raised when LM Studio did not become available within the wait window."""


def _host_of(base_url: str) -> str:
    """Return the hostname of a base URL (best effort, no network)."""
    base = base_url.split("://", 1)[-1]
    base = base.split("/", 1)[0]
    host = base
    if ":" in base and not base.startswith("["):
        host = base.rsplit(":", 1)[0]
    return (host or "").lower().strip("[]")


def _is_local_lmstudio(base_url: str) -> bool:
    """Gate only requests aimed at the local LM Studio server."""
    host = _host_of(base_url)
    if host in _ML_LOCAL_HOSTS:
        return True
    # LM Studio is commonly on port 1234, even when reached via a custom host.
    return ":1234" in base_url


def _server_probe(base_url: str, model: str | None = None) -> bool:
    """True when the server answers /models and (if given) knows the model."""
    base = (base_url or "").rstrip("/")
    if not base:
        return True
    try:
        req = Request(
            f"{base}/models",
            headers={"Accept": "application/json"},
        )
        with urlopen(req, timeout=3) as resp:
            if resp.status != 200:
                return False
            payload = json.loads(resp.read().decode("utf-8"))
    except (URLError, OSError, ValueError):
        return False

    ids = [m.get("id", "") for m in payload.get("data", [])]
    if model in _PLACEHOLDER_MODELS:
        return True

    # LM Studio exposes loaded instances under the native /api/v1/models
    # endpoint; prefer it when the OpenAI-compatible listing is ambiguous.
    api = base
    for suffix in ("/v1",):
        if api.endswith(suffix):
            api = api[: -len(suffix)]
    try:
        req2 = Request(f"{api}/api/v1/models", headers={"Accept": "application/json"})
        with urlopen(req2, timeout=3) as resp2:
            if resp2.status == 200:
                native = json.loads(resp2.read().decode("utf-8"))
                for m in native.get("data", []):
                    if m.get("key") == model and m.get("loaded_instances"):
                        return True
    except (URLError, OSError, ValueError):
        pass

    return model in ids or not ids


class LMStudioGate:
    """Serialize + queue LLM requests until LM Studio is available.

    Thread-safe. At most one request is processed at a time; everyone else
    waits (FIFO-ish via the condition variable). Before a request runs, the
    gate polls ``available_fn`` until the server reports healthy.
    """

    def __init__(
        self,
        *,
        poll_interval: float = DEFAULT_POLL_INTERVAL,
        wait_timeout: float = DEFAULT_WAIT_TIMEOUT,
        available_fn: Callable[[str, str | None], bool] = _server_probe,
    ) -> None:
        self.poll_interval = poll_interval
        self.wait_timeout = wait_timeout
        self._available_fn = available_fn
        self._cond = threading.Condition()
        self._inflight = 0

    def wait_until_available(
        self,
        base_url: str,
        model: str | None = None,
        timeout: float | None = None,
    ) -> None:
        """Block until LM Studio answers, raising :class:`LMStudioBusyError`."""
        wait_for = timeout if timeout is not None else self.wait_timeout
        deadline = time.monotonic() + wait_for
        while True:
            if self._available_fn(base_url, model):
                return
            if time.monotonic() >= deadline:
                raise LMStudioBusyError(
                    f"LM Studio not available after {wait_for:.0f}s wait (base={base_url})"
                )
            time.sleep(self.poll_interval)

    @contextmanager
    def acquire(
        self,
        base_url: str,
        model: str | None = None,
        timeout: float | None = None,
    ) -> Any:
        """Queue the caller behind any in-flight request, then run when ready."""
        with self._cond:
            while self._inflight > 0:
                self._cond.wait()
            self._inflight += 1
        try:
            self.wait_until_available(base_url, model, timeout)
            yield
        finally:
            with self._cond:
                self._inflight -= 1
                self._cond.notify()


_GATE = LMStudioGate()


def _client_base(client: Any) -> str:
    try:
        return str(client.base_url)
    except Exception:
        return ""


def queued_chat_completion(
    client: Any,
    model: str,
    messages: list[dict[str, str]],
    **kwargs: Any,
) -> Any:
    """Run ``client.chat.completions.create`` behind the LM Studio gate.

    Requests to a local LM Studio are serialized and wait for server
    availability before processing; non-LM Studio providers are unchanged.
    """
    base_url = _client_base(client)
    if _is_local_lmstudio(base_url):
        with _GATE.acquire(base_url, model):
            return client.chat.completions.create(model=model, messages=messages, **kwargs)
    return client.chat.completions.create(model=model, messages=messages, **kwargs)