diff --git a/studio/backend/main.py b/studio/backend/main.py index a9f7004df7..321995885c 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -292,6 +292,27 @@ def _desktop_owner() -> dict[str, str] | None: return _DESKTOP_OWNER +def _start_helper_precache_if_enabled() -> None: + """Start optional Helper LLM GGUF pre-cache only after explicit opt-in.""" + try: + from utils.helper_precache_settings import should_preload_helper_on_startup + if not should_preload_helper_on_startup(): + return + except Exception: + return + + import threading + + def _precache(): + try: + from utils.datasets.llm_assist import precache_helper_gguf + precache_helper_gguf() + except Exception: + pass # non-critical + + threading.Thread(target = _precache, daemon = True, name = "helper-gguf-precache").start() + + @asynccontextmanager async def lifespan(app: FastAPI): """Startup: detect hardware, seed default admin if needed. Shutdown: clean up compiled cache.""" @@ -350,18 +371,7 @@ async def lifespan(app: FastAPI): import structlog structlog.get_logger(__name__).warning("cleanup_orphaned_runs failed at startup: %s", exc) - # Pre-cache the helper GGUF model for LLM-assisted dataset detection, - # in a background thread so it doesn't block server startup. - import threading - - def _precache(): - try: - from utils.datasets.llm_assist import precache_helper_gguf - precache_helper_gguf() - except Exception: - pass # non-critical - - threading.Thread(target = _precache, daemon = True).start() + _start_helper_precache_if_enabled() # Initialize RSA key pair for API key encryption (external providers) from core.inference.key_exchange import init_key_pair diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 94ea4c7b6c..516502c1e5 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -16,6 +16,12 @@ from utils.upload_limits import ( 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() @@ -35,6 +41,16 @@ class UploadLimitResponse(BaseModel): 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, @@ -44,6 +60,13 @@ def _upload_limit_response(limit_mb: int) -> UploadLimitResponse: ) +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()) @@ -64,3 +87,27 @@ def update_upload_limit( 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) diff --git a/studio/backend/tests/test_llm_assist_startup_opt_in.py b/studio/backend/tests/test_llm_assist_startup_opt_in.py new file mode 100644 index 0000000000..e81b1d3775 --- /dev/null +++ b/studio/backend/tests/test_llm_assist_startup_opt_in.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for Helper LLM startup pre-cache opt-in behavior.""" + +from __future__ import annotations + +import sys +import types +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from models.datasets import AiAssistMappingRequest +from routes import datasets as datasets_route +from routes import settings as settings_route +from utils import helper_precache_settings + + +def _install_fake_studio_db(monkeypatch, *, stored = None): + storage_pkg = types.ModuleType("storage") + studio_db = types.ModuleType("storage.studio_db") + values: dict[str, object] = {} + if stored is not None: + values[helper_precache_settings.HELPER_PRECACHE_SETTING_KEY] = stored + + def get_app_setting(key, fallback = None): + return values.get(key, fallback) + + def upsert_app_settings(settings): + values.update(settings) + return dict(values) + + studio_db.get_app_setting = get_app_setting + studio_db.upsert_app_settings = upsert_app_settings + monkeypatch.setitem(sys.modules, "storage", storage_pkg) + monkeypatch.setitem(sys.modules, "storage.studio_db", studio_db) + return values + + +def test_helper_precache_defaults_off_when_setting_missing(monkeypatch): + monkeypatch.delenv("UNSLOTH_HELPER_MODEL_DISABLE", raising = False) + _install_fake_studio_db(monkeypatch) + + assert helper_precache_settings.get_helper_precache_enabled() is False + assert helper_precache_settings.should_preload_helper_on_startup() is False + + +def test_helper_precache_opt_in_is_blocked_by_existing_disable_env(monkeypatch): + _install_fake_studio_db(monkeypatch, stored = True) + monkeypatch.setenv("UNSLOTH_HELPER_MODEL_DISABLE", "true") + + assert helper_precache_settings.get_helper_precache_enabled() is True + assert helper_precache_settings.should_preload_helper_on_startup() is False + + +def test_settings_route_persists_helper_precache_toggle(monkeypatch): + values = _install_fake_studio_db(monkeypatch) + monkeypatch.delenv("UNSLOTH_HELPER_MODEL_DISABLE", raising = False) + + response = settings_route.update_helper_precache( + settings_route.HelperPrecachePayload(enabled = True), + current_subject = "test-user", + ) + + assert response.enabled is True + assert response.default_enabled is False + assert response.disabled_by_env is False + assert values[helper_precache_settings.HELPER_PRECACHE_SETTING_KEY] is True + + +def test_main_startup_uses_helper_precache_gate_instead_of_unconditional_precache(): + source = (Path(__file__).resolve().parent.parent / "main.py").read_text(encoding = "utf-8") + startup_section = source[ + source.index("cleanup_orphaned_runs") : source.index("# Initialize RSA key pair") + ] + + assert "_start_helper_precache_if_enabled()" in startup_section + assert "precache_helper_gguf" not in startup_section + assert "threading.Thread(target = _precache" not in startup_section + + +def test_ai_assist_route_still_calls_on_demand_advisor(monkeypatch): + calls: list[dict] = [] + llm_assist = types.ModuleType("utils.datasets.llm_assist") + + def fake_llm_conversion_advisor(**kwargs): + calls.append(kwargs) + return { + "success": True, + "suggested_mapping": {"prompt": "user", "answer": "assistant"}, + "system_prompt": "Answer carefully.", + "dataset_type": "question_answering", + "is_conversational": False, + "user_notification": "Columns mapped by AI Assist.", + } + + llm_assist.llm_conversion_advisor = fake_llm_conversion_advisor + monkeypatch.setitem(sys.modules, "utils.datasets.llm_assist", llm_assist) + + response = datasets_route.ai_assist_mapping( + AiAssistMappingRequest( + columns = ["prompt", "answer"], + samples = [{"prompt": "x" * 250, "answer": "ok", "extra": "ignored"}], + dataset_name = "owner/dataset", + hf_token = "hf_test", + model_name = "unsloth/test", + model_type = "text", + ), + current_subject = "test-user", + ) + + assert response.success is True + assert response.suggested_mapping == {"prompt": "user", "answer": "assistant"} + assert response.system_prompt == "Answer carefully." + assert calls == [ + { + "column_names": ["prompt", "answer"], + "samples": [{"prompt": "x" * 200, "answer": "ok"}], + "dataset_name": "owner/dataset", + "hf_token": "hf_test", + "model_name": "unsloth/test", + "model_type": "text", + } + ] diff --git a/studio/backend/utils/helper_precache_settings.py b/studio/backend/utils/helper_precache_settings.py new file mode 100644 index 0000000000..db19a2d028 --- /dev/null +++ b/studio/backend/utils/helper_precache_settings.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Persisted opt-in controls for Helper LLM startup pre-cache.""" + +from __future__ import annotations + +import os +from typing import Any + +HELPER_PRECACHE_SETTING_KEY = "helper_model_preload_on_startup" +DEFAULT_HELPER_PRECACHE_ENABLED = False + + +def _coerce_bool(value: Any) -> bool | None: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"0", "false", "no", "off", ""}: + return False + return None + + +def helper_model_disabled_by_env() -> bool: + """Return True when existing broad helper-disable env var is active.""" + return os.environ.get("UNSLOTH_HELPER_MODEL_DISABLE", "").strip() in {"1", "true"} + + +def get_helper_precache_enabled() -> bool: + """Read the persisted startup pre-cache preference. + + Missing or unreadable settings default to False so Studio startup never + performs optional network work unless the user explicitly opted in. + """ + try: + from storage.studio_db import get_app_setting + stored = get_app_setting(HELPER_PRECACHE_SETTING_KEY, None) + except Exception: + stored = None + parsed = _coerce_bool(stored) + return parsed if parsed is not None else DEFAULT_HELPER_PRECACHE_ENABLED + + +def set_helper_precache_enabled(value: Any) -> bool: + """Persist whether Studio should pre-cache the Helper LLM at startup.""" + parsed = _coerce_bool(value) + if parsed is None: + raise ValueError("Helper LLM startup pre-cache must be true or false.") + + from storage.studio_db import upsert_app_settings + + upsert_app_settings({HELPER_PRECACHE_SETTING_KEY: parsed}) + return parsed + + +def should_preload_helper_on_startup() -> bool: + """Gate the startup pre-cache thread. + + The persisted setting is opt-in and the existing broad disable env var wins. + Explicit AI Assist calls do not use this gate; they remain user-triggered. + """ + return get_helper_precache_enabled() and not helper_model_disabled_by_env() diff --git a/studio/frontend/src/features/settings/api/helper-precache.ts b/studio/frontend/src/features/settings/api/helper-precache.ts new file mode 100644 index 0000000000..434a6c7d03 --- /dev/null +++ b/studio/frontend/src/features/settings/api/helper-precache.ts @@ -0,0 +1,88 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { authFetch } from "@/features/auth"; +import { readFastApiError } from "@/lib/format-fastapi-error"; + +const HELPER_PRECACHE_EVENT = "unsloth-helper-precache-change"; + +export type HelperPrecacheSettings = { + enabled: boolean; + defaultEnabled: boolean; + disabledByEnv: boolean; +}; + +type ApiHelperPrecacheSettings = { + enabled: boolean; + // biome-ignore lint/style/useNamingConvention: API schema + default_enabled: boolean; + // biome-ignore lint/style/useNamingConvention: API schema + disabled_by_env: boolean; +}; + +let cachedHelperPrecache: HelperPrecacheSettings | null = null; +let inFlightHelperPrecache: Promise | null = null; + +export function subscribeHelperPrecacheSettings( + listener: (settings: HelperPrecacheSettings) => void, +) { + const handleChange = (event: Event) => { + listener((event as CustomEvent).detail); + }; + window.addEventListener(HELPER_PRECACHE_EVENT, handleChange); + return () => window.removeEventListener(HELPER_PRECACHE_EVENT, handleChange); +} + +function fromApi(settings: ApiHelperPrecacheSettings): HelperPrecacheSettings { + return { + enabled: settings.enabled, + defaultEnabled: settings.default_enabled, + disabledByEnv: settings.disabled_by_env, + }; +} + +function cacheHelperPrecache(settings: HelperPrecacheSettings) { + cachedHelperPrecache = settings; + window.dispatchEvent( + new CustomEvent(HELPER_PRECACHE_EVENT, { detail: settings }), + ); + return settings; +} + +async function fetchHelperPrecacheSettings(): Promise { + const res = await authFetch("/api/settings/helper-precache"); + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to load Helper LLM settings"), + ); + } + return fromApi(await res.json()); +} + +export async function loadHelperPrecacheSettings() { + if (cachedHelperPrecache) { + return cachedHelperPrecache; + } + inFlightHelperPrecache ??= fetchHelperPrecacheSettings() + .then(cacheHelperPrecache) + .finally(() => { + inFlightHelperPrecache = null; + }); + return inFlightHelperPrecache; +} + +export async function updateHelperPrecacheSettings( + enabled: boolean, +): Promise { + const res = await authFetch("/api/settings/helper-precache", { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ enabled }), + }); + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to update Helper LLM settings"), + ); + } + return cacheHelperPrecache(fromApi(await res.json())); +} diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index 57a52b83d1..8a15eb4985 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -15,6 +15,11 @@ import { Switch } from "@/components/ui/switch"; import { usePlatformStore } from "@/config/env"; import { resetOnboardingDone } from "@/features/auth"; import { useChatRuntimeStore } from "@/features/chat"; +import { + loadHelperPrecacheSettings, + updateHelperPrecacheSettings, + type HelperPrecacheSettings, +} from "../api/helper-precache"; import { DEFAULT_UPLOAD_LIMIT_MB, loadUploadLimitSettings, @@ -114,6 +119,12 @@ export function GeneralTab() { ); const [uploadLimitError, setUploadLimitError] = useState(null); const [isSavingUploadLimit, setIsSavingUploadLimit] = useState(false); + const [helperPrecache, setHelperPrecache] = + useState(null); + const [helperPrecacheError, setHelperPrecacheError] = useState( + null, + ); + const [isSavingHelperPrecache, setIsSavingHelperPrecache] = useState(false); const draftRef = useRef(draftToken); useEffect(() => { @@ -158,6 +169,44 @@ export function GeneralTab() { }; }, []); + useEffect(() => { + let cancelled = false; + void loadHelperPrecacheSettings() + .then((settings) => { + if (cancelled) return; + setHelperPrecache(settings); + setHelperPrecacheError(null); + }) + .catch((error) => { + if (cancelled) return; + setHelperPrecacheError( + error instanceof Error + ? error.message + : t("settings.general.helperLlm.loadError"), + ); + }); + return () => { + cancelled = true; + }; + }, [t]); + + const saveHelperPrecache = async (enabled: boolean) => { + setIsSavingHelperPrecache(true); + setHelperPrecacheError(null); + try { + const settings = await updateHelperPrecacheSettings(enabled); + setHelperPrecache(settings); + } catch (error) { + setHelperPrecacheError( + error instanceof Error + ? error.message + : t("settings.general.helperLlm.saveError"), + ); + } finally { + setIsSavingHelperPrecache(false); + } + }; + const saveUploadLimit = async () => { const parsed = Number(draftUploadLimit); if (!Number.isInteger(parsed)) { @@ -236,6 +285,36 @@ export function GeneralTab() { + + +
+ void saveHelperPrecache(enabled)} + /> + {helperPrecache?.disabledByEnv ? ( + + {t("settings.general.helperLlm.disabledByEnv")} + + ) : helperPrecacheError ? ( + + {helperPrecacheError} + + ) : null} +
+
+
+