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

"""Facade LLM pour l'analyse de texte (anciennement `Gemma`).

Charge un backend pluggable (mlx par defaut, ou LM Studio) selon les reglages et
expose `generate` / `generate_json` consommes par tout le pipeline. Toute la
logique agnostique du moteur vit ici : calcul des parametres depuis les Settings,
retrait de la chaine de pensee (modeles a raisonnement) et `generate_json`
tolerant qui extrait le premier objet/array JSON valide de la sortie du modele.
"""
from __future__ import annotations

from typing import Any, Optional

from ...settings import Settings, get_settings
from ._text import _extract_json, _strip_reasoning
from .factory import get_llm_backend

# Hook de streaming optionnel. Si defini, `generate()` diffuse chaque morceau de
# texte AU FIL de la generation (pensee comprise, avant tout nettoyage) en
# appelant ce callback. Utilise par `inkflow benchmark --stream` pour voir les
# tokens en temps reel. None -> generation par lot classique (plus rapide).
_TOKEN_SINK: Optional[Any] = None


def set_token_sink(callback) -> None:
    """Definit (ou retire avec None) le callback de streaming des tokens."""
    global _TOKEN_SINK
    _TOKEN_SINK = callback


def _model_ref_for(backend: str, settings: Settings) -> str:
    """Reference de modele par defaut pour un backend donne."""
    if backend == "lmstudio":
        return settings.lmstudio_model
    return settings.gemma_model


class LLM:
    """Petite facade multi-backend pour piloter le LLM d'analyse."""

    def __init__(self, model_id: Optional[str] = None, backend: Optional[str] = None):
        settings = get_settings()
        self.backend_name = backend or settings.gemma_backend
        self.model_ref = model_id or _model_ref_for(self.backend_name, settings)
        self._backend = None

    def _ensure_loaded(self) -> None:
        if self._backend is None:
            self._backend = get_llm_backend(self.backend_name, self.model_ref)

    def generate(
        self,
        prompt: str,
        *,
        system: Optional[str] = None,
        max_tokens: Optional[int] = None,
        temperature: Optional[float] = None,
    ) -> str:
        """Genere une reponse texte a partir d'un prompt (template de chat).

        `max_tokens`/`temperature` non fournis -> valeurs des reglages courants.
        """
        self._ensure_loaded()
        settings = get_settings()
        if max_tokens is None:
            max_tokens = settings.gemma_max_tokens
            # En mode raisonnement, plafond dedie (garde-fou anti-boucle) ; la
            # generation s'arrete de toute facon des que le JSON post-pensee est
            # complet (cf. arret anticipe des backends).
            if settings.gemma_reasoning:
                max_tokens = max(max_tokens, settings.gemma_reasoning_max_tokens)
        if temperature is None:
            temperature = settings.gemma_temperature
        # Decodage glouton (temp 0) + raisonnement = boucles de pensee sans fin.
        # On force un echantillonnage minimal en mode raisonnement.
        if settings.gemma_reasoning and temperature == 0.0:
            temperature = settings.gemma_reasoning_temperature

        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})

        raw = self._backend.complete(
            messages,
            max_tokens=max_tokens,
            temperature=temperature,
            reasoning=settings.gemma_reasoning,
            token_sink=_TOKEN_SINK,
        )
        # Retire la chaine de pensee des modeles a raisonnement (sinon des
        # fragments de la "pensee" parasitent l'extraction JSON en aval).
        if settings.gemma_reasoning:
            return _strip_reasoning(raw)
        return raw

    def generate_json(
        self,
        prompt: str,
        *,
        system: Optional[str] = None,
        max_tokens: Optional[int] = None,
        temperature: Optional[float] = None,
        retries: int = 1,
    ) -> Any:
        """Genere puis parse un JSON. Reessaie en cas d'echec de parsing.

        `max_tokens`/`temperature` non fournis -> valeurs des reglages courants.
        """
        last_err: Optional[Exception] = None
        for attempt in range(retries + 1):
            raw = self.generate(
                prompt, system=system, max_tokens=max_tokens,
                temperature=temperature if attempt == 0 else 0.0,
            )
            try:
                return _extract_json(raw)
            except Exception as exc:  # noqa: BLE001
                last_err = exc
        raise ValueError(f"Reponse JSON invalide apres {retries + 1} essais: {last_err}")