From bdb958e052eca6a47c17410558562f00481f71b8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 04:16:57 -0700 Subject: [PATCH 1/7] Guard RoPE scaling against the transformers v5 buffer blank; honor extended RoPE factor (#6925) * Guard RoPE scaling against the transformers v5 buffer blank; honor extended factor Add a family-agnostic guard that builds each rotary from a scaled config, blanks its non-persistent buffers (what transformers v5 does on load), runs loader._fix_rope_inv_freq, and asserts every buffer is restored to its scaled value (llama3 and longrope). This catches the whole bug class, not just the one call site, and is validated to fail on the pre-fix repair. Also make LlamaExtendedRotaryEmbedding read the llama3 factor from the config instead of hardcoding 8 (wrong for Llama-3.2, factor 32), falling back to the Llama-3.1 defaults when built without a config. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Pass config into extended rotary codegen; skip v5 round-trip on transformers 4.x - patch_llama_rope_scaling now builds the llama3 extended rotary with config=self.config so it reads the real factor (32 for Llama-3.2) instead of falling back to 8; the template already references self.config. - test_v5_blank_repair_roundtrip now skips when loader._NEEDS_ROPE_FIX is False, since _fix_rope_inv_freq is a no-op on transformers 4.x and cannot restore the blanked buffers there. * Raise stream deadlock-guard timeouts from 0.2s to 5.0s in passthrough tests These asyncio.wait_for guards bound test setup and cross-task event signaling that complete near-instantly on success; the 0.2s budget is a latency assertion in disguise and times out under CI scheduling load (seen on the 3.11 matrix leg while 3.10/3.12/3.13 pass the same commit). 5.0s matches the timeout used elsewhere in the suite and still fails fast on a real hang. No test relies on the guard expiring. * Extended rotary reads rope_parameters as well as rope_scaling transformers v5 stores llama3 scaling under config.rope_parameters and exposes rope_scaling only as a back-compat property. Reading that property works on 5.0-5.13 (verified: factor resolves to 32 for Llama-3.2), but a future release may drop the shim, after which the subclass path would fall back to factor 8. Read either field so the factor survives the rename. Adds test_extended_rotary_reads_rope_parameters_v5 (fails on the old single-field read: rope_parameters-only config resolves to 8, not 32). * [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> --- .../tests/test_openai_tool_passthrough.py | 18 +-- tests/utils/test_rope_scaling_drift.py | 138 ++++++++++++++++++ unsloth/models/_utils.py | 1 + unsloth/models/llama.py | 17 ++- 4 files changed, 160 insertions(+), 14 deletions(-) diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index ccbd78e2b1..05e017ba7f 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -1900,7 +1900,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) @@ -2044,7 +2044,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) gate.set() @@ -2107,7 +2107,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) @@ -2190,7 +2190,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) @@ -2252,7 +2252,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) assert cancel_id in inf_mod._CANCEL_REGISTRY @@ -2323,13 +2323,13 @@ class TestApiMonitorProviderAndCompletionStreams: monitor_id = monitor_id, ) ) - await asyncio.wait_for(entered.wait(), timeout = 0.2) + await asyncio.wait_for(entered.wait(), timeout = 5.0) assert cancel_id in inf_mod._CANCEL_REGISTRY task.cancel() with pytest.raises(asyncio.CancelledError): await task - await asyncio.wait_for(cancelled.wait(), timeout = 0.2) + await asyncio.wait_for(cancelled.wait(), timeout = 5.0) assert cancel_id not in inf_mod._CANCEL_REGISTRY asyncio.run(_run()) @@ -2389,13 +2389,13 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) assert cancel_id in inf_mod._CANCEL_REGISTRY gate.set() - await asyncio.wait_for(returned.wait(), timeout = 0.2) + await asyncio.wait_for(returned.wait(), timeout = 5.0) await asyncio.sleep(0) await response._unstarted_cleanup() assert upstream_response.is_closed diff --git a/tests/utils/test_rope_scaling_drift.py b/tests/utils/test_rope_scaling_drift.py index 98f7e2db62..7a738e236c 100644 --- a/tests/utils/test_rope_scaling_drift.py +++ b/tests/utils/test_rope_scaling_drift.py @@ -257,6 +257,63 @@ def test_recompute_helper_scales_on_cpu(): ), "_unsloth_recompute_inv_freq must return vanilla inv_freq when unscaled." +def test_extended_rotary_reads_config_factor(): + # LlamaExtendedRotaryEmbedding must honor the config factor, not hardcode 8 + # (Llama-3.2 uses 32); otherwise the subclass path re-drops scaling (#2405). + from types import SimpleNamespace + + from unsloth.models.llama import LlamaExtendedRotaryEmbedding + + rot = object.__new__(LlamaExtendedRotaryEmbedding) + rot.base = ROPE_THETA + rot.dim = HEAD_DIM + rot._unsloth_rope_config = SimpleNamespace( + rope_scaling = { + "rope_type": "llama3", + "factor": 32.0, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": 8192, + } + ) + vanilla = _vanilla_inv_freq() + scaled = rot._apply_inv_freq_scaling(vanilla).reshape(-1) + ratio = float(vanilla[-1]) / float(scaled[-1]) + assert abs(ratio - 32.0) < 1e-3, ( + f"LlamaExtendedRotaryEmbedding ignored config factor 32 (ratio {ratio}); the " + "low-frequency band must be divided by the config factor (issue #2405)." + ) + + +def test_extended_rotary_reads_rope_parameters_v5(): + # transformers v5 stores scaling under rope_parameters (rope_scaling is a + # back-compat shim that may be removed); the factor must still be read. + from types import SimpleNamespace + + from unsloth.models.llama import LlamaExtendedRotaryEmbedding + + rot = object.__new__(LlamaExtendedRotaryEmbedding) + rot.base = ROPE_THETA + rot.dim = HEAD_DIM + rot._unsloth_rope_config = SimpleNamespace( + rope_scaling = None, + rope_parameters = { + "rope_type": "llama3", + "factor": 32.0, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": 8192, + }, + ) + vanilla = _vanilla_inv_freq() + scaled = rot._apply_inv_freq_scaling(vanilla).reshape(-1) + ratio = float(vanilla[-1]) / float(scaled[-1]) + assert abs(ratio - 32.0) < 1e-3, ( + f"Extended rotary ignored rope_parameters factor 32 (ratio {ratio}); v5 " + "keeps the factor under rope_parameters, not rope_scaling." + ) + + def _cos_at_position(rot, position): """cos row at one position, built like _set_cos_sin_cache but CPU-only.""" inv_freq = rot.inv_freq.float().cpu() @@ -324,6 +381,87 @@ def test_extended_cache_keeps_scaling_after_growth(): ) +def _blank_nonpersistent_buffers(module): + """Mimic transformers v5 meta-load: overwrite non-persistent buffers with garbage.""" + for name, buf in list(module.named_buffers()): + leaf = module + *parents, attr = name.split(".") + for part in parents: + leaf = getattr(leaf, part) + if attr in getattr(leaf, "_non_persistent_buffers_set", set()): + setattr(leaf, attr, torch.rand_like(buf)) + + +def _build_llama3_rotary(): + from unsloth.models import llama as llama_mod + config = _make_config(LLAMA3_ROPE_SCALING) + return llama_mod.LlamaRotaryEmbedding(config = config), config + + +def _build_longrope_rotary(): + from types import SimpleNamespace + + from unsloth.models import llama as llama_mod + + short_factor, long_factor = [1.05] * 48, [1.3] * 48 + rot = llama_mod.LongRopeRotaryEmbedding( + dim = 96, + max_position_embeddings = 131072, + original_max_position_embeddings = 4096, + base = ROPE_THETA, + short_factor = short_factor, + long_factor = long_factor, + ) + config = SimpleNamespace( + rope_scaling = { + "rope_type": "longrope", + "short_factor": short_factor, + "long_factor": long_factor, + "original_max_position_embeddings": 4096, + } + ) + return rot, config + + +@requires_cuda +@pytest.mark.parametrize( + "build", [_build_llama3_rotary, _build_longrope_rotary], ids = ["llama3", "longrope"] +) +def test_v5_blank_repair_roundtrip(build): + # Build scaled -> blank non-persistent buffers (what transformers v5 does on + # load) -> run the repair -> every buffer must return to its scaled value. + # Family-agnostic: encodes no scaling math, so it guards any rotary that + # keeps scaling in a buffer (issue #2405 / PR #6907). + from unsloth.models import loader + + # The repair only runs on transformers v5 (it is what blanks the buffers); + # on v4 _fix_rope_inv_freq is a no-op, so the round-trip cannot restore. + if not loader._NEEDS_ROPE_FIX: + pytest.skip("transformers < 5 does not blank rope buffers; repair is a no-op") + + rot, config = build() + snapshot = {name: buf.detach().clone() for name, buf in rot.named_buffers()} + assert snapshot, "rotary registers no buffers; nothing to guard" + + _blank_nonpersistent_buffers(rot) + assert any( + not torch.equal(rot.get_buffer(name), snapshot[name]) for name in snapshot + ), "blanking changed no buffer; the round-trip would be vacuous" + + wrapper = torch.nn.Module() + wrapper.add_module("rotary_emb", rot) + wrapper.config = config + loader._fix_rope_inv_freq(wrapper) + + for name in snapshot: + assert torch.allclose( + rot.get_buffer(name).cpu(), snapshot[name].cpu(), rtol = 1e-4, atol = 1e-6 + ), ( + f"{name} was not restored to its scaled value by loader._fix_rope_inv_freq " + "after the transformers v5 buffer blank (issue #2405 / PR #6907)." + ) + + def test_object_style_rope_scaling_does_not_crash(): # Object-style rope_scaling must be normalized, not .get()'d directly. from dataclasses import dataclass diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 169b610988..1aa2c6e820 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -2834,6 +2834,7 @@ def patch_llama_rope_scaling( dim = self.head_dim, max_position_embeddings=self.max_position_embeddings, base=self.rope_theta, + config=self.config, ) elif scaling_type == "longrope": self.rotary_emb = {longrope_rope_function}( diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index c25a031b82..a1da099758 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1930,11 +1930,18 @@ class LlamaExtendedRotaryEmbedding(LlamaRotaryEmbedding): # From https://github.com/meta-llama/llama-models/blob/main/models/llama3_1/api/model.py#L41 def _apply_inv_freq_scaling(self, freqs: torch.Tensor): - # Values obtained from grid search - scale_factor = 8 - low_freq_factor = 1 - high_freq_factor = 4 - old_context_len = 8192 # original llama3 length + # llama3 factors from config; Llama-3.1 defaults when built without one + # (legacy codegen path). Hardcoding 8 is wrong for e.g. Llama-3.2 (32). + # v5 renames rope_scaling -> rope_parameters; read either so the factor + # survives even if the rope_scaling back-compat shim is dropped. + config = getattr(self, "_unsloth_rope_config", None) + rope_scaling = _rope_scaling_as_dict( + getattr(config, "rope_scaling", None) or getattr(config, "rope_parameters", None) or {} + ) + scale_factor = rope_scaling.get("factor", 8) + low_freq_factor = rope_scaling.get("low_freq_factor", 1) + high_freq_factor = rope_scaling.get("high_freq_factor", 4) + old_context_len = rope_scaling.get("original_max_position_embeddings", 8192) low_freq_wavelen = old_context_len / low_freq_factor high_freq_wavelen = old_context_len / high_freq_factor From 414503745e71b52c37117d7a96894ea46fc13ec5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 04:30:21 -0700 Subject: [PATCH 2/7] Run the malware gate on the RAG embedding model before it loads (#6887) * Run the malware gate on the RAG embedding model before it loads Setting the RAG embedding model through PUT /api/settings/embedding-model persisted an arbitrary repo and later handed it straight to SentenceTransformer, which deserializes pickle weights. Unlike the normal model-load paths, this route never ran evaluate_file_security, and force skipped verification entirely, so a repo Hugging Face flags as unsafe (or any repo under force) could be downloaded and loaded in the backend process without a scan. Run the malware/pickle scan at both ends: the settings endpoint now scans before persisting and returns 409 on a flagged repo even under force (force still only skips the is-embedding-model type check for offline or local repos), and the embedder scans again at the load sink so a name that arrives via env or default is covered too. Local paths and unreachable scans fail open inside evaluate_file_security, and the sink never bricks the embedder on a gate error. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Thread the load token into the embedding scan and hard-fail on a block The load-sink scan ran without a token, so evaluate_file_security (which passes token=False when none is given) could not reach a gated or private repo and failed open for exactly the model SentenceTransformer would still load. Resolve the loader's own token (HF_TOKEN env or the cached login) and pass it to the sink scan, and fall back to it in the settings endpoint when the request omits one. The sink previously raised a plain RuntimeError, which the llama-server fallback in encode() and _build_st_backend_or_fallback() swallowed as a routine ST failure, silently switching backends instead of blocking. Raise a distinct UnsafeEmbeddingModelError that both fallback paths re-raise, so a flagged model hard-fails. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Scan sentence-transformers module dirs and scope the embedding pickle gate to the ST backend Extend the RAG embedding malware gate so a poisoned pickle under a SentenceTransformer module dir (for example 0_Transformer/pytorch_model.bin) blocks. Those dirs are read from the repo's modules.json and passed as load roots to evaluate_file_security at both the settings endpoint and the load sink, so such a pickle is treated as root-level there instead of an unreferenced nested shard that was previously allowed. Scope the ST pickle scan to the sentence-transformers backend. On the llama-server backend the embedder loads GGUF files (inert) from the -GGUF companion repo, never the ST repo's pickle, so a custom ST repo with a flagged pickle and a clean GGUF companion is no longer rejected. The existing GGUF availability checks already cover that path. Return 403 for the hard security block instead of 409. The settings UI routes every 409 into the forceable save-anyway flow, but this block cannot be bypassed by force, so it now uses a distinct status the client treats as non-forceable. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Base the embedding pickle scan on the actual backend, not just the resolver _llama_backend_active only consulted the auto resolver, so on a GPU box where auto resolves to sentence-transformers but the process already fell back to the llama-server backend at runtime (a torch or CUDA load/encode failure), it returned False and the settings endpoint hard-blocked a save whose ST pickle is flagged even though the process loads only inert GGUF. Add active_backend_is_llama, which reflects the actual built backend (True when the cached backend is a LlamaServerBackend, including a runtime fallback) and otherwise defers to the resolver as a fresh process would, and delegate _llama_backend_active to it. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Report the cached embedding backend verbatim, not the resolver active_backend_is_llama() fell through to the config resolver whenever a backend was already built but was not llama-server, so a live sentence-transformers backend could report llama=True once the resolver picked llama (GPU heuristic or a runtime config change) and wrongly skip its pickle scan. Once a backend exists, return isinstance(backend, LlamaServerBackend) directly; only defer to the resolver before any backend is built. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/rag/embeddings.py | 117 ++++++ studio/backend/routes/settings.py | 78 +++- .../test_embedding_model_security_gate.py | 365 ++++++++++++++++++ .../tests/test_security_gate_consistency.py | 13 + .../features/settings/api/embedding-model.ts | 9 + .../features/settings/tabs/general-tab.tsx | 6 +- 6 files changed, 573 insertions(+), 15 deletions(-) create mode 100644 studio/backend/tests/test_embedding_model_security_gate.py diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 345b4dd853..47d26209b4 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -63,6 +63,87 @@ def _install_torchao_stub_once() -> None: install_torchao_windows_rocm_stub() +class UnsafeEmbeddingModelError(RuntimeError): + """Raised when the embedding model repo is flagged unsafe. A distinct type so the + llama-server fallback paths re-raise it instead of masking a security block as a + routine ST failure.""" + + +def _ambient_hf_token() -> str | None: + """The HF token the loader itself would use (HF_TOKEN env or the cached login), so + the scan can reach a gated/private repo instead of failing open. None if unavailable.""" + try: + from huggingface_hub import get_token + return get_token() + except Exception: + return None + + +def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]: + """The module directories a SentenceTransformer load reads weights from, taken from + the repo's ``modules.json`` (each module's non-empty ``path``, e.g. ``0_Transformer``). + ST deserializes ``pytorch_model.bin`` from these dirs, so they are load roots for the + security scan: a flagged pickle directly under one must block. Returns () on any + failure (no modules.json, offline, malformed) so the guard never bricks the embedder. + """ + try: + import json + + from utils.paths import is_local_path + + if is_local_path(name): + from pathlib import Path + from utils.paths import normalize_path + + path = Path(normalize_path(name)).expanduser() / "modules.json" + if not path.is_file(): + return () + data = json.loads(path.read_text()) + else: + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import EntryNotFoundError + + try: + local = hf_hub_download(name, "modules.json", token = token or None) + except EntryNotFoundError: + return () + data = json.loads(open(local).read()) + subdirs = [] + for module in data or (): + sub = str((module or {}).get("path", "")).strip().strip("/") + if sub: + subdirs.append(sub) + return tuple(dict.fromkeys(subdirs)) + except Exception: + return () + + +def _guard_model_security(name: str) -> None: + """Refuse to load a repo HF flagged as unsafe: a poisoned pickle deserializes inside + SentenceTransformer regardless of trust_remote_code. Defense in depth behind the + /settings gate (a name can also arrive via env/default); local paths and unreachable + scans fail open inside evaluate_file_security. Never bricks the embedder on a gate error. + """ + try: + from utils.security import evaluate_file_security, security_load_subdirs + + token = _ambient_hf_token() + # Union the audio-model load roots with the ST module dirs so a flagged pickle + # directly under a Transformer module dir (0_Transformer/) blocks instead of + # passing as an unreferenced nested shard. + load_subdirs = tuple( + dict.fromkeys((*security_load_subdirs(name, token), *_st_module_subdirs(name, token))) + ) + blocked = evaluate_file_security(name, hf_token = token, load_subdirs = load_subdirs).blocked + except Exception: + return + if blocked: + raise UnsafeEmbeddingModelError( + f"Embedding model {name!r} is flagged as unsafe by Hugging Face's security " + "scan; refusing to load. Set a different RAG embedding model." + ) + + def _get(model_name: str | None = None): """Cached SentenceTransformer, (re)loading on a name change. Loaded in fp16 for a ~1.5x speedup at negligible accuracy loss.""" @@ -75,6 +156,7 @@ def _get(model_name: str | None = None): device = _device() logger.info("loading embedding model %s on %s", name, device) + _guard_model_security(name) _model = SentenceTransformer( name, device = device, model_kwargs = {"torch_dtype": "float16"} ) @@ -159,6 +241,8 @@ class _SentenceTransformersBackend: ): try: return _st_encode(texts, model_name = model_name, normalize = normalize) + except UnsafeEmbeddingModelError: + raise # a security block must hard-fail, not fall back to llama-server except Exception as st_err: # noqa: BLE001 - runtime ST/CUDA encode failure # ST loaded but this encode blew up; swap the process to the llama-server # embedder (so later encodes stay in one space) and retry. @@ -222,6 +306,8 @@ def _build_st_backend_or_fallback(): try: backend.warm(model_name = None) return backend + except UnsafeEmbeddingModelError: + raise # a security block must hard-fail, not fall back to llama-server except Exception as st_err: # noqa: BLE001 - any ST/torch import or load failure fallback = _try_make_llama_backend() if fallback is None: @@ -290,6 +376,37 @@ def _reset_backend() -> None: _backend_key = None +def active_backend_is_llama() -> bool: + """True when this process actually embeds via the llama-server (GGUF) backend. + + Reflects the ACTUAL built backend once one exists: an ``auto`` install that + resolves to sentence-transformers but then falls back to llama-server at + runtime (``_build_st_backend_or_fallback`` on a torch/CUDA load failure, or + ``_switch_to_llama_fallback`` on an encode failure) loads only inert GGUF, so + callers gating on the ST pickle must see llama here. Before any backend is + built, defers to the resolver (``auto`` -> ``_resolve_auto()``, else the raw + key) exactly as a fresh process would. Never raises: a backend probe must not + block saving a model.""" + try: + with _backend_lock: + backend = _backend + if backend is not None: + # A backend exists: report what it ACTUALLY is. A concrete + # sentence-transformers backend must return False even if the + # resolver would now pick llama, so its pickle stays gated. If the + # llama import fails we cannot be llama, so fall to the safe False. + try: + from .embed_llama_server import LlamaServerBackend + except Exception: # noqa: BLE001 - llama plumbing import must never block + return False + return isinstance(backend, LlamaServerBackend) + raw = (config.EMBED_BACKEND or "auto").strip().lower() + key = _resolve_auto() if raw in _AUTO_ALIASES else raw + return key in _LLAMA_ALIASES + except Exception: # noqa: BLE001 - a backend probe must never block saving + return False + + def warm(model_name: str | None = None) -> None: """Eagerly load the embedder so the first real request isn't slow.""" _get_backend().warm(model_name = model_name) diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 862bce8be8..914699f540 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -260,17 +260,29 @@ def _embedding_model_response() -> EmbeddingModelResponse: ) -def _llama_backend_active() -> bool: - """True when this install embeds via the llama-server (GGUF) backend.""" - from core.rag import config as rag_config - from core.rag import embeddings - +def _ambient_hf_token() -> Optional[str]: + """The HF token the loader would use (HF_TOKEN env or the cached login), so a gated + repo is scanned rather than failing open. None if unavailable.""" try: - raw = (rag_config.EMBED_BACKEND or "auto").strip().lower() - key = embeddings._resolve_auto() if raw in embeddings._AUTO_ALIASES else raw + from huggingface_hub import get_token + return get_token() + except Exception: + return None + + +def _llama_backend_active() -> bool: + """True when this install actually embeds via the llama-server (GGUF) backend. + + Delegates to the embeddings module so a runtime fallback from + sentence-transformers to llama-server (after a torch/CUDA load or encode + failure) is honored: in that state the process loads only inert GGUF, so the + ST pickle gate below must not hard-block a repo whose GGUF companion is clean. + Before any backend is built this still reflects the resolver.""" + from core.rag import embeddings + try: + return embeddings.active_backend_is_llama() except Exception: # noqa: BLE001 - backend probe must never block saving return False - return key in embeddings._LLAMA_ALIASES def _resolves_as_local_gguf(model: str) -> bool: @@ -357,6 +369,8 @@ def update_embedding_model( """Set the RAG embedding model. Unless ``force`` is set, the repo is verified to be an embedding model via HF metadata; an unverifiable model (wrong type, typo, gated repo, or no network) returns 409 so the UI can offer "save anyway". + A repo flagged unsafe by HF's security scan returns 403 instead: a hard block + that ``force`` cannot bypass, so the UI must not offer "save anyway". Documents indexed under the previous model must be re-uploaded.""" from utils.models import is_embedding_model @@ -370,15 +384,51 @@ def update_embedding_model( event = "settings.update_embedding_model_failed", log = logger, ) from exc + hf_token = (payload.hf_token or "").strip() or None # The env/default model needs no verification; saving it is a no-op override. # A local GGUF on the llama-server backend is accepted as-is: it is exactly # what the backend loads, and HF metadata cannot verify a local path. - if ( - model != default_embedding_model() - and not payload.force - and not (_llama_backend_active() and _resolves_as_local_gguf(model)) - ): - hf_token = (payload.hf_token or "").strip() or None + is_local_gguf = _llama_backend_active() and _resolves_as_local_gguf(model) + # The pickle gate only matters for the sentence-transformers backend, which is what + # deserializes pickles. On the llama-server backend the embedder loads GGUF files + # (inert) from effective_gguf_repo(), so scanning the ST repo's pickle here would + # wrongly reject a custom repo whose GGUF companion is clean; the GGUF availability + # checks below cover that path instead. + scan_st_pickle = ( + model != default_embedding_model() and not is_local_gguf and not _llama_backend_active() + ) + if scan_st_pickle: + # Malware/pickle gate before we persist a repo the embedder later loads with + # SentenceTransformer. Runs even under force (force only skips the is-embedding + # type check for offline/local repos HF cannot verify); local paths and + # unreachable scans fail open inside evaluate_file_security. + from utils.security import evaluate_file_security, security_load_subdirs + from core.rag.embeddings import _st_module_subdirs + + # Fall back to the loader's own token so a gated/private repo is actually scanned + # (a token-less scan fails open for exactly the repo that would still load). + scan_token = hf_token or _ambient_hf_token() + # Include the ST module dirs (0_Transformer/) so a flagged pickle directly under + # one blocks instead of passing as an unreferenced nested shard. + load_subdirs = tuple( + dict.fromkeys( + ( + *security_load_subdirs(model, scan_token), + *_st_module_subdirs(model, scan_token), + ) + ) + ) + if evaluate_file_security(model, hf_token = scan_token, load_subdirs = load_subdirs).blocked: + # 403, not 409: the client routes every 409 into the forceable "save anyway" + # flow, but this block is a hard, non-forceable security refusal. + raise HTTPException( + status_code = 403, + detail = ( + f"{model!r} is flagged as unsafe by Hugging Face's security scan and " + "cannot be used as the embedding model." + ), + ) + if model != default_embedding_model() and not payload.force and not is_local_gguf: from core.rag import config as rag_config # A GGUF-named repo on the llama-server backend is loaded from its .gguf diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py new file mode 100644 index 0000000000..940b35d7ba --- /dev/null +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -0,0 +1,365 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""The RAG embedding model must pass the malware/pickle gate before it is persisted or +loaded. A flagged repo (or any repo saved with force) previously reached +SentenceTransformer unscanned, bypassing the normal model-load protections.""" + +from pathlib import Path +import sys +import types as _types + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +import routes.settings as settings + + +class _Decision: + def __init__(self, blocked): + self.blocked = blocked + + +def _security_stub(blocked): + mod = _types.ModuleType("utils.security") + mod.evaluate_file_security = lambda *a, **k: _Decision(blocked) + mod.security_load_subdirs = lambda *a, **k: () + return mod + + +@pytest.fixture +def client(monkeypatch): + # The settings scan unions in the ST module dirs read from modules.json; keep it + # offline and deterministic for the endpoint tests that use this fixture. + import core.rag.embeddings as embeddings + + monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ()) + saved: dict = {} + monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) + monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + monkeypatch.setattr(settings, "_llama_backend_active", lambda: False) + monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) + monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) + monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + return TestClient(app, raise_server_exceptions = False), saved + + +def test_flagged_repo_is_blocked_even_with_force(client, monkeypatch): + c, saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) + r = c.put( + "/embedding-model", json = {"embedding_model": "attacker/malicious-embed", "force": True} + ) + # 403, not the forceable 409, so the client does not offer "save anyway". + assert r.status_code == 403 + assert "model" not in saved # force must not persist a flagged repo + + +def test_flagged_repo_is_blocked_without_force(client, monkeypatch): + c, saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) + r = c.put("/embedding-model", json = {"embedding_model": "attacker/malicious-embed"}) + assert r.status_code == 403 + assert "model" not in saved + + +def test_hard_block_uses_non_forceable_status(client, monkeypatch): + # The forceable verification path uses 409; the hard security block must be distinct + # (403) so the frontend never routes it into the "save anyway" force flow. + c, _saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) + blocked = c.put("/embedding-model", json = {"embedding_model": "attacker/malicious-embed"}) + assert blocked.status_code == 403 + + # A verification failure (not-an-embedding-model) stays forceable at 409. + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) + monkeypatch.setattr(settings, "is_embedding_model", lambda *a, **k: False, raising = False) + import utils.models as _models + + monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: False) + unverified = c.put("/embedding-model", json = {"embedding_model": "acme/not-an-embedder"}) + assert unverified.status_code == 409 + + +def test_llama_backend_skips_the_st_pickle_scan(monkeypatch): + # On the llama-server backend the embedder loads GGUF (inert), not the ST repo's + # pickle, so a flagged ST repo with a clean GGUF companion must not be rejected here. + saved: dict = {} + monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) + monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + monkeypatch.setattr(settings, "_llama_backend_active", lambda: True) + monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) + monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) + monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + # force skips the GGUF availability checks; the ST pickle gate is what we assert is skipped. + called = {"scanned": False} + mod = _types.ModuleType("utils.security") + + def _fail(*a, **k): + called["scanned"] = True + return _Decision(True) + + mod.evaluate_file_security = _fail + mod.security_load_subdirs = lambda *a, **k: () + monkeypatch.setitem(sys.modules, "utils.security", mod) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + c = TestClient(app, raise_server_exceptions = False) + r = c.put( + "/embedding-model", + json = {"embedding_model": "attacker/flagged-st-clean-gguf", "force": True}, + ) + assert r.status_code == 200 + assert called["scanned"] is False # the ST pickle scan never ran on the llama path + assert saved.get("model") == "attacker/flagged-st-clean-gguf" + + +def test_runtime_llama_fallback_skips_the_st_pickle_scan(monkeypatch): + # auto resolves to sentence-transformers (GPU present) but the embedder fell back to + # llama-server at runtime (torch/CUDA load or encode failure), so the process now loads + # only inert GGUF. The real _llama_backend_active() must reflect that cached fallback, + # so a flagged ST repo with a clean GGUF companion must not be hard-blocked here. + import core.rag.embeddings as embeddings + from core.rag.embed_llama_server import LlamaServerBackend + + # Simulate the runtime fallback: the process-wide backend is a LlamaServerBackend even + # though the auto resolver would still say sentence-transformers. + monkeypatch.setattr(embeddings, "_backend", LlamaServerBackend()) + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers") + monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ()) + + saved: dict = {} + monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) + monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + # Deliberately do NOT monkeypatch settings._llama_backend_active: this test exercises the + # real delegation to embeddings.active_backend_is_llama() so the cached fallback is honored. + monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) + monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) + monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + + called = {"scanned": False} + mod = _types.ModuleType("utils.security") + + def _fail(*a, **k): + called["scanned"] = True + return _Decision(True) + + mod.evaluate_file_security = _fail + mod.security_load_subdirs = lambda *a, **k: () + monkeypatch.setitem(sys.modules, "utils.security", mod) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + c = TestClient(app, raise_server_exceptions = False) + r = c.put( + "/embedding-model", + json = {"embedding_model": "attacker/flagged-st-clean-gguf", "force": True}, + ) + assert r.status_code == 200 + assert called["scanned"] is False # the ST pickle scan never ran on the llama fallback + assert saved.get("model") == "attacker/flagged-st-clean-gguf" + + +def test_active_backend_is_llama_reflects_cache_and_resolver(monkeypatch): + # active_backend_is_llama() reports the ACTUAL built backend when one exists, and defers + # to the resolver (fresh-process behavior) when none has been built yet. + import core.rag.embeddings as embeddings + import core.rag.config as rag_config + from core.rag.embed_llama_server import LlamaServerBackend + + # A cached llama backend wins even when auto would resolve to sentence-transformers. + monkeypatch.setattr(rag_config, "EMBED_BACKEND", "auto") + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers") + monkeypatch.setattr(embeddings, "_backend", LlamaServerBackend()) + assert embeddings.active_backend_is_llama() is True + + # A cached ST backend reports False even when the resolver now picks llama, so its + # pickle stays gated (the cached backend, not the resolver, is what actually embeds). + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "llama-server") + monkeypatch.setattr(embeddings, "_backend", embeddings._SentenceTransformersBackend()) + assert embeddings.active_backend_is_llama() is False + + # No cached backend -> the resolver decides, unchanged from before. + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers") + monkeypatch.setattr(embeddings, "_backend", None) + assert embeddings.active_backend_is_llama() is False # auto -> sentence-transformers + + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "llama-server") + assert embeddings.active_backend_is_llama() is True # auto -> llama-server + + # An explicit (non-auto) key is honored verbatim without a cached backend. + monkeypatch.setattr(rag_config, "EMBED_BACKEND", "llama-server") + assert embeddings.active_backend_is_llama() is True + + +def test_settings_scan_scopes_module_subdirs(monkeypatch): + # The settings scan must pass the ST module dirs (0_Transformer/) as load roots so a + # pickle directly under one blocks; assert those subdirs reach evaluate_file_security. + saved: dict = {} + monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) + monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + monkeypatch.setattr(settings, "_llama_backend_active", lambda: False) + monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) + monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) + monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + + import core.rag.embeddings as embeddings + + monkeypatch.setattr( + embeddings, "_st_module_subdirs", lambda name, token = None: ("0_Transformer",) + ) + seen = {} + + def _capture(*a, **k): + seen["subdirs"] = tuple(k.get("load_subdirs") or ()) + return _Decision(False) + + mod = _types.ModuleType("utils.security") + mod.security_load_subdirs = lambda *a, **k: () + mod.evaluate_file_security = _capture + monkeypatch.setitem(sys.modules, "utils.security", mod) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + c = TestClient(app, raise_server_exceptions = False) + r = c.put( + "/embedding-model", json = {"embedding_model": "acme/embed-with-module-dir", "force": True} + ) + assert r.status_code == 200 + assert "0_Transformer" in seen["subdirs"] + + +def test_clean_repo_saves_under_force(client, monkeypatch): + c, saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) + r = c.put("/embedding-model", json = {"embedding_model": "acme/clean-embed", "force": True}) + assert r.status_code == 200 + assert saved.get("model") == "acme/clean-embed" + + +def test_load_sink_refuses_flagged_model(monkeypatch): + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) + import core.rag.embeddings as embeddings + with pytest.raises(embeddings.UnsafeEmbeddingModelError): + embeddings._guard_model_security("attacker/malicious-embed") + + +def test_load_sink_allows_clean_model(monkeypatch): + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) + import core.rag.embeddings as embeddings + embeddings._guard_model_security("acme/clean-embed") # no raise + + +def test_sink_threads_ambient_token_into_scan(monkeypatch): + # A gated repo set via env/default has no request token; the guard must feed the + # loader's own token to the scan, or it fails open for the repo that still loads. + seen = {} + mod = _types.ModuleType("utils.security") + mod.security_load_subdirs = ( + lambda name, token = None: seen.setdefault("subdirs_token", token) or () + ) + mod.evaluate_file_security = lambda *a, **k: seen.setdefault( + "scan_token", k.get("hf_token") + ) or _Decision(False) + monkeypatch.setitem(sys.modules, "utils.security", mod) + import core.rag.embeddings as embeddings + + monkeypatch.setattr(embeddings, "_ambient_hf_token", lambda: "hf_ambient") + embeddings._guard_model_security("acme/gated-embed") + assert seen["scan_token"] == "hf_ambient" + assert seen["subdirs_token"] == "hf_ambient" + + +def test_sink_scopes_st_module_subdirs_into_scan(monkeypatch): + # A flagged pickle directly under a Transformer module dir (0_Transformer/) must + # reach the scan as a load root; assert the guard unions the module dirs into + # load_subdirs so evaluate_file_security treats such a pickle as root-level. + seen = {} + + def _capture(*a, **k): + seen["subdirs"] = tuple(k.get("load_subdirs") or ()) + return _Decision(False) + + mod = _types.ModuleType("utils.security") + mod.security_load_subdirs = lambda name, token = None: () + mod.evaluate_file_security = _capture + monkeypatch.setitem(sys.modules, "utils.security", mod) + import core.rag.embeddings as embeddings + + monkeypatch.setattr(embeddings, "_ambient_hf_token", lambda: None) + monkeypatch.setattr( + embeddings, "_st_module_subdirs", lambda name, token = None: ("0_Transformer",) + ) + embeddings._guard_model_security("acme/embed-with-module-dir") + assert "0_Transformer" in seen["subdirs"] + + +def test_st_module_subdirs_reads_local_modules_json(tmp_path, monkeypatch): + # The helper must parse each module's non-empty "path" from a local repo's + # modules.json and drop the root-level ("") Transformer entry. + import json + import core.rag.embeddings as embeddings + + (tmp_path / "modules.json").write_text( + json.dumps( + [ + {"idx": 0, "name": "0", "path": "0_Transformer", "type": "..."}, + {"idx": 1, "name": "1", "path": "1_Pooling", "type": "..."}, + {"idx": 2, "name": "2", "path": "", "type": "..."}, + ] + ) + ) + subdirs = embeddings._st_module_subdirs(str(tmp_path), None) + assert subdirs == ("0_Transformer", "1_Pooling") + + +def test_st_module_subdirs_swallows_errors(monkeypatch): + # Any failure (no modules.json, offline, malformed) returns () so the guard never + # bricks the embedder. + import huggingface_hub + import core.rag.embeddings as embeddings + + def _boom(*a, **k): + raise RuntimeError("offline") + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", _boom) + assert embeddings._st_module_subdirs("acme/no-such-repo-xyz", None) == () + + +def test_security_block_is_not_swallowed_by_llama_fallback(monkeypatch): + # The ST encode fallback must re-raise a security block, not swap to llama-server. + import core.rag.embeddings as embeddings + + def _boom(*a, **k): + raise embeddings.UnsafeEmbeddingModelError("flagged") + + monkeypatch.setattr(embeddings, "_st_encode", _boom) + monkeypatch.setattr( + embeddings, + "_switch_to_llama_fallback", + lambda err: pytest.fail("security block must not fall back to llama-server"), + ) + with pytest.raises(embeddings.UnsafeEmbeddingModelError): + embeddings._SentenceTransformersBackend().encode(["hi"]) diff --git a/studio/backend/tests/test_security_gate_consistency.py b/studio/backend/tests/test_security_gate_consistency.py index db66df8a30..b5f1069f12 100644 --- a/studio/backend/tests/test_security_gate_consistency.py +++ b/studio/backend/tests/test_security_gate_consistency.py @@ -99,3 +99,16 @@ def test_malware_and_consent_gates_cover_the_lora_base(): if runs_gate and not resolves_base: offenders.append(f"{rel} runs a load gate but never resolves the LoRA base") assert not offenders, "\n".join(offenders) + + +def test_rag_embedding_path_runs_the_malware_gate(): + """The RAG embedding model is set through /settings and later loaded by + SentenceTransformer, which deserializes pickles; both sites must run the malware gate + or a flagged repo loads unscanned (bypassing the normal model-load protections).""" + offenders = [] + for rel in ("routes/settings.py", "core/rag/embeddings.py"): + if "evaluate_file_security(" not in (_BACKEND / rel).read_text(): + offenders.append( + f"{rel} loads/persists an embedding model without evaluate_file_security" + ) + assert not offenders, "\n".join(offenders) diff --git a/studio/frontend/src/features/settings/api/embedding-model.ts b/studio/frontend/src/features/settings/api/embedding-model.ts index 8b6bc7ee7f..9a61142f73 100644 --- a/studio/frontend/src/features/settings/api/embedding-model.ts +++ b/studio/frontend/src/features/settings/api/embedding-model.ts @@ -23,6 +23,10 @@ type ApiEmbeddingModelSettings = { * (wrong type, gated repo, or offline). Retry with force to save anyway. */ export class EmbeddingModelVerificationError extends Error {} +/** 403 from the backend: the repo is flagged unsafe by Hugging Face's security scan. + * A hard block; force cannot bypass it, so it must not enter the "save anyway" flow. */ +export class EmbeddingModelBlockedError extends Error {} + function fromApi(settings: ApiEmbeddingModelSettings): EmbeddingModelSettings { return { embeddingModel: settings.embedding_model, @@ -56,6 +60,11 @@ export async function updateEmbeddingModelSettings( force: options?.force ?? false, }), }); + if (res.status === 403) { + throw new EmbeddingModelBlockedError( + await readFastApiError(res, "This model is blocked by a security scan"), + ); + } if (res.status === 409) { throw new EmbeddingModelVerificationError( await readFastApiError(res, "Could not verify the embedding model"), diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index 7670aae5fa..8fd70cc1b7 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -42,6 +42,7 @@ import { updatePreviewSharing, } from "../api/preview-sharing"; import { + EmbeddingModelBlockedError, type EmbeddingModelSettings, EmbeddingModelVerificationError, loadEmbeddingModelSettings, @@ -410,7 +411,10 @@ export function GeneralTab() { description: t("settings.general.rag.reindexWarning"), }); } catch (error) { - if (error instanceof EmbeddingModelVerificationError) { + // A hard security block cannot be forced; keep the "save anyway" action hidden. + if (error instanceof EmbeddingModelBlockedError) { + setEmbeddingModelNeedsForce(false); + } else if (error instanceof EmbeddingModelVerificationError) { setEmbeddingModelNeedsForce(true); } setEmbeddingModelError( From d79495dc96cf7c6ab0d60585e232fa381549c14e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 04:41:37 -0700 Subject: [PATCH 3/7] Add RDNA 2/3/4 ROCm routing tests via a CPU-only torch spoof (#6935) * Add RDNA 2/3/4 ROCm routing tests via a CPU-only torch spoof Introduces tests/_zoo_rocm_spoof.py, the ROCm sibling of _zoo_aggressive_cuda_spoof.py: it reuses the CUDA spoof's torch.cuda no-op machinery and overlays an AMD Radeon identity (torch.version.hip, gcnArchName, capability) for any RDNA 2/3/4 gfx target, so hip code paths run on CPU-only CI with no AMD hardware. tests/studio/install/test_rocm_rdna_routing.py then asserts unsloth_zoo routes every RDNA arch (gfx1030/1031/1032/1034, gfx1100/1101/1102, gfx1150/1151, gfx1200/1201) correctly: device_type resolves to hip, llama.cpp target resolves to (rocm, gfx), and the per-family ROCm bundle suffix (gfx103X/gfx110X/gfx120X, or self for gfx1150/1151) is picked. The torch-facing checks run in a subprocess so the spoof never leaks into sibling tests and DEVICE_TYPE (cached at import) resolves from a clean process; the pure gfx-family mapping runs in-process. Guarded by importorskip so it runs where torch and unsloth_zoo are installed (the Repo tests CPU job) and skips elsewhere. * [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> --- tests/_zoo_rocm_spoof.py | 84 +++++++++++++++++++ .../studio/install/test_rocm_rdna_routing.py | 84 +++++++++++++++++++ 2 files changed, 168 insertions(+) create mode 100644 tests/_zoo_rocm_spoof.py create mode 100644 tests/studio/install/test_rocm_rdna_routing.py diff --git a/tests/_zoo_rocm_spoof.py b/tests/_zoo_rocm_spoof.py new file mode 100644 index 0000000000..050191e9d1 --- /dev/null +++ b/tests/_zoo_rocm_spoof.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. +"""ROCm/RDNA spoof: present torch as an AMD Radeon (RDNA 2/3/4) card on a +GPU-less host, so hip paths (device_type -> "hip", llama.cpp ROCm bundle) are +testable in CPU-only CI with no AMD hardware. The ROCm sibling of +_zoo_aggressive_cuda_spoof.py: it reuses that spoof's torch.cuda no-op machinery +and overlays the AMD identity (torch.version.hip, gcnArchName, Radeon name). +Apply BEFORE importing unsloth/unsloth_zoo, since DEVICE_TYPE is cached there. +""" + +from __future__ import annotations + +import importlib.util +import os +import sys + +# gfx -> (marketing name, (capability major, minor), torch.version.hip). hip is +# the ROCm build torch was made against (RDNA2/3 ship 6.x; gfx1102/115x/RDNA4 7.2). +_PROFILES: dict[str, tuple[str, tuple[int, int], str]] = { + "gfx1030": ("AMD Radeon RX 6900 XT", (10, 3), "6.4.43483"), # RDNA2 + "gfx1031": ("AMD Radeon RX 6700 XT", (10, 3), "6.4.43483"), + "gfx1032": ("AMD Radeon RX 6600", (10, 3), "6.4.43483"), + "gfx1034": ("AMD Radeon RX 6400", (10, 3), "6.4.43483"), + "gfx1100": ("AMD Radeon RX 7900 XTX", (11, 0), "6.4.43483"), # RDNA3 + "gfx1101": ("AMD Radeon RX 7800 XT", (11, 0), "6.4.43483"), + "gfx1102": ("AMD Radeon RX 7600", (11, 0), "7.2.1"), + "gfx1150": ("AMD Radeon 890M", (11, 5), "7.2.1"), # RDNA3.5 APU + "gfx1151": ("AMD Radeon 8060S", (11, 5), "7.2.1"), + "gfx1200": ("AMD Radeon RX 9060 XT", (12, 0), "7.2.1"), # RDNA4 + "gfx1201": ("AMD Radeon RX 9070 XT", (12, 0), "7.2.1"), +} + + +def _cuda_spoof(): + """Load the sibling CUDA spoof by path (robust to sys.path), so we reuse its + torch.cuda machinery instead of duplicating it.""" + if "_zoo_aggressive_cuda_spoof" in sys.modules: + return sys.modules["_zoo_aggressive_cuda_spoof"] + path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "_zoo_aggressive_cuda_spoof.py") + spec = importlib.util.spec_from_file_location("_zoo_aggressive_cuda_spoof", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + sys.modules["_zoo_aggressive_cuda_spoof"] = mod + return mod + + +def apply(gfx: str = "gfx1100", device_count: int = 1) -> None: + """Present torch as `gfx`. Re-callable to switch arch (identity is overlaid; + the underlying no-op machinery is applied once).""" + import torch + + if gfx not in _PROFILES: + raise KeyError(f"Unknown gfx {gfx!r}; known: {', '.join(_PROFILES)}") + name, cap, hip = _PROFILES[gfx] + + _cuda_spoof().apply() # is_available/device_count/streams/rng/amp/... + + # Overlay the AMD identity on top of the (NVIDIA-shaped) CUDA spoof. + torch.version.hip = hip + torch.version.cuda = None + torch.cuda.device_count = lambda: device_count + torch.cuda.get_device_name = lambda *a, **k: name + torch.cuda.get_device_capability = lambda *a, **k: cap + torch.cuda.get_arch_list = lambda: [gfx] + + class _Props: + pass + + _p = _Props() + _p.name = name + _p.gcnArchName = f"{gfx}:sramecc-:xnack-" # ROCm advertises feature flags + _p.major, _p.minor = cap + _p.total_memory = 16 * 1024**3 + _p.multi_processor_count = 40 + _p.warp_size = 32 # RDNA wavefront (CDNA is 64) + _p.is_integrated = gfx in ("gfx1150", "gfx1151") + _p.is_multi_gpu_board = False + torch.cuda.get_device_properties = lambda *a, **k: _p + + +if __name__ == "__main__": + apply() + import torch + print("ROCm spoof applied:", torch.version.hip, torch.cuda.get_device_properties(0).gcnArchName) diff --git a/tests/studio/install/test_rocm_rdna_routing.py b/tests/studio/install/test_rocm_rdna_routing.py new file mode 100644 index 0000000000..b4aeafb7e4 --- /dev/null +++ b/tests/studio/install/test_rocm_rdna_routing.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. +"""RDNA 2/3/4 routing, validated on CPU-only CI with no AMD hardware. + +tests/_zoo_rocm_spoof.py presents torch as each Radeon gfx arch, then we assert +unsloth_zoo routes it: device_type -> "hip", llama.cpp target -> ("rocm", gfx), +and the per-family ROCm bundle suffix. The torch-facing checks run in a +subprocess so the spoof never leaks into sibling tests and DEVICE_TYPE (cached +at import) resolves from a clean process. +""" + +from __future__ import annotations + +import json +import subprocess +import sys +from pathlib import Path + +import pytest + +pytest.importorskip("torch") +pytest.importorskip("unsloth_zoo") + +_TESTS_DIR = Path(__file__).resolve().parents[2] # tests/ + +# gfx -> (expected llama.cpp target, expected ROCm bundle family). +_ARCHES = { + "gfx1030": (("rocm", "gfx1030"), "gfx103X"), # RDNA2 + "gfx1031": (("rocm", "gfx1031"), "gfx103X"), + "gfx1032": (("rocm", "gfx1032"), "gfx103X"), + "gfx1034": (("rocm", "gfx1034"), "gfx103X"), + "gfx1100": (("rocm", "gfx1100"), "gfx110X"), # RDNA3 + "gfx1101": (("rocm", "gfx1101"), "gfx110X"), + "gfx1102": (("rocm", "gfx1102"), "gfx110X"), + "gfx1150": (("rocm", "gfx1150"), "gfx1150"), # RDNA3.5 APU (self-family) + "gfx1151": (("rocm", "gfx1151"), "gfx1151"), + "gfx1200": (("rocm", "gfx1200"), "gfx120X"), # RDNA4 + "gfx1201": (("rocm", "gfx1201"), "gfx120X"), +} + +# Child: spoof each arch, then record device_type once (fresh import) and the +# live llama.cpp target per arch. Emits one JSON line the parent parses. +_CHILD = """ +import json, sys +sys.path.insert(0, {tests!r}) +import _zoo_rocm_spoof as spoof +arches = {arches!r} +spoof.apply(arches[0]) +from unsloth_zoo.device_type import get_device_type, is_hip +device_type = [get_device_type(), is_hip()] +from unsloth_zoo import llama_cpp as lc +targets = {{}} +for gfx in arches: + spoof.apply(gfx) + targets[gfx] = list(lc._detect_gpu_target()) +print("RESULT " + json.dumps({{"device_type": device_type, "targets": targets}})) +""" + + +@pytest.fixture(scope = "module") +def routed(): + code = _CHILD.format(tests = str(_TESTS_DIR), arches = list(_ARCHES)) + proc = subprocess.run([sys.executable, "-c", code], capture_output = True, text = True) + line = next((l for l in proc.stdout.splitlines() if l.startswith("RESULT ")), None) + assert line, f"child produced no result.\nstdout:\n{proc.stdout}\nstderr:\n{proc.stderr}" + return json.loads(line[len("RESULT ") :]) + + +@pytest.mark.parametrize("gfx", list(_ARCHES)) +def test_detect_gpu_target(routed, gfx): + # RDNA card is routed to its ROCm gfx target (drives the llama.cpp bundle). + assert tuple(routed["targets"][gfx]) == _ARCHES[gfx][0] + + +def test_device_type_is_hip(routed): + # An RDNA card must resolve the compute device_type to "hip". + assert routed["device_type"] == ["hip", True] + + +@pytest.mark.parametrize("gfx", list(_ARCHES)) +def test_rocm_gfx_family(gfx): + # Pure mapping (no torch): each gfx picks the right per-family ROCm bundle. + from unsloth_zoo import llama_cpp as lc + assert lc._rocm_gfx_family(gfx) == _ARCHES[gfx][1] From 59977f95c318c1ba81b53c5b50d4a0ef9c342fea Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 05:49:24 -0700 Subject: [PATCH 4/7] GRPO: default router_aux_loss_coef to 0 on TRL >= 1.7.0 (#6938) TRL 1.7.0 enables the MoE router load-balancing aux loss by default (router_aux_loss_coef = 0.001). Unsloth's optimized GRPO forward does not compute it, so default the coefficient to 0, matching pre-1.7.0 behaviour. Users can still opt in with router_aux_loss_coef > 0. No-op on TRL < 1.7.0. --- unsloth/models/rl.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 62ef9e916a..b5cadf2dea 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -1370,6 +1370,9 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): # [TODO] See https://fengyao.notion.site/off-policy-rl # https://github.com/huggingface/trl/pull/3867 (August 7th) "vllm_importance_sampling_correction": False, + # TRL >= 1.7.0 enables the MoE router aux loss by default (0.001); the optimized + # GRPO forward does not compute it, so default off. Opt in via router_aux_loss_coef > 0. + "router_aux_loss_coef": 0.0, } for k, v in replacements.items(): x = f"{k}( = [^,\n]{{1,}})?,\n" From 411c4d1e50362c0fe6c27d4eea0c29171f79a27a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 06:13:43 -0700 Subject: [PATCH 5/7] Add DeepSeek-V4-Flash-GGUF to Studio with none/high/max reasoning (#6908) * Add DeepSeek-V4-Flash-GGUF to Studio with none/high/max reasoning Adds unsloth/DeepSeek-V4-Flash-GGUF as a default selectable model with the recommended decoding defaults (temperature 1.0, top_p 1.0 from the official generation_config.json) and its three tier reasoning control. The high/max ladder is surfaced for deepseek-v4 model ids and flows through the existing enable_thinking_effort reasoning style via chat_template_kwargs, so no frontend changes are needed. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio DeepSeek-V4: segment-scope high, enable thinking for lone effort, render tests Match deepseek-v4 on whole repo-name segments so a future deepseek-v40 or deepseek40 cannot false-match the synthetic 'high'. In _request_reasoning_kwargs, emit enable_thinking when a named effort level is sent without it, so the newly exposed High mode renders thinking-on over the API (the UI already sent it explicitly). Add a none/high/max render-path test file (jinja behind importorskip) with a lone-high regression. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: danielhanchen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../assets/configs/inference_defaults.json | 9 +- studio/backend/core/inference/defaults.py | 2 + studio/backend/core/inference/llama_cpp.py | 18 +- .../tests/test_deepseek_v4_thinking_effort.py | 181 ++++++++++++++++++ .../test_safetensors_capability_advertise.py | 38 ++++ 5 files changed, 245 insertions(+), 3 deletions(-) create mode 100644 studio/backend/tests/test_deepseek_v4_thinking_effort.py diff --git a/studio/backend/assets/configs/inference_defaults.json b/studio/backend/assets/configs/inference_defaults.json index 1c7a409bc1..0633f80bbc 100644 --- a/studio/backend/assets/configs/inference_defaults.json +++ b/studio/backend/assets/configs/inference_defaults.json @@ -235,6 +235,13 @@ "min_p": 0.1, "repetition_penalty": 1.0 }, + "deepseek-v4": { + "temperature": 1.0, + "top_p": 1.0, + "top_k": -1, + "min_p": 0.0, + "repetition_penalty": 1.0 + }, "deepseek-r1": { "temperature": 0.6, "top_p": 0.95, @@ -394,7 +401,7 @@ "phi-4", "phi-3", "mistral-nemo", "mistral-small", "mistral-large", "magistral", "ministral", "devstral", "pixtral", - "deepseek-r1", "deepseek-v3", "deepseek-ocr", + "deepseek-v4", "deepseek-r1", "deepseek-v3", "deepseek-ocr", "glm-5", "glm-4", "nemotron", "minimax-m2.7", "minimax-m2.5", "minimax", diff --git a/studio/backend/core/inference/defaults.py b/studio/backend/core/inference/defaults.py index b64605e16f..a1d03c03e0 100644 --- a/studio/backend/core/inference/defaults.py +++ b/studio/backend/core/inference/defaults.py @@ -8,6 +8,7 @@ import utils.hardware.hardware as hw DEFAULT_MODELS_GGUF = [ "unsloth/Qwen3.6-27B-MTP-GGUF", "unsloth/Qwen3.6-35B-A3B-MTP-GGUF", + "unsloth/DeepSeek-V4-Flash-GGUF", "unsloth/gemma-4-E2B-it-GGUF", "unsloth/gemma-4-E4B-it-GGUF", "unsloth/gemma-4-31B-it-GGUF", @@ -27,6 +28,7 @@ DEFAULT_MODELS_GGUF = [ DEFAULT_MODELS_STANDARD = [ "unsloth/Qwen3.6-27B-MTP-GGUF", "unsloth/Qwen3.6-35B-A3B-MTP-GGUF", + "unsloth/DeepSeek-V4-Flash-GGUF", "unsloth/gemma-4-E2B-it-GGUF", "unsloth/gemma-4-E4B-it-GGUF", "unsloth/gemma-4-31B-it-GGUF", diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 5e6287f528..8b40f5fccd 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -686,6 +686,16 @@ def detect_reasoning_flags( else [] ) if effort_levels: + # DeepSeek-V4's encoder accepts reasoning_effort {'high', 'max'} but its + # template only branches on 'max', so the literal scan misses 'high'. Add it + # (matched on whole repo-name segments, so 'deepseek-v40' won't false-match) + # to expose the full none/high/max ladder instead of none/max. + segments = re.split(r"[-_.]", (model_identifier or "").lower().split("/")[-1]) + is_dsv4 = "deepseek4" in segments or any( + a == "deepseek" and b == "v4" for a, b in zip(segments, segments[1:]) + ) + if is_dsv4 and "high" not in effort_levels: + effort_levels = sorted(set(effort_levels) | {"high"}, key = _REASONING_EFFORT_SCALE.index) # GLM-5.2-style: an enable_thinking on/off gate PLUS a reasoning_effort # level among a discrete set (e.g. 'high' | 'max'). Distinct from # gpt-oss (reasoning_effort only, no on/off gate) and Qwen @@ -1741,9 +1751,13 @@ class LlamaCppBackend: # 'low' effort the way gpt-oss does (those models genuinely # cannot disable). thinking_off = enable_thinking is False or reasoning_effort == "none" - if enable_thinking is not None or reasoning_effort == "none": + # A named effort level implies thinking on, so emit enable_thinking + # even if the caller sent only reasoning_effort (else the template + # defaults it off and the requested level never renders). + effort_on = reasoning_effort in self._reasoning_effort_levels + if enable_thinking is not None or reasoning_effort == "none" or effort_on: kwargs["enable_thinking"] = not thinking_off - if not thinking_off and reasoning_effort in self._reasoning_effort_levels: + if not thinking_off and effort_on: kwargs["reasoning_effort"] = reasoning_effort elif self._reasoning_style == "reasoning_effort": if reasoning_effort in ("none", "low", "medium", "high"): diff --git a/studio/backend/tests/test_deepseek_v4_thinking_effort.py b/studio/backend/tests/test_deepseek_v4_thinking_effort.py new file mode 100644 index 0000000000..19808ad0d7 --- /dev/null +++ b/studio/backend/tests/test_deepseek_v4_thinking_effort.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""DeepSeek-V4-Flash reasoning toggle: None / High / Max. + +The GGUF template gates thinking with ``enable_thinking`` and only branches +``reasoning_effort`` on ``'max'`` (an escalation layered over plain thinking). +Detection used to return the single level ``['max']``, so the UI collapsed to +None / Max and the plain-thinking tier was unreachable. Detection now surfaces +``'high'`` as that plain tier, giving None / High / Max. These tests pin the +classifier, the GLM-style parity case, and the full request-kwargs -> rendered +prompt path for each state (the model itself is too large to load here). +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +_backend_root = Path(__file__).resolve().parent.parent +if str(_backend_root) not in sys.path: + sys.path.insert(0, str(_backend_root)) + + +# Faithful slice of the DeepSeek-V4-Flash GGUF template: the enable_thinking +# gate, the sole ``reasoning_effort == 'max'`` escalation, and the plain-think +# fallback. Any non-'max' effort renders as ordinary thinking. +DEEPSEEK_V4_TEMPLATE = """ +{%- if not thinking is defined -%} + {%- if enable_thinking is defined -%} + {%- set thinking = enable_thinking -%} + {%- else -%} + {%- set thinking = false -%} + {%- endif -%} +{%- endif -%} +{%- if not reasoning_effort is defined -%} + {%- set reasoning_effort = none -%} +{%- endif -%} +{{- bos_token -}} +{%- if thinking and reasoning_effort == 'max' -%} + {{- 'Reasoning Effort: Absolute maximum with no shortcuts permitted.\\n\\n' -}} +{%- endif -%} +{%- for message in messages -%} + {{- '<|User|>' + (message['content'] or '') -}} +{%- endfor -%} +{%- if add_generation_prompt -%} + {{- '<|Assistant|>' -}} + {%- if thinking -%}{{- '' -}}{%- else -%}{{- '' -}}{%- endif -%} +{%- endif -%} +""" + + +# GLM-5.2-style: branches on two effort literals, so 'high' already exists as +# the sub-'max' tier and detection must leave the pair untouched. +GLM_STYLE_TEMPLATE = """ +{%- if enable_thinking -%} + {%- if reasoning_effort == 'high' -%}{{- 'H' -}} + {%- elif reasoning_effort == 'max' -%}{{- 'M' -}} + {%- endif -%} +{%- endif -%} +""" + + +# A ['max']-only template under a non-deepseek id: the synthetic 'high' is scoped +# to deepseek-v4, so this must stay ['max'] (no phantom 'high'). +NON_DEEPSEEK_MAX_ONLY_TEMPLATE = DEEPSEEK_V4_TEMPLATE + + +# A template whose sole effort literal is a sub-'max' level: the guard targets +# only the ['max']-alone case, so a lone 'high' stays a singleton. +HIGH_ONLY_TEMPLATE = """ +{%- if enable_thinking and reasoning_effort == 'high' -%}{{- 'H' -}}{%- endif -%} +""" + + +def _render(template: str, **kwargs) -> str: + jinja2 = pytest.importorskip("jinja2") + env = jinja2.Environment() + tmpl = env.from_string(template) + return tmpl.render(bos_token = "", add_generation_prompt = True, **kwargs) + + +# -- Classifier ------------------------------------------------------- + + +def test_deepseek_v4_surfaces_high_as_plain_tier(): + """Sole 'max' escalation expands to ['high', 'max'] so None/High/Max show.""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash") + assert flags["supports_reasoning"] is True + assert flags["reasoning_style"] == "enable_thinking_effort" + assert flags["reasoning_effort_levels"] == ["high", "max"] + + +def test_glm_style_two_level_template_unchanged(): + """A template that already names a sub-'max' tier is left as-is.""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(GLM_STYLE_TEMPLATE, "unsloth/GLM-5.2") + assert flags["reasoning_style"] == "enable_thinking_effort" + assert flags["reasoning_effort_levels"] == ["high", "max"] + + +def test_synthetic_high_scoped_to_deepseek_v4(): + """The same ['max']-only template under a non-deepseek id keeps ['max'].""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(NON_DEEPSEEK_MAX_ONLY_TEMPLATE, "vendor/OtherHybrid-GGUF") + assert flags["reasoning_effort_levels"] == ["max"] + + +def test_guard_does_not_fire_for_sub_max_singleton(): + """The expansion targets only ['max']; a lone 'high' stays a singleton.""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(HIGH_ONLY_TEMPLATE, "custom/high-only") + assert flags["reasoning_effort_levels"] == ["high"] + + +# -- Request kwargs -> rendered prompt, for each state ---------------- + + +def _kwargs_for(flags: dict, enable_thinking, reasoning_effort): + """Drive the real backend method with a shim carrying the detected flags.""" + from core.inference.llama_cpp import LlamaCppBackend + + shim = SimpleNamespace( + _supports_reasoning = flags["supports_reasoning"], + _reasoning_always_on = flags["reasoning_always_on"], + _reasoning_style = flags["reasoning_style"], + _reasoning_effort_levels = flags["reasoning_effort_levels"], + _supports_preserve_thinking = flags["supports_preserve_thinking"], + ) + build = LlamaCppBackend._request_reasoning_kwargs.__get__(shim) + return build(enable_thinking, reasoning_effort, None) or {} + + +def _flags(): + from core.inference.llama_cpp import detect_reasoning_flags + return detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash") + + +def test_none_state_renders_non_thinking(): + """UI 'None' -> enable_thinking=false -> closed , no preamble.""" + kwargs = _kwargs_for(_flags(), enable_thinking = False, reasoning_effort = None) + assert kwargs == {"enable_thinking": False} + out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + assert out.endswith("") + assert "Absolute maximum" not in out + + +def test_high_state_renders_plain_thinking(): + """UI 'High' -> et=true, effort=high -> open , no max preamble.""" + kwargs = _kwargs_for(_flags(), enable_thinking = True, reasoning_effort = "high") + assert kwargs == {"enable_thinking": True, "reasoning_effort": "high"} + out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + assert out.endswith("") + assert "Absolute maximum" not in out + + +def test_max_state_injects_max_preamble(): + """UI 'Max' -> et=true, effort=max -> open plus the max preamble.""" + kwargs = _kwargs_for(_flags(), enable_thinking = True, reasoning_effort = "max") + assert kwargs == {"enable_thinking": True, "reasoning_effort": "max"} + out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + assert out.endswith("") + assert "Absolute maximum" in out + + +def test_high_effort_alone_enables_thinking(): + """API caller sending only reasoning_effort='high' (no enable_thinking) still + gets thinking on, so the newly exposed High mode renders correctly.""" + kwargs = _kwargs_for(_flags(), enable_thinking = None, reasoning_effort = "high") + assert kwargs == {"enable_thinking": True, "reasoning_effort": "high"} + out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + assert out.endswith("") + assert "Absolute maximum" not in out diff --git a/studio/backend/tests/test_safetensors_capability_advertise.py b/studio/backend/tests/test_safetensors_capability_advertise.py index 9fd1535f22..0ed670ac01 100644 --- a/studio/backend/tests/test_safetensors_capability_advertise.py +++ b/studio/backend/tests/test_safetensors_capability_advertise.py @@ -48,6 +48,21 @@ reasoning_effort: {{ reasoning_effort }} """ +# DeepSeek-V4-Flash: an enable_thinking on/off gate PLUS a reasoning_effort +# 'max' preamble. The shipped template only *branches* on 'max' ('high' renders +# identically to thinking-on-without-the-preamble), so the literal scan alone +# would surface only ['max']; the classifier adds 'high' for deepseek-v4 to +# expose the encoder's full none/high/max ladder. +DEEPSEEK_V4_TEMPLATE = ( + "{%- if not thinking is defined %}" + "{%- if enable_thinking is defined %}{%- set thinking = enable_thinking %}" + "{%- else %}{%- set thinking = false %}{%- endif %}{%- endif %}\n" + "{%- if thinking and reasoning_effort == 'max' %}" + "{{- 'Reasoning Effort: Absolute maximum' }}{%- endif %}\n" + "{%- for message in messages %}{{- message.content }}{%- endfor %}" +) + + PLAIN_TEMPLATE = """ {%- for message in messages %} {{- message.role + ': ' + message.content + '\\n' }} @@ -90,6 +105,29 @@ def test_detect_reasoning_flags_none_template_returns_all_false(): assert flags["reasoning_style"] == "enable_thinking" +def test_detect_reasoning_flags_deepseek_v4_exposes_none_high_max(): + """DeepSeek-V4-Flash: enable_thinking gate + reasoning_effort 'max' preamble. + Classified as the hybrid style with the full none/high/max ladder even + though the template only branches on 'max'.""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash-GGUF") + assert flags["supports_reasoning"] is True + assert flags["reasoning_style"] == "enable_thinking_effort" + assert flags["reasoning_effort_levels"] == ["high", "max"] + assert flags["reasoning_always_on"] is False + + +def test_detect_reasoning_flags_non_deepseek_v4_effort_only_max_not_injected(): + """The 'high' injection is scoped to deepseek-v4: a different model whose + template only branches on 'max' keeps ['max'] (no phantom 'high').""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "vendor/OtherHybrid-GGUF") + assert flags["reasoning_style"] == "enable_thinking_effort" + assert flags["reasoning_effort_levels"] == ["max"] + + def test_detect_safetensors_features_passes_template_through_to_classifier(): """Route wrapper forwards a real template to the inner classifier.""" from routes.inference import _detect_safetensors_features From 10d8f985a270cd4cacf2518e9e895874da201f3e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 06:19:25 -0700 Subject: [PATCH 6/7] Versioning --- pyproject.toml | 6 +++--- unsloth/models/_utils.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 844ead2454..80b3d757e3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,7 +73,7 @@ triton = [ ] huggingfacenotorch = [ - "unsloth_zoo>=2026.6.7", + "unsloth_zoo>=2026.7.1", "wheel>=0.42.0", "packaging", "numpy", @@ -94,7 +94,7 @@ huggingfacenotorch = [ ] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.6.7", + "unsloth_zoo>=2026.7.1", "torchvision", "unsloth[triton]", ] @@ -579,7 +579,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2026.6.7", + "unsloth_zoo>=2026.7.1", "packaging", "tyro", "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0", diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 1aa2c6e820..1c75f8ce66 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "2026.6.9" +__version__ = "2026.7.1" __all__ = [ "SUPPORTS_BFLOAT16", From ba450b437eb34ee476df7ffb61bb25db365c7353 Mon Sep 17 00:00:00 2001 From: Etherll <61019402+Etherll@users.noreply.github.com> Date: Tue, 7 Jul 2026 16:26:00 +0300 Subject: [PATCH 7/7] Studio: add assistant response details panel (#6842) * Studio: add assistant response details panel * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Hide model badge by default, show on hover/focus Wrap MessageResponseModelBadge in a span with hidden/group-hover visibility classes to reduce visual clutter. The badge now only displays when hovering or focusing on the assistant message, improving the UI presentation. Updated corresponding tests to verify the new CSS classes. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../message-response-details-sheet.tsx | 483 ++++++++++++++++++ .../src/components/assistant-ui/reasoning.tsx | 10 +- .../src/components/assistant-ui/thread.tsx | 118 +++-- .../src/features/chat/api/chat-adapter.ts | 92 +++- studio/frontend/src/features/chat/index.ts | 7 +- .../chat/stores/chat-preferences-store.ts | 7 + .../src/features/settings/tabs/chat-tab.tsx | 15 + .../test_chat_response_details_ui_contract.py | 93 ++++ 8 files changed, 776 insertions(+), 49 deletions(-) create mode 100644 studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx create mode 100644 tests/studio/test_chat_response_details_ui_contract.py diff --git a/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx new file mode 100644 index 0000000000..823696693a --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx @@ -0,0 +1,483 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"use client"; + +import { + Sheet, + SheetContent, + SheetDescription, + SheetHeader, + SheetTitle, +} from "@/components/ui/sheet"; +import { + customProviderDisplayName, + parseExternalModelId, + useChatPreferencesStore, + useChatRuntimeStore, + useExternalProvidersStore, +} from "@/features/chat"; +import { cn } from "@/lib/utils"; +import { FileDatabaseIcon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { useMessage, useMessageTiming } from "@assistant-ui/react"; +import type { FC, ReactNode } from "react"; + +type ResponseDetailsMetadata = { + modelId?: string; + modelLabel?: string; + responseModelId?: string; + providerId?: string; + providerName?: string; + providerType?: string; + startedAt?: number; + finishedAt?: number; + durationMs?: number; + sessionId?: string | null; + cancelId?: string; + toolCalls?: string[]; + tools?: Record; +}; + +type ContextUsageMetadata = { + promptTokens?: number; + completionTokens?: number; + totalTokens?: number; + cachedTokens?: number; + cacheWriteTokens?: number; + modelId?: string; +}; + +type MessageCustomMetadata = { + responseDetails?: ResponseDetailsMetadata; + contextUsage?: ContextUsageMetadata; + serverTimings?: Record; + reasoningDuration?: number; +}; + +function asNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) + ? value + : undefined; +} + +function formatNumber(value: number | undefined): string | null { + return value == null ? null : value.toLocaleString(); +} + +function formatMs(value: number | undefined): string | null { + if (value == null) return null; + if (value < 1000) return `${Math.round(value)}ms`; + return `${(value / 1000).toFixed(2)}s`; +} + +function formatRate(value: number | undefined): string | null { + if (value == null) return null; + return `${value.toFixed(1)} tok/s`; +} + +function formatDate(value: Date | number | string | undefined): string | null { + if (value == null) return null; + const date = value instanceof Date ? value : new Date(value); + if (Number.isNaN(date.getTime())) return null; + return new Intl.DateTimeFormat(undefined, { + dateStyle: "medium", + timeStyle: "medium", + }).format(date); +} + +const TOOL_CATEGORY_LABELS: Record = { + search: "Search", + fetch: "Fetch", + code: "Code", + images: "Images", + mcp: "MCP", + docs: "Docs", + artifacts: "Canvas", +}; + +const TOOL_CALL_LABELS: Record = { + web_search: "Search", + web_fetch: "Fetch", + code_execution: "Code", + python: "Python", + terminal: "Terminal", + image_generation: "Images", + search_knowledge_base: "Docs", + render_html: "Canvas", +}; + +function uniqueValues(values: string[]): string[] { + return Array.from(new Set(values)); +} + +function toolCategoryFromCall(toolName: string): string | null { + const normalized = toolName.toLowerCase(); + if (normalized === "web_search") return "search"; + if (normalized === "web_fetch") return "fetch"; + if ( + normalized === "code_execution" || + normalized === "python" || + normalized === "terminal" + ) { + return "code"; + } + if (normalized === "image_generation") return "images"; + if (normalized === "search_knowledge_base") return "docs"; + if (normalized === "render_html") return "artifacts"; + if (normalized.startsWith("mcp__")) return "mcp"; + return null; +} + +function formatToolCallName(toolName: string): string { + const normalized = toolName.toLowerCase(); + if (TOOL_CALL_LABELS[normalized]) return TOOL_CALL_LABELS[normalized]; + if (normalized.startsWith("mcp__")) return `MCP: ${toolName.slice(5)}`; + return toolName + .replace(/[_-]+/g, " ") + .replace(/\b\w/g, (letter) => letter.toUpperCase()); +} + +function toolCallsFromContent(content: unknown): string[] { + if (!Array.isArray(content)) return []; + return uniqueValues( + content + .map((part) => + part && typeof part === "object" && "type" in part + ? (part as { type?: unknown; toolName?: unknown }) + : null, + ) + .filter( + (part): part is { type: "tool-call"; toolName: string } => + part?.type === "tool-call" && + typeof part.toolName === "string" && + part.toolName.length > 0, + ) + .map((part) => part.toolName), + ); +} + +function enabledTools( + tools: Record | undefined, + toolCalls: string[], +): string | null { + if (!tools && toolCalls.length === 0) return null; + const activeKeys = new Set(); + for (const key of Object.keys(TOOL_CATEGORY_LABELS)) { + if (tools?.[key] === true) activeKeys.add(key); + } + for (const toolName of toolCalls) { + const key = toolCategoryFromCall(toolName); + if (key) activeKeys.add(key); + } + const active = Object.keys(TOOL_CATEGORY_LABELS) + .filter((key) => activeKeys.has(key)) + .map((key) => TOOL_CATEGORY_LABELS[key]); + return active.length > 0 ? active.join(", ") : "None"; +} + +function calledTools(toolCalls: string[]): string | null { + if (toolCalls.length === 0) return null; + return uniqueValues(toolCalls.map(formatToolCallName)).join(", "); +} + +function DetailSection({ + title, + children, +}: { + title: string; + children: ReactNode; +}) { + return ( +
+

{title}

+
{children}
+
+ ); +} + +function DetailRow({ + label, + value, + mono = false, +}: { + label: string; + value: ReactNode | null | undefined; + mono?: boolean; +}) { + if (value == null || value === "") return null; + return ( +
+ {label} + + {value} + +
+ ); +} + +function useResponseModelDisplay() { + const message = useMessage(); + const models = useChatRuntimeStore((s) => s.models); + const providers = useExternalProvidersStore((s) => s.providers); + + const custom = ( + message.metadata as Record | undefined + )?.custom as MessageCustomMetadata | undefined; + const responseDetails = custom?.responseDetails; + const usage = custom?.contextUsage; + const serverTimings = custom?.serverTimings; + + const recordedModelId = + responseDetails?.responseModelId ?? + responseDetails?.modelId ?? + usage?.modelId; + const parsedExternal = parseExternalModelId(recordedModelId); + const provider = parsedExternal + ? providers.find((candidate) => candidate.id === parsedExternal.providerId) + : null; + const modelSummary = models.find( + (candidate) => candidate.id === recordedModelId, + ); + const modelLabel = + responseDetails?.modelLabel ?? + responseDetails?.responseModelId ?? + parsedExternal?.modelId ?? + modelSummary?.name ?? + recordedModelId ?? + "Not recorded"; + const providerLabel = + responseDetails?.providerName ?? + provider?.name ?? + (responseDetails?.providerType + ? customProviderDisplayName(responseDetails.providerType) + : parsedExternal + ? customProviderDisplayName(provider?.providerType) + : recordedModelId + ? "Local model" + : null); + + return { + message, + custom, + responseDetails, + usage, + serverTimings, + modelLabel, + providerLabel, + }; +} + +export const MessageResponseModelBadge: FC<{ className?: string }> = ({ + className, +}) => { + const showResponseModel = useChatPreferencesStore( + (state) => state.showResponseModel, + ); + const { modelLabel, providerLabel } = useResponseModelDisplay(); + + if (!showResponseModel || modelLabel === "Not recorded") { + return null; + } + + return ( + + {modelLabel} + + ); +}; + +export const MessageResponseDetailsSheet: FC<{ + open: boolean; + onOpenChange: (open: boolean) => void; +}> = ({ open, onOpenChange }) => { + const timing = useMessageTiming(); + const { + message, + responseDetails, + usage, + serverTimings, + modelLabel, + providerLabel, + } = useResponseModelDisplay(); + const promptTokens = + usage?.promptTokens ?? asNumber(serverTimings?.prompt_n); + const completionTokens = + usage?.completionTokens ?? + timing?.tokenCount ?? + asNumber(serverTimings?.predicted_n); + const totalTokens = + usage?.totalTokens ?? + (promptTokens != null && completionTokens != null + ? promptTokens + completionTokens + : undefined); + const totalTime = + responseDetails?.durationMs ?? timing?.totalStreamTime ?? undefined; + const summaryLabel = + modelLabel === "Not recorded" ? "Model not recorded" : `Used ${modelLabel}`; + const messageToolCalls = toolCallsFromContent(message.content); + const toolCalls = + responseDetails?.toolCalls && responseDetails.toolCalls.length > 0 + ? responseDetails.toolCalls + : messageToolCalls; + + return ( + + + + + + Response details + + + Timing, model, token, and tool details for this response. + + + +
+
+

+ {summaryLabel} +

+ {providerLabel ? ( +

+ {providerLabel} +

+ ) : null} +
+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
+
+
+ ); +}; diff --git a/studio/frontend/src/components/assistant-ui/reasoning.tsx b/studio/frontend/src/components/assistant-ui/reasoning.tsx index c5b53577cb..96d21d6fe7 100644 --- a/studio/frontend/src/components/assistant-ui/reasoning.tsx +++ b/studio/frontend/src/components/assistant-ui/reasoning.tsx @@ -6,6 +6,7 @@ /* eslint-disable react-refresh/only-export-components */ import { MarkdownText } from "@/components/assistant-ui/markdown-text"; +import { MessageResponseModelBadge } from "@/components/assistant-ui/message-response-details-sheet"; import { Collapsible, CollapsibleContent, @@ -390,14 +391,17 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({ onOpenChange={handleOpenChange} variant={variant} > -
+
-
+ + + +
{isOpen && !isReasoningStreaming && ( )} diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 7fe98b701f..09551cd413 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -12,6 +12,10 @@ import { import { downloadImagePart } from "@/components/assistant-ui/image"; import { MarkdownText } from "@/components/assistant-ui/markdown-text"; import { MessageHtmlArtifacts } from "@/components/assistant-ui/message-html-artifacts"; +import { + MessageResponseDetailsSheet, + MessageResponseModelBadge, +} from "@/components/assistant-ui/message-response-details-sheet"; import { MessageTiming } from "@/components/assistant-ui/message-timing"; import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning"; import { RagSourcesGroup } from "@/components/assistant-ui/rag-sources"; @@ -3564,6 +3568,9 @@ const AssistantMessage: FC = () => { const aui = useAui(); const messageId = useAuiState(({ message }) => message.id); const messageContent = useAuiState(({ message }) => message.content); + const hasReasoningParts = useAuiState(({ message }) => + message.parts.some((part) => part.type === "reasoning"), + ); const incognito = useChatRuntimeStore((s) => s.incognito); // Use global store for editing state to ensure a single source of truth @@ -3620,7 +3627,7 @@ const AssistantMessage: FC = () => { return (
@@ -3649,6 +3656,11 @@ const AssistantMessage: FC = () => {
) : ( <> + {!hasReasoningParts ? ( +
+ +
+ ) : null} @@ -3893,58 +3905,76 @@ const EditAssistantMessageButton: FC = () => { const AssistantActionBar: FC = () => { const { forkMessage, forkDisabled } = useForkMessageAction(); + const [detailsOpen, setDetailsOpen] = useState(false); return ( - - - - - - - - - - - - - - + <> + + + + + + - - e.preventDefault()} - className="aui-action-bar-more-content z-50 min-w-32 overflow-hidden rounded-[21px] bg-popover px-[9px] py-2 text-popover-foreground shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:shadow-none" - > - void forkMessage()} - className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground data-[disabled]:pointer-events-none data-[disabled]:opacity-50" + + + + + + + + + + e.preventDefault()} + className="aui-action-bar-more-content z-50 min-w-32 overflow-hidden rounded-[21px] bg-popover px-[9px] py-2 text-popover-foreground shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:shadow-none" > - - Fork in new chat - - - + void forkMessage()} + className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground data-[disabled]:pointer-events-none data-[disabled]:opacity-50" + > + + Fork in new chat + + + + + Export as Markdown + + + setDetailsOpen(true)} + className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground" + > - Export as Markdown + See response details - - - - - + + + + + + ); }; diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index df266fd749..2ca6ceddf7 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -140,6 +140,32 @@ interface ServerTimings { diffusion_steps_per_second?: number; } +interface ResponseDetailsMetadata { + modelId: string; + modelLabel: string; + responseModelId: string; + providerId?: string; + providerName: string; + providerType: string; + startedAt: number; + finishedAt: number; + durationMs: number; + sessionId?: string; + cancelId: string; + toolCalls: string[]; + tools: { + search: boolean; + fetch: boolean; + code: boolean; + images: boolean; + mcp: boolean; + docs: boolean; + artifacts: boolean; + confirmToolCalls: boolean; + bypassPermissions: boolean; + }; +} + type RunMessages = Parameters[0]["messages"]; type RunMessage = RunMessages[number]; @@ -1769,6 +1795,9 @@ export function createOpenAIStreamAdapter( (provider) => provider.id === externalSelection.providerId, ) : null; + const selectedModelSummary = runtime.models.find( + (model) => model.id === params.checkpoint, + ); const externalApiKey = externalProvider ? getExternalProviderApiKey(externalProvider.id).trim() : ""; @@ -2151,6 +2180,7 @@ export function createOpenAIStreamAdapter( let waitingFirstChunk = true; let firstTokenSettled = false; const streamStartTime = Date.now(); + let responseModelId = externalSelection?.modelId ?? params.checkpoint; let firstTokenTime: number | undefined; let totalChunks = 0; let resolveFirstToken: (() => void) | null = null; @@ -2372,6 +2402,59 @@ export function createOpenAIStreamAdapter( const externalBackendProviderType = toExternalBackendProviderType( externalProvider?.providerType, ); + const buildResponseDetails = ( + finishedAt: number, + ): ResponseDetailsMetadata => ({ + modelId: params.checkpoint, + modelLabel: + (isExternalRequest || responseModelId !== params.checkpoint + ? responseModelId + : selectedModelSummary?.name || responseModelId) || + params.checkpoint || + "Unknown model", + responseModelId: + responseModelId || + externalSelection?.modelId || + params.checkpoint, + ...(externalProvider?.id ? { providerId: externalProvider.id } : {}), + providerName: + externalProvider?.name ?? + (isExternalRequest ? "External provider" : "Local model"), + providerType: externalProvider?.providerType ?? "local", + startedAt: streamStartTime, + finishedAt, + durationMs: finishedAt - streamStartTime, + ...(sandboxSessionId ? { sessionId: sandboxSessionId } : {}), + cancelId, + toolCalls: Array.from( + new Set( + toolCallParts + .map((part) => part.toolName) + .filter( + (toolName): toolName is string => + typeof toolName === "string" && toolName.length > 0, + ), + ), + ), + tools: { + search: + webSearchEnabledForThisTurn || + (!isExternalRequest && supportsTools && toolsEnabled), + fetch: webFetchEnabledForThisTurn, + code: + codeExecEnabledForThisTurn || + (!isExternalRequest && supportsTools && codeToolsEnabled), + images: imageGenerationEnabledForThisTurn, + mcp: !isExternalRequest && supportsTools && mcpEnabledForChat, + docs: + !isExternalRequest && + supportsTools && + (ragEnabled || projectRagEnabled), + artifacts: renderHtmlToolEnabledForThisTurn, + confirmToolCalls, + bypassPermissions, + }, + }); const externalCapabilities = getProviderCapabilities( externalProvider?.providerType, ); @@ -2768,6 +2851,11 @@ export function createOpenAIStreamAdapter( const stream = streamChatCompletions(requestPayload, abortSignal); for await (const chunk of stream) { + const chunkModel = (chunk as { model?: unknown }).model; + if (typeof chunkModel === "string" && chunkModel.length > 0) { + responseModelId = chunkModel; + } + // Handle tool status events const toolStatusText = ( chunk as unknown as { _toolStatus?: string } @@ -3435,11 +3523,12 @@ export function createOpenAIStreamAdapter( }); } + const finishedAt = Date.now(); const finalTiming = buildTiming( streamStartTime, totalChunks, serverPromptEvalTime ?? firstTokenTime, - Date.now() - streamStartTime, + finishedAt - streamStartTime, finalTokenCount, toolCallParts.length, finalTokPerSec, @@ -3475,6 +3564,7 @@ export function createOpenAIStreamAdapter( modelId: params.checkpoint, } : undefined, + responseDetails: buildResponseDetails(finishedAt), timing: finalTiming, }, }, diff --git a/studio/frontend/src/features/chat/index.ts b/studio/frontend/src/features/chat/index.ts index 75749b0040..7cd9611c71 100644 --- a/studio/frontend/src/features/chat/index.ts +++ b/studio/frontend/src/features/chat/index.ts @@ -26,7 +26,12 @@ export { type PlusMenuItemId, } from "./stores/plus-menu-prefs-store"; export { useChatModelRuntime } from "./hooks/use-chat-model-runtime"; -export { isExternalModelId } from "./external-providers"; +export { + customProviderDisplayName, + isExternalModelId, + parseExternalModelId, +} from "./external-providers"; +export { useExternalProvidersStore } from "./stores/external-providers-store"; export { ChatSearchDialog } from "./components/chat-search-dialog"; export { setTrainingCompareHandoff } from "./lib/training-compare-handoff"; export type { ProjectRecord } from "./types"; diff --git a/studio/frontend/src/features/chat/stores/chat-preferences-store.ts b/studio/frontend/src/features/chat/stores/chat-preferences-store.ts index f019b7dc3d..5b5011a78f 100644 --- a/studio/frontend/src/features/chat/stores/chat-preferences-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-preferences-store.ts @@ -7,11 +7,14 @@ import { persist } from "zustand/middleware"; // Client-side chat UI prefs kept in localStorage, not the chat DB. // confirmDeleteChats: when off, deleting a chat skips the confirm dialog. // showModelDisclaimer: when off, hide the "LLMs can make mistakes" footer note. +// showResponseModel: when on, assistant responses show the producing model. export interface ChatPreferencesState { confirmDeleteChats: boolean; setConfirmDeleteChats: (value: boolean) => void; showModelDisclaimer: boolean; setShowModelDisclaimer: (value: boolean) => void; + showResponseModel: boolean; + setShowResponseModel: (value: boolean) => void; } export const useChatPreferencesStore = create()( @@ -23,6 +26,9 @@ export const useChatPreferencesStore = create()( showModelDisclaimer: true, setShowModelDisclaimer: (showModelDisclaimer) => set({ showModelDisclaimer }), + showResponseModel: false, + setShowResponseModel: (showResponseModel) => + set({ showResponseModel }), }), { name: "unsloth_chat_preferences", @@ -32,6 +38,7 @@ export const useChatPreferencesStore = create()( ...current, confirmDeleteChats: saved?.confirmDeleteChats ?? true, showModelDisclaimer: saved?.showModelDisclaimer ?? true, + showResponseModel: saved?.showResponseModel ?? false, }; }, }, diff --git a/studio/frontend/src/features/settings/tabs/chat-tab.tsx b/studio/frontend/src/features/settings/tabs/chat-tab.tsx index b26b09e022..9519898b21 100644 --- a/studio/frontend/src/features/settings/tabs/chat-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/chat-tab.tsx @@ -213,6 +213,12 @@ export function ChatTab() { const setShowModelDisclaimer = useChatPreferencesStore( (state) => state.setShowModelDisclaimer, ); + const showResponseModel = useChatPreferencesStore( + (state) => state.showResponseModel, + ); + const setShowResponseModel = useChatPreferencesStore( + (state) => state.setShowResponseModel, + ); useEffect(() => { void countAllChats().then(setCount); @@ -412,6 +418,15 @@ export function ChatTab() { onCheckedChange={setShowModelDisclaimer} /> + + + diff --git a/tests/studio/test_chat_response_details_ui_contract.py b/tests/studio/test_chat_response_details_ui_contract.py new file mode 100644 index 0000000000..04301a0de6 --- /dev/null +++ b/tests/studio/test_chat_response_details_ui_contract.py @@ -0,0 +1,93 @@ +"""Static contract for the chat response-details action and metadata.""" + +from __future__ import annotations + +import re +from pathlib import Path + +REPO = Path(__file__).resolve().parents[2] +THREAD_TSX = REPO / "studio/frontend/src/components/assistant-ui/thread.tsx" +DETAILS_TSX = ( + REPO / "studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx" +) +REASONING_TSX = REPO / "studio/frontend/src/components/assistant-ui/reasoning.tsx" +ADAPTER_TS = REPO / "studio/frontend/src/features/chat/api/chat-adapter.ts" +CHAT_PREFS_TS = REPO / "studio/frontend/src/features/chat/stores/chat-preferences-store.ts" +CHAT_TAB_TSX = REPO / "studio/frontend/src/features/settings/tabs/chat-tab.tsx" + + +def test_assistant_more_menu_exposes_response_details_action(): + src = THREAD_TSX.read_text() + assert "MessageResponseDetailsSheet" in src + assert "See response details" in src + assert "setDetailsOpen(true)" in src + + +def test_response_details_sheet_uses_unsloth_sheet_and_key_sections(): + src = DETAILS_TSX.read_text() + assert "SheetContent" in src + assert "Response details" in src + assert "MessageResponseModelBadge" in src + assert "showResponseModel" in src + assert "ChipIcon" not in src + assert "s.params.checkpoint" not in src + assert "Not recorded" in src + assert "min-w-0 break-words font-heading" in src + assert "toolCallsFromContent(message.content)" in src + assert 'label="Called"' in src + for section in ["Response", "Tokens", "Timing", "Tools"]: + assert f'title="{section}"' in src + for field in ["Model", "Provider", "Total", "Cache hits", "Enabled", "Called"]: + assert f'label="{field}"' in src + + +def test_response_model_chip_is_user_configurable_and_rendered_in_metadata_rows(): + prefs_src = CHAT_PREFS_TS.read_text() + chat_tab_src = CHAT_TAB_TSX.read_text() + thread_src = THREAD_TSX.read_text() + reasoning_src = REASONING_TSX.read_text() + + assert "showResponseModel: boolean" in prefs_src + assert "showResponseModel: false" in prefs_src + assert "showResponseModel: saved?.showResponseModel ?? false" in prefs_src + assert "Show response model" in chat_tab_src + assert "setShowResponseModel" in chat_tab_src + assert "aui-response-model-badge inline-flex min-h-5" in DETAILS_TSX.read_text() + assert "leading-5" in DETAILS_TSX.read_text() + assert "group-hover/assistant-message:opacity-100" in DETAILS_TSX.read_text() + assert "MessageResponseModelBadge" in thread_src + assert "hasReasoningParts" in thread_src + assert "group/assistant-message aui-assistant-message-root" in thread_src + assert "pointer-events-none relative h-0" in thread_src + assert "MessageResponseModelBadge" in reasoning_src + assert 'className="min-w-0 flex-none"' in reasoning_src + assert "hidden min-w-0 max-w-[12rem]" in reasoning_src + assert "group-hover/assistant-message:inline-flex" in reasoning_src + + +def test_response_details_metadata_is_persisted_without_backend_schema_change(): + src = ADAPTER_TS.read_text() + assert "interface ResponseDetailsMetadata" in src + assert "buildResponseDetails" in src + assert "responseDetails: buildResponseDetails(finishedAt)" in src + assert "toolCalls: Array.from(" in src + assert "!isExternalRequest && supportsTools && toolsEnabled" in src + assert "!isExternalRequest && supportsTools && codeToolsEnabled" in src + assert re.search(r"selectedModelSummary\?\.name\s*\|\|\s*responseModelId", src) + assert "providerName" in src + assert "cancelId" in src + metadata_block = src[ + src.find("interface ResponseDetailsMetadata") : src.find("type RunMessages") + ] + builder_block = src[ + src.find("const buildResponseDetails") : src.find("const externalCapabilities") + ] + for forbidden in [ + "encrypted_api_key", + "externalApiKey", + "apiKey", + "providerKey", + "secret", + ]: + assert forbidden not in metadata_block + assert forbidden not in builder_block