Initial commit: LocalPilot：本地模型运行时，复用 llama-server 并提供 Ollama 兼容 provider
This commit is contained in:
WpyQwq
@@ -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
|
||||
Reference in New Issue
Block a user