Trixy7/plugins/nlp_llm/backends/llama_cpp.py

# -*- coding: utf-8 -*-
"""
llama-cpp-python Backend für NLP.

Verwendet llama-cpp-python für lokale LLM-Inferenz.
Optimiert für Raspberry Pi 5 mit 4GB RAM.
"""

from pathlib import Path
from typing import Any

from trixy_core.nlp import NLPConfig, NLPState
from trixy_core.utils.debug import pinfo, pdebug, perror, pwarn

from plugins.nlp_llm.provider import LLMNLPProvider


class LlamaCppNLPBackend(LLMNLPProvider):
    """
    llama-cpp-python Backend.

    Lädt und verwaltet lokale GGUF-Modelle.
    Unterstützt CPU und GPU (CUDA/Metal) Inferenz.
    """

    # Standard-Modellpfade
    DEFAULT_MODEL_PATHS = [
        Path("models/nlp"),
        Path("~/.local/share/trixy/models").expanduser(),
        Path("/usr/share/trixy/models"),
    ]

    def __init__(self) -> None:
        super().__init__()
        self._llm = None

    async def _initialize_backend(self, config: NLPConfig) -> bool:
        """
        Initialisiert das llama-cpp-python Backend.

        Args:
            config: NLP-Konfiguration

        Returns:
            True bei Erfolg
        """
        try:
            from llama_cpp import Llama
        except ImportError:
            perror("llama-cpp-python nicht installiert. Installieren mit: pip install llama-cpp-python")
            return False

        # Modellpfad ermitteln
        model_path = self._find_model_path(config)

        if model_path is None:
            perror(f"Modell nicht gefunden: {config.model_name}")
            pwarn("Verfügbare Modelle können mit 'trixy model list' angezeigt werden")
            return False

        pinfo(f"Lade Modell: {model_path}")

        try:
            # LLM laden mit optimierten Einstellungen für Pi 5
            self._llm = Llama(
                model_path=str(model_path),
                n_ctx=config.context_window,
                n_threads=config.num_threads,
                n_gpu_layers=-1 if config.use_gpu else 0,
                verbose=False,
                # Speicher-Optimierungen für Pi 5
                use_mmap=True,
                use_mlock=False,
            )

            pinfo(f"Modell geladen: {config.model_name}")
            return True

        except Exception as e:
            perror(f"Fehler beim Laden des Modells: {e}")
            return False

    def _find_model_path(self, config: NLPConfig) -> Path | None:
        """
        Findet den Pfad zum Modell.

        Args:
            config: NLP-Konfiguration

        Returns:
            Modellpfad oder None
        """
        # Expliziter Pfad in Konfiguration
        if config.model_path:
            path = Path(config.model_path)
            if path.exists():
                return path
            # Versuche mit .gguf Endung
            if not path.suffix:
                for suffix in [".gguf", ".bin"]:
                    with_suffix = path.with_suffix(suffix)
                    if with_suffix.exists():
                        return with_suffix

        # Suche in Standard-Pfaden
        model_name = config.model_name
        possible_names = [
            model_name,
            f"{model_name}.gguf",
            f"{model_name}.Q4_K_M.gguf",
            f"{model_name}.Q5_K_M.gguf",
            f"{model_name}.Q8_0.gguf",
        ]

        for base_path in self.DEFAULT_MODEL_PATHS:
            if not base_path.exists():
                continue

            for name in possible_names:
                full_path = base_path / name
                if full_path.exists():
                    return full_path

                # Suche in Unterverzeichnissen
                for subdir in base_path.iterdir():
                    if subdir.is_dir():
                        sub_path = subdir / name
                        if sub_path.exists():
                            return sub_path

        return None

    async def _generate(self, messages: list[dict[str, str]]) -> str:
        """
        Generiert eine Antwort vom Modell.

        Args:
            messages: Nachrichtenliste

        Returns:
            Generierte Antwort
        """
        if self._llm is None:
            raise RuntimeError("Modell nicht geladen")

        # Chat-Completion verwenden
        response = self._llm.create_chat_completion(
            messages=messages,
            temperature=self._config.temperature if self._config else 0.1,
            max_tokens=self._config.max_tokens if self._config else 256,
            stop=["</s>", "<|end|>", "<|eot_id|>"],
        )

        # Antwort extrahieren
        if response and "choices" in response and len(response["choices"]) > 0:
            content = response["choices"][0].get("message", {}).get("content", "")
            return content.strip()

        return ""

    async def _shutdown_backend(self) -> None:
        """Gibt Modell-Ressourcen frei."""
        if self._llm is not None:
            # llama-cpp-python hat keinen expliziten close()
            # Setze Referenz auf None für Garbage Collection
            self._llm = None
            pdebug("llama-cpp Modell entladen")

    async def health_check(self) -> bool:
        """Prüft ob das Modell funktionsfähig ist."""
        if self._llm is None:
            return False

        try:
            # Kurzer Test-Prompt
            response = self._llm.create_chat_completion(
                messages=[{"role": "user", "content": "Hi"}],
                max_tokens=5,
            )
            return response is not None
        except Exception:
            return False