NexusOS/synapse/ollama_manager.py

import asyncio
import json
import logging
import re
import subprocess
import time
import httpx
import os
import signal
from pathlib import Path

from .nexus_config import settings

OLLAMA_PORT = 11434

_log = logging.getLogger("nexus.ollama")

# --- Global Singleton Instance ---
_ollama_manager = None

# Bundled binary ships alongside the project; fall back to system PATH
_BUNDLED_OLLAMA = Path(__file__).resolve().parent.parent / "ollama" / "bin" / "ollama"


def _ollama_bin() -> str:
    """Return path to the Ollama executable, preferring the bundled copy."""
    if _BUNDLED_OLLAMA.exists():
        return str(_BUNDLED_OLLAMA)
    return "ollama"


def _best_vulkan_device() -> tuple[int, str]:
    """
    Parse `vulkaninfo --summary` and return (device_index, device_name) for the
    best Vulkan compute device. Prefers discrete GPUs over integrated ones, and
    AMD/NVIDIA vendor IDs over Intel — so a Radeon is chosen over an Intel iGPU
    even when the iGPU appears first in the device list.
    """
    try:
        r = subprocess.run(
            ["vulkaninfo", "--summary"], capture_output=True, text=True, timeout=5,
        )
        if r.returncode != 0:
            return 0, "GPU"

        devices: list[dict] = []
        current: dict = {}
        for line in r.stdout.splitlines():
            line = line.strip()
            if line.startswith("GPU") and line.endswith(":"):
                if current:
                    devices.append(current)
                raw_idx = line[3:-1]
                current = {
                    "index": int(raw_idx) if raw_idx.isdigit() else len(devices),
                    "name": "GPU",
                    "type": "",
                    "vendor": "",
                }
            elif "=" in line:
                key, _, val = line.partition("=")
                key, val = key.strip(), val.strip()
                if key == "deviceName":
                    current["name"] = val
                elif key == "deviceType":
                    current["type"] = val.upper()
                elif key == "vendorID":
                    current["vendor"] = val.lower()

        if current:
            devices.append(current)

        if not devices:
            return 0, "GPU"

        def _score(d: dict) -> tuple:
            # Discrete beats everything; integrated is last resort
            type_score = 2 if "DISCRETE" in d["type"] else (0 if "INTEGRATED" in d["type"] else 1)
            # AMD (0x1002) and NVIDIA (0x10de) preferred over Intel (0x8086)
            vendor_score = 1 if any(v in d["vendor"] for v in ("0x1002", "0x10de")) else 0
            return (type_score, vendor_score)

        best = max(devices, key=_score)
        return best["index"], best["name"]

    except Exception:
        return 0, "GPU"


def _detect_gpu_backend() -> tuple[str, dict]:
    """
    Probe available GPU compute backends.
    Returns (label, env_overrides) where env_overrides is merged into the
    Ollama subprocess environment before launch.
    Priority: CUDA > Vulkan > ROCm > CPU.
    """
    # NVIDIA CUDA — preferred when both GPU and CUDA drivers are present
    try:
        r = subprocess.run(
            ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
            capture_output=True, text=True, timeout=5,
        )
        if r.returncode == 0:
            name = r.stdout.strip().splitlines()[0]
            _log.info("GPU backend: CUDA (%s)", name)
            return f"cuda ({name})", {}  # Ollama auto-detects CUDA
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass

    # Vulkan — works on AMD, Intel, and NVIDIA without a full CUDA/ROCm stack
    vulkan_ok = False
    try:
        r = subprocess.run(
            ["vulkaninfo", "--summary"], capture_output=True, text=True, timeout=5,
        )
        vulkan_ok = r.returncode == 0
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass

    if not vulkan_ok:
        # Fall back to checking for ICD loader files without the vulkaninfo tool
        icd_dirs = [
            Path("/usr/share/vulkan/icd.d"),
            Path("/etc/vulkan/icd.d"),
            Path(os.path.expanduser("~/.local/share/vulkan/icd.d")),
        ]
        try:
            vulkan_ok = any(p.is_dir() and any(p.iterdir()) for p in icd_dirs)
        except PermissionError:
            pass

    if vulkan_ok:
        idx, name = _best_vulkan_device()
        env_overrides: dict = {"OLLAMA_GPU": "vulkan"}
        if idx > 0:
            # Explicitly route Ollama to the discrete GPU when it isn't device 0
            env_overrides["GGML_VK_VISIBLE_DEVICES"] = str(idx)
        _log.info("GPU backend: Vulkan device %d (%s)", idx, name)
        return f"vulkan ({name})", env_overrides

    # AMD ROCm — fallback when Vulkan ICD is absent but ROCm stack is installed
    try:
        r = subprocess.run(
            ["rocm-smi", "--showproductname"], capture_output=True, text=True, timeout=5,
        )
        if r.returncode == 0:
            _log.info("GPU backend: ROCm")
            return "rocm", {}  # Ollama auto-detects ROCm
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass

    _log.info("GPU backend: CPU (no GPU acceleration detected)")
    return "cpu", {}


def _model_score(name: str) -> tuple:
    """Score a model name for auto-selection. Higher tuple = better."""
    lower = name.lower()
    match = re.search(r'(\d+(?:\.\d+)?)[bB]', lower)
    params = float(match.group(1)) if match else 7.0
    # Tier-break: known quality models get a small bonus
    quality = next(
        (i for i, prefix in enumerate(("llama3", "llama2", "mistral", "gemma", "phi", "qwen"), 1)
         if prefix in lower),
        0,
    )
    return (params, quality)


class OllamaManager:
    def __init__(self, runtime_dir=None):
        self.process = None
        self.running = False

        self.runtime_dir = Path(runtime_dir) if runtime_dir else Path(__file__).resolve().parent.parent / "runtime"
        (self.runtime_dir / "logs").mkdir(parents=True, exist_ok=True)

        self.log_file = self.runtime_dir / "logs" / "ollama.log"

        # Resolved at construction so host changes in settings take effect
        self._api_base = settings.ollama_host.rstrip("/")

        # Model selection cache
        self._model_cache: str | None = None
        self._model_cache_ts: float = 0.0

    def is_available(self):
        bin_path = _ollama_bin()
        try:
            subprocess.run([bin_path, "--version"], capture_output=True, check=True, timeout=5)
            return True
        except Exception:
            return False

    def is_running(self):
        try:
            r = httpx.get(f"{self._api_base}/api/tags", timeout=2.0)
            return r.status_code == 200
        except httpx.RequestError:
            return False

    def start(self):
        if not self.is_available():
            _log.warning("Ollama not found at %s; skipping startup", _ollama_bin())
            return False

        if self.is_running():
            _log.info("Ollama already running (%s)", self._api_base)
            self.running = True
            return True

        try:
            backend, gpu_env = _detect_gpu_backend()
            _log.info("Starting Ollama service via %s (backend: %s)...", _ollama_bin(), backend)
            env = os.environ.copy()
            env["OLLAMA_HOST"] = self._api_base
            env["OLLAMA_MODELS"] = str(Path(__file__).resolve().parent.parent / "models")
            env.update(gpu_env)

            with open(self.log_file, "w") as log:
                self.process = subprocess.Popen(
                    [_ollama_bin(), "serve"],
                    stdout=log,
                    stderr=subprocess.STDOUT,
                    start_new_session=True,
                    env=env,
                )

            for attempt in range(30):
                if self.is_running():
                    _log.info("Ollama service started (%s)", self._api_base)
                    self.running = True
                    return True
                time.sleep(1)
                if attempt % 5 == 0:
                    _log.info("Waiting for Ollama... (%ds)", attempt)

            _log.error("Ollama failed to start: timeout")
            return False

        except Exception as e:
            _log.exception("Ollama failed to start: %s", e)
            return False

    async def start_async(self):
        """Async-safe version of start() for use inside async startup handlers."""
        if not self.is_available():
            _log.warning("Ollama not found at %s; skipping startup", _ollama_bin())
            return False

        if self.is_running():
            _log.info("Ollama already running (%s)", self._api_base)
            self.running = True
            return True

        try:
            backend, gpu_env = _detect_gpu_backend()
            _log.info("Starting Ollama service via %s (backend: %s)...", _ollama_bin(), backend)
            env = os.environ.copy()
            env["OLLAMA_HOST"] = self._api_base
            env["OLLAMA_MODELS"] = str(Path(__file__).resolve().parent.parent / "models")
            env.update(gpu_env)

            with open(self.log_file, "w") as log:
                self.process = subprocess.Popen(
                    [_ollama_bin(), "serve"],
                    stdout=log,
                    stderr=subprocess.STDOUT,
                    start_new_session=True,
                    env=env,
                )

            for attempt in range(30):
                if self.is_running():
                    _log.info("Ollama service started (%s)", self._api_base)
                    self.running = True
                    return True
                await asyncio.sleep(1)
                if attempt % 5 == 0:
                    _log.info("Waiting for Ollama... (%ds)", attempt)

            _log.error("Ollama failed to start: timeout")
            return False

        except Exception as e:
            _log.exception("Ollama failed to start: %s", e)
            return False

    def stop(self):
        if self.process and self.running:
            try:
                _log.info("Stopping Ollama service...")
                os.killpg(os.getpgid(self.process.pid), signal.SIGTERM)
                self.process.wait(timeout=10)
                _log.info("Ollama service stopped")
            except Exception:
                try:
                    os.killpg(os.getpgid(self.process.pid), signal.SIGKILL)
                except Exception:
                    pass
            finally:
                self.running = False

    def get_status(self):
        if self.is_running():
            return "running"
        elif self.is_available():
            return "available"
        else:
            return "unavailable"

    async def generate(
        self,
        prompt: str,
        model: str = "mistral",
        stream: bool = False,
        system: str = "",
        **kwargs
    ):
        start = time.perf_counter()

        _log.debug("generate model=%s system=%r prompt=%.120s", model, system, prompt)

        try:
            if stream:
                return self._stream(prompt=prompt, model=model, system=system, start=start)
            else:
                async with httpx.AsyncClient(timeout=300.0) as client:
                    r = await client.post(
                        f"{self._api_base}/api/generate",
                        json={
                            "model": model,
                            "prompt": prompt,
                            "system": system,
                            "stream": False,
                        },
                    )

                elapsed = time.perf_counter() - start
                _log.info("generate completed model=%s status=%d elapsed=%.3fs", model, r.status_code, elapsed)
                r.raise_for_status()
                return r.json().get("response", "")

        except Exception as e:
            elapsed = time.perf_counter() - start
            _log.exception("generate error after %.3fs: %s", elapsed, e)
            return None

    async def _stream(self, prompt: str, model: str, system: str, start: float):
        """
        Async generator that streams token chunks from Ollama.
        Yields string chunks as they arrive.
        """
        try:
            async with httpx.AsyncClient(timeout=300.0) as client:
                async with client.stream(
                    "POST",
                    f"{self._api_base}/api/generate",
                    json={
                        "model": model,
                        "prompt": prompt,
                        "system": system,
                        "stream": True,
                    },
                ) as response:
                    response.raise_for_status()
                    async for line in response.aiter_lines():
                        if not line.strip():
                            continue
                        try:
                            data = json.loads(line)
                            token = data.get("response", "")
                            if token:
                                yield token
                            if data.get("done", False):
                                elapsed = time.perf_counter() - start
                                _log.info("generate stream completed model=%s elapsed=%.3fs", model, elapsed)
                                break
                        except Exception:
                            continue

        except Exception as e:
            elapsed = time.perf_counter() - start
            _log.exception("generate stream error after %.3fs: %s", elapsed, e)
            return


    async def chat(
        self,
        messages: list,
        model: str = "mistral",
        stream: bool = False,
        temperature: float | None = None,
        **kwargs,
    ):
        """Multi-turn chat via /api/chat (accepts a messages array with roles)."""
        start = time.perf_counter()
        try:
            if stream:
                return self._chat_stream(messages=messages, model=model, temperature=temperature, start=start)
            else:
                body: dict = {"model": model, "messages": messages, "stream": False}
                if temperature is not None:
                    body["options"] = {"temperature": temperature}
                async with httpx.AsyncClient(timeout=300.0) as client:
                    r = await client.post(f"{self._api_base}/api/chat", json=body)
                elapsed = time.perf_counter() - start
                r.raise_for_status()
                return r.json().get("message", {}).get("content", "")
        except Exception as e:
            elapsed = time.perf_counter() - start
            _log.exception("chat error after %.3fs: %s", elapsed, e)
            return None

    async def list_models(self) -> list[str]:
        """Return names of all locally installed Ollama models."""
        try:
            async with httpx.AsyncClient(timeout=5.0) as client:
                r = await client.get(f"{self._api_base}/api/tags")
                r.raise_for_status()
                return [m["name"] for m in r.json().get("models", [])]
        except Exception:
            return []

    async def select_best_model(self) -> str:
        """Pick the highest-scoring available model, falling back to 'mistral'.

        Result is cached for 60 seconds so rapid chat requests don't each
        hit the Ollama API to build the model list.
        """
        now = time.monotonic()
        if self._model_cache and (now - self._model_cache_ts) < 60:
            return self._model_cache

        models = await self.list_models()
        best = max(models, key=_model_score) if models else "mistral"

        self._model_cache = best
        self._model_cache_ts = now
        return best

    def invalidate_model_cache(self):
        """Force next select_best_model() to re-query (e.g. after pull/delete)."""
        self._model_cache = None
        self._model_cache_ts = 0.0

    async def _chat_stream(self, messages: list, model: str, start: float, temperature: float | None = None):
        """Async generator streaming tokens, then a final __meta__ stats sentinel."""
        try:
            body: dict = {"model": model, "messages": messages, "stream": True}
            if temperature is not None:
                body["options"] = {"temperature": temperature}
            async with httpx.AsyncClient(timeout=300.0) as client:
                async with client.stream(
                    "POST",
                    f"{self._api_base}/api/chat",
                    json=body,
                ) as response:
                    response.raise_for_status()
                    async for line in response.aiter_lines():
                        if not line.strip():
                            continue
                        try:
                            data = json.loads(line)
                            token = data.get("message", {}).get("content", "")
                            if token:
                                yield token
                            if data.get("done", False):
                                elapsed = time.perf_counter() - start
                                _log.info("chat stream completed model=%s elapsed=%.3fs", model, elapsed)
                                eval_count = data.get("eval_count", 0)
                                eval_ns = data.get("eval_duration", 0)
                                tokens_per_s = round(eval_count / (eval_ns / 1e9), 1) if eval_ns else 0
                                stats = json.dumps({
                                    "model": model,
                                    "tokens": eval_count,
                                    "elapsed_s": round(elapsed, 2),
                                    "tokens_per_s": tokens_per_s,
                                })
                                yield f"__meta__{stats}"
                                break
                        except Exception:
                            continue
        except Exception as e:
            elapsed = time.perf_counter() - start
            _log.exception("chat stream error after %.3fs: %s", elapsed, e)
            return


def initialize_ollama() -> OllamaManager:
    global _ollama_manager

    if _ollama_manager is None:
        manager = OllamaManager()

        if not manager.is_running():
            manager.start()

        if not manager.is_running():
            raise RuntimeError("Ollama API is not reachable after start().")

        _ollama_manager = manager

    return _ollama_manager


async def initialize_ollama_async() -> OllamaManager:
    """Async-safe initializer — uses asyncio.sleep so the event loop stays live."""
    global _ollama_manager

    if _ollama_manager is None:
        manager = OllamaManager()

        if not manager.is_running():
            await manager.start_async()

        if not manager.is_running():
            raise RuntimeError("Ollama API is not reachable after start().")

        _ollama_manager = manager

    return _ollama_manager


def get_ollama_manager() -> OllamaManager:
    global _ollama_manager
    if _ollama_manager is None:
        _ollama_manager = OllamaManager()
    return _ollama_manager


def shutdown_ollama() -> None:
    global _ollama_manager
    if _ollama_manager is not None:
        _ollama_manager.stop()
        _ollama_manager = None