Files
llm-backend/localpilot/server.py
T

635 lines
26 KiB
Python

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