ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge branch 'main' into dh/test-5106-windows-gpu-ci-mock

Browse source 
Resolve conflict in studio/backend/core/inference/llama_cpp.py: take main's
load_model body wrapped in the serial_load_lock (#5401 / #5161) and the
#5347 mmproj guard, layered with this PR's two intended edits:

- New @staticmethod _build_windows_path_dirs(binary_dir, prefix, cuda_path)
  next to _windows_pip_nvidia_dll_dirs, returning the win32 PATH ordering
  binary_dir, pip-nvidia wheels, CUDA_PATH/bin, CUDA_PATH/bin/x64.
- start_llama_server's win32 branch now calls the helper instead of
  re-implementing the same ordering inline.

The PR's own install_llama_prebuilt.py nit (f"llama-" -> "llama-") and the
new test file under studio/backend/tests/ auto-merged cleanly.

Picks up PR #5423's transformers 5.x drift-detector predicate fixes via the
merge, which is what was failing on the three Core CI jobs.
This commit is contained in:
danielhanchen 2026-05-17 12:54:48 +00:00
commit 9bf8d3aecc
104 changed files with 13002 additions and 2417 deletions
Download patch file Download diff file Expand all files Collapse all files

107
.github/scripts/hf-download-with-retry.sh vendored Executable file
View file

@ -0,0 +1,107 @@
#!/usr/bin/env bash
#
# Download a single file from a Hugging Face repo with a stall-retry
# watchdog. Used by the Studio CI workflows so a hung hf-xet transfer
# kills + retries instead of silently consuming the job's timeout.
#
# Usage: hf-download-with-retry.sh REPO FILE LOCAL_DIR
#
# Why this exists
# ---------------
# huggingface_hub 1.15+ deprecated `hf_transfer` and routes every
# transfer through the `hf-xet` binary package. In CI we observed
# `hf download` on a 3 GB GGUF (gemma-4-E2B-it-UD-Q4_K_XL) progress
# to ~46% via Xet, then go completely silent for the remainder of
# the 30-min job timeout -- no progress bytes, no error, no exit.
# A sibling 940 MB mmproj on the same step downloaded in ~21s
# moments earlier, so the hang is per-file inside hf-xet rather
# than a network outage. The Xet env-vars below put hf-xet into
# its highest-throughput mode and force a 500 s client-read
# timeout; the watchdog loop ensures a stall does not eat the
# whole job: if the hf process has not exited after STALL_S
# seconds (default 180 = 3 min), we SIGTERM, then SIGKILL, then
# start a fresh attempt. Retries are unbounded -- the enclosing
# GitHub Actions job's `timeout-minutes` is the real bound.
#
# See https://huggingface.co/docs/huggingface_hub/package_reference/environment_variables
# for the HF_XET_* documentation, and npm/cli#7308's pattern (silent
# CI hang with no error) for prior art on this class of failure.
set -uo pipefail
REPO="${1:?usage: hf-download-with-retry.sh REPO FILE [LOCAL_DIR]}"
FILE="${2:?usage: hf-download-with-retry.sh REPO FILE [LOCAL_DIR]}"
# LOCAL_DIR is optional. If empty, hf falls back to HF_HUB_CACHE
# (~/.cache/huggingface/hub) which is the desired path for callers
# that populate HF_HOME for a downstream Studio model load.
LOCAL_DIR="${3:-}"
# Stall threshold per attempt, in seconds. Override with
# HF_DOWNLOAD_STALL_SECONDS in the workflow env if 3 min is too tight
# for a specific runner / file. The script keeps retrying past this
# until the job timeout fires.
STALL_S="${HF_DOWNLOAD_STALL_SECONDS:-180}"
# hf-xet tuning. HF_HUB_ENABLE_HF_TRANSFER is deliberately NOT set --
# it is a no-op on huggingface_hub>=1.15 and only emits a deprecation
# FutureWarning. The five HF_XET_* knobs below mirror the settings
# Daniel asked for: max bandwidth + 64 parallel range gets, no chunk
# cache (download-once usage pattern), parallel disk writes (SSD/NVMe
# runners), and a generous 500 s read timeout so individual chunk
# requests fail loudly instead of stalling forever.
export HF_XET_HIGH_PERFORMANCE=1
export HF_XET_CHUNK_CACHE_SIZE_BYTES=0
export HF_XET_NUM_CONCURRENT_RANGE_GETS=64
export HF_XET_RECONSTRUCT_WRITE_SEQUENTIALLY=0
export HF_XET_CLIENT_READ_TIMEOUT=500
if [ -n "$LOCAL_DIR" ]; then
mkdir -p "$LOCAL_DIR"
fi
attempt=1
while : ; do
log="$(mktemp -t hf-download.XXXXXX)"
echo "[hf-download] $FILE attempt $attempt (stall threshold ${STALL_S}s, log=$log)"
if [ -n "$LOCAL_DIR" ]; then
hf download "$REPO" "$FILE" --local-dir "$LOCAL_DIR" > "$log" 2>&1 &
else
hf download "$REPO" "$FILE" > "$log" 2>&1 &
fi
pid=$!
elapsed=0
while kill -0 "$pid" 2>/dev/null && [ "$elapsed" -lt "$STALL_S" ]; do
sleep 5
elapsed=$((elapsed + 5))
done
if kill -0 "$pid" 2>/dev/null; then
echo "[hf-download] $FILE attempt $attempt exceeded ${STALL_S}s -- killing PID $pid and retrying"
kill -TERM "$pid" 2>/dev/null || true
sleep 2
kill -KILL "$pid" 2>/dev/null || true
wait "$pid" 2>/dev/null || true
echo "[hf-download] $FILE attempt $attempt log tail (last 40 lines):"
tail -40 "$log" || true
attempt=$((attempt + 1))
continue
fi
if wait "$pid"; then
rc=0
else
rc=$?
fi
if [ "$rc" -eq 0 ]; then
echo "[hf-download] $FILE attempt $attempt succeeded"
tail -20 "$log" || true
exit 0
fi
echo "[hf-download] $FILE attempt $attempt failed (exit $rc) -- retrying"
tail -40 "$log" || true
attempt=$((attempt + 1))
done

136
.github/workflows/consolidated-tests-ci.yml vendored
View file

@ -206,7 +206,8 @@ jobs:
'numpy<3' pytest==9.0.3 pytest-asyncio httpx \
protobuf sentencepiece triton \
psutil packaging tqdm safetensors datasets \
'peft>=0.18,<0.20' 'accelerate>=0.34,<2'
'peft>=0.18,<0.20' 'accelerate>=0.34,<2' \
ipython
# torchvision: unsloth_zoo.vision_utils imports it at module scope.
pip install --index-url https://download.pytorch.org/whl/cpu \
'torch>=2.4,<2.11' 'torchvision<0.26'
@ -304,6 +305,17 @@ jobs:
run: |
python -m pytest -v --tb=short tests/test_import_fixes_drift.py
- name: public-api surface drift detectors (9 tests, HARD GATE)
# Companion to test_import_fixes_drift.py: that file catches
# third-party drift; this one catches drift in unsloth's OWN
# public surface (FastLanguageModel / FastVisionModel /
# FastModel + their classmethods + is_bf16_supported). A
# rename here would silently break the unslothai/notebooks tree
# one PR cycle later -- this gate catches it BEFORE the
# breakage reaches users.
run: |
python -m pytest -v --tb=short tests/test_public_api_surface.py
- name: unsloth Bucket-A — CPU tests not in Repo tests (CPU)
# 16 tests across 5 files. They live inside tests/saving/ and
# tests/utils/, both of which Repo tests (CPU) excludes via --ignore
@ -875,14 +887,23 @@ jobs:
import _zoo_aggressive_cuda_spoof as _spoof
_spoof.apply()
# Hermetic cache dir + force compile path BEFORE importing
# unsloth_zoo.compiler (its globals capture env at module load).
# Hermetic cache dir + force compile path. The compiler's
# globals (UNSLOTH_COMPILE_LOCATION, UNSLOTH_COMPILE_USE_TEMP)
# are captured at module load; an earlier conftest `import
# unsloth` may have already imported unsloth_zoo.compiler with
# the default "unsloth_compiled_cache" path. Mutate the live
# module globals after import so this shim is robust to that
# ordering. Otherwise the compiler silently writes to the
# default cache and the per-model file assertion fails.
_CACHE = pathlib.Path(tempfile.mkdtemp(prefix="unsloth_cache_"))
os.environ["UNSLOTH_COMPILE_LOCATION"] = str(_CACHE)
os.environ["UNSLOTH_COMPILE_OVERWRITE"] = "1"
os.environ.pop("UNSLOTH_COMPILE_DISABLE", None)
import pytest
import unsloth_zoo.compiler as _zoo_compiler
_zoo_compiler.UNSLOTH_COMPILE_LOCATION = str(_CACHE)
_zoo_compiler.UNSLOTH_COMPILE_USE_TEMP = False
from unsloth_zoo.compiler import unsloth_compile_transformers
@ -941,6 +962,12 @@ jobs:
# Category E: undefined name in emitted file.
"perceiver": "name 'AbstractPreprocessor' is not defined",
"sam3_lite_text": "name 'Sam3LiteTextLayerScaledResidual' is not defined",
# Category F: compile exceeds 60s budget on the runner.
# First seen on transformers >=5,<6; each represents a slow
# or recursive source-rewriter path the zoo can address.
"beit": "TimeoutError: compile exceeds per-model budget",
"sam": "TimeoutError: compile exceeds per-model budget",
"sam_hq": "TimeoutError: compile exceeds per-model budget",
}
@ -956,40 +983,59 @@ jobs:
skipped -> no `modeling_<x>.py` file (expected for some
umbrella packages like `auto`, `deprecated`)
known -> in KNOWN_BROKEN_COMPILE; tracked for follow-up.
Any uncaught failure fails the cell."""
Any uncaught failure fails the cell.
Per-model SIGALRM cap so one infinite-looping model_type
cannot wedge the whole sweep + nuke the job timeout
(observed on transformers >=5,<6 -- 30+ min hang before
this guard landed)."""
import importlib as _il
import signal
ok = 0
skipped = []
known = []
new_failures = []
for model_type in _all_model_types():
modeling_path = f"transformers.models.{model_type}.modeling_{model_type}"
try:
_il.import_module(modeling_path)
except (ModuleNotFoundError, ImportError):
skipped.append((model_type, "no modeling file"))
continue
try:
unsloth_compile_transformers(
model_type=model_type, fast_lora_forwards=False,
)
except Exception as e:
msg = f"{type(e).__name__}: {str(e)[:200]}"
models = _all_model_types()
def _on_timeout(signum, frame):
raise TimeoutError("compile exceeded per-model budget")
prev_handler = signal.signal(signal.SIGALRM, _on_timeout)
try:
for i, model_type in enumerate(models):
if i % 25 == 0:
print(f" sweep progress: {i}/{len(models)} -> {model_type}", flush=True)
modeling_path = f"transformers.models.{model_type}.modeling_{model_type}"
try:
_il.import_module(modeling_path)
except (ModuleNotFoundError, ImportError):
skipped.append((model_type, "no modeling file"))
continue
signal.alarm(60)
try:
unsloth_compile_transformers(
model_type=model_type, fast_lora_forwards=False,
)
except Exception as e:
signal.alarm(0)
msg = f"{type(e).__name__}: {str(e)[:200]}"
if model_type in KNOWN_BROKEN_COMPILE:
known.append((model_type, msg))
else:
new_failures.append((model_type, msg))
continue
signal.alarm(0)
if model_type in KNOWN_BROKEN_COMPILE:
known.append((model_type, msg))
else:
new_failures.append((model_type, msg))
continue
if model_type in KNOWN_BROKEN_COMPILE:
# Came back green unexpectedly -- that's GOOD news,
# the bug was fixed. Surface it so we can drop the
# entry from KNOWN_BROKEN_COMPILE.
print(
f" UNEXPECTED-OK {model_type}: was in "
"KNOWN_BROKEN_COMPILE, now compiles cleanly. "
"Drop the entry."
)
ok += 1
# Came back green unexpectedly -- that's GOOD news,
# the bug was fixed. Surface it so we can drop the
# entry from KNOWN_BROKEN_COMPILE.
print(
f" UNEXPECTED-OK {model_type}: was in "
"KNOWN_BROKEN_COMPILE, now compiles cleanly. "
"Drop the entry."
)
ok += 1
finally:
signal.alarm(0)
signal.signal(signal.SIGALRM, prev_handler)
print(f"\nCompile sweep: ok={ok} skipped={len(skipped)} "
f"known-broken={len(known)} new-failures={len(new_failures)}")
for m, r in known:
@ -1020,24 +1066,34 @@ jobs:
"""Spot-check on the three production-relevant families that
the compile_every sweep also covers; this case verifies the
emitted cache file has the model-specific RMSNorm class
attribute, not just that the file parses + imports."""
attribute, not just that the file parses + imports.
``unsloth_compile_transformers`` is not idempotent in-
process: calling it twice on the same modeling module
after rewriting class attributes corrupts the inspect
source/line cache and the second emitted file is malformed
Python. The sweep above already produced a valid cache
file for every non-KNOWN_BROKEN model_type, so just verify
that artefact here. Trigger a compile only when running
this test in isolation (no sweep preceded)."""
import importlib as _il
try:
_il.import_module(
modeling = _il.import_module(
f"transformers.models.{model_type}.modeling_{model_type}"
)
except ModuleNotFoundError:
pytest.skip(
f"transformers build lacks model_type={model_type}"
)
unsloth_compile_transformers(
model_type=model_type, fast_lora_forwards=False,
)
modeling = _il.import_module(
f"transformers.models.{model_type}.modeling_{model_type}"
)
assert getattr(modeling, "__UNSLOTH_PATCHED__", False) is True
combined = _CACHE / f"unsloth_compiled_module_{model_type}.py"
if not combined.exists():
unsloth_compile_transformers(
model_type=model_type, fast_lora_forwards=False,
)
modeling = _il.import_module(
f"transformers.models.{model_type}.modeling_{model_type}"
)
assert getattr(modeling, "__UNSLOTH_PATCHED__", False) is True
_verify_file(combined, must_expose=[rms_class])

13
.github/workflows/mlx-ci.yml vendored
View file

@ -302,15 +302,10 @@ jobs:
"$LLAMA_QUANT" --help >/dev/null && echo " llama-quantize loads OK"
mkdir -p /tmp/ggufs
python -c "
from huggingface_hub import hf_hub_download
p = hf_hub_download(
'unsloth/gemma-3-270m-it-GGUF',
'gemma-3-270m-it-Q4_K_M.gguf',
local_dir = '/tmp/ggufs',
)
print('downloaded:', p)
"
bash .github/scripts/hf-download-with-retry.sh \
'unsloth/gemma-3-270m-it-GGUF' \
'gemma-3-270m-it-Q4_K_M.gguf' \
/tmp/ggufs
PORT=18080
echo "=== starting llama-server on 127.0.0.1:$PORT ==="

6
.github/workflows/notebooks-ci.yml vendored
View file

@ -200,7 +200,7 @@ jobs:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
persist-credentials: false
with: { path: unsloth }
path: unsloth
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
repository: unslothai/notebooks
@ -246,7 +246,7 @@ jobs:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
persist-credentials: false
with: { path: unsloth }
path: unsloth
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
repository: unslothai/notebooks
@ -352,7 +352,7 @@ jobs:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
persist-credentials: false
with: { path: unsloth }
path: unsloth
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
repository: unslothai/notebooks

4
.github/workflows/security-audit.yml vendored
View file

@ -137,8 +137,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
- uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable @ 2026-03-27
@ -1066,8 +1064,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:

7
.github/workflows/studio-api-smoke.yml vendored
View file

@ -63,8 +63,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
@ -87,10 +85,9 @@ jobs:
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
python -m pip install --upgrade huggingface_hub hf_transfer
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
HF_HUB_ENABLE_HF_TRANSFER=1 \
hf download "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
- name: Save HF_HOME for ${{ env.GGUF_REPO }}
if: always() && steps.prime-hf.outcome == 'success'

24
.github/workflows/studio-frontend-ci.yml vendored
View file

@ -15,6 +15,8 @@ on:
pull_request:
paths:
- 'studio/frontend/**'
- 'scripts/check_frontend_dep_removal.py'
- 'tests/studio/test_frontend_dep_removal.py'
- '.github/workflows/studio-frontend-ci.yml'
push:
branches: [main, pip]
@ -57,8 +59,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
# Run the structural lockfile scan BEFORE npm ci. A compromised
# tarball runs its `prepare` / `postinstall` during `npm ci`,
@ -86,6 +86,26 @@ jobs:
exit 1
fi
# Catch the common foot-gun: a dep dropped from package.json that is
# still imported somewhere. The script walks the lockfile dep graph
# from the new top-level deps and only counts top-level node_modules
# paths as valid resolution targets for bare src/ imports.
#
# actions/checkout uses fetch-depth: 1 by default, so the base branch
# is not available locally. Fetch the single base commit with an
# explicit refspec so origin/<base> is reliably created (a bare
# `git fetch origin <ref>` only updates FETCH_HEAD in some configs).
- name: Dependency removal safety check
if: github.event_name == 'pull_request'
working-directory: ${{ github.workspace }}
run: |
git fetch --no-tags --depth=1 origin \
"${{ github.base_ref }}:refs/remotes/origin/${{ github.base_ref }}"
python3 scripts/check_frontend_dep_removal.py \
--base "origin/${{ github.base_ref }}" \
--enumerate-dead
python3 tests/studio/test_frontend_dep_removal.py
- name: Typecheck
run: npm run typecheck

24
.github/workflows/studio-inference-smoke.yml vendored
View file

@ -79,8 +79,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
@ -101,10 +99,9 @@ jobs:
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
python -m pip install --upgrade huggingface_hub hf_transfer
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
HF_HUB_ENABLE_HF_TRANSFER=1 \
hf download "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
- name: Save HF_HOME for ${{ env.GGUF_REPO }}
if: always() && steps.prime-hf.outcome == 'success'
@ -329,8 +326,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
@ -351,10 +346,9 @@ jobs:
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
python -m pip install --upgrade huggingface_hub hf_transfer
python -m pip install --upgrade huggingface_hub
mkdir -p gguf-cache
HF_HUB_ENABLE_HF_TRANSFER=1 \
hf download "$GGUF_REPO" "$GGUF_FILE" --local-dir gguf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache
- name: Save GGUF model file
if: always() && steps.download-gguf.outcome == 'success'
@ -648,8 +642,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
@ -670,12 +662,10 @@ jobs:
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
python -m pip install --upgrade huggingface_hub hf_transfer
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
HF_HUB_ENABLE_HF_TRANSFER=1 \
hf download "$GGUF_REPO" "$GGUF_FILE"
HF_HUB_ENABLE_HF_TRANSFER=1 \
hf download "$GGUF_REPO" "$MMPROJ_FILE"
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$MMPROJ_FILE"
- name: Save HF_HOME for ${{ env.GGUF_REPO }} (model + mmproj)
if: always() && steps.prime-hf.outcome == 'success'

7
.github/workflows/studio-mac-api-smoke.yml vendored
View file

@ -50,8 +50,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
@ -72,10 +70,9 @@ jobs:
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
python -m pip install --upgrade huggingface_hub hf_transfer
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
HF_HUB_ENABLE_HF_TRANSFER=1 \
hf download "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
- name: Save HF_HOME for ${{ env.GGUF_REPO }}
if: always() && steps.prime-hf.outcome == 'success'

49
.github/workflows/studio-mac-inference-smoke.yml vendored
View file

@ -73,8 +73,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
@ -95,10 +93,9 @@ jobs:
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
python -m pip install --upgrade huggingface_hub hf_transfer
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
HF_HUB_ENABLE_HF_TRANSFER=1 \
hf download "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
# Save partial caches on cancel/timeout -- hf download resumes by
# content hash. `outcome != skipped` keeps cache-hit a no-op.
@ -325,8 +322,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
@ -347,10 +342,9 @@ jobs:
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
python -m pip install --upgrade huggingface_hub hf_transfer
python -m pip install --upgrade huggingface_hub
mkdir -p gguf-cache
HF_HUB_ENABLE_HF_TRANSFER=1 \
hf download "$GGUF_REPO" "$GGUF_FILE" --local-dir gguf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache
# Save partial caches on cancel; next run resumes via content hash.
- name: Save GGUF model file
@ -692,8 +686,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
@ -710,23 +702,32 @@ jobs:
continue-on-error: true
with:
path: gguf-cache
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-${{ env.MMPROJ_FILE }}-v1
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-${{ env.MMPROJ_FILE }}-v2
- name: Download GGUF + mmproj if cache miss
- name: Verify cache contains BOTH gguf + mmproj
id: verify-cache
if: steps.cache-gguf.outputs.cache-hit == 'true'
run: |
if [[ -f "gguf-cache/$GGUF_FILE" && -f "gguf-cache/$MMPROJ_FILE" ]]; then
echo "ok=true" >> "$GITHUB_OUTPUT"
else
echo "Partial cache hit -- forcing re-download."
echo "ok=false" >> "$GITHUB_OUTPUT"
fi
- name: Download GGUF + mmproj if cache miss or partial
id: download-gguf
if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success'
if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.verify-cache.outputs.ok != 'true'
# Authenticated + parallel: shared macos-14 NAT egress stalls
# multi-GB anonymous downloads.
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
python -m pip install --upgrade huggingface_hub hf_transfer
python -m pip install --upgrade huggingface_hub
mkdir -p gguf-cache
HF_HUB_ENABLE_HF_TRANSFER=1 \
hf download "$GGUF_REPO" "$GGUF_FILE" --local-dir gguf-cache &
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache &
MODEL_PID=$!
HF_HUB_ENABLE_HF_TRANSFER=1 \
hf download "$GGUF_REPO" "$MMPROJ_FILE" --local-dir gguf-cache &
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$MMPROJ_FILE" gguf-cache &
MMPROJ_PID=$!
wait "$MODEL_PID"
wait "$MMPROJ_PID"
@ -734,13 +735,15 @@ jobs:
ls -lh "gguf-cache/$GGUF_FILE" "gguf-cache/$MMPROJ_FILE"
# Save partial caches on cancel. hashFiles guard avoids a hard
# save failure when the download step exits with no files.
# save failure when the download step exits with no files. The
# additional mmproj-presence check stops a partial save from
# poisoning the cache for the next run.
- name: Save GGUF + mmproj files
if: always() && steps.download-gguf.outcome != 'skipped' && hashFiles('gguf-cache/**/*.gguf') != ''
if: always() && steps.download-gguf.outcome != 'skipped' && hashFiles('gguf-cache/**/*.gguf') != '' && hashFiles(format('gguf-cache/{0}', env.MMPROJ_FILE)) != ''
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: gguf-cache
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-${{ env.MMPROJ_FILE }}-v1
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-${{ env.MMPROJ_FILE }}-v2
- name: Install Studio (--local, --no-torch)
env:

7
.github/workflows/studio-mac-ui-smoke.yml vendored
View file

@ -50,8 +50,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
@ -72,10 +70,9 @@ jobs:
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
python -m pip install --upgrade huggingface_hub hf_transfer
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
HF_HUB_ENABLE_HF_TRANSFER=1 \
hf download "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
- name: Save HF_HOME for ${{ env.GGUF_REPO }}
if: always() && steps.prime-hf.outcome == 'success'

2
.github/workflows/studio-mac-update-smoke.yml vendored
View file

@ -52,8 +52,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:

2
.github/workflows/studio-tauri-smoke.yml vendored
View file

@ -53,8 +53,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '24'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
- uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable @ 2026-03-27

62
.github/workflows/studio-ui-smoke.yml vendored
View file

@ -64,8 +64,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
@ -86,10 +84,9 @@ jobs:
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
python -m pip install --upgrade huggingface_hub hf_transfer
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
HF_HUB_ENABLE_HF_TRANSFER=1 \
hf download "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
- name: Save HF_HOME for ${{ env.GGUF_REPO }}
if: always() && steps.prime-hf.outcome == 'success'
@ -232,12 +229,55 @@ jobs:
kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true
sleep 2
# IME + multilingual paste regression (issue #5318 / PR #5327).
# Third Studio on its own port so a hang here cannot poison the
# earlier UI tests. No GGUF -- the bug surface is the composer.
- name: Reset auth + boot Studio for IME / i18n tests (port 18896)
run: |
unsloth studio reset-password
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18896 \
> logs/studio_ime.log 2>&1 &
echo "STUDIO_IME_PID=$!" >> "$GITHUB_ENV"
- name: Wait for /api/health on 18896
run: |
for i in $(seq 1 180); do
if curl -fs "http://127.0.0.1:18896/api/health" > /tmp/health3.json; then
jq -e '.status == "healthy"' /tmp/health3.json && break
fi
sleep 1
done
jq -e '.status == "healthy"' /tmp/health3.json
- name: Pass bootstrap pw for IME / i18n test
# IME smoke does the change-password against the bootstrap that
# Studio's frontend injects into the page, so it only needs the
# NEW password.
run: |
NEW="CIIme-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')"
echo "::add-mask::$NEW"
echo "STUDIO_IME_NEW_PW=$NEW" >> "$GITHUB_ENV"
- name: Drive IME + multilingual paste regression with Playwright
env:
BASE_URL: http://127.0.0.1:18896
STUDIO_NEW_PW: ${{ env.STUDIO_IME_NEW_PW }}
PW_ART_DIR: logs/playwright_ime
STUDIO_UI_STRICT: '1'
run: |
mkdir -p logs/playwright_ime
python tests/studio/playwright_chat_ime_i18n.py
- name: Stop third Studio
if: always()
run: |
kill "${STUDIO_IME_PID}" 2>/dev/null || true
sleep 2
- name: Upload Playwright artifacts
# Always upload (not just failure) so a green run's screenshots
# are reviewable in the Actions UI -- catches "passed but the
# UI is silently broken" regressions that would be invisible
# otherwise. Both Studio's logs (chat + extra) and BOTH
# Playwright artifact dirs are bundled.
# Always upload so a green run's screenshots stay reviewable --
# catches "passed but the UI is silently broken" regressions.
if: always()
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
@ -245,7 +285,9 @@ jobs:
path: |
logs/studio.log
logs/studio_extra.log
logs/studio_ime.log
logs/install.log
logs/playwright
logs/playwright_extra
logs/playwright_ime
retention-days: 7

2
.github/workflows/studio-update-smoke.yml vendored
View file

@ -52,8 +52,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:

7
.github/workflows/studio-windows-api-smoke.yml vendored
View file

@ -58,8 +58,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
@ -79,10 +77,9 @@ jobs:
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
python -m pip install --upgrade huggingface_hub hf_transfer
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
HF_HUB_ENABLE_HF_TRANSFER=1 \
hf download "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
- name: Save HF_HOME for ${{ env.GGUF_REPO }}
if: always() && steps.prime-hf.outcome == 'success'

54
.github/workflows/studio-windows-inference-smoke.yml vendored
View file

@ -68,8 +68,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
@ -101,10 +99,9 @@ jobs:
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
python -m pip install --upgrade huggingface_hub hf_transfer
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
HF_HUB_ENABLE_HF_TRANSFER=1 \
hf download "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
- name: Save HF_HOME cache for ${{ env.GGUF_REPO }}
# Only write a fresh cache entry when we actually rebuilt the
@ -345,9 +342,13 @@ jobs:
- name: Stop Studio
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
sleep 2
# Run as cmd so we are not running through the Git Bash shell;
# Git Bash on windows-latest has been observed to exit 143
# (SIGTERM) from any inline kill/sleep block, masking a green
# test run. The runner reclaims the Studio child process at
# job end either way, so just emit a marker and exit 0.
shell: cmd
run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end)
- name: Upload logs
if: always()
@ -396,8 +397,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
@ -420,10 +419,9 @@ jobs:
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
python -m pip install --upgrade huggingface_hub hf_transfer
python -m pip install --upgrade huggingface_hub
mkdir -p gguf-cache
HF_HUB_ENABLE_HF_TRANSFER=1 \
hf download "$GGUF_REPO" "$GGUF_FILE" --local-dir gguf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache
- name: Save GGUF model cache
if: always() && steps.download-gguf.outcome == 'success'
@ -762,9 +760,13 @@ jobs:
- name: Stop Studio
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
sleep 2
# Run as cmd so we are not running through the Git Bash shell;
# Git Bash on windows-latest has been observed to exit 143
# (SIGTERM) from any inline kill/sleep block, masking a green
# test run. The runner reclaims the Studio child process at
# job end either way, so just emit a marker and exit 0.
shell: cmd
run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end)
- name: Upload logs
if: always()
@ -806,8 +808,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
@ -833,12 +833,10 @@ jobs:
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
python -m pip install --upgrade huggingface_hub hf_transfer
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
HF_HUB_ENABLE_HF_TRANSFER=1 \
hf download "$GGUF_REPO" "$GGUF_FILE"
HF_HUB_ENABLE_HF_TRANSFER=1 \
hf download "$GGUF_REPO" "$MMPROJ_FILE"
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$MMPROJ_FILE"
- name: Save HF_HOME cache for ${{ env.GGUF_REPO }} (model + mmproj)
if: always() && steps.prime-hf.outcome == 'success'
@ -1150,9 +1148,13 @@ jobs:
- name: Stop Studio
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
sleep 2
# Run as cmd so we are not running through the Git Bash shell;
# Git Bash on windows-latest has been observed to exit 143
# (SIGTERM) from any inline kill/sleep block, masking a green
# test run. The runner reclaims the Studio child process at
# job end either way, so just emit a marker and exit 0.
shell: cmd
run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end)
- name: Upload logs
if: always()

14
.github/workflows/studio-windows-ui-smoke.yml vendored
View file

@ -63,8 +63,13 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
# No `cache: 'npm'`. setup-node's npm cache restore silently
# aborts the entire job on Windows runners when the npm cache
# path (`C:\npm\cache` per `npm config get cache`) doesn't yet
# exist on a fresh runner -- the step exits without an error
# message and every following step gets skipped. See
# npm/cli#7308. The frontend `npm ci` is fast enough without
# the cache that the reliability gain is worth the ~30s.
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
@ -88,10 +93,9 @@ jobs:
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
python -m pip install --upgrade huggingface_hub hf_transfer
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
HF_HUB_ENABLE_HF_TRANSFER=1 \
hf download "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
- name: Save HF_HOME for ${{ env.GGUF_REPO }}
if: always() && steps.prime-hf.outcome == 'success'

2
.github/workflows/studio-windows-update-smoke.yml vendored
View file

@ -64,8 +64,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:

3
.github/workflows/version-compat-ci.yml vendored
View file

@ -127,6 +127,7 @@ jobs:
run: |
PYTHONPATH=. python -m pytest \
tests/version_compat/test_peft_pinned_symbols.py \
tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py \
-v --tb=short
st-pinned-symbols:
@ -214,7 +215,7 @@ jobs:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
persist-credentials: false
with: { path: unsloth }
path: unsloth
- name: Clone unsloth-zoo @ main
run: |
# github.com occasionally 500s on the git fetch; retry so a

2
.github/workflows/wheel-smoke.yml vendored
View file

@ -48,8 +48,6 @@ jobs:
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
cache: 'npm'
cache-dependency-path: studio/frontend/package-lock.json
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:

151
pyproject.toml
View file

@ -69,6 +69,7 @@ triton = [
]
huggingfacenotorch = [
"unsloth_zoo>=2026.5.2",
"wheel>=0.42.0",
"packaging",
"numpy",
@ -1017,7 +1018,44 @@ intelgputorch290 = [
intel-gpu-torch290 = [
"unsloth[intelgputorch290]"
]
intelgputorch210 = [
intelgputorch271 = [
"unsloth_zoo[intelgpu]",
"unsloth[huggingfacenotorch]",
"pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=663ce21364096b268c6687f26f22862cb1001cae0c4ec9f98a0998415f99e2b0 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'",
"pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=dd92cc17000bad19f213b6a877d7f10cd71341b703cd188513ce9fff8d42e3dd ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'",
"pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=aa5c3ec21a89e967d1dfe61e3d5b1c1ae9620c871ed804771d3378d6a44066f2 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'",
"pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=d1c6f522e11112a311b1a61ba7b40b43ad8305675fa29153017ccb1ad0b6816d ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'",
"pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp310-cp310-win_amd64.whl#sha256=a5c16dcf449a9cb62bc3788f7ec45782bb3ead6edc2637a12b60ef0f8f45dc55 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp311-cp311-win_amd64.whl#sha256=bc2d76ffa4ceed5b38ae34b52dbff643442e1a44d52ca72d7cb520ca1950e9ae ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp312-cp312-win_amd64.whl#sha256=b09ca59ce52d6d27b1510df783cde222b703a71857a6fa953f1f155f9f50811a ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp313-cp313-win_amd64.whl#sha256=1260c4a4bad426b6cd3c8f3e1a21835381c6f217bf434bcb55fedec08a206dea ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp310-cp310-linux_x86_64.whl#sha256=231c3fbd88a75d94de5ccbbb7f4f9a96cb3c58b3d891c2a1b469d38df95f9be6 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp311-cp311-linux_x86_64.whl#sha256=78edcc27709dd819fc820f5eb9421bd10d3f3dcb14adb25ee60766c76f0e67f3 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp312-cp312-linux_x86_64.whl#sha256=b443df40bc9cb7d648a9f8f9ed1d5c3a1203e561ebd0a61dd55fb8a58833d5ec ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp313-cp313-linux_x86_64.whl#sha256=412b58ffcceebea399c9a1bcdb22896aa10385c2650a8c4f8a677fb11c49b448 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp310-cp310-win_amd64.whl#sha256=2591228dc2cb73c78daf24277c4449ba9474f94cd31938147249269fe89d05d6 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp311-cp311-win_amd64.whl#sha256=1aacb86e9a9684ffc8bde3db14b251d00df7019a9a434ec99a59076a2696325d ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp312-cp312-win_amd64.whl#sha256=9b65dc8562521b60d77aa653132bc03a19da0291318fcf919faa3f03080d8f7e ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp313-cp313-win_amd64.whl#sha256=cd3669fee311bc3ee5501d696bf989226a6f2bf957d120a04881a07af05526d6 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp310-cp310-manylinux_2_28_x86_64.whl#sha256=f8cdf6889c02b3166679eef661b68757ea7e99c314432c3d41dac3d2ed4a59d4 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp311-cp311-manylinux_2_28_x86_64.whl#sha256=f7d15b65d52809745992e0001c25034f33ac01f2dff5248614e07b5d009a59b7 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp312-cp312-manylinux_2_28_x86_64.whl#sha256=1ff1f98d70846352c7f56833bedab1a055ead27b11c120b8c719063ee0383554 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp313-cp313-manylinux_2_28_x86_64.whl#sha256=f46945344ea911a70309231eaaf3b80c96f6646ce5515dc89aa94f94144e310e ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp310-cp310-win_amd64.whl#sha256=ecae9a02de769e2070d37388116beb407c3f0d60b8e65c1da1423f4eafee361a ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp311-cp311-win_amd64.whl#sha256=2914e62782431bebd6ad9a3b98a2b7311e448e84a7534bb7f35874b9279a17de ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp312-cp312-win_amd64.whl#sha256=5b462c156f4e2097e1e53649d3f298ce352fa4c5d1e6addd360375b10ebd6c67 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp313-cp313-win_amd64.whl#sha256=fa87b3677cd1af67ce423004283c1bde80e3571f391182a3e89b485e18e3c70f ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
]
intel-gpu-torch271 = [
"unsloth[intelgputorch271]"
]
intelgputorch291 = [
"unsloth_zoo[intelgpu]",
"unsloth[huggingfacenotorch]",
@ -1030,6 +1068,43 @@ intelgputorch210 = [
"pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp312-cp312-win_amd64.whl#sha256=97337a47425f1963a723475bd61037460e84ba01db4f87a1d662c3718ff6c47e ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp313-cp313-win_amd64.whl#sha256=2caf8138695f6abb023ecd02031a2611ba1bf8fff2f19802567cb2fadefe9e87 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp310-cp310-linux_x86_64.whl#sha256=fb7895c744132d6a8e56ce8434ae1d8355c9bda4e9f58832744ff742d6268eaf ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp311-cp311-linux_x86_64.whl#sha256=da2604a9114a28de71ce654819424d20a246adf644d191ae160837df9731b79e ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp312-cp312-linux_x86_64.whl#sha256=d5968d78d81c1d01efc1b3bf83d7da3d83161dcc3a9fcf91f500591db1c6c75d ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp313-cp313-linux_x86_64.whl#sha256=b56d6b0d65863f370527e971dbfa046a5dd2a1f61cc95071db26c764f36e4dce ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp310-cp310-win_amd64.whl#sha256=2f318fb6a4bf1101cc17f35a5371f7c1768b41fceed03628397834e85b3edfdd ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp311-cp311-win_amd64.whl#sha256=c9cedc3fb099366b2e6c563df6578e323564b1b5d40ac27be73c674755343a1d ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp312-cp312-win_amd64.whl#sha256=bee9623254d0f95a1ca115dbd17e9a9d966fdb8ae123e2ada4a9eb2fb8d38db8 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp313-cp313-win_amd64.whl#sha256=cd5c857da52a63c121561b30b0979e69ade70b575fd74e389787bc7c1ee2ac11 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp310-cp310-manylinux_2_28_x86_64.whl#sha256=cc5272da2cb4554edf059eedd6d1f5ef2859033b0fb79d5dcb8e99a0697f3325 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp311-cp311-manylinux_2_28_x86_64.whl#sha256=3c80d6a068c32fc4ebddb27953e03a0141bd0f10ca8730417cbc0e0748158285 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp312-cp312-manylinux_2_28_x86_64.whl#sha256=8cf640a867cf270b3fda7a10002c29d3fc2ad6dfbd76404a8cdd820489adb04c ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp313-cp313-manylinux_2_28_x86_64.whl#sha256=d9c59ee5ae3d0560f02401c8dfd8054d50813a8dbb5d33a8777de7d02f6fcb7b ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp310-cp310-win_amd64.whl#sha256=843ea7fcd8f5a22ebbc20d2d61d9eec7593821a0372eb8cabb73953d12ef6acf ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp311-cp311-win_amd64.whl#sha256=e5ff8a31d3c700f8dbac59697c8e32298a43ec059609ebc6ea7bab3eff6384e1 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp312-cp312-win_amd64.whl#sha256=8bae6d4c042f8d20818da4a5aa9109c6fbd6ec11bc422be152ce8adf9a7095bf ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp313-cp313-win_amd64.whl#sha256=47059e290fc2a41ba78666ffcde102c436abf7ff8a34d200268b48c4fa0f9c45 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
]
intel-gpu-torch291 = [
"unsloth[intelgputorch291]"
]
intelgputorch210 = [
"unsloth_zoo[intelgpu]",
"unsloth[huggingfacenotorch]",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp310-cp310-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp311-cp311-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp312-cp312-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp313-cp313-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.10.0%2Bxpu-cp310-cp310-linux_x86_64.whl#sha256=abb1d1ec1ac672bac0ff35420c965f2df0c636ef9d94e2a830e34578489d0a57 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.10.0%2Bxpu-cp311-cp311-linux_x86_64.whl#sha256=71ad2f82da0f41eaec159f39fc85854e27c2391efa91b373e550648a6f4aaad3 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.10.0%2Bxpu-cp312-cp312-linux_x86_64.whl#sha256=b473571d478912f92881cc13f15fa18f8463fb0fb8a068c96ed47a7d45a4da0a ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'",
@ -1054,6 +1129,80 @@ intelgputorch210 = [
intel-gpu-torch210 = [
"unsloth[intelgputorch210]"
]
intelgputorch2110 = [
"unsloth_zoo[intelgpu]",
"unsloth[huggingfacenotorch]",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=2a1841138750f708ec017becbf8d357526f3fa350deee6553be5735ad66160a3 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=e85378f1fc1ea002271de2a35475b75008fa554b86ef9d3bc55be9c513a63b51 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=a6663ebe43e3c0d560ff774708632d7a75208ee64a291c1724ed5c16a92d1c72 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=08c8d43b2831faf9d6799480df2b45dde58102257aebd810d07a2ce18cd4e5df ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp310-cp310-win_amd64.whl#sha256=90fb8f767950a4ffca627faa7f86d9c697237ea4352d7e23505c5c9ed8e72216 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp311-cp311-win_amd64.whl#sha256=aa7de82f4265089e74f25a2701b7532e5c47d74224d877b61da1d66156e3f0c1 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp312-cp312-win_amd64.whl#sha256=5ba3a31c6e1b259ad2d924e1b50f72a78c6ebd7eb4f364473bbf93e144734e80 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp313-cp313-win_amd64.whl#sha256=e8b4caba9b2399ea4c7f9a2777042564dea5d6f9e586a2dcb015a4ce20f000f7 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp310-cp310-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp310-cp310-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp311-cp311-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp312-cp312-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp313-cp313-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp310-cp310-manylinux_2_28_x86_64.whl#sha256=6e634354b752b7366e8ad16b84f3e7e5863776a7ab448bbabae4fd36668dee7a ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp311-cp311-manylinux_2_28_x86_64.whl#sha256=293169899f562ce473a58836dd024f0b1e72a347400278287ab393d1b04991e4 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp312-cp312-manylinux_2_28_x86_64.whl#sha256=e204d14be6f0f84d5f0e6e9213556e80326c3ab682cac108bcbef340bf45297b ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp313-cp313-manylinux_2_28_x86_64.whl#sha256=f134344006f0989a2d771554b7905fb05bd93d63b195e64626fde3495ec6f287 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp310-cp310-win_amd64.whl#sha256=7e52729cb9736c66dc79a7f42de6b31db93b9161d3357fd34cfa33f5fe32b8ea ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp311-cp311-win_amd64.whl#sha256=83a6130100c6b6750d8aa9fd29e5d0c53b1c85b1153b8ed4139aea54fc1892cc ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp312-cp312-win_amd64.whl#sha256=03788e0e5a5b85a2f09d11f0263d579fcb0cf5623d8810149be0e37836c2738c ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp313-cp313-win_amd64.whl#sha256=cb1da1d378ce440f7d1e0ed8cf21bd280d904ab25a55c9453f8377825818df74 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
]
intel-gpu-torch2110 = [
"unsloth[intelgputorch2110]"
]
intelgputorch2120 = [
"unsloth_zoo[intelgpu]",
"unsloth[huggingfacenotorch]",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=844d981cb1b3948085e8cfa62c74de9f100259f6131959aa70be49123b88ae81 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=a16b1d00e94ad87d62af3512e390348b8656419598004100c56028bf494f086b ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=4e46e71e077cf483404a4c17ce40d71c5f0e13a81459139d4346ca427b1dd455 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=4fdaed1bafc51d3a2834656a3420a6686a74ea226508765a49bf15d58ff3a930 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-win_amd64.whl#sha256=2778b46b22e9fa0916398db299a125027a1b2331c1173b3dd2b9e2cab6263a31 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-win_amd64.whl#sha256=ad5b147d04ee0d40f3d4d32f85f5aa3a3beb6cd5799ca026d3d7f4afa3d9e24f ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-win_amd64.whl#sha256=d9482063af2a308543f23333e32edd738ea87cbb33ade68afda9ae0fd704ccd9 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-win_amd64.whl#sha256=5d4d67f0deb1e851c01b293e602b8dcddad26ca2be61221cee3dc0e1aa0cdefd ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp310-cp310-linux_x86_64.whl#sha256=e8923cd1fe560472904b1461b745d2f1826bb9c1bc0808225d5f28a450e4d553 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp311-cp311-linux_x86_64.whl#sha256=f7c082b2fc9b61def594d30ea57762dc4a8bc7111a9a9593953ed948de242e28 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp312-cp312-linux_x86_64.whl#sha256=f59decc04bec27862ed0197554a52370dbcba3e6892616d1fbce450e402bf2d5 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp313-cp313-linux_x86_64.whl#sha256=56f74e7c6c096e1a7ac215eb79ee590b764be3fbba8f4febc145bca47194a083 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp310-cp310-win_amd64.whl#sha256=b9779b71457b5a916ae052ed2467c10273cae4862d469b191359173b2038c53e ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp311-cp311-win_amd64.whl#sha256=7ef8e776c992e4e3ae007ebc108eb4f36b1d1dd9da97ecb308ab7fded89a2659 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp312-cp312-win_amd64.whl#sha256=7f1d40febf2b8724adf4ff23866897d87478cc43de2a20f7776dc00be334c464 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp313-cp313-win_amd64.whl#sha256=32770e2613df26e2c81ae64ea001b2ca12b8d152231285caff9b5f963a21ad75 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp310-cp310-manylinux_2_28_x86_64.whl#sha256=0d517462caf6f5201c0d7c880f4ac431783c88fcc59b4587836da6c72a89509c ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp311-cp311-manylinux_2_28_x86_64.whl#sha256=4b6feada86aa0bd606904b05898b33538106120d8ed706ba11d0011046534cb8 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp312-cp312-manylinux_2_28_x86_64.whl#sha256=e231819be0f87829c2344c909c1f0db9d6ae7d6faefe644a526a1a01d0c18d98 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp313-cp313-manylinux_2_28_x86_64.whl#sha256=8bc7d37515cea18af4c389d5fde58b1a9d76b015f2d87e4a7dc62ad50b1cc200 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp310-cp310-win_amd64.whl#sha256=65dbb041057dddfe369f29cfaab63f75563621779a23a7b1e2c0ff8a84d4376a ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp311-cp311-win_amd64.whl#sha256=df647445365924d69fe3bb2a15a7edfe5b63ef91e4ae69af11d93582985237a4 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp312-cp312-win_amd64.whl#sha256=b0db3df0d0d154d18ba988ab420f1da2549f9372113ff54ff66e4ae3c7fe3bd0 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp313-cp313-win_amd64.whl#sha256=c70850842068c43a0d50eaf139c25b6f6cc9b17a0dae70218c7e69edbee0bc80 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
]
intel-gpu-torch2120 = [
"unsloth[intelgputorch2120]"
]
intel = [
"unsloth[intelgputorch280]",
]

1195
scripts/check_frontend_dep_removal.py Normal file
View file

File diff suppressed because it is too large Load diff

1825
studio/backend/core/inference/external_provider.py
View file

File diff suppressed because it is too large Load diff

1329
studio/backend/core/inference/llama_cpp.py
View file

File diff suppressed because it is too large Load diff

100
studio/backend/core/inference/llama_server_args.py
View file

@ -69,7 +69,15 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
# Single-model server -- Studio runs one model per llama-server
# process and serves its own UI. Enabling multi-model loading or
# llama-server's built-in web UI changes the surface clients see.
# ``--webui``/``--no-webui`` are the legacy spelling; current
# upstream uses ``--ui``/``--no-ui`` + ``--ui-*`` companions.
# Keep both so the denylist matches old and new llama-server
# binaries (Studio's prebuilt vs system-llama.cpp).
frozenset({"--webui", "--no-webui"}),
frozenset({"--ui", "--no-ui"}),
frozenset({"--ui-config"}),
frozenset({"--ui-config-file"}),
frozenset({"--ui-mcp-proxy", "--no-ui-mcp-proxy"}),
frozenset({"--models-dir"}),
frozenset({"--models-preset"}),
frozenset({"--models-max"}),
@ -118,3 +126,95 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]:
def is_managed_flag(flag: str) -> bool:
"""True if ``flag`` is a Studio-managed llama-server flag."""
return flag in _DENYLIST
# Pass-through flags that shadow first-class ``LoadRequest`` fields
# (max_seq_length, cache_type_kv, speculative_type,
# chat_template_override). Stripped from inherited extras so they
# can't last-wins-override an Apply that re-sets the same first-class
# field.
_CONTEXT_FLAGS: frozenset[str] = frozenset({"-c", "--ctx-size"})
_CACHE_FLAGS: frozenset[str] = frozenset(
{"-ctk", "--cache-type-k", "-ctv", "--cache-type-v"}
)
_SPEC_FLAGS: frozenset[str] = frozenset(
{
"--spec-default",
"--spec-type",
"--spec-ngram-size-n",
"--spec-ngram-size",
"--draft-min",
"--draft-max",
}
)
_TEMPLATE_FLAGS: frozenset[str] = frozenset(
{
"--chat-template",
"--chat-template-file",
"--chat-template-kwargs",
"--jinja",
"--no-jinja",
}
)
_SHADOWING_FLAGS: frozenset[str] = (
_CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS
)
# Boolean flags inside _SHADOWING_FLAGS that take no value. The
# value-consuming heuristic in strip_shadowing_flags must skip just the
# flag for these, never the following token.
_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset(
{"--spec-default", "--jinja", "--no-jinja"}
)
def strip_shadowing_flags(
args: Iterable[str],
*,
strip_context: bool = True,
strip_cache: bool = True,
strip_spec: bool = True,
strip_template: bool = True,
) -> list[str]:
"""Strip flags that shadow first-class Studio settings.
Used when the route inherits a previous load's ``llama_extra_args``
so that an inherited ``-c 4096`` cannot override the current
request's ``max_seq_length`` (and equivalents for cache /
speculative / chat template). Each ``strip_*`` flag controls one
group; the route only strips groups whose corresponding first-class
field was actually supplied by the caller, so an inherited
``--chat-template-file`` survives an Apply that omits both
``llama_extra_args`` and ``chat_template_override``.
"""
shadowing: set[str] = set()
if strip_context:
shadowing |= _CONTEXT_FLAGS
if strip_cache:
shadowing |= _CACHE_FLAGS
if strip_spec:
shadowing |= _SPEC_FLAGS
if strip_template:
shadowing |= _TEMPLATE_FLAGS
tokens = [str(a) for a in (args or [])]
out: list[str] = []
i, n = 0, len(tokens)
while i < n:
tok = tokens[i]
flag = _flag_name(tok)
if flag is None or flag not in shadowing:
out.append(tok)
i += 1
continue
# Drop this token. Boolean shadowing flags never carry a value;
# other shadowing flags consume the next token when it isn't a
# flag and the value isn't already packed as ``--key=value``.
if flag in _BOOLEAN_SHADOWING_FLAGS or "=" in tok:
i += 1
elif i + 1 < n and _flag_name(tokens[i + 1]) is None:
i += 2
else:
i += 1
return out

4
studio/backend/core/inference/mlx_inference.py
View file

@ -116,11 +116,11 @@ class MLXInferenceBackend:
)
try:
from unsloth_zoo.mlx_loader import FastMLXModel
from unsloth_zoo.mlx.loader import FastMLXModel
except ImportError as e:
raise ImportError(
"Unsloth: MLX inference requires unsloth-zoo with the MLX modules "
"(unsloth_zoo.mlx_loader). Reinstall via install.sh on Apple Silicon."
"(unsloth_zoo.mlx.loader). Reinstall via install.sh on Apple Silicon."
) from e
model, tokenizer_or_processor = FastMLXModel.from_pretrained(

32
studio/backend/core/inference/providers.py
View file

@ -218,6 +218,28 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
# are always among the top regardless of the API's order.
"model_id_limit": 15,
},
"vllm": {
"display_name": "vLLM",
# User-supplied via provider_base_url; the route layer already falls
# back to the payload's base_url when the registry entry has none.
"base_url": "",
"default_models": [],
"supports_streaming": True,
"supports_vision": True,
"supports_tool_calling": True,
"auth_header": "Authorization",
"auth_prefix": "Bearer ",
# Force /v1/chat/completions in stream_chat_completion — vLLM's
# /v1/responses rebuilds messages and runs them through the loaded
# model's chat template, which 400s on strict-alternation templates
# (Gemma 3 raises "Conversation roles must alternate user/assistant
# /user/assistant/..."). The chat-completions path takes messages
# verbatim and avoids that template gauntlet.
"notes": "Self-hosted vLLM server. Always routed to /v1/chat/completions.",
# Surfaced through the frontend's CUSTOM_PROVIDER_PRESETS, not the
# /api/providers/registry dropdown — see list_available_providers.
"hidden": True,
},
"openrouter": {
"display_name": "OpenRouter",
"base_url": "https://openrouter.ai/api/v1",
@ -269,9 +291,17 @@ def get_base_url(provider_type: str) -> str | None:
def list_available_providers() -> list[dict[str, Any]]:
"""Return all registered providers (for the /registry endpoint)."""
"""Return all registered providers (for the /registry endpoint).
Hidden entries (``"hidden": True``) are filtered out — they exist in the
registry only for backend lookups (e.g. ``supports_vision`` for vLLM) and
are surfaced in the frontend via ``CUSTOM_PROVIDER_PRESETS`` instead of
the cloud-provider dropdown.
"""
result = []
for provider_type, info in PROVIDER_REGISTRY.items():
if info.get("hidden"):
continue
result.append(
{
"provider_type": provider_type,

3
studio/backend/core/training/trainer.py
View file

@ -3208,6 +3208,9 @@ class UnslothTrainer:
if eval_steps_val > 0:
config_args["eval_strategy"] = "steps"
config_args["eval_steps"] = eval_steps_val
config_args["per_device_eval_batch_size"] = config_args[
"per_device_train_batch_size"
]
logger.info(
f"✅ Evaluation enabled: eval_steps={eval_steps_val} (fraction of total steps)\n"
)

1
studio/backend/core/training/training.py
View file

@ -214,6 +214,7 @@ class TrainingBackend:
"max_steps": kwargs.get("max_steps", 0),
"save_steps": kwargs.get("save_steps", 0),
"weight_decay": kwargs.get("weight_decay", 0.001),
"max_grad_norm": kwargs.get("max_grad_norm", 0.0),
"random_seed": kwargs.get("random_seed", 3407),
"packing": kwargs.get("packing", False),
"optim": kwargs.get("optim", "adamw_8bit"),

96
studio/backend/core/training/worker.py
View file

@ -30,6 +30,7 @@ from utils.hardware import apply_gpu_ids
from utils.wheel_utils import (
direct_wheel_url,
flash_attn_wheel_url,
has_blackwell_gpu,
install_wheel,
probe_torch_wheel_env,
url_exists,
@ -313,6 +314,12 @@ def _should_try_runtime_flash_attn_install(max_seq_length: int) -> bool:
def _ensure_flash_attn_for_long_context(event_queue: Any, max_seq_length: int) -> None:
if not _should_try_runtime_flash_attn_install(max_seq_length):
return
if has_blackwell_gpu():
_send_status(
event_queue,
"Skipping flash-attn install: Blackwell GPU detected (sm_100+); no compatible prebuilt wheel",
)
return
installed = _install_package_wheel_first(
event_queue = event_queue,
@ -417,6 +424,55 @@ def _normalize_mlx_studio_scheduler(value):
return raw
def _resolve_mlx_local_dataset_files(file_paths: list) -> list[str]:
"""Resolve Studio local dataset uploads without importing the GPU trainer."""
from utils.paths import resolve_dataset_path
all_files: list[str] = []
for dataset_file in file_paths or []:
file_path = (
dataset_file
if os.path.isabs(dataset_file)
else str(resolve_dataset_path(dataset_file))
)
file_path_obj = Path(file_path)
if file_path_obj.is_dir():
parquet_dir = (
file_path_obj / "parquet-files"
if (file_path_obj / "parquet-files").exists()
else file_path_obj
)
parquet_files = sorted(parquet_dir.glob("*.parquet"))
if parquet_files:
all_files.extend(str(p) for p in parquet_files)
continue
candidates: list[Path] = []
for ext in (".json", ".jsonl", ".csv", ".parquet"):
candidates.extend(sorted(file_path_obj.glob(f"*{ext}")))
if candidates:
all_files.extend(str(c) for c in candidates)
continue
raise ValueError(f"No supported data files in directory: {file_path_obj}")
all_files.append(str(file_path_obj))
return all_files
def _mlx_local_dataset_loader_for_files(files: list[str]) -> str:
first_ext = Path(files[0]).suffix.lower()
if first_ext in (".json", ".jsonl"):
return "json"
if first_ext == ".csv":
return "csv"
if first_ext == ".parquet":
return "parquet"
raise ValueError(f"Unsupported dataset format: {files[0]}")
def _run_mlx_training(event_queue, stop_queue, config):
"""Self-contained MLX training path for Apple Silicon.
@ -442,8 +498,8 @@ def _run_mlx_training(event_queue, stop_queue, config):
import mlx.core as mx
try:
from unsloth_zoo.mlx_loader import FastMLXModel
from unsloth_zoo.mlx_trainer import (
from unsloth_zoo.mlx.loader import FastMLXModel
from unsloth_zoo.mlx.trainer import (
MLXTrainer,
MLXTrainingConfig,
train_on_responses_only,
@ -451,7 +507,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
except ImportError as e:
raise ImportError(
"Unsloth: MLX training requires unsloth-zoo with the MLX modules "
"(unsloth_zoo.mlx_loader / unsloth_zoo.mlx_trainer). Reinstall via "
"(unsloth_zoo.mlx.loader / unsloth_zoo.mlx.trainer). Reinstall via "
"install.sh on Apple Silicon."
) from e
from datasets import load_dataset
@ -572,7 +628,6 @@ def _run_mlx_training(event_queue, stop_queue, config):
return ds
def _load_local(file_paths):
from core.training.trainer import UnslothTrainer
from datasets import load_from_disk
if len(file_paths) == 1:
@ -581,10 +636,10 @@ def _run_mlx_training(event_queue, stop_queue, config):
(p / "dataset_info.json").exists() or (p / "state.json").exists()
):
return load_from_disk(str(p))
all_files = UnslothTrainer._resolve_local_files(file_paths)
all_files = _resolve_mlx_local_dataset_files(file_paths)
if not all_files:
raise ValueError("No local dataset files found")
loader = UnslothTrainer._loader_for_files(all_files)
loader = _mlx_local_dataset_loader_for_files(all_files)
return load_dataset(loader, data_files = all_files, split = "train")
if hf_dataset:
@ -718,6 +773,12 @@ def _run_mlx_training(event_queue, stop_queue, config):
else:
eval_steps_val = int(eval_steps_val)
# MLX: per-element clip to [-1, 1]; norm clip disabled (it needs a
# global reduction that breaks MLX's eager pipeline). 1.0 (not 5.0):
# |g_i| > 5 rarely fires, so the historical 5.0 was effectively no-op.
max_grad_norm = 0.0
max_grad_value = 1.0 # TODO: expose MLX grad-clip in Studio UI for power users
trainer = MLXTrainer(
model = model,
tokenizer = tokenizer,
@ -732,6 +793,8 @@ def _run_mlx_training(event_queue, stop_queue, config):
lr_scheduler_type = lr_scheduler_type,
optim = optim_name,
weight_decay = float(config.get("weight_decay", 0.001) or 0.001),
max_grad_norm = max_grad_norm,
max_grad_value = max_grad_value,
logging_steps = 1,
max_seq_length = max_seq_length,
seed = config.get("random_seed", 3407),
@ -820,7 +883,17 @@ def _run_mlx_training(event_queue, stop_queue, config):
# ── 9. Real-time progress callback ──
_send("status", status_message = f"Training {model_name}...")
def _on_step(step, total, loss, lr, tok_s, peak_gb, elapsed, num_tokens):
def _on_step(
step,
total,
loss,
lr,
tok_s,
peak_gb,
elapsed,
num_tokens,
grad_norm = None,
):
eta = (elapsed / step * (total - step)) if step > 0 else 0
_send(
"progress",
@ -831,7 +904,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
total_steps = total,
elapsed_seconds = elapsed,
eta_seconds = max(0, eta),
grad_norm = None,
grad_norm = grad_norm,
num_tokens = num_tokens,
eval_loss = None,
status_message = None,
@ -846,6 +919,11 @@ def _run_mlx_training(event_queue, stop_queue, config):
"train/tokens_per_sec": tok_s,
"train/peak_gb": peak_gb,
"train/num_tokens": num_tokens,
**(
{"train/grad_norm": grad_norm}
if grad_norm is not None
else {}
),
},
step = step,
)
@ -857,6 +935,8 @@ def _run_mlx_training(event_queue, stop_queue, config):
tb_writer.add_scalar("train/learning_rate", lr, step)
tb_writer.add_scalar("train/tokens_per_sec", tok_s, step)
tb_writer.add_scalar("train/peak_gb", peak_gb, step)
if grad_norm is not None:
tb_writer.add_scalar("train/grad_norm", grad_norm, step)
except Exception:
pass

86
studio/backend/models/inference.py
View file

@ -593,6 +593,92 @@ class ChatCompletionRequest(BaseModel):
None,
description = "[x-unsloth] Override base URL for the external provider.",
)
enable_prompt_caching: Optional[bool] = Field(
None,
description = (
"[x-unsloth] Opt in to provider-side prompt caching. On Anthropic, "
"attaches cache_control={type:ephemeral} to the system block so the "
"static prefix is reused across turns. On OpenAI cloud, caching is "
"automatic for prompts >=1024 tokens and this flag is informational. "
"Ignored for every other provider (mistral, gemini, kimi, openrouter, "
"vllm, local, etc.). Treated as enabled when omitted."
),
)
openai_code_exec_container_id: Optional[str] = Field(
None,
description = (
"[x-unsloth] OpenAI shell-tool container id from the prior response "
"in the same chat thread. When set and `code_execution` is in "
"`enabled_tools`, the next /v1/responses call uses "
"environment.type='container_reference' so filesystem state "
"persists across turns. Unset → environment.type='container_auto' "
"and OpenAI creates a fresh container. Only meaningful for the "
"OpenAI cloud + gpt-5.5 family path; ignored otherwise."
),
)
# ── OpenAI shell-tool container management ─────────────────────
class OpenAIContainerRequest(BaseModel):
"""
Shared body for the three OpenAI container endpoints (list / create
/ delete). Carries the encrypted API key + base URL so the route
handler can decrypt it and proxy to the user's OpenAI account.
Same pattern as the inference proxy endpoints — keeps the key off
persistent storage on the backend.
"""
encrypted_api_key: str = Field(
...,
description = "[x-unsloth] RSA-encrypted, base64-encoded OpenAI API key.",
)
provider_base_url: Optional[str] = Field(
None,
description = "[x-unsloth] OpenAI base URL. Only api.openai.com is supported; non-cloud bases are rejected with 400.",
)
class CreateOpenAIContainerBody(OpenAIContainerRequest):
name: str = Field(
...,
min_length = 1,
max_length = 256,
description = "Human-readable container name. Surfaces in the picker UI.",
)
ttl_minutes: int = Field(
20,
ge = 1,
le = 20,
description = (
"Idle-timeout TTL the new container will inherit (anchor="
"last_active_at). OpenAI hard-caps this at 20 minutes and "
"rejects larger values with integer_above_max_value."
),
)
class DeleteOpenAIContainerBody(OpenAIContainerRequest):
container_id: str = Field(
...,
description = "OpenAI container id (cntr_...) to delete.",
)
class OpenAIContainerSummary(BaseModel):
"""One row from GET /v1/containers, reshaped for the UI."""
id: str
name: Optional[str] = None
created_at: Optional[int] = None
last_active_at: Optional[int] = None
expires_after_minutes: Optional[int] = None
status: Optional[str] = None
class ListOpenAIContainersResponse(BaseModel):
containers: list[OpenAIContainerSummary]
# ── Streaming response chunks ────────────────────────────────────

10
studio/backend/models/providers.py
View file

@ -95,8 +95,9 @@ class ProviderModelsRequest(BaseModel):
"""Request to list models from an external provider."""
provider_type: str = Field(..., description = "Provider type from the registry")
encrypted_api_key: str = Field(
..., description = "RSA-encrypted, base64-encoded API key"
encrypted_api_key: Optional[str] = Field(
None,
description = "RSA-encrypted, base64-encoded API key (optional for local providers)",
)
base_url: Optional[str] = Field(
None, description = "Custom base URL (overrides registry default)"
@ -110,8 +111,9 @@ class ProviderTestRequest(BaseModel):
"""Request to test connectivity to an external provider."""
provider_type: str = Field(..., description = "Provider type from the registry")
encrypted_api_key: str = Field(
..., description = "RSA-encrypted, base64-encoded API key"
encrypted_api_key: Optional[str] = Field(
None,
description = "RSA-encrypted, base64-encoded API key (optional for local providers)",
)
base_url: Optional[str] = Field(
None, description = "Custom base URL (overrides registry default)"

5
studio/backend/models/training.py
View file

@ -262,6 +262,11 @@ class TrainingStartRequest(BaseModel):
max_steps: Optional[int] = Field(None, description = "Maximum training steps")
save_steps: int = Field(100, description = "Steps between checkpoints")
weight_decay: float = Field(0.001, description = "Weight decay")
max_grad_norm: float = Field(
0.0,
ge = 0,
description = "Global gradient norm clipping threshold. Set 0 to disable.",
)
random_seed: int = Field(42, description = "Random seed")
packing: bool = Field(False, description = "Enable sequence packing")
optim: str = Field("adamw_8bit", description = "Optimizer")

347
studio/backend/routes/inference.py
View file

@ -119,7 +119,10 @@ try:
_DEFAULT_T_MAX_PREDICT_MS,
detect_reasoning_flags,
)
from core.inference.llama_server_args import validate_extra_args
from core.inference.llama_server_args import (
strip_shadowing_flags,
validate_extra_args,
)
from utils.models import ModelConfig
from utils.inference import load_inference_config
from utils.models.model_config import load_model_defaults
@ -141,7 +144,10 @@ except ImportError:
_DEFAULT_T_MAX_PREDICT_MS,
detect_reasoning_flags,
)
from core.inference.llama_server_args import validate_extra_args
from core.inference.llama_server_args import (
strip_shadowing_flags,
validate_extra_args,
)
from utils.models import ModelConfig
from utils.inference import load_inference_config
from utils.models.model_config import load_model_defaults
@ -194,6 +200,11 @@ from models.inference import (
AnthropicResponseTextBlock,
AnthropicResponseToolUseBlock,
AnthropicUsage,
CreateOpenAIContainerBody,
DeleteOpenAIContainerBody,
ListOpenAIContainersResponse,
OpenAIContainerRequest,
OpenAIContainerSummary,
)
from core.inference.anthropic_compat import (
anthropic_messages_to_openai,
@ -401,6 +412,57 @@ def _validate_native_mmproj_companion(
) from exc
def _normalise_settings_str(value: Optional[str]) -> Optional[str]:
"""Lowercase + strip a settings string, mapping blank/None to None."""
if value is None:
return None
if isinstance(value, str):
stripped = value.strip().lower()
return stripped or None
return value
def _request_matches_loaded_settings(
request: LoadRequest, llama_backend: LlamaCppBackend
) -> bool:
"""True iff every runtime setting on the request matches the loaded
server. Caller has already checked model+variant+is_loaded. See #5401."""
# Compare requested n_ctx (not effective) so VRAM-cap doesn't mask
# an Auto-vs-explicit slider flip.
if request.max_seq_length != llama_backend.requested_n_ctx:
return False
if _normalise_settings_str(request.cache_type_kv) != _normalise_settings_str(
llama_backend.cache_type_kv
):
return False
# Vision loads silently drop speculative decoding (llama_cpp.py gates
# spec on ``not is_vision``), so treat the request as ``off`` against
# the backend's ``None`` to avoid forcing a redundant reload.
if llama_backend.is_vision:
req_spec = "off"
else:
req_spec = _normalise_settings_str(request.speculative_type) or "off"
backend_spec = _normalise_settings_str(llama_backend.speculative_type) or "off"
if req_spec != backend_spec:
return False
if (request.chat_template_override or None) != (
llama_backend.chat_template_override or None
):
return False
# llama_extra_args=None means "inherit"; only an explicit list that
# differs forces a reload. On the inherit path, refuse to match if
# stored extras contain any shadow flag, so the reload path can
# strip them instead of leaving a stale override in effect.
backend_extra = list(llama_backend.extra_args) if llama_backend.extra_args else []
if request.llama_extra_args is None:
if backend_extra and strip_shadowing_flags(backend_extra) != backend_extra:
return False
else:
if list(request.llama_extra_args) != backend_extra:
return False
return True
def _resolve_model_identifier_for_request(
request: LoadRequest | ValidateModelRequest,
*,
@ -456,6 +518,11 @@ async def load_model(
extra_llama_args = validate_extra_args(request.llama_extra_args)
except ValueError as exc:
raise HTTPException(status_code = 400, detail = str(exc))
# Re-narrow []-from-None back to None so the inheritance path
# below can tell "caller omitted" from "caller explicit []".
extra_llama_args: Optional[list[str]] = (
None if request.llama_extra_args is None else extra_llama_args
)
model_identifier, model_log_label, native_grant_backed = (
_resolve_model_identifier_for_request(request, operation = "load-model")
@ -474,6 +541,9 @@ async def load_model(
and llama_backend.hf_variant.lower() == request.gguf_variant.lower()
and llama_backend.model_identifier
and llama_backend.model_identifier.lower() == model_identifier.lower()
# Also require runtime settings to match so Apply changes
# aren't silently dropped (#5401).
and _request_matches_loaded_settings(request, llama_backend)
):
logger.info(
f"Model already loaded (GGUF): {model_log_label} variant={request.gguf_variant}, skipping reload"
@ -608,6 +678,70 @@ async def load_model(
)
unsloth_backend.unload_model(unsloth_backend.active_model_name)
# Inherit llama_extra_args from the previous load when the
# request omits the field (the chat-settings Apply path
# does not round-trip them; explicit [] still clears).
# Inheritance is gated on (model_identifier, hf_variant)
# to refuse cross-model pickup, and shadowing flags are
# stripped so an inherited override can't win the last-wins
# CLI parse against a freshly-supplied first-class field.
if request.llama_extra_args is None and llama_backend.extra_args:
source = llama_backend.extra_args_source
# Compare against the resolved variant, not the request
# field: callers commonly omit gguf_variant for local
# ``.gguf`` paths and HF auto-pick flows. ``config.gguf_
# variant`` is the variant load_model was actually
# invoked with (see the HF / local branches below), so
# both sides of the comparison key off the same string.
resolved_variant = config.gguf_variant
same_source = bool(
source
and source[0]
and source[0].lower() == model_identifier.lower()
and (source[1] or "").lower() == (resolved_variant or "").lower()
)
if not same_source:
logger.info(
"Not inheriting llama_extra_args: stored args came "
"from %s, loading %s",
source,
(model_identifier, resolved_variant),
)
# Cross-model: clear explicitly so the backend
# doesn't inherit via "no opinion" semantics.
extra_llama_args = []
else:
# Strip only the groups whose first-class field
# was actually set by the caller, so an inherited
# --chat-template-file survives an Apply that omits
# chat_template_override.
fields_set = getattr(request, "model_fields_set", set())
stripped = strip_shadowing_flags(
llama_backend.extra_args,
strip_context = "max_seq_length" in fields_set,
strip_cache = "cache_type_kv" in fields_set,
strip_spec = "speculative_type" in fields_set,
strip_template = "chat_template_override" in fields_set,
)
try:
extra_llama_args = validate_extra_args(stripped)
except ValueError:
# Should not happen on already-validated args; degrade
# to no-extras rather than 400 if managed flags changed.
logger.warning(
"Stored llama_extra_args failed revalidation; "
"loading without them: %s",
stripped,
)
extra_llama_args = []
else:
if extra_llama_args:
logger.info(
"Inheriting llama_extra_args from previous "
"load (same model, shadow-stripped): %s",
extra_llama_args,
)
# Route to HF mode or local mode based on config
# Run in a thread so the event loop stays free for progress
# polling and other requests during the (potentially long)
@ -640,6 +774,10 @@ async def load_model(
llama_backend.load_model,
gguf_path = config.gguf_file,
mmproj_path = config.gguf_mmproj_file,
# Pass the resolved variant so _extra_args_source
# is keyed off the same string the inheritance
# check at the top of /load uses (#5401 followup).
hf_variant = config.gguf_variant,
model_identifier = config.identifier,
is_vision = config.is_vision,
n_ctx = request.max_seq_length,
@ -1554,15 +1692,16 @@ async def _proxy_to_external_provider(
detail = f"Unknown provider type: {provider_type}",
)
# Decrypt the API key
try:
api_key = decrypt_api_key(payload.encrypted_api_key)
except Exception as exc:
logger.warning("external_provider.decrypt_failed", error = str(exc))
raise HTTPException(
status_code = 400,
detail = "Failed to decrypt API key. The server key may have changed — try refreshing the page.",
)
api_key = ""
if payload.encrypted_api_key:
try:
api_key = decrypt_api_key(payload.encrypted_api_key)
except Exception as exc:
logger.warning("external_provider.decrypt_failed", error = str(exc))
raise HTTPException(
status_code = 400,
detail = "Failed to decrypt API key. The server key may have changed — try refreshing the page.",
)
model = payload.external_model or payload.model
if model == "default":
@ -1595,6 +1734,9 @@ async def _proxy_to_external_provider(
top_k = payload.top_k,
enable_thinking = payload.enable_thinking,
reasoning_effort = payload.reasoning_effort,
enabled_tools = payload.enabled_tools,
enable_prompt_caching = payload.enable_prompt_caching,
openai_code_exec_container_id = payload.openai_code_exec_container_id,
stream = payload.stream,
)
try:
@ -1624,6 +1766,186 @@ async def _proxy_to_external_provider(
)
# ── OpenAI shell-tool container management ───────────────────────
def _resolve_openai_cloud_client(
body: OpenAIContainerRequest,
) -> ExternalProviderClient:
"""
Decrypt the API key + validate the base URL points at OpenAI cloud,
then build an ExternalProviderClient for the three container CRUD
endpoints below. The shell tool only exists on api.openai.com, so
rejecting non-cloud bases up front prevents confusing 404s on
ollama / llama.cpp / vLLM / custom presets.
"""
base_url = body.provider_base_url or get_base_url("openai")
if not base_url or "api.openai.com" not in base_url:
raise HTTPException(
status_code = 400,
detail = (
"OpenAI container management is only available on the "
"managed cloud (api.openai.com). The provider's base URL "
f"points at {base_url!r}."
),
)
try:
api_key = decrypt_api_key(body.encrypted_api_key)
except Exception as exc:
logger.warning("external_provider.decrypt_failed", error = str(exc))
raise HTTPException(
status_code = 400,
detail = "Failed to decrypt API key. The server key may have changed — try refreshing the page.",
)
return ExternalProviderClient(
provider_type = "openai",
base_url = base_url,
api_key = api_key,
)
def _summarize_container(raw: dict) -> OpenAIContainerSummary:
expires = raw.get("expires_after")
expires_minutes: Optional[int] = None
if isinstance(expires, dict):
minutes = expires.get("minutes")
if isinstance(minutes, int):
expires_minutes = minutes
return OpenAIContainerSummary(
id = str(raw.get("id") or ""),
name = raw.get("name"),
created_at = raw.get("created_at")
if isinstance(raw.get("created_at"), int)
else None,
last_active_at = raw.get("last_active_at")
if isinstance(raw.get("last_active_at"), int)
else None,
expires_after_minutes = expires_minutes,
status = raw.get("status") if isinstance(raw.get("status"), str) else None,
)
@router.post(
"/external/openai/containers/list",
response_model = ListOpenAIContainersResponse,
)
async def list_openai_containers(
body: OpenAIContainerRequest,
current_subject: str = Depends(get_current_subject),
) -> ListOpenAIContainersResponse:
"""List the user's OpenAI shell-tool containers."""
client = _resolve_openai_cloud_client(body)
try:
try:
raw = await client.list_openai_containers()
except httpx.HTTPStatusError as exc:
detail = exc.response.text[:500] if exc.response is not None else str(exc)
raise HTTPException(
status_code = exc.response.status_code if exc.response else 502,
detail = f"OpenAI rejected /containers list: {detail}",
)
except httpx.HTTPError as exc:
raise HTTPException(
status_code = 502,
detail = f"Failed to reach OpenAI: {exc}",
)
# OpenAI keeps expired containers in /v1/containers indefinitely
# with status="expired" — they're effectively dead but still
# listed. Hide them so the picker only shows usable containers.
return ListOpenAIContainersResponse(
containers = [
_summarize_container(c)
for c in raw
if isinstance(c, dict) and c.get("status") != "expired"
],
)
finally:
await client.close()
@router.post(
"/external/openai/containers/create",
response_model = OpenAIContainerSummary,
)
async def create_openai_container(
body: CreateOpenAIContainerBody,
current_subject: str = Depends(get_current_subject),
) -> OpenAIContainerSummary:
"""Create a named container with the user-chosen idle TTL."""
client = _resolve_openai_cloud_client(body)
try:
try:
raw = await client.create_openai_container(
name = body.name,
ttl_minutes = body.ttl_minutes,
)
except httpx.HTTPStatusError as exc:
detail = exc.response.text[:500] if exc.response is not None else str(exc)
raise HTTPException(
status_code = exc.response.status_code if exc.response else 502,
detail = f"OpenAI rejected /containers create: {detail}",
)
except httpx.HTTPError as exc:
raise HTTPException(
status_code = 502,
detail = f"Failed to reach OpenAI: {exc}",
)
if not isinstance(raw, dict):
raise HTTPException(
status_code = 502,
detail = "OpenAI returned an unexpected container payload.",
)
return _summarize_container(raw)
finally:
await client.close()
@router.post("/external/openai/containers/delete", status_code = 204)
async def delete_openai_container(
body: DeleteOpenAIContainerBody,
current_subject: str = Depends(get_current_subject),
) -> None:
"""Delete a named container by id."""
logger.info(
"openai_container_delete.request subject=%s container_id=%s base_url=%s",
current_subject,
body.container_id,
body.provider_base_url,
)
client = _resolve_openai_cloud_client(body)
try:
try:
await client.delete_openai_container(body.container_id)
logger.info(
"openai_container_delete.success container_id=%s",
body.container_id,
)
except httpx.HTTPStatusError as exc:
detail = exc.response.text[:500] if exc.response is not None else str(exc)
logger.warning(
"openai_container_delete.openai_rejected container_id=%s status=%s body=%s",
body.container_id,
exc.response.status_code if exc.response else None,
detail,
)
raise HTTPException(
status_code = exc.response.status_code if exc.response else 502,
detail = f"OpenAI rejected /containers delete: {detail}",
)
except httpx.HTTPError as exc:
logger.warning(
"openai_container_delete.transport_error container_id=%s error=%s",
body.container_id,
exc,
)
raise HTTPException(
status_code = 502,
detail = f"Failed to reach OpenAI: {exc}",
)
finally:
await client.close()
@router.post("/chat/completions")
async def openai_chat_completions(
payload: ChatCompletionRequest,
@ -1644,7 +1966,8 @@ async def openai_chat_completions(
- Other models → Unsloth/transformers via InferenceBackend
"""
# ── External provider routing ────────────────────────────────
if payload.encrypted_api_key and (payload.provider_id or payload.provider_type):
# encrypted_api_key is optional — local providers (llama.cpp / vLLM / Ollama) may run without auth.
if payload.provider_id or payload.provider_type:
return await _proxy_to_external_provider(payload, request)
llama_backend = get_llama_cpp_backend()

40
studio/backend/routes/providers.py
View file

@ -200,14 +200,18 @@ async def test_provider(
detail = f"Unknown provider type: {payload.provider_type}",
)
try:
api_key = decrypt_api_key(payload.encrypted_api_key)
except Exception as exc:
logger.warning("Failed to decrypt API key (%s): %s", type(exc).__name__, exc)
raise HTTPException(
status_code = 400,
detail = "Failed to decrypt API key. The public key may have changed — try refreshing the page.",
)
api_key = ""
if payload.encrypted_api_key:
try:
api_key = decrypt_api_key(payload.encrypted_api_key)
except Exception as exc:
logger.warning(
"Failed to decrypt API key (%s): %s", type(exc).__name__, exc
)
raise HTTPException(
status_code = 400,
detail = "Failed to decrypt API key. The public key may have changed — try refreshing the page.",
)
base_url = payload.base_url or info["base_url"]
client = ExternalProviderClient(
@ -265,14 +269,18 @@ async def list_provider_models(
detail = f"Unknown provider type: {payload.provider_type}",
)
try:
api_key = decrypt_api_key(payload.encrypted_api_key)
except Exception as exc:
logger.warning("Failed to decrypt API key (%s): %s", type(exc).__name__, exc)
raise HTTPException(
status_code = 400,
detail = "Failed to decrypt API key. The public key may have changed — try refreshing the page.",
)
api_key = ""
if payload.encrypted_api_key:
try:
api_key = decrypt_api_key(payload.encrypted_api_key)
except Exception as exc:
logger.warning(
"Failed to decrypt API key (%s): %s", type(exc).__name__, exc
)
raise HTTPException(
status_code = 400,
detail = "Failed to decrypt API key. The public key may have changed — try refreshing the page.",
)
if info.get("model_list_mode") == "curated":
return [

1
studio/backend/routes/training.py
View file

@ -215,6 +215,7 @@ async def start_training(
"max_steps": request.max_steps,
"save_steps": request.save_steps,
"weight_decay": request.weight_decay,
"max_grad_norm": request.max_grad_norm,
"random_seed": request.random_seed,
"packing": request.packing,
"optim": request.optim,

419
studio/backend/tests/test_anthropic_code_execution.py Normal file
View file

@ -0,0 +1,419 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Unit tests for Anthropic's server-side `code_execution_20250825` tool
translation in `_stream_anthropic`.
Covers:
- Request body: when ``enabled_tools=["code_execution"]``, the outbound
``tools`` array carries ``{"type": "code_execution_20250825", "name":
"code_execution"}`` and the ``anthropic-beta`` header includes
``code-execution-2025-08-25``.
- Combined request: ``enabled_tools=["web_search", "code_execution"]``
sends both tool entries; the beta header still merges the code-exec
flag onto whatever the registry contributed.
- SSE translation: a `bash_code_execution` server_tool_use +
`bash_code_execution_tool_result` pair emits one tool_start and one
tool_end ``_toolEvent`` chunk with the expected arguments and result.
- SSE translation: a `text_editor_code_execution` create + result emits
a tool_start with ``kind="text_editor"`` + parsed args, and tool_end
with ``"Created"`` (or ``"Updated"``) based on the ``is_file_update``
flag.
- Error path: a ``bash_code_execution_tool_result_error`` with
``error_code="container_expired"`` renders as ``"Error:
container_expired"`` in the tool_end ``result``.
"""
import asyncio
import json
import httpx
from core.inference import external_provider as ep_mod
from core.inference.external_provider import ExternalProviderClient
def _drive(coro):
return asyncio.new_event_loop().run_until_complete(coro)
async def _collect(agen):
out = []
async for line in agen:
out.append(line)
return out
def _mock_http_client(monkeypatch, handler):
transport = httpx.MockTransport(handler)
monkeypatch.setattr(ep_mod, "_http_client", httpx.AsyncClient(transport = transport))
def _make_client() -> ExternalProviderClient:
return ExternalProviderClient(
provider_type = "anthropic",
base_url = "https://api.anthropic.com/v1",
api_key = "sk-ant-test",
)
def _anthropic_sse(events: list[dict]) -> bytes:
chunks: list[str] = []
for event in events:
chunks.append(f"event: {event['type']}")
chunks.append(f"data: {json.dumps(event)}")
chunks.append("")
return ("\n".join(chunks) + "\n").encode("utf-8")
def _tool_events(lines: list[str]) -> list[dict]:
"""Extract `_toolEvent` payloads from emitted SSE data lines."""
out: list[dict] = []
for line in lines:
if not line.startswith("data:"):
continue
raw = line[len("data:") :].strip()
if not raw or raw == "[DONE]":
continue
try:
parsed = json.loads(raw)
except json.JSONDecodeError:
continue
if isinstance(parsed, dict) and "_toolEvent" in parsed:
out.append(parsed["_toolEvent"])
return out
def test_code_execution_tool_appended_to_request_body(monkeypatch):
captured: dict = {}
def handler(request: httpx.Request) -> httpx.Response:
captured["body"] = json.loads(request.content.decode("utf-8"))
captured["headers"] = dict(request.headers)
return httpx.Response(
200,
content = _anthropic_sse([{"type": "message_stop"}]),
headers = {"content-type": "text/event-stream"},
)
_mock_http_client(monkeypatch, handler)
async def run():
client = _make_client()
async for _ in client._stream_anthropic(
messages = [{"role": "user", "content": "compute 2 + 2"}],
model = "claude-opus-4-7",
temperature = 0.7,
top_p = 0.95,
max_tokens = 4096,
enabled_tools = ["code_execution"],
):
pass
await client.close()
_drive(run())
body = captured["body"]
tools = body.get("tools") or []
assert {
"type": "code_execution_20250825",
"name": "code_execution",
} in tools
# No web_search entry when only code_execution is enabled.
assert all(t.get("type") != "web_search_20250305" for t in tools)
# Beta header carries the documented flag.
beta_header = captured["headers"].get("anthropic-beta", "")
assert "code-execution-2025-08-25" in beta_header
def test_code_execution_with_web_search_sends_both_tools(monkeypatch):
captured: dict = {}
def handler(request: httpx.Request) -> httpx.Response:
captured["body"] = json.loads(request.content.decode("utf-8"))
captured["headers"] = dict(request.headers)
return httpx.Response(
200,
content = _anthropic_sse([{"type": "message_stop"}]),
headers = {"content-type": "text/event-stream"},
)
_mock_http_client(monkeypatch, handler)
async def run():
client = _make_client()
async for _ in client._stream_anthropic(
messages = [{"role": "user", "content": "look it up and chart it"}],
model = "claude-opus-4-7",
temperature = 0.7,
top_p = 0.95,
max_tokens = 4096,
enabled_tools = ["web_search", "code_execution"],
):
pass
await client.close()
_drive(run())
tools = captured["body"].get("tools") or []
tool_types = {t.get("type") for t in tools if isinstance(t, dict)}
assert "web_search_20250305" in tool_types
assert "code_execution_20250825" in tool_types
assert "code-execution-2025-08-25" in captured["headers"].get("anthropic-beta", "")
def test_no_code_execution_tool_when_pill_off(monkeypatch):
captured: dict = {}
def handler(request: httpx.Request) -> httpx.Response:
captured["body"] = json.loads(request.content.decode("utf-8"))
captured["headers"] = dict(request.headers)
return httpx.Response(
200,
content = _anthropic_sse([{"type": "message_stop"}]),
headers = {"content-type": "text/event-stream"},
)
_mock_http_client(monkeypatch, handler)
async def run():
client = _make_client()
async for _ in client._stream_anthropic(
messages = [{"role": "user", "content": "hi"}],
model = "claude-opus-4-7",
temperature = 0.7,
top_p = 0.95,
max_tokens = 4096,
):
pass
await client.close()
_drive(run())
tools = captured["body"].get("tools") or []
assert all(t.get("type") != "code_execution_20250825" for t in tools)
# Beta header must NOT mention code-execution when the tool isn't on
# — that flag is opt-in only.
assert "code-execution-2025-08-25" not in captured["headers"].get(
"anthropic-beta", ""
)
def test_bash_code_execution_emits_tool_start_and_end(monkeypatch):
sse_events = [
{"type": "message_start", "message": {"usage": {}}},
{
"type": "content_block_start",
"index": 0,
"content_block": {
"type": "server_tool_use",
"id": "srvtoolu_1",
"name": "bash_code_execution",
},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {
"type": "input_json_delta",
"partial_json": '{"command": "ls -la"}',
},
},
{"type": "content_block_stop", "index": 0},
{
"type": "content_block_start",
"index": 1,
"content_block": {
"type": "bash_code_execution_tool_result",
"tool_use_id": "srvtoolu_1",
"content": {
"type": "bash_code_execution_result",
"stdout": "total 24\ndrwxr-xr-x 2 user user 4096 Jan 1 12:00 .",
"stderr": "",
"return_code": 0,
},
},
},
{"type": "content_block_stop", "index": 1},
{"type": "message_stop"},
]
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(
200,
content = _anthropic_sse(sse_events),
headers = {"content-type": "text/event-stream"},
)
_mock_http_client(monkeypatch, handler)
async def run():
client = _make_client()
return await _collect(
client._stream_anthropic(
messages = [{"role": "user", "content": "list files"}],
model = "claude-opus-4-7",
temperature = 0.7,
top_p = 0.95,
max_tokens = 4096,
enabled_tools = ["code_execution"],
)
)
lines = _drive(run())
events = _tool_events(lines)
assert len(events) == 2
start, end = events
assert start["type"] == "tool_start"
assert start["tool_name"] == "code_execution"
assert start["tool_call_id"] == "srvtoolu_1"
assert start["arguments"] == {"kind": "bash", "command": "ls -la"}
assert end["type"] == "tool_end"
assert end["tool_call_id"] == "srvtoolu_1"
assert "total 24" in end["result"]
# Non-zero return_code not present, so no return_code line.
assert "return_code:" not in end["result"]
def test_text_editor_create_emits_kind_and_status(monkeypatch):
sse_events = [
{"type": "message_start", "message": {"usage": {}}},
{
"type": "content_block_start",
"index": 0,
"content_block": {
"type": "server_tool_use",
"id": "srvtoolu_2",
"name": "text_editor_code_execution",
},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {
"type": "input_json_delta",
"partial_json": (
'{"command": "create", "path": "new_file.txt", '
'"file_text": "hi"}'
),
},
},
{"type": "content_block_stop", "index": 0},
{
"type": "content_block_start",
"index": 1,
"content_block": {
"type": "text_editor_code_execution_tool_result",
"tool_use_id": "srvtoolu_2",
"content": {
"type": "text_editor_code_execution_result",
"is_file_update": False,
},
},
},
{"type": "content_block_stop", "index": 1},
{"type": "message_stop"},
]
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(
200,
content = _anthropic_sse(sse_events),
headers = {"content-type": "text/event-stream"},
)
_mock_http_client(monkeypatch, handler)
async def run():
client = _make_client()
return await _collect(
client._stream_anthropic(
messages = [{"role": "user", "content": "write a file"}],
model = "claude-opus-4-7",
temperature = 0.7,
top_p = 0.95,
max_tokens = 4096,
enabled_tools = ["code_execution"],
)
)
lines = _drive(run())
events = _tool_events(lines)
assert len(events) == 2
start, end = events
assert start["arguments"]["kind"] == "text_editor"
assert start["arguments"]["command"] == "create"
assert start["arguments"]["path"] == "new_file.txt"
assert end["result"] == "Created"
def test_code_execution_error_renders_error_code(monkeypatch):
sse_events = [
{"type": "message_start", "message": {"usage": {}}},
{
"type": "content_block_start",
"index": 0,
"content_block": {
"type": "server_tool_use",
"id": "srvtoolu_3",
"name": "bash_code_execution",
},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {
"type": "input_json_delta",
"partial_json": '{"command": "echo broken"}',
},
},
{"type": "content_block_stop", "index": 0},
{
"type": "content_block_start",
"index": 1,
"content_block": {
"type": "bash_code_execution_tool_result",
"tool_use_id": "srvtoolu_3",
"content": {
"type": "bash_code_execution_tool_result_error",
"error_code": "container_expired",
},
},
},
{"type": "content_block_stop", "index": 1},
{"type": "message_stop"},
]
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(
200,
content = _anthropic_sse(sse_events),
headers = {"content-type": "text/event-stream"},
)
_mock_http_client(monkeypatch, handler)
async def run():
client = _make_client()
return await _collect(
client._stream_anthropic(
messages = [{"role": "user", "content": "run it"}],
model = "claude-opus-4-7",
temperature = 0.7,
top_p = 0.95,
max_tokens = 4096,
enabled_tools = ["code_execution"],
)
)
lines = _drive(run())
events = _tool_events(lines)
assert len(events) == 2
end = events[1]
assert end["type"] == "tool_end"
assert end["result"] == "Error: container_expired"

326
studio/backend/tests/test_detect_mmproj_file.py Normal file
View file

@ -0,0 +1,326 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for :func:`utils.models.model_config.detect_mmproj_file` (#5347)."""
from __future__ import annotations
from pathlib import Path
import struct
from utils.models.model_config import (
_detect_family_token,
detect_mmproj_file,
mmproj_matches_model_family,
)
_GGUF_MAGIC = 0x46554747
def _gguf_with_general(path: Path, fields: dict) -> Path:
"""Write a minimal GGUF with only ``general.*`` string KVs."""
body = b""
for k, v in fields.items():
kb = k.encode("utf-8")
vb = v.encode("utf-8")
body += struct.pack("<Q", len(kb)) + kb
body += struct.pack("<I", 8) # STRING vtype
body += struct.pack("<Q", len(vb)) + vb
header = struct.pack("<IIQQ", _GGUF_MAGIC, 3, 0, len(fields))
path.parent.mkdir(parents = True, exist_ok = True)
path.write_bytes(header + body)
return path
def _touch(path: Path) -> Path:
path.parent.mkdir(parents = True, exist_ok = True)
path.write_bytes(b"")
return path
def test_returns_none_when_no_mmproj(tmp_path: Path):
model = _touch(tmp_path / "Qwen3.5-9B-Q4_K_M.gguf")
assert detect_mmproj_file(str(model)) is None
def test_single_matching_family_mmproj_picked(tmp_path: Path):
"""Single same-family projector: returned (historical behaviour)."""
model = _touch(tmp_path / "Qwen3.5-9B-Q4_K_M.gguf")
mmproj = _touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf")
assert detect_mmproj_file(str(model)) == str(mmproj.resolve())
def test_hf_style_unprefixed_mmproj_still_works(tmp_path: Path):
"""HF convention: weight + ``mmproj-F16.gguf`` sibling."""
model = _touch(tmp_path / "model.gguf")
mmproj = _touch(tmp_path / "mmproj-F16.gguf")
assert detect_mmproj_file(str(model)) == str(mmproj.resolve())
def test_blocks_single_cross_family_projector(tmp_path: Path):
"""#5347 core: Qwen weight + lone Gemma mmproj returns None."""
model = _touch(tmp_path / "Qwen3.5-9B-Q4_K_M.gguf")
_touch(tmp_path / "gemma-4-26B-A4B-it.mmproj-q8_0.gguf")
assert detect_mmproj_file(str(model)) is None
def test_picks_matching_family_among_mixed_candidates(tmp_path: Path):
"""Mixed Qwen + Gemma projectors: pick Qwen, drop Gemma."""
model = _touch(tmp_path / "Qwen3.5-9B-Q4_K_M.gguf")
qwen_mm = _touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf")
_touch(tmp_path / "gemma-4-26B-A4B-it.mmproj-q8_0.gguf")
assert detect_mmproj_file(str(model)) == str(qwen_mm.resolve())
def test_prefers_longest_prefix_within_same_family(tmp_path: Path):
"""Same family, different sizes: longest shared stem prefix wins."""
model = _touch(tmp_path / "Qwen3.5-35B-A3B-UD-Q4_K_L.gguf")
_touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf")
big_mm = _touch(tmp_path / "Qwen3.5-35B-A3B-BF16-mmproj.gguf")
assert detect_mmproj_file(str(model)) == str(big_mm.resolve())
def test_unrecognised_family_does_not_break_detection(tmp_path: Path):
"""Unknown model family must not return None on a sole candidate."""
model = _touch(tmp_path / "MyCustomBrand-7B-Q4_K_M.gguf")
mmproj = _touch(tmp_path / "MyCustomBrand-7B-BF16-mmproj.gguf")
assert detect_mmproj_file(str(model)) == str(mmproj.resolve())
def test_directory_path_returns_first_candidate(tmp_path: Path):
"""Directory path: no model stem to compare; legacy first-candidate."""
_touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf")
_touch(tmp_path / "gemma-4-26B-A4B-it.mmproj-q8_0.gguf")
result = detect_mmproj_file(str(tmp_path))
assert result is not None
assert "mmproj" in Path(result).name.lower()
def test_search_root_walk_still_works(tmp_path: Path):
"""Snapshot layout: weight in quant subdir, mmproj at snapshot root."""
snapshot = tmp_path / "snapshot"
weight = _touch(snapshot / "BF16" / "Qwen3.5-9B-BF16.gguf")
mmproj = _touch(snapshot / "Qwen3.5-9B-BF16-mmproj.gguf")
result = detect_mmproj_file(str(weight), search_root = str(snapshot))
assert result == str(mmproj.resolve())
# -- Family token detection: word-bounded matching ----------------------
def test_family_token_phi_does_not_match_sapphire():
"""``phi`` substring inside ``sapphire`` must not tag Phi."""
assert _detect_family_token("sapphire-7b-q4_k_m.gguf") is None
def test_family_token_yi_does_not_match_tinyish_names():
"""``yi`` must not cross letter boundaries (``yip``)."""
assert _detect_family_token("yip-7b.gguf") is None
assert _detect_family_token("yi-vl-6b.gguf") == "yi"
def test_family_token_mimo_does_not_match_mimosa():
"""``mimo`` must not tag ``mimosa``."""
assert _detect_family_token("mimosa-rosa-7b.gguf") is None
assert _detect_family_token("MiMo-VL-7B-RL-BF16.gguf") == "mimo"
def test_family_token_mistral_does_not_match_ministral():
"""Pin Mistral-derivative tagging."""
assert _detect_family_token("Ministral-3-8B-Instruct-2512-BF16.gguf") == "ministral"
assert _detect_family_token("Mistral-7B-Instruct-v0.3.gguf") == "mistral"
assert _detect_family_token("Magistral-Small-2506-BF16.gguf") == "magistral"
assert (
_detect_family_token("Devstral-Small-2-24B-Instruct-2512-BF16.gguf")
== "devstral"
)
def test_family_token_picks_leftmost_when_multiple_present():
"""Leftmost family token wins, not tuple order."""
assert _detect_family_token("llama-phi-merge.gguf") == "llama"
assert _detect_family_token("phi-llama-merge.gguf") == "phi"
assert _detect_family_token("llama3-3b-instruct.gguf") == "llama"
def test_family_token_new_families_recognised():
"""Catalogue-audit additions tag correctly."""
assert _detect_family_token("NVIDIA-Nemotron-3-Nano-Omni-30B.gguf") == "nemotron"
assert _detect_family_token("Kimi-K2.6-BF16.gguf") == "kimi"
assert _detect_family_token("Nanonets-OCR-s-BF16.gguf") == "nanonets"
assert _detect_family_token("Cosmos-Reason1-7B-BF16.gguf") == "cosmos"
assert _detect_family_token("Apriel-1.5-15b-Thinker-BF16.gguf") == "apriel"
assert _detect_family_token("LFM2.5-VL-1.6B-BF16.gguf") == "lfm"
# -- Cross-family rejection with the expanded token list ----------------
def test_blocks_cross_family_for_new_token_pair(tmp_path: Path):
"""Nemotron weight + lone Gemma projector returns None."""
model = _touch(
tmp_path / "NVIDIA-Nemotron-3-Nano-Omni-30B-A3B-Reasoning-MXFP4_MOE.gguf"
)
_touch(tmp_path / "gemma-4-26B-A4B-it.mmproj-q8_0.gguf")
assert detect_mmproj_file(str(model)) is None
def test_picks_devstral_mmproj_in_mixed_dir(tmp_path: Path):
"""Devstral weight + Devstral mmproj + a Qwen mmproj: pick Devstral."""
model = _touch(tmp_path / "Devstral-Small-2-24B-Instruct-2512-BF16.gguf")
dev_mm = _touch(tmp_path / "Devstral-Small-2-mmproj-bf16.gguf")
_touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf")
assert detect_mmproj_file(str(model)) == str(dev_mm.resolve())
# -- Launcher-level family guard ----------------------------------------
def test_mmproj_family_guard_blocks_cross_family():
assert (
mmproj_matches_model_family(
"/models/Qwen3.5-9B-Q4_K_M.gguf",
"/models/gemma-4-26B-A4B-it.mmproj-q8_0.gguf",
)
is False
)
def test_mmproj_family_guard_allows_same_family():
assert (
mmproj_matches_model_family(
"/models/Qwen3.5-9B-Q4_K_M.gguf",
"/models/Qwen3.5-9B-BF16-mmproj.gguf",
)
is True
)
def test_mmproj_family_guard_allows_generic_hf_mmproj():
"""No family token on the projector: wildcard."""
assert (
mmproj_matches_model_family(
"/models/Qwen3.5-9B-Q4_K_M.gguf",
"/models/mmproj-F16.gguf",
)
is True
)
def test_mmproj_family_guard_allows_unrecognised_model_family():
"""No family token on the model: wildcard."""
assert (
mmproj_matches_model_family(
"/models/Apriel-1.5-15b-Thinker-BF16.gguf",
"/models/mmproj-F16.gguf",
)
is True
)
# -- Metadata-primary pairing in detect_mmproj_file ---------------------
def test_metadata_url_match_picked_over_filename_lookalike(tmp_path: Path):
"""URL match beats a longer-prefix sibling."""
weight = _gguf_with_general(
tmp_path / "Qwen3.5-9B-Q4_K_M.gguf",
{
"general.architecture": "qwen2vl",
"general.type": "model",
"general.basename": "Qwen3.5",
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
},
)
# Closer filename prefix, wrong upstream.
_gguf_with_general(
tmp_path / "Qwen3.5-9B-mmproj-bf16.gguf",
{
"general.architecture": "clip",
"general.type": "mmproj",
"general.basename": "Qwen3.5",
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-1.5B",
},
)
# Matching upstream.
correct = _gguf_with_general(
tmp_path / "mmproj-BF16.gguf",
{
"general.architecture": "clip",
"general.type": "mmproj",
"general.basename": "Qwen3.5",
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
},
)
assert detect_mmproj_file(str(weight)) == str(correct.resolve())
def test_metadata_url_mismatch_dropped(tmp_path: Path):
"""Filenames match family but metadata disagrees: returns None."""
weight = _gguf_with_general(
tmp_path / "qwen-9b.gguf",
{
"general.architecture": "qwen2vl",
"general.type": "model",
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
},
)
_gguf_with_general(
tmp_path / "qwen-9b-mmproj.gguf",
{
"general.architecture": "clip",
"general.type": "mmproj",
"general.base_model.0.repo_url": "https://huggingface.co/google/gemma-3-9B",
},
)
assert detect_mmproj_file(str(weight)) is None
def test_metadata_identifies_mmproj_without_filename_hint(tmp_path: Path):
"""Projector named ``vision-projector.gguf`` discovered via header."""
weight = _gguf_with_general(
tmp_path / "Qwen3.5-9B.gguf",
{
"general.architecture": "qwen2vl",
"general.type": "model",
"general.basename": "Qwen3.5",
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
},
)
projector = _gguf_with_general(
tmp_path / "vision-projector.gguf",
{
"general.architecture": "clip",
"general.type": "mmproj",
"general.basename": "Qwen3.5",
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
},
)
assert detect_mmproj_file(str(weight)) == str(projector.resolve())
def test_metadata_score_outranks_filename_prefix(tmp_path: Path):
"""Score 100 (URL match) beats score 0 (long filename prefix)."""
weight = _gguf_with_general(
tmp_path / "Qwen3.5-9B-Q4_K_M.gguf",
{
"general.architecture": "qwen2vl",
"general.type": "model",
"general.basename": "Qwen3.5",
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
},
)
# Headerless: long shared stem, score 0.
_touch(tmp_path / "Qwen3.5-9B-Q4_K_M-mmproj.gguf")
# Headered: generic name, score 100.
correct = _gguf_with_general(
tmp_path / "mmproj-BF16.gguf",
{
"general.architecture": "clip",
"general.type": "mmproj",
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
},
)
assert detect_mmproj_file(str(weight)) == str(correct.resolve())

216
studio/backend/tests/test_gguf_metadata.py Normal file
View file

@ -0,0 +1,216 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for :mod:`utils.models.gguf_metadata`. Synthesise small GGUF
headers in tmp dirs so we never depend on real model files."""
from __future__ import annotations
import struct
from pathlib import Path
from typing import Iterable, Mapping
from utils.models.gguf_metadata import (
is_mmproj_by_metadata,
pairing_score,
read_gguf_general_metadata,
)
_GGUF_MAGIC = 0x46554747
_VTYPE_STRING = 8
_VTYPE_UINT32 = 4
_VTYPE_ARRAY = 9
def _enc_string(s: str) -> bytes:
b = s.encode("utf-8")
return struct.pack("<Q", len(b)) + b
def _enc_kv_string(key: str, value: str) -> bytes:
return _enc_string(key) + struct.pack("<I", _VTYPE_STRING) + _enc_string(value)
def _enc_kv_uint32(key: str, value: int) -> bytes:
return (
_enc_string(key) + struct.pack("<I", _VTYPE_UINT32) + struct.pack("<I", value)
)
def _enc_kv_string_array(key: str, values: Iterable[str]) -> bytes:
vals = list(values)
out = _enc_string(key) + struct.pack("<I", _VTYPE_ARRAY)
out += struct.pack("<I", _VTYPE_STRING) + struct.pack("<Q", len(vals))
for v in vals:
out += _enc_string(v)
return out
def _write_synthetic_gguf(
path: Path,
general_strings: Mapping[str, str],
*,
extra_uint32: Mapping[str, int] | None = None,
extra_string_arrays: Mapping[str, Iterable[str]] | None = None,
) -> Path:
"""Minimal GGUF: header + KV body, no tensors."""
extra_uint32 = extra_uint32 or {}
extra_string_arrays = extra_string_arrays or {}
kv_count = len(general_strings) + len(extra_uint32) + len(extra_string_arrays)
body = b""
for k, v in general_strings.items():
body += _enc_kv_string(k, v)
for k, v in extra_uint32.items():
body += _enc_kv_uint32(k, v)
for k, v in extra_string_arrays.items():
body += _enc_kv_string_array(k, v)
header = struct.pack(
"<IIQQ",
_GGUF_MAGIC,
3, # version
0, # tensor_count
kv_count,
)
path.parent.mkdir(parents = True, exist_ok = True)
path.write_bytes(header + body)
return path
# --- read_gguf_general_metadata ----------------------------------------
def test_returns_none_for_missing_file(tmp_path: Path):
assert read_gguf_general_metadata(str(tmp_path / "nope.gguf")) is None
def test_returns_none_for_non_gguf(tmp_path: Path):
p = tmp_path / "garbage.gguf"
p.write_bytes(b"not a gguf file at all, just bytes")
assert read_gguf_general_metadata(str(p)) is None
def test_extracts_general_string_fields(tmp_path: Path):
p = _write_synthetic_gguf(
tmp_path / "model.gguf",
{
"general.architecture": "qwen2vl",
"general.type": "model",
"general.basename": "Qwen3.5",
"general.organization": "Qwen",
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
"general.base_model.0.name": "Qwen3.5 9B",
"general.base_model.0.organization": "Qwen",
},
)
meta = read_gguf_general_metadata(str(p))
assert meta is not None
assert meta["general.architecture"] == "qwen2vl"
assert meta["general.basename"] == "Qwen3.5"
assert (
meta["general.base_model.0.repo_url"]
== "https://huggingface.co/Qwen/Qwen3.5-9B"
)
def test_skips_unrelated_fields_without_breaking(tmp_path: Path):
"""Skip unwanted arrays and uint32s without losing position."""
p = _write_synthetic_gguf(
tmp_path / "model.gguf",
{"general.basename": "Foo"},
extra_uint32 = {"qwen2vl.context_length": 32768},
extra_string_arrays = {"tokenizer.ggml.tokens": ["a", "bc", "def"]},
)
meta = read_gguf_general_metadata(str(p))
assert meta == {"general.basename": "Foo"}
def test_metadata_is_cached(tmp_path: Path):
"""Cache invalidates on size change."""
p = _write_synthetic_gguf(
tmp_path / "model.gguf",
{"general.basename": "First"},
)
first = read_gguf_general_metadata(str(p))
assert first == {"general.basename": "First"}
# Force size change so the (path, mtime, size) key invalidates.
_write_synthetic_gguf(
tmp_path / "model.gguf",
{"general.basename": "Second", "general.organization": "X"},
)
second = read_gguf_general_metadata(str(p))
assert second == {"general.basename": "Second", "general.organization": "X"}
# --- is_mmproj_by_metadata --------------------------------------------
def test_is_mmproj_by_metadata_signals():
assert is_mmproj_by_metadata({"general.type": "mmproj"}) is True
assert is_mmproj_by_metadata({"general.type": "MMProj"}) is True
assert is_mmproj_by_metadata({"general.type": "model"}) is False
assert is_mmproj_by_metadata({"general.basename": "foo"}) is None
assert is_mmproj_by_metadata({}) is None
assert is_mmproj_by_metadata(None) is None
# --- pairing_score -----------------------------------------------------
def test_pairing_score_base_model_url_match():
weight = {
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
}
mmproj = {
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
}
assert pairing_score(weight, mmproj) == 100
def test_pairing_score_base_model_url_mismatch():
weight = {
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
}
mmproj = {
"general.base_model.0.repo_url": "https://huggingface.co/google/gemma-3-9B",
}
assert pairing_score(weight, mmproj) == -1
def test_pairing_score_base_model_url_trailing_slash_normalised():
weight = {
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B/",
}
mmproj = {
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
}
assert pairing_score(weight, mmproj) == 100
def test_pairing_score_basename_plus_org_fallback():
weight = {
"general.basename": "Nanonets-Ocr-S",
"general.base_model.0.organization": "Nanonets",
}
mmproj = {
"general.basename": "Nanonets-Ocr-S",
"general.base_model.0.organization": "Nanonets",
}
assert pairing_score(weight, mmproj) == 80
def test_pairing_score_basename_only_fallback():
assert (
pairing_score(
{"general.basename": "Nanonets-Ocr-S"},
{"general.basename": "Nanonets-Ocr-S"},
)
== 60
)
def test_pairing_score_no_overlap_returns_zero():
"""One side empty: scorer punts to filename fallback."""
assert pairing_score({"general.basename": "Foo"}, {}) == 0
assert pairing_score({}, {"general.basename": "Foo"}) == 0
assert pairing_score(None, {"general.basename": "Foo"}) == 0

237
studio/backend/tests/test_gguf_reload_inheritance.py Normal file
View file

@ -0,0 +1,237 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Backend contract for the GGUF reload duplicate-load guard.
``LlamaCppBackend._already_in_target_state`` is the in-process
short-circuit that prevents a serialised duplicate /load from killing
the just-spawned llama-server. These tests pin the local-file
identity, the HF-mode hf_variant fallback, and the ``extra_args``
None-vs-[] inherit semantics so the guard cannot silently regress.
"""
from __future__ import annotations
import sys
import types as _types
from pathlib import Path
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
_structlog_stub = _types.ModuleType("structlog")
_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub")
sys.modules.setdefault("structlog", _structlog_stub)
_httpx_stub = _types.ModuleType("httpx")
for _exc in (
"ConnectError",
"TimeoutException",
"ReadTimeout",
"ReadError",
"RemoteProtocolError",
"CloseError",
):
setattr(_httpx_stub, _exc, type(_exc, (Exception,), {}))
_httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None})
_httpx_stub.Client = type(
"C",
(),
{
"__init__": lambda s, **kw: None,
"__enter__": lambda s: s,
"__exit__": lambda s, *a: None,
},
)
sys.modules.setdefault("httpx", _httpx_stub)
from core.inference.llama_cpp import LlamaCppBackend
class _FakeProcess:
"""Stand-in for subprocess.Popen so atexit cleanup doesn't crash."""
def terminate(self):
pass
def wait(self, timeout = None):
return 0
def kill(self):
pass
def poll(self):
return 0
def _loaded_backend(**overrides):
backend = LlamaCppBackend()
backend._process = _FakeProcess() # is_loaded only checks "is not None"
backend._healthy = True
backend._model_identifier = "owner/repo"
backend._hf_variant = "Q4_K_M"
backend._requested_n_ctx = 8192
backend._cache_type_kv = None
backend._speculative_type = None
backend._chat_template_override = None
backend._is_vision = False
backend._extra_args = None
backend._extra_args_source = None
backend._gguf_path = None
for key, value in overrides.items():
setattr(backend, key, value)
return backend
# ── Local-file identity via gguf_path ────────────────────────────────
def test_already_in_target_state_uses_gguf_path_when_present(tmp_path):
gguf_file = tmp_path / "model.Q4_K_M.gguf"
gguf_file.write_bytes(b"")
backend = _loaded_backend(
_hf_variant = "Q4_K_M",
_gguf_path = str(gguf_file),
)
assert (
backend._already_in_target_state(
gguf_path = str(gguf_file),
model_identifier = "owner/repo",
hf_variant = None,
n_ctx = 8192,
cache_type_kv = None,
speculative_type = None,
chat_template_override = None,
extra_args = None,
is_vision = False,
)
is True
)
def test_already_in_target_state_rejects_different_gguf_path(tmp_path):
a = tmp_path / "a.gguf"
a.write_bytes(b"")
b = tmp_path / "b.gguf"
b.write_bytes(b"")
backend = _loaded_backend(_gguf_path = str(a))
assert (
backend._already_in_target_state(
gguf_path = str(b),
model_identifier = "owner/repo",
hf_variant = None,
n_ctx = 8192,
cache_type_kv = None,
speculative_type = None,
chat_template_override = None,
extra_args = None,
is_vision = False,
)
is False
)
# ── HF mode falls back to hf_variant comparison ──────────────────────
def test_already_in_target_state_falls_back_to_hf_variant_for_hf_loads():
backend = _loaded_backend(_hf_variant = "Q4_K_M", _gguf_path = None)
assert (
backend._already_in_target_state(
gguf_path = None,
model_identifier = "owner/repo",
hf_variant = "Q8_0",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = None,
chat_template_override = None,
extra_args = None,
is_vision = False,
)
is False
)
def test_already_in_target_state_hf_same_variant_matches():
backend = _loaded_backend(_hf_variant = "Q4_K_M", _gguf_path = None)
assert (
backend._already_in_target_state(
gguf_path = None,
model_identifier = "owner/repo",
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = None,
chat_template_override = None,
extra_args = None,
is_vision = False,
)
is True
)
# ── extra_args: None inherits, [] forces reload, list enforces ───────
def test_already_in_target_state_none_extras_inherits_stored():
backend = _loaded_backend(_extra_args = ["--top-k", "20"])
assert (
backend._already_in_target_state(
gguf_path = None,
model_identifier = "owner/repo",
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = None,
chat_template_override = None,
extra_args = None,
is_vision = False,
)
is True
)
def test_already_in_target_state_empty_extras_forces_reload_when_stored():
backend = _loaded_backend(_extra_args = ["--top-k", "20"])
assert (
backend._already_in_target_state(
gguf_path = None,
model_identifier = "owner/repo",
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = None,
chat_template_override = None,
extra_args = [],
is_vision = False,
)
is False
)
def test_already_in_target_state_explicit_extras_match():
backend = _loaded_backend(_extra_args = ["--top-k", "20"])
assert (
backend._already_in_target_state(
gguf_path = None,
model_identifier = "owner/repo",
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = None,
chat_template_override = None,
extra_args = ["--top-k", "20"],
is_vision = False,
)
is True
)
def test_extra_args_source_default_is_none():
backend = LlamaCppBackend()
assert backend.extra_args_source is None

118
studio/backend/tests/test_llama_server_args.py
View file

@ -15,6 +15,7 @@ import pytest
from core.inference.llama_server_args import (
is_managed_flag,
strip_shadowing_flags,
validate_extra_args,
)
@ -187,3 +188,120 @@ def test_is_managed_flag_false_for_pass_through():
assert is_managed_flag("--flash-attn") is False
assert is_managed_flag("-ngl") is False
assert is_managed_flag("--threads") is False
# ── strip_shadowing_flags ─────────────────────────────────────────────
def test_strip_shadowing_flags_drops_context_when_requested():
out = strip_shadowing_flags(
["-c", "4096", "--top-k", "20"],
strip_context = True,
strip_cache = False,
strip_spec = False,
strip_template = False,
)
assert out == ["--top-k", "20"]
def test_strip_shadowing_flags_keeps_context_when_not_requested():
out = strip_shadowing_flags(
["-c", "4096", "--top-k", "20"],
strip_context = False,
strip_cache = False,
strip_spec = False,
strip_template = False,
)
assert out == ["-c", "4096", "--top-k", "20"]
def test_strip_shadowing_flags_keeps_chat_template_when_template_disabled():
# Caller did not supply chat_template_override; the inherited
# --chat-template-file must survive the strip.
out = strip_shadowing_flags(
["--chat-template-file", "/tmp/custom.jinja", "--top-k", "20"],
strip_context = True,
strip_cache = True,
strip_spec = True,
strip_template = False,
)
assert out == ["--chat-template-file", "/tmp/custom.jinja", "--top-k", "20"]
def test_strip_shadowing_flags_drops_template_when_requested():
out = strip_shadowing_flags(
["--chat-template-file", "/tmp/custom.jinja", "--top-k", "20"],
strip_template = True,
)
assert out == ["--top-k", "20"]
def test_strip_shadowing_flags_keeps_cache_when_cache_disabled():
out = strip_shadowing_flags(
["--cache-type-k", "q8_0", "--cache-type-v", "q8_0", "--top-k", "20"],
strip_cache = False,
)
assert out == [
"--cache-type-k",
"q8_0",
"--cache-type-v",
"q8_0",
"--top-k",
"20",
]
def test_strip_shadowing_flags_keeps_spec_when_spec_disabled():
out = strip_shadowing_flags(
["--spec-type", "ngram-mod", "--draft-min", "48", "--top-k", "20"],
strip_spec = False,
)
assert out == [
"--spec-type",
"ngram-mod",
"--draft-min",
"48",
"--top-k",
"20",
]
def test_strip_shadowing_flags_boolean_does_not_consume_next_token():
# --spec-default is a boolean shadowing flag; the value-skipping
# heuristic must skip just the flag, not the following positional.
out = strip_shadowing_flags(["--spec-default", "ngram-mod"], strip_spec = True)
assert out == ["ngram-mod"]
def test_strip_shadowing_flags_jinja_boolean_preserves_positional():
out = strip_shadowing_flags(["--jinja", "trailing-positional"], strip_template = True)
assert out == ["trailing-positional"]
def test_strip_shadowing_flags_no_jinja_boolean_preserves_positional():
out = strip_shadowing_flags(
["--no-jinja", "trailing-positional"], strip_template = True
)
assert out == ["trailing-positional"]
def test_strip_shadowing_flags_equals_form_drops_only_the_flag():
out = strip_shadowing_flags(["--ctx-size=4096", "--seed", "-1"], strip_context = True)
assert out == ["--seed", "-1"]
def test_strip_shadowing_flags_handles_none_input():
assert strip_shadowing_flags(None) == []
def test_strip_shadowing_flags_handles_empty_input():
assert strip_shadowing_flags([]) == []
def test_strip_shadowing_flags_defaults_strip_everything():
# The route's already-loaded comparator calls strip_shadowing_flags
# with no kwargs to detect ANY shadowing flag in stored extras.
out = strip_shadowing_flags(
["-c", "4096", "--cache-type-k", "q8_0", "--spec-default", "--jinja"]
)
assert out == []

9
studio/backend/tests/test_mlx_inference_backend.py
View file

@ -56,11 +56,14 @@ def _install_fake_fast_mlx(monkeypatch, calls):
return _DummyModel(), _DummyTokenizer()
unsloth_zoo_pkg = types.ModuleType("unsloth_zoo")
mlx_loader = types.ModuleType("unsloth_zoo.mlx_loader")
mlx_pkg = types.ModuleType("unsloth_zoo.mlx")
mlx_loader = types.ModuleType("unsloth_zoo.mlx.loader")
mlx_loader.FastMLXModel = _FastMLXModel
unsloth_zoo_pkg.mlx_loader = mlx_loader
unsloth_zoo_pkg.mlx = mlx_pkg
mlx_pkg.loader = mlx_loader
monkeypatch.setitem(sys.modules, "unsloth_zoo", unsloth_zoo_pkg)
monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx_loader", mlx_loader)
monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx", mlx_pkg)
monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.loader", mlx_loader)
def test_mlx_inference_text_load_forwards_studio_settings(monkeypatch):

1
studio/backend/tests/test_mlx_training_worker_config.py
View file

@ -37,6 +37,7 @@ def _load_worker_module():
for name in (
"direct_wheel_url",
"flash_attn_wheel_url",
"has_blackwell_gpu",
"install_wheel",
"probe_torch_wheel_env",
"url_exists",

391
studio/backend/tests/test_openai_code_execution.py Normal file
View file

@ -0,0 +1,391 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Unit tests for OpenAI's server-side `shell` tool translation in
`_stream_openai_responses`.
Covers:
- Request body: ``enabled_tools=["code_execution"]`` on the OpenAI
cloud base_url appends ``{"type": "shell", "environment": {"type":
"container_auto"}}`` to ``tools``.
- Container reuse: when ``openai_code_exec_container_id`` is provided,
the outgoing ``environment.type`` flips to ``"container_reference"``
and the id propagates.
- Cloud guard: code_execution on a non-cloud base_url (e.g. a local
OpenAI-compat preset / ollama / llama.cpp / vLLM) does NOT add the
shell tool, preventing a guaranteed 400 from those servers.
- SSE translation: a `shell_call` + `shell_call_output` pair emits one
``_toolEvent`` `tool_start` (`tool_name="code_execution"`,
`arguments.kind="bash"`) and one `tool_end` whose `result` contains
the joined stdout from the shell_call_output entries.
- Container surfacing: container_id captured from
`response.completed.container_id` is emitted as a synthetic
`container_ready` `_toolEvent` (only when it differs from the
inbound id).
- Stale-container handling: 400 with "container expired" body emits a
`container_invalidated` event before propagating the error.
"""
import asyncio
import json
import httpx
from core.inference import external_provider as ep_mod
from core.inference.external_provider import ExternalProviderClient
def _drive(coro):
return asyncio.new_event_loop().run_until_complete(coro)
async def _collect(agen):
out = []
async for line in agen:
out.append(line)
return out
def _mock_http_client(monkeypatch, handler):
transport = httpx.MockTransport(handler)
monkeypatch.setattr(ep_mod, "_http_client", httpx.AsyncClient(transport = transport))
def _make_client(base_url: str = "https://api.openai.com/v1") -> ExternalProviderClient:
return ExternalProviderClient(
provider_type = "openai",
base_url = base_url,
api_key = "sk-test",
)
def _openai_sse(events: list[dict]) -> bytes:
chunks: list[str] = []
for event in events:
chunks.append(f"event: {event['type']}")
chunks.append(f"data: {json.dumps(event)}")
chunks.append("")
return ("\n".join(chunks) + "\n").encode("utf-8")
def _tool_events(lines: list[str]) -> list[dict]:
out: list[dict] = []
for line in lines:
if not line.startswith("data:"):
continue
raw = line[len("data:") :].strip()
if not raw or raw == "[DONE]":
continue
try:
parsed = json.loads(raw)
except json.JSONDecodeError:
continue
if isinstance(parsed, dict) and "_toolEvent" in parsed:
out.append(parsed["_toolEvent"])
return out
def test_shell_tool_added_on_cloud_with_container_auto(monkeypatch):
captured: dict = {}
def handler(request: httpx.Request) -> httpx.Response:
captured["body"] = json.loads(request.content.decode("utf-8"))
return httpx.Response(
200,
content = _openai_sse([{"type": "response.completed", "response": {}}]),
headers = {"content-type": "text/event-stream"},
)
_mock_http_client(monkeypatch, handler)
async def run():
client = _make_client()
async for _ in client._stream_openai_responses(
messages = [{"role": "user", "content": "compute 2+2"}],
model = "gpt-5.5",
temperature = 0.7,
top_p = 0.95,
max_tokens = 4096,
enable_thinking = None,
reasoning_effort = None,
enabled_tools = ["code_execution"],
):
pass
await client.close()
_drive(run())
tools = captured["body"].get("tools") or []
assert {
"type": "shell",
"environment": {"type": "container_auto"},
} in tools
def test_shell_tool_uses_container_reference_when_id_supplied(monkeypatch):
captured: dict = {}
def handler(request: httpx.Request) -> httpx.Response:
captured["body"] = json.loads(request.content.decode("utf-8"))
return httpx.Response(
200,
content = _openai_sse([{"type": "response.completed", "response": {}}]),
headers = {"content-type": "text/event-stream"},
)
_mock_http_client(monkeypatch, handler)
async def run():
client = _make_client()
async for _ in client._stream_openai_responses(
messages = [{"role": "user", "content": "what did i write earlier"}],
model = "gpt-5.5",
temperature = 0.7,
top_p = 0.95,
max_tokens = 4096,
enable_thinking = None,
reasoning_effort = None,
enabled_tools = ["code_execution"],
openai_code_exec_container_id = "cntr_abc123",
):
pass
await client.close()
_drive(run())
tools = captured["body"].get("tools") or []
assert {
"type": "shell",
"environment": {
"type": "container_reference",
"container_id": "cntr_abc123",
},
} in tools
def test_shell_tool_refused_for_non_cloud_base_url(monkeypatch):
captured: dict = {}
def handler(request: httpx.Request) -> httpx.Response:
captured["body"] = json.loads(request.content.decode("utf-8"))
return httpx.Response(
200,
content = _openai_sse([{"type": "response.completed", "response": {}}]),
headers = {"content-type": "text/event-stream"},
)
_mock_http_client(monkeypatch, handler)
async def run():
client = _make_client(base_url = "http://localhost:11434/v1")
async for _ in client._stream_openai_responses(
messages = [{"role": "user", "content": "hi"}],
model = "gpt-5.5",
temperature = 0.7,
top_p = 0.95,
max_tokens = 4096,
enable_thinking = None,
reasoning_effort = None,
enabled_tools = ["code_execution"],
):
pass
await client.close()
_drive(run())
tools = captured["body"].get("tools") or []
# Shell tool must NOT leak to local OpenAI-compat servers — those
# 400 on the unknown tool type.
assert all(t.get("type") != "shell" for t in tools)
def test_shell_call_emits_tool_start_and_end(monkeypatch):
sse_events = [
{
"type": "response.output_item.added",
"item": {
"type": "shell_call",
"id": "scall_1",
"action": {"commands": ["ls -la"]},
},
},
{
"type": "response.output_item.done",
"item": {
"type": "shell_call",
"id": "scall_1",
"action": {"commands": ["ls -la"]},
"status": "completed",
},
},
{
"type": "response.output_item.done",
"item": {
"type": "shell_call_output",
"id": "scout_1",
"call_id": "scall_1",
"output": [
{
"stdout": "total 24\ndrwxr-xr-x .",
"stderr": "",
"outcome": {"type": "exit", "exit_code": 0},
}
],
},
},
{"type": "response.completed", "response": {}},
]
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(
200,
content = _openai_sse(sse_events),
headers = {"content-type": "text/event-stream"},
)
_mock_http_client(monkeypatch, handler)
async def run():
client = _make_client()
return await _collect(
client._stream_openai_responses(
messages = [{"role": "user", "content": "list files"}],
model = "gpt-5.5",
temperature = 0.7,
top_p = 0.95,
max_tokens = 4096,
enable_thinking = None,
reasoning_effort = None,
enabled_tools = ["code_execution"],
)
)
lines = _drive(run())
events = _tool_events(lines)
starts = [e for e in events if e["type"] == "tool_start"]
ends = [e for e in events if e["type"] == "tool_end"]
assert len(starts) == 1
assert len(ends) == 1
assert starts[0]["tool_name"] == "code_execution"
assert starts[0]["tool_call_id"] == "scall_1"
assert starts[0]["arguments"] == {"kind": "bash", "command": "ls -la"}
assert ends[0]["tool_call_id"] == "scall_1"
assert "total 24" in ends[0]["result"]
def test_container_ready_emitted_when_new_id_surfaces(monkeypatch):
sse_events = [
{
"type": "response.completed",
"response": {"container_id": "cntr_new_456"},
},
]
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(
200,
content = _openai_sse(sse_events),
headers = {"content-type": "text/event-stream"},
)
_mock_http_client(monkeypatch, handler)
async def run():
client = _make_client()
return await _collect(
client._stream_openai_responses(
messages = [{"role": "user", "content": "do stuff"}],
model = "gpt-5.5",
temperature = 0.7,
top_p = 0.95,
max_tokens = 4096,
enable_thinking = None,
reasoning_effort = None,
enabled_tools = ["code_execution"],
)
)
lines = _drive(run())
events = _tool_events(lines)
ready = [e for e in events if e["type"] == "container_ready"]
assert len(ready) == 1
assert ready[0]["container_id"] == "cntr_new_456"
def test_container_ready_not_emitted_when_id_unchanged(monkeypatch):
sse_events = [
{
"type": "response.completed",
"response": {"container_id": "cntr_same_789"},
},
]
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(
200,
content = _openai_sse(sse_events),
headers = {"content-type": "text/event-stream"},
)
_mock_http_client(monkeypatch, handler)
async def run():
client = _make_client()
return await _collect(
client._stream_openai_responses(
messages = [{"role": "user", "content": "do stuff"}],
model = "gpt-5.5",
temperature = 0.7,
top_p = 0.95,
max_tokens = 4096,
enable_thinking = None,
reasoning_effort = None,
enabled_tools = ["code_execution"],
openai_code_exec_container_id = "cntr_same_789",
)
)
lines = _drive(run())
events = _tool_events(lines)
# No churn — id matches the one already on the thread record.
assert not any(e["type"] == "container_ready" for e in events)
def test_stale_container_emits_invalidated(monkeypatch):
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(
400,
content = json.dumps(
{
"error": {
"message": "container has expired",
"type": "invalid_request_error",
}
}
).encode("utf-8"),
headers = {"content-type": "application/json"},
)
_mock_http_client(monkeypatch, handler)
async def run():
client = _make_client()
return await _collect(
client._stream_openai_responses(
messages = [{"role": "user", "content": "hi"}],
model = "gpt-5.5",
temperature = 0.7,
top_p = 0.95,
max_tokens = 4096,
enable_thinking = None,
reasoning_effort = None,
enabled_tools = ["code_execution"],
openai_code_exec_container_id = "cntr_stale_999",
)
)
lines = _drive(run())
events = _tool_events(lines)
invalidated = [e for e in events if e["type"] == "container_invalidated"]
assert len(invalidated) == 1

201
studio/backend/tests/test_openai_container_crud.py Normal file
View file

@ -0,0 +1,201 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Unit tests for the /v1/containers CRUD client methods.
Covers:
- All three calls (list / create / delete) send
``OpenAI-Beta: containers=v1``. Without it, OpenAI silently no-ops
the DELETE while still returning 200 ``{"deleted": true}``.
- ``delete_openai_container`` raises when the response body does not
report ``{"deleted": true}``, even on a 2xx response.
"""
from __future__ import annotations
import asyncio
import json
import httpx
import pytest
from core.inference import external_provider as ep_mod
from core.inference.external_provider import ExternalProviderClient
def _drive(coro):
return asyncio.new_event_loop().run_until_complete(coro)
def _mock_http_client(monkeypatch, handler):
"""Wire `handler` for both the shared `_http_client` AND any
per-call `httpx.AsyncClient(...)` instances. delete_openai_container
intentionally creates a fresh AsyncClient (see comment in
external_provider.delete_openai_container) so the test must
also intercept that constructor."""
transport = httpx.MockTransport(handler)
monkeypatch.setattr(ep_mod, "_http_client", httpx.AsyncClient(transport = transport))
real_async_client = httpx.AsyncClient
def _patched_async_client(*args, **kwargs):
kwargs["transport"] = transport
return real_async_client(*args, **kwargs)
monkeypatch.setattr(ep_mod.httpx, "AsyncClient", _patched_async_client)
def _make_client() -> ExternalProviderClient:
return ExternalProviderClient(
provider_type = "openai",
base_url = "https://api.openai.com/v1",
api_key = "sk-test",
)
def test_list_sends_openai_beta_header(monkeypatch):
seen: dict = {}
def handler(request: httpx.Request) -> httpx.Response:
seen["headers"] = dict(request.headers)
seen["url"] = str(request.url)
return httpx.Response(
200,
json = {"data": [{"id": "cntr_x", "name": "auto"}]},
)
_mock_http_client(monkeypatch, handler)
result = _drive(_make_client().list_openai_containers())
assert result == [{"id": "cntr_x", "name": "auto"}]
assert seen["headers"].get("openai-beta") == "containers=v1"
assert seen["url"] == "https://api.openai.com/v1/containers"
def test_create_sends_openai_beta_header(monkeypatch):
seen: dict = {}
def handler(request: httpx.Request) -> httpx.Response:
seen["headers"] = dict(request.headers)
seen["body"] = json.loads(request.content.decode("utf-8"))
return httpx.Response(200, json = {"id": "cntr_new", "name": "analysis"})
_mock_http_client(monkeypatch, handler)
result = _drive(
_make_client().create_openai_container(name = "analysis", ttl_minutes = 30)
)
assert result == {"id": "cntr_new", "name": "analysis"}
assert seen["headers"].get("openai-beta") == "containers=v1"
assert seen["body"]["name"] == "analysis"
assert seen["body"]["expires_after"] == {
"anchor": "last_active_at",
"minutes": 30,
}
def test_delete_sends_openai_beta_header_and_accepts_confirmation(monkeypatch):
seen: dict = {}
def handler(request: httpx.Request) -> httpx.Response:
seen["headers"] = dict(request.headers)
seen["url"] = str(request.url)
seen["method"] = request.method
return httpx.Response(
200,
json = {"id": "cntr_x", "object": "container.deleted", "deleted": True},
)
_mock_http_client(monkeypatch, handler)
_drive(_make_client().delete_openai_container("cntr_x"))
assert seen["method"] == "DELETE"
assert seen["url"] == "https://api.openai.com/v1/containers/cntr_x"
assert seen["headers"].get("openai-beta") == "containers=v1"
def test_delete_raises_when_response_lacks_deleted_true(monkeypatch):
"""OpenAI returns 200 ``{"deleted": true}`` even when the request is
silently rejected (e.g. before we started sending OpenAI-Beta).
Defensive guard: when the body omits ``deleted: true``, surface it
as an error so the UI can report the failure instead of falsely
reporting success."""
def handler(request: httpx.Request) -> httpx.Response:
# 200 but no deleted flag — simulate an unexpected payload shape.
return httpx.Response(200, json = {"id": "cntr_x", "object": "container"})
_mock_http_client(monkeypatch, handler)
with pytest.raises(httpx.HTTPError, match = "did not confirm container deletion"):
_drive(_make_client().delete_openai_container("cntr_x"))
def test_delete_raises_when_deleted_is_false(monkeypatch):
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(
200,
json = {"id": "cntr_x", "object": "container.deleted", "deleted": False},
)
_mock_http_client(monkeypatch, handler)
with pytest.raises(httpx.HTTPError, match = "did not confirm container deletion"):
_drive(_make_client().delete_openai_container("cntr_x"))
def test_delete_raises_when_body_is_not_json(monkeypatch):
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(200, content = b"<html>OK</html>")
_mock_http_client(monkeypatch, handler)
with pytest.raises(httpx.HTTPError, match = "did not confirm container deletion"):
_drive(_make_client().delete_openai_container("cntr_x"))
def test_delete_propagates_openai_4xx(monkeypatch):
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(404, json = {"error": {"message": "not found"}})
_mock_http_client(monkeypatch, handler)
with pytest.raises(httpx.HTTPStatusError):
_drive(_make_client().delete_openai_container("cntr_missing"))
def test_list_route_filters_expired_containers(monkeypatch):
"""OpenAI keeps containers in /v1/containers indefinitely with
status="expired" after their idle TTL passes — they can't be
used but still show up. The list route must drop them so the
picker only surfaces usable containers."""
from routes import inference as inf_mod
from models.inference import OpenAIContainerRequest
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(
200,
json = {
"data": [
{"id": "cntr_active", "name": "live", "status": "running"},
{"id": "cntr_dead", "name": "old", "status": "expired"},
{"id": "cntr_unknown", "name": "no-status"},
],
},
)
_mock_http_client(monkeypatch, handler)
def fake_resolve(_body):
return _make_client()
monkeypatch.setattr(inf_mod, "_resolve_openai_cloud_client", fake_resolve)
body = OpenAIContainerRequest(
encrypted_api_key = "enc",
provider_base_url = "https://api.openai.com/v1",
)
response = _drive(inf_mod.list_openai_containers(body, current_subject = "u"))
ids = [c.id for c in response.containers]
assert "cntr_active" in ids
assert "cntr_unknown" in ids # missing status is treated as usable
assert "cntr_dead" not in ids

62
studio/backend/tests/test_openai_responses_translation.py
View file

@ -286,6 +286,68 @@ def test_responses_reasoning_effort_included_when_requested(monkeypatch):
assert captured["body"]["reasoning"] == {"effort": "high", "summary": "auto"}
def test_responses_reasoning_summary_omitted_for_o3(monkeypatch):
captured: dict = {}
def handler(request: httpx.Request) -> httpx.Response:
captured["body"] = json.loads(request.content.decode("utf-8"))
return httpx.Response(
200,
content = _responses_sse([{"type": "response.completed", "response": {}}]),
headers = {"content-type": "text/event-stream"},
)
_mock_http_client(monkeypatch, handler)
async def run():
client = _make_client()
async for _ in client._stream_openai_responses(
messages = [{"role": "user", "content": "hi"}],
model = "o3",
temperature = 0.7,
top_p = 0.95,
max_tokens = None,
enable_thinking = None,
reasoning_effort = "high",
):
pass
await client.close()
_drive(run())
assert captured["body"]["reasoning"] == {"effort": "high"}
def test_responses_reasoning_summary_omitted_for_o3_with_enable_thinking(monkeypatch):
captured: dict = {}
def handler(request: httpx.Request) -> httpx.Response:
captured["body"] = json.loads(request.content.decode("utf-8"))
return httpx.Response(
200,
content = _responses_sse([{"type": "response.completed", "response": {}}]),
headers = {"content-type": "text/event-stream"},
)
_mock_http_client(monkeypatch, handler)
async def run():
client = _make_client()
async for _ in client._stream_openai_responses(
messages = [{"role": "user", "content": "hi"}],
model = "o3",
temperature = 0.7,
top_p = 0.95,
max_tokens = None,
enable_thinking = True,
reasoning_effort = None,
):
pass
await client.close()
_drive(run())
assert captured["body"]["reasoning"] == {"effort": "medium"}
def test_responses_reasoning_effort_none_omits_summary(monkeypatch):
captured: dict = {}

42
studio/backend/tests/test_training_raw_support.py
View file

@ -70,6 +70,48 @@ class TestTrainingRawSupport(unittest.TestCase):
self.assertTrue(config["load_in_4bit"])
self.assertEqual(config["embedding_learning_rate"], 1e-5)
def test_training_backend_forwards_grad_clipping_controls(self):
backend = TrainingBackend()
class DummyProcess:
pid = 12345
def start(self):
return None
class DummyThread:
def start(self):
return None
dummy_queue = object()
with (
patch(
"core.training.training.prepare_gpu_selection",
return_value = ([0], {"selection_mode": "auto"}),
),
patch(
"core.training.training._CTX.Queue",
side_effect = [dummy_queue, dummy_queue],
),
patch(
"core.training.training._CTX.Process", return_value = DummyProcess()
) as mock_process,
patch(
"core.training.training.threading.Thread",
return_value = DummyThread(),
),
):
backend.start_training(
job_id = "test-grad-clip",
model_name = "unsloth/test",
training_type = "LoRA/QLoRA",
max_grad_norm = 0.7,
)
config = mock_process.call_args.kwargs["kwargs"]["config"]
self.assertEqual(config["max_grad_norm"], 0.7)
def test_training_route_forwards_embedding_learning_rate(self):
training_route = _load_route_module(
"training_route_module_raw_support",

25
studio/backend/tests/test_training_worker_flash_attn.py
View file

@ -37,6 +37,7 @@ def test_runtime_flash_attn_prefers_prebuilt_wheel(monkeypatch):
statuses: list[str] = []
monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False)
monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: False)
monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import())
monkeypatch.setattr(
worker,
@ -65,6 +66,7 @@ def test_runtime_flash_attn_falls_back_to_pypi(monkeypatch):
statuses: list[str] = []
monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False)
monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: False)
monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import())
monkeypatch.setattr(
worker,
@ -112,6 +114,29 @@ def test_runtime_flash_attn_skip_env_avoids_all_install_work(monkeypatch):
worker._sp.run.assert_not_called()
def test_runtime_flash_attn_skips_on_blackwell(monkeypatch):
statuses: list[str] = []
install_mock = mock.Mock()
monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False)
monkeypatch.setattr(
worker, "_should_try_runtime_flash_attn_install", lambda max_seq: True
)
monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: True)
monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock)
monkeypatch.setattr(
worker,
"_send_status",
lambda queue, message: statuses.append(message),
)
worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 65536)
install_mock.assert_not_called()
assert len(statuses) == 1
assert "Blackwell" in statuses[0]
def test_causal_conv1d_fast_path_preserves_wheel_first_install_args(monkeypatch):
install_mock = mock.Mock(return_value = True)
monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock)

25
studio/backend/utils/datasets/chat_templates.py
View file

@ -28,6 +28,23 @@ DEFAULT_ALPACA_TEMPLATE = """Below is an instruction that describes a task, pair
{}"""
def _is_mlx_runtime() -> bool:
try:
from unsloth_zoo.mlx import is_mlx_available
except ImportError:
return False
return is_mlx_available()
def _chat_template_kwargs() -> dict:
if not _is_mlx_runtime():
return {}
return {
"patch_saving": False,
"use_zoo_tokenizer_patch": True,
}
def get_tokenizer_chat_template(tokenizer, model_name):
"""
Gets appropriate chat template for tokenizer based on model.
@ -60,6 +77,7 @@ def get_tokenizer_chat_template(tokenizer, model_name):
tokenizer = get_chat_template(
tokenizer,
chat_template = matched_template,
**_chat_template_kwargs(),
)
except Exception as e:
logger.info(f"⚠️ Failed to apply Unsloth template '{matched_template}': {e}")
@ -79,6 +97,7 @@ def get_tokenizer_chat_template(tokenizer, model_name):
tokenizer = get_chat_template(
tokenizer,
chat_template = "chatml",
**_chat_template_kwargs(),
)
except Exception as e:
logger.info(f"⚠️ Failed to apply default ChatML template: {e}")
@ -255,7 +274,11 @@ def apply_chat_template_to_dataset(
if not (hasattr(tokenizer, 'chat_template') and tokenizer.chat_template):
try:
from unsloth.chat_templates import get_chat_template
tokenizer = get_chat_template(tokenizer, chat_template = "alpaca")
tokenizer = get_chat_template(
tokenizer,
chat_template = "alpaca",
**_chat_template_kwargs(),
)
logger.info(f"📝 Set alpaca chat template on tokenizer for model saving")
except Exception as e:
logger.info(f"⚠️ Could not set alpaca template on tokenizer: {e}")

236
studio/backend/utils/models/gguf_metadata.py Normal file
View file

@ -0,0 +1,236 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Free-function ``general.*`` reader for GGUF headers, used by
``detect_mmproj_file`` to pair weights and projectors via
``general.base_model.0.repo_url``. ~30 ms per file, cached by
(path, mtime, size)."""
from __future__ import annotations
import os
import struct
import threading
from pathlib import Path
from typing import Dict, Optional, Tuple
from loggers import get_logger
logger = get_logger(__name__)
_GGUF_MAGIC = 0x46554747 # b"GGUF" LE u32
_WANTED_GENERAL_KEYS: frozenset[str] = frozenset(
{
"general.architecture",
"general.type",
"general.name",
"general.basename",
"general.organization",
"general.size_label",
"general.finetune",
"general.base_model.0.name",
"general.base_model.0.organization",
"general.base_model.0.repo_url",
"general.repo_url",
"general.source.url",
"general.source.repo_url",
"general.source.huggingface.repository",
}
)
# Cache failed parses too so a broken file is not retried each scan.
_CacheKey = Tuple[str, int, int]
_METADATA_CACHE: Dict[_CacheKey, Optional[Dict[str, str]]] = {}
_CACHE_LOCK = threading.Lock()
_CACHE_MAX_ENTRIES = 4096
def _cache_key(path: str) -> Optional[_CacheKey]:
try:
st = os.stat(path)
except OSError:
return None
try:
resolved = str(Path(path).resolve())
except OSError:
resolved = str(path)
return (resolved, st.st_mtime_ns, st.st_size)
def read_gguf_general_metadata(path: str) -> Optional[Dict[str, str]]:
"""Return ``general.*`` strings from a GGUF header, or ``None`` if
the file is missing, unreadable, or not a GGUF. ``{}`` means the
file is valid but carries none of the wanted keys."""
key = _cache_key(path)
if key is None:
return None
with _CACHE_LOCK:
if key in _METADATA_CACHE:
return _METADATA_CACHE[key]
result = _parse_gguf_header(path)
with _CACHE_LOCK:
# Arbitrary eviction; header reads are cheap so true LRU is overkill.
while len(_METADATA_CACHE) >= _CACHE_MAX_ENTRIES:
try:
_METADATA_CACHE.pop(next(iter(_METADATA_CACHE)))
except StopIteration:
break
_METADATA_CACHE[key] = result
return result
def _parse_gguf_header(path: str) -> Optional[Dict[str, str]]:
out: Dict[str, str] = {}
try:
with open(path, "rb") as f:
head = f.read(24)
if len(head) < 24:
return None
magic, _version, _tcount, kv_count = struct.unpack("<IIQQ", head)
if magic != _GGUF_MAGIC:
return None
for _ in range(kv_count):
try:
klen_bytes = f.read(8)
if len(klen_bytes) < 8:
break
klen = struct.unpack("<Q", klen_bytes)[0]
if klen > 1 << 20: # 1 MB sanity bound
break
kbytes = f.read(klen)
if len(kbytes) < klen:
break
key = kbytes.decode("utf-8", "replace")
vt_bytes = f.read(4)
if len(vt_bytes) < 4:
break
vtype = struct.unpack("<I", vt_bytes)[0]
if vtype == 8 and key in _WANTED_GENERAL_KEYS:
slen_bytes = f.read(8)
if len(slen_bytes) < 8:
break
slen = struct.unpack("<Q", slen_bytes)[0]
if slen > 1 << 22: # 4 MB sanity bound
break
sbytes = f.read(slen)
if len(sbytes) < slen:
break
out[key] = sbytes.decode("utf-8", "replace")
else:
if not _skip_gguf_value(f, vtype):
break
except (struct.error, UnicodeDecodeError):
break
except OSError as e:
logger.debug(f"read_gguf_general_metadata: cannot open {path}: {e}")
return None
except Exception as e:
logger.debug(f"read_gguf_general_metadata: parse failure on {path}: {e}")
return None
return out
# Strings (8) and arrays (9) are handled inline.
_FIXED_VTYPE_SIZES: Dict[int, int] = {
0: 1, # uint8
1: 1, # int8
2: 2, # uint16
3: 2, # int16
4: 4, # uint32
5: 4, # int32
6: 4, # float32
7: 1, # bool
10: 8, # uint64
11: 8, # int64
12: 8, # float64
}
def _skip_gguf_value(f, vtype: int) -> bool:
"""Advance past one GGUF value. ``f.seek(.., 1)`` past EOF is legal
on a regular file so truncation is detected on the next read; we
only return False for unknown types or sanity-bound overflow."""
if vtype == 8: # STRING
slen_bytes = f.read(8)
if len(slen_bytes) < 8:
return False
slen = struct.unpack("<Q", slen_bytes)[0]
if slen > 1 << 30: # 1 GB sanity bound
return False
f.seek(slen, 1)
return True
if vtype == 9: # ARRAY
head = f.read(12)
if len(head) < 12:
return False
atype, alen = struct.unpack("<IQ", head)
if alen > 1 << 30:
return False
if atype == 8:
for _ in range(alen):
slen_bytes = f.read(8)
if len(slen_bytes) < 8:
return False
slen = struct.unpack("<Q", slen_bytes)[0]
if slen > 1 << 30:
return False
f.seek(slen, 1)
return True
sz = _FIXED_VTYPE_SIZES.get(atype)
if sz is None:
return False
f.seek(sz * alen, 1)
return True
sz = _FIXED_VTYPE_SIZES.get(vtype)
if sz is None:
return False
f.seek(sz, 1)
return True
def is_mmproj_by_metadata(meta: Optional[Dict[str, str]]) -> Optional[bool]:
"""True/False from ``general.type``; None means fall back to filename."""
if not meta:
return None
t = meta.get("general.type")
if t is None:
return None
return t.lower() == "mmproj"
def pairing_score(
weight_meta: Optional[Dict[str, str]],
mmproj_meta: Optional[Dict[str, str]],
) -> int:
"""Pairing confidence: 100 = base_model URL match, 80 = basename + org,
60 = basename, -1 = definitive mismatch, 0 = decide from filename."""
if not weight_meta or not mmproj_meta:
return 0
w_url = weight_meta.get("general.base_model.0.repo_url")
p_url = mmproj_meta.get("general.base_model.0.repo_url")
if w_url and p_url:
return 100 if w_url.strip().rstrip("/") == p_url.strip().rstrip("/") else -1
w_base = weight_meta.get("general.basename")
p_base = mmproj_meta.get("general.basename")
w_org = weight_meta.get("general.base_model.0.organization") or weight_meta.get(
"general.organization"
)
p_org = mmproj_meta.get("general.base_model.0.organization") or mmproj_meta.get(
"general.organization"
)
if w_base and p_base and w_org and p_org:
if w_base.lower() == p_base.lower() and w_org.lower() == p_org.lower():
return 80
return -1
if w_base and p_base:
return 60 if w_base.lower() == p_base.lower() else -1
return 0

198
studio/backend/utils/models/model_config.py
View file

@ -19,6 +19,11 @@ from utils.paths import (
resolve_export_dir,
)
from utils.utils import without_hf_auth
from utils.models.gguf_metadata import (
is_mmproj_by_metadata,
pairing_score,
read_gguf_general_metadata,
)
import structlog
from loggers import get_logger
import os
@ -801,12 +806,15 @@ _AUDIO_TOKEN_PATTERNS = {
"whisper": lambda tokens: "<|startoftranscript|>" in tokens,
"audio_vlm": lambda tokens: "<audio_soft_token>" in tokens,
"bicodec": lambda tokens: any(t.startswith("<|bicodec_") for t in tokens),
"dac": lambda tokens: "<|audio_start|>" in tokens
and "<|audio_end|>" in tokens
and "<|text_start|>" in tokens
and "<|text_end|>" in tokens,
"snac": lambda tokens: sum(1 for t in tokens if t.startswith("<custom_token_"))
> 10000,
"dac": lambda tokens: (
"<|audio_start|>" in tokens
and "<|audio_end|>" in tokens
and "<|text_start|>" in tokens
and "<|text_end|>" in tokens
),
"snac": lambda tokens: (
sum(1 for t in tokens if t.startswith("<custom_token_")) > 10000
),
}
@ -913,6 +921,85 @@ def _is_mmproj(filename: str) -> bool:
return "mmproj" in filename.lower()
# Family tokens for #5347's filename fallback. Lowercase. Order does not
# matter (see ``_detect_family_token``).
_MODEL_FAMILY_TOKENS: tuple[str, ...] = (
"qwen",
"gemma",
"llama",
"mistral",
"ministral",
"magistral",
"devstral",
"phi",
"deepseek",
"internvl",
"minicpm",
"llava",
"glm",
"yi",
"command-r",
"molmo",
"pixtral",
"smolvlm",
"moondream",
"granite",
"ovis",
"nemotron",
"kimi",
"nanonets",
"cosmos",
"mimo",
"apriel",
"lfm",
)
# Word-bounded match: any letter on either side disqualifies. Stops
# ``phi`` matching ``sapphire``, ``yi`` matching ``tiny``, etc.
_FAMILY_TOKEN_RE_CACHE: Dict[str, "_re.Pattern[str]"] = {}
def _family_token_re(token: str) -> "_re.Pattern[str]":
pat = _FAMILY_TOKEN_RE_CACHE.get(token)
if pat is None:
pat = _re.compile(rf"(?:^|[^a-z])({_re.escape(token)})(?:[^a-z]|$)")
_FAMILY_TOKEN_RE_CACHE[token] = pat
return pat
def _detect_family_token(filename: str) -> Optional[str]:
"""Leftmost-position match; ties prefer the longer token."""
name = filename.lower()
best: Optional[tuple[int, int, str]] = None # (start, -len, token)
for token in _MODEL_FAMILY_TOKENS:
m = _family_token_re(token).search(name)
if m is None:
continue
key = (m.start(1), -len(token), token)
if best is None or key < best:
best = key
return None if best is None else best[2]
def mmproj_matches_model_family(model_path: str, mmproj_path: str) -> bool:
"""Defense-in-depth guard for the launcher: True unless both filenames
carry recognised family tokens that disagree."""
model_fam = _detect_family_token(Path(model_path).name)
mmproj_fam = _detect_family_token(Path(mmproj_path).name)
if model_fam is None or mmproj_fam is None:
return True
return model_fam == mmproj_fam
def _shared_prefix_len(a: str, b: str) -> int:
n = min(len(a), len(b))
for i in range(n):
if a[i] != b[i]:
return i
return n
def _is_gguf_filename(filename: str) -> bool:
return filename.lower().endswith(".gguf")
@ -927,33 +1014,18 @@ def _iter_gguf_files(directory: Path, recursive: bool = False):
def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional[str]:
"""
Find the mmproj (vision projection) GGUF file for a given model.
"""Find the mmproj GGUF for a model.
Args:
path: Directory to search — or a .gguf file (uses its parent dir
as the starting point).
search_root: Optional outer directory that should also be scanned
(and any directory between it and ``path``). This handles
local layouts where the model weights live in a quant-named
subdir (``snapshot/BF16/foo.gguf``) but the mmproj sits at
the snapshot root (``snapshot/mmproj-BF16.gguf``). When
``None``, only the immediate parent dir is scanned, matching
the historical behavior.
Returns:
Full path to the mmproj .gguf file, or None if not found.
"""
``path``: directory or a .gguf file. ``search_root``: optional ancestor
to also walk (snapshot layouts where the weight is in ``snapshot/BF16/``
but the projector sits at ``snapshot/``). Returns the projector path or
``None``."""
p = Path(path)
start_dir = p.parent if p.is_file() else p
if not start_dir.is_dir():
return None
# Build the list of dirs to scan: immediate dir first, then walk up
# to (and including) ``search_root`` if it is an ancestor. We walk
# incrementally rather than recursing into ``search_root`` so we
# don't accidentally pick up an mmproj from a sibling subdir
# belonging to a different model variant.
# Walk incrementally so a sibling subdir's mmproj cannot leak in.
seen: set[Path] = set()
scan_order: list[Path] = []
@ -969,12 +1041,7 @@ def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional
_add(start_dir)
# When ``path`` is a symlink (e.g. Ollama's ``.studio_links/...gguf``
# -> ``blobs/sha256-...``), the symlink's parent directory rarely
# contains the mmproj sibling; the real mmproj file lives next to
# the symlink target. Add the target's parent to the scan so vision
# GGUFs that are surfaced via symlinks are still recognised as
# vision models.
# Ollama's .studio_links/foo.gguf -> blobs/sha256-...: also scan target dir.
try:
if p.is_symlink() and p.is_file():
target_parent = p.resolve().parent
@ -986,14 +1053,12 @@ def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional
try:
root_resolved = Path(search_root).resolve()
start_resolved = start_dir.resolve()
# Only walk if start_dir is inside (or equal to) search_root.
if root_resolved == start_resolved or (
start_resolved.is_relative_to(root_resolved)
if hasattr(start_resolved, "is_relative_to")
else str(start_resolved).startswith(str(root_resolved) + "/")
):
cur = start_resolved
# Walk up from start_dir to (and including) root_resolved.
while cur != root_resolved and cur.parent != cur:
cur = cur.parent
_add(cur)
@ -1002,11 +1067,66 @@ def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional
except OSError:
pass
candidates: list[Path] = []
seen_resolved: set[Path] = set()
for d in scan_order:
for f in _iter_gguf_files(d):
if _is_mmproj(f.name):
return str(f.resolve())
return None
try:
resolved = f.resolve()
except OSError:
continue
if resolved in seen_resolved:
continue
# Prefer ``general.type=='mmproj'``; fall back to filename.
meta = read_gguf_general_metadata(str(resolved))
by_meta = is_mmproj_by_metadata(meta)
if by_meta is True or (by_meta is None and _is_mmproj(f.name)):
seen_resolved.add(resolved)
candidates.append(resolved)
if not candidates:
return None
# Directory path: no model name to compare against; legacy behaviour.
if not p.is_file():
return str(candidates[0])
# Stage 1: GGUF metadata. Stage 2: filename family token (#5347).
model_stem = p.stem.lower()
model_family = _detect_family_token(p.name)
weight_meta = read_gguf_general_metadata(str(p))
scored: list[tuple[int, Path]] = []
for c in candidates:
cand_meta = read_gguf_general_metadata(str(c))
meta_score = pairing_score(weight_meta, cand_meta)
if meta_score == -1:
logger.info(f"detect_mmproj_file: dropped {c.name} (metadata mismatch)")
continue
if meta_score == 0 and model_family is not None:
# Unrecognised candidate family is a wildcard (``mmproj-F16.gguf``).
cand_family = _detect_family_token(c.name)
if cand_family is not None and cand_family != model_family:
logger.info(
f"detect_mmproj_file: dropped {c.name} "
f"(filename family {cand_family!r} vs model {model_family!r})"
)
continue
scored.append((meta_score, c))
if not scored:
return None
# Score first, then longest shared prefix, then shorter stem.
best = max(
scored,
key = lambda sc: (
sc[0],
_shared_prefix_len(model_stem, sc[1].stem.lower()),
-len(sc[1].stem),
),
)
return str(best[1])
def detect_gguf_model(path: str) -> Optional[str]:
@ -1360,7 +1480,7 @@ def detect_gguf_model_remote(
if attempt < 2:
time.sleep(2**attempt)
logger.warning(
f"Could not check GGUF files for '{repo_id}' after 3 attempts: " f"{last_err}"
f"Could not check GGUF files for '{repo_id}' after 3 attempts: {last_err}"
)
return None

44
studio/backend/utils/wheel_utils.py
View file

@ -3,6 +3,7 @@
from __future__ import annotations
import functools
import json
import logging
import platform
@ -22,6 +23,49 @@ FLASH_ATTN_RELEASE_BASE_URL = (
)
@functools.lru_cache(maxsize = 1)
def has_blackwell_gpu() -> bool:
"""Return True if any visible NVIDIA GPU has compute capability >= 10.0
(Blackwell: sm_100, sm_120, sm_121, ...).
Dao-AILab does not publish prebuilt flash-attention wheels for these
architectures, and the older-arch wheels fail to load on Blackwell, so
callers use this gate to skip the flash-attn install/upgrade path.
Result is cached for the process lifetime since GPU hardware does not
change. Tests that mock subprocess/nvidia-smi must call
``has_blackwell_gpu.cache_clear()`` before each invocation.
"""
exe = shutil.which("nvidia-smi")
if not exe:
return False
try:
result = subprocess.run(
[exe, "--query-gpu=compute_cap", "--format=csv,noheader"],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 10,
env = child_env_without_native_path_secret(),
)
except (OSError, subprocess.TimeoutExpired):
return False
if result.returncode != 0:
return False
for line in result.stdout.splitlines():
cap = line.strip()
if not cap:
continue
major_part = cap.split(".", 1)[0]
try:
major = int(major_part)
except ValueError:
continue
if major >= 10:
return True
return False
def linux_wheel_platform_tag() -> str | None:
machine = platform.machine().lower()
if sys.platform.startswith("linux"):

1231
studio/frontend/package-lock.json generated
View file

File diff suppressed because it is too large Load diff

10
studio/frontend/package.json
View file

@ -18,8 +18,6 @@
"dependencies": {
"@assistant-ui/core": "0.1.17",
"@assistant-ui/react": "0.12.28",
"@assistant-ui/react-markdown": "0.12.11",
"@assistant-ui/react-streamdown": "0.1.11",
"@assistant-ui/tap": "0.5.10",
"@base-ui/react": "^1.2.0",
"@dagrejs/dagre": "^2.0.4",
@ -30,13 +28,11 @@
"@hugeicons/core-free-icons": "^4.1.1",
"@hugeicons/react": "^1.1.5",
"@huggingface/hub": "^2.9.0",
"@langchain/core": "^1.1.27",
"@radix-ui/react-checkbox": "^1.3.3",
"@radix-ui/react-label": "^2.1.8",
"@radix-ui/react-select": "^2.2.6",
"@radix-ui/react-separator": "^1.1.8",
"@radix-ui/react-slot": "^1.2.4",
"@streamdown/cjk": "1.0.3",
"@streamdown/code": "1.1.1",
"@streamdown/math": "1.0.2",
"@streamdown/mermaid": "1.0.2",
@ -50,21 +46,18 @@
"@tauri-apps/plugin-process": "^2.3.1",
"@tauri-apps/plugin-updater": "^2.10.1",
"@toolwind/corner-shape": "^0.0.8-3",
"@types/canvas-confetti": "^1.9.0",
"@xyflow/react": "^12.10.0",
"assistant-stream": "0.3.12",
"canvas-confetti": "^1.9.4",
"class-variance-authority": "^0.7.1",
"clsx": "^2.1.1",
"cmdk": "^1.1.1",
"date-fns": "^4.1.0",
"dexie": "^4.3.0",
"js-yaml": "^4.1.1",
"katex": "^0.16.28",
"lucide-react": "^1.7.0",
"mammoth": "^1.11.0",
"motion": "^12.34.0",
"next": "^16.1.6",
"next-themes": "^0.4.6",
"node-forge": "^1.4.0",
"radix-ui": "^1.4.3",
@ -73,7 +66,6 @@
"react-dom": "^19.2.4",
"react-resizable-panels": "^4.6.4",
"recharts": "3.7.0",
"remark-gfm": "^4.0.1",
"shadcn": "^4.2.0",
"sonner": "^2.0.7",
"streamdown": "2.5.0",
@ -92,6 +84,7 @@
"devDependencies": {
"@biomejs/biome": "^1.9.4",
"@eslint/js": "^9.39.1",
"@types/canvas-confetti": "^1.9.0",
"@types/js-yaml": "^4.0.9",
"@types/node-forge": "^1.3.14",
"@types/node": "^25.5.2",
@ -102,7 +95,6 @@
"eslint-plugin-react-hooks": "^7.0.1",
"eslint-plugin-react-refresh": "^0.5.2",
"globals": "^17.4.0",
"playwright": "^1.59.1",
"typescript": "~5.9.3",
"typescript-eslint": "^8.55.0",
"vite": "^8.0.1"

1
studio/frontend/public/provider-logos/llama_cpp.svg Normal file
View file

@ -0,0 +1 @@
<svg xmlns="http://www.w3.org/2000/svg" xml:space="preserve" viewBox="0 0 512 512"><path d="m356.4 201.3-32.8 58.3c-43.3-33.3-107.4-38.2-150.7-2.4-69.8 57.6-64.9 190.8 43.7 191.6 30.4 0 56.2-14.3 83.9-23.8l14.6 58.1c-24.6 11.4-49.6 23.1-76.6 26.7-246 33.5-231.9-321.6-9.5-340.1 46.7-3.9 87.8 8.3 127.6 31.6zm-169.9-55.9c-37.4 11.2-72.2 31.8-98.5 60.8-4.9-58.8 8.3-177.7 73.7-201 9.7-3.4 43-11.9 42.1 5.3-1 17.3-24.1 46.9-29.7 63-9.7 28.2-.7 47.6 12.6 72.2zm92.4 252.8h-36.5v-41.3h-41.3v-34h37.7l3.6-3.6v-40.1h36.5V323h38.9v34h-38.9zm133.7-41.3v41.3h-36.5v-41.3h-38.9v-34h38.9v-43.8h36.5v40.1l3.6 3.6h37.7v34h-41.3zM305.4 31.4c4.9 7.3-22.6 38.7-27 46.7-12.6 23.8-4.1 37.4 5.3 60-27.5-4.1-53-.7-80.2 2.4C209.6 88.3 239 12.2 305.4 31.4" style="fill:#ff8236"/></svg>
Side by side

After

Width:  |  Height:  |  Size: 763 B

14
studio/frontend/public/provider-logos/ollama.svg Normal file
View file

File diff suppressed because one or more lines are too long
Side by side

After

Width:  |  Height:  |  Size: 8.6 KiB

1
studio/frontend/public/provider-logos/vllm.svg Normal file
View file

@ -0,0 +1 @@
<svg version="1.1" viewBox="0.0 0.0 96.0 96.0" fill="none" stroke="none" stroke-linecap="square" stroke-miterlimit="10" xmlns:xlink="http://www.w3.org/1999/xlink" xmlns="http://www.w3.org/2000/svg"><clipPath id="g31e21232314_0_33.0"><path d="m0 0l96.0 0l0 96.0l-96.0 0l0 -96.0z" clip-rule="nonzero"/></clipPath><g clip-path="url(#g31e21232314_0_33.0)"><path fill="#d9d9d9" d="m41.04961 80.271324l1.8897629 0l0 2.3307114l-1.8897629 0z" fill-rule="evenodd"/><path fill="#d9d9d9" d="m42.221855 81.45145l1.8897629 0l0 2.3307037l-1.8897629 0z" fill-rule="evenodd"/><g filter="url(#shadowFilter-g31e21232314_0_33.1)"><use xlink:href="#g31e21232314_0_33.1" transform="matrix(1.0 0.0 0.0 1.0 0.0 2.0)"/></g><defs><filter id="shadowFilter-g31e21232314_0_33.1" filterUnits="userSpaceOnUse"><feGaussianBlur in="SourceAlpha" stdDeviation="2.0" result="blur"/><feComponentTransfer in="blur" color-interpolation-filters="sRGB"><feFuncR type="linear" slope="0" intercept="0.0"/><feFuncG type="linear" slope="0" intercept="0.0"/><feFuncB type="linear" slope="0" intercept="0.0"/><feFuncA type="linear" slope="0.0" intercept="0"/></feComponentTransfer></filter></defs><g id="g31e21232314_0_33.1"><path fill="#d9d9d9" d="m42.22417 28.470434l0 55.307083l-27.653543 -55.307083z" fill-rule="evenodd"/></g><g filter="url(#shadowFilter-g31e21232314_0_33.2)"><use xlink:href="#g31e21232314_0_33.2" transform="matrix(1.0 0.0 0.0 1.0 0.0 2.0)"/></g><defs><filter id="shadowFilter-g31e21232314_0_33.2" filterUnits="userSpaceOnUse"><feGaussianBlur in="SourceAlpha" stdDeviation="2.0" result="blur"/><feComponentTransfer in="blur" color-interpolation-filters="sRGB"><feFuncR type="linear" slope="0" intercept="0.0"/><feFuncG type="linear" slope="0" intercept="0.0"/><feFuncB type="linear" slope="0" intercept="0.0"/><feFuncA type="linear" slope="0.0" intercept="0"/></feComponentTransfer></filter></defs><g id="g31e21232314_0_33.2"><path fill="#d9d9d9" d="m42.223038 83.77752l21.729656 0l18.653545 -70.385826l-25.574802 13.461943z" fill-rule="evenodd"/></g><g filter="url(#shadowFilter-g31e21232314_0_33.3)"><use xlink:href="#g31e21232314_0_33.3" transform="matrix(1.0 0.0 0.0 1.0 0.0 2.0)"/></g><defs><filter id="shadowFilter-g31e21232314_0_33.3" filterUnits="userSpaceOnUse"><feGaussianBlur in="SourceAlpha" stdDeviation="2.0" result="blur"/><feComponentTransfer in="blur" color-interpolation-filters="sRGB"><feFuncR type="linear" slope="0" intercept="0.0"/><feFuncG type="linear" slope="0" intercept="0.0"/><feFuncB type="linear" slope="0" intercept="0.0"/><feFuncA type="linear" slope="0.0" intercept="0"/></feComponentTransfer></filter></defs><g id="g31e21232314_0_33.3"><path fill="#fdb515" d="m41.0477 27.293962l0 55.30709l-27.653542 -55.30709z" fill-rule="evenodd"/></g><g filter="url(#shadowFilter-g31e21232314_0_33.4)"><use xlink:href="#g31e21232314_0_33.4" transform="matrix(1.0 0.0 0.0 1.0 0.0 2.0)"/></g><defs><filter id="shadowFilter-g31e21232314_0_33.4" filterUnits="userSpaceOnUse"><feGaussianBlur in="SourceAlpha" stdDeviation="2.0" result="blur"/><feComponentTransfer in="blur" color-interpolation-filters="sRGB"><feFuncR type="linear" slope="0" intercept="0.0"/><feFuncG type="linear" slope="0" intercept="0.0"/><feFuncB type="linear" slope="0" intercept="0.0"/><feFuncA type="linear" slope="0.0" intercept="0"/></feComponentTransfer></filter></defs><g id="g31e21232314_0_33.4"><path fill="#30a2ff" d="m41.046566 82.60105l21.72966 0l18.653545 -70.385826l-25.574806 13.461943z" fill-rule="evenodd"/></g></g></svg>
Side by side

After

Width:  |  Height:  |  Size: 3.4 KiB

16
studio/frontend/src/components/assistant-ui/model-selector.tsx
View file

@ -10,10 +10,12 @@ import {
} from "@/components/ui/popover";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { usePlatformStore } from "@/config/env";
import { isCustomProviderType } from "@/features/chat/external-providers";
import { cn } from "@/lib/utils";
import {
ArrowDown01Icon,
CloudIcon,
DashboardSquare01Icon,
FolderSearchIcon,
Logout01Icon,
Search01Icon,
@ -40,6 +42,9 @@ const PROVIDER_LOGO_EXT: Record<string, "svg" | "png" | "jpg"> = {
kimi: "jpg",
qwen: "png",
openrouter: "svg",
vllm: "svg",
ollama: "svg",
llama_cpp: "svg",
};
function providerLogoSrc(providerType: string | undefined): string | undefined {
@ -59,6 +64,17 @@ function ExternalProviderLogo({
title?: string;
}) {
const src = providerLogoSrc(providerType);
if (!src && isCustomProviderType(providerType)) {
return (
<span title={title} aria-hidden={true} className="inline-flex shrink-0">
<HugeiconsIcon
icon={DashboardSquare01Icon}
className={cn("shrink-0", className)}
/>
</span>
);
}
if (!src) return null;
return (
<img

67
studio/frontend/src/components/assistant-ui/thread.tsx
View file

@ -13,6 +13,7 @@ import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning";
import { Sources, SourcesGroup } from "@/components/assistant-ui/sources";
import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
import { ToolGroup } from "@/components/assistant-ui/tool-group";
import { CodeExecutionToolUI } from "@/components/assistant-ui/tool-ui-code-execution";
import { PythonToolUI } from "@/components/assistant-ui/tool-ui-python";
import { TerminalToolUI } from "@/components/assistant-ui/tool-ui-terminal";
import { WebSearchToolUI } from "@/components/assistant-ui/tool-ui-web-search";
@ -327,6 +328,9 @@ const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => {
autoFocus={!disabled}
disabled={disabled}
aria-label="Message input"
// dir="auto": browser picks LTR/RTL from the first strong char;
// no effect on Latin / CJK / Devanagari.
dir="auto"
{...inputProps}
/>
<ComposerAction
@ -496,6 +500,9 @@ const ReasoningToggle: FC = () => {
externalSelection != null
? externalProviders.find((p) => p.id === externalSelection.providerId)
: undefined;
const isKimiExternal = selectedExternalProvider?.providerType === "kimi";
const toolsEnabled = useChatRuntimeStore((s) => s.toolsEnabled);
const setToolsEnabled = useChatRuntimeStore((s) => s.setToolsEnabled);
const effectiveExternalModelId =
selectedExternalProvider?.providerType === "openrouter" &&
externalSelection?.modelId === "openrouter/free" &&
@ -507,6 +514,10 @@ const ReasoningToggle: FC = () => {
? getExternalReasoningCapabilities(
selectedExternalProvider?.providerType,
effectiveExternalModelId,
{
isReasoningProvider:
selectedExternalProvider?.isReasoningModel === true,
},
)
: null;
const effectiveReasoningStyle =
@ -587,6 +598,11 @@ const ReasoningToggle: FC = () => {
setReasoningEffort(level);
setReasoningEnabled(true);
applyQwenThinkingParams(true);
// Kimi's $web_search builtin forbids thinking, so
// enabling thinking flips the Search pill off.
if (isKimiExternal && toolsEnabled) {
setToolsEnabled(false);
}
}}
>
{formatEffortLabel(level)}
@ -613,6 +629,11 @@ const ReasoningToggle: FC = () => {
const next = !reasoningEnabled;
setReasoningEnabled(next);
applyQwenThinkingParams(next);
// Mutual exclusion with the Search pill on Kimi — see the
// dropdown branch above and shared-composer for the same rule.
if (isKimiExternal && next && toolsEnabled) {
setToolsEnabled(false);
}
}}
className="composer-pill-btn"
data-active={
@ -680,16 +701,44 @@ const WebSearchToggle: FC = () => {
const modelLoaded = useChatRuntimeStore(
(s) => !!s.params.checkpoint && !s.modelLoading,
);
const checkpoint = useChatRuntimeStore((s) => s.params.checkpoint);
const supportsTools = useChatRuntimeStore((s) => s.supportsTools);
// External providers (OpenAI today) expose a server-side web_search tool
// even when the local tool runtime is unavailable — gate the Search pill
// on either source so it lights up on external models too. Mirror of
// shared-composer's searchDisabled.
const supportsBuiltinWebSearch = useChatRuntimeStore(
(s) => s.supportsBuiltinWebSearch,
);
const toolsEnabled = useChatRuntimeStore((s) => s.toolsEnabled);
const setToolsEnabled = useChatRuntimeStore((s) => s.setToolsEnabled);
const disabled = !(modelLoaded && supportsTools);
const setReasoningEnabled = useChatRuntimeStore((s) => s.setReasoningEnabled);
const externalProviders = useExternalProvidersStore((s) => s.providers);
const externalSelection = parseExternalModelId(checkpoint);
const selectedExternalProvider =
externalSelection != null
? externalProviders.find((p) => p.id === externalSelection.providerId)
: undefined;
const isKimiExternal = selectedExternalProvider?.providerType === "kimi";
const disabled =
!modelLoaded || !(supportsTools || supportsBuiltinWebSearch);
return (
<button
type="button"
disabled={disabled}
onClick={() => setToolsEnabled(!toolsEnabled)}
onClick={() => {
const next = !toolsEnabled;
setToolsEnabled(next);
// Kimi's $web_search builtin requires thinking=disabled (see
// https://platform.kimi.ai/docs/guide/use-web-search). Keep
// the two pills mutually exclusive so the visible state always
// matches what the backend ends up sending.
if (isKimiExternal) {
setReasoningEnabled(!next);
applyQwenThinkingParams(!next);
}
}}
className="composer-pill-btn"
data-active={toolsEnabled && !disabled ? "true" : "false"}
aria-label={toolsEnabled ? "Disable web search" : "Enable web search"}
@ -705,9 +754,18 @@ const CodeToolsToggle: FC = () => {
(s) => !!s.params.checkpoint && !s.modelLoading,
);
const supportsTools = useChatRuntimeStore((s) => s.supportsTools);
// External providers have no local tool runtime, but Anthropic's
// Claude 4.x dispatches code_execution_20250825 server-side. The
// chat-page resolver stashes that capability in the runtime store
// (next to supportsBuiltinWebSearch). Mirror of shared-composer's
// codeDisabled so this pill lights up in active threads too.
const supportsBuiltinCodeExecution = useChatRuntimeStore(
(s) => s.supportsBuiltinCodeExecution,
);
const codeToolsEnabled = useChatRuntimeStore((s) => s.codeToolsEnabled);
const setCodeToolsEnabled = useChatRuntimeStore((s) => s.setCodeToolsEnabled);
const disabled = !(modelLoaded && supportsTools);
const disabled =
!modelLoaded || !(supportsTools || supportsBuiltinCodeExecution);
return (
<button
@ -901,6 +959,7 @@ const AssistantMessage: FC = () => {
web_search: WebSearchToolUI,
python: PythonToolUI,
terminal: TerminalToolUI,
code_execution: CodeExecutionToolUI,
},
Fallback: ToolFallback,
},
@ -1105,6 +1164,8 @@ const EditComposer: FC = () => {
<ComposerPrimitive.Input
className="aui-edit-composer-input min-h-14 w-full resize-none bg-transparent p-4 text-foreground text-sm font-[450] outline-none"
autoFocus={true}
// See main composer above for the dir="auto" rationale.
dir="auto"
{...inputProps}
/>
<div className="aui-edit-composer-footer mx-3 mb-3 flex items-center gap-2 self-end">

126
studio/frontend/src/components/assistant-ui/tool-ui-code-execution.tsx Normal file
View file

@ -0,0 +1,126 @@
// 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 { type ToolCallMessagePartComponent, useAuiState } from "@assistant-ui/react";
import { FileTextIcon, LoaderIcon, TerminalIcon } from "lucide-react";
import { memo, useEffect, useState } from "react";
import {
ToolFallbackContent,
ToolFallbackRoot,
ToolFallbackTrigger,
} from "./tool-fallback";
/**
* Renders the synthetic `_toolEvent` chunks emitted by
* `_stream_anthropic` when Anthropic's `code_execution_20250825` tool
* fires. The backend collapses Anthropic's two sub-tools
* (`bash_code_execution`, `text_editor_code_execution`) into a single
* `tool_name: "code_execution"`, with `arguments.kind` ("bash" or
* "text_editor") and a per-kind argument shape:
*
* kind=bash: { command: "<shell command>" }
* kind=text_editor: { command: "view"|"create"|"str_replace", path, ... }
*
* The `result` payload is preformatted text:
* - bash: stdout, then "--- stderr ---" block + return_code if non-zero
* - text_editor view: file contents verbatim
* - text_editor create: "Created <path>" / "Updated <path>"
* - text_editor str_replace: unified-diff `lines` joined with "\n"
* - error: "Error: <error_code>"
*/
interface CodeExecutionArgs {
kind?: "bash" | "text_editor";
command?: string;
path?: string;
}
const CodeExecutionToolUIImpl: ToolCallMessagePartComponent = ({
args,
result,
status,
}) => {
const parsedArgs = (args as CodeExecutionArgs) ?? {};
const kind = parsedArgs.kind ?? "bash";
const command = parsedArgs.command ?? "";
const path = parsedArgs.path ?? "";
const isRunning = status?.type === "running";
let runningLabel: string;
let completedLabel: string;
let Icon = TerminalIcon;
if (kind === "text_editor") {
Icon = FileTextIcon;
if (command === "view") {
runningLabel = path ? `Viewing ${path}…` : "Viewing file…";
completedLabel = path ? `Viewed ${path}` : "Viewed file";
} else if (command === "create") {
runningLabel = path ? `Writing ${path}…` : "Writing file…";
completedLabel = path ? `Wrote ${path}` : "Wrote file";
} else if (command === "str_replace") {
runningLabel = path ? `Editing ${path}…` : "Editing file…";
completedLabel = path ? `Edited ${path}` : "Edited file";
} else {
runningLabel = "Running file operation…";
completedLabel = "File operation";
}
} else {
runningLabel = "Running command…";
completedLabel = command ? `Ran \`${command}\`` : "Ran command";
}
// Collapse the card once the model has resumed streaming prose after
// the tool call. Mirrors WebSearchToolUI's behavior so the tool-card
// doesn't crowd the final answer once the run is done.
const hasText = useAuiState(({ message }) =>
message.content.some(
(p) =>
p.type === "text" &&
"text" in p &&
(p as { text: string }).text.length > 0,
),
);
const [open, setOpen] = useState(isRunning);
useEffect(() => {
if (isRunning) {
setOpen(true);
} else if (hasText) {
setOpen(false);
}
}, [isRunning, hasText]);
const resultText =
typeof result === "string"
? result
: result != null
? JSON.stringify(result, null, 2)
: "";
return (
<ToolFallbackRoot open={open} onOpenChange={setOpen}>
<ToolFallbackTrigger
toolName={isRunning ? runningLabel : completedLabel}
status={status}
icon={Icon}
/>
<ToolFallbackContent>
{isRunning ? (
<div className="flex items-center gap-2 text-sm text-muted-foreground">
<LoaderIcon className="size-3.5 animate-spin" />
<span>{runningLabel}</span>
</div>
) : resultText ? (
<pre className="max-h-64 overflow-auto whitespace-pre-wrap break-words rounded bg-muted/50 p-2 text-xs">
{resultText}
</pre>
) : null}
</ToolFallbackContent>
</ToolFallbackRoot>
);
};
export const CodeExecutionToolUI = memo(
CodeExecutionToolUIImpl,
) as unknown as ToolCallMessagePartComponent;
CodeExecutionToolUI.displayName = "CodeExecutionToolUI";

8
studio/frontend/src/features/chat/api-provider-logo.tsx
View file

@ -4,6 +4,7 @@
import { cn } from "@/lib/utils";
import { DashboardSquare01Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { isCustomProviderType } from "./external-providers";
/**
* Registry logos live at `public/provider-logos/{provider_type}.{ext}` where `provider_type`
@ -19,6 +20,9 @@ const PROVIDER_LOGO_EXT: Record<string, "svg" | "png" | "jpg"> = {
kimi: "jpg",
qwen: "png",
openrouter: "svg",
vllm: "svg",
ollama: "svg",
llama_cpp: "svg",
};
export function apiProviderLogoSrc(
@ -42,7 +46,8 @@ interface ApiProviderLogoProps {
* OpenAI's asset is black-on-transparent; it is inverted in dark mode for contrast.
*/
export function ApiProviderLogo({ providerType, className, title }: ApiProviderLogoProps) {
if (providerType === "custom") {
const src = apiProviderLogoSrc(providerType);
if (!src && isCustomProviderType(providerType)) {
return (
<span title={title} aria-hidden className="inline-flex shrink-0">
<HugeiconsIcon icon={DashboardSquare01Icon} className={cn("shrink-0", className)} />
@ -50,7 +55,6 @@ export function ApiProviderLogo({ providerType, className, title }: ApiProviderL
);
}
const src = apiProviderLogoSrc(providerType);
if (!src) return null;
return (
<img

215
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -15,6 +15,8 @@ import {
streamChatCompletions,
validateModel,
} from "./chat-api";
import { pickFriendlyContainerName } from "../lib/friendly-names";
import { createOpenAIContainer } from "./openai-containers";
import {
encryptProviderApiKey,
isProviderKeyRotationError,
@ -26,8 +28,11 @@ import type {
} from "../types/api";
import {
getExternalProviderApiKey,
isCustomProviderType,
loadExternalProviders,
parseExternalModelId,
supportsProviderPromptCaching,
toExternalBackendProviderType,
} from "../external-providers";
import {
EXTERNAL_MAX_OUTPUT_TOKENS,
@ -35,6 +40,8 @@ import {
getExternalMinOutputTokens,
getExternalReasoningCapabilities,
getProviderCapabilities,
providerSupportsBuiltinCodeExecution,
providerSupportsBuiltinWebSearch,
} from "../provider-capabilities";
import { useChatRuntimeStore } from "../stores/chat-runtime-store";
import { isMultimodalResponse } from "../types/api";
@ -741,13 +748,17 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
if (isExternalRequest && !externalProvider) {
toast.error("External provider not found.", {
description: "Open API Providers and re-add this provider.",
description: "Open Connections and re-add this provider.",
});
throw new Error("External provider not found.");
}
if (isExternalRequest && !externalApiKey) {
// Local providers (llama.cpp / vLLM / Ollama) allow an empty key — only block hosted providers.
const externalProviderIsCustom = externalProvider
? isCustomProviderType(externalProvider.providerType)
: false;
if (isExternalRequest && !externalApiKey && !externalProviderIsCustom) {
toast.error("Missing API key for selected external provider.", {
description: "Open API Providers and set the API key again.",
description: "Open Connections and set the API key again.",
});
throw new Error("Missing external provider API key.");
}
@ -928,10 +939,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
supportsPreserveThinking,
preserveThinking,
} = runtime;
const externalBackendProviderType =
externalProvider?.providerType === "custom"
? "openai"
: externalProvider?.providerType;
const externalBackendProviderType = toExternalBackendProviderType(
externalProvider?.providerType,
);
const externalCapabilities = getProviderCapabilities(
externalProvider?.providerType,
);
@ -942,6 +952,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
? getExternalReasoningCapabilities(
externalProvider.providerType,
externalSelection.modelId,
{
isReasoningProvider:
externalProvider.isReasoningModel === true,
},
)
: {
supportsReasoning,
@ -972,6 +986,106 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
forceRefreshPublicKey = false,
): Promise<OpenAIChatCompletionsRequest> => {
if (externalSelection && externalProvider) {
// OpenAI shell-tool container reuse: pull the per-thread
// container_id (if any) so subsequent turns in the same
// thread reference the existing container instead of
// auto-creating a fresh one. Empty string / undefined →
// backend falls back to container_auto. Anthropic doesn't
// use this (server-side per-turn container).
let openaiCodeExecContainerId: string | null = null;
const codeExecEnabledForThisTurn =
codeToolsEnabled &&
providerSupportsBuiltinCodeExecution(
externalProvider.providerType,
externalSelection.modelId,
externalProvider.baseUrl,
);
if (codeExecEnabledForThisTurn && resolvedThreadId) {
try {
const thread = await db.threads.get(resolvedThreadId);
openaiCodeExecContainerId =
thread?.openaiCodeExecContainerId ?? null;
} catch {
openaiCodeExecContainerId = null;
}
// Cross-thread inheritance: when the active thread has
// no container yet, default to the one most recently
// used on *any* other thread (provider-scoped).
// Matches what the Code Execution settings section
// shows in the picker, and keeps the user from getting
// a fresh container on every new thread. The picker
// can still be set to "Auto-create per thread"
// explicitly to opt into a fresh container — but
// that's done via the dropdown, not silently.
if (
!openaiCodeExecContainerId &&
externalProvider.providerType === "openai"
) {
try {
const others = await db.threads
.orderBy("createdAt")
.reverse()
.toArray();
for (const t of others) {
if (t.id === resolvedThreadId) continue;
if (t.openaiCodeExecContainerId) {
openaiCodeExecContainerId = t.openaiCodeExecContainerId;
void db.threads
.update(resolvedThreadId, {
openaiCodeExecContainerId,
})
.catch(() => {});
break;
}
}
} catch {
/* fall through to lazy-create below */
}
}
// Lazy pre-create when there's no inherited container.
// We always POST /v1/containers ourselves (rather than
// letting the backend send container_auto) so every
// container shows up in the picker with a friendly
// English-word name and the user's configured TTL.
// Falls back to container_auto only if the POST fails
// — keeps the chat moving in that case.
if (
!openaiCodeExecContainerId &&
externalProvider.providerType === "openai"
) {
const ttl = externalProvider.openaiContainerTtlMinutes;
const ttlToUse =
typeof ttl === "number" && ttl >= 1 ? ttl : 20;
try {
const created = await createOpenAIContainer(
{
apiKey: externalApiKey,
baseUrl: externalProvider.baseUrl || null,
},
{
// Friendly English-word name so the container
// is human-readable in the picker list (e.g.
// "kestrel-3f9c") instead of a thread-id slug
// or OpenAI's default blank name.
name: pickFriendlyContainerName(),
ttlMinutes: ttlToUse,
},
);
openaiCodeExecContainerId = created.id;
void db.threads
.update(resolvedThreadId, {
openaiCodeExecContainerId: created.id,
})
.catch(() => {});
} catch {
// Fall back to backend's container_auto path on
// failure — keeps the chat moving; the next turn
// can retry. The auto-created container will be
// unnamed, but the chat doesn't break.
openaiCodeExecContainerId = null;
}
}
}
return {
model: externalSelection.modelId,
messages: outboundMessages,
@ -1005,14 +1119,65 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
...(externalCapabilities?.presencePenalty
? { presence_penalty: params.presencePenalty }
: {}),
// Built-in tools: Search pill maps to provider-side
// web_search (currently OpenAI / Anthropic / OpenRouter /
// Kimi); Code pill maps to Anthropic's server-side
// code_execution_20250825 tool (Anthropic is the only
// external provider that ships one today). Backend
// translates enabled_tools into each provider's tool
// schema — for Anthropic that's the entries appended to
// body["tools"] inside _stream_anthropic.
...((toolsEnabled &&
providerSupportsBuiltinWebSearch(externalProvider.providerType)) ||
(codeToolsEnabled &&
providerSupportsBuiltinCodeExecution(
externalProvider.providerType,
externalSelection.modelId,
externalProvider.baseUrl,
))
? {
enable_tools: true,
enabled_tools: [
...(toolsEnabled &&
providerSupportsBuiltinWebSearch(
externalProvider.providerType,
)
? ["web_search"]
: []),
...(codeToolsEnabled &&
providerSupportsBuiltinCodeExecution(
externalProvider.providerType,
externalSelection.modelId,
externalProvider.baseUrl,
)
? ["code_execution"]
: []),
],
}
: {}),
provider_id: externalProvider.id,
provider_type: externalBackendProviderType,
external_model: externalSelection.modelId,
encrypted_api_key: await encryptProviderApiKey(
externalApiKey,
forceRefreshPublicKey,
),
...(externalApiKey
? {
encrypted_api_key: await encryptProviderApiKey(
externalApiKey,
forceRefreshPublicKey,
),
}
: {}),
provider_base_url: externalProvider.baseUrl || null,
...(openaiCodeExecContainerId
? {
openai_code_exec_container_id: openaiCodeExecContainerId,
}
: {}),
...(supportsProviderPromptCaching(externalProvider.providerType)
? {
enable_prompt_caching:
externalProvider.enablePromptCaching ?? true,
}
: {}),
...(externalReasoningCaps.supportsReasoning
? externalReasoningCaps.reasoningStyle === "reasoning_effort"
? externalReasoningEnabled
@ -1090,6 +1255,34 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
// On tool_end: set result on the existing part (transitions to "complete").
const toolEvent = (chunk as unknown as { _toolEvent?: Record<string, unknown> })._toolEvent;
if (toolEvent !== undefined) {
// OpenAI shell-tool container persistence — see
// ThreadRecord.openaiCodeExecContainerId. The backend
// emits these synthetic events on the OpenAI Responses
// SSE stream after capturing the container_id from a
// response, or detecting an expired-container error.
if (toolEvent.type === "container_ready") {
const newContainerId = toolEvent.container_id as
| string
| undefined;
if (newContainerId && resolvedThreadId) {
void db.threads
.update(resolvedThreadId, {
openaiCodeExecContainerId: newContainerId,
})
.catch(() => {});
}
continue;
}
if (toolEvent.type === "container_invalidated") {
if (resolvedThreadId) {
void db.threads
.update(resolvedThreadId, {
openaiCodeExecContainerId: null,
})
.catch(() => {});
}
continue;
}
if (toolEvent.type === "tool_start") {
const id = (toolEvent.tool_call_id as string) || `${toolEvent.tool_name}_${Date.now()}`;
const toolArgs = (toolEvent.arguments ?? {}) as ToolCallMessagePart["args"];

124
studio/frontend/src/features/chat/api/openai-containers.ts Normal file
View file

@ -0,0 +1,124 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
/**
* Wrappers for the three OpenAI shell-tool container management
* endpoints exposed by the backend (studio/backend/routes/inference.py).
* Each one proxies to OpenAI's /v1/containers REST surface using the
* user's encrypted API key. Backend rejects any base URL that isn't
* api.openai.com — the shell tool only exists on the managed cloud.
*/
import { authFetch } from "@/features/auth";
import { encryptProviderApiKey } from "./providers-api";
export interface OpenAIContainerSummary {
id: string;
name?: string | null;
createdAt?: number | null;
lastActiveAt?: number | null;
expiresAfterMinutes?: number | null;
status?: string | null;
}
interface RawSummary {
id: string;
name?: string | null;
created_at?: number | null;
last_active_at?: number | null;
expires_after_minutes?: number | null;
status?: string | null;
}
function fromRaw(raw: RawSummary): OpenAIContainerSummary {
return {
id: raw.id,
name: raw.name ?? null,
createdAt: raw.created_at ?? null,
lastActiveAt: raw.last_active_at ?? null,
expiresAfterMinutes: raw.expires_after_minutes ?? null,
status: raw.status ?? null,
};
}
async function parseError(response: Response): Promise<string> {
try {
const body = (await response.json()) as { detail?: string };
if (body && typeof body.detail === "string") return body.detail;
} catch {
/* fall through */
}
return `HTTP ${response.status}`;
}
interface AuthInputs {
apiKey: string;
baseUrl: string | null;
}
async function buildAuthBody(auth: AuthInputs) {
return {
encrypted_api_key: await encryptProviderApiKey(auth.apiKey),
provider_base_url: auth.baseUrl,
};
}
export async function listOpenAIContainers(
auth: AuthInputs,
): Promise<OpenAIContainerSummary[]> {
const response = await authFetch(
"/api/inference/external/openai/containers/list",
{
method: "POST",
headers: { "content-type": "application/json" },
body: JSON.stringify(await buildAuthBody(auth)),
},
);
if (!response.ok) throw new Error(await parseError(response));
const body = (await response.json()) as { containers?: RawSummary[] };
return (body.containers ?? []).map(fromRaw);
}
export async function createOpenAIContainer(
auth: AuthInputs,
params: { name: string; ttlMinutes: number },
): Promise<OpenAIContainerSummary> {
const response = await authFetch(
"/api/inference/external/openai/containers/create",
{
method: "POST",
headers: { "content-type": "application/json" },
body: JSON.stringify({
...(await buildAuthBody(auth)),
name: params.name,
ttl_minutes: params.ttlMinutes,
}),
},
);
if (!response.ok) throw new Error(await parseError(response));
const raw = (await response.json()) as RawSummary;
return fromRaw(raw);
}
export async function deleteOpenAIContainer(
auth: AuthInputs,
containerId: string,
): Promise<void> {
const response = await authFetch(
"/api/inference/external/openai/containers/delete",
{
method: "POST",
headers: { "content-type": "application/json" },
body: JSON.stringify({
...(await buildAuthBody(auth)),
container_id: containerId,
}),
},
);
// 404 = container already gone (deleted elsewhere, or expired-then-purged).
// Treat as idempotent success so a stale list entry doesn't surface as a
// confusing error — the caller will refresh and the entry will disappear.
if (!response.ok && response.status !== 204 && response.status !== 404) {
throw new Error(await parseError(response));
}
}

6
studio/frontend/src/features/chat/api/providers-api.ts
View file

@ -176,8 +176,12 @@ export async function updateProviderConfig(
async function withApiKeyEncryptionRetry<T>(
plaintextApiKey: string,
call: (encryptedApiKey: string) => Promise<T>,
call: (encryptedApiKey: string | null) => Promise<T>,
): Promise<T> {
// Empty key (local providers): skip RSA round-trip and let the backend omit auth.
if (!plaintextApiKey) {
return await call(null);
}
try {
const encrypted = await encryptProviderApiKey(plaintextApiKey, false);
return await call(encrypted);

153
studio/frontend/src/features/chat/chat-page.tsx
View file

@ -50,6 +50,8 @@ import {
clampReasoningEffortToLevels,
getExternalReasoningCapabilities,
getProviderCapabilities,
providerSupportsBuiltinCodeExecution,
providerSupportsBuiltinWebSearch,
} from "./provider-capabilities";
import { useChatModelRuntime } from "./hooks/use-chat-model-runtime";
import {
@ -549,6 +551,7 @@ export function ChatPage(): ReactElement {
const settingsOpen = useChatRuntimeStore((s) => s.settingsPanelOpen);
const setSettingsOpen = useChatRuntimeStore((s) => s.setSettingsPanelOpen);
const externalProviders = useExternalProvidersStore((s) => s.providers);
const setExternalProviders = useExternalProvidersStore((s) => s.setProviders);
useEffect(() => {
const threadId = search.thread;
@ -628,14 +631,16 @@ export function ChatPage(): ReactElement {
const reasoningStyle = useChatRuntimeStore((s) => s.reasoningStyle);
const reasoningEffort = useChatRuntimeStore((s) => s.reasoningEffort);
const supportsReasoningOff = useChatRuntimeStore((s) => s.supportsReasoningOff);
const activeExternalProviderType = useMemo(() => {
const activeExternalProvider = useMemo(() => {
const selection = parseExternalModelId(inferenceParams.checkpoint);
if (!selection) return null;
const provider = externalProviders.find(
(p) => p.id === selection.providerId,
return (
externalProviders.find(
(p) => p.id === selection.providerId,
) ?? null
);
return provider?.providerType ?? null;
}, [externalProviders, inferenceParams.checkpoint]);
const activeExternalProviderType = activeExternalProvider?.providerType ?? null;
const activeProviderCapabilities = useMemo(() => {
const selection = parseExternalModelId(inferenceParams.checkpoint);
if (!selection) return null;
@ -670,6 +675,7 @@ export function ChatPage(): ReactElement {
const reasoningCaps = getExternalReasoningCapabilities(
provider?.providerType,
selection.modelId,
{ isReasoningProvider: provider?.isReasoningModel === true },
);
const state = useChatRuntimeStore.getState();
const preferredEffort = state.reasoningEffort;
@ -678,9 +684,61 @@ export function ChatPage(): ReactElement {
preferredEffort,
effortLevels,
);
// Per-provider default effort. Anthropic gets the highest available
// level (xhigh on 4.6/4.7, high on 4.5) since Claude's adaptive
// thinking adjusts cost per turn — sitting at the top of the dial
// gives users the strongest answers and the model can still skip
// thinking when the turn is trivial. OpenAI gets "high" by default
// — the gpt-5.x reasoning models accept high across the board and
// it's the right cost/quality sweet spot for Responses-API tools
// (web search included). Everyone else gets "medium" as a balanced
// default. Users can pick another level via the Think dropdown.
const isAnthropic = provider?.providerType === "anthropic";
const isOpenAI = provider?.providerType === "openai";
const anthropicTopEffort = effortLevels.includes("xhigh")
? "xhigh"
: effortLevels.includes("high")
? "high"
: clampedEffort;
const openaiDefaultEffort = effortLevels.includes("high")
? "high"
: effortLevels.includes("medium")
? "medium"
: clampedEffort;
const nextReasoningEffort = reasoningCaps.supportsReasoning
? clampedEffort
? isAnthropic
? anthropicTopEffort
: isOpenAI
? openaiDefaultEffort
: effortLevels.includes("medium")
? "medium"
: clampedEffort
: state.reasoningEffort;
const supportsBuiltinWebSearch = providerSupportsBuiltinWebSearch(
provider?.providerType,
);
const supportsBuiltinCodeExecution = providerSupportsBuiltinCodeExecution(
provider?.providerType,
selection.modelId,
provider?.baseUrl,
);
// Kimi's k2.6/k2.5 default to thinking enabled on the server side
// (per https://platform.kimi.ai/docs/models). Mirror that default
// in the UI so the Think pill comes up clicked when the user picks
// a Kimi model. The Search pill stays off by default; the mutual-
// exclusion handlers in the composer flip the two when needed.
const isKimi = provider?.providerType === "kimi";
// Web search is on by default for the two providers we trust most
// for it: Anthropic (web_search_20250305 server tool, structured
// citations) and OpenAI (/v1/responses web_search, structured
// citations). Other providers stay off-by-default — OpenRouter's
// plugins shape and Kimi's $web_search builtin still work when the
// user opts in via the pill, but they're a notch less reliable so
// we don't pre-enable them.
const searchOnByDefault =
supportsBuiltinWebSearch &&
(provider?.providerType === "anthropic" ||
provider?.providerType === "openai");
useChatRuntimeStore.setState({
supportsReasoning: reasoningCaps.supportsReasoning,
reasoningAlwaysOn: reasoningCaps.reasoningAlwaysOn,
@ -690,10 +748,24 @@ export function ChatPage(): ReactElement {
reasoningEffort: nextReasoningEffort,
reasoningEnabled: reasoningCaps.supportsReasoning
? reasoningCaps.supportsReasoningOff
? state.reasoningEnabled
? isKimi
? true
: state.reasoningEnabled
: true
: state.reasoningEnabled,
supportsPreserveThinking: false,
// External models never give us a local tool runtime (no
// python sandbox), so `supportsTools` must be false. The two
// `supportsBuiltin*` flags pick up the slack for providers that
// run the tool server-side: `supportsBuiltinWebSearch` lights
// up the Search pill (OpenAI / Anthropic / OpenRouter / Kimi),
// `supportsBuiltinCodeExecution` lights up the Code pill
// (Anthropic Claude 4.x only, today).
supportsTools: false,
supportsBuiltinWebSearch,
supportsBuiltinCodeExecution,
toolsEnabled: searchOnByDefault,
codeToolsEnabled: false,
});
}, [externalProviders, inferenceParams.checkpoint]);
const canCompare = useMemo(() => {
@ -803,6 +875,10 @@ export function ChatPage(): ReactElement {
const reasoningCaps = getExternalReasoningCapabilities(
selectedProvider?.providerType,
selectedExternal?.modelId,
{
isReasoningProvider:
selectedProvider?.isReasoningModel === true,
},
);
const preferredEffort = store.reasoningEffort;
const effortLevels = reasoningCaps.reasoningEffortLevels;
@ -810,8 +886,29 @@ export function ChatPage(): ReactElement {
preferredEffort,
effortLevels,
);
// Same per-provider default policy as the useEffect path above:
// Anthropic picks the highest available level, OpenAI picks
// "high", everyone else picks "medium".
const isAnthropic = selectedProvider?.providerType === "anthropic";
const isOpenAI = selectedProvider?.providerType === "openai";
const anthropicTopEffort = effortLevels.includes("xhigh")
? "xhigh"
: effortLevels.includes("high")
? "high"
: clampedEffort;
const openaiDefaultEffort = effortLevels.includes("high")
? "high"
: effortLevels.includes("medium")
? "medium"
: clampedEffort;
const nextReasoningEffort = reasoningCaps.supportsReasoning
? clampedEffort
? isAnthropic
? anthropicTopEffort
: isOpenAI
? openaiDefaultEffort
: effortLevels.includes("medium")
? "medium"
: clampedEffort
: store.reasoningEffort;
// Clear any cached router-picked openrouter/free model unless the
// user is staying on openrouter/free — otherwise the chip would
@ -823,6 +920,25 @@ export function ChatPage(): ReactElement {
...store.params,
checkpoint: value,
});
const supportsBuiltinWebSearch = providerSupportsBuiltinWebSearch(
selectedProvider?.providerType,
);
const supportsBuiltinCodeExecution = providerSupportsBuiltinCodeExecution(
selectedProvider?.providerType,
selectedExternal?.modelId,
selectedProvider?.baseUrl,
);
// See sibling useEffect above: Kimi's k2.x default to thinking
// enabled, so the Think pill comes up clicked. Search pill stays
// off by default; mutual exclusion flips them via the composer.
const isKimi = selectedProvider?.providerType === "kimi";
// Mirror of sibling useEffect: Anthropic and OpenAI get Search
// on-by-default since their server tools emit structured
// citations end-to-end. OpenRouter and Kimi stay off-by-default.
const searchOnByDefault =
supportsBuiltinWebSearch &&
(selectedProvider?.providerType === "anthropic" ||
selectedProvider?.providerType === "openai");
useChatRuntimeStore.setState({
activeGgufVariant: null,
ggufContextLength: null,
@ -837,10 +953,23 @@ export function ChatPage(): ReactElement {
reasoningEffort: nextReasoningEffort,
reasoningEnabled: reasoningCaps.supportsReasoning
? reasoningCaps.supportsReasoningOff
? store.reasoningEnabled
? isKimi
? true
: store.reasoningEnabled
: true
: store.reasoningEnabled,
supportsPreserveThinking: false,
// External models have no local tool runtime → supportsTools
// stays false. The two supportsBuiltin* flags carry the
// server-side capability info for each pill:
// - Search → providerSupportsBuiltinWebSearch
// - Code → providerSupportsBuiltinCodeExecution
// (Anthropic Claude 4.x only, today)
supportsTools: false,
supportsBuiltinWebSearch,
supportsBuiltinCodeExecution,
toolsEnabled: searchOnByDefault,
codeToolsEnabled: false,
...(stillOnOpenRouterFree ? {} : { lastOpenRouterChosenModel: null }),
});
return;
@ -1321,6 +1450,14 @@ export function ChatPage(): ReactElement {
onParamsChange={setInferenceParams}
isExternalModel={isExternalModel}
providerCapabilities={activeProviderCapabilities}
activeExternalProvider={activeExternalProvider}
onExternalProviderChange={(updatedProvider) => {
setExternalProviders(
externalProviders.map((provider) =>
provider.id === updatedProvider.id ? updatedProvider : provider,
),
);
}}
externalProviderType={activeExternalProviderType}
onReloadModel={() => {
const state = useChatRuntimeStore.getState();

331
studio/frontend/src/features/chat/chat-providers-dialog.tsx
View file

@ -15,7 +15,9 @@ import { Label } from "@/components/ui/label";
import {
Select,
SelectContent,
SelectGroup,
SelectItem,
SelectSeparator,
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
@ -23,7 +25,6 @@ import { Spinner } from "@/components/ui/spinner";
import { Textarea } from "@/components/ui/textarea";
import {
ArrowLeft02Icon,
DashboardSquare01Icon,
Delete02Icon,
Edit03Icon,
PlusSignIcon,
@ -47,9 +48,19 @@ import {
} from "./api/providers-api";
import type { ExternalProviderConfig } from "./external-providers";
import {
CUSTOM_BACKEND_PROVIDER_TYPE,
CUSTOM_PROVIDER_PRESETS,
customProviderBaseUrlPlaceholder,
customProviderDisplayName,
customProviderModelIdsPlaceholder,
getExternalProviderApiKey,
isCustomProviderType,
LEGACY_CUSTOM_PROVIDER_TYPE,
removeExternalProviderApiKey,
setExternalProviderApiKey,
supportsProviderPromptCaching,
supportsProviderReasoningToggle,
toExternalBackendProviderType,
} from "./external-providers";
/** Matches navbar / thread layout easing (see index.css --ease-out-quart) */
@ -57,12 +68,11 @@ const PROVIDER_FORM_EASE: [number, number, number, number] = [
0.165, 0.84, 0.44, 1,
];
const PROVIDER_FORM_DURATION = 0.2;
const CUSTOM_PROVIDER_TYPE = "custom";
const CUSTOM_BACKEND_PROVIDER_TYPE = "openai";
const CUSTOM_PROVIDER_MISSING_KEY_MESSAGE =
"No API key found, please make sure API key is added and valid for this provider.";
const ANTHROPIC_DATED_SNAPSHOT_SUFFIX = /-\d{8}$/;
const OPENAI_DEPRECATED_MODELS = new Set(["gpt-5.3"]);
const HIDDEN_PROVIDER_TYPES = new Set(["qwen"]);
const OPENROUTER_EXCLUDED_MODELS = new Set([
"google/chirp-3",
"kwaivgi/kling-v3.0-pro",
@ -82,37 +92,37 @@ function resolveUiProviderTypeFromConfig(
registryRows: ProviderRegistryEntry[],
existingProviderType: string | undefined,
): string {
if (existingProviderType === CUSTOM_PROVIDER_TYPE) {
return CUSTOM_PROVIDER_TYPE;
if (existingProviderType && isCustomProviderType(existingProviderType)) {
return existingProviderType;
}
if (configProviderType !== CUSTOM_BACKEND_PROVIDER_TYPE) {
return configProviderType;
}
const displayName = (configDisplayName ?? "").trim().toLowerCase();
const matchingCustomPreset = CUSTOM_PROVIDER_PRESETS.find(
(preset) => preset.displayName.toLowerCase() === displayName,
);
if (matchingCustomPreset) {
return matchingCustomPreset.providerType;
}
const openAiRegistry = registryRows.find(
(entry) => entry.provider_type === CUSTOM_BACKEND_PROVIDER_TYPE,
);
if (!openAiRegistry) {
return configProviderType;
}
const displayName = (configDisplayName ?? "").trim().toLowerCase();
const openAiDisplayName = openAiRegistry.display_name.trim().toLowerCase();
if (displayName.length > 0 && displayName !== openAiDisplayName) {
return CUSTOM_PROVIDER_TYPE;
return LEGACY_CUSTOM_PROVIDER_TYPE;
}
const configUrl = normalizeUrl(configBaseUrl ?? "");
const defaultUrl = normalizeUrl(openAiRegistry.base_url ?? "");
if (configUrl.length > 0 && configUrl !== defaultUrl) {
return CUSTOM_PROVIDER_TYPE;
return LEGACY_CUSTOM_PROVIDER_TYPE;
}
return configProviderType;
}
function toBackendProviderType(uiProviderType: string): string {
return uiProviderType === CUSTOM_PROVIDER_TYPE
? CUSTOM_BACKEND_PROVIDER_TYPE
: uiProviderType;
}
function parseManualModelIds(text: string): string[] {
const seen = new Set<string>();
const out: string[] = [];
@ -175,22 +185,27 @@ export function ChatProvidersSettings({
const [manualModelIds, setManualModelIds] = useState("");
const [modelSearchQuery, setModelSearchQuery] = useState("");
const [customProviderName, setCustomProviderName] = useState("Custom");
const [isReasoningModel, setIsReasoningModel] = useState(false);
const reduceMotion = useReducedMotion();
const isCustomProvider = providerType === CUSTOM_PROVIDER_TYPE;
const isCustomProvider = isCustomProviderType(providerType);
// Ollama runs locally and does not require an API key. Hide the input
// entirely rather than just marking it optional so users aren't prompted
// for a credential the provider never uses.
const isOllamaProvider = providerType === "ollama";
const showApiKeyField = !isOllamaProvider;
const showReasoningToggle = supportsProviderReasoningToggle(providerType);
const registryByType = useMemo(
() => new Map(registry.map((entry) => [entry.provider_type, entry])),
[registry],
);
const hasCustomInRegistry = registryByType.has(CUSTOM_PROVIDER_TYPE);
const isCuratedModelList = useMemo(() => {
return registryByType.get(providerType)?.model_list_mode === "curated";
}, [registryByType, providerType]);
const isManualModelList = isCustomProvider || isCuratedModelList;
const modelsPanelKey = isCustomProvider
? "custom"
? providerType || "custom"
: isCuratedModelList
? "curated"
: "remote";
@ -225,14 +240,22 @@ export function ChatProvidersSettings({
useEffect(() => {
if (!providerType || editingProviderId) return;
const entry = registryByType.get(providerType);
if (!entry) return;
// Seed the registry's default_models for every provider — curated and
// remote alike. For remote-mode providers, loadModels() will replace
// this with the union of defaults + the live /models response once the
// user clicks "Load Models"; until then (or if the call fails — e.g.
// decryption issues during key rotation) the seeded list ensures
// curated picks like claude-haiku-4-5 are always reachable.
setAvailableModels([...entry.default_models]);
if (!entry) {
if (isCustomProviderType(providerType)) {
setCustomProviderName(customProviderDisplayName(providerType));
}
return;
}
// Seed default_models only when the catalog is not fetched live:
// curated providers (catalog too large to enumerate, defaults are
// the suggestion shortlist) and Ollama (local, no API key — local
// /models stands in). Remote-mode cloud providers stay empty until
// the user clicks "Load available models" with a key, since
// different API tiers expose different catalogs and we don't want
// to advertise models the user can't actually call.
const seedDefaults =
entry.model_list_mode === "curated" || providerType === "ollama";
setAvailableModels(seedDefaults ? [...entry.default_models] : []);
setSelectedModelIds([]);
setManualModelIds("");
setModelSearchQuery("");
@ -297,6 +320,12 @@ export function ChatProvidersSettings({
baseUrl: config.base_url ?? "",
models: existingModels,
availableModels: existing?.availableModels ?? [],
enablePromptCaching: supportsProviderPromptCaching(uiProviderType)
? (existing?.enablePromptCaching ?? true)
: undefined,
isReasoningModel: supportsProviderReasoningToggle(uiProviderType)
? existing?.isReasoningModel === true
: undefined,
createdAt: existing?.createdAt ?? createdAt,
updatedAt,
};
@ -328,15 +357,19 @@ export function ChatProvidersSettings({
setSelectedModelIds([]);
setManualModelIds("");
setModelSearchQuery("");
setCustomProviderName("Custom");
setCustomProviderName(customProviderDisplayName(providerType));
setIsReasoningModel(false);
}
function openAddProvider() {
resetForm();
const entry = providerType ? registryByType.get(providerType) : null;
if (entry) {
// Keep first-open behavior consistent with provider re-selection.
setAvailableModels([...entry.default_models]);
const seedDefaults =
entry.model_list_mode === "curated" || providerType === "ollama";
if (seedDefaults) {
setAvailableModels([...entry.default_models]);
}
}
setPage("form");
}
@ -384,7 +417,7 @@ export function ChatProvidersSettings({
const trimmed = input.trim();
if (!trimmed) {
if (required) {
throw new Error("Base URL is required for custom providers.");
throw new Error("Base URL is required for this connection.");
}
return null;
}
@ -397,7 +430,7 @@ export function ChatProvidersSettings({
return;
}
if (isCustomProvider) {
toast.info("Custom providers use manual model IDs.");
toast.info("This connection uses manual model IDs.");
return;
}
if (isCuratedModelList) {
@ -458,10 +491,10 @@ export function ChatProvidersSettings({
toast.error("Choose a provider first.");
return;
}
const backendProviderType = toBackendProviderType(providerType);
const backendProviderType = toExternalBackendProviderType(providerType);
const selectedRegistryEntry = registryByType.get(backendProviderType);
const displayName = isCustomProvider
? customProviderName.trim() || "Custom"
? customProviderName.trim() || customProviderDisplayName(providerType)
: (selectedRegistryEntry?.display_name ?? providerType);
if (!isCustomProvider && !apiKey.trim()) {
toast.error("API key is required.");
@ -511,17 +544,21 @@ export function ChatProvidersSettings({
const updatedAt = Number.isFinite(Date.parse(created.updated_at))
? Date.parse(created.updated_at)
: Date.now();
const uiProviderType = isCustomProvider
? providerType
: created.provider_type;
const provider: ExternalProviderConfig = {
id: created.id,
providerType: isCustomProvider
? CUSTOM_PROVIDER_TYPE
: created.provider_type,
providerType: uiProviderType,
name: created.display_name,
baseUrl: created.base_url ?? "",
models: modelsToSave,
availableModels: manualModels
? []
: pruneProviderModelIds(providerType, availableModels),
isReasoningModel: supportsProviderReasoningToggle(uiProviderType)
? isReasoningModel
: undefined,
createdAt,
updatedAt,
};
@ -553,7 +590,7 @@ export function ChatProvidersSettings({
return;
}
const isEditingCustomProvider =
existing.providerType === CUSTOM_PROVIDER_TYPE;
isCustomProviderType(existing.providerType);
if (!isEditingCustomProvider && !apiKey.trim()) {
toast.error("API key is required.");
return;
@ -597,7 +634,8 @@ export function ChatProvidersSettings({
);
const updated = await updateProviderConfig(editingProviderId, {
displayName: isEditingCustomProvider
? customProviderName.trim() || "Custom"
? customProviderName.trim() ||
customProviderDisplayName(existing.providerType)
: existing.name,
baseUrl,
});
@ -620,6 +658,11 @@ export function ChatProvidersSettings({
availableModels: manualModels
? []
: pruneProviderModelIds(existing.providerType, availableModels),
isReasoningModel: supportsProviderReasoningToggle(
existing.providerType,
)
? isReasoningModel
: undefined,
updatedAt,
}
: provider,
@ -640,12 +683,19 @@ export function ChatProvidersSettings({
setEditingProviderId(provider.id);
setPage("form");
setProviderType(provider.providerType);
setCustomProviderName(provider.name || "Custom");
setCustomProviderName(
provider.name || customProviderDisplayName(provider.providerType),
);
setApiKey(getExternalProviderApiKey(provider.id));
setShowApiKey(false);
setBaseUrlDraft(provider.baseUrl);
setModelSearchQuery("");
if (provider.providerType === CUSTOM_PROVIDER_TYPE) {
setIsReasoningModel(
supportsProviderReasoningToggle(provider.providerType)
? provider.isReasoningModel === true
: false,
);
if (isCustomProviderType(provider.providerType)) {
setAvailableModels([]);
setSelectedModelIds([]);
setManualModelIds(provider.models.join("\n"));
@ -695,8 +745,11 @@ export function ChatProvidersSettings({
async function testProvider(provider: ExternalProviderConfig) {
const savedKey = getExternalProviderApiKey(provider.id).trim();
if (!savedKey) {
if (provider.providerType === CUSTOM_PROVIDER_TYPE) {
// Ollama runs locally and never requires a key — fall through to the
// real connection check instead of prompting for credentials the form
// no longer exposes.
if (!savedKey && provider.providerType !== "ollama") {
if (isCustomProviderType(provider.providerType)) {
await editProvider(provider);
toast.info(CUSTOM_PROVIDER_MISSING_KEY_MESSAGE);
return;
@ -707,7 +760,9 @@ export function ChatProvidersSettings({
}
try {
const result = await testProviderConnection({
providerType: toBackendProviderType(provider.providerType),
providerType:
toExternalBackendProviderType(provider.providerType) ??
provider.providerType,
apiKey: savedKey,
baseUrl: provider.baseUrl || null,
});
@ -715,7 +770,7 @@ export function ChatProvidersSettings({
toast.success(result.message);
} else {
if (
provider.providerType === CUSTOM_PROVIDER_TYPE &&
isCustomProviderType(provider.providerType) &&
result.message.includes("Illegal header value b'Bearer '")
) {
toast.error(CUSTOM_PROVIDER_MISSING_KEY_MESSAGE);
@ -726,7 +781,7 @@ export function ChatProvidersSettings({
} catch (error) {
const message = error instanceof Error ? error.message : "Unknown error";
if (
provider.providerType === CUSTOM_PROVIDER_TYPE &&
isCustomProviderType(provider.providerType) &&
message.includes("Illegal header value b'Bearer '")
) {
toast.error(CUSTOM_PROVIDER_MISSING_KEY_MESSAGE);
@ -753,7 +808,7 @@ export function ChatProvidersSettings({
</Button>
<div className="flex min-w-0 items-center gap-2 leading-none">
<span className="text-xs font-medium text-muted-foreground">
Cloud
Connections
</span>
<span className="size-1 rounded-full bg-muted-foreground/35" />
<span className="truncate text-xs font-medium text-muted-foreground">
@ -774,7 +829,7 @@ export function ChatProvidersSettings({
Provider
</Label>
<p className="text-xs leading-snug text-muted-foreground">
Supported registry or Custom.
Supported registry or local OpenAI-compatible connection.
</p>
</div>
<Select
@ -786,6 +841,9 @@ export function ChatProvidersSettings({
setSelectedModelIds([]);
setManualModelIds("");
setModelSearchQuery("");
if (isCustomProviderType(value)) {
setCustomProviderName(customProviderDisplayName(value));
}
}}
>
<SelectTrigger
@ -796,72 +854,88 @@ export function ChatProvidersSettings({
<SelectValue placeholder="Choose a provider" />
</SelectTrigger>
<SelectContent>
{registry.map((entry) => (
<SelectItem
key={entry.provider_type}
value={entry.provider_type}
>
<span className="flex items-center gap-2">
<ApiProviderLogo
providerType={entry.provider_type}
className="size-4"
title={entry.display_name}
/>
{entry.display_name}
</span>
</SelectItem>
))}
{hasCustomInRegistry ? null : (
<SelectItem value={CUSTOM_PROVIDER_TYPE}>
<span className="flex items-center gap-2">
<HugeiconsIcon
icon={DashboardSquare01Icon}
className="size-4"
/>
Custom
</span>
</SelectItem>
)}
<SelectGroup>
{CUSTOM_PROVIDER_PRESETS.map((preset) => (
<SelectItem
key={preset.providerType}
value={preset.providerType}
>
<span className="flex items-center gap-2">
<ApiProviderLogo
providerType={preset.providerType}
className="size-4"
title={preset.displayName}
/>
{preset.displayName}
</span>
</SelectItem>
))}
</SelectGroup>
<SelectSeparator />
<SelectGroup>
{registry
.filter(
(entry) =>
!HIDDEN_PROVIDER_TYPES.has(entry.provider_type),
)
.map((entry) => (
<SelectItem
key={entry.provider_type}
value={entry.provider_type}
>
<span className="flex items-center gap-2">
<ApiProviderLogo
providerType={entry.provider_type}
className="size-4"
title={entry.display_name}
/>
{entry.display_name}
</span>
</SelectItem>
))}
</SelectGroup>
</SelectContent>
</Select>
</div>
<div className="grid grid-cols-[minmax(150px,0.8fr)_minmax(260px,1.2fr)] items-center gap-4 px-4 py-3 max-sm:grid-cols-1">
<div className="flex min-w-0 flex-col gap-0.5">
<Label
htmlFor="provider-api-key"
className="text-sm font-medium"
>
API key {isCustomProvider ? "(optional)" : ""}
</Label>
<p className="text-xs leading-snug text-muted-foreground">
Stored locally.
</p>
{showApiKeyField ? (
<div className="grid grid-cols-[minmax(150px,0.8fr)_minmax(260px,1.2fr)] items-center gap-4 px-4 py-3 max-sm:grid-cols-1">
<div className="flex min-w-0 flex-col gap-0.5">
<Label
htmlFor="provider-api-key"
className="text-sm font-medium"
>
API key {isCustomProvider ? "(optional)" : ""}
</Label>
<p className="text-xs leading-snug text-muted-foreground">
Stored locally.
</p>
</div>
<div className="relative min-w-0">
<Input
id="provider-api-key"
type={showApiKey ? "text" : "password"}
value={apiKey}
onChange={(event) => setApiKey(event.target.value)}
placeholder="Enter API key"
className="h-9 pr-9 text-sm"
/>
<button
type="button"
onClick={() => setShowApiKey((visible) => !visible)}
className="absolute top-1/2 right-1.5 flex size-5 -translate-y-1/2 items-center justify-center rounded text-muted-foreground transition-colors hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
aria-label={showApiKey ? "Hide API key" : "Show API key"}
aria-pressed={showApiKey}
>
{showApiKey ? (
<Eye className="size-3.5" />
) : (
<EyeOff className="size-3.5" />
)}
</button>
</div>
</div>
<div className="relative min-w-0">
<Input
id="provider-api-key"
type={showApiKey ? "text" : "password"}
value={apiKey}
onChange={(event) => setApiKey(event.target.value)}
placeholder="Enter API key"
className="h-9 pr-9 text-sm"
/>
<button
type="button"
onClick={() => setShowApiKey((visible) => !visible)}
className="absolute top-1/2 right-1.5 flex size-5 -translate-y-1/2 items-center justify-center rounded text-muted-foreground transition-colors hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
aria-label={showApiKey ? "Hide API key" : "Show API key"}
aria-pressed={showApiKey}
>
{showApiKey ? (
<Eye className="size-3.5" />
) : (
<EyeOff className="size-3.5" />
)}
</button>
</div>
</div>
) : null}
{isCustomProvider ? (
<div className="grid grid-cols-[minmax(150px,0.8fr)_minmax(260px,1.2fr)] items-center gap-4 px-4 py-3 max-sm:grid-cols-1">
@ -902,11 +976,35 @@ export function ChatProvidersSettings({
type="text"
value={baseUrlDraft}
onChange={(event) => setBaseUrlDraft(event.target.value)}
placeholder="https://my-vllm-server.com/v1"
placeholder={customProviderBaseUrlPlaceholder(providerType)}
className="h-9 text-sm"
/>
</div>
) : null}
{showReasoningToggle ? (
<div className="grid grid-cols-[minmax(150px,0.8fr)_minmax(260px,1.2fr)] items-center gap-4 px-4 py-3 max-sm:grid-cols-1">
<Label
htmlFor="provider-is-reasoning"
className="text-sm font-medium"
>
Reasoning model
</Label>
<label
htmlFor="provider-is-reasoning"
className="flex cursor-pointer items-center gap-2 text-sm"
>
<Checkbox
id="provider-is-reasoning"
checked={isReasoningModel}
onCheckedChange={(checked) =>
setIsReasoningModel(checked === true)
}
/>
This server runs a reasoning model
</label>
</div>
) : null}
</div>
</section>
@ -952,7 +1050,7 @@ export function ChatProvidersSettings({
}
title={
isCustomProvider
? "Custom providers use manual model IDs"
? "This connection uses manual model IDs"
: isCuratedModelList
? "Full catalog is not fetched for this provider"
: undefined
@ -986,7 +1084,7 @@ export function ChatProvidersSettings({
onChange={(event) =>
setManualModelIds(event.target.value)
}
placeholder={"gpt-4o-mini\nQwen/Qwen3-14B"}
placeholder={customProviderModelIdsPlaceholder(providerType)}
rows={5}
className="min-h-[100px] resize-y font-mono text-sm"
/>
@ -1197,9 +1295,9 @@ export function ChatProvidersSettings({
<div className="flex min-h-0 flex-col gap-6">
<header className="flex flex-col gap-1 pr-8">
<div className="flex min-w-0 flex-col gap-1">
<h1 className="font-heading text-lg font-semibold">Cloud</h1>
<h1 className="font-heading text-lg font-semibold">Connections</h1>
<p className="text-xs leading-relaxed text-muted-foreground">
Manage cloud provider connections for chat through the Studio proxy.
Manage model provider connections for chat through the Studio proxy.
</p>
</div>
</header>
@ -1237,7 +1335,8 @@ export function ChatProvidersSettings({
const detail =
provider.baseUrl || registryEntry?.base_url || "";
const providerLabel =
registryEntry?.display_name ?? provider.providerType;
registryEntry?.display_name ??
customProviderDisplayName(provider.providerType);
const modelSummary = formatModelSummary(provider.models);
return (
<div
@ -1346,9 +1445,9 @@ export function ChatProvidersDialog({
className="flex max-h-[90dvh] w-[96vw] flex-col gap-0 overflow-y-auto p-8 sm:max-w-none md:max-w-[44rem]"
>
<DialogHeader className="sr-only">
<DialogTitle>Cloud</DialogTitle>
<DialogTitle>Connections</DialogTitle>
<DialogDescription>
Manage external model providers for chat.
Manage external model connections for chat.
</DialogDescription>
</DialogHeader>
<ChatProvidersSettings

71
studio/frontend/src/features/chat/chat-settings-sheet.tsx
View file

@ -64,6 +64,12 @@ import { Fragment, type ReactNode } from "react";
import { useEffect, useMemo, useRef, useState } from "react";
import { toast } from "sonner";
import { useChatRuntimeStore } from "./stores/chat-runtime-store";
import {
type ExternalProviderConfig,
getExternalProviderApiKey,
parseExternalModelId,
supportsProviderPromptCaching,
} from "./external-providers";
import {
applyPresetParams,
BUILTIN_PRESET_NAMES,
@ -80,9 +86,11 @@ import {
toPresetParams,
type Preset,
} from "./presets/preset-policy";
import { OpenAICodeExecSection } from "./components/openai-code-exec-section";
import {
EXTERNAL_MAX_OUTPUT_TOKENS,
getExternalMinOutputTokens,
providerSupportsBuiltinCodeExecution,
type ProviderCapabilities,
} from "./provider-capabilities";
import type { InferenceParams } from "./types/runtime";
@ -235,7 +243,7 @@ function loadSavedActivePreset(): string {
}
}
function InfoHint({ children }: { children: ReactNode }) {
export function InfoHint({ children }: { children: ReactNode }) {
return (
<Tooltip>
<TooltipTrigger asChild>
@ -517,6 +525,8 @@ interface ChatSettingsPanelProps {
* per-param visibility in the sampling section.
*/
providerCapabilities?: ProviderCapabilities | null;
activeExternalProvider?: ExternalProviderConfig | null;
onExternalProviderChange?: (provider: ExternalProviderConfig) => void;
/**
* Backend provider type for the active external model (e.g. "kimi",
* "anthropic", "openai"), or `null` for local models. Drives the
@ -533,6 +543,8 @@ export function ChatSettingsPanel({
onParamsChange,
isExternalModel = false,
providerCapabilities = null,
activeExternalProvider = null,
onExternalProviderChange,
externalProviderType = null,
onReloadModel,
}: ChatSettingsPanelProps) {
@ -662,6 +674,26 @@ export function ChatSettingsPanel({
Boolean(currentCheckpoint) &&
modelRequiresTrustRemoteCode &&
!(params.trustRemoteCode ?? false);
const showPromptCachingControl =
activeExternalProvider != null &&
supportsProviderPromptCaching(activeExternalProvider.providerType);
const promptCachingEnabled =
activeExternalProvider?.enablePromptCaching !== false;
const externalSelection = currentCheckpoint
? parseExternalModelId(currentCheckpoint)
: null;
const showOpenAICodeExecSection =
activeExternalProvider != null &&
providerSupportsBuiltinCodeExecution(
activeExternalProvider.providerType,
externalSelection?.modelId,
activeExternalProvider.baseUrl,
) &&
activeExternalProvider.providerType === "openai";
const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId);
const openAiApiKeyForSection = activeExternalProvider
? getExternalProviderApiKey(activeExternalProvider.id) || null
: null;
function set<K extends keyof InferenceParams>(key: K) {
return (v: InferenceParams[K]) => {
@ -1145,6 +1177,43 @@ export function ChatSettingsPanel({
</div>
</CollapsibleSection>
{showPromptCachingControl && activeExternalProvider ? (
<CollapsibleSection label="Provider" defaultOpen={true}>
<div className="flex items-center justify-between gap-3 pt-1">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
Prompt caching
</span>
<InfoHint>
Reuse compatible prompt prefixes for lower latency and cost.
</InfoHint>
</div>
<Switch
className="panel-switch shrink-0"
checked={promptCachingEnabled}
onCheckedChange={(checked) => {
onExternalProviderChange?.({
...activeExternalProvider,
enablePromptCaching: checked,
});
}}
aria-label="Enable prompt caching"
/>
</div>
</CollapsibleSection>
) : null}
{showOpenAICodeExecSection && activeExternalProvider ? (
<CollapsibleSection label="Code Execution" defaultOpen={false}>
<OpenAICodeExecSection
provider={activeExternalProvider}
apiKey={openAiApiKeyForSection}
activeThreadId={activeThreadId}
onProviderChange={(p) => onExternalProviderChange?.(p)}
/>
</CollapsibleSection>
) : null}
<CollapsibleSection label="System Prompt" defaultOpen={true}>
<button
type="button"

697
studio/frontend/src/features/chat/components/openai-code-exec-section.tsx Normal file
View file

@ -0,0 +1,697 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
/**
* Settings-sheet section for OpenAI shell-tool container management.
* Renders only when:
* - active provider is OpenAI cloud (api.openai.com base URL), AND
* - the active model is gpt-5.5 or gpt-5.5-pro (the only families
* where the shell tool is wired through today).
*
* Surfaces three controls:
* 1. Default container idle-timeout (minutes). Persists on the
* provider record; pre-fills the create dialog and is used by
* the chat-adapter's lazy-create path on the first turn of a
* thread.
* 2. Container picker for the *active thread* — pick any of the
* user's existing OpenAI containers, or "Auto-create per thread"
* (default; lets the auto-create path manage it).
* 3. Create-new-container inline form. Refresh + delete actions
* per row.
*
* State persistence:
* - TTL → ExternalProviderConfig.openaiContainerTtlMinutes
* - Active container for this thread → ThreadRecord.openaiCodeExecContainerId
*
* No new global stores — list is fetched on open / refresh and held
* in component state.
*/
"use client";
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { toast } from "sonner";
import {
AlertDialog,
AlertDialogAction,
AlertDialogCancel,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { TrashIcon, RefreshCwIcon, PlusIcon } from "lucide-react";
import {
createOpenAIContainer,
deleteOpenAIContainer,
listOpenAIContainers,
type OpenAIContainerSummary,
} from "../api/openai-containers";
import { db } from "../db";
import type { ExternalProviderConfig } from "../external-providers";
import { useLiveQuery } from "../db";
import { ensureThreadRecord } from "../runtime-provider";
import { InfoHint } from "../chat-settings-sheet";
const AUTO_OPTION_VALUE = "__auto__";
const DEFAULT_TTL_MINUTES = 20;
const TTL_MIN = 1;
const TTL_MAX = 20; // OpenAI hard cap on expires_after.minutes
// Cadence for re-fetching the container list while the section is
// mounted. OpenAI's container TTL flips at minute granularity, so 30s
// is fast enough that an expired container loses its ACTIVE pill within
// half a minute without hammering /v1/containers.
const REFRESH_POLL_MS = 30_000;
function ageLabel(epochSeconds: number | null | undefined): string {
if (!epochSeconds) return "";
const ageSec = Math.max(0, Math.floor(Date.now() / 1000) - epochSeconds);
if (ageSec < 60) return `${ageSec}s ago`;
const ageMin = Math.floor(ageSec / 60);
if (ageMin < 60) return `${ageMin}m ago`;
const ageHr = Math.floor(ageMin / 60);
if (ageHr < 48) return `${ageHr}h ago`;
const ageDay = Math.floor(ageHr / 24);
return `${ageDay}d ago`;
}
function shortContainerId(id: string): string {
// Mid-truncate keeps the "cntr_" prefix readable and still surfaces the
// tail digits users sometimes copy off OpenAI's dashboard.
if (id.length <= 18) return id;
return `${id.slice(0, 12)}…${id.slice(-4)}`;
}
function isContainerRunning(c: OpenAIContainerSummary): boolean {
// OpenAI's containers API reports `status: "running"` while idle TTL is
// valid and `status: "expired"` once the idle window has passed. Treat
// a missing status as running so we don't false-positive on any older
// payloads that didn't include the field.
return c.status == null || c.status === "running";
}
interface OpenAICodeExecSectionProps {
provider: ExternalProviderConfig;
apiKey: string | null;
activeThreadId: string | null;
onProviderChange: (provider: ExternalProviderConfig) => void;
}
export function OpenAICodeExecSection({
provider,
apiKey,
activeThreadId,
onProviderChange,
}: OpenAICodeExecSectionProps) {
const [containers, setContainers] = useState<OpenAIContainerSummary[]>([]);
const [isLoading, setIsLoading] = useState(false);
const [creating, setCreating] = useState(false);
const [createOpen, setCreateOpen] = useState(false);
const [createName, setCreateName] = useState("");
// Ids that have been deleted in this session. Once tombstoned, an id
// stays hidden from the picker for the lifetime of the page — OpenAI's
// /containers list can keep returning a freshly-deleted id for an
// undocumented and variable amount of time, and an automatic re-show
// creates more confusion than it solves. Refreshing the page resets
// the tombstone naturally.
const [tombstones, setTombstones] = useState<Set<string>>(() => new Set());
// Ids optimistically inserted after a successful create but not yet
// confirmed by a /v1/containers list response. OpenAI's list endpoint
// is eventually consistent — a freshly-created container can be absent
// for several seconds. We render the row immediately with a "Creating"
// pill, then drop it from this set once a refresh sees the id.
const [pendingIds, setPendingIds] = useState<Set<string>>(() => new Set());
// Ref mirror so `refresh()` can read the current pending set without
// re-binding when it changes (the callback is in a useEffect dep).
const pendingIdsRef = useRef<Set<string>>(pendingIds);
useEffect(() => {
pendingIdsRef.current = pendingIds;
}, [pendingIds]);
// One-shot follow-up refresh scheduled after a create, to catch the
// common case where the server list lags the create response by a few
// seconds. Tracked so we can clear it on unmount.
const pendingRetryRef = useRef<number | null>(null);
// Target row for the destructive confirmation dialog. Held in state
// (rather than blocking with window.confirm) so the dialog sits inside
// the settings sheet instead of a native browser alert.
const [pendingDelete, setPendingDelete] =
useState<OpenAIContainerSummary | null>(null);
const [deleting, setDeleting] = useState(false);
const thread = useLiveQuery(
async () => (activeThreadId ? db.threads.get(activeThreadId) : undefined),
[activeThreadId],
);
const activeContainerId = thread?.openaiCodeExecContainerId ?? null;
// Hide just-deleted containers even if OpenAI's list still returns them.
// This is the single chokepoint — every downstream view (sorted picker,
// auto-bind candidate, all-containers list) derives from visibleContainers.
const visibleContainers = useMemo(() => {
if (tombstones.size === 0) return containers;
return containers.filter((c) => !tombstones.has(c.id));
}, [containers, tombstones]);
// Containers sorted newest-first by lastActiveAt so the dropdown's
// default (auto-bind target) shows up first.
const sortedContainers = useMemo(
() =>
[...visibleContainers].sort(
(a, b) => (b.lastActiveAt ?? 0) - (a.lastActiveAt ?? 0),
),
[visibleContainers],
);
// First running container by lastActiveAt — the auto-bind target and
// also what we surface visually before Dexie catches up.
const firstRunningContainer = useMemo(
() => sortedContainers.find(isContainerRunning) ?? null,
[sortedContainers],
);
// What the picker should treat as "active" right now. We decouple
// this from `activeContainerId` (Dexie state) so the user immediately
// sees the most-recent running container while the auto-bind effect's
// async write propagates. If the Dexie-bound container has since
// expired, fall back to the first running candidate — the stale-bind
// sweeper below will clear Dexie shortly after.
const boundContainer = useMemo(
() => sortedContainers.find((c) => c.id === activeContainerId) ?? null,
[sortedContainers, activeContainerId],
);
const displayedContainerId =
(boundContainer && isContainerRunning(boundContainer)
? boundContainer.id
: firstRunningContainer?.id) ?? null;
const refresh = useCallback(async () => {
if (!apiKey) return;
setIsLoading(true);
try {
const list = await listOpenAIContainers({
apiKey,
baseUrl: provider.baseUrl || null,
});
const serverIds = new Set(list.map((c) => c.id));
setContainers((prev) => {
// Preserve optimistic inserts the server hasn't acknowledged
// yet so they don't disappear on the reconciling refresh.
const orphans = prev.filter(
(c) => !serverIds.has(c.id) && pendingIdsRef.current.has(c.id),
);
return orphans.length > 0 ? [...orphans, ...list] : list;
});
setPendingIds((prev) => {
if (prev.size === 0) return prev;
const next = new Set(prev);
let changed = false;
for (const id of serverIds) {
if (next.delete(id)) changed = true;
}
return changed ? next : prev;
});
} catch (err) {
toast.error(
`Failed to list containers: ${err instanceof Error ? err.message : "Unknown"}`,
);
} finally {
setIsLoading(false);
}
}, [apiKey, provider.baseUrl]);
// Fetch once when the section mounts (or provider changes), then
// poll on a low cadence so an expired container's ACTIVE pill clears
// without the user clicking the refresh button. Also re-fetch when
// the tab regains visibility — covers the common case of leaving the
// sheet open across a long idle period.
useEffect(() => {
void refresh();
const interval = window.setInterval(() => {
if (document.visibilityState === "visible") {
void refresh();
}
}, REFRESH_POLL_MS);
const onVisibility = () => {
if (document.visibilityState === "visible") {
void refresh();
}
};
document.addEventListener("visibilitychange", onVisibility);
return () => {
window.clearInterval(interval);
document.removeEventListener("visibilitychange", onVisibility);
if (pendingRetryRef.current != null) {
window.clearTimeout(pendingRetryRef.current);
pendingRetryRef.current = null;
}
};
}, [refresh]);
// Auto-bind the active thread to the most-recently-active container
// whenever the thread has none set and at least one container exists
// on the user's OpenAI account. Sorting by `lastActiveAt` matches
// what feels "most recent" from the user's perspective.
//
// We eagerly materialize the thread row via `ensureThreadRecord` so
// the bind actually lands in Dexie before the user has sent a first
// message. This does NOT create anything at OpenAI — only a local
// ThreadRecord — so it does not bypass the user's expectation that
// a fresh OpenAI container is not created until first send.
//
// If `containers` is empty (no OpenAI containers exist yet), this
// effect short-circuits: the picker renders an empty-state hint and
// the chat-adapter's lazy-create path will mint the first container
// on first send.
useEffect(() => {
if (
!activeThreadId ||
activeContainerId ||
visibleContainers.length === 0
) {
return;
}
const sorted = [...visibleContainers].sort(
(a, b) => (b.lastActiveAt ?? 0) - (a.lastActiveAt ?? 0),
);
const candidate = sorted[0];
if (!candidate) return;
void (async () => {
try {
await ensureThreadRecord({
threadId: activeThreadId,
modelType: "base",
});
await db.threads.update(activeThreadId, {
openaiCodeExecContainerId: candidate.id,
});
} catch {
// Best-effort; the chat-adapter will inherit/create on send.
}
})();
}, [activeThreadId, activeContainerId, visibleContainers]);
const ttlValue = provider.openaiContainerTtlMinutes ?? DEFAULT_TTL_MINUTES;
const onTtlChange = (raw: string) => {
const n = parseInt(raw, 10);
if (Number.isNaN(n)) return;
const clamped = Math.min(Math.max(n, TTL_MIN), TTL_MAX);
onProviderChange({ ...provider, openaiContainerTtlMinutes: clamped });
};
const onPick = async (value: string) => {
if (!activeThreadId || !value) return;
// value is always a container id now — the "Auto-create per thread"
// option has been removed in favour of always defaulting to the
// most-recently-active container. The chat-adapter still handles
// the no-containers-exist case (lazy-create on first send).
//
// ensureThreadRecord materializes the thread row eagerly (modelType
// "base" — settings sheet is single-thread-mode only) so the update
// actually lands when the user hasn't sent a message yet.
try {
await ensureThreadRecord({ threadId: activeThreadId, modelType: "base" });
const affected = await db.threads.update(activeThreadId, {
openaiCodeExecContainerId: value,
});
if (affected === 0) {
toast.error("Could not update thread.");
}
} catch (err) {
toast.error(
`Could not update thread: ${err instanceof Error ? err.message : "Unknown"}`,
);
}
};
const onCreate = async () => {
if (!apiKey) return;
const name = createName.trim();
if (!name) {
toast.error("Container name is required");
return;
}
// TTL inherits from the section-level "Idle timeout" control —
// there is no per-container override on the form. Read it at
// submit time so a last-second change to the TTL row applies.
const ttlMinutes =
provider.openaiContainerTtlMinutes ?? DEFAULT_TTL_MINUTES;
setCreating(true);
try {
const created = await createOpenAIContainer(
{ apiKey, baseUrl: provider.baseUrl || null },
{ name, ttlMinutes },
);
toast.success(`Created container ${name}`);
setCreateName("");
setCreateOpen(false);
// Optimistic insert + "Creating" pill. OpenAI's /v1/containers
// list endpoint is eventually consistent and can omit the new
// container for several seconds — without this, the row only
// shows up on the next 30s poll or a manual refresh.
setContainers((prev) =>
prev.some((c) => c.id === created.id) ? prev : [created, ...prev],
);
setPendingIds((prev) => {
if (prev.has(created.id)) return prev;
const next = new Set(prev);
next.add(created.id);
return next;
});
// Follow-up refresh ~5s later to reconcile the optimistic row
// with the server's view once /v1/containers catches up. One
// shot; the regular poll covers any longer tail.
if (pendingRetryRef.current != null) {
window.clearTimeout(pendingRetryRef.current);
}
pendingRetryRef.current = window.setTimeout(() => {
pendingRetryRef.current = null;
void refresh();
}, 5000);
// Auto-bind the just-created container to the active thread.
// ensureThreadRecord first so the bind lands even when the user
// creates a container before sending the first message — without
// it, db.threads.update silently affects 0 rows and the chat
// adapter falls back to cross-thread inheritance / lazy-create,
// which can pick a stale container that fails with "container
// does not exist" on the first turn.
if (activeThreadId) {
try {
await ensureThreadRecord({
threadId: activeThreadId,
modelType: "base",
});
await db.threads.update(activeThreadId, {
openaiCodeExecContainerId: created.id,
});
} catch {
/* best-effort; toast above already confirmed creation */
}
}
} catch (err) {
toast.error(
`Create failed: ${err instanceof Error ? err.message : "Unknown"}`,
);
} finally {
setCreating(false);
// Refresh even on failure: the request may have partially succeeded
// server-side (created container, lost response), and a re-fetch
// keeps the picker in sync with OpenAI's actual state.
await refresh();
}
};
const confirmDelete = async () => {
if (!apiKey || !pendingDelete) return;
const { id, name } = pendingDelete;
setDeleting(true);
try {
await deleteOpenAIContainer(
{ apiKey, baseUrl: provider.baseUrl || null },
id,
);
// Tombstone the id so the picker hides it immediately even if
// OpenAI's list keeps returning it for a while.
setTombstones((prev) => {
if (prev.has(id)) return prev;
const next = new Set(prev);
next.add(id);
return next;
});
// Clear any thread bindings pointing at the now-deleted id.
const affected = await db.threads
.filter((t) => t.openaiCodeExecContainerId === id)
.toArray();
await Promise.all(
affected.map((t) =>
db.threads.update(t.id, { openaiCodeExecContainerId: null }),
),
);
toast.success(`Deleted container ${name || id}`);
} catch (err) {
toast.error(
`Delete failed: ${err instanceof Error ? err.message : "Unknown"}`,
);
} finally {
setDeleting(false);
setPendingDelete(null);
// Always refresh so a stale list entry (e.g. container deleted
// elsewhere, or already expired) is purged from the UI even when
// the delete call itself errored.
await refresh();
}
};
const displayActiveId = displayedContainerId;
return (
<div className="flex flex-col gap-3 pt-1">
{/* TTL */}
<div className="flex items-center justify-between gap-3">
<div className="flex min-w-0 items-center gap-1.5">
<label
htmlFor="openai-container-ttl"
className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg"
>
Idle timeout
</label>
<InfoHint>
Minutes a newly-created container stays alive between calls.
OpenAI caps this at 20.
</InfoHint>
</div>
<Input
id="openai-container-ttl"
type="number"
min={TTL_MIN}
max={TTL_MAX}
value={ttlValue}
onChange={(e) => onTtlChange(e.target.value)}
className="h-8 w-14 px-2 text-center text-sm tabular-nums"
/>
</div>
{/* Single container list. Clicking a row binds it to the active
thread and the ACTIVE pill marks which one — no separate
picker needed. */}
<div className="flex flex-col gap-1.5">
<div className="flex items-center justify-between gap-2">
<span className="text-[11px] uppercase tracking-wider text-muted-foreground">
Containers
</span>
<Button
size="sm"
variant="ghost"
className="-mr-1 h-6 w-6 p-0 text-muted-foreground"
onClick={() => void refresh()}
disabled={isLoading || !apiKey}
aria-label="Refresh container list"
>
<RefreshCwIcon
className={`size-3.5 ${isLoading ? "animate-spin" : ""}`}
/>
</Button>
</div>
{sortedContainers.length === 0 ? (
// Quiet placeholder with the same muted border as row cards
// so an empty section doesn't masquerade as an active control.
// The first container is minted by the chat-adapter on first
// send (lazy-create) and appears here after the next refresh.
<div className="flex h-9 w-full items-center rounded-md border border-dashed border-border/60 bg-muted/20 px-2 text-xs text-muted-foreground">
None yet - one will be created on first send.
</div>
) : (
<ul className="flex max-h-52 flex-col gap-1 overflow-auto">
{sortedContainers.map((c) => {
const running = isContainerRunning(c);
const isActive = running && c.id === displayActiveId;
const isPending = pendingIds.has(c.id);
const ttlMinutes = c.expiresAfterMinutes ?? DEFAULT_TTL_MINUTES;
const canActivate =
activeThreadId != null && !isActive && running;
const statusLabel = !running ? (c.status ?? "expired") : null;
return (
<li
key={c.id}
className={`flex items-center gap-2 rounded-md border px-2 py-1.5 text-xs transition-colors ${
isActive
? "border-primary/30 bg-primary/5"
: "border-border/60 hover:bg-muted/40"
} ${canActivate ? "cursor-pointer" : ""} ${
running ? "" : "opacity-60"
}`}
onClick={() => {
if (canActivate) void onPick(c.id);
}}
onKeyDown={(e) => {
if (!canActivate) return;
if (e.key === "Enter" || e.key === " ") {
e.preventDefault();
void onPick(c.id);
}
}}
tabIndex={canActivate ? 0 : undefined}
role={canActivate ? "button" : undefined}
aria-pressed={isActive}
title={
canActivate
? "Use this container for the active thread"
: !running
? `Container is ${statusLabel}`
: undefined
}
>
{/* min-w-0 + truncate keeps long OpenAI container ids
from spilling under the trash button on narrow
settings sheets. */}
<div className="flex min-w-0 flex-1 flex-col gap-0.5">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 truncate font-medium">
{c.name ?? "(unnamed)"}
</span>
{isPending ? (
<span className="shrink-0 rounded-sm bg-muted px-1 py-px text-[9px] font-medium uppercase tracking-wider text-muted-foreground">
Creating
</span>
) : isActive ? (
<span className="shrink-0 rounded-sm bg-primary/15 px-1 py-px text-[9px] font-medium uppercase tracking-wider text-primary">
Active
</span>
) : statusLabel ? (
<span className="shrink-0 rounded-sm bg-muted px-1 py-px text-[9px] font-medium uppercase tracking-wider text-muted-foreground">
{statusLabel}
</span>
) : null}
</div>
<div
className="flex min-w-0 items-center gap-1.5 text-muted-foreground"
title={c.id}
>
<span className="min-w-0 truncate font-mono text-[11px]">
{shortContainerId(c.id)}
</span>
<span className="shrink-0 text-[10px] uppercase tracking-wider">
· {ttlMinutes}m
</span>
</div>
</div>
<Button
size="sm"
variant="ghost"
className="h-6 w-6 shrink-0 p-0 text-muted-foreground hover:text-destructive"
onClick={(e) => {
e.stopPropagation();
setPendingDelete(c);
}}
aria-label={`Delete container ${c.name ?? c.id}`}
>
<TrashIcon className="size-3.5" />
</Button>
</li>
);
})}
</ul>
)}
</div>
{/* Create new — inline single-row edit that visually echoes a
container card. TTL is inherited from the section's top
"Idle timeout" control (no per-container override), which
keeps the form light and avoids a duplicated input. */}
{createOpen ? (
<div className="flex items-center gap-1 rounded-md border border-border/60 bg-muted/20 px-1.5 py-1">
<Input
autoFocus
placeholder="Name"
value={createName}
onChange={(e) => setCreateName(e.target.value)}
onKeyDown={(e) => {
if (e.key === "Enter") {
e.preventDefault();
if (createName.trim() && !creating && apiKey) {
void onCreate();
}
} else if (e.key === "Escape") {
e.preventDefault();
setCreateOpen(false);
setCreateName("");
}
}}
className="h-7 min-w-0 flex-1 border-0 bg-transparent px-1.5 text-xs shadow-none focus-visible:ring-0"
/>
<Button
size="sm"
variant="ghost"
className="h-7 shrink-0 px-2 text-xs"
onClick={() => {
setCreateOpen(false);
setCreateName("");
}}
disabled={creating}
>
Cancel
</Button>
<Button
size="sm"
className="h-7 shrink-0 px-3 text-xs"
onClick={() => void onCreate()}
disabled={creating || !createName.trim() || !apiKey}
>
{creating ? "Creating…" : "Create"}
</Button>
</div>
) : (
<Button
size="sm"
variant="outline"
className="h-8"
onClick={() => setCreateOpen(true)}
disabled={!apiKey}
>
<PlusIcon className="size-3.5 mr-1" />
New container
</Button>
)}
<AlertDialog
open={pendingDelete !== null}
onOpenChange={(nextOpen) => {
if (!nextOpen && deleting) return;
if (!nextOpen) setPendingDelete(null);
}}
>
<AlertDialogContent size="sm">
<AlertDialogHeader>
<AlertDialogTitle>
Delete{" "}
<span className="font-mono">
{pendingDelete?.name ?? "container"}
</span>
?
</AlertDialogTitle>
<AlertDialogDescription>
Threads using this container will fall back to auto-create on
their next turn. This cannot be undone.
</AlertDialogDescription>
</AlertDialogHeader>
<AlertDialogFooter>
<AlertDialogCancel disabled={deleting}>Cancel</AlertDialogCancel>
<AlertDialogAction
variant="destructive"
disabled={deleting}
onClick={(e) => {
e.preventDefault();
void confirmDelete();
}}
>
{deleting ? "Deleting…" : "Delete"}
</AlertDialogAction>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
</div>
);
}

159
studio/frontend/src/features/chat/external-providers.ts
View file

@ -14,10 +14,154 @@ export interface ExternalProviderConfig {
models: string[];
/** Cached available model ids from the provider's /models response. */
availableModels?: string[];
/** Whether to ask supported hosted providers to use prompt caching. */
enablePromptCaching?: boolean;
/** User-pinned: the loaded vLLM model supports `enable_thinking`. */
isReasoningModel?: boolean;
/**
* Default idle-timeout (in minutes) for newly created OpenAI shell
* containers. Pre-fills the "Create container" dialog and is the
* TTL the auto-create-per-thread path POSTs to /v1/containers with.
* OpenAI's hard default is 20. Only meaningful for OpenAI cloud.
*/
openaiContainerTtlMinutes?: number;
createdAt: number;
updatedAt: number;
}
const PROMPT_CACHING_PROVIDER_TYPES = new Set(["openai", "anthropic"]);
export function supportsProviderPromptCaching(
providerType: string | null | undefined,
): boolean {
return providerType != null && PROMPT_CACHING_PROVIDER_TYPES.has(providerType);
}
// Provider types that expose the connection-level "reasoning model"
// toggle. vLLM's OpenAI-compat endpoint doesn't advertise this per model.
const REASONING_TOGGLE_PROVIDER_TYPES = new Set(["vllm"]);
export function supportsProviderReasoningToggle(
providerType: string | null | undefined,
): boolean {
return (
providerType != null && REASONING_TOGGLE_PROVIDER_TYPES.has(providerType)
);
}
export const CUSTOM_BACKEND_PROVIDER_TYPE = "openai";
export const LEGACY_CUSTOM_PROVIDER_TYPE = "custom";
export const CUSTOM_PROVIDER_PRESETS = [
{
providerType: "llama_cpp",
displayName: "llama.cpp",
baseUrlPlaceholder: "http://localhost:8080/v1",
modelIdsPlaceholder: "gpt-oss-20b\nqwen3-14b",
},
{
providerType: "vllm",
displayName: "vLLM",
baseUrlPlaceholder: "https://my-vllm-server.com/v1",
modelIdsPlaceholder: "openai/gpt-oss-20b\nQwen/Qwen3-14B",
},
{
providerType: "ollama",
displayName: "Ollama",
baseUrlPlaceholder: "http://localhost:11434/v1",
modelIdsPlaceholder: "gpt-oss:20b\nqwen3:14b",
},
] as const;
const CUSTOM_PROVIDER_LABELS: Record<string, string> = {
[LEGACY_CUSTOM_PROVIDER_TYPE]: "Custom",
...Object.fromEntries(
CUSTOM_PROVIDER_PRESETS.map((preset) => [
preset.providerType,
preset.displayName,
]),
),
};
const CUSTOM_PROVIDER_BASE_URL_PLACEHOLDERS: Record<string, string> = {
[LEGACY_CUSTOM_PROVIDER_TYPE]: "https://my-vllm-server.com/v1",
...Object.fromEntries(
CUSTOM_PROVIDER_PRESETS.map((preset) => [
preset.providerType,
preset.baseUrlPlaceholder,
]),
),
};
const CUSTOM_PROVIDER_MODEL_IDS_PLACEHOLDERS: Record<string, string> = {
[LEGACY_CUSTOM_PROVIDER_TYPE]: "openai/gpt-oss-20b\nQwen/Qwen3-14B",
...Object.fromEntries(
CUSTOM_PROVIDER_PRESETS.map((preset) => [
preset.providerType,
preset.modelIdsPlaceholder,
]),
),
};
export function isCustomProviderType(
providerType: string | null | undefined,
): boolean {
if (!providerType) return false;
return providerType in CUSTOM_PROVIDER_LABELS;
}
export function customProviderDisplayName(
providerType: string | null | undefined,
): string {
if (!providerType) return "Custom";
return CUSTOM_PROVIDER_LABELS[providerType] ?? providerType;
}
export function customProviderBaseUrlPlaceholder(
providerType: string | null | undefined,
): string {
if (!providerType) {
return CUSTOM_PROVIDER_BASE_URL_PLACEHOLDERS[LEGACY_CUSTOM_PROVIDER_TYPE];
}
return (
CUSTOM_PROVIDER_BASE_URL_PLACEHOLDERS[providerType] ??
CUSTOM_PROVIDER_BASE_URL_PLACEHOLDERS[LEGACY_CUSTOM_PROVIDER_TYPE]
);
}
export function customProviderModelIdsPlaceholder(
providerType: string | null | undefined,
): string {
if (!providerType) {
return CUSTOM_PROVIDER_MODEL_IDS_PLACEHOLDERS[LEGACY_CUSTOM_PROVIDER_TYPE];
}
return (
CUSTOM_PROVIDER_MODEL_IDS_PLACEHOLDERS[providerType] ??
CUSTOM_PROVIDER_MODEL_IDS_PLACEHOLDERS[LEGACY_CUSTOM_PROVIDER_TYPE]
);
}
export function toExternalBackendProviderType(providerType: string): string;
export function toExternalBackendProviderType(
providerType: null | undefined,
): undefined;
export function toExternalBackendProviderType(
providerType: string | null | undefined,
): string | undefined;
export function toExternalBackendProviderType(
providerType: string | null | undefined,
): string | undefined {
if (!providerType) return undefined;
// vLLM's /v1/responses applies the loaded model's chat template, which
// 400s on strict-alternation templates (e.g. Gemma 3). Pass the actual
// type through so the backend routes vLLM to /v1/chat/completions instead
// of the OpenAI Responses path used for gpt-5.x.
if (providerType === "vllm") return "vllm";
return isCustomProviderType(providerType)
? CUSTOM_BACKEND_PROVIDER_TYPE
: providerType;
}
const EXTERNAL_PROVIDERS_KEY = "unsloth_chat_external_providers";
const EXTERNAL_PROVIDER_KEYS_KEY = "unsloth_chat_external_provider_keys";
const EXTERNAL_MODEL_PREFIX = "external::";
@ -71,9 +215,10 @@ function mapLegacyPresetToProviderType(presetId: string): string {
}
function normalizeProvider(raw: ExternalProviderConfig): ExternalProviderConfig {
const providerType = raw.providerType.trim();
return {
...raw,
providerType: raw.providerType.trim(),
providerType,
name: raw.name.trim(),
baseUrl: raw.baseUrl.trim(),
models: raw.models
@ -82,6 +227,18 @@ function normalizeProvider(raw: ExternalProviderConfig): ExternalProviderConfig
availableModels: (raw.availableModels ?? [])
.map((model) => model.trim())
.filter((model) => model.length > 0),
enablePromptCaching: supportsProviderPromptCaching(providerType)
? raw.enablePromptCaching !== false
: undefined,
isReasoningModel: supportsProviderReasoningToggle(providerType)
? raw.isReasoningModel === true
: undefined,
openaiContainerTtlMinutes:
providerType === "openai" &&
typeof raw.openaiContainerTtlMinutes === "number" &&
raw.openaiContainerTtlMinutes >= 1
? Math.min(raw.openaiContainerTtlMinutes, 20)
: undefined,
};
}

244
studio/frontend/src/features/chat/lib/friendly-names.ts Normal file
View file

@ -0,0 +1,244 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
/**
* Friendly default names for auto-created OpenAI shell containers.
* Used by the chat-adapter when the lazy-create path fires (Code pill
* on, no thread container yet, user has set a non-default TTL). The
* goal is a human-memorable label like "otter" or "harbor" instead of
* "chat-abc12345" — the user can still rename via the Studio-side
* alias map.
*
* The list is curated to:
* - Be unambiguous, non-offensive nouns from natural categories
* (animals, plants, geography, materials, weather).
* - Avoid technical / political / brand words that might read as
* odd in a chat UI.
* - Stay reasonably small so the bundle cost is negligible (~200
* entries × ~7 bytes ≈ 1.5 KB).
*
* Collisions are tolerated — the container's real unique key is its
* ``cntr_*`` id, not its name. A short random hex suffix is appended
* to make accidental same-name collisions visually distinct in the
* picker list.
*/
const WORDS = [
// animals
"otter",
"falcon",
"heron",
"lynx",
"marten",
"stoat",
"raven",
"magpie",
"salmon",
"trout",
"perch",
"tortoise",
"gecko",
"iguana",
"axolotl",
"narwhal",
"manatee",
"dolphin",
"porpoise",
"octopus",
"cuttlefish",
"nautilus",
"starfish",
"urchin",
"anemone",
"coral",
"puffin",
"kestrel",
"osprey",
"buzzard",
"kingfisher",
"robin",
"wren",
"finch",
"sparrow",
"thrush",
"siskin",
"warbler",
"tanager",
"oriole",
"hare",
"badger",
"weasel",
"ferret",
"polecat",
"civet",
"tapir",
"okapi",
"ibex",
"chamois",
// plants & trees
"alder",
"aspen",
"birch",
"cedar",
"cypress",
"elder",
"elm",
"fir",
"ginkgo",
"hawthorn",
"hazel",
"hemlock",
"holly",
"juniper",
"larch",
"linden",
"maple",
"oak",
"olive",
"pine",
"rowan",
"spruce",
"sycamore",
"willow",
"yew",
"thistle",
"fern",
"moss",
"ivy",
"clover",
"heather",
"lavender",
"rosemary",
"sage",
"thyme",
"myrtle",
"laurel",
"magnolia",
// geography / landscape
"harbor",
"atoll",
"lagoon",
"estuary",
"fjord",
"delta",
"isthmus",
"mesa",
"plateau",
"valley",
"ridge",
"summit",
"glade",
"meadow",
"moor",
"heath",
"tundra",
"savanna",
"prairie",
"steppe",
"bayou",
"marsh",
"fen",
"grotto",
"cavern",
"canyon",
"ravine",
"gorge",
"knoll",
"dell",
"vale",
"coast",
// materials / minerals / colors
"amber",
"agate",
"onyx",
"opal",
"jade",
"quartz",
"obsidian",
"basalt",
"granite",
"marble",
"slate",
"flint",
"lapis",
"topaz",
"garnet",
"pearl",
"coral",
"ivory",
"ebony",
"copper",
"cobalt",
"indigo",
"saffron",
"vermilion",
"ochre",
"umber",
"sienna",
"russet",
// weather / sky / time
"aurora",
"comet",
"ember",
"frost",
"gale",
"harvest",
"monsoon",
"nebula",
"solstice",
"twilight",
"zephyr",
"drizzle",
"tempest",
"halcyon",
"equinox",
"rainbow",
"horizon",
"meridian",
"zenith",
"comet",
// misc tactile / cozy nouns
"lantern",
"kettle",
"compass",
"anchor",
"beacon",
"harbor",
"voyage",
"trellis",
"cottage",
"thicket",
"orchard",
"bramble",
"haystack",
"snowfall",
"campfire",
];
/** RFC 4122-ish 4-character lowercase hex suffix using crypto.randomUUID. */
function randomHexSuffix(): string {
if (
typeof crypto !== "undefined" &&
typeof crypto.randomUUID === "function"
) {
return crypto.randomUUID().replace(/-/g, "").slice(0, 4);
}
// Older browser fallback. Math.random is fine here — this is a
// display suffix, not a security token.
return Math.floor(Math.random() * 0xffff)
.toString(16)
.padStart(4, "0");
}
/**
* Returns a single English-word name with a short random hex suffix.
*
* Example output: "kestrel-3f9c", "harbor-a012".
*
* The suffix keeps containers visually distinguishable in the picker
* when the same word recurs across creations.
*/
export function pickFriendlyContainerName(): string {
const word = WORDS[Math.floor(Math.random() * WORDS.length)] ?? "container";
return `${word}-${randomHexSuffix()}`;
}

158
studio/frontend/src/features/chat/provider-capabilities.ts
View file

@ -83,6 +83,126 @@ export function clampReasoningEffortToLevels(
*/
export const EXTERNAL_MAX_OUTPUT_TOKENS = 32768;
/**
* Whether the external provider offers a built-in web-search tool that the
* model invokes server-side. When `true`, the chat composer's Search button
* is available for that provider and the chat-adapter forwards
* `enable_tools: true, enabled_tools: ["web_search"]` on the request — the
* backend routes the call through the provider's tool schema:
* - OpenAI: `tools: [{type: "web_search"}]` on /v1/responses
* - Anthropic: `tools: [{type: "web_search_20250305", name: "web_search",
* max_uses: 5}]` on /v1/messages
* - OpenRouter: `plugins: [{id: "web"}]` on /v1/chat/completions (the
* router's universal web-search shape; works for every
* underlying model including the `openrouter/free` router).
* - Kimi: `tools: [{type: "builtin_function", function: {name:
* "$web_search"}}]` with `thinking: {type:
* "disabled"}`. Requires a client round-trip:
* the first call returns the search args; the backend
* echoes them back as a role=tool message; the second
* call streams the answer. Handled in
* _stream_kimi_web_search on the backend.
*
* Mistral is intentionally excluded: their `web_search` connector lives on
* the Agents API (`/v1/agents` + `/v1/conversations`), not chat completions,
* and returns `"WebSearchTool connector is not supported"` if injected into
* /v1/chat/completions. Wiring it would require a dedicated Agents streaming
* path. Gemini's grounded-search can be added with the same pattern when
* matching backend translation lands.
*/
export function providerSupportsBuiltinWebSearch(
providerType: string | null | undefined,
): boolean {
return (
providerType === "openai" ||
providerType === "anthropic" ||
providerType === "openrouter" ||
providerType === "kimi"
);
}
/**
* Whether the selected external provider/model exposes a server-side
* code-execution tool. Two providers ship one today:
*
* - **Anthropic** (`code_execution_20250825`): Python + bash +
* str_replace-based file edits inside a 5 GB sandboxed container
* per request. Documented at
* https://platform.claude.com/docs/en/agents-and-tools/tool-use/code-execution-tool
*
* - **OpenAI cloud** (`shell` on /v1/responses): bash inside a
* reusable container; we auto-create one on the first turn of a
* chat thread and reference it on subsequent turns via the
* thread's stored `openaiCodeExecContainerId`. Documented at
* https://developers.openai.com/api/docs/guides/tools-shell
*
* Returns false for every other provider. The backend additionally
* gates the OpenAI shell tool on `is_openai_cloud` so custom
* OpenAI-compat servers (ollama / llama.cpp / vLLM) that also report
* `provider_type="openai"` never receive the tool — but in practice
* none of those catalogs surface the `gpt-5.5` ids anyway, so the
* frontend prefix match is enough.
*
* v1 wires the tools themselves; file uploads (Anthropic
* `container_upload` / OpenAI `input_file`) are a deliberate follow-up.
*/
const ANTHROPIC_CODE_EXECUTION_MODEL_PREFIXES = [
"claude-opus-4-7",
"claude-opus-4-6",
"claude-sonnet-4-6",
"claude-opus-4-5",
"claude-sonnet-4-5",
"claude-haiku-4-5",
// Deprecated upstream but the registry still exposes the ids, so the
// pill should remain functional for users on those snapshots.
"claude-opus-4-1",
"claude-opus-4",
"claude-sonnet-4",
] as const;
// OpenAI cloud shell-tool gating. Docs only explicitly demonstrate
// gpt-5.5; gpt-5.5-pro is included because the family share the same
// /v1/responses contract. `gpt-5.5-pro` is checked first so the prefix
// match doesn't collide with a hypothetical `gpt-5.5-turbo` etc.
const OPENAI_CODE_EXECUTION_MODEL_PREFIXES = [
"gpt-5.5-pro",
"gpt-5.5",
] as const;
/**
* Strict check that a provider configuration points at OpenAI's
* managed cloud (api.openai.com), as opposed to a custom OpenAI-compat
* backend (ollama / llama.cpp / vLLM / generic "custom" preset). The
* shell tool ONLY exists on OpenAI cloud; sending it to anything else
* 400s the request. Mirror of the backend's
* `is_openai_cloud = "api.openai.com" in self.base_url` guard.
*/
function isOpenAICloudBaseUrl(baseUrl: string | null | undefined): boolean {
if (!baseUrl) return true; // No override → uses the default openai.com base.
return baseUrl.trim().toLowerCase().includes("api.openai.com");
}
export function providerSupportsBuiltinCodeExecution(
providerType: string | null | undefined,
modelId: string | null | undefined,
baseUrl?: string | null,
): boolean {
const normalized = modelId?.trim().toLowerCase() ?? "";
if (!normalized) return false;
if (providerType === "anthropic") {
return ANTHROPIC_CODE_EXECUTION_MODEL_PREFIXES.some((prefix) =>
normalized.startsWith(prefix),
);
}
if (providerType === "openai") {
if (!isOpenAICloudBaseUrl(baseUrl)) return false;
return OPENAI_CODE_EXECUTION_MODEL_PREFIXES.some((prefix) =>
normalized.startsWith(prefix),
);
}
return false;
}
/**
* Per-provider minimum on the outbound max_tokens. Kimi's docs require
* `max_tokens >= 16000` whenever a thinking model is in use so the
@ -183,10 +303,13 @@ const PROVIDER_CAPABILITIES: Record<string, ProviderCapabilities> = {
// OpenRouter silently drops params the target model does not support, so we
// surface every knob and let the gateway handle the per-model fan-out.
openrouter: ALL_SUPPORTED,
// Custom providers are assumed OpenAI-compatible by the backend; users who
// point at vLLM/Ollama backends often want top_k / min_p / repetition,
// so be permissive.
// Local OpenAI-compatible connections are proxied through the OpenAI backend
// path, but vLLM/Ollama/llama.cpp users often want top_k / min_p /
// repetition controls, so be permissive.
custom: ALL_SUPPORTED,
vllm: ALL_SUPPORTED,
ollama: ALL_SUPPORTED,
llama_cpp: ALL_SUPPORTED,
};
const DEFAULT_EXTERNAL_CAPABILITIES = OPENAI_COMPAT_BASE;
@ -239,7 +362,7 @@ const NO_REASONING_CAPS: ReasoningCaps = {
const ANTHROPIC_REASONING_MODELS = [
{
prefixes: ["claude-opus-4-7"],
levels: ["none", "low", "medium", "high", "xhigh"],
levels: ["none", "low", "medium", "high", "xhigh", "max"],
},
{
prefixes: ["claude-opus-4-6", "claude-sonnet-4-6"],
@ -382,6 +505,25 @@ function resolveMistralReasoningCapabilities(modelId: string): ExternalReasoning
return withEnableThinkingStyle();
}
export interface ExternalReasoningResolveOptions {
/** vLLM connection flagged as a reasoning model in provider config. */
isReasoningProvider?: boolean;
}
// vLLM has no per-model reasoning signal on OpenAI-compat — pin via user toggle.
function resolveConnectionLevelReasoning(
normalizedProvider: string,
options: ExternalReasoningResolveOptions | undefined,
): ExternalReasoningCapabilities | null {
if (normalizedProvider === "vllm" && options?.isReasoningProvider) {
return withEnableThinkingStyle({
supportsReasoning: true,
supportsReasoningOff: true,
});
}
return null;
}
/**
* resolve external-model thinking capabilities.
* provider-specific matching lives in the OpenAI/Anthropic resolvers.
@ -390,9 +532,17 @@ function resolveMistralReasoningCapabilities(modelId: string): ExternalReasoning
export function getExternalReasoningCapabilities(
providerType: string | null | undefined,
modelId: string | null | undefined,
options?: ExternalReasoningResolveOptions,
): ExternalReasoningCapabilities {
const normalizedModel = modelId?.trim().toLowerCase() ?? "";
const normalizedProvider = providerType?.trim().toLowerCase() ?? "";
const connectionLevel = resolveConnectionLevelReasoning(
normalizedProvider,
options,
);
if (connectionLevel) {
return connectionLevel;
}
if (!normalizedModel) {
return withEnableThinkingStyle();
}

2
studio/frontend/src/features/chat/runtime-provider.tsx
View file

@ -396,7 +396,7 @@ function toThreadMessage(m: MessageRecord): ThreadMessage {
};
}
async function ensureThreadRecord({
export async function ensureThreadRecord({
threadId,
modelType,
pairId,

76
studio/frontend/src/features/chat/shared-composer.tsx
View file

@ -24,7 +24,10 @@ import {
type ReasoningEffort,
useChatRuntimeStore,
} from "./stores/chat-runtime-store";
import { getExternalReasoningCapabilities } from "./provider-capabilities";
import {
getExternalReasoningCapabilities,
providerSupportsBuiltinCodeExecution,
} from "./provider-capabilities";
import {
type CompositionEvent,
type KeyboardEvent,
@ -304,6 +307,9 @@ export function SharedComposer({
const preserveThinking = useChatRuntimeStore((s) => s.preserveThinking);
const setPreserveThinking = useChatRuntimeStore((s) => s.setPreserveThinking);
const supportsTools = useChatRuntimeStore((s) => s.supportsTools);
const supportsBuiltinWebSearch = useChatRuntimeStore(
(s) => s.supportsBuiltinWebSearch,
);
const toolsEnabled = useChatRuntimeStore((s) => s.toolsEnabled);
const setToolsEnabled = useChatRuntimeStore((s) => s.setToolsEnabled);
const codeToolsEnabled = useChatRuntimeStore((s) => s.codeToolsEnabled);
@ -327,6 +333,10 @@ export function SharedComposer({
? getExternalReasoningCapabilities(
selectedExternalProvider?.providerType,
effectiveExternalModelId,
{
isReasoningProvider:
selectedExternalProvider?.isReasoningModel === true,
},
)
: null;
const isExternalOpenAIReasoning =
@ -345,13 +355,39 @@ export function SharedComposer({
const reasoningLockedOn =
effectiveSupportsReasoning &&
(effectiveReasoningAlwaysOn || !effectiveSupportsReasoningOff);
// Kimi's $web_search builtin mandates thinking=disabled per the docs at
// https://platform.kimi.ai/docs/guide/use-web-search. Both pills stay
// clickable for Kimi, but turning one on flips the other off — the
// click handlers below enforce this mutual exclusion so the visible
// state always matches what the backend actually sends.
const isKimiExternal = selectedExternalProvider?.providerType === "kimi";
const effectiveReasoningEnabled = reasoningLockedOn ? true : reasoningEnabled;
const effectiveReasoningVisualEnabled =
effectiveReasoningEnabled && reasoningEffort !== "none";
const reasoningDisabled = !modelLoaded || !effectiveSupportsReasoning;
const showReasoningControl =
effectiveSupportsReasoning || effectiveReasoningAlwaysOn;
const toolsDisabled = !modelLoaded || !supportsTools;
// Two-pill gating: Search pill lights up when the runtime has either
// a local tool runtime (supportsTools, gives us our Code/python + local
// web_search) OR a server-side web_search the provider runs for us
// (supportsBuiltinWebSearch, currently OpenAI / Anthropic / OpenRouter
// / Kimi). Code pill lights up on the local runtime OR when Anthropic
// is selected with a model that accepts the server-side
// code_execution_20250825 tool — see
// providerSupportsBuiltinCodeExecution. Anthropic is the only external
// provider that ships a code-execution tool today.
const supportsBuiltinCodeExecution = providerSupportsBuiltinCodeExecution(
selectedExternalProvider?.providerType,
effectiveExternalModelId,
selectedExternalProvider?.baseUrl,
);
const searchDisabled =
!modelLoaded || !(supportsTools || supportsBuiltinWebSearch);
const codeDisabled =
!modelLoaded || !(supportsTools || supportsBuiltinCodeExecution);
// Backwards-compatible alias for any other call site that may still
// reference `toolsDisabled` (rare; both pills used it before).
const toolsDisabled = codeDisabled;
const setPendingAudioStore = useChatRuntimeStore((s) => s.setPendingAudio);
const clearPendingAudioStore = useChatRuntimeStore((s) => s.clearPendingAudio);
@ -654,6 +690,9 @@ export function SharedComposer({
placeholder="Send to both models..."
className="composer-input"
rows={1}
// dir="auto" auto-detects RTL (Arabic / Hebrew / Persian / Urdu)
// from the first strong character; no effect on LTR scripts.
dir="auto"
/>
<div className="composer-action-wrapper">
<div className="flex items-center gap-1">
@ -766,6 +805,11 @@ export function SharedComposer({
setReasoningEffort(level);
setReasoningEnabled(true);
applyQwenThinkingParams(true);
// Mutual exclusion: turning thinking on for a
// Kimi model forces the web_search builtin off.
if (isKimiExternal && toolsEnabled) {
setToolsEnabled(false);
}
}}
>
{formatReasoningEffortLabel(level, externalSelection?.modelId)}
@ -789,6 +833,12 @@ export function SharedComposer({
const next = !reasoningEnabled;
setReasoningEnabled(next);
applyQwenThinkingParams(next);
// Mutual exclusion: Kimi's $web_search builtin
// requires thinking off, so turning thinking on flips
// the Search pill off (and vice versa).
if (isKimiExternal && next && toolsEnabled) {
setToolsEnabled(false);
}
}}
className={cn(
"flex items-center gap-1.5 rounded-full px-2.5 py-1 text-xs font-medium transition-colors",
@ -845,10 +895,22 @@ export function SharedComposer({
)}
<button
type="button"
disabled={toolsDisabled}
onClick={() => setToolsEnabled(!toolsEnabled)}
disabled={searchDisabled}
onClick={() => {
const next = !toolsEnabled;
setToolsEnabled(next);
// Kimi's $web_search builtin requires thinking=disabled
// (https://platform.kimi.ai/docs/guide/use-web-search).
// Toggle the Think pill off when Search comes on, and
// back on when Search goes off — mutual exclusion that
// mirrors what the backend enforces.
if (isKimiExternal) {
setReasoningEnabled(!next);
applyQwenThinkingParams(!next);
}
}}
className="composer-pill-btn"
data-active={toolsEnabled && !toolsDisabled ? "true" : "false"}
data-active={toolsEnabled && !searchDisabled ? "true" : "false"}
aria-label={toolsEnabled ? "Disable web search" : "Enable web search"}
>
<GlobeIcon className="size-3.5" />
@ -856,10 +918,10 @@ export function SharedComposer({
</button>
<button
type="button"
disabled={toolsDisabled}
disabled={codeDisabled}
onClick={() => setCodeToolsEnabled(!codeToolsEnabled)}
className="composer-pill-btn"
data-active={codeToolsEnabled && !toolsDisabled ? "true" : "false"}
data-active={codeToolsEnabled && !codeDisabled ? "true" : "false"}
aria-label={codeToolsEnabled ? "Disable code execution" : "Enable code execution"}
>
<CodeToggleIcon className="size-3.5" />

22
studio/frontend/src/features/chat/stores/chat-runtime-store.ts
View file

@ -229,6 +229,24 @@ type ChatRuntimeStore = {
supportsPreserveThinking: boolean;
preserveThinking: boolean;
supportsTools: boolean;
/**
* Whether the active external provider exposes a server-side
* web_search tool (OpenAI's /v1/responses today). Distinct from
* `supportsTools` — that flag governs the local tool runtime (Code,
* python sandbox, our DuckDuckGo web_search). This one only enables
* the chat composer's Search pill for external models. Local models
* keep `supportsTools` only.
*/
supportsBuiltinWebSearch: boolean;
/**
* Whether the active external provider exposes a server-side
* code-execution tool (Anthropic's `code_execution_20250825` on the
* Claude 4.x family). Distinct from `supportsTools` for the same
* reason as `supportsBuiltinWebSearch`: external providers don't
* give us a local tool runtime, but Anthropic dispatches code
* execution server-side. Read by both composers' Code pill gate.
*/
supportsBuiltinCodeExecution: boolean;
toolsEnabled: boolean;
codeToolsEnabled: boolean;
toolStatus: string | null;
@ -320,6 +338,8 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
supportsPreserveThinking: false,
preserveThinking: loadBool(PRESERVE_THINKING_KEY, false),
supportsTools: false,
supportsBuiltinWebSearch: false,
supportsBuiltinCodeExecution: false,
toolsEnabled: false,
codeToolsEnabled: false,
toolStatus: null,
@ -430,6 +450,8 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
reasoningEffortLevels: ["low", "medium", "high"],
supportsPreserveThinking: false,
supportsTools: false,
supportsBuiltinWebSearch: false,
supportsBuiltinCodeExecution: false,
toolsEnabled: false,
codeToolsEnabled: false,
toolStatus: null,

16
studio/frontend/src/features/chat/types.ts
View file

@ -15,6 +15,22 @@ export interface ThreadRecord {
pairId?: string;
archived: boolean;
createdAt: number;
/**
* OpenAI shell tool container id captured from a prior response on
* this thread. When set, the next turn reuses it via
* `environment.type="container_reference"` so the model can read
* files it wrote earlier in the conversation. When null/undefined,
* the next turn auto-creates a fresh container.
*
* OpenAI containers expire after ~20 min of inactivity by default;
* if a stale id is sent, the backend surfaces an
* `_toolEvent.type="container_invalidated"` and the chat-adapter
* clears this field so the following turn falls back to auto-create.
*
* Anthropic's code-execution path doesn't need this — each turn
* gets a fresh container server-side.
*/
openaiCodeExecContainerId?: string | null;
}
export interface MessageRecord {

11
studio/frontend/src/features/chat/types/api.ts
View file

@ -224,6 +224,17 @@ export interface OpenAIChatCompletionsRequest {
external_model?: string;
encrypted_api_key?: string;
provider_base_url?: string | null;
enable_prompt_caching?: boolean | null;
/**
* OpenAI shell-tool container id captured from the prior response in
* this chat thread. When set and the Code pill is on, the backend
* routes the next /v1/responses call with
* `environment.type="container_reference"` so filesystem state
* persists across turns. Unset → backend uses
* `environment.type="container_auto"` and OpenAI creates a fresh
* container. Only meaningful for OpenAI cloud + gpt-5.5 family.
*/
openai_code_exec_container_id?: string | null;
}
export interface OpenAIChatDelta {

2
studio/frontend/src/features/settings/settings-dialog.tsx
View file

@ -45,7 +45,7 @@ const TABS: TabDef[] = [
{ id: "profile", label: "Profile", icon: UserIcon },
{ id: "appearance", label: "Appearance", icon: PaintBrush02Icon },
{ id: "chat", label: "Chat", icon: Message01Icon },
{ id: "connections", label: "Cloud", icon: CloudIcon, badge: "New" },
{ id: "connections", label: "Connections", icon: CloudIcon, badge: "New" },
{ id: "api-keys", label: "API", icon: Globe02Icon, badge: "New" },
{ id: "about", label: "Help", icon: HelpCircleIcon },
];

1
studio/frontend/src/features/training/api/mappers.ts
View file

@ -83,6 +83,7 @@ export function buildTrainingStartPayload(
save_steps: config.saveSteps,
eval_steps: config.evalSteps,
weight_decay: config.weightDecay,
max_grad_norm: 0.0,
random_seed: config.randomSeed,
packing: isEmbedding ? false : config.packing,
optim: config.optimizerType,

1
studio/frontend/src/features/training/types/api.ts
View file

@ -31,6 +31,7 @@ export interface TrainingStartRequest {
save_steps: number;
eval_steps: number;
weight_decay: number;
max_grad_norm: number;
random_seed: number;
packing: boolean;
optim: string;

55
studio/install_llama_prebuilt.py
View file

@ -3433,10 +3433,59 @@ def extract_archive(archive_path: Path, destination: Path) -> None:
) from exc
return target
def _try_repair_missing_slash(
member_name: str, link_name: str, archive_names: set[str]
) -> str | None:
"""Some upstream llama.cpp Mac releases (e.g. b9165, b9169) ship
symlinks whose linkname is missing the directory separator AND
the leading character of the file basename between the
top-level dir and the rest of the path:
llama-b9165/libggml-rpc.0.dylib -> llama-b9165ibggml-rpc.0.11.1.dylib
That cannot be resolved as written. Detect the pattern
(linkname starts with the top-level dir name but no following
slash) and search archive entries under that dir for a real
file whose basename ends with the mangled suffix. Only accept
when the suffix uniquely identifies a real archive entry.
Returns the corrected linkname expressed relative to the
member's parent directory -- callers join it with
`target.parent`, so a full `top/file` path would double the
prefix into `top/top/file`."""
if "/" not in member_name or "/" in link_name:
return None
top, _, _ = member_name.partition("/")
if not link_name.startswith(top) or len(link_name) <= len(top):
return None
bad_suffix = link_name[len(top) :]
if not bad_suffix or bad_suffix.startswith("/"):
return None
prefix = f"{top}/"
candidates = [
name
for name in archive_names
if name.startswith(prefix)
and "/" not in name[len(prefix) :]
and name[len(prefix) :].endswith(bad_suffix)
]
if len(candidates) != 1:
return None
# Strip the top-level dir so the caller's `target.parent / Path(...)`
# composition resolves inside the staging dir, not into a duplicate
# `top/top/...` path.
return candidates[0][len(prefix) :]
def safe_link_target(
base: Path, member_name: str, link_name: str, target: Path
base: Path,
member_name: str,
link_name: str,
target: Path,
archive_names: set[str],
) -> tuple[str, Path]:
normalized = link_name.replace("\\", "/")
repaired = _try_repair_missing_slash(member_name, normalized, archive_names)
if repaired is not None:
normalized = repaired
link_path = Path(normalized)
if link_path.is_absolute():
raise PrebuiltFallback(
@ -3473,8 +3522,10 @@ def extract_archive(archive_path: Path, destination: Path) -> None:
def extract_tar_safely(source: Path, base: Path) -> None:
pending_links: list[tuple[tarfile.TarInfo, Path]] = []
archive_names: set[str] = set()
with tarfile.open(source, "r:gz") as archive:
for member in archive.getmembers():
archive_names.add(member.name)
target = safe_extract_path(base, member.name)
if member.isdir():
target.mkdir(parents = True, exist_ok = True)
@ -3501,7 +3552,7 @@ def extract_archive(archive_path: Path, destination: Path) -> None:
progressed = False
for member, target in unresolved:
normalized_link, resolved_target = safe_link_target(
base, member.name, member.linkname, target
base, member.name, member.linkname, target, archive_names
)
if not resolved_target.exists() and not resolved_target.is_symlink():
next_round.append((member, target))

14
studio/install_python_stack.py
View file

@ -28,6 +28,7 @@ if str(_BACKEND_DIR) not in sys.path:
from backend.utils.wheel_utils import (
flash_attn_package_version,
flash_attn_wheel_url,
has_blackwell_gpu,
install_wheel,
probe_torch_wheel_env,
url_exists,
@ -628,10 +629,19 @@ def _flash_attn_install_disabled() -> bool:
def _ensure_flash_attn() -> None:
if NO_TORCH or IS_WINDOWS or IS_MACOS:
return
if _flash_attn_install_disabled():
return
if NO_TORCH:
return
if has_blackwell_gpu():
_step(
"warning",
"Skipping flash-attn: Blackwell GPU detected (sm_100+); no compatible prebuilt wheel",
_cyan,
)
return
if IS_WINDOWS or IS_MACOS:
return
if (
subprocess.run(
[sys.executable, "-c", "import flash_attn"],

195
tests/python/test_flash_attn_install_python_stack.py
View file

@ -10,8 +10,133 @@ from unittest import mock
STUDIO_DIR = Path(__file__).resolve().parents[2] / "studio"
sys.path.insert(0, str(STUDIO_DIR))
sys.path.insert(0, str(STUDIO_DIR / "backend"))
import install_python_stack as ips
from backend.utils import wheel_utils
def _smi_result(stdout: str, returncode: int = 0) -> subprocess.CompletedProcess:
return subprocess.CompletedProcess(["nvidia-smi"], returncode, stdout, "")
class TestHasBlackwellGpu:
def setup_method(self):
wheel_utils.has_blackwell_gpu.cache_clear()
def teardown_method(self):
wheel_utils.has_blackwell_gpu.cache_clear()
def test_returns_false_when_nvidia_smi_missing(self):
with mock.patch.object(wheel_utils.shutil, "which", return_value = None):
assert wheel_utils.has_blackwell_gpu() is False
def test_returns_true_for_sm_100(self):
with (
mock.patch.object(
wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
),
mock.patch.object(
wheel_utils.subprocess, "run", return_value = _smi_result("10.0\n")
),
):
assert wheel_utils.has_blackwell_gpu() is True
def test_returns_true_for_sm_120(self):
with (
mock.patch.object(
wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
),
mock.patch.object(
wheel_utils.subprocess, "run", return_value = _smi_result("12.0\n")
),
):
assert wheel_utils.has_blackwell_gpu() is True
def test_returns_true_for_sm_121(self):
with (
mock.patch.object(
wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
),
mock.patch.object(
wheel_utils.subprocess, "run", return_value = _smi_result("12.1\n")
),
):
assert wheel_utils.has_blackwell_gpu() is True
def test_returns_false_for_sm_90(self):
with (
mock.patch.object(
wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
),
mock.patch.object(
wheel_utils.subprocess, "run", return_value = _smi_result("9.0\n")
),
):
assert wheel_utils.has_blackwell_gpu() is False
def test_returns_false_for_sm_89(self):
with (
mock.patch.object(
wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
),
mock.patch.object(
wheel_utils.subprocess, "run", return_value = _smi_result("8.9\n")
),
):
assert wheel_utils.has_blackwell_gpu() is False
def test_mixed_gpus_with_one_blackwell_returns_true(self):
with (
mock.patch.object(
wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
),
mock.patch.object(
wheel_utils.subprocess,
"run",
return_value = _smi_result("8.0\n10.0\n"),
),
):
assert wheel_utils.has_blackwell_gpu() is True
def test_returns_false_when_nvidia_smi_fails(self):
with (
mock.patch.object(
wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
),
mock.patch.object(
wheel_utils.subprocess,
"run",
return_value = _smi_result("", returncode = 1),
),
):
assert wheel_utils.has_blackwell_gpu() is False
def test_returns_false_on_subprocess_timeout(self):
with (
mock.patch.object(
wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
),
mock.patch.object(
wheel_utils.subprocess,
"run",
side_effect = subprocess.TimeoutExpired(cmd = "nvidia-smi", timeout = 10),
),
):
assert wheel_utils.has_blackwell_gpu() is False
def test_returns_false_on_malformed_output(self):
with (
mock.patch.object(
wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
),
mock.patch.object(
wheel_utils.subprocess,
"run",
return_value = _smi_result("not-a-number\n\n"),
),
):
assert wheel_utils.has_blackwell_gpu() is False
class TestFlashAttnWheelSelection:
@ -234,6 +359,76 @@ class TestEnsureFlashAttn:
mock_probe.assert_not_called()
mock_install_wheel.assert_not_called()
def test_blackwell_gpu_skips_install_with_warning(self):
step_messages: list[tuple[str, str]] = []
def fake_step(label: str, value: str, color_fn = None):
step_messages.append((label, value))
with (
mock.patch.object(ips, "NO_TORCH", False),
mock.patch.object(ips, "IS_WINDOWS", False),
mock.patch.object(ips, "IS_MACOS", False),
mock.patch.object(ips, "has_blackwell_gpu", return_value = True),
mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe,
mock.patch.object(ips, "install_wheel") as mock_install_wheel,
mock.patch.object(ips, "_step", side_effect = fake_step),
mock.patch("subprocess.run", return_value = self._import_check()),
):
ips._ensure_flash_attn()
mock_probe.assert_not_called()
mock_install_wheel.assert_not_called()
assert any(
label == "warning" and "Blackwell" in msg for label, msg in step_messages
)
def test_blackwell_gpu_on_windows_emits_blackwell_warning(self):
step_messages: list[tuple[str, str]] = []
def fake_step(label: str, value: str, color_fn = None):
step_messages.append((label, value))
with (
mock.patch.object(ips, "NO_TORCH", False),
mock.patch.object(ips, "IS_WINDOWS", True),
mock.patch.object(ips, "IS_MACOS", False),
mock.patch.object(ips, "has_blackwell_gpu", return_value = True),
mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe,
mock.patch.object(ips, "install_wheel") as mock_install_wheel,
mock.patch.object(ips, "_step", side_effect = fake_step),
mock.patch("subprocess.run", return_value = self._import_check()),
):
ips._ensure_flash_attn()
mock_probe.assert_not_called()
mock_install_wheel.assert_not_called()
assert any(
label == "warning" and "Blackwell" in msg for label, msg in step_messages
)
def test_non_blackwell_windows_does_not_emit_blackwell_warning(self):
step_messages: list[tuple[str, str]] = []
def fake_step(label: str, value: str, color_fn = None):
step_messages.append((label, value))
with (
mock.patch.object(ips, "NO_TORCH", False),
mock.patch.object(ips, "IS_WINDOWS", True),
mock.patch.object(ips, "IS_MACOS", False),
mock.patch.object(ips, "has_blackwell_gpu", return_value = False),
mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe,
mock.patch.object(ips, "install_wheel") as mock_install_wheel,
mock.patch.object(ips, "_step", side_effect = fake_step),
mock.patch("subprocess.run", return_value = self._import_check()),
):
ips._ensure_flash_attn()
mock_probe.assert_not_called()
mock_install_wheel.assert_not_called()
assert not any("Blackwell" in msg for _, msg in step_messages)
class TestInstallPythonStackFlashAttnIntegration:
def _run_install(self, *, no_torch: bool, is_macos: bool, is_windows: bool) -> int:

13
tests/studio/install/test_llama_pr_force_and_source.py
View file

@ -35,9 +35,11 @@ requires_pwsh = pytest.mark.skipif(not PWSH_AVAILABLE, reason = "pwsh not availa
# Helpers
# ---------------------------------------------------------------------------
def run_bash(
script: str, *, timeout: int = 10, env: dict | None = None
script: str, *, timeout: int = 60, env: dict | None = None
) -> subprocess.CompletedProcess:
"""Run a bash script fragment and return the CompletedProcess."""
"""Run a bash script fragment and return the CompletedProcess.
60s default tolerates slow shell startup on heavily-loaded CI
runners; the scripts themselves run in well under a second."""
run_env = os.environ.copy()
if env:
run_env.update(env)
@ -51,9 +53,12 @@ def run_bash(
def run_pwsh(
script: str, *, timeout: int = 10, env: dict | None = None
script: str, *, timeout: int = 60, env: dict | None = None
) -> subprocess.CompletedProcess:
"""Run a PowerShell script fragment and return the CompletedProcess."""
"""Run a PowerShell script fragment and return the CompletedProcess.
60s default tolerates slow pwsh startup on heavily-loaded CI
runners; the scripts themselves run in well under a second.
A 10s budget previously surfaced as a flaky TimeoutExpired."""
run_env = os.environ.copy()
run_env["NO_COLOR"] = "1"
if env:

457
tests/studio/playwright_chat_ime_i18n.py Normal file
View file

@ -0,0 +1,457 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Studio chat composer IME + multilingual regression smoke.
Covers two surfaces:
A. Stuck IME composition (issue #5318 / PR #5327): duplicate
compositionstart with no compositionend left isComposing=true,
dropping all subsequent keystrokes including ASCII.
B. Multilingual paste round-trip across 31 scripts -- guards the
controlled-textarea / React state plumbing against Unicode mangling.
Model-free; the bug surface is the composer, not inference.
Env contract matches playwright_chat_ui.py:
BASE_URL, STUDIO_NEW_PW, PW_ART_DIR, STUDIO_UI_STRICT.
"""
import os
import sys
from pathlib import Path
from playwright.sync_api import expect, sync_playwright
sys.path.insert(0, str(Path(__file__).resolve().parent))
from _playwright_robust import ( # noqa: E402
chromium_launch_args,
click_and_wait_for_response,
install_view_transition_killer,
install_wall_clock_watchdog,
is_benign_console_error,
is_benign_page_error,
recover_or_replace_page,
wait_for_health,
)
BASE = os.environ["BASE_URL"]
NEW = os.environ["STUDIO_NEW_PW"]
ART_DIR = os.environ.get("PW_ART_DIR", "logs/playwright_ime")
ART = Path(ART_DIR)
ART.mkdir(parents = True, exist_ok = True)
STRICT = os.environ.get("STUDIO_UI_STRICT", "0") == "1"
# Wall-clock cap. Realistic run is 30-60s; 5 min leaves cold-launch headroom.
WALL_TIMEOUT_S = float(os.environ.get("STUDIO_IME_WALL_TIMEOUT_S", "300"))
# One short greeting + arithmetic per script (ordered by speaker count) --
# each entry catches a distinct class of Unicode regression.
I18N_SAMPLES = [
("en", "English", "Hello, 1+1=2"),
("zh-CN", "Chinese (Simplified)", "你好，1+1=2"),
("es", "Spanish", "Hola, 1+1=2"),
("hi", "Hindi (Devanagari)", "नमस्ते, 1+1=2"),
("ar", "Arabic (RTL)", "مرحبا، ١+١=٢"),
("bn", "Bengali", "নমস্কার, ১+১=২"),
("pt", "Portuguese", "Olá, 1+1=2"),
("ru", "Russian (Cyrillic)", "Привет, 1+1=2"),
("ja", "Japanese", "こんにちは、1+1=2"),
("pa", "Punjabi (Gurmukhi)", "ਸਤ ਸ੍ਰੀ ਅਕਾਲ, 1+1=2"),
("de", "German", "Hallo, 1+1=2"),
("jv", "Javanese", "Halo, 1+1=2"),
("ko", "Korean (Hangul)", "안녕하세요, 1+1=2"),
("fr", "French", "Bonjour, 1+1=2"),
("tr", "Turkish", "Merhaba, 1+1=2"),
("vi", "Vietnamese (diacritics)", "Xin chào, 1+1=2"),
("ur", "Urdu (Arabic-Naskh)", "ہیلو، 1+1=2"),
("ta", "Tamil", "வணக்கம், 1+1=2"),
("te", "Telugu", "నమస్తే, 1+1=2"),
("mr", "Marathi (Devanagari)", "नमस्कार, 1+1=2"),
("it", "Italian", "Ciao, 1+1=2"),
("th", "Thai", "สวัสดี, ๑+๑=๒"),
("pl", "Polish", "Cześć, 1+1=2"),
("uk", "Ukrainian (Cyrillic)", "Привіт, 1+1=2"),
("fa", "Persian (RTL)", "سلام، ۱+۱=۲"),
("nl", "Dutch", "Hallo, 1+1=2"),
("he", "Hebrew (RTL)", "שלום, 1+1=2"),
("el", "Greek", "Γειά, 1+1=2"),
("id", "Indonesian", "Halo, 1+1=2"),
("sw", "Swahili", "Habari, 1+1=2"),
("emoji", "Emoji + ZWJ + flag", "👋 🇺🇳 👨‍👩‍👧‍👦 1+1=2"),
]
_n = [0]
def step(s):
print(f"[ime] STEP {s}", flush = True)
def info(s):
print(f"[ime] {s}", flush = True)
def fail(m):
raise AssertionError(f"[ime] FAIL: {m}")
def soft_fail(m):
"""Hard fail in STRICT mode, info-warn otherwise. Mirrors playwright_chat_ui.py."""
if STRICT:
fail(m)
info(f"WARN (strict-off): {m}")
with sync_playwright() as p:
_watchdog = install_wall_clock_watchdog(
WALL_TIMEOUT_S,
label = "ime",
info = info,
)
wait_for_health(BASE, timeout = 30.0, info = info)
browser = p.chromium.launch(
headless = True,
args = chromium_launch_args(),
)
ctx = browser.new_context(
viewport = {"width": 1280, "height": 900},
reduced_motion = "reduce",
)
install_view_transition_killer(ctx)
page = ctx.new_page()
page.set_default_timeout(60_000)
page_errors: list[str] = []
console_errors: list[str] = []
def _on_console(m):
if m.type != "error":
return
try:
console_errors.append(m.text)
except Exception:
return
def _attach_listeners(target):
target.on("pageerror", lambda e: page_errors.append(str(e)))
target.on("console", _on_console)
_attach_listeners(page)
def shoot(name):
_n[0] += 1
try:
page.screenshot(
path = str(ART / f"{_n[0]:02d}-{name}.png"),
full_page = True,
timeout = 90_000,
animations = "disabled",
)
except Exception as _shoot_err:
info(f"WARN: screenshot {name} failed: {_shoot_err}")
# 1. Bootstrap auth via /change-password (mirrors playwright_chat_ui.py
# retry-on-rerender to absorb React form-detach races).
step("change-password through UI (Setup your account)")
form_err: Exception | None = None
for _form_attempt in range(3):
try:
page.goto(
f"{BASE}/change-password",
wait_until = "domcontentloaded",
timeout = 60_000,
)
try:
page.wait_for_load_state("networkidle", timeout = 30_000)
except Exception:
pass
pw_field = page.locator("#new-password")
pw_field.wait_for(state = "visible", timeout = 60_000)
pw_field.fill(NEW, timeout = 60_000)
page.fill("#confirm-password", NEW, timeout = 60_000)
shoot("01-change-password-filled")
status, _ = click_and_wait_for_response(
page,
url_substr = "/api/auth/change-password",
method = "POST",
do_click = lambda: page.locator('button[type="submit"]').click(),
timeout_ms = 30_000,
info = lambda m: print(f"[ime] {m}", flush = True),
)
if status is not None and status >= 400:
raise AssertionError(f"change-password POST returned {status}")
form_err = None
break
except Exception as e:
form_err = e
info(
f"change-password attempt {_form_attempt + 1} failed: "
f"{type(e).__name__}: {str(e)[:200]}"
)
if _form_attempt < 2:
page = recover_or_replace_page(
page,
ctx,
default_timeout_ms = 60_000,
info = lambda m: print(f"[ime] recovery: {m}", flush = True),
)
_attach_listeners(page)
if form_err is not None:
raise form_err
# 2. Wait for composer mount. No GGUF: the bug surface is React state, not inference.
step("wait for composer to mount")
try:
page.wait_for_load_state("networkidle", timeout = 30_000)
except Exception:
pass
composer = page.locator('textarea[aria-label="Message input"]')
_mount_err: Exception | None = None
for _mount_attempt in range(2):
try:
composer.wait_for(state = "visible", timeout = 60_000)
_mount_err = None
break
except Exception as e:
_mount_err = e
info(
f"composer.wait_for attempt {_mount_attempt + 1} failed: "
f"{type(e).__name__}: {str(e)[:200]}"
)
try:
shoot(f"02-composer-wait-attempt-{_mount_attempt + 1}-fail")
except Exception:
pass
if _mount_attempt == 0:
page = recover_or_replace_page(
page,
ctx,
default_timeout_ms = 60_000,
info = lambda m: print(f"[ime] recovery: {m}", flush = True),
)
_attach_listeners(page)
composer = page.locator('textarea[aria-label="Message input"]')
if _mount_err is not None:
raise _mount_err
composer.click()
shoot("02-composer-focused")
# Main composer must carry dir="auto" so RTL flows right-to-left.
dir_attr = composer.evaluate("(el) => el.getAttribute('dir')")
if dir_attr != "auto":
soft_fail(
f'composer is missing dir="auto" (got {dir_attr!r}); RTL '
"languages will render LTR."
)
else:
info('composer dir="auto" present')
# Source-level guard for the edit and compare composers (neither
# is mounted here): grep the JSX for dir="auto" inside each block.
_repo_root = Path(__file__).resolve().parents[2]
_thread_src = (
_repo_root / "studio/frontend/src/components/assistant-ui/thread.tsx"
).read_text()
_shared_src = (
_repo_root / "studio/frontend/src/features/chat/shared-composer.tsx"
).read_text()
_edit_idx = _thread_src.find("aui-edit-composer-input")
if _edit_idx == -1 or 'dir="auto"' not in _thread_src[_edit_idx : _edit_idx + 600]:
soft_fail('edit composer source is missing dir="auto"')
else:
info('edit composer dir="auto" present (source)')
_compare_idx = _shared_src.find("Send to both models")
if (
_compare_idx == -1
or 'dir="auto"'
not in _shared_src[max(_compare_idx - 400, 0) : _compare_idx + 400]
):
soft_fail('compare composer source is missing dir="auto"')
else:
info('compare composer dir="auto" present (source)')
def read_value() -> str:
return composer.evaluate("(el) => el.value")
def set_value_via_setter(s: str) -> str:
"""Write via React's monkey-patched setter + paste input event,
then await two rAFs so the controlled value is committed before
readback (plain `.value=s` would be overwritten on next render)."""
return composer.evaluate(
"""async (el, v) => {
const setter = Object.getOwnPropertyDescriptor(
window.HTMLTextAreaElement.prototype, 'value'
).set;
setter.call(el, v);
el.dispatchEvent(new InputEvent('input', {
bubbles: true,
inputType: 'insertFromPaste',
data: v,
}));
await new Promise((r) => requestAnimationFrame(r));
await new Promise((r) => requestAnimationFrame(r));
return el.value;
}""",
s,
)
def clear() -> None:
set_value_via_setter("")
# 3. Baseline: ASCII keyboard typing works. Bail fast if not.
step("baseline ASCII keyboard typing")
clear()
composer.click()
for ch in "hello world":
page.keyboard.type(ch)
got = read_value()
if got != "hello world":
fail(f"ASCII typing readback {got!r} != 'hello world'")
info("baseline ASCII OK")
shoot("03-baseline-ascii")
clear()
# 4. Multilingual paste round-trip; byte-for-byte readback required.
step(f"multilingual paste round-trip ({len(I18N_SAMPLES)} samples)")
paste_failures: list[tuple[str, str, str, str]] = []
for code, label, text in I18N_SAMPLES:
got = set_value_via_setter(text)
if got != text:
paste_failures.append((code, label, text, got))
info(f" {code:>6} ({label}): FAIL -- got {got!r}")
else:
info(f" {code:>6} ({label}): OK")
clear()
if paste_failures:
shoot("04-paste-failures")
lines = [
f" {code} ({label}): want={want!r} got={got!r}"
for code, label, want, got in paste_failures
]
fail(
f"{len(paste_failures)}/{len(I18N_SAMPLES)} languages failed paste round-trip:\n"
+ "\n".join(lines)
)
info(f"all {len(I18N_SAMPLES)} multilingual paste samples OK")
shoot("04-paste-all-ok")
# 5. Healthy IME composition (compositionstart/update/end + insert events).
step("normal IME composition (compose 你好)")
clear()
composer.click()
composer.evaluate(
"""(el) => {
el.focus();
el.dispatchEvent(new CompositionEvent('compositionstart', {bubbles:true, data:''}));
el.dispatchEvent(new CompositionEvent('compositionupdate', {bubbles:true, data:'你'}));
el.dispatchEvent(new CompositionEvent('compositionupdate', {bubbles:true, data:'你好'}));
const setter = Object.getOwnPropertyDescriptor(
window.HTMLTextAreaElement.prototype, 'value'
).set;
setter.call(el, el.value + '你好');
el.dispatchEvent(new InputEvent('input', {
bubbles:true, inputType:'insertCompositionText',
data:'你好', isComposing:true,
}));
el.dispatchEvent(new CompositionEvent('compositionend', {bubbles:true, data:'你好'}));
el.dispatchEvent(new InputEvent('input', {
bubbles:true, inputType:'insertFromComposition', data:'你好',
}));
}"""
)
got = read_value()
if "你好" not in got:
shoot("05-normal-composition-FAIL")
fail(f"normal composition readback {got!r} missing '你好'")
info(f"normal composition OK: ta.value={got!r}")
shoot("05-normal-composition")
clear()
# 6. Stuck IME repro for issue #5318: duplicate compositionstart with
# no compositionend wedged isComposing=true and dropped ASCII keys.
# PR #5327 cleared the stale state on non-composing input.
step("BUG REPRO: stuck IME composition recovery (issue #5318)")
clear()
composer.click()
composer.evaluate(
"""(el) => {
el.focus();
el.dispatchEvent(new CompositionEvent('compositionstart', {bubbles:true, data:''}));
// Duplicate compositionstart with NO matching compositionend.
// This is exactly the event sequence observed from the IMEs
// in issue #5318 (kei-yamazaki / langxiaopiao030 / PapyrusNotes).
el.dispatchEvent(new CompositionEvent('compositionstart', {bubbles:true, data:''}));
}"""
)
# Drive the real keyboard path; on the broken build React drops
# 'abcd' and reconciles el.value back to ''. wait_for_function
# crosses the microtask boundary so we see committed React state.
page.keyboard.type("abcd")
try:
page.wait_for_function(
"""(el) => el.value === 'abcd'""",
composer.element_handle(),
timeout = 5_000,
)
except Exception:
pass
after_key = read_value()
info(f"after_key='abcd' readback={after_key!r}")
shoot("06-stuck-composition-recovery")
if after_key != "abcd":
fail(
"stuck-composition repro: keyboard 'abcd' was not preserved after "
f"duplicate compositionstart; readback {after_key!r}. React state "
"likely still stuck in isComposing=true (issue #5318 / before "
"PR #5327)."
)
# Cross-check React's view of isComposing via the Send button:
# ComposerAction stays disabled while isComposing is true (PR #5327).
send_btn = page.locator('button[aria-label="Send message"]')
if send_btn.count() == 0:
soft_fail("Send button not found after stuck-composition recovery")
else:
try:
expect(send_btn).not_to_be_disabled(timeout = 5_000)
info("Send button correctly enabled after stuck-composition recovery")
except Exception:
soft_fail(
"Send button still disabled after stuck-composition recovery -- "
"React isComposing state likely never cleared"
)
info("stuck-composition recovery PASS")
clear()
# 7. Final state. The change-password redirect emits benign 401 noise,
# so we filter via is_benign_* and only fail on real errors.
shoot("07-final")
real_page_errors = [e for e in page_errors if not is_benign_page_error(e)]
real_console_errors = [e for e in console_errors if not is_benign_console_error(e)]
info(
f"page_errors={len(page_errors)} ({len(real_page_errors)} non-benign); "
f"console_errors={len(console_errors)} "
f"({len(real_console_errors)} non-benign)"
)
if page_errors:
info(f"first page error: {page_errors[0][:200]!r}")
if console_errors:
info(f"first console error: {console_errors[0][:200]!r}")
if real_page_errors:
fail(
f"{len(real_page_errors)} non-benign pageerror events; "
f"first={real_page_errors[0][:200]!r}"
)
if real_console_errors:
fail(
f"{len(real_console_errors)} non-benign console.error events; "
f"first={real_console_errors[0][:200]!r}"
)
info(
f"DONE: ascii=OK paste={len(I18N_SAMPLES)}/{len(I18N_SAMPLES)} "
f"normal_composition=OK stuck_recovery=OK"
)
_watchdog.cancel()
browser.close()

6
tests/studio/run_real_mlx_smoke.py
View file

@ -211,8 +211,8 @@ def cmd_train(args) -> int:
workdir.mkdir(parents = True, exist_ok = True)
import mlx.core as mx
from unsloth_zoo.mlx_loader import FastMLXModel
from unsloth_zoo.mlx_trainer import MLXTrainer, MLXTrainingConfig
from unsloth_zoo.mlx.loader import FastMLXModel
from unsloth_zoo.mlx.trainer import MLXTrainer, MLXTrainingConfig
hf_token = os.environ.get("HF_TOKEN") or None
@ -440,7 +440,7 @@ def cmd_reload(args) -> int:
return _reload_gguf(save_dir, metrics)
import mlx.core as mx
from unsloth_zoo.mlx_loader import FastMLXModel
from unsloth_zoo.mlx.loader import FastMLXModel
from mlx_lm import generate
hf_token = os.environ.get("HF_TOKEN") or None

73
tests/studio/test_composer_rtl_bidi_attribute.py Normal file
View file

@ -0,0 +1,73 @@
"""Lock down the RTL bidi auto-detection contract on the chat composers.
The browser's Unicode bidi algorithm only flows Arabic / Hebrew / Persian /
Urdu right-to-left when the textarea carries `dir="auto"`. The three
composer surfaces (main chat, inline edit, compare mode) each need the
attribute, and the IME / i18n Playwright smoke must keep its env contract
minimal (no dead `STUDIO_OLD_PW`).
"""
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"
SHARED_TSX = REPO / "studio/frontend/src/features/chat/shared-composer.tsx"
WORKFLOW_YML = REPO / ".github/workflows/studio-ui-smoke.yml"
IME_PY = REPO / "tests/studio/playwright_chat_ime_i18n.py"
def _block_around(src: str, anchor: str, radius: int = 600) -> str:
idx = src.find(anchor)
assert idx != -1, f"anchor {anchor!r} not found"
return src[max(idx - radius, 0) : idx + radius]
def test_main_composer_has_dir_auto():
block = _block_around(THREAD_TSX.read_text(), 'aria-label="Message input"')
assert 'dir="auto"' in block, 'main composer is missing dir="auto"'
def test_edit_composer_has_dir_auto():
block = _block_around(THREAD_TSX.read_text(), "aui-edit-composer-input")
assert 'dir="auto"' in block, 'edit composer is missing dir="auto"'
def test_compare_composer_has_dir_auto():
block = _block_around(SHARED_TSX.read_text(), "Send to both models")
assert 'dir="auto"' in block, 'compare composer is missing dir="auto"'
def test_ime_workflow_step_does_not_set_studio_old_pw():
yml = WORKFLOW_YML.read_text()
drive_idx = yml.find("Drive IME + multilingual paste regression")
assert drive_idx != -1, "IME drive step not found in workflow"
next_step_idx = yml.find("- name:", drive_idx + 1)
drive_block = yml[drive_idx : next_step_idx if next_step_idx != -1 else None]
assert (
"STUDIO_OLD_PW" not in drive_block
), "IME drive step still passes dead STUDIO_OLD_PW env var"
assert "STUDIO_NEW_PW" in drive_block, "IME drive step missing STUDIO_NEW_PW"
def test_ime_pass_password_step_does_not_export_old_pw():
yml = WORKFLOW_YML.read_text()
pass_idx = yml.find("Pass bootstrap pw for IME / i18n test")
assert pass_idx != -1, "IME password setup step not found"
next_step_idx = yml.find("- name:", pass_idx + 1)
pass_block = yml[pass_idx : next_step_idx if next_step_idx != -1 else None]
assert (
"STUDIO_IME_OLD_PW" not in pass_block
), "IME password setup still exports dead STUDIO_IME_OLD_PW"
assert "STUDIO_IME_NEW_PW" in pass_block
def test_ime_playwright_script_does_not_read_studio_old_pw():
src = IME_PY.read_text()
code_only = re.sub(r'""".*?"""', "", src, flags = re.DOTALL)
assert (
"STUDIO_OLD_PW" not in code_only
), "IME Playwright script still references dead STUDIO_OLD_PW env var"
assert 'os.environ["STUDIO_NEW_PW"]' in code_only

1628
tests/studio/test_frontend_dep_removal.py Normal file
View file

File diff suppressed because it is too large Load diff

66
tests/studio/test_is_mlx_dispatch_gate.py
View file

@ -7,8 +7,9 @@ Two gates drive every dispatch decision in Studio's MLX path:
1. ``unsloth._IS_MLX`` at the top of ``unsloth/__init__.py`` -- evaluated
once at import time and read by Studio worker code to choose between
the GPU and MLX trainer / inference / export paths. Defined as
``Darwin AND arm64 AND find_spec("mlx") is not None``.
the GPU and MLX trainer / inference / export paths. It delegates to
the shared zoo MLX runtime gate, with a local import barrier while the
paired unsloth-zoo runtime rollout is in flight.
2. ``utils.hardware.detect_hardware()`` -- runtime probe in the Studio
backend. Priority order: CUDA -> XPU -> MLX -> CPU. The MLX branch is
@ -18,8 +19,8 @@ Two gates drive every dispatch decision in Studio's MLX path:
These gates are the canaries for "MLX support accidentally hijacks
CUDA/AMD/Intel users". The tests here:
* verify the source-level structure of the ``_IS_MLX`` expression so an
accidental rewrite (e.g. dropping the ``arm64`` check) is caught,
* verify the source-level structure of the ``_IS_MLX`` helper so an
accidental rewrite importing zoo before the local MLX precheck is caught,
* exercise the runtime gate logic under a spoofed Darwin+arm64 platform
with a fake ``mlx`` module in ``sys.modules`` to confirm both gates
flip True together,
@ -64,20 +65,36 @@ def test_is_mlx_gate_uses_three_required_predicates():
target = node.value
break
assert target is not None, "_IS_MLX assignment not found in unsloth/__init__.py"
assert isinstance(target, ast.BoolOp) and isinstance(
target.op, ast.And
), "_IS_MLX must be a BoolOp(And) of platform + mlx checks"
assert isinstance(target, ast.Call), "_IS_MLX must call the shared MLX helper"
expr_src = ast.unparse(target)
assert (
"platform.system()" in expr_src and "Darwin" in expr_src
), "_IS_MLX must check platform.system() == 'Darwin'"
expr_src == "_is_mlx_available()"
), "_IS_MLX must delegate to the shared MLX runtime gate"
helper = None
for node in ast.walk(tree):
if isinstance(node, ast.FunctionDef) and node.name == "_is_mlx_available":
helper = node
break
assert helper is not None, "_is_mlx_available helper not found"
helper_src = ast.unparse(helper)
assert (
"platform.machine()" in expr_src and "arm64" in expr_src
), "_IS_MLX must check platform.machine() == 'arm64'"
"platform.system()" in helper_src
and "'Darwin'" in helper_src
and "platform.machine()" in helper_src
and "'arm64'" in helper_src
and "find_spec" in helper_src
and "'mlx'" in helper_src
and "from unsloth_zoo.mlx import is_mlx_available" in helper_src
), "_IS_MLX helper must precheck local MLX predicates before importing zoo"
assert (
"find_spec" in expr_src and "'mlx'" in expr_src
), "_IS_MLX must check importlib.util.find_spec('mlx')"
"from unsloth_zoo.mlx import is_mlx_available" in helper_src
and "return is_mlx_available()" in helper_src
), "_IS_MLX helper must delegate final detection to the shared zoo MLX runtime gate"
assert helper_src.index("UNSLOTH_FORCE_GPU_PATH") < helper_src.index(
"from unsloth_zoo.mlx import is_mlx_available"
), "_IS_MLX helper must run the local MLX precheck before importing zoo"
# ---------------------------------------------------------------------------
@ -87,13 +104,14 @@ def test_is_mlx_gate_uses_three_required_predicates():
# ---------------------------------------------------------------------------
def _evaluate_is_mlx_gate(platform_module, importlib_util):
"""Re-evaluate the _IS_MLX expression using injected dependencies.
def _evaluate_is_mlx_precheck(platform_module, importlib_util, os_module):
"""Re-evaluate the local _is_mlx_available precheck using injected dependencies.
Mirrors the assignment in unsloth/__init__.py exactly.
Mirrors only the cheap import barrier before unsloth imports unsloth_zoo.
"""
return (
platform_module.system() == "Darwin"
os_module.environ.get("UNSLOTH_FORCE_GPU_PATH", "0") != "1"
and platform_module.system() == "Darwin"
and platform_module.machine() == "arm64"
and importlib_util.find_spec("mlx") is not None
)
@ -112,7 +130,9 @@ def test_is_mlx_gate_true_on_apple_silicon_with_mlx_present(monkeypatch):
monkeypatch.setattr(platform, "system", lambda: "Darwin")
monkeypatch.setattr(platform, "machine", lambda: "arm64")
assert _evaluate_is_mlx_gate(platform, importlib.util) is True
import os
assert _evaluate_is_mlx_precheck(platform, importlib.util, os) is True
def test_is_mlx_gate_false_when_mlx_missing(monkeypatch):
@ -133,7 +153,9 @@ def test_is_mlx_gate_false_when_mlx_missing(monkeypatch):
monkeypatch.setattr(importlib.util, "find_spec", _no_mlx)
assert _evaluate_is_mlx_gate(platform, importlib.util) is False
import os
assert _evaluate_is_mlx_precheck(platform, importlib.util, os) is False
def test_is_mlx_gate_false_on_non_apple_silicon():
@ -147,7 +169,9 @@ def test_is_mlx_gate_false_on_non_apple_silicon():
pytest.skip("Test host is Apple Silicon; CUDA-side canary doesn't apply.")
assert _evaluate_is_mlx_gate(platform, importlib.util) is False
import os
assert _evaluate_is_mlx_precheck(platform, importlib.util, os) is False
# ---------------------------------------------------------------------------

8
tests/studio/test_mlx_training_worker_behaviors.py
View file

@ -46,8 +46,8 @@ def test_wandb_init_strips_secret_keys():
def test_local_dataset_loader_uses_load_dataset_path():
src = WORKER.read_text()
assert "_resolve_local_files" in src
assert "_loader_for_files" in src
assert "_resolve_mlx_local_dataset_files" in src
assert "_mlx_local_dataset_loader_for_files" in src
assert "data_files = all_files" in src or "data_files=all_files" in src
@ -84,7 +84,7 @@ def test_poll_stop_returns_on_broken_pipe():
def test_unsloth_zoo_mlx_imports_have_friendly_error():
src = WORKER.read_text()
assert "from unsloth_zoo.mlx_loader import FastMLXModel" in src
assert "from unsloth_zoo.mlx_trainer import" in src
assert "from unsloth_zoo.mlx.loader import FastMLXModel" in src
assert "from unsloth_zoo.mlx.trainer import" in src
assert "raise ImportError" in src
assert "install.sh" in src

216
tests/test_public_api_surface.py Normal file
View file

@ -0,0 +1,216 @@
# Unsloth - 2x faster, 60% less VRAM LLM training and finetuning
# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
"""Public-API surface drift detectors for unsloth itself.
Companion to tests/test_import_fixes_drift.py: that file catches drift
in THIRD-PARTY libraries (transformers / trl / triton / peft / etc.)
that unsloth's import_fixes patches around. This file catches drift in
unsloth's OWN public-surface API -- the top-10 symbols and classmethods
that the unslothai/notebooks tree (and therefore every user on Colab)
calls. If a refactor on this repo renames FastLanguageModel.from_pretrained
or drops one of the documented kwargs, the test fires DRIFT DETECTED
here BEFORE the breakage reaches users.
Call-site counts measured against unslothai/notebooks @ main:
FastLanguageModel.from_pretrained 506
FastLanguageModel.for_inference 370
FastLanguageModel.get_peft_model 304
FastVisionModel.for_inference 183
FastVisionModel.from_pretrained 176
FastVisionModel.get_peft_model 99
FastVisionModel.for_training 60
FastModel.from_pretrained 103
FastModel.get_peft_model 67
Mirrors the unsloth-zoo / unsloth drift-detector skeleton:
``pytest.importorskip("unsloth")`` to gate, assert the healthy upstream
shape, ``pytest.fail("DRIFT DETECTED: ...")`` (never ``pytest.skip``) on
regression so the matrix cell goes red.
"""
from __future__ import annotations
import inspect
import pytest
def _signature_param_names(callable_obj) -> set[str]:
try:
sig = inspect.signature(callable_obj)
except (TypeError, ValueError):
return set()
return set(sig.parameters)
def _accepts(callable_obj, kwargs: set[str]) -> tuple[bool, set[str]]:
"""True if every name in ``kwargs`` is either a named parameter on
``callable_obj`` OR the callable's signature has a ``**kwargs``
catch-all. Returns (ok, missing_set)."""
try:
sig = inspect.signature(callable_obj)
except (TypeError, ValueError):
return True, set()
params = sig.parameters
has_var_kw = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values())
if has_var_kw:
return True, set()
missing = kwargs - set(params)
return (not missing), missing
# ===========================================================================
# FastLanguageModel: the headline class. 506 from_pretrained + 370
# for_inference + 304 get_peft_model call sites across the notebooks.
# ===========================================================================
def test_fast_language_model_class_present():
unsloth = pytest.importorskip("unsloth")
if not hasattr(unsloth, "FastLanguageModel"):
pytest.fail(
"DRIFT DETECTED: unsloth.FastLanguageModel is missing; every "
"LoRA notebook fails at the first import cell."
)
def test_fast_language_model_from_pretrained_kwargs():
"""from_pretrained must accept the canonical kwargs the notebooks pass."""
unsloth = pytest.importorskip("unsloth")
required = {"model_name", "max_seq_length", "dtype", "load_in_4bit"}
ok, missing = _accepts(unsloth.FastLanguageModel.from_pretrained, required)
if not ok:
pytest.fail(
f"DRIFT DETECTED: FastLanguageModel.from_pretrained dropped "
f"kwargs {sorted(missing)}; 506 notebook call sites would "
f"crash with TypeError."
)
def test_fast_language_model_get_peft_model_kwargs():
unsloth = pytest.importorskip("unsloth")
required = {
"r",
"lora_alpha",
"lora_dropout",
"target_modules",
"bias",
"use_gradient_checkpointing",
"random_state",
}
ok, missing = _accepts(unsloth.FastLanguageModel.get_peft_model, required)
if not ok:
pytest.fail(
f"DRIFT DETECTED: FastLanguageModel.get_peft_model dropped "
f"kwargs {sorted(missing)}; 304 notebook call sites would crash."
)
def test_fast_language_model_for_inference_callable():
unsloth = pytest.importorskip("unsloth")
if not callable(getattr(unsloth.FastLanguageModel, "for_inference", None)):
pytest.fail(
"DRIFT DETECTED: FastLanguageModel.for_inference is missing; "
"370 inference-cell call sites would crash."
)
# ===========================================================================
# FastVisionModel: 183 + 176 + 99 + 60 call sites across vision notebooks.
# ===========================================================================
def test_fast_vision_model_class_and_methods():
unsloth = pytest.importorskip("unsloth")
if not hasattr(unsloth, "FastVisionModel"):
pytest.fail(
"DRIFT DETECTED: unsloth.FastVisionModel is missing; every "
"vision fine-tuning notebook fails at import."
)
cls = unsloth.FastVisionModel
missing = [
m
for m in ("from_pretrained", "get_peft_model", "for_inference", "for_training")
if not callable(getattr(cls, m, None))
]
if missing:
pytest.fail(f"DRIFT DETECTED: FastVisionModel is missing methods {missing}.")
def test_fast_vision_model_get_peft_model_vision_kwargs():
"""Vision-specific kwargs the notebooks pass on the vision LoRA path."""
unsloth = pytest.importorskip("unsloth")
required = {
"finetune_vision_layers",
"finetune_language_layers",
"finetune_attention_modules",
"finetune_mlp_modules",
}
ok, missing = _accepts(unsloth.FastVisionModel.get_peft_model, required)
if not ok:
pytest.fail(
f"DRIFT DETECTED: FastVisionModel.get_peft_model dropped "
f"vision kwargs {sorted(missing)}."
)
# ===========================================================================
# FastModel: the modern unified entry point. 103 + 67 call sites.
# ===========================================================================
def test_fast_model_class_and_methods():
unsloth = pytest.importorskip("unsloth")
if not hasattr(unsloth, "FastModel"):
pytest.fail(
"DRIFT DETECTED: unsloth.FastModel is missing; the modern "
"unified entry point used by 100+ notebooks would crash."
)
missing = [
m
for m in ("from_pretrained", "get_peft_model")
if not callable(getattr(unsloth.FastModel, m, None))
]
if missing:
pytest.fail(f"DRIFT DETECTED: FastModel is missing methods {missing}.")
def test_fast_model_from_pretrained_kwargs():
unsloth = pytest.importorskip("unsloth")
required = {"model_name", "max_seq_length", "dtype", "load_in_4bit"}
ok, missing = _accepts(unsloth.FastModel.from_pretrained, required)
if not ok:
pytest.fail(
f"DRIFT DETECTED: FastModel.from_pretrained dropped kwargs "
f"{sorted(missing)}; 103 notebook call sites would crash."
)
# ===========================================================================
# Bf16 helper alias (renamed once already; keep both accepted).
# ===========================================================================
def test_is_bf16_supported_or_alias_callable():
"""48 notebook import sites for is_bf16_supported plus 8 for the
legacy is_bfloat16_supported alias. Either must remain importable."""
unsloth = pytest.importorskip("unsloth")
has_new = callable(getattr(unsloth, "is_bf16_supported", None))
has_old = callable(getattr(unsloth, "is_bfloat16_supported", None))
if not (has_new or has_old):
pytest.fail(
"DRIFT DETECTED: neither unsloth.is_bf16_supported nor "
"unsloth.is_bfloat16_supported is callable; dtype probing "
"in 50+ notebooks fails."
)

128
tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py Normal file
View file

@ -0,0 +1,128 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team.
"""Pinned-symbol canary for unsloth-zoo save_pretrained_merged guards
(unslothai/unsloth-zoo#647 / unslothai/unsloth#5410). Skips until #647
lands, then becomes a hard gate. CPU-only static fetch."""
from __future__ import annotations
import re
import pytest
from tests.version_compat._fetch import fetch_text
ZOO_TAG = "main"
def _fetch_saving_utils() -> str:
src = fetch_text("unslothai/unsloth-zoo", ZOO_TAG, "unsloth_zoo/saving_utils.py")
if src is None:
pytest.skip("unsloth_zoo/saving_utils.py not fetchable")
return src
def _fetch_merge_tests() -> str:
src = fetch_text(
"unslothai/unsloth-zoo",
ZOO_TAG,
"tests/test_unsloth_zoo_lora_merge.py",
)
if src is None:
pytest.skip("tests/test_unsloth_zoo_lora_merge.py not fetchable")
return src
def _skip_until_pr_647_lands(src: str) -> None:
if not any(
m in src
for m in (
"_MOE_MERGE_STATE",
"_detect_moe_lora_layout",
"_resolve_num_experts_from_lora_stats",
)
):
pytest.skip(
"unslothai/unsloth-zoo#647 has not yet merged into main; "
"tests auto-promote to hard gates once it lands."
)
def test_zoo_saving_utils_has_moe_merge_state():
src = _fetch_saving_utils()
_skip_until_pr_647_lands(src)
for sym in (
"_MOE_MERGE_STATE",
"_reset_moe_merge_state",
"_record_moe_merge_fallback",
):
assert sym in src, f"{sym} missing from saving_utils.py (issue #5410 guard)."
# zoo#647 wraps the fallback guard's message onto a second line;
# allow the regex to span newlines via re.DOTALL.
assert re.search(
r"raise\s+RuntimeError\b.*?MoE", src, re.IGNORECASE | re.DOTALL
), "no `raise RuntimeError(...MoE...)`; post-loop guard weakened."
def test_zoo_saving_utils_has_layout_detector():
src = _fetch_saving_utils()
_skip_until_pr_647_lands(src)
assert (
"_detect_moe_lora_layout" in src
), "_detect_moe_lora_layout removed (issue #5410)."
assert (
'"swapped"' in src and '"standard"' in src
), "one of the layout labels removed."
def test_zoo_saving_utils_has_num_experts_resolver():
src = _fetch_saving_utils()
_skip_until_pr_647_lands(src)
assert "_resolve_num_experts_from_lora_stats" in src, "resolver removed (#5410)."
assert re.search(
r"for\s+_\s+in\s+range\s*\(\s*\d+\s*\)", src
), "resolver walk no longer bounded by `for _ in range(N):`."
def test_zoo_saving_utils_writes_generation_config():
src = _fetch_saving_utils()
_skip_until_pr_647_lands(src)
# zoo#647 binds the generation_config attr to a local var
# (`gen_cfg = getattr(model, "generation_config", ...); ...
# gen_cfg.save_pretrained(save_directory)`) so an exact
# `generation_config.save_pretrained(` substring no longer
# matches. Anchor on the conceptual operation: a `generation_config`
# mention plus a `.save_pretrained(` call nearby, which is what
# the canary actually cares about.
assert re.search(
r"generation_config[\s\S]{0,400}?\.save_pretrained\s*\(", src
), "generation_config.json no longer saved (#5410)."
def test_zoo_lora_merge_tests_have_standard_layout_coverage():
src = _fetch_merge_tests()
if "test_merge_moe_gate_expert_standard_layout" not in src:
pytest.skip("unslothai/unsloth-zoo#647 not yet merged; coverage appears later.")
for name in (
"test_merge_moe_gate_expert_standard_layout",
"test_merge_moe_up_expert_standard_layout",
"test_merge_moe_down_proj_expert_standard_layout",
"test_detect_moe_lora_layout_classifies_both_conventions",
"test_moe_merge_fallback_counter_records_bad_layout",
"test_resolve_num_experts_walks_base_layer_chain",
):
assert name in src, f"regression test `{name}` removed."
def test_unsloth_save_pretrained_merged_entry_point_exists():
import pathlib
save_py = pathlib.Path(__file__).resolve().parents[2] / "unsloth" / "save.py"
if not save_py.is_file():
pytest.skip(f"{save_py} not present")
text = save_py.read_text(encoding = "utf-8", errors = "replace")
assert "save_pretrained_merged" in text, "entry point removed from unsloth/save.py."
assert (
"merge_and_overwrite_lora" in text
), "no dispatch into unsloth_zoo merge; #647 bypassed."

37
unsloth/__init__.py
View file

@ -12,16 +12,33 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import os, platform, importlib.util
import os, importlib.util, platform
os.environ["UNSLOTH_IS_PRESENT"] = "1"
def _is_mlx_available():
# Transitional import barrier: while the paired unsloth-zoo MLX runtime
# rollout is in flight, keep non-Apple-Silicon imports from touching
# unsloth_zoo here. After both PRs are released together and
# unsloth_zoo.mlx is guaranteed to be import-safe on GPU hosts,
# this helper can collapse back to the centralized zoo runtime call below.
if (
os.environ.get("UNSLOTH_FORCE_GPU_PATH", "0") == "1"
or platform.system() != "Darwin"
or platform.machine() != "arm64"
or importlib.util.find_spec("mlx") is None
):
return False
try:
from unsloth_zoo.mlx import is_mlx_available
except ImportError:
return False
return is_mlx_available()
# Detect Apple Silicon + MLX before any torch/numpy imports
_IS_MLX = (
platform.system() == "Darwin"
and platform.machine() == "arm64"
and importlib.util.find_spec("mlx") is not None
)
_IS_MLX = _is_mlx_available()
if _IS_MLX:
try:
@ -31,18 +48,18 @@ if _IS_MLX:
"Unsloth: MLX support requires `unsloth-zoo` with MLX modules. "
"Reinstall with `pip install unsloth-zoo` or rerun install.sh."
) from _e
# The mlx_trainer / mlx_loader submodules ship with unsloth-zoo's MLX
# The mlx.trainer / mlx.loader submodules ship with unsloth-zoo's MLX
# support. An older installed unsloth-zoo (e.g. from PyPI before the
# MLX release lands) will satisfy `import unsloth_zoo` but be missing
# these submodules. Surface the same friendly install hint instead of
# a raw ImportError on the submodule path.
try:
from unsloth_zoo.mlx_trainer import MLXTrainer, MLXTrainingConfig
from unsloth_zoo.mlx_loader import FastMLXModel
from unsloth_zoo.mlx.trainer import MLXTrainer, MLXTrainingConfig
from unsloth_zoo.mlx.loader import FastMLXModel
except ImportError as _e:
raise ImportError(
"Unsloth: MLX support requires an unsloth-zoo build that includes "
"`unsloth_zoo.mlx_trainer` and `unsloth_zoo.mlx_loader`. Upgrade with "
"`unsloth_zoo.mlx.trainer` and `unsloth_zoo.mlx.loader`. Upgrade with "
"`pip install -U unsloth-zoo` or rerun install.sh."
) from _e

9
unsloth/_gpu_init.py
View file

@ -349,5 +349,10 @@ from unsloth_zoo.rl_environments import (
launch_openenv,
)
# Patch TRL trainers for backwards compatibility
_patch_trl_trainer()
# Patch TRL trainers for backwards compatibility.
# Skipped under UNSLOTH_ALLOW_CPU=1 (CPU-only CI) because rebinding
# trl.SFTTrainer.__init__ to a generic wrapper changes
# inspect.getsource(SFTTrainer.__init__) and corrupts downstream
# drift detectors that anchor on the pristine upstream source.
if os.environ.get("UNSLOTH_ALLOW_CPU", "0") != "1":
_patch_trl_trainer()

16
unsloth/chat_templates.py
View file

@ -30,11 +30,9 @@ __all__ = [
from transformers import StoppingCriteria, StoppingCriteriaList
from torch import LongTensor, FloatTensor
from transformers.models.llama.modeling_llama import logger
from .save import patch_saving_functions
import os
import shutil
from .tokenizer_utils import *
from .models._utils import patch_tokenizer
import re
from .ollama_template_mappers import OLLAMA_TEMPLATES
from unsloth_zoo.dataset_utils import (
@ -213,7 +211,7 @@ vicuna_ollama = _ollama_template("vicuna")
vicuna_eos_token = "eos_token"
CHAT_TEMPLATES["vicuna"] = (vicuna_template, vicuna_eos_token, False, vicuna_ollama,)
DEFAULT_SYSTEM_MESSAGE["vicuna"] = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions."
DEFAULT_SYSTEM_MESSAGE["vicuna"] = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user\\'s questions."
# =========================================== Vicuna Old
# https://github.com/lm-sys/FastChat/blob/main/docs/vicuna_weights_version.md#prompt-template
@ -1844,6 +1842,8 @@ def get_chat_template(
mapping = {"role" : "role", "content" : "content", "user" : "user", "assistant" : "assistant"},
map_eos_token = True,
system_message = None,
patch_saving = True,
use_zoo_tokenizer_patch = False,
):
assert(type(map_eos_token) is bool)
old_tokenizer = tokenizer
@ -2026,6 +2026,12 @@ def get_chat_template(
.replace("'user'", "'" + mapping["user"] + "'")\
.replace("'assistant'", "'" + mapping["assistant"] + "'")
if use_zoo_tokenizer_patch:
# Studio MLX avoids the model-utils tokenizer wrapper because that
# import path pulls in Torch/GPU-specific modules before MLX training.
from unsloth_zoo.tokenizer_utils import patch_tokenizer
else:
from .models._utils import patch_tokenizer
_, tokenizer = patch_tokenizer(model = None, tokenizer = tokenizer)
tokenizer.padding_side = old_padding_side
@ -2059,7 +2065,9 @@ def get_chat_template(
# stopping_criteria = create_stopping_criteria(tokenizer, stop_word)
# Patch saving functions
tokenizer = patch_saving_functions(tokenizer)
if patch_saving:
from .save import patch_saving_functions
tokenizer = patch_saving_functions(tokenizer)
# Add Ollama
tokenizer._ollama_modelfile = ollama_modelfile

30
unsloth/device_type.py
View file

@ -20,21 +20,47 @@ __all__ = [
"DEVICE_COUNT",
"ALLOW_PREQUANTIZED_MODELS",
"ALLOW_BITSANDBYTES",
"is_mlx_available",
]
import torch
import functools
import inspect
import os
from unsloth_zoo.utils import Version
def is_mlx_available():
try:
from unsloth_zoo.mlx import is_mlx_available as _is_mlx_available
except ImportError:
return False
return _is_mlx_available()
_IS_MLX = is_mlx_available()
if not _IS_MLX:
import torch
@functools.cache
def is_hip():
if _IS_MLX:
return False
return bool(getattr(getattr(torch, "version", None), "hip", None))
@functools.cache
def get_device_type():
# Test-only CPU fallback. Short-circuits the detection chain so the
# rest of the function -- and every DEVICE_TYPE == "cuda" branch in
# the codebase -- behaves identically to a real CUDA host. The env
# var is read exactly once per process because get_device_type is
# @functools.cache'd, so production hosts pay no runtime cost.
if os.environ.get("UNSLOTH_ALLOW_CPU", "0") == "1":
return "cuda"
if _IS_MLX:
return "mlx"
if hasattr(torch, "cuda") and torch.cuda.is_available():
if is_hip():
return "hip"
@ -64,6 +90,8 @@ DEVICE_TYPE: str = get_device_type()
DEVICE_TYPE_TORCH = DEVICE_TYPE
if DEVICE_TYPE_TORCH == "hip":
DEVICE_TYPE_TORCH = "cuda"
elif DEVICE_TYPE_TORCH == "mlx":
DEVICE_TYPE_TORCH = "mps"
@functools.cache

72
unsloth/import_fixes.py
View file

@ -1289,53 +1289,61 @@ def patch_torchcodec_audio_decoder():
def disable_torchcodec_if_broken():
"""Disable torchcodec in transformers if it cannot actually load.
"""Make broken torchcodec behave as if uninstalled (#5446).
transformers checks if torchcodec is installed via importlib.util.find_spec(),
but this returns True even when torchcodec cannot load its native libraries
(e.g., when FFmpeg is missing). This causes runtime errors when transformers
tries to use torchcodec for audio loading.
This function tests if torchcodec can actually load and if not, patches
transformers to think torchcodec is unavailable so it falls back to librosa.
Two shapes to cover:
* transformers < 5: a module-level ``_torchcodec_available`` flag
cached in ``transformers.utils.import_utils``; flip it to False.
* transformers >= 5: a public ``is_torchcodec_available()`` callable
wrapped with ``functools.lru_cache``; replace it with a stub that
returns False and clear the cache so subsequent callers see it.
transformers and datasets both detect torchcodec via find_spec, which
returns True even when the native libs cannot dlopen. We flip their
flags and seat a sys.modules sentinel so downstream imports fall through
their existing except ImportError handlers cleanly.
"""
try:
import importlib.util
if importlib.util.find_spec("torchcodec") is None:
return # torchcodec not installed, nothing to do
return # absent or already disabled
# Test if torchcodec can actually load
# RuntimeError on dlopen failure; OSError covers chained libavutil.so misses.
from torchcodec.decoders import AudioDecoder
except (ImportError, RuntimeError, OSError):
# torchcodec cannot load - disable it in transformers
# transformers: flip flag (<5) and/or rebind lru_cache'd func (>=5).
try:
import transformers.utils.import_utils as tf_import_utils
except ImportError:
return
# transformers < 5 path: module-level cached flag.
try:
tf_import_utils._torchcodec_available = False
except AttributeError:
pass
# transformers >= 5 path: public lru_cache'd function. Clear any
# cached True result then rebind to a stub that returns False.
is_avail = getattr(tf_import_utils, "is_torchcodec_available", None)
if is_avail is not None:
try:
is_avail.cache_clear()
tf_import_utils._torchcodec_available = False
except AttributeError:
pass
tf_import_utils.is_torchcodec_available = lambda: False
is_avail = getattr(tf_import_utils, "is_torchcodec_available", None)
if is_avail is not None:
try:
is_avail.cache_clear()
except AttributeError:
pass
tf_import_utils.is_torchcodec_available = lambda: False
except ImportError:
pass
# datasets >= 4.0: own flag gating audio/video/features/formatters.
try:
import datasets.config as datasets_config
if hasattr(datasets_config, "TORCHCODEC_AVAILABLE"):
datasets_config.TORCHCODEC_AVAILABLE = False
except ImportError:
pass
# Drop half-loaded entries and seat the absence sentinel. After this,
# import torchcodec raises ModuleNotFoundError and find_spec returns None.
for _stale in [
n
for n in list(sys.modules)
if n == "torchcodec"
or n.startswith("torchcodec.")
or n == "datasets.features._torchcodec"
]:
sys.modules.pop(_stale, None)
sys.modules["torchcodec"] = None
def disable_broken_wandb():

4
unsloth/kernels/moe/autotune_cache.py
View file

@ -71,7 +71,7 @@ def load_cached_config(cache_key: str) -> Optional[Dict[str, Any]]:
return None
try:
with open(cache_file, "r") as f:
with open(cache_file, "r", encoding = "utf-8") as f:
cached_data = json.load(f)
# Verify cache is still valid (same device, etc.)
@ -118,7 +118,7 @@ def save_cached_config(
}
try:
with open(cache_file, "w") as f:
with open(cache_file, "w", encoding = "utf-8") as f:
json.dump(cache_data, f, indent = 2)
logger.info(f"Saved MoE kernel config cache: {cache_key}")
except Exception as e:

2
unsloth/kernels/moe/benchmark/utils.py
View file

@ -183,7 +183,7 @@ def save_autotune_results(autotune_cache, mode, ref_time, fused_time, results_di
filename = "_".join(key)
save_path = f"{save_dir}/{filename}.json"
print(f"Saving autotune results to {save_path}")
with open(save_path, "w") as f:
with open(save_path, "w", encoding = "utf-8") as f:
result = {
**config.all_kwargs(),
"ref_time": ref_time,

9
unsloth/kernels/utils.py
View file

@ -160,6 +160,10 @@ else:
# INTEL GPU Specific Logic
if DEVICE_TYPE == "xpu":
_gpu_getCurrentRawStream = torch._C._xpu_getCurrentRawStream
elif DEVICE_TYPE == "mlx":
def _gpu_getCurrentRawStream(_index = 0):
return 0
# NVIDIA GPU Default Logic
elif hasattr(torch._C, "_cuda_getCurrentRawStream"):
_gpu_getCurrentRawStream = torch._C._cuda_getCurrentRawStream
@ -206,6 +210,11 @@ if DEVICE_TYPE == "xpu":
XPU_STREAMS = ()
WEIGHT_BUFFERS = []
ABSMAX_BUFFERS = []
elif DEVICE_TYPE == "mlx":
CUDA_STREAMS = ()
XPU_STREAMS = ()
WEIGHT_BUFFERS = []
ABSMAX_BUFFERS = []
else:
# NVIDIA GPU Default Logic
if DEVICE_COUNT > 0:

Some files were not shown because too many files have changed in this diff Show more