From d159212416e4d6eb64120d86390bc9dc8f6f485b Mon Sep 17 00:00:00 2001 From: WpyQwq <3911625973@qq.com> Date: Sat, 19 Sep 2026 11:57:59 +0800 Subject: [PATCH] =?UTF-8?q?Initial=20commit:=20LocalPilot=EF=BC=9A?= =?UTF-8?q?=E6=9C=AC=E5=9C=B0=E6=A8=A1=E5=9E=8B=E8=BF=90=E8=A1=8C=E6=97=B6?= =?UTF-8?q?=EF=BC=8C=E5=A4=8D=E7=94=A8=20llama-server=20=E5=B9=B6=E6=8F=90?= =?UTF-8?q?=E4=BE=9B=20Ollama=20=E5=85=BC=E5=AE=B9=20provider?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 11 + README.md | 83 +++ config/config.example.json | 56 ++ config/config.json | 78 +++ localpilot/__init__.py | 4 + localpilot/__main__.py | 4 + localpilot/backends/__init__.py | 8 + localpilot/backends/base.py | 39 ++ localpilot/backends/gguf.py | 241 ++++++++ localpilot/backends/ollama.py | 127 ++++ localpilot/backends/onnx.py | 178 ++++++ localpilot/backends/remote.py | 121 ++++ localpilot/backends/transformers_backend.py | 330 ++++++++++ localpilot/cli.py | 111 ++++ localpilot/config.py | 135 +++++ localpilot/manager.py | 369 ++++++++++++ localpilot/modelfile.py | 107 ++++ localpilot/server.py | 634 ++++++++++++++++++++ localpilot/tui.py | 169 ++++++ localpilot/types.py | 50 ++ pyproject.toml | 27 + scripts/start-server.ps1 | 9 + scripts/start-tui.ps1 | 9 + tests/test_config.py | 26 + tests/test_manager.py | 28 + tests/test_modelfile.py | 27 + tests/test_server_helpers.py | 41 ++ 27 files changed, 3022 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 config/config.example.json create mode 100644 config/config.json create mode 100644 localpilot/__init__.py create mode 100644 localpilot/__main__.py create mode 100644 localpilot/backends/__init__.py create mode 100644 localpilot/backends/base.py create mode 100644 localpilot/backends/gguf.py create mode 100644 localpilot/backends/ollama.py create mode 100644 localpilot/backends/onnx.py create mode 100644 localpilot/backends/remote.py create mode 100644 localpilot/backends/transformers_backend.py create mode 100644 localpilot/cli.py create mode 100644 localpilot/config.py create mode 100644 localpilot/manager.py create mode 100644 localpilot/modelfile.py create mode 100644 localpilot/server.py create mode 100644 localpilot/tui.py create mode 100644 localpilot/types.py create mode 100644 pyproject.toml create mode 100644 scripts/start-server.ps1 create mode 100644 scripts/start-tui.ps1 create mode 100644 tests/test_config.py create mode 100644 tests/test_manager.py create mode 100644 tests/test_modelfile.py create mode 100644 tests/test_server_helpers.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..376aac4 --- /dev/null +++ b/.gitignore @@ -0,0 +1,11 @@ +__pycache__/ +*.py[cod] +.pytest_cache/ +.mypy_cache/ +.ruff_cache/ +*.egg-info/ +runtime/*.log +cache/ +models/ +secrets.json + diff --git a/README.md b/README.md new file mode 100644 index 0000000..a2c6f3e --- /dev/null +++ b/README.md @@ -0,0 +1,83 @@ +# LocalPilot + +LocalPilot 是一个放在 E 盘的本地模型运行时，目标是把 Ollama 的易用生命周期和更开放的模型适配组合起来： + +- GGUF：复用本机 CUDA 版 `llama-server.exe`，自动分配端口、启停、GPU offload、Flash Attention、continuous batching。 +- Ollama 兼容 provider：可以直接调用本机 Ollama 已管理的模型，不复制权重；适合 Ollama blob 与独立 llama.cpp runner 版本不兼容时使用。 +- safetensors / Transformers：在 Conda `LLM` 环境中使用 PyTorch CUDA，支持 `device_map`、FP16/BF16、4-bit bitsandbytes、CPU 动态 INT8。 +- ONNX：提供 ONNX Runtime 检查和 Optimum 接入点；没有 Optimum 时会给出清晰提示。 +- 本地中转站：FastAPI 同时提供 OpenAI 路由和 Ollama 风格的 `/api/chat`、`/api/generate`、`/api/tags`、`/api/show`、`/api/ps`、`/api/embed`。 +- 云端代理：把 OpenAI/DeepSeek/其他 OpenAI-compatible endpoint 配进 `cloud_profiles`，API key 只从环境变量读取。 +- TUI：Textual 实现，支持切模型、加载/卸载、流式对话和动态调整生成参数。 +- 生命周期：支持 `keep_alive`、显存模型 LRU、活动请求计数、模型注册/删除（只改配置，不删除权重文件）。 +- 模型定义：支持 Modelfile 的 `FROM`、`PARAMETER`、`SYSTEM`、`TEMPLATE`、`MESSAGE`、`ADAPTER` 子集，以及 `/api/create`、`/api/copy`、本地路径导入和 Hugging Face `/api/pull`。 + +## 启动 + +当前机器已有 Conda 环境 `LLM` 和 CUDA 版 `llama-server.exe`，项目配置已指向它： + +```powershell +Set-Location E:\llm-backend +conda run -n LLM python -m localpilot doctor +conda run -n LLM python -m localpilot models +conda run -n LLM python -m localpilot tui +``` + +启动本地 OpenAI 兼容服务： + +```powershell +Set-Location E:\llm-backend +.\scripts\start-server.ps1 +``` + +服务默认监听 `127.0.0.1:8787`。客户端示例： + +```powershell +Invoke-RestMethod http://127.0.0.1:8787/v1/chat/completions ` + -Method Post -ContentType 'application/json' ` + -Body (@{ + model='apex-qwen3' + messages=@(@{role='user';content='你好，请用一句话介绍自己'}) + stream=$false + max_tokens=128 + temperature=0.4 + } | ConvertTo-Json -Depth 8) +``` + +## TUI 命令 + +```text +/models +/model apex-qwen3 +/load apex-qwen3 +/unload +/params temperature=0.2 top_p=0.9 top_k=40 max_tokens=512 enable_thinking=false +/clear +``` + +## 与 Ollama 的边界 + +这不是把 Ollama 的多年工程压缩成几百行代码。当前已经补上了协议、生命周期、配置模型、本地仓库导入和代理这些高价值差距，但以下能力仍明确保留为下一阶段： + +- `/api/pull` 支持本地路径导入和 Hugging Face snapshot 下载，但不实现 Ollama registry 的分层 blob、digest 校验、push。 +- GGUF 直连依赖兼容的 llama.cpp runner；Ollama 的内部 blob 不保证能被任意版本的独立 runner 读取。遇到这种情况，用 `kind: "ollama"` 的兼容 provider，或换成标准 GGUF 文件。 +- Transformers 的工具调用和 JSON Schema 目前是参数透传/提示约束，不等价于 Ollama runner 对每个模型的原生模板执行。 +- 目前有显存槽位和请求排队，但还没有 Ollama 那种完整的跨 GPU 显存规划、连续批处理调度、模型分层缓存和自动下载器。 + +## 配置模型 + +编辑 `config/config.json`。GGUF 的 `path` 指向单个 `.gguf` 文件；safetensors 的 `path` 指向包含 `config.json`、tokenizer 和权重的 Transformers 目录。项目默认已登记本机的 `E:\ollama\Apex_AI-catmodel`，但不会启动时自动加载，首次选择或请求时才加载。 + +需要节省显存时，可在 Transformers 模型的 `options` 中设置 `"quantization": "4bit"` 或 `"quantization": "8bit"`；CPU 场景可设置 `"quantization": "dynamic-int8"`。量化发生在模型加载阶段，切换量化配置后重新加载模型即可。 + +云端模型只配置环境变量名，例如： + +```powershell +$env:DEEPSEEK_API_KEY = 'your-key' +``` + +然后把 `deepseek-cloud` 的 `enabled` 改为 `true`，并请求 `model: "deepseek-cloud"`。 + +## 设计边界 + +权重量化通常是加载阶段能力，不能安全地在每个请求中重新量化；LocalPilot 将它放在模型 `options` 中。Transformers 后端支持 FP16/BF16、bitsandbytes 4/8-bit，以及 CPU 动态 INT8；GGUF 的量化策略由 GGUF 文件和 llama.cpp 决定。对话参数（temperature、top_p、top_k、min_p、repeat_penalty、seed、stop、max_tokens、enable_thinking）支持按请求动态覆盖，OpenAI 请求中的 `extra_body` 会透传给后端以支持模型特有参数。后续可以继续接入 AWQ/GPTQ、MLC、TensorRT-LLM、更多 ONNX CausalLM 结构和多模型并发调度，而不改 OpenAI API 层。 diff --git a/config/config.example.json b/config/config.example.json new file mode 100644 index 0000000..13723a0 --- /dev/null +++ b/config/config.example.json @@ -0,0 +1,56 @@ +{ + "runtime": { + "llama_server": "E:\\path\\to\\llama-server.exe", + "ollama_url": "http://127.0.0.1:11434", + "host": "127.0.0.1", + "port": 8787, + "default_model": "my-model", + "runtime_dir": "E:\\llm-backend\\runtime", + "cache_dir": "E:\\llm-backend\\cache", + "n_gpu_layers": "auto", + "ctx_size": 32768, + "flash_attn": "auto", + "fit_vram": true, + "auto_load": false, + "max_loaded_models": 1, + "keep_alive_seconds": 300, + "api_key": null + }, + "models": [ + { + "id": "my-gguf", + "kind": "gguf", + "path": "E:\\model\\my-model.Q4_K_M.gguf", + "enabled": true, + "options": {} + }, + { + "id": "my-safetensors", + "kind": "transformers", + "path": "E:\\model\\my-model", + "enabled": true, + "options": { + "torch_dtype": "float16", + "device": "cuda", + "device_map": "auto", + "quantization": "none", + "load_in_4bit": false, + "dynamic_int8_cpu": false + } + }, + { + "id": "my-ollama-model", + "kind": "ollama", + "remote_model": "mistral-nemo:latest", + "enabled": false + } + ], + "cloud_profiles": [ + { + "id": "openai-compatible", + "base_url": "https://api.openai.com/v1", + "api_key_env": "OPENAI_API_KEY", + "default_model": "gpt-4o-mini" + } + ] +} diff --git a/config/config.json b/config/config.json new file mode 100644 index 0000000..66005e4 --- /dev/null +++ b/config/config.json @@ -0,0 +1,78 @@ +{ + "runtime": { + "llama_server": "E:\\ollama\\模型\\llama-b8642-bin-win-cuda-13.1-x64\\llama-server.exe", + "ollama_url": "http://127.0.0.1:11434", + "host": "127.0.0.1", + "port": 8787, + "default_model": "apex-qwen3", + "model_dir": "E:\\model", + "runtime_dir": "E:\\llm-backend\\runtime", + "cache_dir": "E:\\llm-backend\\cache", + "n_gpu_layers": "auto", + "ctx_size": 32768, + "batch_size": 2048, + "ubatch_size": 512, + "flash_attn": "auto", + "fit_vram": true, + "auto_load": false, + "api_key": null, + "max_loaded_models": 1, + "keep_alive_seconds": 300 + }, + "models": [ + { + "id": "apex-qwen3", + "kind": "transformers", + "path": "E:\\ollama\\Apex_AI-catmodel", + "enabled": true, + "description": "本机 Qwen3 safetensors 模型", + "options": { + "torch_dtype": "float16", + "device": "cuda", + "device_map": "auto", + "load_in_4bit": false, + "attn_implementation": "sdpa" + } + }, + { + "id": "gguf-local", + "kind": "gguf", + "path": "E:\\model\\example.gguf", + "enabled": false, + "description": "将 path 改成真实 GGUF 文件后启用", + "options": {} + }, + { + "id": "onnx-local", + "kind": "onnx", + "path": "E:\\model\\example.onnx", + "enabled": false, + "description": "ONNX Runtime/Optimum 扩展入口", + "options": {} + }, + { + "id": "mistral-nemo-ollama", + "kind": "ollama", + "remote_model": "mistral-nemo-12b:latest", + "enabled": false, + "description": "兼容本机 Ollama 模型仓库；不复制权重" + }, + { + "id": "deepseek-cloud", + "kind": "cloud", + "cloud_profile": "deepseek", + "remote_model": "deepseek-chat", + "enabled": false, + "description": "OpenAI 兼容云端模型" + } + ], + "cloud_profiles": [ + { + "id": "deepseek", + "base_url": "https://api.deepseek.com", + "api_key_env": "DEEPSEEK_API_KEY", + "default_model": "deepseek-chat", + "timeout": 120 + } + ] +} diff --git a/localpilot/__init__.py b/localpilot/__init__.py new file mode 100644 index 0000000..9f992e3 --- /dev/null +++ b/localpilot/__init__.py @@ -0,0 +1,4 @@ +"""LocalPilot: a local model runtime and OpenAI-compatible gateway.""" + +__version__ = "0.1.0" + diff --git a/localpilot/__main__.py b/localpilot/__main__.py new file mode 100644 index 0000000..e55e7d3 --- /dev/null +++ b/localpilot/__main__.py @@ -0,0 +1,4 @@ +from .cli import main + +main() + diff --git a/localpilot/backends/__init__.py b/localpilot/backends/__init__.py new file mode 100644 index 0000000..7b1d4f7 --- /dev/null +++ b/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"] diff --git a/localpilot/backends/base.py b/localpilot/backends/base.py new file mode 100644 index 0000000..d080b49 --- /dev/null +++ b/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") diff --git a/localpilot/backends/gguf.py b/localpilot/backends/gguf.py new file mode 100644 index 0000000..ca05908 --- /dev/null +++ b/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, + } diff --git a/localpilot/backends/ollama.py b/localpilot/backends/ollama.py new file mode 100644 index 0000000..5d68f09 --- /dev/null +++ b/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, + } diff --git a/localpilot/backends/onnx.py b/localpilot/backends/onnx.py new file mode 100644 index 0000000..50bfc20 --- /dev/null +++ b/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, + } diff --git a/localpilot/backends/remote.py b/localpilot/backends/remote.py new file mode 100644 index 0000000..8013ddb --- /dev/null +++ b/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 ''}。" + ) + 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, + } diff --git a/localpilot/backends/transformers_backend.py b/localpilot/backends/transformers_backend.py new file mode 100644 index 0000000..c6b9224 --- /dev/null +++ b/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() diff --git a/localpilot/cli.py b/localpilot/cli.py new file mode 100644 index 0000000..ee91512 --- /dev/null +++ b/localpilot/cli.py @@ -0,0 +1,111 @@ +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import sys +from pathlib import Path + +from .config import infer_kind, load_config + + +PROJECT_DIR = Path(__file__).resolve().parent.parent +DEFAULT_CONFIG = PROJECT_DIR / "config" / "config.json" + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(prog="localpilot", description="LocalPilot 本地模型后端") + parser.add_argument("--config", default=str(DEFAULT_CONFIG), help="配置 JSON 路径") + sub = parser.add_subparsers(dest="command", required=True) + sub.add_parser("models", help="列出配置模型") + sub.add_parser("doctor", help="检查本机运行环境") + serve = sub.add_parser("serve", help="启动 OpenAI 兼容 API") + serve.add_argument("--host") + serve.add_argument("--port", type=int) + sub.add_parser("tui", help="启动 Textual TUI") + load = sub.add_parser("load", help="预加载模型") + load.add_argument("model_id") + inspect = sub.add_parser("inspect", help="推断模型格式") + inspect.add_argument("path") + return parser + + +def _doctor() -> dict[str, object]: + result: dict[str, object] = {"python": sys.version, "executable": sys.executable} + try: + import torch + + result["torch"] = torch.__version__ + result["cuda_available"] = bool(torch.cuda.is_available()) + result["cuda_device"] = torch.cuda.get_device_name(0) if torch.cuda.is_available() else None + result["cuda_version"] = torch.version.cuda + except Exception as exc: + result["torch_error"] = str(exc) + try: + import onnxruntime as ort + + result["onnxruntime"] = ort.__version__ + result["onnx_providers"] = ort.get_available_providers() + except Exception as exc: + result["onnx_error"] = str(exc) + result["conda_prefix"] = os.getenv("CONDA_PREFIX") + return result + + +def main() -> None: + args = _parser().parse_args() + config_path = Path(args.config) + config = load_config(config_path) + + if args.command == "models": + print(json.dumps([ + {"id": model.id, "kind": infer_kind(model), "path": model.path, "enabled": model.enabled} + for model in config.models + if model.enabled + ], ensure_ascii=False, indent=2)) + return + if args.command == "doctor": + print(json.dumps(_doctor(), ensure_ascii=False, indent=2)) + return + if args.command == "inspect": + from .config import ModelConfig + + spec = ModelConfig(id=Path(args.path).stem, path=args.path) + print(json.dumps({"path": args.path, "kind": infer_kind(spec)}, ensure_ascii=False, indent=2)) + return + if args.command == "serve": + import uvicorn + + from .server import create_app + + app = create_app(config_path) + uvicorn.run( + app, + host=args.host or config.runtime.host, + port=args.port or config.runtime.port, + log_level="info", + ) + return + if args.command == "tui": + from .tui import LocalPilotTUI + + LocalPilotTUI(config).run() + return + if args.command == "load": + from .manager import ModelManager + + async def run() -> None: + manager = ModelManager(config) + try: + backend = await manager.load(args.model_id) + print(json.dumps(backend.info(), ensure_ascii=False, indent=2)) + finally: + await manager.shutdown() + + asyncio.run(run()) + + +if __name__ == "__main__": + main() + diff --git a/localpilot/config.py b/localpilot/config.py new file mode 100644 index 0000000..8f4b30c --- /dev/null +++ b/localpilot/config.py @@ -0,0 +1,135 @@ +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field + + +class RuntimeConfig(BaseModel): + model_config = ConfigDict(extra="allow") + + llama_server: str = r"E:\ollama\模型\llama-b8642-bin-win-cuda-13.1-x64\llama-server.exe" + ollama_url: str = "http://127.0.0.1:11434" + host: str = "127.0.0.1" + port: int = 8787 + default_model: str | None = "apex-qwen3" + model_dir: str = r"E:\model" + runtime_dir: str = r"E:\llm-backend\runtime" + cache_dir: str = r"E:\llm-backend\cache" + n_gpu_layers: str | int = "auto" + ctx_size: int = 32768 + batch_size: int = 2048 + ubatch_size: int = 512 + flash_attn: str = "auto" + fit_vram: bool = True + auto_load: bool = False + api_key: str | None = None + max_loaded_models: int = 1 + keep_alive_seconds: float = 300.0 + + +class ModelConfig(BaseModel): + model_config = ConfigDict(extra="allow") + + id: str + kind: str = "auto" + path: str | None = None + enabled: bool = True + description: str = "" + cloud_profile: str | None = None + remote_model: str | None = None + template: str | None = None + system: str | None = None + messages: list[dict[str, Any]] = Field(default_factory=list) + parameters: dict[str, Any] = Field(default_factory=dict) + options: dict[str, Any] = Field(default_factory=dict) + + +class CloudProfile(BaseModel): + model_config = ConfigDict(extra="allow") + + id: str + base_url: str + api_key_env: str | None = None + api_key: str | None = None + default_model: str | None = None + timeout: float = 120.0 + headers: dict[str, str] = Field(default_factory=dict) + + def resolved_api_key(self) -> str | None: + if self.api_key: + return self.api_key + if self.api_key_env: + return os.getenv(self.api_key_env) + return None + + +class AppConfig(BaseModel): + model_config = ConfigDict(extra="allow") + + runtime: RuntimeConfig = Field(default_factory=RuntimeConfig) + models: list[ModelConfig] = Field(default_factory=list) + cloud_profiles: list[CloudProfile] = Field(default_factory=list) + + def model_by_id(self, model_id: str) -> ModelConfig: + for model in self.models: + if model.id == model_id and model.enabled: + return model + raise KeyError(f"模型未配置或已禁用: {model_id}") + + def cloud_by_id(self, profile_id: str) -> CloudProfile: + for profile in self.cloud_profiles: + if profile.id == profile_id: + return profile + raise KeyError(f"云端配置不存在: {profile_id}") + + +def infer_kind(model: ModelConfig) -> str: + if model.kind != "auto": + return { + "safetensors": "transformers", + "hf": "transformers", + "onnxruntime": "onnx", + }.get(model.kind, model.kind) + if not model.path: + return "cloud" if model.cloud_profile else "unknown" + path = Path(model.path) + suffix = path.suffix.lower() + if suffix == ".gguf": + return "gguf" + if suffix == ".onnx": + return "onnx" + if suffix in {".safetensors", ".bin", ".pt", ".pth"}: + return "transformers" + if path.is_dir(): + if (path / "config.json").exists() and ( + list(path.glob("*.safetensors")) or (path / "model.safetensors.index.json").exists() + ): + return "transformers" + if list(path.glob("*.onnx")): + return "onnx" + return "unknown" + + +def load_config(path: str | Path) -> AppConfig: + config_path = Path(path) + if not config_path.exists(): + return AppConfig() + with config_path.open("r", encoding="utf-8") as handle: + return AppConfig.model_validate(json.load(handle)) + + +def save_config(path: str | Path, config: AppConfig) -> None: + config_path = Path(path) + config_path.parent.mkdir(parents=True, exist_ok=True) + with config_path.open("w", encoding="utf-8") as handle: + json.dump(config.model_dump(exclude_none=True), handle, ensure_ascii=False, indent=2) + handle.write("\n") + + +def ensure_runtime_dirs(config: AppConfig) -> None: + Path(config.runtime.runtime_dir).mkdir(parents=True, exist_ok=True) + Path(config.runtime.cache_dir).mkdir(parents=True, exist_ok=True) diff --git a/localpilot/manager.py b/localpilot/manager.py new file mode 100644 index 0000000..25dde46 --- /dev/null +++ b/localpilot/manager.py @@ -0,0 +1,369 @@ +from __future__ import annotations + +import asyncio +import re +import time +from collections.abc import AsyncIterator +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + +from .backends import GGUFBackend, ONNXBackend, OllamaBackend, RemoteBackend, TransformersBackend +from .config import AppConfig, ModelConfig, infer_kind, save_config +from .types import BackendError, ChatChunk, ChatParams + + +@dataclass +class ModelRuntimeState: + loaded_at: float = 0.0 + last_used: float = 0.0 + keep_alive_seconds: float = 300.0 + last_load_duration_seconds: float = 0.0 + active_requests: int = 0 + total_requests: int = 0 + last_duration_seconds: float = 0.0 + idle_task: asyncio.Task[None] | None = None + + +def parse_keep_alive(value: str | int | float | None, default: float) -> float: + """Parse Ollama-style keep_alive values such as 5m, 30s, 1h, 0, and -1.""" + if value is None: + return default + if isinstance(value, (int, float)): + return float("inf") if value < 0 else max(0.0, float(value)) + raw = str(value).strip().lower() + if raw in {"", "default"}: + return default + if raw in {"-1", "forever", "inf", "infinite"}: + return float("inf") + match = re.fullmatch(r"(-?\d+(?:\.\d+)?)\s*([smhd]?)", raw) + if not match: + return default + number = float(match.group(1)) + if number < 0: + return float("inf") + multiplier = {"": 1.0, "s": 1.0, "m": 60.0, "h": 3600.0, "d": 86400.0}[match.group(2)] + return number * multiplier + + +class ModelManager: + def __init__(self, config: AppConfig, config_path: str | Path | None = None) -> None: + self.config = config + self.config_path = Path(config_path) if config_path else None + self.backends: dict[str, Any] = {} + self.states: dict[str, ModelRuntimeState] = {} + self.active_model_id: str | None = None + self._load_lock = asyncio.Lock() + + def _state(self, model_id: str) -> ModelRuntimeState: + if model_id not in self.states: + self.states[model_id] = ModelRuntimeState( + keep_alive_seconds=self.config.runtime.keep_alive_seconds + ) + return self.states[model_id] + + @staticmethod + def _model_size(path: str | None) -> int: + if not path: + return 0 + target = Path(path) + if target.is_file(): + try: + return target.stat().st_size + except OSError: + return 0 + if not target.is_dir(): + return 0 + total = 0 + try: + for item in target.rglob("*"): + if item.is_file() and item.suffix.lower() in {".gguf", ".safetensors", ".bin", ".pt", ".onnx"}: + total += item.stat().st_size + except OSError: + return total + return total + + def _model_record(self, spec: ModelConfig) -> dict[str, Any]: + backend = self.backends.get(spec.id) + state = self._state(spec.id) + size = self._model_size(spec.path) + modified = 0.0 + if spec.path: + try: + modified = Path(spec.path).stat().st_mtime + except OSError: + pass + record: dict[str, Any] = { + "id": spec.id, + "name": spec.id, + "model": spec.id, + "object": "model", + "owned_by": "localpilot" if infer_kind(spec) not in {"cloud", "ollama"} else infer_kind(spec), + "kind": infer_kind(spec), + "path": spec.path, + "size": size, + "digest": f"local-{int(modified)}-{size}" if size else None, + "modified_at": modified, + "loaded": bool(backend and backend.loaded), + "active": spec.id == self.active_model_id, + "active_requests": state.active_requests, + "total_requests": state.total_requests, + "keep_alive_seconds": state.keep_alive_seconds, + "load_duration_seconds": state.last_load_duration_seconds, + "last_request_duration_seconds": state.last_duration_seconds, + "description": spec.description, + "capabilities": ["completion"], + } + if infer_kind(spec) in {"transformers", "ollama"}: + record["capabilities"].append("embedding") + if spec.options.get("mmproj"): + record["capabilities"].append("vision") + if backend: + record.update({key: value for key, value in backend.info().items() if key not in record}) + return record + + def list_models(self) -> list[dict[str, Any]]: + return [self._model_record(spec) for spec in self.config.models if spec.enabled] + + def running_models(self) -> list[dict[str, Any]]: + return [ + self._model_record(spec) + for spec in self.config.models + if (backend := self.backends.get(spec.id)) and backend.loaded + ] + + def show_model(self, model_id: str) -> dict[str, Any]: + spec = self.config.model_by_id(model_id) + return { + "model": model_id, + "details": { + "format": infer_kind(spec), + "family": spec.options.get("family", "unknown"), + "parameter_size": spec.options.get("parameter_size"), + "quantization_level": spec.options.get("quantization_level"), + }, + "template": spec.template, + "system": spec.system, + "messages": spec.messages, + "parameters": spec.parameters, + "options": spec.options, + "runtime": self._model_record(spec), + } + + def _make_backend(self, spec: ModelConfig) -> Any: + kind = infer_kind(spec) + if kind == "gguf": + return GGUFBackend(spec, self.config.runtime) + if kind == "transformers": + return TransformersBackend(spec, self.config.runtime) + if kind == "onnx": + return ONNXBackend(spec, self.config.runtime) + if kind == "ollama": + return OllamaBackend(spec, self.config.runtime) + if kind == "cloud": + if not spec.cloud_profile: + raise BackendError(f"云端模型 {spec.id} 缺少 cloud_profile") + return RemoteBackend(spec, self.config.cloud_by_id(spec.cloud_profile), self.config.runtime) + raise BackendError(f"无法识别模型格式: {spec.id} ({spec.path})") + + async def _evict_if_needed(self, requested_id: str) -> None: + spec = self.config.model_by_id(requested_id) + if infer_kind(spec) in {"cloud", "ollama"}: + return + limit = max(1, self.config.runtime.max_loaded_models) + local_loaded = [ + (model_id, self._state(model_id)) + for model_id, backend in self.backends.items() + if backend.loaded and infer_kind(self.config.model_by_id(model_id)) not in {"cloud", "ollama"} + ] + while len(local_loaded) >= limit: + candidates = [(mid, state) for mid, state in local_loaded if mid != requested_id and state.active_requests == 0] + if not candidates: + raise BackendError("模型加载槽位都在使用中，请等待当前请求结束或提高 runtime.max_loaded_models") + oldest_id = min(candidates, key=lambda item: item[1].last_used or item[1].loaded_at)[0] + await self.unload(oldest_id) + local_loaded = [item for item in local_loaded if item[0] != oldest_id] + + async def load(self, model_id: str) -> Any: + async with self._load_lock: + if model_id in self.backends and self.backends[model_id].loaded: + self.active_model_id = model_id + self._state(model_id).last_used = time.time() + return self.backends[model_id] + spec = self.config.model_by_id(model_id) + await self._evict_if_needed(model_id) + backend = self.backends.get(model_id) or self._make_backend(spec) + started = time.perf_counter() + await backend.load() + state = self._state(model_id) + state.loaded_at = time.time() + state.last_used = state.loaded_at + state.last_load_duration_seconds = time.perf_counter() - started + self.backends[model_id] = backend + self.active_model_id = model_id + return backend + + async def unload(self, model_id: str) -> None: + backend = self.backends.get(model_id) + state = self.states.get(model_id) + if state and state.idle_task and state.idle_task is not asyncio.current_task(): + state.idle_task.cancel() + state.idle_task = None + if backend: + await backend.unload() + if self.active_model_id == model_id: + self.active_model_id = None + + def _schedule_idle_unload(self, model_id: str, seconds: float) -> None: + state = self._state(model_id) + if state.idle_task: + state.idle_task.cancel() + state.keep_alive_seconds = seconds + state.last_used = time.time() + if seconds == float("inf"): + state.idle_task = None + return + + async def expire() -> None: + try: + await asyncio.sleep(seconds) + current = self.states.get(model_id) + if not current or current.active_requests: + return + if time.time() - current.last_used >= seconds: + await self.unload(model_id) + except asyncio.CancelledError: + return + + state.idle_task = asyncio.create_task(expire()) + + def _effective_params(self, spec: ModelConfig, params: ChatParams) -> ChatParams: + values = asdict(params) + defaults = ChatParams() + for key, value in spec.parameters.items(): + if key in values and getattr(params, key) == getattr(defaults, key): + values[key] = value + return ChatParams(**values) + + def _prepare_messages( + self, + spec: ModelConfig, + messages: list[dict[str, Any]], + params: ChatParams, + ) -> list[dict[str, Any]]: + prepared = [dict(message) for message in spec.messages] + [dict(message) for message in messages] + if spec.system and not any(message.get("role") == "system" for message in prepared): + prepared.insert(0, {"role": "system", "content": spec.system}) + if params.response_format and infer_kind(spec) in {"transformers", "onnx"}: + if isinstance(params.response_format, dict): + hint = "只输出符合以下 JSON Schema 的有效 JSON，不要输出 Markdown：" + str(params.response_format) + else: + hint = "只输出有效 JSON，不要输出 Markdown 或额外解释。" + prepared.insert(0, {"role": "system", "content": hint}) + return prepared + + async def chat( + self, + model_id: str | None, + messages: list[dict[str, Any]], + params: ChatParams, + stream: bool = False, + keep_alive: str | int | float | None = None, + ) -> AsyncIterator[ChatChunk]: + chosen = model_id or self.active_model_id or self.config.runtime.default_model + if not chosen: + raise BackendError("没有可用模型，请在 config.json 设置 runtime.default_model 或请求中传 model") + spec = self.config.model_by_id(chosen) + effective_params = self._effective_params(spec, params) + effective_params.keep_alive = keep_alive + prepared_messages = self._prepare_messages(spec, messages, effective_params) + backend = await self.load(chosen) + state = self._state(chosen) + state.active_requests += 1 + state.total_requests += 1 + started = time.perf_counter() + try: + async for chunk in backend.chat(prepared_messages, effective_params, stream=stream): + yield chunk + finally: + state.active_requests = max(0, state.active_requests - 1) + state.last_duration_seconds = time.perf_counter() - started + self._schedule_idle_unload(chosen, parse_keep_alive(keep_alive, self.config.runtime.keep_alive_seconds)) + + async def embed(self, model_id: str | None, inputs: list[str], keep_alive: str | int | float | None = None) -> list[list[float]]: + chosen = model_id or self.active_model_id or self.config.runtime.default_model + if not chosen: + raise BackendError("没有可用 embedding 模型") + backend = await self.load(chosen) + state = self._state(chosen) + state.active_requests += 1 + state.total_requests += 1 + started = time.perf_counter() + try: + return await backend.embed(inputs) + finally: + state.active_requests = max(0, state.active_requests - 1) + state.last_duration_seconds = time.perf_counter() - started + self._schedule_idle_unload(chosen, parse_keep_alive(keep_alive, self.config.runtime.keep_alive_seconds)) + + def register_model(self, spec: ModelConfig) -> ModelConfig: + if any(item.id == spec.id for item in self.config.models): + raise BackendError(f"模型 ID 已存在: {spec.id}") + self.config.models.append(spec) + if self.config_path: + save_config(self.config_path, self.config) + return spec + + def copy_model(self, source_id: str, destination_id: str) -> ModelConfig: + source = self.config.model_by_id(source_id) + copied = source.model_copy(update={"id": destination_id}) + return self.register_model(copied) + + async def pull_huggingface( + self, + model_id: str, + repo_id: str, + destination: str | Path | None = None, + revision: str | None = None, + allow_patterns: list[str] | None = None, + ) -> ModelConfig: + try: + from huggingface_hub import snapshot_download + except ImportError as exc: + raise BackendError("当前 Conda LLM 环境缺少 huggingface_hub") from exc + if destination is None: + safe_repo = re.sub(r"[^A-Za-z0-9._-]+", "_", repo_id) + destination = Path(self.config.runtime.cache_dir) / "models" / safe_repo + target = Path(destination) + target.parent.mkdir(parents=True, exist_ok=True) + try: + downloaded = await asyncio.to_thread( + snapshot_download, + repo_id=repo_id, + revision=revision, + local_dir=str(target), + allow_patterns=allow_patterns, + ) + except Exception as exc: + raise BackendError(f"Hugging Face 下载失败: {exc}") from exc + spec = ModelConfig(id=model_id, kind="auto", path=downloaded, description=f"Hugging Face: {repo_id}") + return self.register_model(spec) + + async def remove_model(self, model_id: str) -> None: + await self.unload(model_id) + before = len(self.config.models) + self.config.models = [item for item in self.config.models if item.id != model_id] + if len(self.config.models) == before: + raise BackendError(f"模型不存在: {model_id}") + if self.config_path: + save_config(self.config_path, self.config) + + async def shutdown(self) -> None: + for state in self.states.values(): + if state.idle_task: + state.idle_task.cancel() + for backend in list(self.backends.values()): + if backend.loaded: + await backend.unload() + self.backends.clear() + self.active_model_id = None diff --git a/localpilot/modelfile.py b/localpilot/modelfile.py new file mode 100644 index 0000000..53a2517 --- /dev/null +++ b/localpilot/modelfile.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +import json +import re +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from .config import ModelConfig + + +@dataclass +class ModelfileSpec: + source: str = "" + parameters: dict[str, Any] = field(default_factory=dict) + system: str | None = None + template: str | None = None + messages: list[dict[str, str]] = field(default_factory=list) + adapters: list[str] = field(default_factory=list) + licenses: list[str] = field(default_factory=list) + + +def _value(raw: str) -> Any: + try: + return json.loads(raw) + except json.JSONDecodeError: + return raw + + +def parse_modelfile(text: str) -> ModelfileSpec: + spec = ModelfileSpec() + block_pattern = re.compile(r"(?ms)^\s*(SYSTEM|TEMPLATE|LICENSE)\s+\"\"\"(.*?)\"\"\"\s*$") + + def take_block(match: re.Match[str]) -> str: + key, value = match.group(1), match.group(2).strip("\r\n") + if key == "SYSTEM": + spec.system = value + elif key == "TEMPLATE": + spec.template = value + else: + spec.licenses.append(value) + return "" + + remaining = block_pattern.sub(take_block, text) + for line in remaining.splitlines(): + line = line.strip() + if not line or line.startswith("#"): + continue + parts = line.split(maxsplit=2) + instruction = parts[0].upper() + if instruction == "FROM" and len(parts) >= 2: + source = line[len(parts[0]):].strip() + spec.source = source.strip('"\'') + elif instruction == "PARAMETER" and len(parts) >= 3: + key, raw = parts[1], parts[2] + value = _value(raw) + if key == "stop": + spec.parameters.setdefault("stop", []).append(str(value)) + else: + spec.parameters[key] = value + elif instruction == "MESSAGE" and len(parts) >= 3: + spec.messages.append({"role": parts[1].lower(), "content": parts[2]}) + elif instruction == "ADAPTER" and len(parts) >= 2: + spec.adapters.append(parts[1]) + return spec + + +def model_config_from_modelfile(model_id: str, text: str) -> ModelConfig: + parsed = parse_modelfile(text) + if not parsed.source: + raise ValueError("Modelfile 缺少 FROM") + source_path = Path(parsed.source) + if source_path.exists() or source_path.suffix.lower() in {".gguf", ".onnx", ".safetensors", ".bin", ".pt", ".pth"}: + path: str | None = parsed.source + kind = "auto" + remote_model = None + else: + path = None + kind = "ollama" + remote_model = parsed.source + parameters: dict[str, Any] = {} + options: dict[str, Any] = {} + aliases = { + "num_predict": "max_tokens", + "num_ctx": "ctx_size", + "num_batch": "batch_size", + "num_gpu": "n_gpu_layers", + } + for key, value in parsed.parameters.items(): + target = aliases.get(key, key) + if target in {"ctx_size", "batch_size", "n_gpu_layers"}: + options[target] = value + else: + parameters[target] = value + if parsed.adapters: + options["adapters"] = parsed.adapters + return ModelConfig( + id=model_id, + kind=kind, + path=path, + remote_model=remote_model, + template=parsed.template, + system=parsed.system, + messages=parsed.messages, + parameters=parameters, + options=options, + ) diff --git a/localpilot/server.py b/localpilot/server.py new file mode 100644 index 0000000..65944d7 --- /dev/null +++ b/localpilot/server.py @@ -0,0 +1,634 @@ +from __future__ import annotations + +import json +import time +import uuid +from contextlib import asynccontextmanager +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Literal + +from fastapi import Depends, FastAPI, Header, HTTPException, Request +from fastapi.responses import PlainTextResponse, StreamingResponse +from pydantic import BaseModel, ConfigDict, Field + +from .config import AppConfig, ModelConfig, load_config +from .manager import ModelManager +from .modelfile import model_config_from_modelfile +from .types import BackendError, ChatParams + + +class Message(BaseModel): + model_config = ConfigDict(extra="allow") + + role: Literal["system", "user", "assistant", "tool", "function", "developer"] + content: Any + + def payload(self) -> dict[str, Any]: + return self.model_dump(exclude_none=True) + + +class ChatCompletionRequest(BaseModel): + model_config = ConfigDict(extra="allow") + + model: str | None = None + messages: list[Message] + stream: bool = False + max_tokens: int | None = Field(default=None, ge=1, le=131072) + max_completion_tokens: int | None = Field(default=None, ge=1, le=131072) + temperature: float | None = Field(default=None, ge=0, le=2) + top_p: float | None = Field(default=None, ge=0, le=1) + top_k: int | None = Field(default=None, ge=0, le=100000) + min_p: float | None = Field(default=None, ge=0, le=1) + repeat_penalty: float | None = Field(default=None, ge=0) + seed: int | None = None + stop: str | list[str] | None = None + enable_thinking: bool | None = None + response_format: str | dict[str, Any] | None = None + tools: list[dict[str, Any]] = Field(default_factory=list) + tool_choice: Any = None + keep_alive: str | int | float | None = None + extra_body: dict[str, Any] = Field(default_factory=dict) + + def params(self, config: AppConfig) -> ChatParams: + del config + default = ChatParams() + stop = [self.stop] if isinstance(self.stop, str) else (self.stop or []) + extra = dict(self.extra_body) + known = { + "model", "messages", "stream", "max_tokens", "max_completion_tokens", "temperature", + "top_p", "top_k", "min_p", "repeat_penalty", "seed", "stop", "enable_thinking", + "response_format", "tools", "tool_choice", "keep_alive", "extra_body", + } + extra.update({key: value for key, value in (self.model_extra or {}).items() if key not in known}) + return ChatParams( + max_tokens=self.max_completion_tokens or self.max_tokens or default.max_tokens, + temperature=default.temperature if self.temperature is None else self.temperature, + top_p=default.top_p if self.top_p is None else self.top_p, + top_k=default.top_k if self.top_k is None else self.top_k, + min_p=default.min_p if self.min_p is None else self.min_p, + repeat_penalty=default.repeat_penalty if self.repeat_penalty is None else self.repeat_penalty, + seed=self.seed, + stop=stop, + enable_thinking=self.enable_thinking, + response_format=self.response_format, + tools=self.tools, + tool_choice=self.tool_choice, + extra=extra, + ) + + +class OllamaChatRequest(BaseModel): + model_config = ConfigDict(extra="allow") + + model: str + messages: list[Message] + stream: bool = True + format: str | dict[str, Any] | None = None + options: dict[str, Any] = Field(default_factory=dict) + keep_alive: str | int | float | None = "5m" + think: bool | None = None + tools: list[dict[str, Any]] = Field(default_factory=list) + + +class OllamaGenerateRequest(BaseModel): + model_config = ConfigDict(extra="allow") + + model: str + prompt: str = "" + suffix: str | None = None + system: str | None = None + template: str | None = None + context: list[int] | None = None + stream: bool = True + raw: bool = False + format: str | dict[str, Any] | None = None + options: dict[str, Any] = Field(default_factory=dict) + keep_alive: str | int | float | None = "5m" + think: bool | None = None + + +class OllamaEmbedRequest(BaseModel): + model: str + input: str | list[str] + keep_alive: str | int | float | None = "5m" + + +class OpenAIEmbeddingRequest(BaseModel): + model: str + input: str | list[str] + encoding_format: str | None = None + + +class RegisterModelRequest(BaseModel): + id: str + path: str | None = None + kind: str = "auto" + enabled: bool = True + description: str = "" + cloud_profile: str | None = None + remote_model: str | None = None + template: str | None = None + system: str | None = None + messages: list[dict[str, Any]] = Field(default_factory=list) + parameters: dict[str, Any] = Field(default_factory=dict) + options: dict[str, Any] = Field(default_factory=dict) + + +class CreateModelRequest(BaseModel): + model: str + modelfile: str + stream: bool = False + + +class CopyModelRequest(BaseModel): + source: str + destination: str + + +class PullModelRequest(BaseModel): + model: str + path: str | None = None + kind: str = "auto" + stream: bool = True + repo_id: str | None = None + revision: str | None = None + destination: str | None = None + allow_patterns: list[str] | None = None + alias: str | None = None + + +def _content_to_text(content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + return "".join( + str(item.get("text", "")) if isinstance(item, dict) and item.get("type") == "text" else "[image]" + for item in content + ) + return str(content) + + +def _sse(data: dict[str, Any]) -> str: + return f"data: {json.dumps(data, ensure_ascii=False)}\n\n" + + +def _iso_now() -> str: + return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") + + +def _ollama_params(options: dict[str, Any], fmt: str | dict[str, Any] | None, think: bool | None, tools: list[dict[str, Any]]) -> ChatParams: + stop = options.get("stop", []) + if isinstance(stop, str): + stop = [stop] + standard = { + "num_predict": "max_tokens", + "temperature": "temperature", + "top_p": "top_p", + "top_k": "top_k", + "min_p": "min_p", + "repeat_penalty": "repeat_penalty", + "seed": "seed", + } + values: dict[str, Any] = {key: options[source] for source, key in standard.items() if source in options} + extra = {key: value for key, value in options.items() if key not in standard and key not in {"stop"}} + return ChatParams( + max_tokens=int(values.get("max_tokens", 512)), + temperature=float(values.get("temperature", 0.7)), + top_p=float(values.get("top_p", 0.95)), + top_k=int(values.get("top_k", 40)), + min_p=float(values.get("min_p", 0.05)), + repeat_penalty=float(values.get("repeat_penalty", 1.05)), + seed=values.get("seed"), + stop=stop, + enable_thinking=think, + response_format=fmt, + tools=tools, + extra=extra, + ) + + +def _usage_from_metadata(metadata: dict[str, Any]) -> dict[str, int]: + usage = metadata.get("usage") or {} + if usage: + prompt = int(usage.get("prompt_tokens", usage.get("prompt_eval_count", 0)) or 0) + completion = int(usage.get("completion_tokens", usage.get("eval_count", 0)) or 0) + return { + "prompt_tokens": prompt, + "completion_tokens": completion, + "total_tokens": int(usage.get("total_tokens", prompt + completion)), + } + prompt = int(metadata.get("prompt_eval_count", metadata.get("prompt_tokens", 0)) or 0) + completion = int(metadata.get("eval_count", metadata.get("completion_tokens", 0)) or 0) + return {"prompt_tokens": prompt, "completion_tokens": completion, "total_tokens": prompt + completion} + + +def _ollama_usage(metadata: dict[str, Any]) -> dict[str, int]: + usage = metadata.get("usage") or {} + prompt = int(usage.get("prompt_eval_count", usage.get("prompt_tokens", metadata.get("prompt_eval_count", 0))) or 0) + completion = int(usage.get("eval_count", usage.get("completion_tokens", metadata.get("eval_count", 0))) or 0) + return {"prompt_eval_count": prompt, "eval_count": completion, "eval_duration": 0, "prompt_eval_duration": 0} + + +def _ollama_runtime_fields(manager: ModelManager, model_id: str) -> dict[str, int]: + state = manager.states.get(model_id) + if not state: + return {} + return { + "total_duration": int(state.last_duration_seconds * 1_000_000_000), + "load_duration": int(state.last_load_duration_seconds * 1_000_000_000), + } + + +def create_app(config_path: str | Path | None = None) -> FastAPI: + resolved = Path(config_path or r"E:\llm-backend\config\config.json") + config = load_config(resolved) + manager = ModelManager(config, resolved) + + @asynccontextmanager + async def lifespan(_: FastAPI): + if config.runtime.auto_load and config.runtime.default_model: + try: + await manager.load(config.runtime.default_model) + except Exception: + pass + yield + await manager.shutdown() + + app = FastAPI(title="LocalPilot", version="0.2.0", lifespan=lifespan) + app.state.config = config + app.state.manager = manager + + async def authorize(request: Request, authorization: str | None = Header(default=None)) -> None: + del request + expected = config.runtime.api_key + if expected and (authorization or "").removeprefix("Bearer ") != expected: + raise HTTPException(status_code=401, detail="invalid API key") + + @app.get("/health") + async def health() -> dict[str, Any]: + return { + "status": "ok", + "version": "0.2.0", + "active_model": manager.active_model_id, + "running": manager.running_models(), + } + + @app.get("/v1/models") + async def models(_: None = Depends(authorize)) -> dict[str, Any]: + data = [] + for item in manager.list_models(): + data.append({**item, "created": int(item.get("modified_at") or time.time())}) + return {"object": "list", "data": data} + + @app.get("/api/version") + async def api_version() -> dict[str, str]: + return {"version": "0.2.0"} + + @app.get("/api/tags") + async def tags(_: None = Depends(authorize)) -> dict[str, Any]: + models_data = [] + for item in manager.list_models(): + details = manager.show_model(item["id"])["details"] + models_data.append({ + "name": item["name"], + "model": item["model"], + "modified_at": datetime.fromtimestamp(item["modified_at"], timezone.utc).isoformat() if item["modified_at"] else _iso_now(), + "size": item["size"], + "digest": item["digest"], + "details": details, + "capabilities": item["capabilities"], + }) + return {"models": models_data} + + @app.get("/api/ps") + async def ps(_: None = Depends(authorize)) -> dict[str, Any]: + models_data = [] + for item in manager.running_models(): + state = manager.states[item["id"]] + expires_at = None + if state.keep_alive_seconds != float("inf"): + expires_at = datetime.fromtimestamp( + state.last_used + state.keep_alive_seconds, timezone.utc + ).isoformat() + models_data.append({ + "name": item["name"], + "model": item["model"], + "size": item["size"], + "size_vram": item.get("vram_allocated_bytes", 0), + "expires_at": expires_at, + "active_requests": item["active_requests"], + "kind": item["kind"], + }) + return {"models": models_data} + + @app.get("/metrics", response_class=PlainTextResponse) + async def metrics(_: None = Depends(authorize)) -> str: + records = manager.list_models() + loaded = sum(1 for item in records if item["loaded"]) + active = sum(int(item["active_requests"]) for item in records) + total = sum(int(item["total_requests"]) for item in records) + lines = [ + "# TYPE localpilot_models_loaded gauge", + f"localpilot_models_loaded {loaded}", + "# TYPE localpilot_active_requests gauge", + f"localpilot_active_requests {active}", + "# TYPE localpilot_requests_total counter", + f"localpilot_requests_total {total}", + ] + return "\n".join(lines) + "\n" + + @app.post("/api/show") + async def show(body: dict[str, Any], _: None = Depends(authorize)) -> dict[str, Any]: + model_id = body.get("model") or body.get("name") + if not model_id: + raise HTTPException(status_code=400, detail="missing model/name") + try: + return manager.show_model(model_id) + except KeyError as exc: + raise HTTPException(status_code=404, detail=str(exc)) from exc + + @app.post("/api/models/register") + async def register(body: RegisterModelRequest, _: None = Depends(authorize)) -> dict[str, Any]: + try: + spec = manager.register_model(ModelConfig(**body.model_dump())) + return {"status": "registered", "model": spec.model_dump(exclude_none=True)} + except BackendError as exc: + raise HTTPException(status_code=409, detail=str(exc)) from exc + + @app.delete("/api/models/{model_id}") + async def delete_model(model_id: str, _: None = Depends(authorize)) -> dict[str, Any]: + try: + await manager.remove_model(model_id) + return {"status": "deleted", "model": model_id, "files_removed": False} + except BackendError as exc: + raise HTTPException(status_code=404, detail=str(exc)) from exc + + @app.post("/api/create") + async def create_model(body: CreateModelRequest, _: None = Depends(authorize)) -> Any: + try: + modelfile_text = body.modelfile + modelfile_path = Path(modelfile_text) + if modelfile_path.is_file(): + modelfile_text = modelfile_path.read_text(encoding="utf-8") + spec = model_config_from_modelfile(body.model, modelfile_text) + manager.register_model(spec) + except (BackendError, ValueError) as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + if body.stream: + async def create_stream(): + yield json.dumps({"status": "success", "model": body.model}, ensure_ascii=False) + "\n" + + return StreamingResponse(create_stream(), media_type="application/x-ndjson") + return {"status": "success", "model": body.model} + + @app.post("/api/copy") + async def copy_model(body: CopyModelRequest, _: None = Depends(authorize)) -> dict[str, Any]: + try: + spec = manager.copy_model(body.source, body.destination) + return {"status": "success", "model": spec.id} + except (BackendError, KeyError) as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + + @app.post("/api/pull") + async def pull_model(body: PullModelRequest, _: None = Depends(authorize)) -> Any: + try: + if body.path: + spec = manager.register_model(ModelConfig(id=body.model, kind=body.kind, path=body.path)) + elif body.repo_id or "/" in body.model: + repo_id = body.repo_id or body.model + model_id = body.alias or repo_id.replace("/", "-") + spec = await manager.pull_huggingface( + model_id=model_id, + repo_id=repo_id, + destination=body.destination, + revision=body.revision, + allow_patterns=body.allow_patterns, + ) + else: + raise HTTPException( + status_code=501, + detail="请提供 path 导入本地模型，或使用 repo_id/repo 触发 Hugging Face 下载。", + ) + except BackendError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + if body.stream: + async def pull_stream(): + yield json.dumps({"status": "success", "model": spec.id, "path": spec.path}, ensure_ascii=False) + "\n" + + return StreamingResponse(pull_stream(), media_type="application/x-ndjson") + return {"status": "success", "model": spec.id, "path": spec.path} + + @app.post("/api/models/{model_id}/load") + async def load_model(model_id: str, _: None = Depends(authorize)) -> dict[str, Any]: + try: + backend = await manager.load(model_id) + return backend.info() + except (BackendError, KeyError) as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + + @app.post("/api/models/{model_id}/unload") + async def unload_model(model_id: str, _: None = Depends(authorize)) -> dict[str, Any]: + await manager.unload(model_id) + return {"status": "unloaded", "model": model_id} + + @app.post("/v1/chat/completions") + async def chat_completions(body: ChatCompletionRequest, _: None = Depends(authorize)) -> Any: + messages = [item.payload() for item in body.messages] + params = body.params(config) + request_id = f"chatcmpl-localpilot-{uuid.uuid4().hex[:12]}" + model_id = body.model or config.runtime.default_model or manager.active_model_id or "localpilot" + try: + if body.stream: + async def event_stream(): + yield _sse({ + "id": request_id, + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": model_id, + "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}], + }) + async for chunk in manager.chat(model_id, messages, params, stream=True, keep_alive=body.keep_alive): + delta = dict(chunk.delta) + if chunk.text and "content" not in delta: + delta["content"] = chunk.text + if delta: + yield _sse({ + "id": request_id, + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": model_id, + "choices": [{"index": 0, "delta": delta, "finish_reason": None}], + }) + if chunk.done: + yield _sse({ + "id": request_id, + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": model_id, + "choices": [{"index": 0, "delta": {}, "finish_reason": chunk.finish_reason or "stop"}], + }) + yield "data: [DONE]\n\n" + + return StreamingResponse(event_stream(), media_type="text/event-stream") + + parts: list[str] = [] + finish_reason = "stop" + metadata: dict[str, Any] = {} + async for chunk in manager.chat(model_id, messages, params, stream=False, keep_alive=body.keep_alive): + parts.append(chunk.text) + finish_reason = chunk.finish_reason or finish_reason + metadata = chunk.metadata or metadata + message: dict[str, Any] = {"role": "assistant", "content": "".join(parts)} + if metadata: + original_message = ((metadata.get("choices") or [{}])[0].get("message") or {}) + for key in ("tool_calls", "function_call", "refusal"): + if key in original_message: + message[key] = original_message[key] + return { + "id": request_id, + "object": "chat.completion", + "created": int(time.time()), + "model": model_id, + "choices": [{ + "index": 0, + "message": message, + "finish_reason": finish_reason, + }], + "usage": _usage_from_metadata(metadata), + } + except (BackendError, KeyError) as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + + @app.post("/v1/completions") + async def completions(body: dict[str, Any], _: None = Depends(authorize)) -> Any: + prompt = body.get("prompt", "") + if isinstance(prompt, list): + prompt = "\n".join(str(item) for item in prompt) + payload = ChatCompletionRequest( + model=body.get("model"), + messages=[Message(role="user", content=str(prompt))], + stream=bool(body.get("stream", False)), + max_tokens=body.get("max_tokens"), + temperature=body.get("temperature"), + top_p=body.get("top_p"), + stop=body.get("stop"), + keep_alive=body.get("keep_alive"), + ) + return await chat_completions(payload) + + @app.post("/v1/embeddings") + async def openai_embeddings(body: OpenAIEmbeddingRequest, _: None = Depends(authorize)) -> dict[str, Any]: + inputs = [body.input] if isinstance(body.input, str) else body.input + try: + vectors = await manager.embed(body.model, inputs) + except (BackendError, KeyError) as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + return { + "object": "list", + "data": [{"object": "embedding", "index": index, "embedding": vector} for index, vector in enumerate(vectors)], + "model": body.model, + "usage": {"prompt_tokens": 0, "total_tokens": 0}, + } + + @app.post("/api/embed") + async def ollama_embed(body: OllamaEmbedRequest, _: None = Depends(authorize)) -> dict[str, Any]: + inputs = [body.input] if isinstance(body.input, str) else body.input + try: + vectors = await manager.embed(body.model, inputs, body.keep_alive) + except (BackendError, KeyError) as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + return {"model": body.model, "embeddings": vectors} + + @app.post("/api/embeddings") + async def ollama_embeddings_legacy(body: OllamaEmbedRequest, _: None = Depends(authorize)) -> dict[str, Any]: + return await ollama_embed(body) + + @app.post("/api/chat") + async def ollama_chat(body: OllamaChatRequest, _: None = Depends(authorize)) -> Any: + messages = [item.payload() for item in body.messages] + params = _ollama_params(body.options, body.format, body.think, body.tools) + try: + if body.stream: + async def event_stream(): + answer: list[str] = [] + async for chunk in manager.chat(body.model, messages, params, stream=True, keep_alive=body.keep_alive): + answer.append(chunk.text) + message = {"role": "assistant", "content": chunk.text, **chunk.delta} + yield json.dumps({ + "model": body.model, + "created_at": _iso_now(), + "message": message, + "done": False, + }, ensure_ascii=False) + "\n" + yield json.dumps({ + "model": body.model, + "created_at": _iso_now(), + "message": {"role": "assistant", "content": ""}, + "done": True, + "done_reason": "stop", + **_ollama_runtime_fields(manager, body.model), + }, ensure_ascii=False) + "\n" + + return StreamingResponse(event_stream(), media_type="application/x-ndjson") + + parts: list[str] = [] + metadata: dict[str, Any] = {} + async for chunk in manager.chat(body.model, messages, params, stream=False, keep_alive=body.keep_alive): + parts.append(chunk.text) + metadata = chunk.metadata or metadata + return { + "model": body.model, + "created_at": _iso_now(), + "message": {"role": "assistant", "content": "".join(parts)}, + "done": True, + "done_reason": "stop", + **_ollama_runtime_fields(manager, body.model), + **_ollama_usage(metadata), + } + except (BackendError, KeyError) as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + + @app.post("/api/generate") + async def ollama_generate(body: OllamaGenerateRequest, _: None = Depends(authorize)) -> Any: + del body.suffix, body.template, body.context, body.raw + messages: list[dict[str, Any]] = [] + if body.system: + messages.append({"role": "system", "content": body.system}) + messages.append({"role": "user", "content": body.prompt}) + params = _ollama_params(body.options, body.format, body.think, []) + try: + if body.stream: + async def event_stream(): + async for chunk in manager.chat(body.model, messages, params, stream=True, keep_alive=body.keep_alive): + yield json.dumps({ + "model": body.model, + "created_at": _iso_now(), + "response": chunk.text, + "done": False, + }, ensure_ascii=False) + "\n" + yield json.dumps({ + "model": body.model, + "created_at": _iso_now(), + "response": "", + "done": True, + "done_reason": "stop", + **_ollama_runtime_fields(manager, body.model), + }, ensure_ascii=False) + "\n" + + return StreamingResponse(event_stream(), media_type="application/x-ndjson") + + parts: list[str] = [] + async for chunk in manager.chat(body.model, messages, params, stream=False, keep_alive=body.keep_alive): + parts.append(chunk.text) + return { + "model": body.model, + "created_at": _iso_now(), + "response": "".join(parts), + "done": True, + **_ollama_runtime_fields(manager, body.model), + } + except (BackendError, KeyError) as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + + return app diff --git a/localpilot/tui.py b/localpilot/tui.py new file mode 100644 index 0000000..f237c22 --- /dev/null +++ b/localpilot/tui.py @@ -0,0 +1,169 @@ +from __future__ import annotations + +import asyncio +from dataclasses import asdict +import shlex +from typing import Any + +from rich.markup import escape +from textual.app import App, ComposeResult +from textual.containers import Horizontal, Vertical +from textual.widgets import Footer, Header, Input, RichLog, Static + +from .config import AppConfig +from .manager import ModelManager +from .types import BackendError, ChatParams + + +class LocalPilotTUI(App[None]): + CSS = """ + Screen { background: $surface; } + #body { height: 1fr; } + #sidebar { width: 30; border: round $panel; padding: 1; } + #main { width: 1fr; padding: 0 1; } + #chat { height: 1fr; border: round $panel; padding: 1; } + #status { height: 1; color: $text-muted; padding: 0 1; } + #prompt { height: 3; border: round $accent; } + """ + BINDINGS = [("ctrl+c", "quit", "退出")] + + def __init__(self, config: AppConfig) -> None: + super().__init__() + self.config = config + self.manager = ModelManager(config) + self.messages: list[dict[str, str]] = [] + self.params = ChatParams() + self.current_model = config.runtime.default_model + + def compose(self) -> ComposeResult: + yield Header(show_clock=False) + with Horizontal(id="body"): + with Vertical(id="sidebar"): + yield Static("LocalPilot\n\n模型", id="model_title") + yield Static(id="models") + yield Static(id="commands") + with Vertical(id="main"): + yield RichLog(id="chat", markup=True, wrap=True, highlight=False) + yield Static("就绪", id="status") + yield Input(placeholder="输入消息，/help 查看命令", id="prompt") + yield Footer() + + async def on_mount(self) -> None: + self.query_one("#commands", Static).update( + "/model ID\n/load ID\n/unload\n/params k=v\n/clear\n/help" + ) + self._refresh_model_panel() + self.query_one("#prompt", Input).focus() + + def _refresh_model_panel(self) -> None: + rows = [] + for item in self.manager.list_models(): + mark = "*" if item["id"] == self.current_model else " " + state = "已加载" if item.get("loaded") else "未加载" + rows.append(f"{mark} {item['id']}\n {item['kind']} | {state}") + self.query_one("#models", Static).update("\n".join(rows) or "没有启用模型") + + def _set_status(self, text: str) -> None: + self.query_one("#status", Static).update(text) + + async def on_input_submitted(self, event: Input.Submitted) -> None: + value = event.value.strip() + event.input.value = "" + if not value: + return + if value.startswith("/"): + self.run_worker(self._command(value), exclusive=True) + else: + self.run_worker(self._chat(value), exclusive=True) + + async def _command(self, raw: str) -> None: + try: + parts = shlex.split(raw) + except ValueError as exc: + self._set_status(f"命令解析失败: {exc}") + return + command = parts[0].lower() + if command == "/help": + self.query_one("#chat", RichLog).write( + "[bold]命令[/bold]\n/model ID 选择模型\n/load ID 加载模型\n/unload 卸载当前模型\n" + "/params temperature=0.7 top_p=0.95 max_tokens=512\n/clear 清空对话" + ) + elif command == "/models": + self._refresh_model_panel() + elif command == "/model" and len(parts) > 1: + self.current_model = parts[1] + self._refresh_model_panel() + self._set_status(f"当前模型: {self.current_model}") + elif command == "/load" and len(parts) > 1: + await self._load_model(parts[1]) + elif command == "/unload": + if self.current_model: + await self.manager.unload(self.current_model) + self._refresh_model_panel() + self._set_status("模型已卸载") + elif command == "/params": + self._update_params(parts[1:]) + elif command == "/clear": + self.messages.clear() + self.query_one("#chat", RichLog).clear() + self._set_status("对话已清空") + else: + self._set_status("未知命令，输入 /help") + + async def _load_model(self, model_id: str) -> None: + self._set_status(f"加载 {model_id} ...") + try: + await self.manager.load(model_id) + self.current_model = model_id + self._refresh_model_panel() + self._set_status(f"已加载: {model_id}") + except (BackendError, KeyError) as exc: + self._set_status(str(exc)) + + def _update_params(self, assignments: list[str]) -> None: + updates: dict[str, Any] = {} + for assignment in assignments: + if "=" not in assignment: + continue + key, raw_value = assignment.split("=", 1) + if not hasattr(self.params, key): + continue + old = getattr(self.params, key) + try: + if isinstance(old, bool) or key == "enable_thinking": + updates[key] = raw_value.lower() in {"1", "true", "yes", "on"} + elif isinstance(old, int): + updates[key] = int(raw_value) + elif isinstance(old, float): + updates[key] = float(raw_value) + else: + updates[key] = raw_value.split(",") if key == "stop" else raw_value + except ValueError: + self._set_status(f"参数无效: {assignment}") + return + self.params = ChatParams(**{**asdict(self.params), **updates}) + self._set_status(f"参数已更新: {', '.join(f'{k}={v}' for k, v in updates.items()) or '无变化'}") + + async def _chat(self, prompt: str) -> None: + log = self.query_one("#chat", RichLog) + log.write(f"[bold cyan]你[/bold cyan] {escape(prompt)}") + self.messages.append({"role": "user", "content": prompt}) + self._set_status(f"生成中 | {self.current_model or 'default'}") + pieces: list[str] = [] + try: + async for chunk in self.manager.chat(self.current_model, self.messages, self.params, stream=True): + if chunk.text: + pieces.append(chunk.text) + self._set_status(f"生成中 | {len(''.join(pieces))} 字符") + answer = "".join(pieces) + self.messages.append({"role": "assistant", "content": answer}) + log.write(f"[bold green]模型[/bold green] {escape(answer)}") + self._refresh_model_panel() + self._set_status("就绪") + except (BackendError, KeyError) as exc: + self._set_status(str(exc)) + log.write(f"[bold red]错误[/bold red] {escape(str(exc))}") + + async def action_quit(self) -> None: + await self.manager.shutdown() + self.exit() diff --git a/localpilot/types.py b/localpilot/types.py new file mode 100644 index 0000000..77e0dca --- /dev/null +++ b/localpilot/types.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + + +@dataclass(slots=True) +class ChatParams: + max_tokens: int = 512 + temperature: float = 0.7 + top_p: float = 0.95 + top_k: int = 40 + min_p: float = 0.05 + repeat_penalty: float = 1.05 + seed: int | None = None + stop: list[str] = field(default_factory=list) + enable_thinking: bool | None = None + extra: dict[str, Any] = field(default_factory=dict) + response_format: str | dict[str, Any] | None = None + tools: list[dict[str, Any]] = field(default_factory=list) + tool_choice: Any = None + keep_alive: str | int | float | None = None + + def as_generation_kwargs(self) -> dict[str, Any]: + result: dict[str, Any] = { + "max_tokens": self.max_tokens, + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": self.top_k, + "min_p": self.min_p, + "repeat_penalty": self.repeat_penalty, + } + if self.seed is not None: + result["seed"] = self.seed + if self.stop: + result["stop"] = self.stop + return result + + +@dataclass(slots=True) +class ChatChunk: + text: str = "" + done: bool = False + finish_reason: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + delta: dict[str, Any] = field(default_factory=dict) + + +class BackendError(RuntimeError): + """An expected model backend failure with a user-facing message.""" diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..e3de5bd --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,27 @@ +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.build_meta" + +[project] +name = "localpilot" +version = "0.1.0" +description = "A fast local model runtime and OpenAI-compatible gateway" +requires-python = ">=3.11" +dependencies = [ + "fastapi>=0.110", + "httpx>=0.27", + "pydantic>=2.0", + "textual>=0.50", + "uvicorn[standard]>=0.29", +] + +[project.optional-dependencies] +transformers = ["accelerate>=0.30", "safetensors>=0.4", "torch>=2.2", "transformers>=4.45"] +onnx = ["onnxruntime>=1.18", "optimum[onnxruntime]>=1.20"] + +[project.scripts] +localpilot = "localpilot.cli:main" + +[tool.setuptools.packages.find] +where = ["."] + diff --git a/scripts/start-server.ps1 b/scripts/start-server.ps1 new file mode 100644 index 0000000..db9352c --- /dev/null +++ b/scripts/start-server.ps1 @@ -0,0 +1,9 @@ +$ErrorActionPreference = "Stop" +$projectDir = "E:\llm-backend" +$env:HF_HOME = "$projectDir\cache\huggingface" +$env:TRANSFORMERS_CACHE = "$projectDir\cache\transformers" +$env:TORCH_HOME = "$projectDir\cache\torch" +New-Item -ItemType Directory -Force -Path $env:HF_HOME,$env:TRANSFORMERS_CACHE,$env:TORCH_HOME | Out-Null +Set-Location $projectDir +conda run -n LLM python -m localpilot --config "$projectDir\config\config.json" serve @args + diff --git a/scripts/start-tui.ps1 b/scripts/start-tui.ps1 new file mode 100644 index 0000000..a3835f2 --- /dev/null +++ b/scripts/start-tui.ps1 @@ -0,0 +1,9 @@ +$ErrorActionPreference = "Stop" +$projectDir = "E:\llm-backend" +$env:HF_HOME = "$projectDir\cache\huggingface" +$env:TRANSFORMERS_CACHE = "$projectDir\cache\transformers" +$env:TORCH_HOME = "$projectDir\cache\torch" +New-Item -ItemType Directory -Force -Path $env:HF_HOME,$env:TRANSFORMERS_CACHE,$env:TORCH_HOME | Out-Null +Set-Location $projectDir +conda run -n LLM python -m localpilot --config "$projectDir\config\config.json" tui @args + diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..a39a024 --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,26 @@ +import json +import tempfile +import unittest +from pathlib import Path + +from localpilot.config import AppConfig, ModelConfig, infer_kind, load_config + + +class ConfigTests(unittest.TestCase): + def test_infer_extensions(self) -> None: + self.assertEqual(infer_kind(ModelConfig(id="x", path="E:/x/model.gguf")), "gguf") + self.assertEqual(infer_kind(ModelConfig(id="x", path="E:/x/model.onnx")), "onnx") + self.assertEqual(infer_kind(ModelConfig(id="x", kind="cloud")), "cloud") + + def test_load_config(self) -> None: + with tempfile.TemporaryDirectory() as folder: + path = Path(folder) / "config.json" + path.write_text(json.dumps({"models": [{"id": "x", "kind": "gguf", "path": "x.gguf"}]}), encoding="utf-8") + config = load_config(path) + self.assertIsInstance(config, AppConfig) + self.assertEqual(config.models[0].id, "x") + + +if __name__ == "__main__": + unittest.main() + diff --git a/tests/test_manager.py b/tests/test_manager.py new file mode 100644 index 0000000..e5b8e49 --- /dev/null +++ b/tests/test_manager.py @@ -0,0 +1,28 @@ +import tempfile +import unittest +from pathlib import Path + +from localpilot.config import AppConfig, ModelConfig, RuntimeConfig +from localpilot.manager import ModelManager, parse_keep_alive + + +class ManagerTests(unittest.TestCase): + def test_parse_keep_alive(self) -> None: + self.assertEqual(parse_keep_alive("5m", 30), 300) + self.assertEqual(parse_keep_alive("30s", 30), 30) + self.assertEqual(parse_keep_alive(0, 30), 0) + self.assertEqual(parse_keep_alive(-1, 30), float("inf")) + self.assertEqual(parse_keep_alive("bad", 30), 30) + + def test_register_and_remove_only_changes_registry(self) -> None: + with tempfile.TemporaryDirectory() as folder: + path = Path(folder) / "config.json" + manager = ModelManager(AppConfig(runtime=RuntimeConfig()), path) + manager.register_model(ModelConfig(id="demo", kind="gguf", path="E:/model/demo.gguf")) + self.assertEqual(manager.config.models[0].id, "demo") + self.assertTrue(path.exists()) + + +if __name__ == "__main__": + unittest.main() + diff --git a/tests/test_modelfile.py b/tests/test_modelfile.py new file mode 100644 index 0000000..b4b5c5f --- /dev/null +++ b/tests/test_modelfile.py @@ -0,0 +1,27 @@ +import unittest + +from localpilot.modelfile import model_config_from_modelfile, parse_modelfile + + +class ModelfileTests(unittest.TestCase): + def test_parse_and_convert(self) -> None: + text = ''' +FROM E:/models/demo.gguf +PARAMETER temperature 0.2 +PARAMETER top_p 0.9 +PARAMETER num_ctx 8192 +SYSTEM """You are concise.""" +MESSAGE user hello +''' + parsed = parse_modelfile(text) + self.assertEqual(parsed.source, "E:/models/demo.gguf") + self.assertEqual(parsed.system, "You are concise.") + model = model_config_from_modelfile("demo", text) + self.assertEqual(model.parameters["temperature"], 0.2) + self.assertEqual(model.options["ctx_size"], 8192) + self.assertEqual(model.messages[0]["role"], "user") + + +if __name__ == "__main__": + unittest.main() + diff --git a/tests/test_server_helpers.py b/tests/test_server_helpers.py new file mode 100644 index 0000000..ebc001f --- /dev/null +++ b/tests/test_server_helpers.py @@ -0,0 +1,41 @@ +import unittest +import json +import tempfile +from pathlib import Path + +from fastapi.testclient import TestClient + +from localpilot.server import create_app +from localpilot.server import _content_to_text, _sse + + +class ServerHelperTests(unittest.TestCase): + def test_content_to_text(self) -> None: + self.assertEqual(_content_to_text("hello"), "hello") + self.assertEqual(_content_to_text([{"type": "text", "text": "a"}, {"type": "text", "text": "b"}]), "ab") + + def test_sse(self) -> None: + value = _sse({"ok": True}) + self.assertTrue(value.startswith("data: ")) + self.assertTrue(value.endswith("\n\n")) + + def test_ollama_surface_and_registry(self) -> None: + with tempfile.TemporaryDirectory() as folder: + path = Path(folder) / "config.json" + path.write_text(json.dumps({"models": []}), encoding="utf-8") + client = TestClient(create_app(path)) + self.assertEqual(client.get("/api/version").status_code, 200) + self.assertEqual(client.get("/api/tags").json(), {"models": []}) + registered = client.post("/api/models/register", json={"id": "demo", "kind": "gguf", "path": "E:/demo.gguf"}) + self.assertEqual(registered.status_code, 200) + self.assertEqual(client.get("/api/tags").json()["models"][0]["name"], "demo") + created = client.post("/api/create", json={"model": "created", "modelfile": "FROM E:/models/created.gguf\nPARAMETER temperature 0.2"}) + self.assertEqual(created.status_code, 200) + copied = client.post("/api/copy", json={"source": "created", "destination": "created-copy"}) + self.assertEqual(copied.status_code, 200) + removed = client.delete("/api/models/demo") + self.assertEqual(removed.status_code, 200) + + +if __name__ == "__main__": + unittest.main()