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

localpilot/backends/__init__.py

@@ -0,0 +1,8 @@
+from .base import BaseBackend
+from .gguf import GGUFBackend
+from .onnx import ONNXBackend
+from .ollama import OllamaBackend
+from .remote import RemoteBackend
+from .transformers_backend import TransformersBackend
+
+__all__ = ["BaseBackend", "GGUFBackend", "ONNXBackend", "OllamaBackend", "RemoteBackend", "TransformersBackend"]

localpilot/backends/base.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator
+from typing import Any
+
+from ..config import ModelConfig, RuntimeConfig
+from ..types import BackendError, ChatChunk, ChatParams
+
+
+class BaseBackend(ABC):
+    def __init__(self, model: ModelConfig, runtime: RuntimeConfig) -> None:
+        self.model = model
+        self.runtime = runtime
+        self.loaded = False
+
+    @abstractmethod
+    async def load(self) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def unload(self) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def chat(
+        self,
+        messages: list[dict[str, Any]],
+        params: ChatParams,
+        stream: bool = False,
+    ) -> AsyncIterator[ChatChunk]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def info(self) -> dict[str, Any]:
+        raise NotImplementedError
+
+    async def embed(self, inputs: list[str]) -> list[list[float]]:
+        raise BackendError(f"后端 {self.model.id} 不支持 embeddings")

localpilot/backends/gguf.py

@@ -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,
+        }

localpilot/backends/ollama.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+import json
+from collections.abc import AsyncIterator
+from typing import Any
+
+import httpx
+
+from ..config import ModelConfig, RuntimeConfig
+from ..types import BackendError, ChatChunk, ChatParams
+from .base import BaseBackend
+
+
+class OllamaBackend(BaseBackend):
+    """Compatibility provider for models already managed by a local Ollama daemon."""
+
+    def __init__(self, model: ModelConfig, runtime: RuntimeConfig) -> None:
+        super().__init__(model, runtime)
+        self.base_url = runtime.ollama_url.rstrip("/")
+        self.remote_model = model.remote_model or model.id
+
+    async def load(self) -> None:
+        if self.loaded:
+            return
+        try:
+            async with httpx.AsyncClient(timeout=5.0) as client:
+                response = await client.get(f"{self.base_url}/api/version")
+                response.raise_for_status()
+        except httpx.HTTPError as exc:
+            raise BackendError(f"无法连接本地 Ollama: {self.base_url}: {exc}") from exc
+        self.loaded = True
+
+    async def unload(self) -> None:
+        self.loaded = False
+
+    async def chat(
+        self,
+        messages: list[dict[str, Any]],
+        params: ChatParams,
+        stream: bool = False,
+    ) -> AsyncIterator[ChatChunk]:
+        if not self.loaded:
+            await self.load()
+        options: dict[str, Any] = {
+            "num_predict": 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:
+            options["seed"] = params.seed
+        if params.stop:
+            options["stop"] = params.stop
+        options.update(params.extra)
+        payload: dict[str, Any] = {
+            "model": self.remote_model,
+            "messages": messages,
+            "stream": stream,
+            "options": options,
+            "keep_alive": params.keep_alive if params.keep_alive is not None else "5m",
+        }
+        if params.response_format is not None:
+            payload["format"] = params.response_format
+        if params.enable_thinking is not None:
+            payload["think"] = params.enable_thinking
+        if params.tools:
+            payload["tools"] = params.tools
+        async with httpx.AsyncClient(timeout=httpx.Timeout(600.0, connect=10.0)) as client:
+            if stream:
+                async with client.stream("POST", f"{self.base_url}/api/chat", json=payload) as response:
+                    if response.status_code >= 400:
+                        detail = (await response.aread()).decode("utf-8", errors="replace")
+                        raise BackendError(f"Ollama 请求失败 {response.status_code}: {detail}")
+                    async for line in response.aiter_lines():
+                        if not line.strip():
+                            continue
+                        try:
+                            data = json.loads(line)
+                        except json.JSONDecodeError:
+                            continue
+                        message = data.get("message") or {}
+                        text = message.get("content") or ""
+                        if data.get("done"):
+                            yield ChatChunk(done=True, finish_reason=data.get("done_reason") or "stop", metadata=data)
+                            return
+                        if text or len(message) > 1:
+                            yield ChatChunk(text=text, delta=message, metadata=data)
+                    yield ChatChunk(done=True, finish_reason="stop")
+                return
+
+            payload["stream"] = False
+            response = await client.post(f"{self.base_url}/api/chat", json=payload)
+            if response.status_code >= 400:
+                raise BackendError(f"Ollama 请求失败 {response.status_code}: {response.text}")
+            data = response.json()
+            message = data.get("message") or {}
+            yield ChatChunk(
+                text=message.get("content") or "",
+                delta=message,
+                done=True,
+                finish_reason=data.get("done_reason") or "stop",
+                metadata=data,
+            )
+
+    async def embed(self, inputs: list[str]) -> list[list[float]]:
+        if not self.loaded:
+            await self.load()
+        async with httpx.AsyncClient(timeout=120.0) as client:
+            response = await client.post(
+                f"{self.base_url}/api/embed",
+                json={"model": self.remote_model, "input": inputs},
+            )
+        if response.status_code >= 400:
+            raise BackendError(f"Ollama embeddings 失败 {response.status_code}: {response.text}")
+        data = response.json()
+        return data.get("embeddings", [])
+
+    def info(self) -> dict[str, Any]:
+        return {
+            "id": self.model.id,
+            "kind": "ollama",
+            "remote_model": self.remote_model,
+            "loaded": self.loaded,
+            "endpoint": self.base_url,
+        }

localpilot/backends/onnx.py

@@ -0,0 +1,178 @@
+from __future__ import annotations
+
+import asyncio
+import queue
+import threading
+from collections.abc import AsyncIterator
+from pathlib import Path
+from typing import Any
+
+from ..config import ModelConfig, RuntimeConfig
+from ..types import BackendError, ChatChunk, ChatParams
+from .base import BaseBackend
+
+
+class ONNXBackend(BaseBackend):
+    """ONNX loader with an Optimum path when installed, plus safe inspection fallback."""
+
+    def __init__(self, model: ModelConfig, runtime: RuntimeConfig) -> None:
+        super().__init__(model, runtime)
+        self.session: Any = None
+        self.ort_model: Any = None
+        self.tokenizer: Any = None
+        self.providers: list[str] = []
+
+    async def load(self) -> None:
+        if self.loaded:
+            return
+        if not self.model.path:
+            raise BackendError(f"ONNX 模型缺少 path: {self.model.id}")
+        path = Path(self.model.path)
+        if not path.exists():
+            raise BackendError(f"ONNX 文件或目录不存在: {path}")
+        try:
+            import onnxruntime as ort
+        except ImportError as exc:
+            raise BackendError("当前 Conda LLM 环境没有 onnxruntime") from exc
+
+        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        available = ort.get_available_providers()
+        providers = [provider for provider in providers if provider in available]
+        if not providers:
+            providers = available
+        self.providers = providers
+        try:
+            if path.is_file():
+                self.session = await asyncio.to_thread(ort.InferenceSession, str(path), providers=providers)
+            else:
+                onnx_file = next(path.glob("*.onnx"), None)
+                if onnx_file is None:
+                    raise BackendError(f"目录中没有 .onnx 文件: {path}")
+                self.session = await asyncio.to_thread(
+                    ort.InferenceSession, str(onnx_file), providers=providers
+                )
+        except Exception as exc:
+            raise BackendError(f"ONNX Runtime 加载失败: {exc}") from exc
+
+        try:
+            from optimum.onnxruntime import ORTModelForCausalLM
+            from transformers import AutoTokenizer
+
+            model_dir = path if path.is_dir() else path.parent
+            self.ort_model = await asyncio.to_thread(
+                ORTModelForCausalLM.from_pretrained,
+                model_dir,
+                provider=providers[0] if providers else "CPUExecutionProvider",
+            )
+            self.tokenizer = await asyncio.to_thread(
+                AutoTokenizer.from_pretrained, model_dir, local_files_only=True
+            )
+        except ImportError:
+            self.ort_model = None
+        self.loaded = True
+
+    async def unload(self) -> None:
+        self.session = None
+        self.ort_model = None
+        self.tokenizer = None
+        self.loaded = False
+
+    async def chat(
+        self,
+        messages: list[dict[str, Any]],
+        params: ChatParams,
+        stream: bool = False,
+    ) -> AsyncIterator[ChatChunk]:
+        if not self.loaded:
+            await self.load()
+        if self.ort_model is None or self.tokenizer is None:
+            raise BackendError(
+                "ONNX 文件已加载并可检查，但聊天生成需要安装 optimum[onnxruntime]，"
+                "并使用带 tokenizer/config 的 ONNX CausalLM 目录。"
+            )
+        kwargs: dict[str, Any] = {
+            "add_generation_prompt": True,
+            "tokenize": True,
+            "return_tensors": "pt",
+            "return_dict": True,
+        }
+        if params.enable_thinking is not None:
+            kwargs["enable_thinking"] = params.enable_thinking
+        try:
+            inputs = self.tokenizer.apply_chat_template(messages, **kwargs)
+        except TypeError:
+            kwargs.pop("enable_thinking", None)
+            inputs = self.tokenizer.apply_chat_template(messages, **kwargs)
+        generation_kwargs: dict[str, Any] = {
+            "max_new_tokens": params.max_tokens,
+            "do_sample": params.temperature > 0,
+            "use_cache": True,
+            "repetition_penalty": params.repeat_penalty,
+            "pad_token_id": self.tokenizer.pad_token_id,
+            "eos_token_id": self.tokenizer.eos_token_id,
+        }
+        if params.temperature > 0:
+            generation_kwargs.update({"temperature": params.temperature, "top_p": params.top_p, "top_k": params.top_k})
+        generation_kwargs.update(params.extra)
+
+        if not stream:
+            output = await asyncio.to_thread(self.ort_model.generate, **inputs, **generation_kwargs)
+            prompt_len = int(inputs["input_ids"].shape[-1])
+            text = self.tokenizer.decode(output[0][prompt_len:], skip_special_tokens=True)
+            for stop in params.stop:
+                text = text.split(stop, 1)[0]
+            yield ChatChunk(text=text, done=True, finish_reason="stop")
+            return
+
+        from transformers import TextIteratorStreamer
+
+        streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)
+        events: queue.Queue[tuple[str, Any]] = queue.Queue()
+
+        def worker() -> None:
+            try:
+                self.ort_model.generate(**inputs, streamer=streamer, **generation_kwargs)
+            except BaseException as exc:
+                events.put(("error", exc))
+                try:
+                    streamer.end()
+                except Exception:
+                    pass
+
+        def forward_stream() -> None:
+            try:
+                for item in streamer:
+                    events.put(("text", item))
+                events.put(("done", None))
+            except BaseException as exc:
+                events.put(("error", exc))
+
+        threading.Thread(target=worker, daemon=True).start()
+        threading.Thread(target=forward_stream, daemon=True).start()
+        while True:
+            kind, value = await asyncio.to_thread(events.get)
+            if kind == "text":
+                text = str(value)
+                for stop in params.stop:
+                    text = text.split(stop, 1)[0]
+                if text:
+                    yield ChatChunk(text=text)
+            elif kind == "error":
+                raise BackendError(f"ONNX 推理失败: {value}") from value
+            else:
+                yield ChatChunk(done=True, finish_reason="stop")
+                return
+
+    def info(self) -> dict[str, Any]:
+        inputs: list[str] = []
+        if self.session is not None:
+            inputs = [item.name for item in self.session.get_inputs()]
+        return {
+            "id": self.model.id,
+            "kind": "onnx",
+            "path": self.model.path,
+            "loaded": self.loaded,
+            "providers": self.providers,
+            "inputs": inputs,
+            "generation_ready": self.ort_model is not None and self.tokenizer is not None,
+        }

localpilot/backends/remote.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import json
+from collections.abc import AsyncIterator
+from typing import Any
+
+import httpx
+
+from ..config import CloudProfile, ModelConfig, RuntimeConfig
+from ..types import BackendError, ChatChunk, ChatParams
+from .base import BaseBackend
+
+
+class RemoteBackend(BaseBackend):
+    """Proxy any OpenAI-compatible cloud endpoint without exposing its key to clients."""
+
+    def __init__(self, model: ModelConfig, profile: CloudProfile, runtime: RuntimeConfig) -> None:
+        super().__init__(model, runtime)
+        self.profile = profile
+        self.base_url = profile.base_url.rstrip("/")
+        if not self.base_url.endswith("/v1"):
+            self.base_url += "/v1"
+
+    async def load(self) -> None:
+        key = self.profile.resolved_api_key()
+        if not key:
+            raise BackendError(
+                f"云端配置 {self.profile.id} 没有 API key。请设置环境变量 {self.profile.api_key_env or '<api_key_env>'}。"
+            )
+        self.loaded = True
+
+    async def unload(self) -> None:
+        self.loaded = False
+
+    def _headers(self) -> dict[str, str]:
+        key = self.profile.resolved_api_key()
+        if not key:
+            raise BackendError(f"云端配置 {self.profile.id} 的 API key 不可用")
+        headers = {"Authorization": f"Bearer {key}", "Content-Type": "application/json"}
+        headers.update(self.profile.headers)
+        return headers
+
+    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.remote_model or self.profile.default_model or self.model.id,
+            "messages": messages,
+            "stream": stream,
+            "max_tokens": params.max_tokens,
+            "temperature": params.temperature,
+            "top_p": params.top_p,
+        }
+        if params.stop:
+            payload["stop"] = params.stop
+        if params.seed is not None:
+            payload["seed"] = params.seed
+        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(self.profile.timeout, connect=15.0)
+        async with httpx.AsyncClient(timeout=timeout, headers=self._headers()) as client:
+            if stream:
+                async with client.stream("POST", f"{self.base_url}/chat/completions", json=payload) as response:
+                    if response.status_code >= 400:
+                        detail = (await response.aread()).decode("utf-8", errors="replace")
+                        raise BackendError(f"云端请求失败 {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.base_url}/chat/completions", json=payload)
+            if response.status_code >= 400:
+                raise BackendError(f"云端请求失败 {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": "cloud",
+            "profile": self.profile.id,
+            "remote_model": self.model.remote_model or self.profile.default_model or self.model.id,
+            "loaded": self.loaded,
+            "base_url": self.base_url,
+        }

localpilot/backends/transformers_backend.py

@@ -0,0 +1,330 @@
+from __future__ import annotations
+
+import asyncio
+import gc
+import queue
+import threading
+import time
+from collections.abc import AsyncIterator
+from pathlib import Path
+from typing import Any
+
+from ..config import ModelConfig, RuntimeConfig
+from ..types import BackendError, ChatChunk, ChatParams
+from .base import BaseBackend
+
+
+class TransformersBackend(BaseBackend):
+    """Load a Transformers/safetensors causal LM in the existing Conda LLM env."""
+
+    def __init__(self, model: ModelConfig, runtime: RuntimeConfig) -> None:
+        super().__init__(model, runtime)
+        self.tokenizer: Any = None
+        self.model_object: Any = None
+        self.device: Any = None
+        self._generation_lock = asyncio.Lock()
+
+    async def load(self) -> None:
+        if self.loaded:
+            return
+        if not self.model.path:
+            raise BackendError(f"Transformers 模型缺少 path: {self.model.id}")
+        path = Path(self.model.path)
+        if not path.exists():
+            raise BackendError(f"模型目录不存在: {path}")
+        try:
+            self.tokenizer, self.model_object, self.device = await asyncio.to_thread(self._load_sync, path)
+        except Exception as exc:
+            raise BackendError(f"Transformers 模型加载失败: {exc}") from exc
+        self.loaded = True
+
+    def _load_sync(self, path: Path) -> tuple[Any, Any, Any]:
+        import torch
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
+        options = self.model.options
+        if torch.cuda.is_available():
+            try:
+                torch.backends.cuda.matmul.fp32_precision = "tf32"
+                torch.backends.cudnn.conv.fp32_precision = "tf32"
+            except AttributeError:
+                pass
+        model_dir = path if path.is_dir() else path.parent
+        dtype_name = str(options.get("torch_dtype", "float16")).lower()
+        dtype = {
+            "float16": torch.float16,
+            "fp16": torch.float16,
+            "bfloat16": torch.bfloat16,
+            "bf16": torch.bfloat16,
+            "float32": torch.float32,
+            "fp32": torch.float32,
+        }.get(dtype_name, torch.float16)
+        device_name = str(options.get("device", "cuda" if torch.cuda.is_available() else "cpu"))
+        if device_name == "cuda" and not torch.cuda.is_available():
+            device_name = "cpu"
+        tokenizer = AutoTokenizer.from_pretrained(model_dir, use_fast=True, local_files_only=True)
+        if tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        load_kwargs: dict[str, Any] = {
+            "dtype": dtype,
+            "low_cpu_mem_usage": True,
+        }
+        if options.get("attn_implementation"):
+            load_kwargs["attn_implementation"] = options["attn_implementation"]
+        quantization = str(options.get("quantization", "none")).lower()
+        load_in_4bit = bool(options.get("load_in_4bit", False)) or quantization in {"4bit", "int4", "nf4"}
+        load_in_8bit = bool(options.get("load_in_8bit", False)) or quantization in {"8bit", "int8", "bnb8"}
+        dynamic_int8_cpu = bool(options.get("dynamic_int8_cpu", False)) or (
+            quantization in {"dynamic-int8", "int8-dynamic"} and device_name == "cpu"
+        )
+        if (load_in_4bit or load_in_8bit) and device_name.startswith("cuda"):
+            try:
+                from transformers import BitsAndBytesConfig
+
+                if load_in_4bit:
+                    load_kwargs["quantization_config"] = BitsAndBytesConfig(
+                        load_in_4bit=True,
+                        bnb_4bit_quant_type=str(options.get("bnb_4bit_quant_type", "nf4")),
+                        bnb_4bit_compute_dtype=dtype,
+                        bnb_4bit_use_double_quant=bool(options.get("bnb_4bit_use_double_quant", True)),
+                    )
+                else:
+                    load_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)
+                load_kwargs["device_map"] = options.get("device_map", "auto")
+            except Exception as exc:
+                raise BackendError(f"4-bit 量化加载失败，请检查 bitsandbytes/CUDA: {exc}") from exc
+        else:
+            try:
+                load_kwargs["device_map"] = options.get("device_map", "auto")
+                model_object = AutoModelForCausalLM.from_pretrained(model_dir, **load_kwargs)
+            except TypeError as first_exc:
+                if "dtype" not in str(first_exc):
+                    raise
+                load_kwargs.pop("dtype", None)
+                load_kwargs["torch_dtype"] = dtype
+                model_object = AutoModelForCausalLM.from_pretrained(model_dir, **load_kwargs)
+            except (ImportError, ValueError) as first_exc:
+                load_kwargs.pop("device_map", None)
+                try:
+                    model_object = AutoModelForCausalLM.from_pretrained(model_dir, **load_kwargs)
+                    model_object.to(device_name)
+                except Exception as second_exc:
+                    raise RuntimeError(f"自动设备映射失败: {first_exc}; 手动放置也失败: {second_exc}") from second_exc
+            if dynamic_int8_cpu and device_name == "cpu":
+                model_object = torch.quantization.quantize_dynamic(
+                    model_object, {torch.nn.Linear}, dtype=torch.qint8
+                )
+        if "model_object" not in locals():
+            model_object = AutoModelForCausalLM.from_pretrained(model_dir, **load_kwargs)
+        adapters = options.get("adapters", [])
+        if isinstance(adapters, str):
+            adapters = [adapters]
+        for adapter in adapters:
+            try:
+                from peft import PeftModel
+
+                model_object = PeftModel.from_pretrained(model_object, adapter)
+            except ImportError as exc:
+                raise BackendError("配置了 ADAPTER，但当前环境缺少 peft") from exc
+        model_object.eval()
+        actual_device = next(model_object.parameters()).device
+        return tokenizer, model_object, actual_device
+
+    async def unload(self) -> None:
+        self.loaded = False
+        self.tokenizer = None
+        self.model_object = None
+        self.device = None
+        gc.collect()
+        try:
+            import torch
+
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+        except Exception:
+            pass
+
+    def _prepare_inputs_sync(self, messages: list[dict[str, Any]], params: ChatParams) -> dict[str, Any]:
+        prepared_messages: list[dict[str, Any]] = []
+        for message in messages:
+            item = dict(message)
+            content = item.get("content")
+            if isinstance(content, list):
+                item["content"] = "".join(
+                    str(part.get("text", "")) if isinstance(part, dict) and part.get("type") == "text" else "[image]"
+                    for part in content
+                )
+            elif not isinstance(content, str):
+                item["content"] = str(content)
+            prepared_messages.append(item)
+        kwargs: dict[str, Any] = {
+            "add_generation_prompt": True,
+            "tokenize": True,
+            "return_tensors": "pt",
+            "return_dict": True,
+        }
+        if self.model.template:
+            kwargs["chat_template"] = self.model.template
+        if params.enable_thinking is not None:
+            kwargs["enable_thinking"] = params.enable_thinking
+        try:
+            inputs = self.tokenizer.apply_chat_template(prepared_messages, **kwargs)
+        except TypeError:
+            kwargs.pop("enable_thinking", None)
+            inputs = self.tokenizer.apply_chat_template(prepared_messages, **kwargs)
+        return {key: value.to(self.device) if hasattr(value, "to") else value for key, value in inputs.items()}
+
+    def _generation_kwargs(self, params: ChatParams) -> dict[str, Any]:
+        kwargs: dict[str, Any] = {
+            "max_new_tokens": params.max_tokens,
+            "do_sample": params.temperature > 0,
+            "use_cache": True,
+            "repetition_penalty": params.repeat_penalty,
+            "pad_token_id": self.tokenizer.pad_token_id,
+            "eos_token_id": self.tokenizer.eos_token_id,
+        }
+        if params.temperature > 0:
+            kwargs["temperature"] = params.temperature
+            kwargs["top_k"] = params.top_k
+            kwargs["top_p"] = params.top_p
+        allowed_extra = {
+            "do_sample", "num_beams", "num_return_sequences", "typical_p", "epsilon_cutoff", "eta_cutoff",
+            "diversity_penalty", "length_penalty", "no_repeat_ngram_size", "max_length", "max_new_tokens",
+            "renormalize_logits", "remove_invalid_values", "bad_words_ids", "force_words_ids",
+            "repetition_penalty", "stop_strings",
+        }
+        kwargs.update({key: value for key, value in params.extra.items() if key in allowed_extra})
+        if params.seed is not None:
+            import torch
+
+            torch.manual_seed(params.seed)
+        return kwargs
+
+    def _clean_stop(self, text: str, params: ChatParams) -> str:
+        for stop in params.stop:
+            if stop in text:
+                text = text.split(stop, 1)[0]
+        return text
+
+    async def chat(
+        self,
+        messages: list[dict[str, Any]],
+        params: ChatParams,
+        stream: bool = False,
+    ) -> AsyncIterator[ChatChunk]:
+        if not self.loaded:
+            await self.load()
+        async with self._generation_lock:
+            inputs = await asyncio.to_thread(self._prepare_inputs_sync, messages, params)
+            kwargs = self._generation_kwargs(params)
+            if not stream:
+                started = time.perf_counter()
+                output = await asyncio.to_thread(self.model_object.generate, **inputs, **kwargs)
+                prompt_len = int(inputs["input_ids"].shape[-1])
+                text = self.tokenizer.decode(output[0][prompt_len:], skip_special_tokens=True)
+                output_tokens = max(0, int(output.shape[-1]) - prompt_len)
+                duration = max(time.perf_counter() - started, 1e-9)
+                yield ChatChunk(
+                    text=self._clean_stop(text, params),
+                    done=True,
+                    finish_reason="stop",
+                    metadata={
+                        "usage": {
+                            "prompt_tokens": prompt_len,
+                            "completion_tokens": output_tokens,
+                            "total_tokens": prompt_len + output_tokens,
+                        },
+                        "timings": {
+                            "total_seconds": duration,
+                            "tokens_per_second": output_tokens / duration,
+                        },
+                    },
+                )
+                return
+
+            from transformers import TextIteratorStreamer
+
+            streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)
+            events: queue.Queue[tuple[str, Any]] = queue.Queue()
+
+            def worker() -> None:
+                try:
+                    self.model_object.generate(**inputs, streamer=streamer, **kwargs)
+                except BaseException as exc:
+                    events.put(("error", exc))
+                    try:
+                        streamer.end()
+                    except Exception:
+                        pass
+
+            def forward_stream() -> None:
+                try:
+                    for item in streamer:
+                        events.put(("text", item))
+                    events.put(("done", None))
+                except BaseException as exc:
+                    events.put(("error", exc))
+
+            generation_thread = threading.Thread(target=worker, name=f"localpilot-gen-{self.model.id}", daemon=True)
+            stream_thread = threading.Thread(target=forward_stream, name=f"localpilot-stream-{self.model.id}", daemon=True)
+            generation_thread.start()
+            stream_thread.start()
+            while True:
+                kind, value = await asyncio.to_thread(events.get)
+                if kind == "text":
+                    text = self._clean_stop(str(value), params)
+                    if text:
+                        yield ChatChunk(text=text)
+                elif kind == "error":
+                    raise BackendError(f"Transformers 推理失败: {value}") from value
+                elif kind == "done":
+                    yield ChatChunk(done=True, finish_reason="stop")
+                    return
+
+    def info(self) -> dict[str, Any]:
+        vram: dict[str, int] = {}
+        try:
+            import torch
+
+            if torch.cuda.is_available() and self.loaded:
+                vram = {
+                    "vram_allocated_bytes": int(torch.cuda.memory_allocated()),
+                    "vram_reserved_bytes": int(torch.cuda.memory_reserved()),
+                }
+        except Exception:
+            pass
+        return {
+            "id": self.model.id,
+            "kind": "transformers",
+            "path": self.model.path,
+            "loaded": self.loaded,
+            "device": str(self.device) if self.device is not None else None,
+            "load_in_4bit": bool(self.model.options.get("load_in_4bit", False)),
+            "quantization": self.model.options.get("quantization", "none"),
+            "dynamic_int8_cpu": bool(self.model.options.get("dynamic_int8_cpu", False)),
+            **vram,
+        }
+
+    async def embed(self, inputs: list[str]) -> list[list[float]]:
+        if not self.loaded:
+            await self.load()
+        return await asyncio.to_thread(self._embed_sync, inputs)
+
+    def _embed_sync(self, inputs: list[str]) -> list[list[float]]:
+        import torch
+
+        tokens = self.tokenizer(
+            inputs,
+            padding=True,
+            truncation=True,
+            return_tensors="pt",
+        )
+        tokens = {key: value.to(self.device) if hasattr(value, "to") else value for key, value in tokens.items()}
+        with torch.inference_mode():
+            outputs = self.model_object(**tokens, output_hidden_states=True, use_cache=False)
+        hidden = outputs.hidden_states[-1]
+        mask = tokens["attention_mask"].unsqueeze(-1).to(hidden.dtype)
+        pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1)
+        pooled = torch.nn.functional.normalize(pooled, p=2, dim=1)
+        return pooled.float().cpu().tolist()