Initial commit: LocalPilot：本地模型运行时，复用 llama-server 并提供 Ollama 兼容 provider

2026-09-19 11:57:59 +08:00
commit d159212416
27 changed files with 3022 additions and 0 deletions
@@ -0,0 +1,241 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import socket
+import subprocess
+from collections.abc import AsyncIterator
+from pathlib import Path
+from typing import Any
+
+import httpx
+
+from ..config import ModelConfig, RuntimeConfig
+from ..types import BackendError, ChatChunk, ChatParams
+from .base import BaseBackend
+
+
+def _free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("127.0.0.1", 0))
+        return int(sock.getsockname()[1])
+
+
+class GGUFBackend(BaseBackend):
+    """Run a GGUF model through an external CUDA-enabled llama-server."""
+
+    def __init__(self, model: ModelConfig, runtime: RuntimeConfig) -> None:
+        super().__init__(model, runtime)
+        self.port: int | None = None
+        self.process: subprocess.Popen[str] | None = None
+        self.log_path: Path | None = None
+        self._log_handle: Any = None
+
+    @property
+    def endpoint(self) -> str:
+        if self.port is None:
+            raise BackendError("GGUF 后端尚未启动")
+        return f"http://127.0.0.1:{self.port}"
+
+    async def load(self) -> None:
+        if self.loaded:
+            return
+        if not self.model.path:
+            raise BackendError(f"GGUF 模型缺少 path: {self.model.id}")
+        model_path = Path(self.model.path)
+        if not model_path.exists():
+            raise BackendError(f"GGUF 文件不存在: {model_path}")
+        server = Path(self.runtime.llama_server)
+        if not server.exists():
+            raise BackendError(
+                f"找不到 llama-server: {server}。请在 config.json 的 runtime.llama_server 中配置。"
+            )
+
+        self.port = _free_port()
+        runtime_dir = Path(self.runtime.runtime_dir)
+        runtime_dir.mkdir(parents=True, exist_ok=True)
+        self.log_path = runtime_dir / f"{self.model.id}.llama-server.log"
+        self._log_handle = self.log_path.open("a", encoding="utf-8", errors="replace")
+
+        options = self.model.options
+        args = [
+            str(server),
+            "--model",
+            str(model_path),
+            "--host",
+            "127.0.0.1",
+            "--port",
+            str(self.port),
+            "--alias",
+            self.model.id,
+            "--no-webui",
+            "--jinja",
+            "--cont-batching",
+            "--cache-prompt",
+            "--flash-attn",
+            str(options.get("flash_attn", self.runtime.flash_attn)),
+            "--gpu-layers",
+            str(options.get("n_gpu_layers", self.runtime.n_gpu_layers)),
+            "--ctx-size",
+            str(options.get("ctx_size", self.runtime.ctx_size)),
+            "--batch-size",
+            str(options.get("batch_size", self.runtime.batch_size)),
+            "--ubatch-size",
+            str(options.get("ubatch_size", self.runtime.ubatch_size)),
+        ]
+        if self.runtime.fit_vram and options.get("fit_vram", True):
+            args.extend(["--fit", "on", "--fit-ctx", str(options.get("fit_ctx", 4096))])
+        if options.get("mlock", False):
+            args.append("--mlock")
+        if options.get("no_mmap", False):
+            args.append("--no-mmap")
+        if options.get("mmproj"):
+            args.extend(["--mmproj", str(options["mmproj"])])
+        if options.get("chat_template_file"):
+            args.extend(["--chat-template-file", str(options["chat_template_file"])])
+        elif self.model.template:
+            args.extend(["--chat-template", self.model.template])
+
+        creationflags = getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0)
+        self.process = subprocess.Popen(
+            args,
+            cwd=str(server.parent),
+            stdout=self._log_handle,
+            stderr=subprocess.STDOUT,
+            text=True,
+            creationflags=creationflags,
+        )
+        try:
+            await self._wait_ready()
+        except Exception:
+            await self.unload()
+            raise
+        self.loaded = True
+
+    async def _wait_ready(self) -> None:
+        deadline = asyncio.get_running_loop().time() + 90.0
+        last_error = ""
+        async with httpx.AsyncClient(timeout=2.0) as client:
+            while asyncio.get_running_loop().time() < deadline:
+                if self.process and self.process.poll() is not None:
+                    tail = ""
+                    if self.log_path and self.log_path.exists():
+                        tail = self.log_path.read_text(encoding="utf-8", errors="replace")[-2000:]
+                    if "check_tensor_dims" in tail or "wrong shape" in tail:
+                        tail = (
+                            "检测到模型张量与当前 llama.cpp runner 不兼容。若 path 指向 Ollama 的 "
+                            "content-addressed blob，请改用 kind=ollama；否则换成与 runner 匹配的标准 GGUF 文件。\n"
+                            + tail
+                        )
+                    raise BackendError(f"llama-server 启动失败。日志尾部:\n{tail}")
+                try:
+                    response = await client.get(f"{self.endpoint}/health")
+                    if response.status_code == 200:
+                        return
+                    last_error = f"HTTP {response.status_code}"
+                except httpx.HTTPError as exc:
+                    last_error = str(exc)
+                await asyncio.sleep(0.25)
+        raise BackendError(f"llama-server 在 90 秒内未就绪: {last_error}")
+
+    async def unload(self) -> None:
+        process, self.process = self.process, None
+        if process is not None and process.poll() is None:
+            process.terminate()
+            try:
+                await asyncio.to_thread(process.wait, 5)
+            except subprocess.TimeoutExpired:
+                process.kill()
+                await asyncio.to_thread(process.wait)
+        if self._log_handle is not None:
+            self._log_handle.close()
+            self._log_handle = None
+        self.loaded = False
+        self.port = None
+
+    async def chat(
+        self,
+        messages: list[dict[str, Any]],
+        params: ChatParams,
+        stream: bool = False,
+    ) -> AsyncIterator[ChatChunk]:
+        if not self.loaded:
+            await self.load()
+        payload: dict[str, Any] = {
+            "model": self.model.id,
+            "messages": messages,
+            "stream": stream,
+            "max_tokens": params.max_tokens,
+            "temperature": params.temperature,
+            "top_p": params.top_p,
+            "top_k": params.top_k,
+            "min_p": params.min_p,
+            "repeat_penalty": params.repeat_penalty,
+        }
+        if params.seed is not None:
+            payload["seed"] = params.seed
+        if params.stop:
+            payload["stop"] = params.stop
+        if params.enable_thinking is not None:
+            payload["chat_template_kwargs"] = {"enable_thinking": params.enable_thinking}
+        if params.response_format is not None:
+            payload["response_format"] = (
+                {"type": params.response_format} if isinstance(params.response_format, str) else params.response_format
+            )
+        if params.tools:
+            payload["tools"] = params.tools
+        if params.tool_choice is not None:
+            payload["tool_choice"] = params.tool_choice
+        payload.update(params.extra)
+
+        timeout = httpx.Timeout(600.0, connect=10.0)
+        async with httpx.AsyncClient(timeout=timeout) as client:
+            if stream:
+                async with client.stream("POST", f"{self.endpoint}/v1/chat/completions", json=payload) as response:
+                    if response.status_code >= 400:
+                        detail = (await response.aread()).decode("utf-8", errors="replace")
+                        raise BackendError(f"llama-server 请求失败 {response.status_code}: {detail}")
+                    async for line in response.aiter_lines():
+                        if not line.startswith("data:"):
+                            continue
+                        raw = line[5:].strip()
+                        if raw == "[DONE]":
+                            yield ChatChunk(done=True, finish_reason="stop")
+                            return
+                        try:
+                            data = json.loads(raw)
+                        except json.JSONDecodeError:
+                            continue
+                        choice = (data.get("choices") or [{}])[0]
+                        delta = choice.get("delta") or {}
+                        text = delta.get("content") or ""
+                        finish = choice.get("finish_reason")
+                        if text or finish or len(delta) > 1:
+                            yield ChatChunk(text=text, finish_reason=finish, metadata=data, delta=delta)
+                    yield ChatChunk(done=True, finish_reason="stop")
+                return
+
+            response = await client.post(f"{self.endpoint}/v1/chat/completions", json=payload)
+            if response.status_code >= 400:
+                raise BackendError(f"llama-server 请求失败 {response.status_code}: {response.text}")
+            data = response.json()
+            choice = (data.get("choices") or [{}])[0]
+            message = choice.get("message") or {}
+            yield ChatChunk(
+                text=message.get("content") or "",
+                done=True,
+                finish_reason=choice.get("finish_reason") or "stop",
+                metadata=data,
+                delta=message,
+            )
+
+    def info(self) -> dict[str, Any]:
+        return {
+            "id": self.model.id,
+            "kind": "gguf",
+            "path": self.model.path,
+            "loaded": self.loaded,
+            "endpoint": self.endpoint if self.port else None,
+            "pid": self.process.pid if self.process else None,
+            "log": str(self.log_path) if self.log_path else None,
+        }