136 lines
4.2 KiB
Python
136 lines
4.2 KiB
Python
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)
|