0d6d7dd4b3* Studio: make Helper LLM startup pre-cache opt in * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
113 lines
3.6 KiB
Python
113 lines
3.6 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
from fastapi import APIRouter, Depends
|
|
from pydantic import BaseModel, Field
|
|
|
|
from auth.authentication import get_current_subject
|
|
from loggers import get_logger
|
|
from utils.utils import safe_error_detail, log_and_http_error
|
|
from utils.upload_limits import (
|
|
MAX_UPLOAD_LIMIT_MB,
|
|
MIN_UPLOAD_LIMIT_MB,
|
|
default_upload_limit_mb,
|
|
get_upload_limit_mb,
|
|
set_upload_limit_mb,
|
|
upload_limit_bytes,
|
|
upload_limit_label,
|
|
)
|
|
from utils.helper_precache_settings import (
|
|
DEFAULT_HELPER_PRECACHE_ENABLED,
|
|
get_helper_precache_enabled,
|
|
helper_model_disabled_by_env,
|
|
set_helper_precache_enabled,
|
|
)
|
|
|
|
router = APIRouter()
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class UploadLimitPayload(BaseModel):
|
|
max_upload_size_mb: int = Field(..., ge = MIN_UPLOAD_LIMIT_MB, le = MAX_UPLOAD_LIMIT_MB)
|
|
|
|
|
|
class UploadLimitResponse(BaseModel):
|
|
max_upload_size_mb: int
|
|
max_upload_size_bytes: int
|
|
max_upload_size_label: str
|
|
default_upload_size_mb: int
|
|
min_upload_size_mb: int = MIN_UPLOAD_LIMIT_MB
|
|
max_allowed_upload_size_mb: int = MAX_UPLOAD_LIMIT_MB
|
|
|
|
|
|
class HelperPrecachePayload(BaseModel):
|
|
enabled: bool
|
|
|
|
|
|
class HelperPrecacheResponse(BaseModel):
|
|
enabled: bool
|
|
default_enabled: bool = DEFAULT_HELPER_PRECACHE_ENABLED
|
|
disabled_by_env: bool
|
|
|
|
|
|
def _upload_limit_response(limit_mb: int) -> UploadLimitResponse:
|
|
return UploadLimitResponse(
|
|
max_upload_size_mb = limit_mb,
|
|
max_upload_size_bytes = upload_limit_bytes(limit_mb),
|
|
max_upload_size_label = upload_limit_label(limit_mb),
|
|
default_upload_size_mb = default_upload_limit_mb(),
|
|
)
|
|
|
|
|
|
def _helper_precache_response(enabled: bool | None = None) -> HelperPrecacheResponse:
|
|
return HelperPrecacheResponse(
|
|
enabled = get_helper_precache_enabled() if enabled is None else enabled,
|
|
disabled_by_env = helper_model_disabled_by_env(),
|
|
)
|
|
|
|
|
|
@router.get("/upload-limit", response_model = UploadLimitResponse)
|
|
def get_upload_limit(current_subject: str = Depends(get_current_subject)) -> UploadLimitResponse:
|
|
return _upload_limit_response(get_upload_limit_mb())
|
|
|
|
|
|
@router.put("/upload-limit", response_model = UploadLimitResponse)
|
|
def update_upload_limit(
|
|
payload: UploadLimitPayload, current_subject: str = Depends(get_current_subject)
|
|
) -> UploadLimitResponse:
|
|
try:
|
|
limit_mb = set_upload_limit_mb(payload.max_upload_size_mb)
|
|
except ValueError as exc:
|
|
raise log_and_http_error(
|
|
exc,
|
|
400,
|
|
safe_error_detail(exc, fallback = "Invalid upload limit."),
|
|
event = "settings.update_upload_limit_failed",
|
|
log = logger,
|
|
) from exc
|
|
return _upload_limit_response(limit_mb)
|
|
|
|
|
|
@router.get("/helper-precache", response_model = HelperPrecacheResponse)
|
|
def get_helper_precache(
|
|
current_subject: str = Depends(get_current_subject),
|
|
) -> HelperPrecacheResponse:
|
|
return _helper_precache_response()
|
|
|
|
|
|
@router.put("/helper-precache", response_model = HelperPrecacheResponse)
|
|
def update_helper_precache(
|
|
payload: HelperPrecachePayload, current_subject: str = Depends(get_current_subject)
|
|
) -> HelperPrecacheResponse:
|
|
try:
|
|
enabled = set_helper_precache_enabled(payload.enabled)
|
|
except ValueError as exc:
|
|
raise log_and_http_error(
|
|
exc,
|
|
400,
|
|
safe_error_detail(exc, fallback = "Invalid Helper LLM pre-cache setting."),
|
|
event = "settings.update_helper_precache_failed",
|
|
log = logger,
|
|
) from exc
|
|
return _helper_precache_response(enabled)
|