engine.py

"""
engine.py — llama.cpp 推理引擎 + 模型管理
=========================================
LLamaEngine: 模型加载/卸载/流式推理 (本地 Vulkan GPU + HTTP 双模式)。
业务函数: select_model / load_model_action / get_model_info_display / get_memory_status。

硬编码: 意图匹配的模型上下文值 (来自 HuggingFace 官方数据)
        GPU 检测通过 llama_cpp.llama_max_devices()
"""

import gc
import os
from pathlib import Path

import psutil

from config import load_config, save_config
from discovery import get_all_models


# ---- 模型最大上下文检测 (HuggingFace 已知值) -------------------------------------------------
# 硬编码: 以下值为各模型 HuggingFace 官方标注的最大上下文
def get_model_max_ctx(filepath: str, model_name: str = "") -> int:
    """返回模型真实最大上下文。已知模型用 HF 官方值，未知模型读 GGUF 元数据兜底。"""
    name = (model_name + filepath).lower()
    if "qwen3.6" in name:                     return 262144   # Qwen3.6: native 262K
    if "qwen3.5" in name:                     return 262144   # Qwen3.5: native 262K
    if "qwen3" in name:                       return 32768    # Qwen3: native 32K
    if "gemma4-12b" in name or "gemma4-31b" in name: return 262144  # Gemma4 Dense: 256K
    if "gemma4" in name:                      return 131072   # Gemma4 E2B/E4B: 128K
    if "minicpm" in name:                     return 32768
    if any(x in name for x in ["all-minilm", "nomic", "nanonets"]): return 2048
    # 兜底: 读 GGUF 元数据
    try:
        import llama_cpp
        with open(filepath, "rb") as f:
            meta = llama_cpp.llama_cpp.llama_model_loader(f).get_metadata()
            ctx = meta.get("llama.context_length", 0) or 0
            return max(int(ctx), 2048) if ctx else 4096
    except Exception:
        return 4096


# ---- 推理引擎 --------------------------------------------------------------------------------
class LLamaEngine:
    """llama.cpp 推理引擎。自动检测 Vulkan GPU，支持本地推理和 HTTP 服务器双模式。"""

    def __init__(self) -> None:
        self.model = None
        self.config = load_config()
        self._local_available = False
        try:
            import llama_cpp  # noqa: F401
            self._local_available = True
        except ImportError:
            pass

    # --- 属性 ---
    @property
    def is_local_available(self) -> bool:          return self._local_available
    @property
    def is_model_loaded(self) -> bool:             return self.model is not None
    @property
    def loaded_model_path(self) -> str:
        return getattr(self.model, "model_path", "") if self.model else ""

    # --- 内部 ---
    def _create_llama(self, model_path: str, n_ctx: int):
        """创建 llama-cpp-python 实例，自动检测 Vulkan GPU。"""
        import llama_cpp
        n_gpu = -1  # all layers to GPU
        try:
            llama_cpp.llama_cpp.llama_max_devices()
        except Exception:
            n_gpu = 0  # CPU-only fallback
        return llama_cpp.Llama(
            model_path=model_path, n_ctx=n_ctx,
            n_threads=os.cpu_count() or 4,
            n_gpu_layers=n_gpu, verbose=False,
        )

    # --- 公开 API ---
    def load_model(self, model_path: str, n_ctx: int = 2048) -> str:
        """加载 GGUF 模型到显存。返回状态消息。"""
        if not self._local_available:
            return "⚠️ 请先安装 llama-cpp-python"
        try:
            if self.model:
                del self.model; self.model = None; gc.collect()
            self.model = self._create_llama(model_path, n_ctx)
            return f"✅ 已加载 (ctx={n_ctx})"
        except Exception as e:
            return f"❌ 加载失败: {e}"

    def unload_model(self) -> str:
        """卸载模型释放显存。返回状态消息。"""
        if self.model:
            del self.model; self.model = None; gc.collect()
            return "✅ 模型已卸载"
        return "ℹ️ 无加载模型"

    def chat_server(self, messages: list[dict], params: dict, server_url: str):
        """HTTP SSE 流式调用 llama.cpp server (OpenAI 兼容端点)。"""
        import json as _json
        import httpx
        url = f"{server_url.rstrip('/')}/v1/chat/completions"
        payload = {
            "messages": messages, "stream": True,
            "temperature": params.get("temperature", 0.7),
            "top_p": params.get("top_p", 0.9),
            "top_k": params.get("top_k", 40),
            "repeat_penalty": params.get("repeat_penalty", 1.1),
            "max_tokens": params.get("max_tokens", 512),
        }
        try:
            with httpx.Client(timeout=120) as c:
                with c.stream("POST", url, json=payload) as r:
                    if r.status_code != 200:
                        yield f"⚠️ 服务器错误 ({r.status_code})"
                        return
                    for line in r.iter_lines():
                        if not line or not line.startswith("data: "):
                            continue
                        d = line[6:].strip()
                        if d == "[DONE]":
                            break
                        try:
                            token = _json.loads(d)["choices"][0]["delta"].get("content", "")
                            if token:
                                yield token
                        except _json.JSONDecodeError:
                            continue
        except Exception as e:
            yield f"❌ 连接失败: {e}"

    def chat_local(self, messages: list[dict], params: dict):
        """本地推理: create_chat_completion 流式（自动使用模型内置 chat 模板）。"""
        import llama_cpp
        mp = self.config.get("model_path", "")
        if not mp or not Path(mp).exists():
            yield "⚠️ 请先选择并加载模型"
            return
        # 模型未加载或路径已变 → 自动重载
        if self.model is None or getattr(self.model, "model_path", "") != mp:
            if self.model:
                del self.model; self.model = None; gc.collect()
            try:
                self.model = self._create_llama(mp, params.get("n_ctx", 2048))
            except Exception as e:
                yield f"❌ 模型加载失败: {e}"
                return
        try:
            stream = self.model.create_chat_completion(
                messages=messages, stream=True,
                max_tokens=params.get("max_tokens", 512),
                temperature=params.get("temperature", 0.7),
                top_p=params.get("top_p", 0.9),
                top_k=params.get("top_k", 40),
                repeat_penalty=params.get("repeat_penalty", 1.1),
            )
            for chunk in stream:
                token = chunk["choices"][0]["delta"].get("content", "")
                if token:
                    yield token
        except Exception as e:
            yield f"❌ 推理出错: {e}"

    def generate(self, messages: list[dict], params: dict, server_url: str):
        """统一推理入口: 本地模式 → chat_local, 服务器模式 → chat_server。"""
        if self.config["mode"] == "local" and self._local_available:
            yield from self.chat_local(messages, params)
        else:
            yield from self.chat_server(messages, params, server_url)


# ---- 全局引擎实例 ----------------------------------------------------------------------------
engine = LLamaEngine()


# ---- 业务逻辑 (UI 层直接调用) -----------------------------------------------------------------
def select_model(series_name: str, model_path: str) -> dict:
    """选择模型(不加载)，更新 config.json 和引擎配置。"""
    cfg = load_config()
    cfg["model_path"] = model_path
    cfg["active_series"] = series_name
    for m in get_all_models():
        if m["path"] == model_path:
            cfg["model_name"] = m["name"]
            break
    else:
        cfg["model_name"] = Path(model_path).stem
    cfg.setdefault("series_favorites", {})[series_name] = model_path
    save_config(cfg)
    engine.config = cfg
    return cfg


def load_model_action(series_name: str, model_path: str, n_ctx_val: int | None = None) -> tuple[str, str]:
    """加载模型。n_ctx_val=None 时自动取模型最大上下文的一半。返回 (状态消息, 模型信息)。"""
    if not model_path:
        return "⚠️ 请先选择 variant", get_model_info_display()
    cfg = select_model(series_name, model_path)
    if cfg.get("mode") == "server":
        msg = f"✅ 已选 {cfg['model_name']}（服务器模式）"
    else:
        if n_ctx_val is None:
            max_ctx = get_model_max_ctx(model_path, cfg.get("model_name", ""))
            n_ctx_val = max_ctx // 2
        msg = engine.load_model(model_path, int(n_ctx_val))
        if msg.startswith("✅"):
            msg = f"{cfg['model_name']} {msg}"
    return msg, get_model_info_display()


def get_model_info_display() -> str:
    """返回当前模型加载状态文本。"""
    cfg = load_config()
    mn = cfg.get("model_name", "未选")
    if engine.is_model_loaded:
        return f"{mn} ✅ 已加载"
    return f"{mn} ❌ 未加载" if mn != "未选" else "未加载任何模型"


def get_memory_status() -> str:
    """返回内存状态 HTML。>28GB 红色警告。"""
    mem = psutil.virtual_memory()
    used_gb = mem.used / (1024 ** 3)
    total_gb = mem.total / (1024 ** 3)
    color, icon = ("#dc3545", "⚠️") if used_gb > 28 else ("#28a745", "")
    return (f'<span style="color:{color};font-weight:600;">'
            f'{icon} 内存 {used_gb:.1f}/{total_gb:.0f}GB ({mem.percent:.0f}%)'
            f'</span>')