Trixy7/plugins/nlp_llm/backends/ollama.py

# -*- coding: utf-8 -*-
"""
Ollama Backend für NLP.

Verwendet Ollama für lokale LLM-Inferenz.
Ollama verwaltet Modelle und bietet eine REST-API.
"""

import asyncio
import json
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 OllamaNLPBackend(LLMNLPProvider):
    """
    Ollama Backend.

    Kommuniziert mit dem Ollama-Server über die REST-API.
    Ollama muss separat installiert und gestartet sein.
    """

    DEFAULT_HOST = "http://localhost:11434"

    def __init__(self) -> None:
        super().__init__()
        self._host = self.DEFAULT_HOST
        self._model_name = ""
        self._session = None

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

        Prüft Verbindung zum Ollama-Server und ob das Modell verfügbar ist.

        Args:
            config: NLP-Konfiguration

        Returns:
            True bei Erfolg
        """
        self._host = config.extra.get("ollama_host", self.DEFAULT_HOST)
        self._model_name = config.model_name

        # HTTP-Client erstellen
        try:
            import aiohttp
            self._session = aiohttp.ClientSession()
        except ImportError:
            perror("aiohttp nicht installiert. Installieren mit: pip install aiohttp")
            return False

        # Verbindung testen
        if not await self._check_connection():
            perror(f"Keine Verbindung zum Ollama-Server: {self._host}")
            pwarn("Stellen Sie sicher, dass Ollama läuft: ollama serve")
            return False

        # Prüfen ob Modell verfügbar
        if not await self._check_model():
            pwarn(f"Modell '{self._model_name}' nicht gefunden, versuche Pull...")
            if not await self._pull_model():
                perror(f"Modell konnte nicht geladen werden: {self._model_name}")
                return False

        pinfo(f"Ollama Backend initialisiert: {self._model_name}")
        return True

    async def _check_connection(self) -> bool:
        """Prüft Verbindung zum Ollama-Server."""
        try:
            async with self._session.get(f"{self._host}/api/tags", timeout=5) as resp:
                return resp.status == 200
        except Exception:
            return False

    async def _check_model(self) -> bool:
        """Prüft ob das Modell verfügbar ist."""
        try:
            async with self._session.get(f"{self._host}/api/tags") as resp:
                if resp.status != 200:
                    return False

                data = await resp.json()
                models = data.get("models", [])

                for model in models:
                    name = model.get("name", "")
                    # Prüfe exakten Namen oder ohne Tag
                    if name == self._model_name or name.split(":")[0] == self._model_name:
                        return True

                return False
        except Exception as e:
            pdebug(f"Fehler beim Prüfen des Modells: {e}")
            return False

    async def _pull_model(self) -> bool:
        """Lädt das Modell von Ollama herunter."""
        try:
            pinfo(f"Lade Modell: {self._model_name} (dies kann einige Zeit dauern)")

            async with self._session.post(
                f"{self._host}/api/pull",
                json={"name": self._model_name, "stream": False},
                timeout=None,  # Kein Timeout für Download
            ) as resp:
                if resp.status == 200:
                    pinfo(f"Modell geladen: {self._model_name}")
                    return True

                error = await resp.text()
                perror(f"Fehler beim Laden des Modells: {error}")
                return False

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

    async def _generate(self, messages: list[dict[str, str]]) -> str:
        """
        Generiert eine Antwort über die Ollama API.

        Args:
            messages: Nachrichtenliste

        Returns:
            Generierte Antwort
        """
        if self._session is None:
            raise RuntimeError("Ollama-Session nicht initialisiert")

        try:
            payload = {
                "model": self._model_name,
                "messages": messages,
                "stream": False,
                "options": {
                    "temperature": self._config.temperature if self._config else 0.1,
                    "num_predict": self._config.max_tokens if self._config else 256,
                    "num_ctx": self._config.context_window if self._config else 2048,
                    "num_thread": self._config.num_threads if self._config else 4,
                },
            }

            async with self._session.post(
                f"{self._host}/api/chat",
                json=payload,
                timeout=60,
            ) as resp:
                if resp.status != 200:
                    error = await resp.text()
                    raise RuntimeError(f"Ollama API Fehler: {error}")

                data = await resp.json()
                message = data.get("message", {})
                return message.get("content", "").strip()

        except asyncio.TimeoutError:
            raise RuntimeError("Ollama Timeout - Anfrage dauerte zu lange")
        except Exception as e:
            raise RuntimeError(f"Ollama Fehler: {e}")

    async def _shutdown_backend(self) -> None:
        """Schließt die HTTP-Session."""
        if self._session is not None:
            await self._session.close()
            self._session = None
            pdebug("Ollama-Session geschlossen")

    async def health_check(self) -> bool:
        """Prüft ob Ollama erreichbar ist."""
        return await self._check_connection() and await self._check_model()

    def supports_streaming(self) -> bool:
        """Ollama unterstützt Streaming."""
        return True

    async def generate_stream(self, messages: list[dict[str, str]]):
        """
        Generiert eine Antwort mit Streaming.

        Args:
            messages: Nachrichtenliste

        Yields:
            Text-Chunks
        """
        if self._session is None:
            raise RuntimeError("Ollama-Session nicht initialisiert")

        payload = {
            "model": self._model_name,
            "messages": messages,
            "stream": True,
            "options": {
                "temperature": self._config.temperature if self._config else 0.1,
                "num_predict": self._config.max_tokens if self._config else 256,
            },
        }

        async with self._session.post(
            f"{self._host}/api/chat",
            json=payload,
        ) as resp:
            async for line in resp.content:
                if line:
                    try:
                        data = json.loads(line)
                        content = data.get("message", {}).get("content", "")
                        if content:
                            yield content
                    except json.JSONDecodeError:
                        continue