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

"""Backend LLM mlx-lm (Apple Silicon) — moteur historique d'InkFlow.

Charge un modele mlx-community paresseusement (une fois par process, cache LRU)
et genere via le template de chat du tokenizer. Comportement strictement
identique a l'ancien `Gemma.generate` : controle fin de `enable_thinking`,
streaming avec arret anticipe des que le JSON post-pensee est complet.
"""
from __future__ import annotations

import json
from functools import lru_cache
from typing import Callable, Optional

from .base import LLMBackend
from ._text import _REASONING_END_MARKERS, _has_complete_json, _strip_reasoning


@lru_cache(maxsize=2)
def _load(model_id: str):
    # Import paresseux : evite de charger mlx tant qu'on n'analyse pas.
    from mlx_lm import load
    return load(model_id)


def _resolve_chat_template(model_id: str, tokenizer) -> Optional[str]:
    """Renvoie un template de chat a passer explicitement, ou None.

    Certaines conversions (Mistral recents...) logent leur template dans un
    fichier `chat_template.jinja` que le downloader de mlx-lm n'embarque pas
    toujours : `tokenizer.chat_template` est alors vide et `apply_chat_template`
    echoue. On recupere alors le fichier officiel du repo. None si le tokenizer
    possede deja un template (cas courant) ou si aucun n'est disponible.
    """
    if getattr(tokenizer, "chat_template", None):
        return None
    from pathlib import Path

    from huggingface_hub import hf_hub_download
    # Selon les conversions : fichier Jinja brut, ou JSON {"chat_template": ...}.
    for fname in ("chat_template.jinja", "chat_template.json"):
        try:
            text = Path(hf_hub_download(model_id, fname)).read_text(encoding="utf-8")
        except Exception:  # noqa: BLE001 — fichier absent, on tente le suivant
            continue
        if fname.endswith(".json"):
            data = json.loads(text)
            return data.get("chat_template") if isinstance(data, dict) else None
        return text
    return None  # aucun template dispo -> apply_chat_template levera une erreur claire


class MLXBackend(LLMBackend):
    """Moteur mlx-lm : modeles mlx-community (HuggingFace) sur Apple Silicon."""

    name = "mlx"

    def __init__(self, model_ref: str):
        super().__init__(model_ref)
        self._model = None
        self._tokenizer = None
        self._chat_template = None   # template recupere si absent du tokenizer

    def _ensure_loaded(self) -> None:
        if self._model is None:
            self._model, self._tokenizer = _load(self.model_ref)
            self._chat_template = _resolve_chat_template(
                self.model_ref, self._tokenizer)

    def complete(
        self,
        messages: list[dict],
        *,
        max_tokens: int,
        temperature: float,
        reasoning: bool,
        token_sink: Optional[Callable[[str], None]] = None,
    ) -> str:
        self._ensure_loaded()
        from mlx_lm.sample_utils import make_sampler

        # Modeles hybrides (Qwen3...) : hors mode raisonnement, on DESACTIVE la
        # pensee via enable_thinking=False -> JSON direct, bien plus rapide. Avec
        # raisonnement, on laisse penser puis la facade retire la pensee. Ce
        # kwarg est ignore par les templates qui ne l'utilisent pas (Gemma...).
        template_kwargs = {}
        if not reasoning:
            template_kwargs["enable_thinking"] = False
        formatted = self._tokenizer.apply_chat_template(
            messages, add_generation_prompt=True, tokenize=False,
            chat_template=self._chat_template,  # None -> celui du tokenizer
            **template_kwargs,
        )
        sampler = make_sampler(temp=temperature)
        # On streame (token par token) si : un sink est branche (--stream) OU on
        # est en mode raisonnement (pour pouvoir s'arreter des que la reponse est
        # prete, sans subir les boucles de pensee sans fin). Sinon, lot rapide.
        if token_sink is not None or reasoning:
            from mlx_lm import stream_generate
            parts = []
            seen_end = False  # marqueur de fin de pensee rencontre
            for resp in stream_generate(
                self._model, self._tokenizer, prompt=formatted,
                max_tokens=max_tokens, sampler=sampler,
            ):
                parts.append(resp.text)
                if token_sink is not None:
                    token_sink(resp.text)
                # Arret anticipe : une fois la pensee close, des que le JSON
                # post-pensee est complet, inutile de continuer a generer.
                if reasoning and ("}" in resp.text or "]" in resp.text):
                    buf = "".join(parts)
                    if not seen_end:
                        seen_end = any(mk in buf for mk in _REASONING_END_MARKERS)
                    if seen_end and _has_complete_json(_strip_reasoning(buf)):
                        break
            if token_sink is not None:
                token_sink("\n")  # separe les generations successives
            return "".join(parts)
        from mlx_lm import generate
        return generate(
            self._model,
            self._tokenizer,
            prompt=formatted,
            max_tokens=max_tokens,
            sampler=sampler,
            verbose=False,
        )