Initial commit: LocalPilot：本地模型运行时，复用 llama-server 并提供 Ollama 兼容 provider
This commit is contained in:
WpyQwq
@@ -0,0 +1,4 @@
|
||||
"""LocalPilot: a local model runtime and OpenAI-compatible gateway."""
|
||||
|
||||
__version__ = "0.1.0"
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
from .cli import main
|
||||
|
||||
main()
|
||||
|
||||
@@ -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"]
|
||||
@@ -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")
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from ..config import CloudProfile, ModelConfig, RuntimeConfig
|
||||
from ..types import BackendError, ChatChunk, ChatParams
|
||||
from .base import BaseBackend
|
||||
|
||||
|
||||
class RemoteBackend(BaseBackend):
|
||||
"""Proxy any OpenAI-compatible cloud endpoint without exposing its key to clients."""
|
||||
|
||||
def __init__(self, model: ModelConfig, profile: CloudProfile, runtime: RuntimeConfig) -> None:
|
||||
super().__init__(model, runtime)
|
||||
self.profile = profile
|
||||
self.base_url = profile.base_url.rstrip("/")
|
||||
if not self.base_url.endswith("/v1"):
|
||||
self.base_url += "/v1"
|
||||
|
||||
async def load(self) -> None:
|
||||
key = self.profile.resolved_api_key()
|
||||
if not key:
|
||||
raise BackendError(
|
||||
f"云端配置 {self.profile.id} 没有 API key。请设置环境变量 {self.profile.api_key_env or '<api_key_env>'}。"
|
||||
)
|
||||
self.loaded = True
|
||||
|
||||
async def unload(self) -> None:
|
||||
self.loaded = False
|
||||
|
||||
def _headers(self) -> dict[str, str]:
|
||||
key = self.profile.resolved_api_key()
|
||||
if not key:
|
||||
raise BackendError(f"云端配置 {self.profile.id} 的 API key 不可用")
|
||||
headers = {"Authorization": f"Bearer {key}", "Content-Type": "application/json"}
|
||||
headers.update(self.profile.headers)
|
||||
return headers
|
||||
|
||||
async def chat(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
params: ChatParams,
|
||||
stream: bool = False,
|
||||
) -> AsyncIterator[ChatChunk]:
|
||||
if not self.loaded:
|
||||
await self.load()
|
||||
payload: dict[str, Any] = {
|
||||
"model": self.model.remote_model or self.profile.default_model or self.model.id,
|
||||
"messages": messages,
|
||||
"stream": stream,
|
||||
"max_tokens": params.max_tokens,
|
||||
"temperature": params.temperature,
|
||||
"top_p": params.top_p,
|
||||
}
|
||||
if params.stop:
|
||||
payload["stop"] = params.stop
|
||||
if params.seed is not None:
|
||||
payload["seed"] = params.seed
|
||||
if params.response_format is not None:
|
||||
payload["response_format"] = (
|
||||
{"type": params.response_format} if isinstance(params.response_format, str) else params.response_format
|
||||
)
|
||||
if params.tools:
|
||||
payload["tools"] = params.tools
|
||||
if params.tool_choice is not None:
|
||||
payload["tool_choice"] = params.tool_choice
|
||||
payload.update(params.extra)
|
||||
timeout = httpx.Timeout(self.profile.timeout, connect=15.0)
|
||||
async with httpx.AsyncClient(timeout=timeout, headers=self._headers()) as client:
|
||||
if stream:
|
||||
async with client.stream("POST", f"{self.base_url}/chat/completions", json=payload) as response:
|
||||
if response.status_code >= 400:
|
||||
detail = (await response.aread()).decode("utf-8", errors="replace")
|
||||
raise BackendError(f"云端请求失败 {response.status_code}: {detail}")
|
||||
async for line in response.aiter_lines():
|
||||
if not line.startswith("data:"):
|
||||
continue
|
||||
raw = line[5:].strip()
|
||||
if raw == "[DONE]":
|
||||
yield ChatChunk(done=True, finish_reason="stop")
|
||||
return
|
||||
try:
|
||||
data = json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
choice = (data.get("choices") or [{}])[0]
|
||||
delta = choice.get("delta") or {}
|
||||
text = delta.get("content") or ""
|
||||
finish = choice.get("finish_reason")
|
||||
if text or finish or len(delta) > 1:
|
||||
yield ChatChunk(text=text, finish_reason=finish, metadata=data, delta=delta)
|
||||
yield ChatChunk(done=True, finish_reason="stop")
|
||||
return
|
||||
|
||||
response = await client.post(f"{self.base_url}/chat/completions", json=payload)
|
||||
if response.status_code >= 400:
|
||||
raise BackendError(f"云端请求失败 {response.status_code}: {response.text}")
|
||||
data = response.json()
|
||||
choice = (data.get("choices") or [{}])[0]
|
||||
message = choice.get("message") or {}
|
||||
yield ChatChunk(
|
||||
text=message.get("content") or "",
|
||||
done=True,
|
||||
finish_reason=choice.get("finish_reason") or "stop",
|
||||
metadata=data,
|
||||
delta=message,
|
||||
)
|
||||
|
||||
def info(self) -> dict[str, Any]:
|
||||
return {
|
||||
"id": self.model.id,
|
||||
"kind": "cloud",
|
||||
"profile": self.profile.id,
|
||||
"remote_model": self.model.remote_model or self.profile.default_model or self.model.id,
|
||||
"loaded": self.loaded,
|
||||
"base_url": self.base_url,
|
||||
}
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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."""
|
||||
Reference in New Issue
Block a user