74 lines
2.7 KiB
Python
74 lines
2.7 KiB
Python
import base64
|
|
import hashlib
|
|
import json
|
|
|
|
from cryptography.fernet import Fernet, InvalidToken
|
|
|
|
from .config import get_settings
|
|
from .llm import DeepSeekClient
|
|
from .models import AppSetting
|
|
|
|
|
|
def _cipher():
|
|
secret = get_settings().wisegrade_secret_key
|
|
if len(secret) < 32 or secret == "replace-with-a-random-32-character-secret":
|
|
raise ValueError("请设置至少32位随机WISEGRADE_SECRET_KEY后保存模型配置")
|
|
return Fernet(base64.urlsafe_b64encode(hashlib.sha256(secret.encode()).digest()))
|
|
|
|
|
|
def _get(db, key):
|
|
item = db.get(AppSetting, key)
|
|
if not item:
|
|
return None
|
|
try:
|
|
return _cipher().decrypt(item.encrypted_value).decode()
|
|
except InvalidToken as exc:
|
|
raise ValueError("已保存的模型设置无法解密，请恢复原WISEGRADE_SECRET_KEY") from exc
|
|
|
|
|
|
def _set(db, key, value):
|
|
item = db.get(AppSetting, key)
|
|
if not item:
|
|
item = AppSetting(key=key, encrypted_value=b"")
|
|
db.add(item)
|
|
item.encrypted_value = _cipher().encrypt(value.encode())
|
|
|
|
|
|
def model_config(db):
|
|
settings = get_settings()
|
|
base = {"base_url": settings.deepseek_base_url, "model": settings.deepseek_model,
|
|
"timeout_seconds": settings.model_timeout_seconds, "concurrency": settings.model_concurrency,
|
|
"retries": settings.model_retries}
|
|
stored = _get(db, "model_config") if db.get(AppSetting, "model_config") else None
|
|
return {**base, **(json.loads(stored) if stored else {})}
|
|
|
|
|
|
def key_configured(db):
|
|
return bool(get_settings().deepseek_api_key or db.get(AppSetting, "api_key"))
|
|
|
|
|
|
def save_model_config(db, data: dict):
|
|
config = model_config(db)
|
|
config.update({key: data[key] for key in config if key in data})
|
|
if "api_key" in data and data["api_key"]:
|
|
_set(db, "api_key", data["api_key"])
|
|
_set(db, "model_config", json.dumps(config))
|
|
db.commit()
|
|
return {**config, "api_key_configured": key_configured(db)}
|
|
|
|
|
|
def client_for_batch(db, batch):
|
|
stored_key = _get(db, "api_key") if db.get(AppSetting, "api_key") else None
|
|
config = batch.model_config_snapshot
|
|
return DeepSeekClient(base_url=config["base_url"], model=config["model"],
|
|
api_key=stored_key or get_settings().deepseek_api_key,
|
|
timeout=config["timeout_seconds"], retries=config["retries"])
|
|
|
|
|
|
def client_for_chat(db):
|
|
stored_key = _get(db, "api_key") if db.get(AppSetting, "api_key") else None
|
|
config = model_config(db)
|
|
return DeepSeekClient(base_url=config["base_url"], model=config["model"],
|
|
api_key=stored_key or get_settings().deepseek_api_key,
|
|
timeout=config["timeout_seconds"], retries=config["retries"])
|