app/backend/minicpm_llama_cpp.py

from __future__ import annotations

import ctypes
from dataclasses import dataclass
from functools import lru_cache
import importlib.util
from importlib import metadata
import logging
import os
import site
import threading
import time
from pathlib import Path
from typing import Any

from huggingface_hub import hf_hub_download

try:
    import spaces
except ImportError:
    class _SpacesFallback:
        @staticmethod
        def GPU(*args: Any, **kwargs: Any):
            def decorator(fn):
                return fn

            return decorator

    spaces = _SpacesFallback()


logger = logging.getLogger(__name__)

DEFAULT_MAX_NEW_TOKENS = int(os.getenv("SMOLNALYSIS_MINICPM_MAX_NEW_TOKENS", os.getenv("MAX_TOKENS", "850")))
DEFAULT_TEMPERATURE = float(os.getenv("SMOLNALYSIS_MINICPM_TEMPERATURE", os.getenv("TEMPERATURE", "0.7")))
DEFAULT_TOP_P = float(os.getenv("SMOLNALYSIS_MINICPM_TOP_P", os.getenv("TOP_P", "0.9")))
DEFAULT_N_CTX = int(os.getenv("SMOLNALYSIS_MINICPM_N_CTX", os.getenv("N_CTX", "4096")))
DEFAULT_N_BATCH = int(os.getenv("SMOLNALYSIS_MINICPM_N_BATCH", os.getenv("N_BATCH", "512")))
DEFAULT_N_GPU_LAYERS = int(os.getenv("SMOLNALYSIS_MINICPM_N_GPU_LAYERS", os.getenv("N_GPU_LAYERS", "0")))
ZERO_GPU_DURATION_SECONDS = int(os.getenv("SMOLNALYSIS_MINICPM_ZEROGPU_DURATION_SECONDS", "120"))
EAGER_LOAD_ROLES_ENV = "SMOLNALYSIS_MINICPM_EAGER_LOAD_ROLES"

ROLE_ALIASES = {
    "auto": "auto",
    "router": "auto",
    "base": "general_agent",
    "none": "general_agent",
    "general": "general_agent",
    "general_agent": "general_agent",
    "ckan": "ckan_retrieval",
    "ckan_tool": "ckan_retrieval",
    "retrieval": "ckan_retrieval",
    "ckan_retrieval": "ckan_retrieval",
    "data": "data_analysis",
    "analysis": "data_analysis",
    "data_analysis": "data_analysis",
    "openui": "openui_translator",
    "openui_translator": "openui_translator",
}

ROLE_ENV_KEYS = {
    "general_agent": "GENERAL_AGENT",
    "ckan_retrieval": "CKAN_RETRIEVAL",
    "data_analysis": "DATA_ANALYSIS",
    "openui_translator": "OPENUI_TRANSLATOR",
}

EAGER_LOAD_STATUS: dict[str, Any] = {"enabled": False, "roles": {}, "duration_ms": 0}


@dataclass(frozen=True)
class LlamaCppRoleConfig:
    role: str
    model_path: str
    model_repo_id: str
    model_filename: str
    lora_path: str
    lora_repo_id: str
    lora_filename: str


def _clean_env_value(name: str, default: str = "") -> str:
    raw = os.getenv(name, default)
    lines = []
    for line in str(raw).splitlines():
        value = line.strip().strip('"').strip("'")
        if value and not value.startswith("#"):
            lines.append(value)
    return lines[-1] if lines else default


def _role_env(role: str, suffix: str) -> str:
    return f"SMOLNALYSIS_MINICPM_{ROLE_ENV_KEYS[role]}_{suffix}"


def normalize_role(adapter: str | None) -> str:
    value = (adapter or "auto").strip().casefold()
    return ROLE_ALIASES.get(value, value)


def route_role(messages: list[dict[str, str]], adapter: str | None = "auto") -> str:
    requested = normalize_role(adapter)
    if requested != "auto":
        return requested

    try:
        from . import router_runtime
    except ImportError:
        import router_runtime  # type: ignore

    prediction = router_runtime.predict_role(messages, model_id=router_runtime.router_tokenizer_model_id())
    if prediction and prediction.role in ROLE_ENV_KEYS:
        return prediction.role

    last_user_text = next(
        (message["content"] for message in reversed(messages) if message.get("role") == "user"),
        "",
    ).casefold()
    if any(term in last_user_text for term in ("openui", "component", "render", "ui", "card", "chart")):
        return "openui_translator"
    if any(term in last_user_text for term in ("analy", "quality", "distribution", "trend", "statistics", "missing")):
        return "data_analysis"
    if any(term in last_user_text for term in ("ckan", "dataset", "resource", "search", "retrieve", "catalog")):
        return "ckan_retrieval"
    return "general_agent"


def role_config(role: str) -> LlamaCppRoleConfig:
    if role not in ROLE_ENV_KEYS:
        available = ", ".join(ROLE_ENV_KEYS)
        raise KeyError(f"Unknown MiniCPM llama.cpp role '{role}'. Available roles: {available}")

    model_path = _clean_env_value(_role_env(role, "MODEL_PATH"), _clean_env_value("SMOLNALYSIS_MINICPM_MODEL_PATH", _clean_env_value("MODEL_PATH")))
    model_repo_id = _clean_env_value(
        _role_env(role, "MODEL_REPO_ID"),
        _clean_env_value("SMOLNALYSIS_MINICPM_MODEL_REPO_ID", _clean_env_value("MODEL_REPO_ID")),
    )
    model_filename = _clean_env_value(
        _role_env(role, "MODEL_FILENAME"),
        _clean_env_value("SMOLNALYSIS_MINICPM_MODEL_FILENAME", _clean_env_value("MODEL_FILENAME")),
    )
    lora_path = _clean_env_value(_role_env(role, "LORA_PATH"), "")
    lora_repo_id = _clean_env_value(_role_env(role, "LORA_REPO_ID"), "")
    lora_filename = _clean_env_value(_role_env(role, "LORA_FILENAME"), "")
    return LlamaCppRoleConfig(role, model_path, model_repo_id, model_filename, lora_path, lora_repo_id, lora_filename)


def _resolve_model_path(config: LlamaCppRoleConfig) -> str:
    if config.model_path:
        path = Path(config.model_path).expanduser()
        if not path.exists():
            raise FileNotFoundError(f"MiniCPM GGUF model path does not exist: {path}")
        return str(path)
    if config.model_repo_id and config.model_filename:
        return hf_hub_download(repo_id=config.model_repo_id, filename=config.model_filename)
    raise RuntimeError(
        "MiniCPM llama.cpp model is not configured. Set MODEL_PATH or "
        "MODEL_REPO_ID and MODEL_FILENAME, or use the SMOLNALYSIS_MINICPM_* equivalents."
    )


def _resolve_lora_path(config: LlamaCppRoleConfig) -> str:
    if config.lora_path:
        path = Path(config.lora_path).expanduser()
        if not path.exists():
            raise FileNotFoundError(f"MiniCPM LoRA path does not exist for role {config.role}: {path}")
        return str(path)
    if config.lora_repo_id and config.lora_filename:
        return hf_hub_download(repo_id=config.lora_repo_id, filename=config.lora_filename)
    return ""


def _role_runtime_options(role: str) -> dict[str, Any]:
    options: dict[str, Any] = {
        "n_ctx": int(_clean_env_value(_role_env(role, "N_CTX"), str(DEFAULT_N_CTX))),
        "n_batch": int(_clean_env_value(_role_env(role, "N_BATCH"), str(DEFAULT_N_BATCH))),
        "n_gpu_layers": int(_clean_env_value(_role_env(role, "N_GPU_LAYERS"), str(DEFAULT_N_GPU_LAYERS))),
        "verbose": _clean_env_value("SMOLNALYSIS_MINICPM_VERBOSE", "false").casefold() in {"1", "true", "yes", "on"},
    }
    n_threads = _clean_env_value(_role_env(role, "N_THREADS"), _clean_env_value("SMOLNALYSIS_MINICPM_N_THREADS", _clean_env_value("N_THREADS")))
    if n_threads:
        options["n_threads"] = int(n_threads)
    return options


@lru_cache(maxsize=4)
def _load_llama_cached(
    model_path: str,
    lora_path: str,
    n_ctx: int,
    n_batch: int,
    n_gpu_layers: int,
    n_threads: int | None,
    verbose: bool,
):
    try:
        _preload_cuda_runtime()
        from llama_cpp import Llama
    except ImportError as exc:
        raise RuntimeError("llama-cpp-python is not installed in this runtime.") from exc

    kwargs: dict[str, Any] = {
        "model_path": model_path,
        "n_ctx": n_ctx,
        "n_batch": n_batch,
        "n_gpu_layers": n_gpu_layers,
        "verbose": verbose,
    }
    if n_threads is not None:
        kwargs["n_threads"] = n_threads
    if lora_path:
        kwargs["lora_path"] = lora_path

    logger.info("loading MiniCPM llama.cpp model=%s lora=%s", model_path, lora_path or "none")
    return Llama(**kwargs)


def _preload_cuda_runtime() -> str:
    candidates = _cuda_library_candidates()
    loaded = []
    errors = []
    for candidate in candidates:
        try:
            ctypes.CDLL(str(candidate), mode=ctypes.RTLD_GLOBAL)
            loaded.append(str(candidate))
        except OSError as exc:
            errors.append(f"{candidate}: {exc}")
    if errors:
        logger.debug("CUDA runtime preload attempts failed: %s", " | ".join(errors))
    return os.pathsep.join(loaded)


def _cuda_library_candidates() -> list[Path]:
    candidates: list[Path] = []
    package_libraries = [
        ("nvidia.nvjitlink", "libnvJitLink.so.12"),
        ("nvidia.cuda_runtime", "libcudart.so.12"),
        ("nvidia.cublas", "libcublasLt.so.12"),
        ("nvidia.cublas", "libcublas.so.12"),
    ]
    for package, library in package_libraries:
        spec = importlib.util.find_spec(package)
        if spec and spec.submodule_search_locations:
            for location in spec.submodule_search_locations:
                candidates.append(Path(location) / "lib" / library)
    for root in [*site.getsitepackages(), site.getusersitepackages()]:
        base = Path(root) / "nvidia"
        candidates.extend(
            [
                base / "nvjitlink" / "lib" / "libnvJitLink.so.12",
                base / "cuda_runtime" / "lib" / "libcudart.so.12",
                base / "cublas" / "lib" / "libcublasLt.so.12",
                base / "cublas" / "lib" / "libcublas.so.12",
            ]
        )
    return [candidate for candidate in candidates if candidate.exists()]


def _load_llama(role: str):
    config = role_config(role)
    model_path = _resolve_model_path(config)
    lora_path = _resolve_lora_path(config)
    options = _role_runtime_options(role)
    return _load_llama_cached(
        model_path,
        lora_path,
        options["n_ctx"],
        options["n_batch"],
        options["n_gpu_layers"],
        options.get("n_threads"),
        options["verbose"],
    )


def _load_llama_with_gpu_fallback(role: str):
    config = role_config(role)
    model_path = _resolve_model_path(config)
    lora_path = _resolve_lora_path(config)
    options = _role_runtime_options(role)
    try:
        llm = _load_llama_cached(
            model_path,
            lora_path,
            options["n_ctx"],
            options["n_batch"],
            options["n_gpu_layers"],
            options.get("n_threads"),
            options["verbose"],
        )
        return llm, options, ""
    except Exception as exc:
        if options["n_gpu_layers"] == 0:
            raise
        fallback_options = {**options, "n_gpu_layers": 0}
        logger.exception("MiniCPM llama.cpp GPU load failed for role=%s; retrying with CPU.", role)
        llm = _load_llama_cached(
            model_path,
            lora_path,
            fallback_options["n_ctx"],
            fallback_options["n_batch"],
            fallback_options["n_gpu_layers"],
            fallback_options.get("n_threads"),
            fallback_options["verbose"],
        )
        return llm, fallback_options, f"{type(exc).__name__}: {str(exc).strip() or type(exc).__name__}"


def role_runtime_status(role: str) -> dict[str, Any]:
    config = role_config(role)
    options = _role_runtime_options(role)
    model_path = ""
    lora_path = ""
    model_error = ""
    lora_error = ""
    try:
        model_path = _resolve_model_path(config)
    except Exception as exc:
        model_error = str(exc)
    try:
        lora_path = _resolve_lora_path(config)
    except Exception as exc:
        lora_error = str(exc)
    return {
        "role": role,
        "llama_cpp": llama_cpp_runtime_info(),
        "model_path": model_path or config.model_path,
        "model_repo_id": config.model_repo_id,
        "model_filename": config.model_filename,
        "model_hub_url": _hub_url(config.model_repo_id, config.model_filename),
        "model_error": model_error,
        "lora_path": lora_path or config.lora_path,
        "lora_repo_id": config.lora_repo_id,
        "lora_filename": config.lora_filename,
        "lora_hub_url": _hub_url(config.lora_repo_id, config.lora_filename),
        "lora_error": lora_error,
        "options": options,
        "configured": bool(config.model_path or (config.model_repo_id and config.model_filename)),
        "loaded_models": _load_llama_cached.cache_info().currsize,
    }


def _hub_url(repo_id: str, filename: str = "") -> str:
    if not repo_id or "/" not in repo_id:
        return ""
    clean_repo_id = "/".join(part.strip("/") for part in repo_id.split("/") if part.strip("/"))
    if not clean_repo_id:
        return ""
    clean_filename = filename.strip().lstrip("/")
    if clean_filename:
        return f"https://huggingface.co/{clean_repo_id}/blob/main/{clean_filename}"
    return f"https://huggingface.co/{clean_repo_id}"


def llama_cpp_runtime_info() -> dict[str, Any]:
    info: dict[str, Any] = {
        "installed": False,
        "version": "",
        "supports_gpu_offload": None,
        "backend": "",
        "cuda_runtime_preload": "",
        "error": "",
    }
    try:
        info["version"] = metadata.version("llama-cpp-python")
    except metadata.PackageNotFoundError:
        info["error"] = "llama-cpp-python is not installed."
        return info
    except Exception as exc:
        info["error"] = str(exc)

    try:
        info["cuda_runtime_preload"] = _preload_cuda_runtime()
        import llama_cpp

        info["installed"] = True
        info["version"] = str(getattr(llama_cpp, "__version__", info["version"]))
        low_level = getattr(llama_cpp, "llama_cpp", None)
        supports_gpu_offload = getattr(low_level, "llama_supports_gpu_offload", None)
        if callable(supports_gpu_offload):
            info["supports_gpu_offload"] = bool(supports_gpu_offload())
        supports_mmap = getattr(low_level, "llama_supports_mmap", None)
        if callable(supports_mmap):
            info["supports_mmap"] = bool(supports_mmap())
        supports_mlock = getattr(low_level, "llama_supports_mlock", None)
        if callable(supports_mlock):
            info["supports_mlock"] = bool(supports_mlock())
    except Exception as exc:
        info["error"] = str(exc)
    return info

ROLE_SYSTEM_PROMPTS = {
    "general_agent": "You are smolnalysis, a concise assistant for exploring open data and planning analysis steps.",
    "ckan_retrieval": "You are the smolnalysis CKAN retrieval specialist. Help identify datasets, resources, filters, and catalog search steps.",
    "data_analysis": "You are the smolnalysis data analyst. Focus on columns, quality checks, aggregations, distributions, trends, and clear next analyses.",
    "openui_translator": "You are the smolnalysis OpenUI translator. When asked for UI, return valid OpenUI-Lang only.",
}


def _with_role_system_prompt(messages: list[dict[str, str]], role: str) -> list[dict[str, str]]:
    if any(message.get("role") == "system" for message in messages):
        return messages
    prompt = ROLE_SYSTEM_PROMPTS.get(role)
    if not prompt:
        return messages
    return [{"role": "system", "content": prompt}, *messages]


MODEL_LOCK = threading.Lock()


@spaces.GPU(duration=ZERO_GPU_DURATION_SECONDS)
def generate_chat_response(
    messages: list[dict[str, str]],
    *,
    adapter: str | None = "auto",
    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
    temperature: float = DEFAULT_TEMPERATURE,
    top_p: float = DEFAULT_TOP_P,
    top_k: int | None = None,
) -> str:
    response, _trace = generate_chat_response_with_trace(
        messages,
        adapter=adapter,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_p=top_p,
        top_k=top_k,
    )
    return response


@spaces.GPU(duration=ZERO_GPU_DURATION_SECONDS)
def generate_chat_response_with_trace(
    messages: list[dict[str, str]],
    *,
    adapter: str | None = "auto",
    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
    temperature: float = DEFAULT_TEMPERATURE,
    top_p: float = DEFAULT_TOP_P,
    top_k: int | None = None,
) -> tuple[str, dict[str, Any]]:
    started = time.perf_counter()
    role = route_role(messages, adapter)
    runtime = role_runtime_status(role)
    routed_messages = _with_role_system_prompt(messages, role)
    cache_before = _load_llama_cached.cache_info()
    effective_options: dict[str, Any] = {}
    gpu_fallback_error = ""
    with MODEL_LOCK:
        try:
            llm, effective_options, gpu_fallback_error = _load_llama_with_gpu_fallback(role)
        except Exception as exc:
            raise RuntimeError(f"MiniCPM llama.cpp load failed for role '{role}'.") from exc
        cache_after_load = _load_llama_cached.cache_info()
        payload: dict[str, Any] = {
            "messages": routed_messages,
            "temperature": temperature,
            "top_p": top_p,
            "max_tokens": max_new_tokens,
            "stream": False,
        }
        if top_k is not None:
            payload["top_k"] = top_k
        try:
            response = llm.create_chat_completion(**payload)
        except Exception as exc:
            raise RuntimeError(f"MiniCPM llama.cpp generation failed for role '{role}'.") from exc

    content = response["choices"][0]["message"]["content"]
    elapsed_ms = round((time.perf_counter() - started) * 1000, 1)
    cache_hit = cache_after_load.hits > cache_before.hits
    trace = {
        "backend": "llama.cpp",
        "model_family": "MiniCPM",
        "requested_adapter": adapter or "auto",
        "role": role,
        "message_count": len(messages),
        "routed_message_count": len(routed_messages),
        "sampling": {
            "max_new_tokens": max_new_tokens,
            "temperature": temperature,
            "top_p": top_p,
            "top_k": top_k,
        },
        "runtime": runtime,
        "effective_options": effective_options,
        "gpu_fallback_error": gpu_fallback_error,
        "cache": {
            "hit": cache_hit,
            "loaded_models": cache_after_load.currsize,
            "hits": cache_after_load.hits,
            "misses": cache_after_load.misses,
        },
        "events": [
            {"name": "route_role", "detail": f"{adapter or 'auto'} -> {role}"},
            {"name": "resolve_runtime", "detail": runtime.get("model_path") or runtime.get("model_repo_id") or "unconfigured"},
            {"name": "load_model", "detail": "cache hit" if cache_hit else "cache miss"},
            {"name": "generate", "detail": f"{len(str(content).strip())} chars in {elapsed_ms} ms"},
        ],
        "duration_ms": elapsed_ms,
        "output_chars": len(str(content).strip()),
    }
    logger.info("MiniCPM llama.cpp response generated: role=%s chars=%d", role, len(content))
    return str(content).strip(), trace


def runtime_status() -> dict[str, Any]:
    try:
        from . import router_runtime
    except ImportError:
        import router_runtime  # type: ignore

    roles = {}
    for role in ROLE_ENV_KEYS:
        config = role_config(role)
        status = role_runtime_status(role)
        roles[role] = {
            "model_path": config.model_path,
            "model_repo_id": config.model_repo_id,
            "model_filename": config.model_filename,
            "model_hub_url": status.get("model_hub_url", ""),
            "lora_path": config.lora_path,
            "lora_repo_id": config.lora_repo_id,
            "lora_filename": config.lora_filename,
            "lora_hub_url": status.get("lora_hub_url", ""),
            "configured": bool(config.model_path or (config.model_repo_id and config.model_filename)),
            "loaded": _load_llama_cached.cache_info().currsize > 0,
            "resolved_model_path": status.get("model_path", ""),
            "resolved_lora_path": status.get("lora_path", ""),
            "model_error": status.get("model_error", ""),
            "lora_error": status.get("lora_error", ""),
            "options": status.get("options", {}),
            "llama_cpp": status.get("llama_cpp", {}),
        }
    return {
        "backend": "llama.cpp",
        "model_family": "MiniCPM",
        "llama_cpp": llama_cpp_runtime_info(),
        "eager_load": EAGER_LOAD_STATUS,
        "router": router_runtime.runtime_status(),
        "roles": roles,
        "n_ctx": DEFAULT_N_CTX,
        "n_gpu_layers": DEFAULT_N_GPU_LAYERS,
        "max_new_tokens": DEFAULT_MAX_NEW_TOKENS,
    }


def probe_runtime(role: str = "general_agent") -> dict[str, Any]:
    started = time.perf_counter()
    normalized_role = normalize_role(role)
    if normalized_role == "auto":
        normalized_role = "general_agent"
    result: dict[str, Any] = {
        "role": normalized_role,
        "llama_cpp": llama_cpp_runtime_info(),
        "status": role_runtime_status(normalized_role),
        "load": {},
        "duration_ms": 0,
    }
    config = role_config(normalized_role)
    try:
        model_path = _resolve_model_path(config)
        lora_path = _resolve_lora_path(config)
        options = _role_runtime_options(normalized_role)
        result["load"] = {
            "ok": False,
            "model_path": model_path,
            "lora_path": lora_path,
            "options": options,
        }
        _preload_cuda_runtime()
        from llama_cpp import Llama

        kwargs: dict[str, Any] = {
            "model_path": model_path,
            "n_ctx": min(options["n_ctx"], 512),
            "n_batch": min(options["n_batch"], 128),
            "n_gpu_layers": options["n_gpu_layers"],
            "verbose": True,
        }
        if options.get("n_threads") is not None:
            kwargs["n_threads"] = options["n_threads"]
        if lora_path:
            kwargs["lora_path"] = lora_path
        Llama(**kwargs)
        result["load"]["ok"] = True
    except Exception as exc:
        result["load"]["error_type"] = type(exc).__name__
        result["load"]["error"] = str(exc).strip() or type(exc).__name__
    result["duration_ms"] = round((time.perf_counter() - started) * 1000, 1)
    return result


def _configured_for_role(role: str) -> bool:
    config = role_config(role)
    return bool(config.model_path or (config.model_repo_id and config.model_filename))


def _eager_load_roles() -> list[str]:
    raw = _clean_env_value(EAGER_LOAD_ROLES_ENV, "general_agent").casefold()
    if raw in {"0", "false", "no", "off", "none", "disabled"}:
        return []
    if raw in {"1", "true", "yes", "on", "default"}:
        return ["general_agent"]
    if raw == "all":
        return list(ROLE_ENV_KEYS)
    roles = []
    for item in raw.replace(";", ",").split(","):
        role = normalize_role(item.strip())
        if role == "auto":
            role = "general_agent"
        if role in ROLE_ENV_KEYS and role not in roles:
            roles.append(role)
    return roles


def _eager_load_configured_roles() -> None:
    roles = _eager_load_roles()
    EAGER_LOAD_STATUS["enabled"] = bool(roles)
    if not roles:
        return

    started = time.perf_counter()
    for role in roles:
        role_status: dict[str, Any] = {"configured": False, "loaded": False, "error": "", "effective_options": {}}
        EAGER_LOAD_STATUS["roles"][role] = role_status
        try:
            role_status["configured"] = _configured_for_role(role)
            if not role_status["configured"]:
                role_status["error"] = "Model is not configured for this role."
                continue
            _llm, effective_options, gpu_fallback_error = _load_llama_with_gpu_fallback(role)
            role_status["loaded"] = True
            role_status["effective_options"] = effective_options
            role_status["gpu_fallback_error"] = gpu_fallback_error
        except Exception as exc:
            logger.exception("MiniCPM eager load failed for role=%s.", role)
            role_status["error"] = f"{type(exc).__name__}: {str(exc).strip() or type(exc).__name__}"
    EAGER_LOAD_STATUS["duration_ms"] = round((time.perf_counter() - started) * 1000, 1)


_eager_load_configured_roles()