InkFlow/backend/inkflow/analysis/llm/lmstudio_backend.py

"""Backend LLM via LM Studio (API OpenAI locale).

LM Studio sert indifferemment des modeles GGUF *et* MLX charges depuis sa GUI,
exposes sur un endpoint OpenAI-compatible (`http://127.0.0.1:1234/v1` par
defaut). InkFlow ne fait que parler HTTP : zero dependance native a compiler, et
le modele reste charge entre redemarrages d'InkFlow.

Caveat : `enable_thinking=False` (coupe la pensee des modeles hybrides cote mlx)
n'est pas pilotable de facon fiable via l'API ; le template embarque decide. En
mode non-raisonnement on prend le `content` final et on le strip de toute facon.
"""
from __future__ import annotations

import os
from typing import Callable, Optional

from .base import LLMBackend
from ._text import _has_complete_json, _strip_reasoning
from ...settings import get_settings


def list_models(base_url: str) -> list[dict]:
    """Liste les modeles LLM *telecharges* dans LM Studio (charges ou non).

    Utilise l'API REST native (`/api/v0/models`) et non `/v1/models` (qui ne
    renvoie que les modeles deja charges) : on peut ainsi proposer n'importe quel
    modele telecharge ; LM Studio le charge a la volee (JIT) a la 1re requete.
    Renvoie [{id, state, type}] filtre sur les LLM/VLM. Leve en cas d'echec.
    """
    import httpx

    root = base_url.rstrip("/")
    if root.endswith("/v1"):
        root = root[:-len("/v1")]
    resp = httpx.get(f"{root}/api/v0/models", timeout=5.0)
    resp.raise_for_status()
    data = resp.json().get("data", [])
    return [
        {"id": m.get("id"), "state": m.get("state"), "type": m.get("type")}
        for m in data if m.get("type") in ("llm", "vlm")
    ]


class LMStudioBackend(LLMBackend):
    """Moteur LM Studio : client OpenAI pointe sur le serveur local."""

    name = "lmstudio"

    def __init__(self, model_ref: str):
        super().__init__(model_ref)
        self._client = None
        self._model = None  # resolu paresseusement (model_ref vide -> modele actif)

    def _ensure_client(self):
        if self._client is None:
            try:
                from openai import OpenAI
            except ImportError as exc:  # noqa: BLE001
                raise RuntimeError(
                    "Le paquet `openai` est requis pour le backend LM Studio "
                    "(pip install -e backend)."
                ) from exc
            base_url = get_settings().lmstudio_base_url
            # api_key factice : LM Studio n'authentifie pas, mais le SDK l'exige.
            self._client = OpenAI(base_url=base_url, api_key="lm-studio")
        return self._client

    def _resolve_model(self, client) -> str:
        """Renvoie le nom de modele a utiliser (model_ref, ou 1er modele charge)."""
        if self.model_ref:
            return self.model_ref
        if self._model is None:
            try:
                models = client.models.list()
            except Exception as exc:  # noqa: BLE001
                raise self._connection_error(exc) from exc
            ids = [m.id for m in getattr(models, "data", [])]
            if not ids:
                raise RuntimeError(
                    "Aucun modele charge dans LM Studio : charge un modele "
                    "(GGUF ou MLX) dans l'app avant de lancer l'analyse."
                )
            self._model = ids[0]
        return self._model

    def _connection_error(self, exc: Exception) -> RuntimeError:
        url = get_settings().lmstudio_base_url
        return RuntimeError(
            f"LM Studio injoignable sur {url} — lance l'application et active le "
            f"serveur local (onglet Developer > Start Server). Detail: {exc}"
        )

    def _sampling(self, max_tokens: int, temperature: float) -> dict:
        """Params de sampling a transmettre a LM Studio.

        Par defaut (`lmstudio_defer_config=True`) : dict VIDE -> on delegue
        temperature ET plafond de tokens a la config du modele charge dans LM
        Studio (ne pas imposer `max_tokens` evite de tronquer la reponse, ce qui
        cassait les modeles a raisonnement). Le contexte est de toute facon gere
        au chargement cote LM Studio. Si l'utilisateur desactive la delegation,
        on reimpose les reglages "Generation Gemma" d'InkFlow.
        """
        if get_settings().lmstudio_defer_config:
            return {}
        return {"temperature": temperature, "max_tokens": max_tokens}

    def complete(
        self,
        messages: list[dict],
        *,
        max_tokens: int,
        temperature: float,
        reasoning: bool,
        token_sink: Optional[Callable[[str], None]] = None,
    ) -> str:
        client = self._ensure_client()
        model = self._resolve_model(client)
        sampling = self._sampling(max_tokens, temperature)
        # Prefill optionnel de la reponse assistant (INKFLOW_LMSTUDIO_PREFILL) :
        # ex. "<think></think>" force les modeles distilles a raisonnement (Qwen)
        # a sauter la pensee (seul levier efficace quand enable_thinking/_no_think
        # sont ignores). Le modele continue a partir du prefill -> JSON direct.
        prefill = os.environ.get("INKFLOW_LMSTUDIO_PREFILL")
        if prefill:
            messages = messages + [{"role": "assistant", "content": prefill}]
        from openai import APIConnectionError

        # LM Studio separe la pensee (`reasoning_content`) de la reponse finale
        # (`content`, deja propre). On ne renvoie QUE `content` : la facade en
        # extrait le JSON. La pensee n'est diffusee qu'au `token_sink` (affichage
        # --stream) ; l'inclure dans le retour risquerait de capter un JSON
        # d'exemple present dans le raisonnement. Pour les modeles qui mettent au
        # contraire la pensee INLINE dans `content` (<think>...), la facade la
        # retire via _strip_reasoning quand reasoning=True.
        if token_sink is not None or reasoning:
            content_parts: list[str] = []
            try:
                stream = client.chat.completions.create(
                    model=model, messages=messages, stream=True, **sampling,
                )
                for chunk in stream:
                    if not chunk.choices:
                        continue
                    delta = chunk.choices[0].delta
                    rc = getattr(delta, "reasoning_content", None)
                    if rc and token_sink is not None:
                        token_sink(rc)  # pensee : affichage seulement
                    piece = delta.content or ""
                    if piece:
                        content_parts.append(piece)
                        if token_sink is not None:
                            token_sink(piece)
                    # Arret anticipe en mode raisonnement : des que la reponse
                    # finale (content) contient un JSON complet, inutile de
                    # continuer (certains modeles divaguent ensuite).
                    if reasoning and piece and ("}" in piece or "]" in piece):
                        if _has_complete_json(_strip_reasoning("".join(content_parts))):
                            break
            except APIConnectionError as exc:
                raise self._connection_error(exc) from exc
            if token_sink is not None:
                token_sink("\n")  # separe les generations successives
            return "".join(content_parts)

        try:
            resp = client.chat.completions.create(
                model=model, messages=messages, **sampling,
            )
        except APIConnectionError as exc:
            raise self._connection_error(exc) from exc
        return resp.choices[0].message.content or ""